| FazBrowse GitHub Viewer | Trending | | Home |
| Tools: [Download Repo ZIP] [Original HTTPS Page] |
1 parent 9259644 commit 2572362
2 files changed
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1,25 +1,19 @@ | |||
| 1 | 1 | """Translate the connector's ``AuthProvider`` into ``databricks_sql_kernel`` | |
| 2 | 2 | ``Session`` auth kwargs. | |
| 3 | 3 | ||
| 4 | - The connector already implements every auth flow it supports (PAT, | ||
| 5 | - OAuth M2M, OAuth U2M, external token providers, federation). The | ||
| 6 | - kernel must not re-implement them. Decision D9 in the integration | ||
| 7 | - design: PAT goes through the kernel's PAT path; everything else | ||
| 8 | - delegates back to the connector via the kernel's ``External`` | ||
| 9 | - trampoline, with a Python callback that returns a fresh bearer | ||
| 10 | - token. | ||
| 4 | + This phase ships PAT only. The kernel-side PyO3 binding accepts | ||
| 5 | + ``auth_type='pat'``; OAuth / federation / custom credentials | ||
| 6 | + providers are reserved but not yet wired in either layer. Non-PAT | ||
| 7 | + auth raises ``NotSupportedError`` from this bridge so the failure | ||
| 8 | + surfaces at session-open time with a clear message rather than | ||
| 9 | + deep inside the kernel. | ||
| 11 | 10 | ||
| 12 | 11 | Token extraction goes through ``AuthProvider.add_headers({})`` | |
| 13 | 12 | rather than touching auth-provider-specific attributes, so the | |
| 14 | - bridge works for every subclass — including custom providers a | ||
| 15 | - caller may have wired in. | ||
| 16 | - | ||
| 17 | - End-to-end limitation: the kernel's | ||
| 18 | - ``build_auth_provider`` currently rejects ``AuthConfig::External`` | ||
| 19 | - ("reserved; v0 wires PAT + OAuthM2M + OAuthU2M only"). Until the | ||
| 20 | - kernel-side follow-up PR lands, non-PAT auth surfaces a clear | ||
| 21 | - ``KernelError(code='InvalidArgument', message='AuthConfig::External | ||
| 22 | - is reserved...')`` from ``Session.open_session``. PAT works today. | ||
| 13 | + bridge works uniformly for every PAT shape — including | ||
| 14 | + ``AccessTokenAuthProvider`` wrapped in ``TokenFederationProvider`` | ||
| 15 | + (which ``get_python_sql_connector_auth_provider`` does for every | ||
| 16 | + provider it builds). | ||
| 23 | 17 | """ | |
| 24 | 18 | ||
| 25 | 19 | from __future__ import annotations | |
@@ -29,6 +23,7 @@ | |||
| 29 | 23 | ||
| 30 | 24 | from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider | |
| 31 | 25 | from databricks.sql.auth.token_federation import TokenFederationProvider | |
| 26 | + from databricks.sql.exc import NotSupportedError | ||
| 32 | 27 | ||
| 33 | 28 | logger = logging.getLogger(__name__) | |
| 34 | 29 | ||
@@ -64,8 +59,8 @@ def _extract_bearer_token(auth_provider: AuthProvider) -> Optional[str]: | |||
| 64 | 59 | provider-specific internals. | |
| 65 | 60 | ||
| 66 | 61 | Returns ``None`` if the provider did not write an Authorization | |
| 67 | - header or wrote a non-Bearer scheme — neither shape is | ||
| 68 | - representable in the kernel's auth surface today. | ||
| 62 | + header or wrote a non-Bearer scheme — neither is representable | ||
| 63 | + in the kernel's PAT auth surface. | ||
| 69 | 64 | """ | |
| 70 | 65 | headers: Dict[str, str] = {} | |
| 71 | 66 | auth_provider.add_headers(headers) | |
@@ -80,29 +75,13 @@ def _extract_bearer_token(auth_provider: AuthProvider) -> Optional[str]: | |||
| 80 | 75 | def kernel_auth_kwargs(auth_provider: AuthProvider) -> Dict[str, Any]: | |
| 81 | 76 | """Build the kwargs passed to ``databricks_sql_kernel.Session(...)``. | |
| 82 | 77 | ||
| 83 | - Two routing decisions: | ||
| 84 | - | ||
| 85 | - 1. ``AccessTokenAuthProvider`` → ``auth_type='pat'`` with the | ||
| 86 | - static token. Kernel uses it verbatim for every request. | ||
| 87 | - 2. Anything else → ``auth_type='external'`` with a callback that | ||
| 88 | - calls ``auth_provider.add_headers({})`` and returns the | ||
| 89 | - fresh bearer token. The connector keeps owning the OAuth / | ||
| 90 | - MSAL / federation flow; the kernel asks for a token whenever | ||
| 91 | - it needs one. | ||
| 92 | - | ||
| 93 | - The PAT special-case exists because it's the only path the | ||
| 94 | - kernel actually serves end-to-end today. Once the kernel-side | ||
| 95 | - External enablement lands, PAT could collapse into the | ||
| 96 | - External path too (one callback that returns the static token); | ||
| 97 | - but keeping the explicit ``pat`` route means the kernel does | ||
| 98 | - not pay the GIL-reacquire cost on every HTTP request for PAT | ||
| 99 | - users. | ||
| 78 | + PAT (including ``TokenFederationProvider``-wrapped PAT) routes | ||
| 79 | + through the kernel's PAT path. Anything else raises | ||
| 80 | + ``NotSupportedError`` — the kernel binding doesn't accept OAuth | ||
| 81 | + today, and routing OAuth through PAT would silently break | ||
| 82 | + token refresh during long-running sessions. | ||
| 100 | 83 | """ | |
| 101 | 84 | if _is_pat(auth_provider): | |
| 102 | - # PAT case: pull the static token out and feed the kernel's | ||
| 103 | - # PAT path. We go through ``add_headers`` regardless of | ||
| 104 | - # whether the provider was wrapped in TokenFederation or | ||
| 105 | - # not — both shapes write the same Authorization header. | ||
| 106 | 85 | token = _extract_bearer_token(auth_provider) | |
| 107 | 86 | if not token: | |
| 108 | 87 | raise ValueError( | |
@@ -111,21 +90,8 @@ def kernel_auth_kwargs(auth_provider: AuthProvider) -> Dict[str, Any]: | |||
| 111 | 90 | ) | |
| 112 | 91 | return {"auth_type": "pat", "access_token": token} | |
| 113 | 92 | ||
| 114 | - # Every other provider: trampoline a callback. The callback is | ||
| 115 | - # invoked once per HTTP request that needs auth (the kernel does | ||
| 116 | - # not cache the returned token), so the auth_provider's own | ||
| 117 | - # caching is what keeps this fast. | ||
| 118 | - def token_callback() -> str: | ||
| 119 | - token = _extract_bearer_token(auth_provider) | ||
| 120 | - if not token: | ||
| 121 | - raise RuntimeError( | ||
| 122 | - f"{type(auth_provider).__name__}.add_headers did not produce " | ||
| 123 | - "a Bearer Authorization header; cannot supply a token to the kernel" | ||
| 124 | - ) | ||
| 125 | - return token | ||
| 126 | - | ||
| 127 | - logger.debug( | ||
| 128 | - "Routing %s through kernel External trampoline", | ||
| 129 | - type(auth_provider).__name__, | ||
| 93 | + raise NotSupportedError( | ||
| 94 | + f"The kernel backend (use_sea=True) currently only supports PAT auth, " | ||
| 95 | + f"but got {type(auth_provider).__name__}. Use use_sea=False (Thrift) " | ||
| 96 | + "for OAuth / federation / custom credential providers." | ||
| 130 | 97 | ) | |
| 131 | - return {"auth_type": "external", "token_callback": token_callback} | ||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
@@ -1,9 +1,12 @@ | |||
| 1 | 1 | """Unit tests for the kernel backend's auth bridge. | |
| 2 | 2 | ||
| 3 | - The bridge translates the connector's ``AuthProvider`` hierarchy | ||
| 4 | - into ``databricks_sql_kernel.Session`` auth kwargs. PAT goes through | ||
| 5 | - the kernel's PAT path; everything else trampolines through the | ||
| 6 | - ``External`` path with a Python callback. | ||
| 3 | + Phase 1 ships PAT only. Tests verify: | ||
| 4 | + - PAT routes through ``auth_type='pat'``. | ||
| 5 | + - ``TokenFederationProvider``-wrapped PAT also routes through | ||
| 6 | + PAT (every provider built by ``get_python_sql_connector_auth_provider`` | ||
| 7 | + is federation-wrapped, so the naive isinstance check has to | ||
| 8 | + look through the wrapper). | ||
| 9 | + - Anything else raises ``NotSupportedError`` with a clear message. | ||
| 7 | 10 | """ | |
| 8 | 11 | ||
| 9 | 12 | from __future__ import annotations | |
@@ -12,38 +15,36 @@ | |||
| 12 | 15 | ||
| 13 | 16 | import pytest | |
| 14 | 17 | ||
| 15 | - from databricks.sql.auth.authenticators import AccessTokenAuthProvider, AuthProvider | ||
| 18 | + from databricks.sql.auth.authenticators import ( | ||
| 19 | + AccessTokenAuthProvider, | ||
| 20 | + AuthProvider, | ||
| 21 | + DatabricksOAuthProvider, | ||
| 22 | + ExternalAuthProvider, | ||
| 23 | + ) | ||
| 16 | 24 | from databricks.sql.backend.kernel.auth_bridge import ( | |
| 17 | 25 | _extract_bearer_token, | |
| 18 | 26 | kernel_auth_kwargs, | |
| 19 | 27 | ) | |
| 28 | + from databricks.sql.exc import NotSupportedError | ||
| 20 | 29 | ||
| 21 | 30 | ||
| 22 | 31 | class _FakeOAuthProvider(AuthProvider): | |
| 23 | - """Stand-in for OAuth/MSAL/federation providers — anything that | ||
| 24 | - isn't ``AccessTokenAuthProvider``. Returns a counter-stamped | ||
| 25 | - token so tests can prove the callback is invoked each call.""" | ||
| 26 | - | ||
| 27 | - def __init__(self): | ||
| 28 | - self.calls = 0 | ||
| 32 | + """Stand-in for any non-PAT provider. The bridge should reject | ||
| 33 | + these with NotSupportedError.""" | ||
| 29 | 34 | ||
| 30 | 35 | def add_headers(self, request_headers): | |
| 31 | - self.calls += 1 | ||
| 32 | - request_headers["Authorization"] = f"Bearer token-{self.calls}" | ||
| 36 | + request_headers["Authorization"] = "Bearer oauth-token-xyz" | ||
| 33 | 37 | ||
| 34 | 38 | ||
| 35 | 39 | class _MalformedProvider(AuthProvider): | |
| 36 | - """Provider that returns a non-Bearer Authorization header | ||
| 37 | - (e.g. Basic auth). The bridge should reject this rather than | ||
| 38 | - silently sending the wrong shape to the kernel.""" | ||
| 40 | + """Provider that returns a non-Bearer Authorization header.""" | ||
| 39 | 41 | ||
| 40 | 42 | def add_headers(self, request_headers): | |
| 41 | 43 | request_headers["Authorization"] = "Basic dXNlcjpwYXNz" | |
| 42 | 44 | ||
| 43 | 45 | ||
| 44 | 46 | class _SilentProvider(AuthProvider): | |
| 45 | - """Provider that writes nothing — represents misconfigured | ||
| 46 | - auth or a placeholder. The bridge must surface this clearly.""" | ||
| 47 | + """Provider that writes nothing — misconfigured auth.""" | ||
| 47 | 48 | ||
| 48 | 49 | def add_headers(self, request_headers): | |
| 49 | 50 | pass | |
@@ -74,43 +75,53 @@ def test_federation_wrapped_pat_routes_to_kernel_pat(self): | |||
| 74 | 75 | underlying ``AccessTokenAuthProvider``.""" | |
| 75 | 76 | from databricks.sql.auth.token_federation import TokenFederationProvider | |
| 76 | 77 | ||
| 77 | - # TokenFederationProvider needs an http_client; a MagicMock | ||
| 78 | - # is sufficient because we don't trigger any token exchange | ||
| 79 | - # in the test (the cached-token path is never hit). | ||
| 80 | 78 | base = AccessTokenAuthProvider("dapi-abc") | |
| 79 | + # TokenFederationProvider's __init__ requires an http_client | ||
| 80 | + # to construct cleanly; for this unit test we only exercise | ||
| 81 | + # the add_headers passthrough + the external_provider | ||
| 82 | + # attribute. Bypass __init__ with __new__ and stash just | ||
| 83 | + # the fields the bridge touches. | ||
| 81 | 84 | federated = TokenFederationProvider.__new__(TokenFederationProvider) | |
| 82 | 85 | federated.external_provider = base | |
| 83 | - # The bridge only touches `add_headers` (delegated to the | ||
| 84 | - # base) and `external_provider`. Other attrs would be set | ||
| 85 | - # by __init__ but aren't exercised here. | ||
| 86 | 86 | federated.add_headers = base.add_headers | |
| 87 | 87 | kwargs = kernel_auth_kwargs(federated) | |
| 88 | 88 | assert kwargs == {"auth_type": "pat", "access_token": "dapi-abc"} | |
| 89 | 89 | ||
| 90 | - def test_pat_with_silent_provider_raises(self): | ||
| 90 | + def test_pat_with_silent_provider_raises_value_error(self): | ||
| 91 | 91 | """An AccessTokenAuthProvider that produces no Authorization | |
| 92 | 92 | header is misconfigured; surface that at bridge-build time, | |
| 93 | 93 | not on the first kernel HTTP request.""" | |
| 94 | 94 | broken = AccessTokenAuthProvider("dapi-x") | |
| 95 | - # Force the broken state by monkey-patching add_headers. | ||
| 96 | 95 | broken.add_headers = lambda h: None # type: ignore[method-assign] | |
| 97 | 96 | with pytest.raises(ValueError, match="Bearer"): | |
| 98 | 97 | kernel_auth_kwargs(broken) | |
| 99 | 98 | ||
| 100 | - def test_oauth_routes_to_external_trampoline(self): | ||
| 101 | - provider = _FakeOAuthProvider() | ||
| 102 | - kwargs = kernel_auth_kwargs(provider) | ||
| 103 | - assert kwargs["auth_type"] == "external" | ||
| 104 | - callback = kwargs["token_callback"] | ||
| 105 | - assert callable(callback) | ||
| 106 | - # First call -> token-1, second call -> token-2. Proves the | ||
| 107 | - # callback delegates to the live auth_provider each time | ||
| 108 | - # rather than caching. | ||
| 109 | - assert callback() == "token-1" | ||
| 110 | - assert callback() == "token-2" | ||
| 111 | - assert provider.calls == 2 | ||
| 112 | - | ||
| 113 | - def test_external_callback_raises_on_missing_header(self): | ||
| 114 | - kwargs = kernel_auth_kwargs(_SilentProvider()) | ||
| 115 | - with pytest.raises(RuntimeError, match="Bearer"): | ||
| 116 | - kwargs["token_callback"]() | ||
| 99 | + def test_generic_oauth_provider_raises_not_supported(self): | ||
| 100 | + with pytest.raises(NotSupportedError, match="only supports PAT"): | ||
| 101 | + kernel_auth_kwargs(_FakeOAuthProvider()) | ||
| 102 | + | ||
| 103 | + def test_external_credentials_provider_raises_not_supported(self): | ||
| 104 | + """``ExternalAuthProvider`` wraps user-supplied | ||
| 105 | + credentials_provider — kernel doesn't accept these today, | ||
| 106 | + and the bridge surfaces that explicitly.""" | ||
| 107 | + # ExternalAuthProvider's __init__ calls the credentials | ||
| 108 | + # provider; supply a noop one. | ||
| 109 | + from databricks.sql.auth.authenticators import CredentialsProvider | ||
| 110 | + | ||
| 111 | + class _NoopCreds(CredentialsProvider): | ||
| 112 | + def auth_type(self): | ||
| 113 | + return "noop" | ||
| 114 | + | ||
| 115 | + def __call__(self, *args, **kwargs): | ||
| 116 | + return lambda: {"Authorization": "Bearer noop"} | ||
| 117 | + | ||
| 118 | + ext = ExternalAuthProvider(_NoopCreds()) | ||
| 119 | + with pytest.raises(NotSupportedError, match="only supports PAT"): | ||
| 120 | + kernel_auth_kwargs(ext) | ||
| 121 | + | ||
| 122 | + def test_silent_non_pat_provider_also_raises_not_supported(self): | ||
| 123 | + """Even if a non-PAT provider produces no header, the bridge | ||
| 124 | + rejects the type itself — we don't try to extract a token | ||
| 125 | + from something we already know is unsupported.""" | ||
| 126 | + with pytest.raises(NotSupportedError): | ||
| 127 | + kernel_auth_kwargs(_SilentProvider()) | ||
| Back | FazBrowse Home | New Git URL |
0 commit comments