Fix context loss in polling and connection close operations #295
Conversation
dgiagio
requested review from
deeksha-db,
gopalldb,
jackyhu-db,
jadewang-db,
jayantsing-db,
jprakash-db,
madhav-db,
samikshya-db,
shivam2680 and
vikrantpuppala
as code owners
connection.go
Outdated
| var statusResp *cli_service.TGetOperationStatusResp | ||
| ctx = driverctx.NewContextWithConnId(ctx, c.id) | ||
| newCtx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), c.id), corrId) | ||
| newCtx := context.WithoutCancel(ctx) |
There was a problem hiding this comment.
if context is cancelled, polling will not have cancellation information. Won't this affect cancellation of long running queries?
There was a problem hiding this comment.
I see your concern, but in this specific implementation, cancellation is handled by ctx rather than through the new context.
Even though newCtx doesn't carry the cancellation signal, the pollSentinel is created using the original ctx. The pollSentinel.Watch function explicitly monitors that original context. As soon as it detects a timeout or cancel signal, it invokes OnCancelFn to terminate the operation.
So, effectively, newCtx doesn't need the cancellation info because pollSentinel is already handling that responsibility via ctx.
connection.go
Outdated
| var statusResp *cli_service.TGetOperationStatusResp | ||
| ctx = driverctx.NewContextWithConnId(ctx, c.id) | ||
| newCtx := driverctx.NewContextWithCorrelationId(driverctx.NewContextWithConnId(context.Background(), c.id), corrId) | ||
| newCtx := context.WithoutCancel(ctx) |
There was a problem hiding this comment.
I see that context.WithoutCancel was added in Go 1.21, but go.mod pins to v1.20. This will fail.
Do we need a version change in go.mod as well?
There was a problem hiding this comment.
We would like to be compatible with 1.20 and hence would recommend finding a way to avoid using context.WithoutCancel
There was a problem hiding this comment.
I've updated the PR with a backport for context.WithoutCancel in internal/compat/context/context.go. Could you please review? Thank you cc: @gopalldb @vikrantpuppala
| config *config.Config, | ||
| directResults *cli_service.TSparkDirectResults, | ||
| ) (driver.Rows, dbsqlerr.DBError) { | ||
|
|
There was a problem hiding this comment.
do we need a check for context being nil?
There was a problem hiding this comment.
I don't believe we need a check for nil here. NewRows is receiving the context when called from two places: QueryContext and execStagingOperation. In both places the ctx is expected to be non-nil already and is used a few times before being passed to NewRows.
dgiagio
force-pushed
the
dgiagio/ctx-pass-fix-pr
branch
from
70cc8c2 to
acc6a9c
Compare
There was a problem hiding this comment.
Pull request overview
This PR fixes a critical context loss issue during polling and connection close operations. The driver was previously discarding the context.Context during result fetching and operation closing, which prevented authentication mechanisms like Azure OBO from functioning properly.
Key changes:
- Pass
context.Contextthrough the result page iterator to preserve authentication credentials - Extract connection and correlation IDs from context instead of passing them separately
- Use
WithoutCancelwrapper for polling operations to preserve context values while avoiding cancellation during cleanup
Reviewed changes
Copilot reviewed 7 out of 7 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| internal/rows/rowscanner/resultPageIterator.go | Store context in iterator and use it for all server operations instead of creating fresh contexts |
| internal/rows/rows.go | Accept context as first parameter and derive IDs from it rather than separate parameters |
| internal/rows/rows_test.go | Update test calls to create context with IDs and pass to NewRows/NewResultPageIterator |
| internal/rows/arrowbased/arrowRows_test.go | Update test to create context with IDs for NewResultPageIterator |
| internal/rows/arrowbased/arrowRecordIterator_test.go | Update tests to create context with IDs for NewResultPageIterator |
| internal/compat/context/context.go | Add WithoutCancel implementation to preserve context values without cancellation |
| connection.go | Use WithoutCancel for polling, remove separate correlation ID extraction, and fix error handling on close |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Signed-off-by: Diego Giagio <[email protected]>
dgiagio
force-pushed
the
dgiagio/ctx-pass-fix-pr
branch
from
acc6a9c to
705f064
Compare
3 participants
This PR fixes an issue where the driver discards the
context.Contextduring polling, making it impossible to use authentication mechanisms (like Azure OBO) that rely on passing credentials via the context.Reference: