Merge branch 'main' into add-databricks-query-tags-session-properties

Author: WhammyLeaf

.github/workflows/pr.yaml

@@ -188,15 +188,15 @@ jobs:
 
   ui-style:
     needs: [changes]
-    if:
-      needs.changes.outputs.client == 'true' || needs.changes.outputs.ci ==
-      'true' || github.ref == 'refs/heads/main'
+    if: false
+      # needs.changes.outputs.client == 'true' || needs.changes.outputs.ci ==
+      # 'true' || github.ref == 'refs/heads/main'
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v5
       - uses: actions/setup-node@v6
         with:
-          node-version: '20'
+          node-version: '22'
       - uses: pnpm/action-setup@v4
         with:
           version: latest
@@ -396,6 +396,7 @@ jobs:
     env:
       PLAYWRIGHT_SKIP_BROWSER_DOWNLOAD: 1
     runs-on: ubuntu-latest
+    if: false
     steps:
       - uses: actions/checkout@v5
       - uses: actions/setup-node@v6

sqlmesh/core/config/connection.py

@@ -814,6 +814,8 @@ class DatabricksConnectionConfig(ConnectionConfig):
     DISPLAY_NAME: t.ClassVar[t.Literal["Databricks"]] = "Databricks"
     DISPLAY_ORDER: t.ClassVar[t.Literal[3]] = 3
 
+    shared_connection: t.ClassVar[bool] = True
+
     _concurrent_tasks_validator = concurrent_tasks_validator
     _http_headers_validator = http_headers_validator
 

tests/core/test_connection_config.py

@@ -1422,6 +1422,35 @@ def test_databricks(make_config):
         )
 
 
+def test_databricks_shared_connection(make_config):
+    """Databricks should use a shared connection pool to prevent OAuth CSRF races.
+
+    When concurrent_tasks > 1, ThreadLocalConnectionPool creates one connection per
+    thread. For U2M OAuth, each thread triggers its own 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
+    gets its own cursor, so only one OAuth flow is ever initiated.
+
+    See: https://github.com/tobymao/sqlmesh/issues/5646
+    """
+    from sqlmesh.utils.connection_pool import ThreadLocalSharedConnectionPool
+
+    config = make_config(
+        type="databricks",
+        server_hostname="dbc-test.cloud.databricks.com",
+        http_path="sql/test/foo",
+        access_token="test-token",
+        concurrent_tasks=4,
+    )
+    assert isinstance(config, DatabricksConnectionConfig)
+    assert config.shared_connection is True
+
+    adapter = config.create_engine_adapter()
+    assert isinstance(adapter._connection_pool, ThreadLocalSharedConnectionPool)
+
+
 def test_engine_import_validator():
     with pytest.raises(
         ConfigError,