From ed13752bd143df33a7b231637e87c8e24bf00180 Mon Sep 17 00:00:00 2001 From: Rashi Jaiswal Date: Mon, 2 Nov 2026 14:55:58 +0530 Subject: [PATCH] fix: keep polling on transport errors Signed-off-by: Rashi Jaiswal --- CHANGELOG.md | 1 + internal/backend/thrift/backend.go | 34 +++++++++++++++ internal/backend/thrift/backend_test.go | 58 +++++++++++++++++++++++++ 3 files changed, 93 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 066fac4e..409ec941 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Thrift: polling no longer abandons a running operation on a transport-level GetOperationStatus failure; it retries up to 30 consecutive failures and fails fast on HTTP 401/403/404. - Upgrade go-oidc to v3.21.0 (go-jose v4.1.6) and x/oauth2 to v0.36.0 to address go-jose/v3 security findings (databricks/databricks-sql-go#526). - Reduce kernel-backend JSON rendering allocations for nested Boolean and signed-integer values. - **Behavior change:** Native Arrow TIMESTAMP_NTZ results now preserve their wall-clock fields and report `TIMESTAMP_NTZ` on both backends. Returned `time.Time` values use UTC instead of the session location; their represented instant and `.UTC()` value are unchanged (databricks/databricks-sql-go#506). On Thrift, this requires a server-supplied Arrow schema with native timestamps; without it, NTZ still reports `TIMESTAMP` and uses session-location handling. String timestamps (`UseArrowNativeTimestamp=false`, internal) and non-Arrow values also retain legacy parsing. diff --git a/internal/backend/thrift/backend.go b/internal/backend/thrift/backend.go index 6601fb30..691e8092 100644 --- a/internal/backend/thrift/backend.go +++ b/internal/backend/thrift/backend.go @@ -12,6 +12,7 @@ import ( "context" "fmt" "net/http" + "regexp" "time" "github.com/databricks/databricks-sql-go/driverctx" @@ -355,8 +356,23 @@ func (b *Backend) executeStatement(ctx context.Context, req backend.ExecRequest) return resp, err } +// maxPollFailures caps consecutive transport-level GetOperationStatus failures +// (no response at all) that pollOperation retries before giving up. +const maxPollFailures = 30 + +var clientErrorStatusRe = regexp.MustCompile(`unexpected HTTP status (401|403|404)\b`) + +// isClientError reports whether err is an HTTP 401/403/404, which retrying +// will not fix. +func isClientError(err error) bool { + return clientErrorStatusRe.MatchString(err.Error()) +} + // pollOperation polls the operation status until it reaches a terminal state, // cancelling the operation if the context is done (via the sentinel poll loop). +// A transport-level GetOperationStatus failure (no response) does not abandon +// the operation: it is retried up to maxPollFailures consecutive times, but +// HTTP 401/403/404 fail immediately. func (b *Backend) pollOperation(ctx context.Context, opHandle *cli_service.TOperationHandle) (*cli_service.TGetOperationStatusResp, error) { corrId := driverctx.CorrelationIdFromContext(ctx) opID := client.SprintGuid(opHandle.OperationId.GUID) @@ -364,6 +380,8 @@ func (b *Backend) pollOperation(ctx context.Context, opHandle *cli_service.TOper defer debuglog.Track(ctx, "thrift.Backend.pollOperation", "op=%s", opID)() var statusResp *cli_service.TGetOperationStatusResp + var lastState *cli_service.TOperationState + var pollFailures int ctx = driverctx.NewContextWithConnId(ctx, b.SessionID()) newCtx := context2.WithoutCancel(ctx) pollSentinel := sentinel.Sentinel{ @@ -378,7 +396,23 @@ func (b *Backend) pollOperation(ctx context.Context, opHandle *cli_service.TOper OperationHandle: opHandle, }) + if statusResp == nil && err != nil { + pollFailures++ + if pollFailures > maxPollFailures || isClientError(err) { + state := "unknown" + if lastState != nil { + state = lastState.String() + } + msg := fmt.Sprintf("giving up polling operation %s, last known state %s", opID, state) + return nil, statusResp, dbsqlerrint.NewRequestError(ctx, msg, err) + } + log.Warn().Err(err).Msg("databricks: polling status failed, retrying") + return func() bool { return false }, nil, nil + } + + pollFailures = 0 if statusResp != nil && statusResp.OperationState != nil { + lastState = statusResp.OperationState log.Debug().Msgf("databricks: status %s", statusResp.GetOperationState().String()) } return func() bool { diff --git a/internal/backend/thrift/backend_test.go b/internal/backend/thrift/backend_test.go index ed76c3bc..ba4e95e2 100644 --- a/internal/backend/thrift/backend_test.go +++ b/internal/backend/thrift/backend_test.go @@ -8,6 +8,7 @@ import ( "database/sql/driver" "errors" "fmt" + "net/http" "testing" "time" @@ -602,6 +603,63 @@ func TestBackend_executeStatement_QueryTags(t *testing.T) { func TestBackend_pollOperation(t *testing.T) { t.Parallel() + pollHandle := func(id byte) *cli_service.TOperationHandle { + return &cli_service.TOperationHandle{ + OperationId: &cli_service.THandleIdentifier{GUID: []byte{1, 2, 3, 4, 2, 23, 4, 2, 3, 1, 3, 4, 4, 223, 34, id}, Secret: []byte("b")}, + } + } + pollConfig := func() *config.Config { + cfg := config.WithDefaults() + cfg.PollInterval = time.Millisecond + return cfg + } + + t.Run("pollOperation recovers from transport errors within the cap", func(t *testing.T) { + var count int + testClient := &client.TestClient{FnGetOperationStatus: func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (*cli_service.TGetOperationStatusResp, error) { + count++ + if count <= maxPollFailures { + return nil, errors.New("connection reset by peer") + } + return &cli_service.TGetOperationStatusResp{OperationState: cli_service.TOperationStatePtr(cli_service.TOperationState_FINISHED_STATE)}, nil + }} + be := NewForTest(testClient, getTestSession(), pollConfig()) + res, err := be.pollOperation(context.Background(), pollHandle(70)) + assert.NoError(t, err) + assert.Equal(t, cli_service.TOperationState_FINISHED_STATE, res.GetOperationState()) + }) + + t.Run("pollOperation gives up after exceeding the transport error cap", func(t *testing.T) { + var count int + testClient := &client.TestClient{FnGetOperationStatus: func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (*cli_service.TGetOperationStatusResp, error) { + count++ + return nil, errors.New("connection reset by peer") + }} + be := NewForTest(testClient, getTestSession(), pollConfig()) + res, err := be.pollOperation(context.Background(), pollHandle(71)) + assert.Error(t, err) + assert.Nil(t, res) + assert.Equal(t, maxPollFailures+1, count) + }) + + t.Run("pollOperation fails immediately on a client error", func(t *testing.T) { + for i, status := range []string{"401", "403", "404"} { + t.Run(status, func(t *testing.T) { + var count int + testClient := &client.TestClient{FnGetOperationStatus: func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (*cli_service.TGetOperationStatusResp, error) { + count++ + // Shaped like the url.Error from the transport (see client.go CheckRetry). + return nil, fmt.Errorf("Post \"https://h/sql\": unexpected HTTP status %s %s", status, http.StatusText(map[string]int{"401": 401, "403": 403, "404": 404}[status])) + }} + be := NewForTest(testClient, getTestSession(), pollConfig()) + res, err := be.pollOperation(context.Background(), pollHandle(byte(80+i))) + assert.Error(t, err) + assert.Nil(t, res) + assert.Equal(t, 1, count) + }) + } + }) + t.Run("pollOperation returns finished state response when query finishes", func(t *testing.T) { var getOperationStatusCount int getOperationStatus := func(ctx context.Context, req *cli_service.TGetOperationStatusReq) (r *cli_service.TGetOperationStatusResp, err error) {