[PECOBLR-1384] Complete telemetry implementation: Phases 8-10

.golangci.yml

@@ -63,6 +63,8 @@ linters-settings:
     exclude-generated: true
     severity: "low"
     confidence: "low"
+  nolintlint:
+    allow-unused: true
 
 run:
   timeout: 5m

README.md

@@ -56,6 +56,34 @@ To disable Cloud Fetch (e.g., when handling smaller datasets or to avoid additio
 token:[your token]@[Workspace hostname]:[Port number][Endpoint HTTP Path]?useCloudFetch=false
 ```
 
+### Telemetry Configuration (Optional)
+
+The driver includes optional telemetry to help improve performance and reliability. Telemetry is **disabled by default** and requires explicit opt-in.
+
+**Opt-in to telemetry** (respects server-side feature flags):
+```
+token:[your token]@[Workspace hostname]:[Port number][Endpoint HTTP Path]?enableTelemetry=true
+```
+
+**Opt-out of telemetry** (explicitly disable):
+```
+token:[your token]@[Workspace hostname]:[Port number][Endpoint HTTP Path]?enableTelemetry=false
+```
+
+**What data is collected:**
+- ✅ Query latency and performance metrics
+- ✅ Error codes (not error messages)
+- ✅ Feature usage (CloudFetch, LZ4, etc.)
+- ✅ Driver version and environment info
+
+**What is NOT collected:**
+- ❌ SQL query text
+- ❌ Query results or data values
+- ❌ Table/column names
+- ❌ User identities or credentials
+
+Telemetry has < 1% performance overhead and uses circuit breaker protection to ensure it never impacts your queries. For more details, see `telemetry/DESIGN.md` and `telemetry/TROUBLESHOOTING.md`.
+
 ### Connecting with a new Connector
 
 You can also connect with a new connector object. For example:

auth/oauth/u2m/authenticator.go

@@ -40,16 +40,17 @@ func NewAuthenticator(hostName string, timeout time.Duration) (auth.Authenticato
 	cloud := oauth.InferCloudFromHost(hostName)
 
 	var clientID, redirectURL string
-	if cloud == oauth.AWS {
+	switch cloud {
+	case oauth.AWS:
 		clientID = awsClientId
 		redirectURL = awsRedirectURL
-	} else if cloud == oauth.Azure {
+	case oauth.Azure:
 		clientID = azureClientId
 		redirectURL = azureRedirectURL
-	} else if cloud == oauth.GCP {
+	case oauth.GCP:
 		clientID = gcpClientId
 		redirectURL = gcpRedirectURL
-	} else {
+	default:
 		return nil, errors.New("unhandled cloud type: " + cloud.String())
 	}
 
@@ -147,14 +148,14 @@ func (tsp *tokenSourceProvider) GetTokenSource() (oauth2.TokenSource, error) {
 	if err != nil {
 		return nil, err
 	}
-	defer listener.Close()
+	defer listener.Close() //nolint:errcheck
 
 	srv := &http.Server{
 		ReadHeaderTimeout: 3 * time.Second,
 		WriteTimeout:      30 * time.Second,
 	}
 
-	defer srv.Close()
+	defer srv.Close() //nolint:errcheck
 
 	// Start local server to wait for callback
 	go func() {
@@ -209,7 +210,7 @@ func (tsp *tokenSourceProvider) ServeHTTP(w http.ResponseWriter, r *http.Request
 	if resp.err != "" {
 		log.Error().Msg(resp.err)
 		w.WriteHeader(http.StatusBadRequest)
-		_, err := w.Write([]byte(errorHTML("Identity Provider returned an error: " + resp.err)))
+		_, err := w.Write([]byte(errorHTML("Identity Provider returned an error: " + resp.err))) //nolint:gosec // XSS not a concern for local OAuth callback
 		if err != nil {
 			log.Error().Err(err).Msg("unable to write error response")
 		}

auth/tokenprovider/exchange.go

@@ -138,7 +138,7 @@ func (p *FederationProvider) tryTokenExchange(ctx context.Context, subjectToken
 	}
 
 	// Create request
-	req, err := http.NewRequestWithContext(ctx, "POST", exchangeURL, strings.NewReader(data.Encode()))
+	req, err := http.NewRequestWithContext(ctx, "POST", exchangeURL, strings.NewReader(data.Encode())) //nolint:gosec // URL is from trusted config
 	if err != nil {
 		return nil, fmt.Errorf("failed to create request: %w", err)
 	}
@@ -147,11 +147,11 @@ func (p *FederationProvider) tryTokenExchange(ctx context.Context, subjectToken
 	req.Header.Set("Accept", "*/*")
 
 	// Make request
-	resp, err := p.httpClient.Do(req)
+	resp, err := p.httpClient.Do(req) //nolint:gosec // G704: URL is from trusted configuration
 	if err != nil {
 		return nil, fmt.Errorf("request failed: %w", err)
 	}
-	defer resp.Body.Close()
+	defer resp.Body.Close() //nolint:errcheck
 
 	body, err := io.ReadAll(resp.Body)
 	if err != nil {

auth/tokenprovider/federation_test.go

@@ -108,7 +108,8 @@ func TestFederationProvider_TokenExchangeSuccess(t *testing.T) {
 		assert.Equal(t, "application/x-www-form-urlencoded", r.Header.Get("Content-Type"))
 		assert.Equal(t, "*/*", r.Header.Get("Accept"))
 
-		// Parse form data
+		// Parse form data - limit body size to prevent G120
+		r.Body = http.MaxBytesReader(w, r.Body, 1<<20)
 		err := r.ParseForm()
 		require.NoError(t, err)
 
@@ -155,13 +156,14 @@ func TestFederationProvider_TokenExchangeWithClientID(t *testing.T) {
 
 	// Create mock server that checks for client_id
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		r.Body = http.MaxBytesReader(w, r.Body, 1<<20)
 		err := r.ParseForm()
 		require.NoError(t, err)
 
 		// Verify client_id is present
 		assert.Equal(t, clientID, r.FormValue("client_id"))
 
-		response := map[string]interface{}{
+		response := map[string]interface{}{ //nolint:gosec // G101: test token, not a real credential
 			"access_token": "sp-wide-federation-token",
 			"token_type":   "Bearer",
 			"expires_in":   3600,

auth/tokenprovider/provider_test.go

@@ -146,7 +146,7 @@ func TestExternalTokenProvider(t *testing.T) {
 		callCount := 0
 		tokenFunc := func() (string, error) {
 			callCount++
-			return "external-token-" + string(rune(callCount)), nil
+			return "external-token-" + string(rune(callCount)), nil //nolint:gosec // G115: test counter, values are always small
 		}
 
 		provider := NewExternalTokenProvider(tokenFunc)
@@ -211,7 +211,7 @@ func TestExternalTokenProvider(t *testing.T) {
 		counter := 0
 		tokenFunc := func() (string, error) {
 			counter++
-			return "token-" + string(rune(counter)), nil
+			return "token-" + string(rune(counter)), nil //nolint:gosec // G115: test counter, values are always small
 		}
 
 		provider := NewExternalTokenProvider(tokenFunc)

connection.go

@@ -52,16 +52,19 @@ func (c *conn) Close() error {
 	log := logger.WithContext(c.id, "", "")
 	ctx := driverctx.NewContextWithConnId(context.Background(), c.id)
 
-	// Close telemetry and release resources
+	// Time CloseSession so we can record DELETE_SESSION before flushing telemetry
+	closeStart := time.Now()
+	_, err := c.client.CloseSession(ctx, &cli_service.TCloseSessionReq{
+		SessionHandle: c.session.SessionHandle,
+	})
+
+	// 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.Close(ctx)
 		telemetry.ReleaseForConnection(c.cfg.Host)
 	}
 
-	_, err := c.client.CloseSession(ctx, &cli_service.TCloseSessionReq{
-		SessionHandle: c.session.SessionHandle,
-	})
-
 	if err != nil {
 		log.Err(err).Msg("databricks: failed to close connection")
 		return dbsqlerrint.NewBadConnectionError(err)
@@ -93,7 +96,7 @@ func (c *conn) Ping(ctx context.Context) error {
 		log.Err(err).Msg("databricks: failed to ping")
 		return dbsqlerrint.NewBadConnectionError(err)
 	}
-	defer rows.Close()
+	defer rows.Close() //nolint:errcheck
 
 	log.Debug().Msg("databricks: ping successful")
 	return nil
@@ -155,9 +158,13 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 		alreadyClosed := exStmtResp.DirectResults != nil && exStmtResp.DirectResults.CloseOperation != nil
 		newCtx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), c.id), corrId)
 		if !alreadyClosed && (opStatusResp == nil || opStatusResp.GetOperationState() != cli_service.TOperationState_CLOSED_STATE) {
+			closeOpStart := time.Now()
 			_, err1 := c.client.CloseOperation(newCtx, &cli_service.TCloseOperationReq{
 				OperationHandle: exStmtResp.OperationHandle,
 			})
+			if c.telemetry != nil {
+				c.telemetry.RecordOperation(ctx, c.id, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), err1)
+			}
 			if err1 != nil {
 				log.Err(err1).Msg("databricks: failed to close operation after executing statement")
 				closeOpErr = err1 // Capture for telemetry
@@ -216,7 +223,15 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 		return nil, dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, opStatusResp)
 	}
 
-	rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
+	// Telemetry callback for tracking row fetching metrics
+	telemetryUpdate := func(chunkCount int, bytesDownloaded int64) {
+		if c.telemetry != nil {
+			c.telemetry.AddTag(ctx, "chunk_count", chunkCount)
+			c.telemetry.AddTag(ctx, "bytes_downloaded", bytesDownloaded)
+		}
+	}
+
+	rows, err := rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate)
 	return rows, err
 
 }
@@ -381,7 +396,14 @@ 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)
 
@@ -514,11 +536,11 @@ func (c *conn) handleStagingPut(ctx context.Context, presignedUrl string, header
 	}
 	client := &http.Client{}
 
-	dat, err := os.Open(localFile)
+	dat, err := os.Open(localFile) //nolint:gosec // localFile is provided by the application, not user input
 	if err != nil {
 		return dbsqlerrint.NewDriverError(ctx, "error reading local file", err)
 	}
-	defer dat.Close()
+	defer dat.Close() //nolint:errcheck
 
 	info, err := dat.Stat()
 	if err != nil {
@@ -535,7 +557,7 @@ func (c *conn) handleStagingPut(ctx context.Context, presignedUrl string, header
 	if err != nil {
 		return dbsqlerrint.NewDriverError(ctx, "error sending http request", err)
 	}
-	defer res.Body.Close()
+	defer res.Body.Close() //nolint:errcheck
 	content, err := io.ReadAll(res.Body)
 
 	if err != nil || !Succeeded(res) {
@@ -559,7 +581,7 @@ func (c *conn) handleStagingGet(ctx context.Context, presignedUrl string, header
 	if err != nil {
 		return dbsqlerrint.NewDriverError(ctx, "error sending http request", err)
 	}
-	defer res.Body.Close()
+	defer res.Body.Close() //nolint:errcheck
 	content, err := io.ReadAll(res.Body)
 
 	if err != nil || !Succeeded(res) {
@@ -583,7 +605,7 @@ func (c *conn) handleStagingRemove(ctx context.Context, presignedUrl string, hea
 	if err != nil {
 		return dbsqlerrint.NewDriverError(ctx, "error sending http request", err)
 	}
-	defer res.Body.Close()
+	defer res.Body.Close() //nolint:errcheck
 	content, err := io.ReadAll(res.Body)
 
 	if err != nil || !Succeeded(res) {
@@ -646,11 +668,18 @@ func (c *conn) execStagingOperation(
 	}
 
 	if len(driverctx.StagingPathsFromContext(ctx)) != 0 {
-		row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults)
+		// Telemetry callback for staging operation row fetching
+		telemetryUpdate := func(chunkCount int, bytesDownloaded int64) {
+			if c.telemetry != nil {
+				c.telemetry.AddTag(ctx, "chunk_count", chunkCount)
+				c.telemetry.AddTag(ctx, "bytes_downloaded", bytesDownloaded)
+			}
+		}
+		row, err = rows.NewRows(ctx, exStmtResp.OperationHandle, c.client, c.cfg, exStmtResp.DirectResults, telemetryUpdate)
 		if err != nil {
 			return dbsqlerrint.NewDriverError(ctx, "error reading row.", err)
 		}
-		defer row.Close()
+		defer row.Close() //nolint:errcheck
 
 	} else {
 		return dbsqlerrint.NewDriverError(ctx, "staging ctx must be provided.", nil)
@@ -663,7 +692,7 @@ func (c *conn) execStagingOperation(
 	if err != nil {
 		return dbsqlerrint.NewDriverError(ctx, "error fetching staging operation results", err)
 	}
-	var stringValues []string = make([]string, 4)
+	stringValues := make([]string, 4)
 	for i, val := range sqlRow { // this will either be 3 (remove op) or 4 (put/get) elements
 		if s, ok := val.(string); ok {
 			stringValues[i] = s

connection_test.go

@@ -165,8 +165,8 @@ func TestConn_executeStatement(t *testing.T) {
 		for _, opTest := range operationStateTests {
 			closeOperationCount = 0
 			executeStatementCount = 0
-			executeStatementResp.DirectResults.OperationStatus.OperationState = &opTest.state
-			executeStatementResp.DirectResults.OperationStatus.DisplayMessage = &opTest.err
+			executeStatementResp.DirectResults.OperationStatus.OperationState = &opTest.state //nolint:gosec // G601: pointer is used only within this loop iteration
+			executeStatementResp.DirectResults.OperationStatus.DisplayMessage = &opTest.err   //nolint:gosec // G601: pointer is used only within this loop iteration
 			_, err := testConn.ExecContext(context.Background(), "select 1", []driver.NamedValue{})
 			if opTest.err == "" {
 				assert.NoError(t, err)

connector.go

@@ -55,6 +55,8 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	}
 
 	protocolVersion := int64(c.cfg.ThriftProtocolVersion)
+
+	sessionStart := time.Now()
 	session, err := tclient.OpenSession(ctx, &cli_service.TOpenSessionReq{
 		ClientProtocolI64: &protocolVersion,
 		Configuration:     sessionParams,
@@ -64,6 +66,8 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 		},
 		CanUseMultipleCatalogs: &c.cfg.CanUseMultipleCatalogs,
 	})
+	sessionLatencyMs := time.Since(sessionStart).Milliseconds()
+
 	if err != nil {
 		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)
 	}
@@ -80,11 +84,15 @@ func (c *connector) Connect(ctx context.Context) (driver.Conn, error) {
 	conn.telemetry = telemetry.InitializeForConnection(
 		ctx,
 		c.cfg.Host,
+		c.cfg.DriverVersion,
 		c.client,
 		c.cfg.EnableTelemetry,
+		c.cfg.TelemetryBatchSize,
+		c.cfg.TelemetryFlushInterval,
 	)
 	if conn.telemetry != nil {
 		log.Debug().Msg("telemetry initialized for connection")
+		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)
@@ -290,8 +298,8 @@ func WithTransport(t http.RoundTripper) ConnOption {
 	return func(c *config.Config) {
 		c.Transport = t
 
-		if c.CloudFetchConfig.HTTPClient == nil {
-			c.CloudFetchConfig.HTTPClient = &http.Client{
+		if c.HTTPClient == nil {
+			c.HTTPClient = &http.Client{
 				Transport: t,
 			}
 		}

connector_test.go

@@ -263,8 +263,8 @@ func TestNewConnector(t *testing.T) {
 
 		coni, ok := con.(*connector)
 		require.True(t, ok)
-		assert.NotNil(t, coni.cfg.CloudFetchConfig.HTTPClient)
-		assert.Equal(t, customTransport, coni.cfg.CloudFetchConfig.HTTPClient.Transport)
+		assert.NotNil(t, coni.cfg.HTTPClient)
+		assert.Equal(t, customTransport, coni.cfg.HTTPClient.Transport)
 	})
 }