fix batch.Close and tests

timeplus-io · Oct 26, 2023 · a280d4f · a280d4f
1 parent b0294ab
commit a280d4f
Show file tree

Hide file tree

Showing 9 changed files with 20 additions and 16 deletions.
diff --git a/clickhouse_std.go b/clickhouse_std.go
@@ -163,7 +163,12 @@ func (std *stdDriver) PrepareContext(ctx context.Context, query string) (driver.
 	if err != nil {
 		return nil, err
 	}
-	std.commit = batch.Send
+	std.commit = func() error {
+		if err := batch.Send(); err != nil {
+			return err
+		}
+		return batch.Close()
+	}
 	return &stdBatch{
 		batch: batch,
 	}, nil

diff --git a/conn_batch.go b/conn_batch.go
@@ -133,6 +133,9 @@ func (b *batch) Close() (err error) {
 	if err = b.conn.encoder.Flush(); err != nil {
 		return err
 	}
+	if err = b.conn.process(b.ctx, b.onProcess); err != nil {
+		return err
+	}
 	b.sent = true
 	b.release(err)
 	return err
@@ -156,10 +159,6 @@ func (b *batch) Send() (err error) {
 	if err = b.block.Clear(); err != nil {
 		return err
 	}
-	// disable process recently
-	// if err = b.conn.process(b.ctx, b.onProcess); err != nil {
-	// 	 return err
-	// }
 	return nil
 }
 

diff --git a/tests/abort_test.go b/tests/abort_test.go
@@ -61,7 +61,7 @@ func TestAbort(t *testing.T) {
 				}
 			}
 			if batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_abort (* except _tp_time)"); assert.NoError(t, err) {
-				if assert.NoError(t, batch.Append(uint8(1))) && assert.NoError(t, batch.Send()) {
+				if assert.NoError(t, batch.Append(uint8(1))) && assert.NoError(t, batch.Send()) && assert.NoError(t, batch.Close()) {
 					var col1 uint8
 					if err := conn.QueryRow(ctx, "SELECT (* except _tp_time) FROM test_abort WHERE _tp_time > earliest_ts() LIMIT 1").Scan(&col1); assert.NoError(t, err) {
 						assert.Equal(t, uint8(1), col1)

diff --git a/tests/array_test.go b/tests/array_test.go
@@ -90,7 +90,7 @@ func TestArray(t *testing.T) {
 						return
 					}
 				}
-				if assert.NoError(t, batch.Send()) {
+				if assert.NoError(t, batch.Send()) && assert.NoError(t, batch.Close()) {
 					if rows, err := conn.Query(ctx, "SELECT (* except _tp_time) FROM test_array WHERE _tp_time > earliest_ts() LIMIT 1"); assert.NoError(t, err) {
 						for rows.Next() {
 							var (
@@ -193,7 +193,7 @@ func TestColumnarArray(t *testing.T) {
 				if err := batch.Column(2).Append(col3DataColArr); !assert.NoError(t, err) {
 					return
 				}
-				if assert.NoError(t, batch.Send()) {
+				if assert.NoError(t, batch.Send()) && assert.NoError(t, batch.Close()) {
 					if rows, err := conn.Query(ctx, "SELECT (* except _tp_time) FROM test_array WHERE _tp_time > earliest_ts() LIMIT 1"); assert.NoError(t, err) {
 						for rows.Next() {
 							var (

diff --git a/tests/base_types_test.go b/tests/base_types_test.go
@@ -73,7 +73,7 @@ func TestUInt8(t *testing.T) {
 				if err := batch.Append(uint8(2), data, nil, []uint8{data}, []*uint8{nil, nil, &data}); !assert.NoError(t, err) {
 					return
 				}
-				if err := batch.Send(); assert.NoError(t, err) {
+				if err := batch.Send(); assert.NoError(t, err) && assert.NoError(t, batch.Close()) {
 					var (
 						result1 result
 						result2 result
@@ -170,7 +170,7 @@ func TestColumnarUInt8(t *testing.T) {
 						return
 					}
 				}
-				if assert.NoError(t, batch.Send()) {
+				if assert.NoError(t, batch.Send()) && assert.NoError(t, batch.Close()) {
 					var result struct {
 						Col1 uint8
 						Col2 *uint8

diff --git a/tests/bigint_test.go b/tests/bigint_test.go
@@ -83,7 +83,7 @@ func TestBigInt(t *testing.T) {
 					}
 				)
 				if err := batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data); assert.NoError(t, err) {
-					if err := batch.Send(); assert.NoError(t, err) {
+					if err := batch.Send(); assert.NoError(t, err) && assert.NoError(t, batch.Close()) {
 						var (
 							col1 big.Int
 							col2 []*big.Int
@@ -160,7 +160,7 @@ func TestNullableBigInt(t *testing.T) {
 					}
 				)
 				if err := batch.Append(col1Data, col2Data, col3Data, col4Data, col5Data, col6Data); assert.NoError(t, err) {
-					if err := batch.Send(); assert.NoError(t, err) {
+					if err := batch.Send(); assert.NoError(t, err) && assert.NoError(t, batch.Close()) {
 						var (
 							col1 *big.Int
 							col2 []*big.Int

diff --git a/tests/bool_test.go b/tests/bool_test.go
@@ -62,7 +62,7 @@ func TestBool(t *testing.T) {
 			if batch, err := conn.PrepareBatch(ctx, "INSERT INTO test_bool (* except _tp_time)"); assert.NoError(t, err) {
 				var val bool
 				if err := batch.Append(true, false, []bool{true, false, true}, nil, []*bool{&val, nil, &val}); assert.NoError(t, err) {
-					if err := batch.Send(); assert.NoError(t, err) {
+					if err := batch.Send(); assert.NoError(t, err) && assert.NoError(t, batch.Close()) {
 						var (
 							col1 bool
 							col2 bool
@@ -157,7 +157,7 @@ func TestColumnarBool(t *testing.T) {
 					if err := batch.Column(5).Append(col5); !assert.NoError(t, err) {
 						return
 					}
-					if err := batch.Send(); assert.NoError(t, err) {
+					if err := batch.Send(); assert.NoError(t, err) && assert.NoError(t, batch.Close()) {
 						var (
 							id   uint64
 							col1 bool

diff --git a/tests/columnar_batch_test.go b/tests/columnar_batch_test.go
@@ -76,7 +76,7 @@ func TestColumnarInterface(t *testing.T) {
 				if err := batch.Column(2).Append(col3Data); !assert.NoError(t, err) {
 					return
 				}
-				if assert.NoError(t, batch.Send()) {
+				if assert.NoError(t, batch.Send()) && assert.NoError(t, batch.Close()) {
 					var count uint64
 					if err := conn.QueryRow(ctx, "SELECT count() FROM test_column_interface WHERE _tp_time > earliest_ts() LIMIT 1").Scan(&count); assert.NoError(t, err) {
 						if assert.Equal(t, uint64(150), count) {

diff --git a/tests/struct_map_test.go b/tests/struct_map_test.go
@@ -91,7 +91,7 @@ func TestAppendStruct(t *testing.T) {
 					}
 				}
 
-				if assert.NoError(t, batch.Send()) {
+				if assert.NoError(t, batch.Send()) && assert.NoError(t, batch.Close()) {
 					for i := 0; i < 100; i++ {
 						var result data
 						if err := conn.QueryRow(ctx, "SELECT (* except _tp_time) FROM test_append_struct WHERE _tp_time > earliest_ts() AND HCol1 = $1 LIMIT 1", i).ScanStruct(&result); assert.NoError(t, err) {