fix(databricks): run U2M OAuth queries concurrently with a shared login

sqlmesh/core/config/connection.py

@@ -839,8 +839,15 @@ class DatabricksConnectionConfig(ConnectionConfig):
     @computed_field  # type: ignore[prop-decorator]
     @property
     def shared_connection(self) -> bool:
-        """The connection should only be shared if U2M OAuth is being used"""
-        return self.auth_type is not None and self.oauth_client_secret is None
+        """U2M OAuth connections share one login, so each thread can have its own connection. Connectors too old
+        to support that fall back to sharing a single connection, so only one browser flow is run."""
+        from sqlmesh.utils.databricks import supports_shared_u2m_login
+
+        return self._uses_u2m_oauth and not supports_shared_u2m_login()
+
+    @property
+    def _uses_u2m_oauth(self) -> bool:
+        return bool(self.auth_type and "oauth" in self.auth_type and not self.oauth_client_secret)
 
     @model_validator(mode="before")
     def _databricks_connect_validator(cls, data: t.Any) -> t.Any:
@@ -992,6 +999,19 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
                 else:
                     # if auth_type is set to an 'oauth' type but no client_id/secret are set, then we are using U2M
                     # ref: https://docs.databricks.com/en/dev-tools/python-sql-connector.html#oauth-user-to-machine-u2m-authentication
+                    # All connections share one login so that each thread can have its own connection without
+                    # running its own browser flow. `credentials_provider` takes precedence over `auth_type` on the
+                    # connector's Thrift and SEA backends. `auth_type` is still passed because the connector's kernel
+                    # backend rejects a custom `credentials_provider` and runs its own U2M flow from `auth_type`.
+                    from sqlmesh.utils.databricks import (
+                        get_shared_u2m_credentials_provider,
+                        supports_shared_u2m_login,
+                    )
+
+                    if self.server_hostname and supports_shared_u2m_login():
+                        conn_kwargs["credentials_provider"] = get_shared_u2m_credentials_provider(
+                            self.server_hostname, self.auth_type
+                        )
                     conn_kwargs["auth_type"] = self.auth_type
 
             return conn_kwargs

sqlmesh/utils/databricks.py

