fix(databricks): use shared connection pool to prevent OAuth CSRF race#5746
Open
gabepesco wants to merge 1 commit intoSQLMesh:mainfrom
Open
fix(databricks): use shared connection pool to prevent OAuth CSRF race#5746gabepesco wants to merge 1 commit intoSQLMesh:mainfrom
gabepesco wants to merge 1 commit intoSQLMesh:mainfrom
Conversation
When concurrent_tasks > 1, DatabricksConnectionConfig previously used ThreadLocalConnectionPool, which creates a separate databricks.sql.connect() per thread. For U2M OAuth (databricks-oauth / azure-oauth), each thread triggers an independent browser-based OAuth flow; these race on the CSRF state parameter and cause MismatchingStateError. Setting shared_connection = True causes ThreadLocalSharedConnectionPool to be used instead: a single connection is created (behind a lock) and each thread receives its own cursor, so only one OAuth flow is ever initiated. This mirrors the existing pattern used by DuckDBConnectionConfig. Fixes SQLMesh#5646 Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> Signed-off-by: Gabe Pesco <PescoG@medinsight.milliman.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
When
concurrent_tasks > 1,DatabricksConnectionConfigpreviously usedThreadLocalConnectionPool, which creates a separatedatabricks.sql.connect()call per thread. For U2M OAuth (databricks-oauth/azure-oauth), each thread triggers an independent browser-based OAuth flow. These flows race on the CSRFstateparameter — each generates its own random state, opens a local HTTP callback server, and waits for the browser redirect. When the callback arrives at the wrong thread's server, the state validation fails withMismatchingStateError.Setting
shared_connection = TrueonDatabricksConnectionConfigcausesThreadLocalSharedConnectionPoolto be used instead: a single connection is created (behind a lock) so only one OAuth flow is ever initiated, while each thread still receives its own cursor. PAT auth is also unaffected — a shared connection is always preferable for Databricks since each connection carries a full HTTP session.This mirrors the existing pattern used by
DuckDBConnectionConfig.Fixes #5646
Test Plan
Added
test_databricks_shared_connectiontotests/core/test_connection_config.py, which asserts thatDatabricksConnectionConfigwithconcurrent_tasks > 1produces aThreadLocalSharedConnectionPool. The test was confirmed to fail before the fix and pass after.All existing tests pass:
tests/core/test_connection_config.py,tests/utils/test_connection_pool.py,tests/core/engine_adapter/test_databricks.py.Verified end-to-end locally: SQLMesh
plan --auto-applywithauth_type="databricks-oauth"andconcurrent_tasks=4completes successfully with a single OAuth browser prompt and noMismatchingStateError.Checklist
make styleand fixed any issuesmake fast-test)git commit -s) per the DCO