Tag staging, decompression, and cancellation errors for telemetry

CHANGELOG.md

@@ -1,7 +1,7 @@
 # Release History
 
 ## Unreleased
-- Enrich telemetry error classification with source-declared error categories, so CloudFetch download, Arrow schema parsing, and result-fetch failures now report a specific `error_name` (`chunk_download_error`, `arrow_schema_parsing_error`, `result_set_error`) instead of the generic `error` (databricks/databricks-sql-go#414, #415)
+- 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)
 
 ## 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)

connection.go

@@ -801,6 +801,6 @@ func (c *conn) execStagingOperation(
 	case "REMOVE":
 		return c.handleStagingRemove(ctx, presignedUrl, headers)
 	default:
-		return dbsqlerrint.NewDriverError(ctx, fmt.Sprintf("operation %s is not supported. Supported operations are GET, PUT, and REMOVE", operation), nil)
+		return dbsqlerrint.NewDriverError(ctx, fmt.Sprintf("operation %s is not supported. Supported operations are GET, PUT, and REMOVE", operation), nil).WithCategory(dbsqlerrint.CategoryUnsupportedOperation)
 	}
 }

internal/backend/thrift/backend.go

@@ -408,10 +408,15 @@ func (b *Backend) pollOperation(ctx context.Context, opHandle *cli_service.TOper
 	status, resp, err := pollSentinel.Watch(ctx, b.cfg.PollInterval, 0)
 	if err != nil {
 		log.Err(err).Msg("error polling operation status")
-		if status == sentinel.WatchTimeout {
+		switch {
+		case status == sentinel.WatchTimeout:
 			// Unreachable today (production Watch uses timeout=0); tagged so it
 			// classifies correctly if a nonzero poll timeout is ever enabled.
 			err = dbsqlerrint.NewRequestError(ctx, dbsqlerr.ErrSentinelTimeout, err).WithCategory(dbsqlerrint.CategoryStatementTimeout)
+		case errors.Is(err, context.Canceled):
+			// Caller cancelled during polling. DeadlineExceeded is intentionally
+			// not matched here, so it keeps its existing classification.
+			err = dbsqlerrint.NewExecutionError(ctx, dbsqlerr.ErrQueryExecution, err, nil).WithCategory(dbsqlerrint.CategoryStatementCancelled)
 		}
 		return nil, err
 	}

internal/backend/thrift/category_test.go

@@ -0,0 +1,58 @@
+package thrift
+
+import (
+	"context"
+	"errors"
+	"testing"
+	"time"
+
+	"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/stretchr/testify/assert"
+)
+
+// runningStatusClient keeps the operation in RUNNING_STATE so pollOperation
+// keeps polling until the context fires.
+func runningStatusClient() *client.TestClient {
+	return &client.TestClient{
+		FnGetOperationStatus: func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (*cli_service.TGetOperationStatusResp, error) {
+			return &cli_service.TGetOperationStatusResp{
+				OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_RUNNING_STATE),
+			}, nil
+		},
+	}
+}
+
+func pollTestOpHandle() *cli_service.TOperationHandle {
+	return &cli_service.TOperationHandle{
+		OperationId: &cli_service.THandleIdentifier{GUID: []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}},
+	}
+}
+
+// A query cancelled during status polling must be tagged
+// execute_statement_cancelled, and must still satisfy errors.Is(context.Canceled).
+func TestPollOperationCancelledCarriesCategory(t *testing.T) {
+	be := NewForTest(runningStatusClient(), getTestSession(), config.WithDefaults())
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+
+	_, err := be.pollOperation(ctx, pollTestOpHandle())
+	assert.NotNil(t, err)
+	assert.Equal(t, dbsqlerrint.CategoryStatementCancelled, dbsqlerrint.CategoryFromError(err))
+	assert.True(t, errors.Is(err, context.Canceled), "cancellation identity must be preserved for the thrift layer")
+}
+
+// A deadline exceeded during polling must NOT be tagged cancelled — it keeps its
+// existing (untagged) classification.
+func TestPollOperationDeadlineNotTaggedCancelled(t *testing.T) {
+	be := NewForTest(runningStatusClient(), getTestSession(), config.WithDefaults())
+	ctx, cancel := context.WithDeadline(context.Background(), time.Now().Add(-time.Hour))
+	defer cancel()
+
+	_, err := be.pollOperation(ctx, pollTestOpHandle())
+	assert.NotNil(t, err)
+	assert.Equal(t, dbsqlerrint.ErrorCategory(""), dbsqlerrint.CategoryFromError(err))
+	assert.True(t, errors.Is(err, context.DeadlineExceeded))
+}

internal/rows/arrowbased/batchloader.go

