Fix telemetry: correct EXECUTE_STATEMENT latency, add sql_statement_id, add CLOSE_STATEMENT

- Remove RecordOperation(EXECUTE_STATEMENT) from executeStatement(): it measured
  only the Thrift RPC submission time (~5s), not end-to-end query time. The
  correct end-to-end metric already flows through BeforeExecuteWithTime →
  AfterExecute → CompleteStatement.
- Add AddTag(operation_type=EXECUTE_STATEMENT) after BeforeExecuteWithTime in
  QueryContext and ExecContext so the aggregated statement metric carries the
  right operation_type, sql_statement_id, and accurate latency in one event.
- Add closeCallback func(latencyMs int64, err error) to rows.NewRows() and
  rows.Close() to emit CLOSE_STATEMENT telemetry when rows.Close() is called.

Co-authored-by: Isaac

Author: samikshya-db

connection.go

@@ -139,6 +139,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 		statementID = client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID)
 		// Use BeforeExecuteWithTime to set the correct start time (before execution)
 		ctx = c.telemetry.BeforeExecuteWithTime(ctx, c.id, statementID, executeStart)
+		c.telemetry.AddTag(ctx, "operation_type", telemetry.OperationTypeExecuteStatement)
 		defer func() {
 			finalErr := err
 			if stagingErr != nil {
@@ -212,6 +213,7 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		statementID = client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID)
 		// Use BeforeExecuteWithTime to set the correct start time (before execution)
 		ctx = c.telemetry.BeforeExecuteWithTime(ctx, c.id, statementID, executeStart)
+		c.telemetry.AddTag(ctx, "operation_type", telemetry.OperationTypeExecuteStatement)
 		defer func() {
 			c.telemetry.AfterExecute(ctx, err)
 			c.telemetry.CompleteStatement(ctx, statementID, err != nil)
@@ -231,7 +233,17 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		}
 	}
 
-	rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate)
+	// Telemetry callback for CLOSE_STATEMENT — fired from rows.Close()
+	var closeCallback func(latencyMs int64, err error)
+	if c.telemetry != nil {
+		interceptor := c.telemetry
+		connID := c.id
+		closeCallback = func(latencyMs int64, closeErr error) {
+			interceptor.RecordOperation(ctx, connID, telemetry.OperationTypeCloseStatement, latencyMs, closeErr)
+		}
+	}
+
+	rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate, closeCallback)
 	return rows, err
 
 }
@@ -396,14 +408,7 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
 		}
 	}
 
-	executeStart := time.Now()
 	resp, err := c.client.ExecuteStatement(ctx, &req)
-	// Record the Thrift call latency as a separate operation metric.
-	// This is distinct from the statement-level metric (BeforeExecuteWithTime), which
-	// measures end-to-end latency including polling and row fetching.
-	if c.telemetry != nil {
-		c.telemetry.RecordOperation(ctx, c.id, telemetry.OperationTypeExecuteStatement, time.Since(executeStart).Milliseconds(), err)
-	}
 	var log *logger.DBSQLLogger
 	log, ctx = client.LoggerAndContext(ctx, resp)
 
@@ -675,7 +680,7 @@ func (c *conn) execStagingOperation(
 				c.telemetry.AddTag(ctx, "bytes_downloaded", bytesDownloaded)
 			}
 		}
-		row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate)
+		row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate, nil)
 		if err != nil {
 			return dbsqlerrint.NewDriverError(ctx, "error reading row.", err)
 		}

internal/rows/rows.go

@@ -60,6 +60,7 @@ type rows struct {
 
 	// Telemetry tracking
 	telemetryUpdate func(chunkCount int, bytesDownloaded int64)
+	closeCallback   func(latencyMs int64, err error)
 	chunkCount      int
 	bytesDownloaded int64
 }
@@ -78,6 +79,7 @@ func NewRows(
 	config *config.Config,
 	directResults *cli_service.TSparkDirectResults,
 	telemetryUpdate func(chunkCount int, bytesDownloaded int64),
+	closeCallback func(latencyMs int64, err error),
 ) (driver.Rows, dbsqlerr.DBError) {
 
 	connId := driverctx.ConnIdFromContext(ctx)
@@ -118,6 +120,7 @@ func NewRows(
 		logger_:         logger,
 		ctx:             ctx,
 		telemetryUpdate: telemetryUpdate,
+		closeCallback:   closeCallback,
 		chunkCount:      0,
 		bytesDownloaded: 0,
 	}
@@ -210,7 +213,11 @@ func (r *rows) Close() error {
 
 	if r.ResultPageIterator != nil {
 		r.logger().Debug().Msgf("databricks: closing Rows operation")
+		closeStart := time.Now()
 		err := r.ResultPageIterator.Close()
+		if r.closeCallback != nil {
+			r.closeCallback(time.Since(closeStart).Milliseconds(), err)
+		}
 		if err != nil {
 			r.logger().Err(err).Msg(errRowsCloseFailed)
 			return dbsqlerr_int.NewRequestError(r.ctx, errRowsCloseFailed, err)

internal/rows/rows_test.go

@@ -421,7 +421,7 @@ func TestColumnsWithDirectResults(t *testing.T) {
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "corrId")
 
-	d, err := NewRows(ctx, nil, client, nil, nil, nil)
+	d, err := NewRows(ctx, nil, client, nil, nil, nil, nil)
 	assert.Nil(t, err)
 
 	rowSet := d.(*rows)
@@ -720,7 +720,7 @@ func TestRowsCloseOptimization(t *testing.T) {
 	ctx := driverctx.NewContextWithConnId(context.Background(), "connId")
 	ctx = driverctx.NewContextWithCorrelationId(ctx, "corrId")
 	opHandle := &cli_service.TOperationHandle{OperationId: &cli_service.THandleIdentifier{GUID: []byte{'f', 'o'}}}
-	rowSet, _ := NewRows(ctx, opHandle, client, nil, nil, nil)
+	rowSet, _ := NewRows(ctx, opHandle, client, nil, nil, nil, nil)
 
 	// rowSet has no direct results calling Close should result in call to client to close operation
 	err := rowSet.Close()
@@ -733,7 +733,7 @@ func TestRowsCloseOptimization(t *testing.T) {
 		ResultSet:         &cli_service.TFetchResultsResp{Results: &cli_service.TRowSet{Columns: []*cli_service.TColumn{}}},
 	}
 	closeCount = 0
-	rowSet, _ = NewRows(ctx, opHandle, client, nil, directResults, nil)
+	rowSet, _ = NewRows(ctx, opHandle, client, nil, directResults, nil, nil)
 	err = rowSet.Close()
 	assert.Nil(t, err, "rows.Close should not throw an error")
 	assert.Equal(t, 1, closeCount)
@@ -746,7 +746,7 @@ func TestRowsCloseOptimization(t *testing.T) {
 		ResultSetMetadata: &cli_service.TGetResultSetMetadataResp{Schema: &cli_service.TTableSchema{}},
 		ResultSet:         &cli_service.TFetchResultsResp{Results: &cli_service.TRowSet{Columns: []*cli_service.TColumn{}}},
 	}
-	rowSet, _ = NewRows(ctx, opHandle, client, nil, directResults, nil)
+	rowSet, _ = NewRows(ctx, opHandle, client, nil, directResults, nil, nil)
 	err = rowSet.Close()
 	assert.Nil(t, err, "rows.Close should not throw an error")
 	assert.Equal(t, 0, closeCount)
@@ -816,7 +816,7 @@ func TestGetArrowBatches(t *testing.T) {
 
 		client := getSimpleClient([]cli_service.TFetchResultsResp{fetchResp1, fetchResp2})
 		cfg := config.WithDefaults()
-		rows, err := NewRows(ctx, nil, client, cfg, executeStatementResp.DirectResults, nil)
+		rows, err := NewRows(ctx, nil, client, cfg, executeStatementResp.DirectResults, nil, nil)
 		assert.Nil(t, err)
 
 		rows2, ok := rows.(dbsqlrows.Rows)
@@ -889,7 +889,7 @@ func TestGetArrowBatches(t *testing.T) {
 
 		client := getSimpleClient([]cli_service.TFetchResultsResp{fetchResp1, fetchResp2, fetchResp3})
 		cfg := config.WithDefaults()
-		rows, err := NewRows(ctx, nil, client, cfg, nil, nil)
+		rows, err := NewRows(ctx, nil, client, cfg, nil, nil, nil)
 		assert.Nil(t, err)
 
 		rows2, ok := rows.(dbsqlrows.Rows)
@@ -950,7 +950,7 @@ func TestGetArrowBatches(t *testing.T) {
 
 		client := getSimpleClient([]cli_service.TFetchResultsResp{fetchResp1})
 		cfg := config.WithDefaults()
-		rows, err := NewRows(ctx, nil, client, cfg, nil, nil)
+		rows, err := NewRows(ctx, nil, client, cfg, nil, nil, nil)
 		assert.Nil(t, err)
 
 		rows2, ok := rows.(dbsqlrows.Rows)
@@ -977,7 +977,7 @@ func TestGetArrowBatches(t *testing.T) {
 
 		client := getSimpleClient([]cli_service.TFetchResultsResp{})
 		cfg := config.WithDefaults()
-		rows, err := NewRows(ctx, nil, client, cfg, executeStatementResp.DirectResults, nil)
+		rows, err := NewRows(ctx, nil, client, cfg, executeStatementResp.DirectResults, nil, nil)
 		assert.Nil(t, err)
 
 		rows2, ok := rows.(dbsqlrows.Rows)
@@ -1557,7 +1557,7 @@ func TestFetchResultPage_PropagatesGetNextPageError(t *testing.T) {
 
 	executeStatementResp := cli_service.TExecuteStatementResp{}
 	cfg := config.WithDefaults()
-	rows, _ := NewRows(ctx, nil, client, cfg, executeStatementResp.DirectResults, nil)
+	rows, _ := NewRows(ctx, nil, client, cfg, executeStatementResp.DirectResults, nil, nil)
 	// Call Next and ensure it propagates the error from getNextPage
 	actualErr := rows.Next(nil)