@@ -0,0 +1,137 @@
+from __future__ import annotations
+
+import importlib.util
+import re
+import typing as t
+from functools import lru_cache
+from threading import Lock
+
+HeaderFactory = t.Callable[[], t.Dict[str, str]]
+
+
+class SharedU2MCredentialsProvider:
+    """A Databricks SQL connector `credentials_provider` that shares one U2M OAuth login across connections.
+
+    Without this, every `databricks.sql.connect()` call builds its own OAuth provider. For U2M OAuth that means
+    each thread runs its own browser flow, and these race on the CSRF `state` parameter
+    (https://github.com/tobymao/sqlmesh/issues/5646). Even if the browser flow is only run once, each connection
+    would still hold its own copy of the refresh token. Databricks can issue single-use refresh tokens, so the
+    first connection to refresh would invalidate the token held by every other connection.
+
+    Instead, a single connector OAuth provider is created lazily (behind a lock, so the browser flow only runs once)
+    and every connection reads its headers from it. Reads are also serialized so only one thread can refresh the
+    token at a time. This lets each thread have its own connection while sharing a single login.
+
+    If reading the token fails (for example because the refresh token expired or was revoked) the login is
+    discarded, so the next connection or request logs in again rather than failing for the rest of the process.
+    Network failures while refreshing keep the login, so the next request retries the refresh without a browser flow.
+    """
+
+    def __init__(self, server_hostname: str, auth_type: str):
+        self._server_hostname = server_hostname
+        self._auth_type = auth_type
+        self._auth_provider: t.Optional[t.Any] = None
+        self._lock = Lock()
+        # Incremented when a login attempt finishes, so threads that waited on a failed attempt can tell it happened
+        self._login_attempts_finished = 0
+        self._login_error: t.Optional[Exception] = None
+
+    def auth_type(self) -> str:
+        return self._auth_type
+
+    def __call__(self, *args: t.Any, **kwargs: t.Any) -> HeaderFactory:
+        """Called by the connector once per connection. Logs in if needed and returns that connection's header factory."""
+        login_attempts_seen = self._login_attempts_finished
+        with self._lock:
+            self._ensure_logged_in(login_attempts_seen)
+        return self._headers
+
+    def _headers(self) -> t.Dict[str, str]:
+        """Returns the auth headers for one request, refreshing the token if it has expired."""
+        from databricks.sql.exc import RequestError
+
+        headers: t.Dict[str, str] = {}
+        login_attempts_seen = self._login_attempts_finished
+        with self._lock:
+            auth_provider = self._ensure_logged_in(login_attempts_seen)
+            try:
+                auth_provider.add_headers(headers)
+            except RequestError:
+                # A network failure while refreshing leaves the login intact, since the connector only replaces its
+                # tokens after a successful refresh, so the next request retries the refresh rather than logging in
+                raise
+            except Exception:
+                self._auth_provider = None
+                raise
+        return headers
+
+    def _ensure_logged_in(self, login_attempts_seen: int) -> t.Any:
+        """Returns the shared connector auth provider, logging in if needed. Must be called with the lock held."""
+        if self._auth_provider is not None:
+            return self._auth_provider
+
+        # A login attempt failed while this thread was waiting on the lock. Re-raise its error rather than opening
+        # another browser window for every waiting thread; callers that arrive afterwards will try again.
+        if self._login_attempts_finished != login_attempts_seen and self._login_error is not None:
+            raise self._login_error
+
+        try:
+            self._auth_provider = self._create_auth_provider()
+            self._login_error = None
+        except Exception as e:
+            self._login_error = e
+            raise
+        finally:
+            self._login_attempts_finished += 1
+        return self._auth_provider
+
+    def _create_auth_provider(self) -> t.Any:
+        # These are internal to the connector, but using them means the OAuth flow is identical to the one the
+        # connector runs when `auth_type` is passed to `connect()` (same client ID, redirect ports, and scopes)
+        from databricks.sql import __version__
+        from databricks.sql.auth.auth import get_python_sql_connector_auth_provider
+        from databricks.sql.common.unified_http_client import UnifiedHttpClient
+        from databricks.sql.utils import build_client_context
+
+        # The provider outlives any single connection, so it gets its own HTTP client rather than borrowing one
+        # that would be closed along with its connection
+        http_client = UnifiedHttpClient(
+            build_client_context(self._server_hostname, __version__, user_agent_entry="sqlmesh")
+        )
+        return get_python_sql_connector_auth_provider(
+            self._server_hostname, http_client=http_client, auth_type=self._auth_type
+        )
+
+
+@lru_cache(maxsize=1)
+def supports_shared_u2m_login() -> bool:
+    """Whether the installed connector has the internals `SharedU2MCredentialsProvider` builds on (4.1.0+)."""
+    try:
+        return importlib.util.find_spec("databricks.sql.common.unified_http_client") is not None
+    except ModuleNotFoundError:
+        return False
+
+
+_shared_u2m_credentials_providers: t.Dict[t.Tuple[str, str], SharedU2MCredentialsProvider] = {}
+_shared_u2m_credentials_providers_lock = Lock()
+
+
+def get_shared_u2m_credentials_provider(
+    server_hostname: str, auth_type: str
+) -> SharedU2MCredentialsProvider:
+    """Returns the process-wide U2M credentials provider for the given workspace and auth type.
+
+    Engine adapters and config copies are created several times per process, so the provider is kept at module
+    level to avoid prompting the user to log in more than once. U2M OAuth is an interactive, single-user login, so
+    every connection to a workspace in this process uses the identity of whoever completed the browser flow.
+    """
+    # The same workspace can be written with or without a scheme or trailing slash
+    normalized_hostname = re.sub(r"^https?://", "", server_hostname.strip(), flags=re.I)
+    normalized_hostname = normalized_hostname.rstrip("/").lower()
+    key = (normalized_hostname, auth_type)
+    with _shared_u2m_credentials_providers_lock:
+        if key not in _shared_u2m_credentials_providers:
+            _shared_u2m_credentials_providers[key] = SharedU2MCredentialsProvider(
+                normalized_hostname, auth_type
+            )
+        return _shared_u2m_credentials_providers[key]

