Tag session-close, statement-close, and rate-limit errors for telemetry

CHANGELOG.md

@@ -1,7 +1,7 @@
 # Release History
 
 ## Unreleased
-- Enrich telemetry error classification with source-declared error categories, so CloudFetch download and decompression, Arrow schema parsing, result-fetch, unsupported staging-operation, and query-cancellation failures now report a specific `error_name` (`chunk_download_error`, `decompression_error`, `arrow_schema_parsing_error`, `result_set_error`, `unsupported_operation`, `execute_statement_cancelled`) instead of a generic or message-inferred one (databricks/databricks-sql-go#414, #415)
+- Improve telemetry error reporting: driver failures are now categorized by cause instead of reported as a generic error (databricks/databricks-sql-go#414, #415, #417, #419, #424)
 
 ## v1.14.0 (2026-07-13)
 - **Minimum Go version is now 1.25.0** (previously 1.20): the `go` directive was raised to 1.25.0 while clearing OSV-Scanner findings and updating dependencies. Consumers building with an older toolchain will need to upgrade Go (databricks/databricks-sql-go#368)

category_test.go

@@ -0,0 +1,45 @@
+package dbsql
+
+import (
+	"context"
+	"database/sql/driver"
+	"testing"
+
+	thriftbackend "github.com/databricks/databricks-sql-go/internal/backend/thrift"
+	"github.com/databricks/databricks-sql-go/internal/cli_service"
+	"github.com/databricks/databricks-sql-go/internal/client"
+	"github.com/databricks/databricks-sql-go/internal/config"
+	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+// tagStatementClosed tags a non-nil close error statement_closed; nil stays nil.
+func TestTagStatementClosed(t *testing.T) {
+	assert.Nil(t, tagStatementClosed(context.Background(), nil))
+
+	tagged := tagStatementClosed(context.Background(), errors.New("close rpc failed"))
+	assert.NotNil(t, tagged)
+	assert.Equal(t, dbsqlerrint.CategoryStatementClosed, dbsqlerrint.CategoryFromError(tagged))
+}
+
+// A failed conn.Close must still return a driver.ErrBadConn (pool eviction),
+// even though it tags a telemetry-only session_closed copy.
+func TestConnCloseFailureReturnsBadConn(t *testing.T) {
+	testClient := &client.TestClient{
+		FnCloseSession: func(ctx context.Context, req *cli_service.TCloseSessionReq) (*cli_service.TCloseSessionResp, error) {
+			return nil, errors.New("close session rpc failed")
+		},
+	}
+	c := &conn{
+		cfg:     config.WithDefaults(),
+		backend: thriftbackend.NewForTest(testClient, getTestSession(), config.WithDefaults()),
+		// telemetry nil: the returned-error path is outside the telemetry block.
+	}
+
+	err := c.Close()
+	assert.Error(t, err)
+	assert.True(t, errors.Is(err, driver.ErrBadConn), "a failed close must surface as ErrBadConn")
+	assert.Equal(t, dbsqlerrint.ErrorCategory(""), dbsqlerrint.CategoryFromError(err),
+		"the returned error must not carry the telemetry category")
+}

connection.go

@@ -33,6 +33,16 @@ type conn struct {
 	telemetry *telemetry.Interceptor // Optional telemetry interceptor
 }
 
+// tagStatementClosed returns a telemetry-only copy of a close-RPC error tagged
+// statement_closed (nil stays nil). The raw error still flows to the caller.
+func tagStatementClosed(ctx context.Context, err error) error {
+	if err == nil {
+		return nil
+	}
+	return dbsqlerrint.NewRequestError(ctx, "close statement request error", err).
+		WithCategory(dbsqlerrint.CategoryStatementClosed)
+}
+
 // Prepare prepares a statement with the query bound to this connection.
 func (c *conn) Prepare(query string) (driver.Stmt, error) {
 	return &stmt{conn: c, query: query}, nil
@@ -56,7 +66,14 @@ 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)
+		// Tag a telemetry-only copy; err stays untagged so it is still returned
+		// as a driver.ErrBadConn below (database/sql needs that for pool eviction).
+		telErr := err
+		if err != nil {
+			telErr = dbsqlerrint.NewRequestError(ctx, "close session request error", err).
+				WithCategory(dbsqlerrint.CategorySessionClosed)
+		}
+		c.telemetry.RecordOperation(ctx, c.id, "", telemetry.OperationTypeDeleteSession, time.Since(closeStart).Milliseconds(), telErr)
 		_ = c.telemetry.Close(ctx)
 		telemetry.ReleaseForConnection(c.cfg.Host)
 	}
@@ -173,7 +190,7 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 	closed, err1 := op.Close(newCtx)
 	if closed {
 		if c.telemetry != nil {
-			c.telemetry.RecordOperation(ctx, c.id, statementID, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), err1)
+			c.telemetry.RecordOperation(ctx, c.id, statementID, telemetry.OperationTypeCloseStatement, time.Since(closeOpStart).Milliseconds(), tagStatementClosed(ctx, err1))
 		}
 		if err1 != nil {
 			log.Err(err1).Msg("databricks: failed to close operation after executing statement")
@@ -365,13 +382,13 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 			interceptor.AfterExecute(telemetryCtx, iterErr)
 			interceptor.CompleteStatement(telemetryCtx, stmtID, iterErr != nil)
 			// CLOSE_STATEMENT uses the actual CloseOperation RPC error.
-			interceptor.RecordOperation(telemetryCtx, connID, stmtID, telemetry.OperationTypeCloseStatement, latencyMs, closeErr)
+			interceptor.RecordOperation(telemetryCtx, connID, stmtID, telemetry.OperationTypeCloseStatement, latencyMs, tagStatementClosed(telemetryCtx, closeErr))
 		}
 	} else if c.telemetry != nil {
 		interceptor := c.telemetry
 		connID := c.id
 		closeCallback = func(latencyMs int64, _ int, _ error, closeErr error) {
-			interceptor.RecordOperation(telemetryCtx, connID, "", telemetry.OperationTypeCloseStatement, latencyMs, closeErr)
+			interceptor.RecordOperation(telemetryCtx, connID, "", telemetry.OperationTypeCloseStatement, latencyMs, tagStatementClosed(telemetryCtx, closeErr))
 		}
 	}
 

internal/client/category_test.go

@@ -0,0 +1,47 @@
+package client
+
+import (
+	"context"
+	"net/http"
+	"testing"
+
+	"github.com/databricks/databricks-sql-go/driverctx"
+	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
+	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
+	"github.com/pkg/errors"
+	"github.com/stretchr/testify/assert"
+)
+
+// categoryForStatus returns the category errorHandler assigns for a status code.
+func categoryForStatus(code int) dbsqlerrint.ErrorCategory {
+	var resp *http.Response
+	if code != 0 {
+		resp = &http.Response{StatusCode: code}
+	}
+	//nolint:bodyclose // synthetic response, no body
+	_, werr := errorHandler(resp, errors.New("boom"), 3)
+	return dbsqlerrint.CategoryFromError(werr)
+}
+
+// Only a retry-exhausted 429 is tagged rate_limit_exceeded; other statuses aren't.
+func TestErrorHandlerRateLimitTag(t *testing.T) {
+	assert.Equal(t, dbsqlerrint.CategoryRateLimitExceeded, categoryForStatus(http.StatusTooManyRequests))
+	assert.Equal(t, dbsqlerrint.ErrorCategory(""), categoryForStatus(http.StatusServiceUnavailable), "503 is retryable but not rate limiting")
+	assert.Equal(t, dbsqlerrint.ErrorCategory(""), categoryForStatus(http.StatusInternalServerError))
+	assert.Equal(t, dbsqlerrint.ErrorCategory(""), categoryForStatus(0), "nil response")
+}
+
+// The 429 telemetry error must carry the request's connection id, not empty ids.
+func TestErrorHandlerRateLimitCarriesIDs(t *testing.T) {
+	ctx := driverctx.NewContextWithConnId(context.Background(), "conn-xyz")
+	req, err := http.NewRequestWithContext(ctx, http.MethodGet, "http://example", nil)
+	assert.NoError(t, err)
+	resp := &http.Response{StatusCode: http.StatusTooManyRequests, Request: req}
+
+	//nolint:bodyclose // synthetic response, no body
+	_, werr := errorHandler(resp, errors.New("boom"), 3)
+
+	var dbErr dbsqlerr.DBError
+	assert.True(t, errors.As(werr, &dbErr))
+	assert.Equal(t, "conn-xyz", dbErr.ConnectionId(), "429 telemetry error must carry the request's connection id")
+}

internal/client/client.go

@@ -674,6 +674,18 @@ func errorHandler(resp *http.Response, err error, numTries int) (*http.Response,
 			}
 		}
 
+		// Tag retry-exhausted 429s here: net/http drops the response on error,
+		// so the status is invisible downstream. TODO(PECOBLR-3537): a per-attempt
+		// 429 counter metric could follow; this tags only the exhausted failure.
+		if resp.StatusCode == http.StatusTooManyRequests {
+			rlCtx := context.TODO()
+			if resp.Request != nil { // nil only in tests
+				rlCtx = resp.Request.Context()
+			}
+			werr = dbsqlerrint.NewRequestError(rlCtx, "rate limit exceeded (HTTP 429)", werr).
+				WithCategory(dbsqlerrint.CategoryRateLimitExceeded)
+		}
+
 		logger.Err(werr).Msg(resp.Status)
 	}