Skip to content

Commit e017890

Browse files
samikshya-dbclaude
samikshya-db
and
claude
committed
Implement Phases 8-10: Testing, Launch Prep & Documentation
This commit completes the telemetry implementation with validation testing, launch preparation, and final fixes. Key Changes: 1. Phase 8: Validation Testing - Added operation latency tracking for all 3 operation types - Fixed timing issue where BeforeExecute was called after execution - Added BeforeExecuteWithTime to support correct latency measurement - Verified operation_latency_ms is populated for all events 2. Phase 9: Launch Preparation - Set EnableTelemetry default to true (respects server feature flags) - Removed all debug fmt.Printf statements - Cleaned up unused imports - Verified telemetry can be disabled with enableTelemetry=false 3. Telemetry Event Structure - CREATE_SESSION: Tracks session open latency - EXECUTE_STATEMENT: Tracks query execution with full metrics - DELETE_SESSION: Tracks session close latency - All events include system_configuration and operation_detail 4. Implementation Details - System info collection (OS, runtime, locale, process) - Operation types (CREATE_SESSION, EXECUTE_STATEMENT, DELETE_SESSION) - Metric aggregation with immediate flush for operations - Statement execution hooks for chunk/byte tracking - Feature flag checking with caching Testing: - Manual e2e validation passed with 3 events sent successfully - Default behavior enables telemetry with server feature flag check - Explicit opt-out (enableTelemetry=false) disables telemetry Co-Authored-By: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent 130d30d commit e017890

22 files changed

Lines changed: 604 additions & 165 deletions

‎connection.go‎

Lines changed: 57 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -52,15 +52,18 @@ func (c *conn) Close() error {
5252
ctx := driverctx.NewContextWithConnId(context.Background(), c.id)
5353

5454
// Close telemetry and release resources
55+
closeStart := time.Now()
56+
_, err := c.client.CloseSession(ctx, &cli_service.TCloseSessionReq{
57+
SessionHandle: c.session.SessionHandle,
58+
})
59+
closeLatencyMs := time.Since(closeStart).Milliseconds()
60+
5561
if c.telemetry != nil {
62+
c.telemetry.RecordOperation(ctx, c.id, telemetry.OperationTypeDeleteSession, closeLatencyMs)
5663
_ = c.telemetry.Close(ctx)
5764
telemetry.ReleaseForConnection(c.cfg.Host)
5865
}
5966

60-
_, err := c.client.CloseSession(ctx, &cli_service.TCloseSessionReq{
61-
SessionHandle: c.session.SessionHandle,
62-
})
63-
6467
if err != nil {
6568
log.Err(err).Msg("databricks: failed to close connection")
6669
return dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrCloseConnection, err)
@@ -122,15 +125,16 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
122125

123126
corrId := driverctx.CorrelationIdFromContext(ctx)
124127

125-
exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args)
128+
var pollCount int
129+
exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args, &pollCount)
126130
log, ctx = client.LoggerAndContext(ctx, exStmtResp)
127131
stagingErr := c.execStagingOperation(exStmtResp, ctx)
128132

129133
// Telemetry: track statement execution
130134
var statementID string
131135
if c.telemetry != nil && exStmtResp != nil && exStmtResp.OperationHandle != nil && exStmtResp.OperationHandle.OperationId != nil {
132136
statementID = client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID)
133-
ctx = c.telemetry.BeforeExecute(ctx, statementID)
137+
ctx = c.telemetry.BeforeExecute(ctx, c.id, statementID)
134138
defer func() {
135139
finalErr := err
136140
if stagingErr != nil {
@@ -139,6 +143,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
139143
c.telemetry.AfterExecute(ctx, finalErr)
140144
c.telemetry.CompleteStatement(ctx, statementID, finalErr != nil)
141145
}()
146+
c.telemetry.AddTag(ctx, "poll_count", pollCount)
142147
}
143148

144149
if exStmtResp != nil && exStmtResp.OperationHandle != nil {
@@ -180,21 +185,30 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
180185
log, _ := client.LoggerAndContext(ctx, nil)
181186
msg, start := log.Track("QueryContext")
182187

183-
// first we try to get the results synchronously.
184-
// at any point in time that the context is done we must cancel and return
185-
exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args)
188+
// Capture execution start time for telemetry before running the query
189+
executeStart := time.Now()
190+
var pollCount int
191+
exStmtResp, opStatusResp, pollCount, err := c.runQueryWithTelemetry(ctx, query, args, &pollCount)
186192
log, ctx = client.LoggerAndContext(ctx, exStmtResp)
187193
defer log.Duration(msg, start)
188194

