ClickHouse#100 PingContext is now supported

Author: kshvakov

clickhouse.go

@@ -301,5 +301,5 @@ func (ch *clickhouse) watchCancel(ctx context.Context) func() {
 			}
 		}
 	}
-	return nil
+	return func() {}
 }

clickhouse_ping.go

@@ -8,14 +8,16 @@ import (
 )
 
 func (ch *clickhouse) Ping(ctx context.Context) error {
-	return ch.ping()
+	return ch.ping(ctx)
 }
 
-func (ch *clickhouse) ping() error {
+func (ch *clickhouse) ping(ctx context.Context) error {
 	if ch.conn.closed {
 		return driver.ErrBadConn
 	}
 	ch.logf("-> ping")
+	finish := ch.watchCancel(ctx)
+	defer finish()
 	if err := ch.encoder.Uvarint(protocol.ClientPing); err != nil {
 		return err
 	}

clickhouse_test.go

@@ -1020,6 +1020,36 @@ func Test_Context_Timeout(t *testing.T) {
 	}
 }
 
+func Test_Ping_Context_Timeout(t *testing.T) {
+	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?debug=true"); assert.NoError(t, err) && assert.NoError(t, connect.Ping()) {
+		{
+			ctx, cancel := context.WithTimeout(context.Background(), time.Nanosecond)
+			defer cancel()
+			if err := connect.PingContext(ctx); assert.Error(t, err) {
+				assert.Equal(t, context.DeadlineExceeded, err)
+			}
+		}
+		{
+			ctx, cancel := context.WithTimeout(context.Background(), time.Millisecond*20)
+			defer cancel()
+			if row := connect.QueryRowContext(ctx, "SELECT 1, sleep(2)"); assert.NotNil(t, row) {
+				var a, b int
+				assert.Equal(t, driver.ErrBadConn, row.Scan(&a, &b))
+			}
+		}
+		{
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			if row := connect.QueryRowContext(ctx, "SELECT 1, sleep(0.1)"); assert.NotNil(t, row) {
+				var value, value2 int
+				if assert.NoError(t, row.Scan(&value, &value2)) {
+					assert.Equal(t, int(1), value)
+				}
+			}
+		}
+	}
+}
+
 func Test_Timeout(t *testing.T) {
 	if connect, err := sql.Open("clickhouse", "tcp://127.0.0.1:9000?debug=true&read_timeout=0.2"); assert.NoError(t, err) && assert.NoError(t, connect.Ping()) {
 		{

rows.go

@@ -145,9 +145,7 @@ func (rows *rows) Close() error {
 	rows.columns = nil
 	for range rows.stream {
 	}
-	if rows.finish != nil {
-		rows.finish()
-	}
+	rows.finish()
 	return nil
 }
 

stmt.go

@@ -77,16 +77,12 @@ func (stmt *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (d
 func (stmt *stmt) queryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
 	finish := stmt.ch.watchCancel(ctx)
 	if err := stmt.ch.sendQuery(stmt.bind(args)); err != nil {
-		if finish != nil {
-			finish()
-		}
+		finish()
 		return nil, err
 	}
 	meta, err := stmt.ch.readMeta()
 	if err != nil {
-		if finish != nil {
-			finish()
-		}
+		finish()
 		return nil, err
 	}
 	rows := rows{