diff --git a/CHANGELOG.md b/CHANGELOG.md index 0b736e9a..b20245cd 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,8 @@ # Release History +## Unreleased +- Interrupt in-flight Thrift status polls when the query context is cancelled, so remote query cancellation can proceed promptly. + ## v1.19.0 (2026-11-04) - Upgrade the SQL kernel bindings to v1.4.0. - 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). diff --git a/internal/backend/thrift/backend.go b/internal/backend/thrift/backend.go index 6601fb30..348edafc 100644 --- a/internal/backend/thrift/backend.go +++ b/internal/backend/thrift/backend.go @@ -374,9 +374,13 @@ func (b *Backend) pollOperation(ctx context.Context, opHandle *cli_service.TOper var err error log.Debug().Msg("databricks: polling status") debuglog.Logf(ctx, "thrift.Backend.pollOperation", "GetOperationStatus") - statusResp, err = b.client.GetOperationStatus(newCtx, &cli_service.TGetOperationStatusReq{ + statusResp, err = b.client.GetOperationStatus(ctx, &cli_service.TGetOperationStatusReq{ OperationHandle: opHandle, }) + if err != nil && ctx.Err() != nil { + // Let Watch send remote cancellation instead of returning the interrupted poll error. + return func() bool { return false }, nil, nil + } if statusResp != nil && statusResp.OperationState != nil { log.Debug().Msgf("databricks: status %s", statusResp.GetOperationState().String()) diff --git a/internal/backend/thrift/cancellation_test.go b/internal/backend/thrift/cancellation_test.go new file mode 100644 index 00000000..f26ba90a --- /dev/null +++ b/internal/backend/thrift/cancellation_test.go @@ -0,0 +1,100 @@ +package thrift + +import ( + "context" + "errors" + "sync" + "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" + "github.com/stretchr/testify/require" +) + +type cancellationContextKey struct{} + +type cancellationRequest struct { + handle *cli_service.TOperationHandle + value any + err error +} + +func TestPollOperationCancelsBlockedStatusRequest(t *testing.T) { + ctx, cancel := context.WithCancel(context.WithValue(context.Background(), cancellationContextKey{}, "request-value")) + defer cancel() + assertBlockedPollCancellation(t, ctx, cancel) +} + +func TestPollOperationHonorsDeadlineDuringStatusRequest(t *testing.T) { + ctx, cancel := context.WithTimeout(context.WithValue(context.Background(), cancellationContextKey{}, "request-value"), 250*time.Millisecond) + defer cancel() + assertBlockedPollCancellation(t, ctx, nil) +} + +func assertBlockedPollCancellation(t *testing.T, ctx context.Context, cancel context.CancelFunc) { + t.Helper() + started := make(chan struct{}) + release := make(chan struct{}) + requests := make(chan cancellationRequest, 2) + finished := make(chan error, 1) + cfg := config.WithDefaults() + cfg.PollInterval = time.Millisecond + be := NewForTest(blockedStatusClient(started, release, requests), getTestSession(), cfg) + handle := &cli_service.TOperationHandle{OperationId: &cli_service.THandleIdentifier{GUID: make([]byte, 16)}} + defer close(release) + go func() { + _, err := be.pollOperation(ctx, handle) + finished <- err + }() + select { + case <-started: + case <-time.After(time.Second): + t.Fatal("status request did not start") + } + if cancel != nil { + cancel() + } + select { + case <-ctx.Done(): + case <-time.After(time.Second): + t.Fatal("query context was not cancelled") + } + select { + case request := <-requests: + require.Equal(t, handle, request.handle) + require.Equal(t, "request-value", request.value) + require.NoError(t, request.err, "remote cancellation must use a live context") + case <-time.After(250 * time.Millisecond): + t.Fatal("CancelOperation was not called promptly after context cancellation") + } + select { + case err := <-finished: + require.ErrorIs(t, err, ctx.Err()) + case <-time.After(time.Second): + t.Fatal("pollOperation did not return after cancellation") + } + require.Empty(t, requests, "remote cancellation must be sent once") +} + +func blockedStatusClient(started, release chan struct{}, requests chan cancellationRequest) *client.TestClient { + var start sync.Once + return &client.TestClient{ + FnGetOperationStatus: func(ctx context.Context, _ *cli_service.TGetOperationStatusReq) (*cli_service.TGetOperationStatusResp, error) { + start.Do(func() { close(started) }) + select { + case <-ctx.Done(): + return nil, ctx.Err() + case <-release: + return nil, errors.New("test released blocked status request") + } + }, + FnCancelOperation: func(ctx context.Context, request *cli_service.TCancelOperationReq) (*cli_service.TCancelOperationResp, error) { + requests <- cancellationRequest{request.OperationHandle, ctx.Value(cancellationContextKey{}), ctx.Err()} + return &cli_service.TCancelOperationResp{ + Status: &cli_service.TStatus{StatusCode: cli_service.TStatusCode_SUCCESS_STATUS}, + }, nil + }, + } +}