diff --git a/src/ucode/custom_oauth.py b/src/ucode/custom_oauth.py index c6b8ac7e..1d7d2a7d 100644 --- a/src/ucode/custom_oauth.py +++ b/src/ucode/custom_oauth.py @@ -2,6 +2,9 @@ from __future__ import annotations +import hashlib +import json +import os import platform import shlex import subprocess @@ -13,6 +16,7 @@ from databricks.sdk import oauth +from ucode.config_io import APP_DIR from ucode.constants import LOCALHOST, LOOPBACK_HOST from ucode.databricks import build_auth_token_argv from ucode.ui import err_console, normalize_workspace_url, print_warning_err @@ -20,6 +24,8 @@ DEFAULT_REDIRECT_URL = f"http://{LOCALHOST}:8020" # Custom OAuth may need a human to finish browser consent, not just a token fetch. CUSTOM_OAUTH_TIMEOUT_MS = 180_000 +CUSTOM_OAUTH_CLI_VERSION = (1, 17, 0) +CUSTOM_OAUTH_CONFIG_FILE = APP_DIR / "ug.databrickscfg" class CustomOAuthConfig(TypedDict): @@ -93,6 +99,11 @@ def build_custom_auth_shell_command(workspace: str, config: CustomOAuthConfig) - return shlex.join(argv) +def _custom_oauth_profile(workspace: str, client_id: str, scopes: Sequence[str]) -> str: + key = "\0".join((workspace, client_id, *scopes)).encode() + return f"ug-custom-oauth-{hashlib.sha256(key).hexdigest()[:12]}" + + @contextmanager def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]: """Serialize helpers sharing a callback port with a POSIX file lock. @@ -112,17 +123,100 @@ def _custom_oauth_lock(cache_dir: Path, redirect_url: str) -> Iterator[None]: fcntl.flock(lock_file, fcntl.LOCK_UN) -def get_custom_client_token( +def _require_custom_oauth_cli() -> None: + from ucode.databricks import databricks_cli_version + + version = databricks_cli_version() + if version is None or version < CUSTOM_OAUTH_CLI_VERSION: + required = ".".join(map(str, CUSTOM_OAUTH_CLI_VERSION)) + raise RuntimeError( + f"Custom-client OAuth requires Databricks CLI v{required} or newer. Upgrade the CLI " + "and retry." + ) + + +def _token_from_cli(workspace: str, profile: str, env: dict[str, str], force: bool) -> str: + command = [ + "databricks", + "auth", + "token", + "--host", + workspace, + "--profile", + profile, + "--output", + "json", + ] + if force: + command.append("--force-refresh") + result = subprocess.run(command, capture_output=True, text=True, env=env, timeout=15) + if result.returncode != 0: + return "" + try: + return json.loads(result.stdout or "{}").get("access_token", "") + except json.JSONDecodeError: + return "" + + +def _get_custom_client_token_from_cli( workspace: str, - client_id: str, - redirect_url: str = DEFAULT_REDIRECT_URL, - *, - scopes: Sequence[str], - force_refresh: bool = False, + config: CustomOAuthConfig, + force_refresh: bool, +) -> str: + """Delegate custom-client U2M login, refresh, and caching to Databricks CLI.""" + _require_custom_oauth_cli() + profile = _custom_oauth_profile(workspace, config["client_id"], config["scopes"]) + CUSTOM_OAUTH_CONFIG_FILE.parent.mkdir(parents=True, exist_ok=True) + env = os.environ.copy() + env["DATABRICKS_CONFIG_FILE"] = str(CUSTOM_OAUTH_CONFIG_FILE) + + try: + token = _token_from_cli(workspace, profile, env, force_refresh) + if token: + return token + login = subprocess.run( + [ + "databricks", + "auth", + "login", + "--host", + workspace, + "--profile", + profile, + "--client-id", + config["client_id"], + "--scopes", + ",".join(config["scopes"]), + "--timeout", + "3m", + ], + capture_output=True, + text=True, + check=False, + env=env, + timeout=CUSTOM_OAUTH_TIMEOUT_MS / 1000, + ) + if login.returncode == 0: + token = _token_from_cli(workspace, profile, env, False) + if token: + return token + except (OSError, subprocess.TimeoutExpired) as exc: + raise RuntimeError( + "Databricks CLI custom-client OAuth failed. Check the workspace, client ID, and " + "scopes, then retry." + ) from exc + raise RuntimeError( + "Databricks CLI returned no custom-client OAuth token. Check the workspace, client ID, " + "and scopes, then retry." + ) + + +def _get_custom_client_token_from_sdk( + workspace: str, + config: CustomOAuthConfig, + force_refresh: bool, ) -> str: """Reuse the SDK's PKCE flow and per-workspace/client token cache.""" - config = create_custom_oauth_config(client_id, scopes, redirect_url) - workspace = normalize_workspace_url(workspace) try: endpoints = oauth.get_workspace_endpoints(workspace) cache = oauth.TokenCache( @@ -133,8 +227,6 @@ def get_custom_client_token( scopes=config["scopes"], ) with _custom_oauth_lock(Path(cache.filename).parent, config["redirect_url"]): - # Read only after acquiring the lock: another helper may have just - # completed login or rotated the refresh token while we waited. credentials = cache.load() if credentials is not None: try: @@ -174,3 +266,19 @@ def get_custom_client_token( f"redirect URL ({config['redirect_url']}); ensure its local port is available and " "the SDK token cache is writable, then retry." ) from exc + + +def get_custom_client_token( + workspace: str, + client_id: str, + redirect_url: str = DEFAULT_REDIRECT_URL, + *, + scopes: Sequence[str], + force_refresh: bool = False, +) -> str: + """Use CLI-managed OAuth profiles when enabled; otherwise retain the SDK flow.""" + config = create_custom_oauth_config(client_id, scopes, redirect_url) + workspace = normalize_workspace_url(workspace) + if os.environ.get("CUSTOM_OAUTH_CONFIG_FILE") == "1": + return _get_custom_client_token_from_cli(workspace, config, force_refresh) + return _get_custom_client_token_from_sdk(workspace, config, force_refresh) diff --git a/tests/test_custom_oauth.py b/tests/test_custom_oauth.py index ed0d945a..ca29fe14 100644 --- a/tests/test_custom_oauth.py +++ b/tests/test_custom_oauth.py @@ -2,140 +2,72 @@ from __future__ import annotations -import json -from datetime import UTC, datetime, timedelta -from pathlib import Path +import subprocess from unittest.mock import Mock, patch -from urllib.parse import parse_qs import pytest -from databricks.sdk import oauth from typer.testing import CliRunner import ucode.cli as cli_mod +import ucode.custom_oauth as oauth_mod import ucode.databricks as db_mod from ucode.cli import app -from ucode.custom_oauth import _custom_oauth_lock, get_custom_client_token +from ucode.custom_oauth import get_custom_client_token WS = "https://example.databricks.com" TEST_SCOPES = ("offline_access", "catalog.catalogs:read") runner = CliRunner() -class TestCustomOAuthLock: - def test_releases_lock_when_login_fails(self, tmp_path): - with pytest.raises(ValueError, match="login failed"): - with _custom_oauth_lock(tmp_path, "http://localhost:8020/callback"): - raise ValueError("login failed") - with _custom_oauth_lock(tmp_path, "http://127.0.0.1:8020/other-callback"): - assert len(list(tmp_path.glob("*.lock"))) == 1 - - class TestCustomClientToken: @pytest.fixture(autouse=True) def _setup(self, tmp_path, monkeypatch): - monkeypatch.setattr(oauth.TokenCache, "BASE_PATH", str(tmp_path / "oauth")) - monkeypatch.setenv("DATABRICKS_BEARER", "unrelated-bearer") - monkeypatch.setattr(db_mod, "run", Mock(side_effect=AssertionError("CLI not expected"))) - monkeypatch.setattr( - db_mod, "find_profile_name_for_host", Mock(side_effect=AssertionError("No profile")) - ) - self.endpoints = oauth.OidcEndpoints( - authorization_endpoint=f"{WS}/oidc/v1/authorize", - token_endpoint=f"{WS}/oidc/v1/token", - ) - self.discovery = Mock(return_value=self.endpoints) - monkeypatch.setattr(oauth, "get_workspace_endpoints", self.discovery) - self.browser = Mock(return_value=self._credentials("browser-token", "browser-refresh")) - monkeypatch.setattr(oauth.Consent, "launch_external_browser", self.browser) - self.refresh = Mock(return_value=self._credentials("refreshed", "rotated-refresh").token()) - monkeypatch.setattr(oauth, "retrieve_token", self.refresh) - - def _credentials(self, access_token, refresh_token): - return oauth.SessionCredentials( - token=oauth.Token( - access_token=access_token, - token_type="Bearer", - refresh_token=refresh_token, - expiry=datetime.now(UTC) + timedelta(hours=1), - ), - token_endpoint=self.endpoints.token_endpoint, - client_id="custom-client", - redirect_url="http://localhost:8020", - ) - - def _cache(self, workspace=WS, client_id="custom-client"): - return oauth.TokenCache( - host=workspace, - oidc_endpoints=self.endpoints, - client_id=client_id, - redirect_url="http://localhost:8020", - scopes=list(TEST_SCOPES), - ) - - def test_browser_login_uses_custom_client_and_redirect(self, capsys): - redirect_url = "http://localhost:41735/ai-devtools-workspace-oauth" - token = get_custom_client_token( - WS, client_id="custom-client", redirect_url=redirect_url, scopes=TEST_SCOPES - ) - assert token == "browser-token" - cached = self._cache().load().token() - assert cached.refresh_token == "browser-refresh" - output = capsys.readouterr() - assert output.out == "" - query = parse_qs(output.err.split("?", 1)[1].strip()) - assert query["client_id"] == ["custom-client"] - assert query["redirect_uri"] == [redirect_url] - assert query["scope"][0].split() == list(TEST_SCOPES) - - def test_reuses_cached_token_without_refresh_or_login(self): - self._cache().save(self._credentials("cached", "refresh")) - assert ( - get_custom_client_token(WS + "/", client_id="custom-client", scopes=TEST_SCOPES) - == "cached" - ) - self.browser.assert_not_called() - self.refresh.assert_not_called() - - def test_expired_token_refreshes_with_custom_client_and_saves_rotation(self): - cached = self._credentials("expired", "old-refresh") - self._cache().save(cached) - cache_path = Path(self._cache().filename) - payload = json.loads(cache_path.read_text()) - payload["token"]["expiry"] = (datetime.now(UTC) - timedelta(hours=1)).isoformat() - cache_path.write_text(json.dumps(payload)) - assert ( - get_custom_client_token(WS, client_id="custom-client", scopes=TEST_SCOPES) - == "refreshed" + monkeypatch.setenv("CUSTOM_OAUTH_CONFIG_FILE", "1") + monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 17, 0)) + monkeypatch.setattr(oauth_mod, "CUSTOM_OAUTH_CONFIG_FILE", tmp_path / "oauth.cfg") + + @staticmethod + def _result(returncode=0, stdout=""): + return subprocess.CompletedProcess([], returncode, stdout, "") + + def test_reuses_cli_token_without_login(self, monkeypatch): + run = Mock(return_value=self._result(stdout='{"access_token":"cached"}')) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + + assert get_custom_client_token(WS + "/", "custom-client", scopes=TEST_SCOPES) == "cached" + command = run.call_args.args[0] + assert command[:3] == ["databricks", "auth", "token"] + assert run.call_args.kwargs["env"]["DATABRICKS_CONFIG_FILE"].endswith("oauth.cfg") + + def test_login_then_retries_token(self, monkeypatch): + run = Mock( + side_effect=[ + self._result(1), + self._result(), + self._result(stdout='{"access_token":"new-token"}'), + ] ) - self.refresh.assert_called_once() - self.browser.assert_not_called() - assert self._cache().load().token().refresh_token == "rotated-refresh" - - def test_force_refresh_bypasses_fresh_access_token(self): - self._cache().save(self._credentials("cached", "refresh")) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + + assert get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) == "new-token" + first, login, retry = [call.args[0] for call in run.call_args_list] + assert first[first.index("--profile") + 1] == retry[retry.index("--profile") + 1] + assert login[:3] == ["databricks", "auth", "login"] + assert login[login.index("--client-id") + 1] == "custom-client" + assert login[login.index("--scopes") + 1] == ",".join(TEST_SCOPES) + assert login[login.index("--timeout") + 1] == "3m" + assert login[login.index("--profile") + 1] == first[first.index("--profile") + 1] + assert login[login.index("--profile") + 1].startswith("ug-custom-oauth-") + assert run.call_args_list[1].kwargs["capture_output"] is True + + def test_force_refresh_is_forwarded(self, monkeypatch): + run = Mock(return_value=self._result(stdout='{"access_token":"fresh"}')) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) assert ( - get_custom_client_token( - WS, client_id="custom-client", scopes=TEST_SCOPES, force_refresh=True - ) - == "refreshed" - ) - self.refresh.assert_called_once() - self.browser.assert_not_called() - - def test_refresh_failure_falls_back_to_browser(self, capsys): - self._cache().save(self._credentials("cached", "revoked-refresh")) - self.refresh.side_effect = ValueError("sensitive server response") - assert ( - get_custom_client_token( - WS, client_id="custom-client", scopes=TEST_SCOPES, force_refresh=True - ) - == "browser-token" + get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES, force_refresh=True) + == "fresh" ) - self.browser.assert_called_once() - output = capsys.readouterr() - assert "Sign in again" in output.err - assert "sensitive server response" not in output.err + assert "--force-refresh" in run.call_args.args[0] def test_invalid_redirect_fails_before_network(self): with pytest.raises(RuntimeError, match="--redirect-url must be"): @@ -145,19 +77,42 @@ def test_invalid_redirect_fails_before_network(self): redirect_url="https://example.com/callback", scopes=TEST_SCOPES, ) - self.discovery.assert_not_called() - def test_login_failure_is_actionable_and_does_not_expose_response(self): - self.browser.side_effect = ValueError("sensitive server response") - with pytest.raises(RuntimeError, match="registered redirect URL") as error: - get_custom_client_token(WS, client_id="custom-client", scopes=TEST_SCOPES) + def test_old_cli_is_actionable(self, monkeypatch): + monkeypatch.setattr(db_mod, "databricks_cli_version", lambda: (1, 16, 1)) + with pytest.raises(RuntimeError, match="v1.17.0 or newer"): + get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) + + def test_login_failure_is_actionable_and_sanitized(self, monkeypatch): + run = Mock(side_effect=[self._result(1), self._result(1, "sensitive response")]) + monkeypatch.setattr(oauth_mod.subprocess, "run", run) + with pytest.raises(RuntimeError, match="returned no custom-client OAuth token") as error: + get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) assert "sensitive server response" not in str(error.value) @pytest.mark.parametrize("scopes", [["offline_access"], ["catalog.catalogs:read"]]) def test_api_scopes_are_required(self, scopes): with pytest.raises(RuntimeError, match="OAuth scopes|API OAuth scope"): get_custom_client_token(WS, client_id="custom-client", scopes=scopes) - self.discovery.assert_not_called() + + +@pytest.mark.parametrize("value", [None, "0"]) +def test_cli_profile_is_opt_in(monkeypatch, value): + if value is None: + monkeypatch.delenv("CUSTOM_OAUTH_CONFIG_FILE", raising=False) + else: + monkeypatch.setenv("CUSTOM_OAUTH_CONFIG_FILE", value) + legacy = Mock(return_value="legacy-token") + monkeypatch.setattr(oauth_mod, "_get_custom_client_token_from_sdk", legacy) + cli = Mock(side_effect=AssertionError("CLI profile path must be disabled")) + monkeypatch.setattr(oauth_mod, "_get_custom_client_token_from_cli", cli) + + assert get_custom_client_token(WS, "custom-client", scopes=TEST_SCOPES) == "legacy-token" + legacy.assert_called_once() + + +def test_custom_oauth_config_filename(): + assert oauth_mod.CUSTOM_OAUTH_CONFIG_FILE.name == "ug.databrickscfg" class TestCustomClientCommand: