[PECOBLR-1384] Fix telemetry gaps: CloudFetch S3 timing, chunk totals, flush ordering

- Thread per-S3-file download timing from cloudFetchDownloadTask through
  cloudIPCStreamIterator via onFileDownloaded callback; aggregate into
  initial/slowest/sum chunk timing tags matching JDBC per-chunk HTTP GET model
- Zero FetchResults RPC latency for CloudFetch pages (URLs-only response) so
  Thrift round-trip is not misreported as data download time
- Add statementID param to RecordOperation so CLOSE_STATEMENT entries carry
  sql_statement_id in telemetry tables
- Fix EXECUTE_STATEMENT/CLOSE_STATEMENT loss on shutdown: add inFlight WaitGroup
  to metricsAggregator so close() waits for in-flight worker exports before
  calling cancel(), preventing metrics picked up by workers from being dropped
- Derive total_chunks_present from final chunkCount in closeCallback for
  paginated CloudFetch and inline results where server does not report total
- Remove omitempty from OperationLatencyMs so CLOSE_STATEMENT reports 0ms
  (instant server-side close) rather than null

Co-authored-by: samikshya-chand_data

Author: samikshya-db

connection.go

@@ -60,7 +60,7 @@ func (c *conn) Close() error {
 
 	// Record DELETE_SESSION regardless of error (matches JDBC), then flush and release
 	if c.telemetry != nil {
-		c.telemetry.RecordOperation(ctx, c.id, telemetry.OperationTypeDeleteSession, time.Since(closeStart).Milliseconds(), err)
+		c.telemetry.RecordOperation(ctx, c.id, "", telemetry.OperationTypeDeleteSession, time.Since(closeStart).Milliseconds(), err)
 		_ = c.telemetry.Close(ctx)
 		telemetry.ReleaseForConnection(c.cfg.Host)
 	}
@@ -164,7 +164,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 				OperationHandle: exStmtResp.OperationHandle,
 			})
 			if c.telemetry != nil {
-				c.telemetry.RecordOperation(ctx, c.id, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), err1)
+				c.telemetry.RecordOperation(ctx, c.id, statementID, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), err1)
 			}
 			if err1 != nil {
 				log.Err(err1).Msg("databricks: failed to close operation after executing statement")
@@ -273,30 +273,59 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		}
 	}
 
+	// cloudFetchCallback is invoked per S3 file download for CloudFetch result sets.
+	// It aggregates individual file download times into the same initial/slowest/sum vars
+	// used for inline chunk timing, matching JDBC's per-chunk HTTP GET timing model.
+	// For inline (non-CloudFetch) result sets this is never called.
+	var cloudFetchCallback func(downloadMs int64)
+	if c.telemetry != nil {
+		cloudFetchCallback = func(downloadMs int64) {
+			if downloadMs <= 0 {
+				return
+			}
+			if !chunkTimingInitialSet {
+				chunkTimingInitialMs = downloadMs
+				chunkTimingInitialSet = true
+			}
+			if downloadMs > chunkTimingSlowestMs {
+				chunkTimingSlowestMs = downloadMs
+			}
+			chunkTimingSumMs += downloadMs
+			c.telemetry.AddTag(ctx, "chunk_initial_latency_ms", chunkTimingInitialMs)
+			c.telemetry.AddTag(ctx, "chunk_slowest_latency_ms", chunkTimingSlowestMs)
+			c.telemetry.AddTag(ctx, "chunk_sum_latency_ms", chunkTimingSumMs)
+		}
+	}
+
 	// closeCallback is invoked from rows.Close() after all rows have been consumed.
 	// At that point chunk timing is fully accumulated in ctx tags, so we finalize
 	// EXECUTE_STATEMENT here rather than at QueryContext return time.
-	var closeCallback func(latencyMs int64, err error)
+	var closeCallback func(latencyMs int64, chunkCount int, err error)
 	if c.telemetry != nil && statementID != "" {
 		interceptor := c.telemetry
 		connID := c.id
 		stmtID := statementID
-		closeCallback = func(latencyMs int64, closeErr error) {
+		closeCallback = func(latencyMs int64, chunkCount int, closeErr error) {
+			// If the server never reported total chunks (paginated CloudFetch / inline),
+			// derive it from the final chunk count now that all iteration is complete.
+			if chunkTotalPresent == 0 && chunkCount > 0 {
+				interceptor.AddTag(ctx, "chunk_total_present", chunkCount)
+			}
 			// Emit EXECUTE_STATEMENT with complete chunk data now that iteration is done.
 			interceptor.AfterExecute(ctx, nil)
 			interceptor.CompleteStatement(ctx, stmtID, false)
 			// Emit CLOSE_STATEMENT as a separate operation event.
-			interceptor.RecordOperation(ctx, connID, telemetry.OperationTypeCloseStatement, latencyMs, closeErr)
+			interceptor.RecordOperation(ctx, connID, stmtID, telemetry.OperationTypeCloseStatement, latencyMs, closeErr)
 		}
 	} else if c.telemetry != nil {
 		interceptor := c.telemetry
 		connID := c.id
-		closeCallback = func(latencyMs int64, closeErr error) {
-			interceptor.RecordOperation(ctx, connID, telemetry.OperationTypeCloseStatement, latencyMs, closeErr)
+		closeCallback = func(latencyMs int64, _ int, 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)
+	rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate, closeCallback, cloudFetchCallback)
 	return rows, err
 
 }
@@ -733,7 +762,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, nil)
+		row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate, nil, nil)
 		if err != nil {
 			return dbsqlerrint.NewDriverError(ctx, "error reading row.", err)
 		}