189-
// Telemetry: track statement execution
190195
var statementID string
191196
if c.telemetry != nil && exStmtResp != nil && exStmtResp.OperationHandle != nil && exStmtResp.OperationHandle.OperationId != nil {
192197
statementID = client.SprintGuid(exStmtResp.OperationHandle.OperationId.GUID)
193-
ctx = c.telemetry.BeforeExecute(ctx, statementID)
198+
// Use BeforeExecuteWithTime to set the correct start time (before execution)
199+
ctx = c.telemetry.BeforeExecuteWithTime(ctx, c.id, statementID, executeStart)
194200
defer func() {
195201
c.telemetry.AfterExecute(ctx, err)
196202
c.telemetry.CompleteStatement(ctx, statementID, err != nil)
197203
}()
204+
205+
c.telemetry.AddTag(ctx, "poll_count", pollCount)
206+
c.telemetry.AddTag(ctx, "operation_type", telemetry.OperationTypeExecuteStatement)
207+
208+
if exStmtResp.DirectResults != nil && exStmtResp.DirectResults.ResultSetMetadata != nil {
209+
resultFormat := exStmtResp.DirectResults.ResultSetMetadata.GetResultFormat()
210+
c.telemetry.AddTag(ctx, "result.format", resultFormat.String())
211+
}
198212
}
199213

200214
if err != nil {
@@ -203,13 +217,31 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
203217
}
204218

205219
corrId := driverctx.CorrelationIdFromContext(ctx)
206-
rows, err := rows.NewRows(c.id, corrId, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
220+
221+
var telemetryUpdate func(int, int64)
222+
if c.telemetry != nil {
223+
telemetryUpdate = func(chunkCount int, bytesDownloaded int64) {
224+
c.telemetry.AddTag(ctx, "chunk_count", chunkCount)
225+
c.telemetry.AddTag(ctx, "bytes_downloaded", bytesDownloaded)
226+
}
227+
}
228+
229+
rows, err := rows.NewRows(c.id, corrId, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, ctx, telemetryUpdate)
207230

208231
return rows, err
209232

210233
}
211234

212-
func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedValue) (*cli_service.TExecuteStatementResp, *cli_service.TGetOperationStatusResp, error) {
235+
func (c *conn) runQueryWithTelemetry(ctx context.Context, query string, args []driver.NamedValue, pollCount *int) (*cli_service.TExecuteStatementResp, *cli_service.TGetOperationStatusResp, int, error) {
236+
exStmtResp, opStatusResp, err := c.runQuery(ctx, query, args, pollCount)
237+
count := 0
238+
if pollCount != nil {
239+
count = *pollCount
240+
}
241+
return exStmtResp, opStatusResp, count, err
242+
}
243+
244+
func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedValue, pollCount *int) (*cli_service.TExecuteStatementResp, *cli_service.TGetOperationStatusResp, error) {
213245
// first we try to get the results synchronously.
214246
// at any point in time that the context is done we must cancel and return
215247
exStmtResp, err := c.executeStatement(ctx, query, args)
@@ -241,7 +273,7 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
241273
case cli_service.TOperationState_INITIALIZED_STATE,
242274
cli_service.TOperationState_PENDING_STATE,
243275
cli_service.TOperationState_RUNNING_STATE:
244-
statusResp, err := c.pollOperation(ctx, opHandle)
276+
statusResp, err := c.pollOperationWithCount(ctx, opHandle, pollCount)
245277
if err != nil {
246278
return exStmtResp, statusResp, err
247279
}
@@ -269,7 +301,7 @@ func (c *conn) runQuery(ctx context.Context, query string, args []driver.NamedVa
269301
}
270302

271303
} else {
272-
statusResp, err := c.pollOperation(ctx, opHandle)
304+
statusResp, err := c.pollOperationWithCount(ctx, opHandle, pollCount)
273305
if err != nil {
274306
return exStmtResp, statusResp, err
275307
}
@@ -396,7 +428,7 @@ func (c *conn) executeStatement(ctx context.Context, query string, args []driver
396428
return resp, err
397429
}
398430

399-
func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperationHandle) (*cli_service.TGetOperationStatusResp, error) {
431+
func (c *conn) pollOperationWithCount(ctx context.Context, opHandle *cli_service.TOperationHandle, pollCount *int) (*cli_service.TGetOperationStatusResp, error) {
400432
corrId := driverctx.CorrelationIdFromContext(ctx)
401433
log := logger.WithContext(c.id, corrId, client.SprintGuid(opHandle.OperationId.GUID))
402434
var statusResp *cli_service.TGetOperationStatusResp
@@ -413,6 +445,10 @@ func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperati
413445
OperationHandle: opHandle,
414446
})
415447

