@@ -52,16 +52,19 @@ func (c *conn) Close() error {
5252 log := logger .WithContext (c .id , "" , "" )
5353 ctx := driverctx .NewContextWithConnId (context .Background (), c .id )
5454
55- // Close telemetry and release resources
55+ // Time CloseSession so we can record DELETE_SESSION before flushing telemetry
56+ closeStart := time .Now ()
57+ _ , err := c .client .CloseSession (ctx , & cli_service.TCloseSessionReq {
58+ SessionHandle : c .session .SessionHandle ,
59+ })
60+
61+ // Record DELETE_SESSION regardless of error (matches JDBC), then flush and release
5662 if c .telemetry != nil {
63+ c .telemetry .RecordOperation (ctx , c .id , telemetry .OperationTypeDeleteSession , time .Since (closeStart ).Milliseconds (), err )
5764 _ = c .telemetry .Close (ctx )
5865 telemetry .ReleaseForConnection (c .cfg .Host )
5966 }
6067
61- _ , err := c .client .CloseSession (ctx , & cli_service.TCloseSessionReq {
62- SessionHandle : c .session .SessionHandle ,
63- })
64-
6568 if err != nil {
6669 log .Err (err ).Msg ("databricks: failed to close connection" )
6770 return dbsqlerrint .NewBadConnectionError (err )
@@ -93,7 +96,7 @@ func (c *conn) Ping(ctx context.Context) error {
9396 log .Err (err ).Msg ("databricks: failed to ping" )
9497 return dbsqlerrint .NewBadConnectionError (err )
9598 }
96- defer rows .Close ()
99+ defer rows .Close () //nolint:errcheck
97100
98101 log .Debug ().Msg ("databricks: ping successful" )
99102 return nil
@@ -155,9 +158,13 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
155158 alreadyClosed := exStmtResp .DirectResults != nil && exStmtResp .DirectResults .CloseOperation != nil
156159 newCtx := driverctx .NewContextWithCorrelationId (driverctx .NewContextWithConnId (context .Background (), c .id ), corrId )
157160 if ! alreadyClosed && (opStatusResp == nil || opStatusResp .GetOperationState () != cli_service .TOperationState_CLOSED_STATE ) {
161+ closeOpStart := time .Now ()
158162 _ , err1 := c .client .CloseOperation (newCtx , & cli_service.TCloseOperationReq {
159163 OperationHandle : exStmtResp .OperationHandle ,
160164 })
165+ if c .telemetry != nil {
166+ c .telemetry .RecordOperation (ctx , c .id , telemetry .OperationTypeCloseStatement , time .Since (closeOpStart ).Milliseconds (), err1 )
167+ }
161168 if err1 != nil {
162169 log .Err (err1 ).Msg ("databricks: failed to close operation after executing statement" )
163170 closeOpErr = err1 // Capture for telemetry
@@ -216,7 +223,15 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
216223 return nil , dbsqlerrint .NewExecutionError (ctx , dbsqlerr .ErrQueryExecution , err , opStatusResp )
217224 }
218225
219- rows , err := rows .NewRows (ctx , exStmtResp .OperationHandle , c .client , c .cfg , exStmtResp .DirectResults )
226+ // Telemetry callback for tracking row fetching metrics
227+ telemetryUpdate := func (chunkCount int , bytesDownloaded int64 ) {
228+ if c .telemetry != nil {
229+ c .telemetry .AddTag (ctx , "chunk_count" , chunkCount )
230+ c .telemetry .AddTag (ctx , "bytes_downloaded" , bytesDownloaded )
231+ }
232+ }
233+
234+ rows , err := rows .NewRows (ctx , exStmtResp .OperationHandle , c .client , c .cfg , exStmtResp .DirectResults , telemetryUpdate )
220235 return rows , err
221236
222237}
@@ -381,7 +396,14 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
381396 }
382397 }
383398
399+ executeStart := time .Now ()
384400 resp , err := c .client .ExecuteStatement (ctx , & req )
401+ // Record the Thrift call latency as a separate operation metric.
402+ // This is distinct from the statement-level metric (BeforeExecuteWithTime), which
403+ // measures end-to-end latency including polling and row fetching.
404+ if c .telemetry != nil {
405+ c .telemetry .RecordOperation (ctx , c .id , telemetry .OperationTypeExecuteStatement , time .Since (executeStart ).Milliseconds (), err )
406+ }
385407 var log * logger.DBSQLLogger
386408 log , ctx = client .LoggerAndContext (ctx , resp )
387409
@@ -514,11 +536,11 @@ func (c *conn) handleStagingPut(ctx context.Context, presignedUrl string, header
514536 }
515537 client := & http.Client {}
516538
517- dat , err := os .Open (localFile )
539+ dat , err := os .Open (localFile ) //nolint:gosec // localFile is provided by the application, not user input
518540 if err != nil {
519541 return dbsqlerrint .NewDriverError (ctx , "error reading local file" , err )
520542 }
521- defer dat .Close ()
543+ defer dat .Close () //nolint:errcheck
522544
523545 info , err := dat .Stat ()
524546 if err != nil {
@@ -535,7 +557,7 @@ func (c *conn) handleStagingPut(ctx context.Context, presignedUrl string, header
535557 if err != nil {
536558 return dbsqlerrint .NewDriverError (ctx , "error sending http request" , err )
537559 }
538- defer res .Body .Close ()
560+ defer res .Body .Close () //nolint:errcheck
539561 content , err := io .ReadAll (res .Body )
540562
541563 if err != nil || ! Succeeded (res ) {
@@ -559,7 +581,7 @@ func (c *conn) handleStagingGet(ctx context.Context, presignedUrl string, header
559581 if err != nil {
560582 return dbsqlerrint .NewDriverError (ctx , "error sending http request" , err )
561583 }
562- defer res .Body .Close ()
584+ defer res .Body .Close () //nolint:errcheck
563585 content , err := io .ReadAll (res .Body )
564586
565587 if err != nil || ! Succeeded (res ) {
@@ -583,7 +605,7 @@ func (c *conn) handleStagingRemove(ctx context.Context, presignedUrl string, hea
583605 if err != nil {
584606 return dbsqlerrint .NewDriverError (ctx , "error sending http request" , err )
585607 }
586- defer res .Body .Close ()
608+ defer res .Body .Close () //nolint:errcheck
587609 content , err := io .ReadAll (res .Body )
588610
589611 if err != nil || ! Succeeded (res ) {
@@ -646,11 +668,18 @@ func (c *conn) execStagingOperation(
646668 }
647669
648670 if len (driverctx .StagingPathsFromContext (ctx )) != 0 {
649- row , err = rows .NewRows (ctx , exStmtResp .OperationHandle , c .client , c .cfg , exStmtResp .DirectResults )
671+ // Telemetry callback for staging operation row fetching
672+ telemetryUpdate := func (chunkCount int , bytesDownloaded int64 ) {
673+ if c .telemetry != nil {
674+ c .telemetry .AddTag (ctx , "chunk_count" , chunkCount )
675+ c .telemetry .AddTag (ctx , "bytes_downloaded" , bytesDownloaded )
676+ }
677+ }
678+ row , err = rows .NewRows (ctx , exStmtResp .OperationHandle , c .client , c .cfg , exStmtResp .DirectResults , telemetryUpdate )
650679 if err != nil {
651680 return dbsqlerrint .NewDriverError (ctx , "error reading row." , err )
652681 }
653- defer row .Close ()
682+ defer row .Close () //nolint:errcheck
654683
655684 } else {
656685 return dbsqlerrint .NewDriverError (ctx , "staging ctx must be provided." , nil )
@@ -663,7 +692,7 @@ func (c *conn) execStagingOperation(
663692 if err != nil {
664693 return dbsqlerrint .NewDriverError (ctx , "error fetching staging operation results" , err )
665694 }
666- var stringValues [] string = make ([]string , 4 )
695+ stringValues : = make ([]string , 4 )
667696 for i , val := range sqlRow { // this will either be 3 (remove op) or 4 (put/get) elements
668697 if s , ok := val .(string ); ok {
669698 stringValues [i ] = s
0 commit comments