fix: keep polling on transport errors

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.

internal/backend/thrift/backend.go

@@ -12,6 +12,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"regexp"
 	"time"
 
 	"github.com/databricks/databricks-sql-go/driverctx"
@@ -355,15 +356,32 @@ 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)
 	log := logger.WithContext(b.SessionID(), corrId, opID)
 	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 {

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) {