448+
if pollCount != nil {
449+
*pollCount++
450+
}
451+
416452
if statusResp != nil && statusResp.OperationState != nil {
417453
log.Debug().Msgf("databricks: status %s", statusResp.GetOperationState().String())
418454
}
@@ -455,6 +491,10 @@ func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperati
455491
return statusResp, nil
456492
}
457493

494+
func (c *conn) pollOperation(ctx context.Context, opHandle *cli_service.TOperationHandle) (*cli_service.TGetOperationStatusResp, error) {
495+
return c.pollOperationWithCount(ctx, opHandle, nil)
496+
}
497+
458498
func (c *conn) CheckNamedValue(nv *driver.NamedValue) error {
459499
var err error
460500
if parameter, ok := nv.Value.(Parameter); ok {
@@ -623,7 +663,7 @@ func (c *conn) execStagingOperation(
623663
}
624664

625665
if len(driverctx.StagingPathsFromContext(ctx)) != 0 {
626-
row, err = rows.NewRows(c.id, corrId, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
666+
row, err = rows.NewRows(c.id, corrId, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, nil, nil)
627667
if err != nil {
628668
return dbsqlerrint.NewDriverError(ctx, "error reading row.", err)
629669
}

‎connection_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -833,7 +833,7 @@ func TestConn_runQuery(t *testing.T) {
833833
client: testClient,
834834
cfg: config.WithDefaults(),
835835
}
836-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
836+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
837837
assert.Error(t, err)
838838
assert.Nil(t, exStmtResp)
839839
assert.Nil(t, opStatusResp)
@@ -875,7 +875,7 @@ func TestConn_runQuery(t *testing.T) {
875875
client: testClient,
876876
cfg: config.WithDefaults(),
877877
}
878-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
878+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
879879

880880
assert.Error(t, err)
881881
assert.Equal(t, 1, executeStatementCount)
@@ -921,7 +921,7 @@ func TestConn_runQuery(t *testing.T) {
921921
client: testClient,
922922
cfg: config.WithDefaults(),
923923
}
924-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
924+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
925925

926926
assert.NoError(t, err)
927927
assert.Equal(t, 1, executeStatementCount)
@@ -968,7 +968,7 @@ func TestConn_runQuery(t *testing.T) {
968968
client: testClient,
969969
cfg: config.WithDefaults(),
970970
}
971-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
971+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
972972

973973
assert.Error(t, err)
974974
assert.Equal(t, 1, executeStatementCount)
@@ -1021,7 +1021,7 @@ func TestConn_runQuery(t *testing.T) {
10211021
client: testClient,
10221022
cfg: config.WithDefaults(),
10231023
}
1024-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
1024+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
10251025

10261026
assert.NoError(t, err)
10271027
assert.Equal(t, 1, executeStatementCount)
@@ -1073,7 +1073,7 @@ func TestConn_runQuery(t *testing.T) {
10731073
client: testClient,
10741074
cfg: config.WithDefaults(),
10751075
}
1076-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
1076+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
10771077

10781078
assert.Error(t, err)
10791079
assert.Equal(t, 1, executeStatementCount)
@@ -1126,7 +1126,7 @@ func TestConn_runQuery(t *testing.T) {
11261126
client: testClient,
11271127
cfg: config.WithDefaults(),
11281128
}
1129-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
1129+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
11301130

11311131
assert.NoError(t, err)
11321132
assert.Equal(t, 1, executeStatementCount)
@@ -1179,7 +1179,7 @@ func TestConn_runQuery(t *testing.T) {
11791179
client: testClient,
11801180
cfg: config.WithDefaults(),
11811181
}
1182-
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{})
1182+
exStmtResp, opStatusResp, err := testConn.runQuery(context.Background(), "select 1", []driver.NamedValue{}, nil)
11831183

11841184
assert.Error(t, err)
11851185
assert.Equal(t, 1, executeStatementCount)

‎connector.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,8 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
5555
}
5656

5757
protocolVersion := int64(c.cfg.ThriftProtocolVersion)
58+
59+
sessionStart := time.Now()
5860
session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{
5961
ClientProtocolI64: &protocolVersion,
6062
Configuration: sessionParams,
@@ -64,6 +66,8 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
6466
},
6567
CanUseMultipleCatalogs: &c.cfg.CanUseMultipleCatalogs,
6668
})
69+
sessionLatencyMs := time.Since(sessionStart).Milliseconds()
70+
6771
if err != nil {
6872
return nil, dbsqlerrint.NewRequestError(ctx, fmt.Sprintf("error connecting: host=%s port=%d, httpPath=%s", c.cfg.Host, c.cfg.Port, c.cfg.HTTPPath), err)
6973
}
@@ -81,12 +85,14 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
8185
conn.telemetry = telemetry.InitializeForConnection(
8286
ctx,
8387
c.cfg.Host,
88+
c.cfg.DriverVersion,
8489
c.client,
8590
c.cfg.EnableTelemetry,
8691
c.cfg.ForceEnableTelemetry,
8792
)
8893
if conn.telemetry != nil {
8994
log.Debug().Msg("telemetry initialized for connection")
95+
conn.telemetry.RecordOperation(ctx, conn.id, telemetry.OperationTypeCreateSession, sessionLatencyMs)
9096
}
9197
}
9298

