Skip to content

fix(databricks): use shared connection pool to prevent OAuth CSRF race#5746

Open
gabepesco wants to merge 1 commit intoSQLMesh:mainfrom
gabepesco:fix/databricks-oauth-shared-connection
Open

fix(databricks): use shared connection pool to prevent OAuth CSRF race#5746
gabepesco wants to merge 1 commit intoSQLMesh:mainfrom
gabepesco:fix/databricks-oauth-shared-connection

Conversation

@gabepesco
Copy link
Copy Markdown

Description

When concurrent_tasks > 1, DatabricksConnectionConfig previously used ThreadLocalConnectionPool, which creates a separate databricks.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 CSRF state parameter — 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 with MismatchingStateError.

Setting shared_connection = True on DatabricksConnectionConfig causes ThreadLocalSharedConnectionPool to 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_connection to tests/core/test_connection_config.py, which asserts that DatabricksConnectionConfig with concurrent_tasks > 1 produces a ThreadLocalSharedConnectionPool. 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-apply with auth_type="databricks-oauth" and concurrent_tasks=4 completes successfully with a single OAuth browser prompt and no MismatchingStateError.

Checklist

  • I have run make style and fixed any issues
  • I have added tests for my changes (if applicable)
  • All existing tests pass (make fast-test)
  • My commits are signed off (git commit -s) per the DCO

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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Databricks U2M OAuth CSRF Error when concurrent_tasks > 1

1 participant