diff --git a/sqlmesh/core/config/connection.py b/sqlmesh/core/config/connection.py index 7a65e87d72..63de7986a0 100644 --- a/sqlmesh/core/config/connection.py +++ b/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 diff --git a/sqlmesh/utils/databricks.py b/sqlmesh/utils/databricks.py new file mode 100644 index 0000000000..639e3fa6e0 --- /dev/null +++ b/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] diff --git a/tests/core/test_connection_config.py b/tests/core/test_connection_config.py index c506d401d9..bf00f78bee 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,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") diff --git a/tests/utils/test_databricks.py b/tests/utils/test_databricks.py new file mode 100644 index 0000000000..b0bd6a9eb7 --- /dev/null +++ b/tests/utils/test_databricks.py @@ -0,0 +1,265 @@ +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_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_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" + ) + 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 + ) + + # 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. + + 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}"} + + +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)