Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Batch rows count API #1063

Merged
merged 2 commits into from
Aug 10, 2023
Merged

Batch rows count API #1063

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0b8766e
Add Rows() func to batch
EpicStep Aug 7, 2023
7a91874
Add tests
EpicStep Aug 7, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions conn_batch.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -221,6 +221,10 @@ func (b *batch) Flush() error {
return nil
}

func (b *batch) Rows() int {
return b.block.Rows()
}

type batchColumn struct {
err error
batch driver.Batch
Expand Down
4 changes: 4 additions & 0 deletions conn_http_batch.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,4 +229,8 @@ func (b *httpBatch) Send() (err error) {
return err
}

func (b *httpBatch) Rows() int {
return b.block.Rows()
}

var _ driver.Batch = (*httpBatch)(nil)
1 change: 1 addition & 0 deletions lib/driver/driver.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,7 @@ type (
Flush() error
Send() error
IsSent() bool
Rows() int
}
BatchColumn interface {
Append(any) error
Expand Down
7 changes: 7 additions & 0 deletions tests/array_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,6 +49,7 @@ func TestSimpleArray(t *testing.T) {
)
for i := 0; i < 10; i++ {
require.NoError(t, batch.Append(col1Data))
require.Equal(t, 1, batch.Rows())
batch.Flush()
}
require.NoError(t, batch.Send())
Expand Down Expand Up @@ -92,8 +93,10 @@ func TestCustomArray(t *testing.T) {
)
for i := 0; i < 10; i++ {
require.NoError(t, batch.Append(col1Data, col2Data))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Flush())
}
require.Equal(t, 0, batch.Rows())
require.NoError(t, batch.Send())
rows, err := conn.Query(ctx, "SELECT * FROM test_array")
require.NoError(t, err)
Expand Down Expand Up @@ -133,6 +136,7 @@ func TestInterfaceArray(t *testing.T) {
for i := 0; i < 10; i++ {
require.NoError(t, batch.Append(col1Data))
}
require.Equal(t, 10, batch.Rows())
require.Nil(t, batch.Send())
rows, err := conn.Query(ctx, "SELECT * FROM test_array")
require.NoError(t, err)
Expand Down Expand Up @@ -200,8 +204,10 @@ func TestArray(t *testing.T) {
)
for i := 0; i < 10; i++ {
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data))
require.Equal(t, 1, batch.Rows())
batch.Flush()
}
require.Equal(t, 0, batch.Rows())
require.NoError(t, batch.Send())
rows, err := conn.Query(ctx, "SELECT * FROM test_array")
require.NoError(t, err)
Expand Down Expand Up @@ -289,6 +295,7 @@ func TestColumnarArray(t *testing.T) {
require.NoError(t, batch.Column(1).Append(col2DataColArr))
require.NoError(t, batch.Column(2).Append(col3DataColArr))
require.NoError(t, batch.Column(3).Append(col4DataColArr))
require.Equal(t, 10, batch.Rows())
require.NoError(t, batch.Send())
rows, err := conn.Query(ctx, "SELECT * FROM test_array")
require.NoError(t, err)
Expand Down
9 changes: 9 additions & 0 deletions tests/base_types_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -71,6 +71,7 @@ func TestUInt8(t *testing.T) {
require.NoError(t, err)
require.NoError(t, batch.Append(uint8(1), data, &data, []uint8{data}, []*uint8{&data, nil, &data}, customUint8(data)))
require.NoError(t, batch.Append(uint8(2), data, nil, []uint8{data}, []*uint8{nil, nil, &data}, customUint8(data)))
require.Equal(t, 2, batch.Rows())
require.NoError(t, batch.Send())
var (
result1 result
Expand Down Expand Up @@ -151,6 +152,7 @@ func TestColumnarUInt8(t *testing.T) {
return
}
}
require.Equal(t, 1000, batch.Rows())
require.NoError(t, batch.Send())
var result struct {
Col1 uint8
Expand Down Expand Up @@ -209,6 +211,7 @@ func TestNullableInt(t *testing.T) {
col5Data := sql.NullInt16{Int16: 3, Valid: true}
col6Data := sql.NullInt16{Int16: 0, Valid: false}
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 sql.NullInt64
Expand Down Expand Up @@ -250,9 +253,15 @@ func TestIntFlush(t *testing.T) {
vals[i] = uint8(i)
require.NoError(t, batch.Append(vals[i]))
if i%100 == 0 {
if i == 0 {
require.Equal(t, 1, batch.Rows())
} else {
require.Equal(t, 100, batch.Rows())
}
require.NoError(t, batch.Flush())
}
}
require.Equal(t, 99, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM int_flush")
require.NoError(t, err)
Expand Down
6 changes: 6 additions & 0 deletions tests/bigint_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@ func TestSimpleBigInt(t *testing.T) {
col1Data, ok := new(big.Int).SetString("170141183460469231731687303715884105727", 10)
require.True(t, ok)
require.NoError(t, batch.Append(col1Data))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 big.Int
Expand Down Expand Up @@ -111,6 +112,7 @@ func TestBigInt(t *testing.T) {
}
)
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data, col7Data))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 big.Int
Expand Down Expand Up @@ -178,6 +180,7 @@ func TestNullableBigInt(t *testing.T) {
}
)
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 *big.Int
Expand Down Expand Up @@ -246,6 +249,7 @@ func TestBigIntUIntOverflow(t *testing.T) {
}
)
require.NoError(t, batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 big.Int
Expand Down Expand Up @@ -287,8 +291,10 @@ func TestBigIntFlush(t *testing.T) {
bigUint128Val.SetString(RandIntString(20), 10)
vals[i] = bigUint128Val
batch.Append(vals[i])
require.Equal(t, 1, batch.Rows())
batch.Flush()
}
require.Equal(t, 0, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM big_int_flush")
require.NoError(t, err)
Expand Down
4 changes: 4 additions & 0 deletions tests/bool_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ func TestBool(t *testing.T) {
Bool: false,
Valid: false,
}))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 bool
Expand Down Expand Up @@ -140,6 +141,7 @@ func TestColumnarBool(t *testing.T) {
require.NoError(t, batch.Column(3).Append(col3))
require.NoError(t, batch.Column(4).Append(col4))
require.NoError(t, batch.Column(5).Append(col5))
require.Equal(t, 1000, batch.Rows())
require.NoError(t, batch.Send())
var (
id uint64
Expand Down Expand Up @@ -182,8 +184,10 @@ func TestBoolFlush(t *testing.T) {
for i := 0; i < 1000; i++ {
vals[i] = r.Intn(2) != 0
require.NoError(t, batch.Append(vals[i]))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Flush())
}
require.Equal(t, 0, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM bool_flush")
require.NoError(t, err)
Expand Down
5 changes: 5 additions & 0 deletions tests/columnar_batch_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ func TestColumnarInterface(t *testing.T) {
require.NoError(t, batch.Column(3).Append(col4Data))
require.NoError(t, batch.Column(4).Append(col5Data))
require.NoError(t, batch.Column(5).Append(col6Data))
require.Equal(t, 150, batch.Rows())
require.NoError(t, batch.Send())
var count uint64
require.NoError(t, conn.QueryRow(ctx, "SELECT COUNT() FROM test_column_interface").Scan(&count))
Expand Down Expand Up @@ -151,6 +152,7 @@ func TestNullableColumnarInterface(t *testing.T) {
require.NoError(t, batch.Column(1).Append(col2Data))
require.NoError(t, batch.Column(2).Append(col3Data))
require.NoError(t, batch.Column(3).Append(col4Data))
require.Equal(t, 150, batch.Rows())
require.NoError(t, batch.Send())
var count uint64
require.NoError(t, conn.QueryRow(ctx, "SELECT COUNT() FROM test_column_interface").Scan(&count))
Expand Down Expand Up @@ -210,6 +212,7 @@ func TestNullableColumnarInterface(t *testing.T) {
require.NoError(t, batch.Column(1).Append(col2Data))
require.NoError(t, batch.Column(2).Append(col3Data))
require.NoError(t, batch.Column(3).Append(col4Data))
require.Equal(t, 150, batch.Rows())
require.NoError(t, batch.Send())
var count uint64
require.NoError(t, conn.QueryRow(ctx, "SELECT COUNT() FROM test_column_interface").Scan(&count))
Expand Down Expand Up @@ -283,6 +286,7 @@ func TestColumnarAppendRowInterface(t *testing.T) {
require.NoError(t, batch.Column(4).AppendRow(sql.NullTime{Time: currentTime, Valid: true}))
require.NoError(t, batch.Column(5).AppendRow(sql.NullInt64{Int64: int64(i), Valid: true}))
}
require.Equal(t, 150, batch.Rows())
require.NoError(t, batch.Send())
var count uint64
require.NoError(t, conn.QueryRow(ctx, "SELECT COUNT() FROM test_column_interface").Scan(&count))
Expand Down Expand Up @@ -352,6 +356,7 @@ func TestNullableAppendRowColumnarInterface(t *testing.T) {
require.NoError(t, batch.Column(3).AppendRow(&decimalVal))
}
}
require.Equal(t, 150, batch.Rows())
require.NoError(t, batch.Send())
var count uint64
require.NoError(t, conn.QueryRow(ctx, "SELECT COUNT() FROM test_column_interface").Scan(&count))
Expand Down
8 changes: 8 additions & 0 deletions tests/date32_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -81,6 +81,7 @@ func TestDate32(t *testing.T) {
require.NoError(t, batch.Append(uint8(1), date1, &date2, []time.Time{date2}, []*time.Time{&date2, nil, &date1}, dateStr1, dateStrNil, []string{dateStr1, dateStr2, dateStr3}, []*string{dateStrNil, &dateStr1, dateStrNil}))
require.NoError(t, batch.Append(uint8(2), date2, nil, []time.Time{date1}, []*time.Time{nil, nil, &date2}, &testStr{Col1: dateStr1}, nil, []string{dateStr1, dateStr2, dateStr3}, []*string{nil, &dateStr1, dateStrNil}))
require.NoError(t, batch.Append(uint8(3), date3, nil, []time.Time{date3}, []*time.Time{nil, nil, &date3}, &testStr{Col1: dateStr1}, &dateStr1, []string{dateStr1, dateStr2, dateStr3}, []*string{nil, nil, dateStrNil}))
require.Equal(t, 3, batch.Rows())
require.NoError(t, batch.Send())
var (
result1 result
Expand Down Expand Up @@ -158,6 +159,7 @@ func TestNullableDate32(t *testing.T) {
require.NoError(t, err)
dateStr := "2283-11-11"
require.NoError(t, batch.Append(date, &date, dateStr, &dateStr))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 *time.Time
Expand All @@ -176,6 +178,7 @@ func TestNullableDate32(t *testing.T) {
date, err = time.Parse("2006-01-02 15:04:05", "1925-01-01 00:00:00")
require.NoError(t, err)
require.NoError(t, batch.Append(date, nil, &date, nil))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
col2 = nil
col4 = nil
Expand Down Expand Up @@ -245,6 +248,7 @@ func TestColumnarDate32(t *testing.T) {
require.NoError(t, batch.Column(3).Append(col3Data))
require.NoError(t, batch.Column(4).Append(col4Data))
}
require.Equal(t, 1000, batch.Rows())
require.NoError(t, batch.Send())
var result struct {
Col1 time.Time
Expand Down Expand Up @@ -286,8 +290,10 @@ func TestDate32Flush(t *testing.T) {
for i := 0; i < 1000; i++ {
vals[i] = now.Add(time.Duration(i) * time.Hour)
batch.Append(vals[i])
require.Equal(t, 1, batch.Rows())
batch.Flush()
}
require.Equal(t, 0, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM date_32_flush")
require.NoError(t, err)
Expand Down Expand Up @@ -321,6 +327,7 @@ func TestDate32TZ(t *testing.T) {
"2022-07-20 +08:00",
))
require.NoError(t, err)
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col15, col16 time.Time
Expand Down Expand Up @@ -355,6 +362,7 @@ func TestCustomDateTime32(t *testing.T) {
require.NoError(t, err)
now := time.Now().UTC().Truncate(time.Hour)
require.NoError(t, batch.Append(now))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
row := conn.QueryRow(ctx, "SELECT * FROM date32_custom")
var col1 CustomDateTime
Expand Down
6 changes: 6 additions & 0 deletions tests/datetime_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,7 @@ func TestDateTime(t *testing.T) {
&testStr{Col1: dateTimeStr},
iDateTime,
))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 time.Time
Expand Down Expand Up @@ -141,6 +142,7 @@ func TestNullableDateTime(t *testing.T) {
require.NoError(t, err)
datetime := time.Now().Truncate(time.Second)
require.NoError(t, batch.Append(datetime, datetime, datetime, datetime, datetime, datetime, datetime, datetime, datetime, datetime))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 time.Time
Expand Down Expand Up @@ -181,6 +183,7 @@ func TestNullableDateTime(t *testing.T) {
datetimeNilStr *string = nil
)
require.NoError(t, batch.Append(datetime, nil, datetime, nil, datetime, nil, datetimeStr, nil, datetimeStr, datetimeNilStr))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 time.Time
Expand Down Expand Up @@ -294,6 +297,7 @@ func TestColumnarDateTime(t *testing.T) {
require.NoError(t, batch.Column(6).Append(col6Data))
require.NoError(t, batch.Column(7).Append(col7Data))
}
require.Equal(t, 1000, batch.Rows())
require.NoError(t, batch.Send())
var result struct {
Col1 time.Time
Expand Down Expand Up @@ -338,8 +342,10 @@ func TestDateTimeFlush(t *testing.T) {
for i := 0; i < 1000; i++ {
vals[i] = now.Add(time.Duration(i) * time.Hour).Truncate(time.Second)
batch.Append(vals[i])
require.Equal(t, 1, batch.Rows())
batch.Flush()
}
require.Equal(t, 0, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM datetime_flush")
require.NoError(t, err)
Expand Down
10 changes: 9 additions & 1 deletion tests/decimal_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ func TestDecimal(t *testing.T) {
decimal.New(135, 7),
decimal.New(256, 8),
))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 decimal.Decimal
Expand Down Expand Up @@ -109,6 +110,7 @@ func TestNegativeDecimal(t *testing.T) {
decimal.RequireFromString("-0.01171"),
decimal.RequireFromString("-3.0111"),
decimal.RequireFromString("-21111122.0111111111111111111171")))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 decimal.Decimal
Expand Down Expand Up @@ -149,6 +151,7 @@ func TestNullableDecimal(t *testing.T) {
batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_decimal")
require.NoError(t, err)
require.NoError(t, batch.Append(decimal.New(25, 0), decimal.New(30, 0), decimal.New(35, 0)))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
var (
col1 *decimal.Decimal
Expand All @@ -163,6 +166,7 @@ func TestNullableDecimal(t *testing.T) {
batch, err = conn.PrepareBatch(ctx, "INSERT INTO test_decimal")
require.NoError(t, err)
require.NoError(t, batch.Append(decimal.New(25, 0), nil, decimal.New(35, 0)))
require.Equal(t, 1, batch.Rows())
require.NoError(t, batch.Send())
{
var (
Expand Down Expand Up @@ -198,8 +202,10 @@ func TestDecimalFlush(t *testing.T) {
for i := 0; i < 1000; i++ {
vals[i] = decimal.RequireFromString(fmt.Sprintf("1.%s", RandIntString(5)))
batch.Append(vals[i])
require.Equal(t, 1, batch.Rows())
batch.Flush()
}
require.Equal(t, 0, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM decimal_flush")
require.NoError(t, err)
Expand Down Expand Up @@ -253,9 +259,11 @@ func TestRoundDecimals(t *testing.T) {
decimal.NewFromFloat(601), // this will make decimal 601*e^0
decimal.NewFromFloat(601.21), // check that normal case is working
}
for _, c := range checks {
for i, c := range checks {
batch.Append(c)
require.Equal(t, i+1, batch.Rows())
}
require.Equal(t, 3, batch.Rows())
batch.Send()
rows, err := conn.Query(ctx, "SELECT * FROM decimal_flush ORDER BY Col1 asc")
require.NoError(tt, err)
Expand Down
1 change: 1 addition & 0 deletions tests/empty_query_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -49,5 +49,6 @@ func TestEmptyQuery(t *testing.T) {
defer cancel()
batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_empty_query")
require.NoError(t, err)
require.Equal(t, 0, batch.Rows())
assert.NoError(t, batch.Send())
}
Loading