From c67e4e261a4cdcb27a0831cbc4a0546e9070c6af Mon Sep 17 00:00:00 2001 From: Gabe Pesco Date: Mon, 2 Nov 2026 12:45:46 -0500 Subject: [PATCH 1/3] fix(databricks): run U2M OAuth queries concurrently with a shared login U2M OAuth used a single shared connection so that only one browser flow ran, but the connector serializes every request on a connection, so all queries ran one at a time regardless of concurrent_tasks. Give each thread its own connection again, and pass every connection the same process-wide credentials_provider. It builds one connector OAuth provider behind a lock, so the browser flow still runs once, and serializes header reads so only one thread refreshes the token at a time. Databricks can issue single-use refresh tokens, so each connection holding its own copy would break on the first refresh. Refs #5858, #5646 Co-Authored-By: Claude Opus 5.5 Signed-off-by: Gabe Pesco --- sqlmesh/core/config/connection.py | 17 ++-- sqlmesh/utils/databricks.py | 82 ++++++++++++++++ tests/core/test_connection_config.py | 38 +++++--- tests/utils/test_databricks.py | 134 +++++++++++++++++++++++++++ 4 files changed, 254 insertions(+), 17 deletions(-) create mode 100644 sqlmesh/utils/databricks.py create mode 100644 tests/utils/test_databricks.py diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 7a65e87d72..1b28c5abc2 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -836,12 +836,6 @@ class DatabricksConnectionConfig(ConnectionConfig): _concurrent_tasks_validator = concurrent_tasks_validator _http_headers_validator = http_headers_validator - @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 - @model_validator(mode="before") def _databricks_connect_validator(cls, data: t.Any) -> t.Any: # SQLQueryContextLogger will output any error SQL queries even if they are in a try/except block. @@ -992,6 +986,17 @@ 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` in + # the connector, but `auth_type` is still passed for anything else in the connector that reads it. + from sqlmesh.utils.databricks import get_shared_u2m_credentials_provider + + if t.TYPE_CHECKING: + assert self.server_hostname is not None + + 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 diff --git a/sqlmesh/utils/databricks.py b/sqlmesh/utils/databricks.py new file mode 100644 index 0000000000..448c4ee5e1 --- /dev/null +++ b/sqlmesh/utils/databricks.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import typing as t +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. + """ + + 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() + + def auth_type(self) -> str: + return self._auth_type + + def __call__(self, *args: t.Any, **kwargs: t.Any) -> HeaderFactory: + with self._lock: + if self._auth_provider is None: + self._auth_provider = self._create_auth_provider() + return self._headers + + def _headers(self) -> t.Dict[str, str]: + headers: t.Dict[str, str] = {} + with self._lock: + if t.TYPE_CHECKING: + assert self._auth_provider is not None + self._auth_provider.add_headers(headers) + return headers + + 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 + ) + + +_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. + """ + key = (server_hostname, auth_type) + with _shared_u2m_credentials_providers_lock: + if key not in _shared_u2m_credentials_providers: + _shared_u2m_credentials_providers[key] = SharedU2MCredentialsProvider( + server_hostname, auth_type + ) + return _shared_u2m_credentials_providers[key] diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index c506d401d9..55a5b1b1e4 100644 --- a/tests/core/test_connection_config.py +++ b/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,27 @@ 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 + ) @patch.object(DatabricksConnectionConfig, "_connection_factory_with_kwargs") diff --git a/tests/utils/test_databricks.py b/tests/utils/test_databricks.py new file mode 100644 index 0000000000..a407327201 --- /dev/null +++ b/tests/utils/test_databricks.py @@ -0,0 +1,134 @@ +import time +import typing as t +from concurrent.futures import ThreadPoolExecutor +from threading import Lock +from unittest.mock import patch + +import pytest + +from sqlmesh.utils import databricks +from sqlmesh.utils.databricks import ( + SharedU2MCredentialsProvider, + get_shared_u2m_credentials_provider, +) + +pytest.importorskip("databricks.sql") + + +@pytest.fixture(autouse=True) +def clear_shared_providers(): + databricks._shared_u2m_credentials_providers.clear() + yield + databricks._shared_u2m_credentials_providers.clear() + + +class FakeAuthProvider: + """Records how many threads are reading headers at the same time.""" + + def __init__(self) -> None: + self.active_readers = 0 + self.max_active_readers = 0 + self._counter_lock = Lock() + + def add_headers(self, request_headers: t.Dict[str, str]) -> None: + with self._counter_lock: + self.active_readers += 1 + self.max_active_readers = max(self.max_active_readers, self.active_readers) + # Give other threads a chance to enter if reads are not serialized + time.sleep(0.01) + request_headers["Authorization"] = "Bearer token" + with self._counter_lock: + self.active_readers -= 1 + + +def test_browser_flow_runs_once_across_threads(): + provider = SharedU2MCredentialsProvider("dbc-test.cloud.databricks.com", "databricks-oauth") + fake_auth_provider = FakeAuthProvider() + flow_count = 0 + + def slow_browser_flow() -> FakeAuthProvider: + nonlocal flow_count + flow_count += 1 + time.sleep(0.1) + return fake_auth_provider + + with patch.object(provider, "_create_auth_provider", side_effect=slow_browser_flow): + with ThreadPoolExecutor(max_workers=8) as executor: + header_factories = list(executor.map(lambda _: provider(), range(8))) + + assert flow_count == 1 + assert all(factory() == {"Authorization": "Bearer token"} for factory in header_factories) + + +def test_header_reads_are_serialized(): + provider = SharedU2MCredentialsProvider("dbc-test.cloud.databricks.com", "databricks-oauth") + fake_auth_provider = FakeAuthProvider() + + with patch.object(provider, "_create_auth_provider", return_value=fake_auth_provider): + header_factory = provider() + + with ThreadPoolExecutor(max_workers=8) as executor: + list(executor.map(lambda _: header_factory(), range(32))) + + # Only one thread may read (and so possibly refresh) the token at a time, since refresh tokens can be single-use + assert fake_auth_provider.max_active_readers == 1 + + +def test_failed_browser_flow_is_retried(): + provider = SharedU2MCredentialsProvider("dbc-test.cloud.databricks.com", "databricks-oauth") + fake_auth_provider = FakeAuthProvider() + + with patch.object( + provider, + "_create_auth_provider", + side_effect=[RuntimeError("login cancelled"), fake_auth_provider], + ): + with pytest.raises(RuntimeError, match="login cancelled"): + provider() + assert provider()() == {"Authorization": "Bearer token"} + + +def test_get_shared_u2m_credentials_provider(): + provider = get_shared_u2m_credentials_provider( + "dbc-test.cloud.databricks.com", "databricks-oauth" + ) + assert provider.auth_type() == "databricks-oauth" + assert ( + get_shared_u2m_credentials_provider("dbc-test.cloud.databricks.com", "databricks-oauth") + is provider + ) + assert ( + get_shared_u2m_credentials_provider("dbc-other.cloud.databricks.com", "databricks-oauth") + is not provider + ) + assert ( + get_shared_u2m_credentials_provider("dbc-test.cloud.databricks.com", "azure-oauth") + is not provider + ) + + +def test_uses_connector_oauth_provider(): + """Runs the real connector auth provider with only the browser flow mocked out. + + The provider relies on connector internals, so this fails if a connector upgrade changes them. + """ + import jwt + + server_hostname = "dbc-test.cloud.databricks.com" + access_token = jwt.encode( + {"exp": int(time.time()) + 3600, "iss": f"https://{server_hostname}/oidc"}, + "a-signing-key-that-is-at-least-32-bytes", + algorithm="HS256", + ) + provider = SharedU2MCredentialsProvider(server_hostname, "databricks-oauth") + + with patch( + "databricks.sql.auth.oauth.OAuthManager.get_tokens", + return_value=(access_token, "refresh-token"), + ) as mock_get_tokens: + header_factory = provider() + provider() + + mock_get_tokens.assert_called_once() + assert mock_get_tokens.call_args.kwargs["hostname"] == f"https://{server_hostname}/" + assert header_factory() == {"Authorization": f"Bearer {access_token}"} From b8bdb473cf800a59a8d03276adf25fd12d35a5b7 Mon Sep 17 00:00:00 2001 From: Gabe Pesco Date: Mon, 2 Nov 2026 14:32:58 -0500 Subject: [PATCH 2/3] fix(databricks): recover from failed U2M logins and support older connectors - If reading the token fails (e.g. the refresh token expired or was revoked), discard the shared login so the next request logs in again instead of failing for the rest of the process. - Threads waiting on a login that fails get its error rather than each opening another browser window. - Normalise the hostname so the same workspace written with a scheme, trailing slash or different case shares one login. - Fall back to a single shared connection on connectors without the internals the shared login needs (< 4.1.0). - Test that the connector prefers the shared credentials_provider over auth_type when both are passed, as connect() does. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Gabe Pesco --- sqlmesh/core/config/connection.py | 33 ++++++--- sqlmesh/utils/databricks.py | 59 +++++++++++++-- tests/core/test_connection_config.py | 28 ++++++++ tests/utils/test_databricks.py | 103 +++++++++++++++++++++++++++ 4 files changed, 207 insertions(+), 16 deletions(-) diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 1b28c5abc2..63de7986a0 100644 --- a/sqlmesh/core/config/connection.py +++ b/sqlmesh/core/config/connection.py @@ -836,6 +836,19 @@ class DatabricksConnectionConfig(ConnectionConfig): _concurrent_tasks_validator = concurrent_tasks_validator _http_headers_validator = http_headers_validator + @computed_field # type: ignore[prop-decorator] + @property + def shared_connection(self) -> bool: + """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: # SQLQueryContextLogger will output any error SQL queries even if they are in a try/except block. @@ -987,16 +1000,18 @@ def _static_connection_kwargs(self) -> t.Dict[str, t.Any]: # 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` in - # the connector, but `auth_type` is still passed for anything else in the connector that reads it. - from sqlmesh.utils.databricks import get_shared_u2m_credentials_provider - - if t.TYPE_CHECKING: - assert self.server_hostname is not None - - conn_kwargs["credentials_provider"] = get_shared_u2m_credentials_provider( - self.server_hostname, self.auth_type + # 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 diff --git a/sqlmesh/utils/databricks.py b/sqlmesh/utils/databricks.py index 448c4ee5e1..ae1ffbdfd6 100644 --- a/sqlmesh/utils/databricks.py +++ b/sqlmesh/utils/databricks.py @@ -1,6 +1,9 @@ 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]] @@ -18,6 +21,9 @@ class SharedU2MCredentialsProvider: 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. """ def __init__(self, server_hostname: str, auth_type: str): @@ -25,24 +31,51 @@ def __init__(self, server_hostname: str, auth_type: str): 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: + login_attempts_seen = self._login_attempts_finished with self._lock: - if self._auth_provider is None: - self._auth_provider = self._create_auth_provider() + self._ensure_logged_in(login_attempts_seen) return self._headers def _headers(self) -> t.Dict[str, str]: headers: t.Dict[str, str] = {} + login_attempts_seen = self._login_attempts_finished with self._lock: - if t.TYPE_CHECKING: - assert self._auth_provider is not None - self._auth_provider.add_headers(headers) + auth_provider = self._ensure_logged_in(login_attempts_seen) + try: + auth_provider.add_headers(headers) + 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) @@ -61,6 +94,15 @@ def _create_auth_provider(self) -> t.Any: ) +@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() @@ -73,10 +115,13 @@ def get_shared_u2m_credentials_provider( 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. """ - key = (server_hostname, auth_type) + # 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( - server_hostname, auth_type + normalized_hostname, auth_type ) return _shared_u2m_credentials_providers[key] diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index 55a5b1b1e4..bf00f78bee 100644 --- a/tests/core/test_connection_config.py +++ b/tests/core/test_connection_config.py @@ -1472,6 +1472,34 @@ def test_databricks__u2m_oauth__connection_pool(make_config): ) +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") def test_databricks__m2m_oauth__connection_pool(mock_connection_factory_with_kwargs, make_config): from sqlmesh.utils.connection_pool import ThreadLocalConnectionPool diff --git a/tests/utils/test_databricks.py b/tests/utils/test_databricks.py index a407327201..ea83daef11 100644 --- a/tests/utils/test_databricks.py +++ b/tests/utils/test_databricks.py @@ -88,6 +88,61 @@ def test_failed_browser_flow_is_retried(): assert provider()() == {"Authorization": "Bearer token"} +def test_threads_waiting_on_a_failed_login_do_not_retry_it(): + provider = SharedU2MCredentialsProvider("dbc-test.cloud.databricks.com", "databricks-oauth") + fake_auth_provider = FakeAuthProvider() + flow_count = 0 + + def failing_browser_flow() -> FakeAuthProvider: + nonlocal flow_count + flow_count += 1 + time.sleep(0.1) + raise RuntimeError("login cancelled") + + def call_provider(_: int) -> t.Optional[Exception]: + try: + provider() + except Exception as e: + return e + return None + + with patch.object(provider, "_create_auth_provider", side_effect=failing_browser_flow): + with ThreadPoolExecutor(max_workers=8) as executor: + errors = list(executor.map(call_provider, range(8))) + + # Every waiting thread gets the error from the one browser flow instead of opening its own + assert flow_count == 1 + assert all(isinstance(error, RuntimeError) for error in errors) + + # A later caller still gets to try again + with patch.object(provider, "_create_auth_provider", return_value=fake_auth_provider): + assert provider()() == {"Authorization": "Bearer token"} + + +def test_failed_token_read_discards_login(): + provider = SharedU2MCredentialsProvider("dbc-test.cloud.databricks.com", "databricks-oauth") + fake_auth_provider = FakeAuthProvider() + + class RevokedAuthProvider: + def add_headers(self, request_headers: t.Dict[str, str]) -> None: + raise RuntimeError("refresh token revoked") + + with patch.object( + provider, + "_create_auth_provider", + side_effect=[RevokedAuthProvider(), fake_auth_provider], + ) as mock_create_auth_provider: + header_factory = provider() + with pytest.raises(RuntimeError, match="refresh token revoked"): + header_factory() + + # Both existing connections and new ones log in again instead of failing for the rest of the process + assert header_factory() == {"Authorization": "Bearer token"} + assert provider()() == {"Authorization": "Bearer token"} + + assert mock_create_auth_provider.call_count == 2 + + def test_get_shared_u2m_credentials_provider(): provider = get_shared_u2m_credentials_provider( "dbc-test.cloud.databricks.com", "databricks-oauth" @@ -106,6 +161,14 @@ def test_get_shared_u2m_credentials_provider(): is not provider ) + # The same workspace written with a scheme, trailing slash, or different case shares the login + for hostname in ( + "https://dbc-test.cloud.databricks.com", + "dbc-test.cloud.databricks.com/", + "HTTPS://DBC-TEST.cloud.databricks.com/", + ): + assert get_shared_u2m_credentials_provider(hostname, "databricks-oauth") is provider + def test_uses_connector_oauth_provider(): """Runs the real connector auth provider with only the browser flow mocked out. @@ -132,3 +195,43 @@ def test_uses_connector_oauth_provider(): mock_get_tokens.assert_called_once() assert mock_get_tokens.call_args.kwargs["hostname"] == f"https://{server_hostname}/" assert header_factory() == {"Authorization": f"Bearer {access_token}"} + + +def test_connector_prefers_shared_credentials_provider(): + """Builds per-connection auth providers the way `databricks.sql.connect()` does, passing both `auth_type` and + the shared `credentials_provider`, and checks the connector uses the shared login rather than its own.""" + import jwt + 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 + + server_hostname = "dbc-test.cloud.databricks.com" + access_token = jwt.encode( + {"exp": int(time.time()) + 3600, "iss": f"https://{server_hostname}/oidc"}, + "a-signing-key-that-is-at-least-32-bytes", + algorithm="HS256", + ) + provider = get_shared_u2m_credentials_provider(server_hostname, "databricks-oauth") + + def connection_auth_headers(_: int) -> t.Dict[str, str]: + http_client = UnifiedHttpClient(build_client_context(server_hostname, __version__)) + connection_auth_provider = get_python_sql_connector_auth_provider( + server_hostname, + http_client=http_client, + auth_type="databricks-oauth", + credentials_provider=provider, + ) + headers: t.Dict[str, str] = {} + connection_auth_provider.add_headers(headers) + return headers + + with patch( + "databricks.sql.auth.oauth.OAuthManager.get_tokens", + return_value=(access_token, "refresh-token"), + ) as mock_get_tokens: + with ThreadPoolExecutor(max_workers=8) as executor: + all_headers = list(executor.map(connection_auth_headers, range(8))) + + mock_get_tokens.assert_called_once() + assert all(headers["Authorization"] == f"Bearer {access_token}" for headers in all_headers) From dc9adece012601ff6a111dc378ef97433f2495d8 Mon Sep 17 00:00:00 2001 From: Gabe Pesco Date: Tue, 3 Nov 2026 12:04:18 -0500 Subject: [PATCH 3/3] fix(databricks): keep the U2M login on network failures during refresh A network failure while refreshing the access token surfaces as the connector's RequestError. The connector only replaces its tokens after a successful refresh, so the login is still valid and the next request can retry the refresh. Previously any failure discarded the login, forcing a new browser flow after a transient blip. Other failures, such as a rejected refresh token, still discard it. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Gabe Pesco --- sqlmesh/utils/databricks.py | 12 +++++++++++- tests/utils/test_databricks.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/sqlmesh/utils/databricks.py b/sqlmesh/utils/databricks.py index ae1ffbdfd6..639e3fa6e0 100644 --- a/sqlmesh/utils/databricks.py +++ b/sqlmesh/utils/databricks.py @@ -24,6 +24,7 @@ class SharedU2MCredentialsProvider: 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): @@ -39,18 +40,26 @@ 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 @@ -113,7 +122,8 @@ def get_shared_u2m_credentials_provider( """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. + 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) diff --git a/tests/utils/test_databricks.py b/tests/utils/test_databricks.py index ea83daef11..b0bd6a9eb7 100644 --- a/tests/utils/test_databricks.py +++ b/tests/utils/test_databricks.py @@ -143,6 +143,34 @@ def add_headers(self, request_headers: t.Dict[str, str]) -> None: assert mock_create_auth_provider.call_count == 2 +def test_network_failure_during_token_read_keeps_login(): + from databricks.sql.exc import RequestError + + provider = SharedU2MCredentialsProvider("dbc-test.cloud.databricks.com", "databricks-oauth") + fake_auth_provider = FakeAuthProvider() + refresh_attempts = 0 + + class FlakyAuthProvider: + def add_headers(self, request_headers: t.Dict[str, str]) -> None: + nonlocal refresh_attempts + refresh_attempts += 1 + if refresh_attempts == 1: + raise RequestError("HTTP request failed: connection reset") + fake_auth_provider.add_headers(request_headers) + + with patch.object( + provider, "_create_auth_provider", return_value=FlakyAuthProvider() + ) as mock_create_auth_provider: + header_factory = provider() + with pytest.raises(RequestError): + header_factory() + + # The refresh is retried on the next request without another browser flow + assert header_factory() == {"Authorization": "Bearer token"} + + assert mock_create_auth_provider.call_count == 1 + + def test_get_shared_u2m_credentials_provider(): provider = get_shared_u2m_credentials_provider( "dbc-test.cloud.databricks.com", "databricks-oauth"