connector.go

@@ -94,7 +94,7 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	)
 	if conn.telemetry != nil {
 		log.Debug().Msg("telemetry initialized for connection")
-		conn.telemetry.RecordOperation(ctx, conn.id, telemetry.OperationTypeCreateSession, sessionLatencyMs, nil)
+		conn.telemetry.RecordOperation(ctx, conn.id, "", telemetry.OperationTypeCreateSession, sessionLatencyMs, nil)
 	}
 
 	log.Info().Msgf("connect: host=%s port=%d httpPath=%s serverProtocolVersion=0x%X", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath, session.ServerProtocolVersion)

internal/rows/arrowbased/arrowIPCStreamIterator.go

@@ -138,7 +138,7 @@ func (ri *arrowIPCStreamIterator) fetchNextData() error {
 func (ri *arrowIPCStreamIterator) newIPCStreamIterator(fr *cli_service.TFetchResultsResp) (IPCStreamIterator, error) {
 	rowSet := fr.Results
 	if len(rowSet.ResultLinks) > 0 {
-		return NewCloudIPCStreamIterator(ri.ctx, rowSet.ResultLinks, rowSet.StartRowOffset, &ri.cfg)
+		return NewCloudIPCStreamIterator(ri.ctx, rowSet.ResultLinks, rowSet.StartRowOffset, &ri.cfg, nil)
 	} else {
 		return NewLocalIPCStreamIterator(ri.ctx, rowSet.ArrowBatches, rowSet.StartRowOffset, ri.arrowSchemaBytes, &ri.cfg)
 	}

internal/rows/arrowbased/arrowRecordIterator.go

@@ -169,7 +169,7 @@ func (ri *arrowRecordIterator) getBatchIterator() error {
 func (ri *arrowRecordIterator) newBatchIterator(fr *cli_service.TFetchResultsResp) (BatchIterator, error) {
 	rowSet := fr.Results
 	if len(rowSet.ResultLinks) > 0 {
-		return NewCloudBatchIterator(ri.ctx, rowSet.ResultLinks, rowSet.StartRowOffset, &ri.cfg)
+		return NewCloudBatchIterator(ri.ctx, rowSet.ResultLinks, rowSet.StartRowOffset, &ri.cfg, nil)
 	} else {
 		return NewLocalBatchIterator(ri.ctx, rowSet.ArrowBatches, rowSet.StartRowOffset, ri.arrowSchemaBytes, &ri.cfg)
 	}

internal/rows/arrowbased/arrowRows.go

@@ -80,8 +80,10 @@ type arrowRowScanner struct {
 // Make sure arrowRowScanner fulfills the RowScanner interface
 var _ rowscanner.RowScanner = (*arrowRowScanner)(nil)
 
-// NewArrowRowScanner returns an instance of RowScanner which handles arrow format results
-func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp, rowSet *cli_service.TRowSet, cfg *config.Config, logger *dbsqllog.DBSQLLogger, ctx context.Context) (rowscanner.RowScanner, dbsqlerr.DBError) {
+// NewArrowRowScanner returns an instance of RowScanner which handles arrow format results.
+// onCloudFetchDownload is an optional callback invoked for each CloudFetch S3 file download
+// with the download duration in milliseconds. Pass nil for non-telemetry paths.
+func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp, rowSet *cli_service.TRowSet, cfg *config.Config, logger *dbsqllog.DBSQLLogger, ctx context.Context, onCloudFetchDownload func(downloadMs int64)) (rowscanner.RowScanner, dbsqlerr.DBError) {
 
 	// we take a passed in logger, rather than just using the global from dbsqllog, so that the containing rows
 	// instance can pass in a logger with context such as correlation ID and operation ID
@@ -119,7 +121,7 @@ func NewArrowRowScanner(resultSetMetadata *cli_service.TGetResultSetMetadataResp
 		for _, resultLink := range rowSet.ResultLinks {
 			logger.Debug().Msgf("- start row offset: %d, row count: %d", resultLink.StartRowOffset, resultLink.RowCount)
 		}
-		bi, err2 = NewCloudBatchIterator(context.Background(), rowSet.ResultLinks, rowSet.StartRowOffset, cfg)
+		bi, err2 = NewCloudBatchIterator(context.Background(), rowSet.ResultLinks, rowSet.StartRowOffset, cfg, onCloudFetchDownload)
 	} else {
 		bi, err2 = NewLocalBatchIterator(context.Background(), rowSet.ArrowBatches, rowSet.StartRowOffset, schemaBytes, cfg)
 	}