|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import inspect |
| 4 | +import logging |
| 5 | +import re |
| 6 | +import typing as t |
| 7 | +from functools import lru_cache |
| 8 | +from threading import Lock |
| 9 | + |
| 10 | +logger = logging.getLogger(__name__) |
| 11 | + |
| 12 | +HeaderFactory = t.Callable[[], t.Dict[str, str]] |
| 13 | + |
| 14 | + |
| 15 | +class SharedU2MCredentialsProvider: |
| 16 | + """A Databricks SQL connector `credentials_provider` that shares one U2M OAuth login across connections. |
| 17 | +
|
| 18 | + Without this, every `databricks.sql.connect()` call builds its own OAuth provider. For U2M OAuth that means |
| 19 | + each thread runs its own browser flow, and these race on the CSRF `state` parameter |
| 20 | + (https://github.com/tobymao/sqlmesh/issues/5646). Even if the browser flow is only run once, each connection |
| 21 | + would still hold its own copy of the refresh token. Databricks can issue single-use refresh tokens, so the |
| 22 | + first connection to refresh would invalidate the token held by every other connection. |
| 23 | +
|
| 24 | + Instead, a single connector OAuth provider is created lazily (behind a lock, so the browser flow only runs once) |
| 25 | + and every connection reads its headers from it. Reads are also serialized so only one thread can refresh the |
| 26 | + token at a time. This lets each thread have its own connection while sharing a single login. |
| 27 | +
|
| 28 | + If reading the token fails (for example because the refresh token expired or was revoked) the login is |
| 29 | + discarded, so the next connection or request logs in again rather than failing for the rest of the process. |
| 30 | + Network failures while refreshing keep the login, so the next request retries the refresh without a browser flow. |
| 31 | + """ |
| 32 | + |
| 33 | + def __init__(self, server_hostname: str, auth_type: str): |
| 34 | + self._server_hostname = server_hostname |
| 35 | + self._auth_type = auth_type |
| 36 | + self._auth_provider: t.Optional[t.Any] = None |
| 37 | + self._lock = Lock() |
| 38 | + # Incremented when a login attempt finishes, so threads that waited on a failed attempt can tell it happened |
| 39 | + self._login_attempts_finished = 0 |
| 40 | + self._login_error: t.Optional[Exception] = None |
| 41 | + |
| 42 | + def auth_type(self) -> str: |
| 43 | + return self._auth_type |
| 44 | + |
| 45 | + def __call__(self, *args: t.Any, **kwargs: t.Any) -> HeaderFactory: |
| 46 | + """Called by the connector once per connection. Logs in if needed and returns that connection's header factory.""" |
| 47 | + login_attempts_seen = self._login_attempts_finished |
| 48 | + with self._lock: |
| 49 | + self._ensure_logged_in(login_attempts_seen) |
| 50 | + return self._headers |
| 51 | + |
| 52 | + def _headers(self) -> t.Dict[str, str]: |
| 53 | + """Returns the auth headers for one request, refreshing the token if it has expired.""" |
| 54 | + from databricks.sql.exc import RequestError |
| 55 | + |
| 56 | + headers: t.Dict[str, str] = {} |
| 57 | + login_attempts_seen = self._login_attempts_finished |
| 58 | + with self._lock: |
| 59 | + auth_provider = self._ensure_logged_in(login_attempts_seen) |
| 60 | + try: |
| 61 | + auth_provider.add_headers(headers) |
| 62 | + except RequestError: |
| 63 | + # A network failure while refreshing leaves the login intact, since the connector only replaces its |
| 64 | + # tokens after a successful refresh, so the next request retries the refresh rather than logging in |
| 65 | + raise |
| 66 | + except Exception: |
| 67 | + self._auth_provider = None |
| 68 | + raise |
| 69 | + return headers |
| 70 | + |
| 71 | + def _ensure_logged_in(self, login_attempts_seen: int) -> t.Any: |
| 72 | + """Returns the shared connector auth provider, logging in if needed. Must be called with the lock held.""" |
| 73 | + if self._auth_provider is not None: |
| 74 | + return self._auth_provider |
| 75 | + |
| 76 | + # A login attempt failed while this thread was waiting on the lock. Re-raise its error rather than opening |
| 77 | + # another browser window for every waiting thread; callers that arrive afterwards will try again. |
| 78 | + if self._login_attempts_finished != login_attempts_seen and self._login_error is not None: |
| 79 | + raise self._login_error |
| 80 | + |
| 81 | + try: |
| 82 | + self._auth_provider = self._create_auth_provider() |
| 83 | + self._login_error = None |
| 84 | + except Exception as e: |
| 85 | + self._login_error = e |
| 86 | + raise |
| 87 | + finally: |
| 88 | + self._login_attempts_finished += 1 |
| 89 | + return self._auth_provider |
| 90 | + |
| 91 | + def _create_auth_provider(self) -> t.Any: |
| 92 | + # These are internal to the connector, but using them means the OAuth flow is identical to the one the |
| 93 | + # connector runs when `auth_type` is passed to `connect()` (same client ID, redirect ports, and scopes) |
| 94 | + from databricks.sql import __version__ |
| 95 | + from databricks.sql.auth.auth import get_python_sql_connector_auth_provider |
| 96 | + from databricks.sql.common.unified_http_client import UnifiedHttpClient |
| 97 | + from databricks.sql.utils import build_client_context |
| 98 | + |
| 99 | + # The provider outlives any single connection, so it gets its own HTTP client rather than borrowing one |
| 100 | + # that would be closed along with its connection |
| 101 | + http_client = UnifiedHttpClient( |
| 102 | + build_client_context(self._server_hostname, __version__, user_agent_entry="sqlmesh") |
| 103 | + ) |
| 104 | + return get_python_sql_connector_auth_provider( |
| 105 | + self._server_hostname, http_client=http_client, auth_type=self._auth_type |
| 106 | + ) |
| 107 | + |
| 108 | + |
| 109 | +@lru_cache(maxsize=1) |
| 110 | +def supports_shared_u2m_login() -> bool: |
| 111 | + """Whether the connector internals `SharedU2MCredentialsProvider` builds on accept the arguments it passes them. |
| 112 | +
|
| 113 | + These internals aren't public, so a connector upgrade could change them. If one has, U2M OAuth falls back to |
| 114 | + sharing a single connection rather than failing to connect. |
| 115 | + """ |
| 116 | + try: |
| 117 | + from databricks.sql import __version__ |
| 118 | + except ImportError: |
| 119 | + # The connector isn't installed, so there are no connections to share a login across |
| 120 | + return False |
| 121 | + |
| 122 | + try: |
| 123 | + from databricks.sql.auth.auth import get_python_sql_connector_auth_provider |
| 124 | + from databricks.sql.common.unified_http_client import UnifiedHttpClient |
| 125 | + from databricks.sql.exc import RequestError # noqa: F401 |
| 126 | + from databricks.sql.utils import build_client_context |
| 127 | + |
| 128 | + inspect.signature(build_client_context).bind( |
| 129 | + "hostname", __version__, user_agent_entry="sqlmesh" |
| 130 | + ) |
| 131 | + inspect.signature(UnifiedHttpClient).bind(None) |
| 132 | + inspect.signature(get_python_sql_connector_auth_provider).bind( |
| 133 | + "hostname", http_client=None, auth_type="databricks-oauth" |
| 134 | + ) |
| 135 | + except Exception as e: |
| 136 | + logger.warning( |
| 137 | + "Databricks SQL connector %s is not compatible with SQLMesh's shared U2M OAuth login (%s). U2M OAuth " |
| 138 | + "connections will share a single connection, so queries will run one at a time.", |
| 139 | + __version__, |
| 140 | + e, |
| 141 | + ) |
| 142 | + return False |
| 143 | + return True |
| 144 | + |
| 145 | + |
| 146 | +_shared_u2m_credentials_providers: t.Dict[t.Tuple[str, str], SharedU2MCredentialsProvider] = {} |
| 147 | +_shared_u2m_credentials_providers_lock = Lock() |
| 148 | + |
| 149 | + |
| 150 | +def get_shared_u2m_credentials_provider( |
| 151 | + server_hostname: str, auth_type: str |
| 152 | +) -> SharedU2MCredentialsProvider: |
| 153 | + """Returns the process-wide U2M credentials provider for the given workspace and auth type. |
| 154 | +
|
| 155 | + Engine adapters and config copies are created several times per process, so the provider is kept at module |
| 156 | + level to avoid prompting the user to log in more than once. U2M OAuth is an interactive, single-user login, so |
| 157 | + every connection to a workspace in this process uses the identity of whoever completed the browser flow. |
| 158 | + """ |
| 159 | + # The same workspace can be written with or without a scheme or trailing slash |
| 160 | + normalized_hostname = re.sub(r"^https?://", "", server_hostname.strip(), flags=re.I) |
| 161 | + normalized_hostname = normalized_hostname.rstrip("/").lower() |
| 162 | + key = (normalized_hostname, auth_type) |
| 163 | + with _shared_u2m_credentials_providers_lock: |
| 164 | + if key not in _shared_u2m_credentials_providers: |
| 165 | + _shared_u2m_credentials_providers[key] = SharedU2MCredentialsProvider( |
| 166 | + normalized_hostname, auth_type |
| 167 | + ) |
| 168 | + return _shared_u2m_credentials_providers[key] |
0 commit comments