‎internal/config/config.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -184,6 +184,9 @@ func (ucfg UserConfig) WithDefaults() UserConfig {
184184
ucfg.UseLz4Compression = false
185185
ucfg.CloudFetchConfig = CloudFetchConfig{}.WithDefaults()
186186

187+
// Enable telemetry by default (respects server feature flags)
188+
ucfg.EnableTelemetry = true
189+
187190
return ucfg
188191
}
189192

‎internal/rows/rows.go‎

Lines changed: 44 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,12 @@ type rows struct {
5757
logger_ *dbsqllog.DBSQLLogger
5858

5959
ctx context.Context
60+
61+
// Telemetry tracking
62+
telemetryCtx context.Context
63+
telemetryUpdate func(chunkCount int, bytesDownloaded int64)
64+
chunkCount int
65+
bytesDownloaded int64
6066
}
6167

6268
var _ driver.Rows = (*rows)(nil)
@@ -73,6 +79,8 @@ func NewRows(
7379
client cli_service.TCLIService,
7480
config *config.Config,
7581
directResults *cli_service.TSparkDirectResults,
82+
telemetryCtx context.Context,
83+
telemetryUpdate func(chunkCount int, bytesDownloaded int64),
7684
) (driver.Rows, dbsqlerr.DBError) {
7785

7886
var logger *dbsqllog.DBSQLLogger
@@ -103,14 +111,18 @@ func NewRows(
103111
logger.Debug().Msgf("databricks: creating Rows, pageSize: %d, location: %v", pageSize, location)
104112

105113
r := &rows{
106-
client: client,
107-
opHandle: opHandle,
108-
connId: connId,
109-
correlationId: correlationId,
110-
location: location,
111-
config: config,
112-
logger_: logger,
113-
ctx: ctx,
114+
client: client,
115+
opHandle: opHandle,
116+
connId: connId,
117+
correlationId: correlationId,
118+
location: location,
119+
config: config,
120+
logger_: logger,
121+
ctx: ctx,
122+
telemetryCtx: telemetryCtx,
123+
telemetryUpdate: telemetryUpdate,
124+
chunkCount: 0,
125+
bytesDownloaded: 0,
114126
}
115127

116128
// if we already have results for the query do some additional initialization
@@ -127,6 +139,17 @@ func NewRows(
127139
if err != nil {
128140
return r, err
129141
}
142+
143+
r.chunkCount++
144+
if directResults.ResultSet != nil && directResults.ResultSet.Results != nil && directResults.ResultSet.Results.ArrowBatches != nil {
145+
for _, batch := range directResults.ResultSet.Results.ArrowBatches {
146+
r.bytesDownloaded += int64(len(batch.Batch))
147+
}
148+
}
149+
150+
if r.telemetryUpdate != nil {
151+
r.telemetryUpdate(r.chunkCount, r.bytesDownloaded)
152+
}
130153
}
131154

132155
var d rowscanner.Delimiter
@@ -460,6 +483,19 @@ func (r *rows) fetchResultPage() error {
460483
return err1
461484
}
462485

486+
r.chunkCount++
487+
if fetchResult != nil && fetchResult.Results != nil {
488+
if fetchResult.Results.ArrowBatches != nil {
489+
for _, batch := range fetchResult.Results.ArrowBatches {
490+
r.bytesDownloaded += int64(len(batch.Batch))
491+
}
492+
}
493+
}
494+
495+
if r.telemetryUpdate != nil {
496+
r.telemetryUpdate(r.chunkCount, r.bytesDownloaded)
497+
}
498+
463499
err1 = r.makeRowScanner(fetchResult)
464500
if err1 != nil {
465501
return err1

0 commit comments

Comments
 (0)