Codex/use cli custom oauth

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,13 +16,16 @@
 
 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
 
 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)