Bug: Driver discards context, breaking context-based authentication #288
Description
Hello, we've identified a critical issue where the driver discards the context.Context during polling, making it impossible to use authentication mechanisms (like Azure OBO) that rely on passing credentials via the context. This also appears to be an issue during connection Close(), which leads to intermittent failures.
Bug 1: Polling context loss
This bug causes all long-running queries to fail when using context-based authentication.
When a query is executed (e.g., via QueryContext), it correctly receives a context containing authentication credentials. However, if the query is long-running, the driver uses a polling sentinel to check for results. This sentinel is explicitly created with a new, empty context using context.Background(), discarding the credentials.
Code reference: https://github.com/databricks/databricks-sql-go/blob/main/connection.go#L365-L370
When our authenticator is invoked to sign the poll request, it receives this newCtx empty context, cannot find the required credentials (e.g., OBO token), and authentication fails.
Suggested fix: Create the newCtx based on the passed in ctx, and not context.Background().
Bug 2: Connection Close() context loss
The standard Go database/sql package manages the connection pool. When it needs to refresh a connection (e.g., due to lifetime expiration), it calls the Close() method on the driver connection.
The driver.Conn interface for Close() does not accept a Context.context. If the driver's Close() implementation needs to make an authenticated HTTP call to Databricks (e.g., to gracefully close a session), it has no context to pass to the authenticator. This leads to the same authentication failure as Bug 1.
Suggested fix: This is a more complex issue due to database/sql interface limitations. We would request that the driver team investigate if it's possible to avoid making authenticated HTTP calls within the Close() method.