fix context cancellation

conn_batch.go

@@ -179,7 +179,9 @@ func (b *batch) Column(idx int) driver.BatchColumn {
 
 func (b *batch) Send() (err error) {
 	stopCW := contextWatchdog(b.ctx, func() {
-		_ = b.conn.close()
+		// close TCP connection on context cancel. There is no other way simple way to interrupt underlying operations.
+		// as verified in the test, this is safe to do and cleanups resources later on
+		_ = b.conn.conn.Close()
 	})
 
 	defer func() {
@@ -197,8 +199,11 @@ func (b *batch) Send() (err error) {
 	}
 	if b.block.Rows() != 0 {
 		if err = b.conn.sendData(b.block, ""); err != nil {
-			// todo: return context.DeadlineExceeded if deadline is exceeded
-			// use b.err ?
+			// there might be an error caused by context cancellation
+			// in this case we should return context error instead of net.OpError
+			if ctxErr := b.ctx.Err(); ctxErr != nil {
+				return ctxErr
+			}
 
 			return err
 		}

tests/batch_test.go

@@ -26,16 +26,27 @@ import (
 )
 
 func TestBatchContextCancellation(t *testing.T) {
-	conn, err := GetNativeConnection(nil, nil, &clickhouse.Compression{
-		Method: clickhouse.CompressionLZ4,
-	})
+	te, err := GetTestEnvironment(testSet)
+	require.NoError(t, err)
+	opts := ClientOptionsFromEnv(te, clickhouse.Settings{})
+	opts.MaxOpenConns = 1
+	conn, err := GetConnectionWithOptions(&opts)
 	require.NoError(t, err)
 
-	ctx, cancel := context.WithTimeout(context.Background(), time.Microsecond)
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
 	defer cancel()
 
-	b, err := conn.PrepareBatch(ctx, "INSERT INTO test_batch_ctx_cancellation (x)")
+	require.NoError(t, conn.Exec(context.Background(), "create table if not exists test_batch_cancellation (x String) engine=Memory"))
+	defer conn.Exec(context.Background(), "drop table if exists test_batch_cancellation")
+
+	b, err := conn.PrepareBatch(ctx, "insert into test_batch_cancellation")
 	require.NoError(t, err)
+	for i := 0; i < 1_000_000; i++ {
+		require.NoError(t, b.Append("value"))
+	}
 
 	require.Equal(t, context.DeadlineExceeded, b.Send())
+
+	// assert if connection is properly released after context cancellation
+	require.NoError(t, conn.Exec(context.Background(), "SELECT 1"))
 }