@@ -353,7 +353,7 @@ func (cft *cloudFetchDownloadTask) Run() {
 			// corruption, not a transient network condition.
 			buf, err = io.ReadAll(lz4.NewReader(bytes.NewReader(rawBody)))
 			if err != nil {
-				cft.sendResult(cloudFetchDownloadTaskResult{data: nil, err: err})
+				cft.sendResult(cloudFetchDownloadTaskResult{data: nil, err: dbsqlerrint.NewDriverError(cft.ctx, errArrowRowsCloudFetchDecompression, err).WithCategory(dbsqlerrint.CategoryDecompression)})
 				return
 			}
 		}

internal/rows/arrowbased/category_test.go

@@ -7,6 +7,7 @@ import (
 	"testing"
 	"time"
 
+	dbsqlerr "github.com/databricks/databricks-sql-go/errors"
 	"github.com/databricks/databricks-sql-go/internal/cli_service"
 	"github.com/databricks/databricks-sql-go/internal/config"
 	dbsqlerrint "github.com/databricks/databricks-sql-go/internal/errors"
@@ -91,6 +92,41 @@ func TestCloudFetchDownloadErrorCarriesCategory_RetryBranches(t *testing.T) {
 	})
 }
 
+// A CloudFetch payload that fails LZ4 decompression must surface carrying
+// CategoryDecompression, on the same object path the telemetry hook reads.
+func TestCloudFetchDecompressionErrorCarriesCategory(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Not valid LZ4 — lz4.NewReader will fail to decompress this.
+		_, _ = w.Write([]byte("this is not lz4 compressed data"))
+	}))
+	defer server.Close()
+
+	startRowOffset := int64(0)
+	links := []*cli_service.TSparkArrowResultLink{
+		{
+			FileLink:       server.URL,
+			ExpiryTime:     time.Now().Add(10 * time.Minute).Unix(),
+			StartRowOffset: startRowOffset,
+			RowCount:       1,
+		},
+	}
+
+	cfg := config.WithDefaults()
+	cfg.UseLz4Compression = true
+	cfg.MaxDownloadThreads = 1
+
+	bi, err := NewCloudBatchIterator(context.Background(), links, startRowOffset, nil, cfg, nil)
+	assert.Nil(t, err)
+
+	_, nextErr := bi.Next()
+	assert.NotNil(t, nextErr)
+	assert.Equal(t, dbsqlerrint.CategoryDecompression, dbsqlerrint.CategoryFromError(nextErr))
+	// loadBatchFor (arrowRows.go) uses a direct, non-walking assertion; the
+	// outermost error must itself be a DBError or the tag is re-wrapped away.
+	_, ok := nextErr.(dbsqlerr.DBError)
+	assert.True(t, ok, "outermost error must be a DBError to survive loadBatchFor's direct assertion")
+}
+
 // A schema-parse failure while building the row scanner must carry
 // CategoryArrowSchemaParsing. Covers the schema-convert branch (broken decimal
 // type) and the IPC-read branch (invalid pre-supplied arrow schema bytes).

internal/rows/arrowbased/errors.go

@@ -16,6 +16,7 @@ var errArrowRowsMakeColumnValueContainers = "databricks: failed creating column
 var errArrowRowsNotArrowFormat = "databricks: result set is not in arrow format"
 
 const errArrowRowsCloudFetchDownloadFailure = "cloud fetch batch loader failed to download results"
+const errArrowRowsCloudFetchDecompression = "cloud fetch batch loader failed to decompress results"
 
 func errArrowRowsUnsupportedNativeType(t string) string {
 	return fmt.Sprintf("databricks: arrow native values not yet supported for %s", t)

telemetry/errors_test.go

@@ -62,3 +62,12 @@ func TestClassifyErrorUntaggedDbErrorUsesFallback(t *testing.T) {
 	err := dbsqlerrint.NewRequestError(context.TODO(), "syntax error", errors.New("cause"))
 	assert.Equal(t, "syntax_error", classifyError(err))
 }
+
+// The staging unsupported-operation message classifies as the generic "error"
+// on its own; the source tag is what makes it report "unsupported_operation".
+func TestClassifyErrorUnsupportedOperation(t *testing.T) {
+	msg := "operation COPY is not supported. Supported operations are GET, PUT, and REMOVE"
+	assert.Equal(t, "error", classifyError(dbsqlerrint.NewDriverError(context.TODO(), msg, nil)))
+	assert.Equal(t, "unsupported_operation",
+		classifyError(dbsqlerrint.NewDriverError(context.TODO(), msg, nil).WithCategory(dbsqlerrint.CategoryUnsupportedOperation)))
+}