tests/core/test_connection_config.py

@@ -1426,21 +1426,20 @@ def test_databricks(make_config):
         )
 
 
-def test_databricks__u2m_oauth__shared_connection_pool(make_config):
-    """Databricks should use a shared connection pool when using OAuth to prevent CSRF races.
+def test_databricks__u2m_oauth__connection_pool(make_config):
+    """Databricks should use a connection per thread with U2M OAuth, sharing a single login.
 
-    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.
-
-    For non-U2M OAuth authentication types (e.g. access_token and M2M OAuth) then
-    ThreadLocalConnectionPool should still be used.
+    If each connection ran its own U2M OAuth flow then the browser flows would race on the
+    CSRF state parameter and cause MismatchingStateError. Sharing a single connection avoids
+    that but forces all queries to run sequentially. Instead, every connection is given the
+    same credentials provider so the browser flow only runs once.
 
     See:
         https://github.com/tobymao/sqlmesh/issues/5646
         https://github.com/SQLMesh/sqlmesh/issues/5858
     """
-    from sqlmesh.utils.connection_pool import ThreadLocalSharedConnectionPool
+    from sqlmesh.utils.connection_pool import ThreadLocalConnectionPool
+    from sqlmesh.utils.databricks import SharedU2MCredentialsProvider
 
     config = make_config(
         type="databricks",
@@ -1450,10 +1449,55 @@ def test_databricks__u2m_oauth__shared_connection_pool(make_config):
         concurrent_tasks=4,
     )
     assert isinstance(config, DatabricksConnectionConfig)
-    assert config.shared_connection is True
+    assert config.shared_connection is False
 
     adapter = config.create_engine_adapter()
-    assert isinstance(adapter._connection_pool, ThreadLocalSharedConnectionPool)
+    assert isinstance(adapter._connection_pool, ThreadLocalConnectionPool)
+
+    connection_kwargs = config._connection_factory_with_kwargs.keywords
+    assert connection_kwargs["auth_type"] == "databricks-oauth"
+    credentials_provider = connection_kwargs["credentials_provider"]
+    assert isinstance(credentials_provider, SharedU2MCredentialsProvider)
+
+    # Separate engine adapters (and copies of the config) must reuse the same login
+    other_config = make_config(
+        type="databricks",
+        server_hostname="dbc-test.cloud.databricks.com",
+        http_path="sql/test/bar",
+        auth_type="databricks-oauth",
+    )
+    assert (
+        other_config._connection_factory_with_kwargs.keywords["credentials_provider"]
+        is credentials_provider
+    )
+
+
+def test_databricks__u2m_oauth__old_connector_shares_connection(make_config):
+    """Connectors without the internals the shared login builds on fall back to one shared connection."""
+    from sqlmesh.utils import databricks
+    from sqlmesh.utils.connection_pool import ThreadLocalSharedConnectionPool
+
+    databricks.supports_shared_u2m_login.cache_clear()
+    try:
+        with patch("importlib.util.find_spec", return_value=None):
+            config = make_config(
+                type="databricks",
+                server_hostname="dbc-test.cloud.databricks.com",
+                http_path="sql/test/foo",
+                auth_type="databricks-oauth",
+                concurrent_tasks=4,
+            )
+            assert isinstance(config, DatabricksConnectionConfig)
+            assert config.shared_connection is True
+
+            adapter = config.create_engine_adapter()
+            assert isinstance(adapter._connection_pool, ThreadLocalSharedConnectionPool)
+
+            connection_kwargs = config._connection_factory_with_kwargs.keywords
+            assert connection_kwargs["auth_type"] == "databricks-oauth"
+            assert "credentials_provider" not in connection_kwargs
+    finally:
+        databricks.supports_shared_u2m_login.cache_clear()
 
 
 @patch.object(DatabricksConnectionConfig, "_connection_factory_with_kwargs")