Add DATABRICKS_BEARER_COMMAND for externally brokered tokens

src/ucode/cli.py

@@ -54,6 +54,7 @@
     discover_model_services,
     ensure_databricks_auth,
     ensure_pat_bearer,
+    external_bearer_configured,
     find_profile_name_for_host,
     get_databricks_profiles,
     get_databricks_token,
@@ -517,7 +518,7 @@ def configure_shared_state(
         # token to avoid re-reading ~/.databrickscfg.
         ensure_pat_bearer(profile, pat)
         ensure_databricks_auth(workspace, profile)
-    elif force_login:
+    elif force_login and not external_bearer_configured():
         run_databricks_login(workspace, profile)
     else:
         ensure_databricks_auth(workspace, profile)

src/ucode/databricks.py

@@ -692,7 +692,7 @@ def resolve_sql_warehouse_id(workspace: str, token: str) -> tuple[str | None, st
 
 @overload
 def run(
-    args: list[str],
+    args: list[str] | str,
     *,
     check: bool = True,
     capture_output: bool = False,
@@ -704,7 +704,7 @@ def run(
 
 @overload
 def run(
-    args: list[str],
+    args: list[str] | str,
     *,
     check: bool = True,
     capture_output: bool = False,
@@ -715,7 +715,7 @@ def run(
 
 
 def run(
-    args: list[str],
+    args: list[str] | str,
     *,
     check: bool = True,
     capture_output: bool = False,
@@ -897,11 +897,23 @@ def _profile_args(profile: str | None) -> list[str]:
     return ["--profile", profile] if profile else []
 
 
+def external_bearer_configured() -> bool:
+    """Whether something outside ucode owns auth for this process.
+
+    True when either hatch is set: ``DATABRICKS_BEARER`` (a pre-fetched bearer)
+    or ``DATABRICKS_BEARER_COMMAND`` (one minted on demand). Both make an
+    interactive login pointless, since ``get_databricks_token`` returns before
+    it ever reaches the OAuth path."""
+    return bool(
+        os.environ.get("DATABRICKS_BEARER", "").strip()
+        or os.environ.get("DATABRICKS_BEARER_COMMAND", "").strip()
+    )
+
+
 def has_valid_databricks_auth(workspace: str, profile: str | None = None) -> bool:
-    # Honor the CI short-circuit (see ``get_databricks_token``): if a
-    # pre-fetched bearer is available, treat auth as valid and skip the
-    # `databricks auth token` shell-out (which only knows user-OAuth).
-    if os.environ.get("DATABRICKS_BEARER", "").strip():
+    # Auth owned elsewhere is valid by definition: skip the `databricks auth
+    # token` shell-out (which only knows user-OAuth) and any login it triggers.
+    if external_bearer_configured():
         return True
     _log_auth_diagnostics()
     # Mirror run_databricks_login: when ~/.databrickscfg has multiple
@@ -1128,6 +1140,46 @@ def ensure_databricks_auth(
     run_databricks_login(workspace, profile)
 
 
+def _bearer_from_command(command: str) -> str:
+    """Run ``DATABRICKS_BEARER_COMMAND`` and return the bearer it prints.
+
+    Fails closed instead of falling through to OAuth: a caller that set this
+    owns auth, and its profile usually carries no OAuth cache, so a fallback
+    would report a misleading stale-login error instead of the real cause.
+    Mirrors how ``auth-token --use-pat`` fails closed for the same reason."""
+    _debug("get_databricks_token", "using DATABRICKS_BEARER_COMMAND")
+    try:
+        # Windows takes the command line as one string and lets CreateProcess
+        # split it: shlex's POSIX rules would eat the backslashes in `C:\...`,
+        # and posix=False would keep the quotes around a path with spaces. This
+        # is the inverse of what build_auth_shell_command emits there.
+        argv = command if os.name == "nt" else shlex.split(command)
+        result = run(
+            argv,
+            check=False,
+            capture_output=True,
+            text=True,
+            timeout=15,
+        )
+    except (OSError, ValueError, subprocess.TimeoutExpired) as exc:
+        raise RuntimeError(
+            f"DATABRICKS_BEARER_COMMAND could not be run: {type(exc).__name__}: {exc}. "
+            f"Command: {command}"
+        ) from exc
+    # Deliberately not _format_subprocess_result: that includes stdout on a
+    # non-zero exit, and this command's stdout is the bearer itself.
+    stderr = (result.stderr or "").strip()[:500]
+    _debug("bearer command", f"rc={result.returncode} stderr={stderr!r}")
+    token = (result.stdout or "").strip()
+    if result.returncode == 0 and token:
+        return token
+    # A non-zero exit fails closed even when something reached stdout: that is a
+    # diagnostic, not a bearer, and forwarding it only resurfaces as a 401.
+    reason = f"exited {result.returncode}" if result.returncode else "printed no token"
+    detail = f" Stderr: {stderr}" if stderr else ""
+    raise RuntimeError(f"DATABRICKS_BEARER_COMMAND {reason}. Command: {command}.{detail}")
+
+
 def get_databricks_token(
     workspace: str,
     profile: str | None = None,
@@ -1145,6 +1197,14 @@ def get_databricks_token(
         _debug("get_databricks_token", "using DATABRICKS_BEARER env var")
         return bearer
 
+    # ``DATABRICKS_BEARER_COMMAND`` is the same escape hatch in command form,
+    # for callers whose bearer expires and has to be re-minted (an external
+    # credential broker, a sidecar). A static env var cannot be rewritten in a
+    # running process, so the command is re-run on every fetch instead.
+    command = os.environ.get("DATABRICKS_BEARER_COMMAND", "").strip()
+    if command:
+        return _bearer_from_command(command)
+
     _log_auth_diagnostics()
     # See has_valid_databricks_auth: resolve the profile from the host when
     # the caller didn't supply one, so duplicate-host cfgs don't break us.
@@ -1437,7 +1497,8 @@ def build_auth_token_argv(
     Unlike the previous POSIX `databricks ... | jq` pipeline, this is a single
     executable with plain arguments — no `sh`, no `jq`, no shell quoting — so it
     runs identically on macOS, Linux, and Windows (issue #116). The DATABRICKS_BEARER
-    short-circuit and the PAT path both live inside `auth-token` itself."""
+    short-circuit, its DATABRICKS_BEARER_COMMAND counterpart, and the PAT path all
+    live inside `auth-token` itself."""
     argv = [_ucode_binary(), "auth-token", "--host", workspace.rstrip("/")]
     if profile:
         argv += ["--profile", profile]

tests/test_cli.py

@@ -3719,3 +3719,48 @@ def test_no_flag_and_no_state_forwards_false(self, monkeypatch):
         result, captured = self._invoke(monkeypatch, flag=False, state={"workspace": "https://x"})
         assert result.exit_code == 0, result.output
         assert captured["kwargs"]["use_pat"] is False
+
+
+class TestForcedLoginWithExternalBearer:
+    """`configure --workspaces` forces `databricks auth login`, which cannot help
+    when a bearer (or a command that mints one) is supplied from outside: the
+    login is interactive, and `get_databricks_token` returns before it would ever
+    reach the OAuth path. A sandbox whose credential comes from a broker would
+    otherwise hang on a browser prompt it can never satisfy."""
+
+    _SENTINEL = "stop-after-auth"
+
+    def _run(self, monkeypatch) -> list:
+        """Drive configure_shared_state's auth branch, stopping right after it."""
+        calls: list = []
+        monkeypatch.setattr(
+            cli_mod, "run_databricks_login", lambda ws, profile=None: calls.append(ws)
+        )
+        monkeypatch.setattr(cli_mod, "ensure_databricks_auth", lambda *a, **k: None)
+        monkeypatch.setattr(cli_mod, "find_profile_name_for_host", lambda ws: None)
+
+        def stop(*_a, **_k):
+            raise RuntimeError(self._SENTINEL)
+
+        monkeypatch.setattr(cli_mod, "get_databricks_token", stop)
+        with pytest.raises(RuntimeError, match=self._SENTINEL):
+            cli_mod.configure_shared_state("https://ws.cloud.databricks.com", force_login=True)
+        return calls
+
+    def test_bearer_command_skips_the_interactive_login(self, monkeypatch):
+        monkeypatch.setenv("DATABRICKS_BEARER_COMMAND", "/opt/broker/mint.sh")
+        monkeypatch.delenv("DATABRICKS_BEARER", raising=False)
+
+        assert self._run(monkeypatch) == []
+
+    def test_static_bearer_skips_the_interactive_login(self, monkeypatch):
+        monkeypatch.setenv("DATABRICKS_BEARER", "ci-bearer")
+        monkeypatch.delenv("DATABRICKS_BEARER_COMMAND", raising=False)
+
+        assert self._run(monkeypatch) == []
+
+    def test_still_logs_in_when_nothing_external_is_set(self, monkeypatch):
+        monkeypatch.delenv("DATABRICKS_BEARER", raising=False)
+        monkeypatch.delenv("DATABRICKS_BEARER_COMMAND", raising=False)
+
+        assert self._run(monkeypatch) == ["https://ws.cloud.databricks.com"]

tests/test_databricks.py

@@ -3445,3 +3445,126 @@ def test_requests_the_account_users_principal(self, monkeypatch):
         all_users_can_use_schema("https://ws", "tok", "main.tien_le")
         assert "effective-permissions/schema/main.tien_le" in seen["url"]
         assert "principal=account%20users" in seen["url"]
+
+
+class TestBearerCommand:
+    """``DATABRICKS_BEARER_COMMAND`` is the command form of the static
+    ``DATABRICKS_BEARER`` hatch, for callers whose bearer expires and has to be
+    re-minted mid-session (an external credential broker, a sidecar)."""
+
+    def _env(self, tmp_path, monkeypatch, command: str | None):
+        """Patch in an env with a recording fake `databricks` on PATH.
+
+        The fake would happily serve a token, so any test asserting the marker
+        is absent is asserting the OAuth path was never reached.
+        """
+        marker = tmp_path / "cli-calls"
+        fake = tmp_path / "databricks"
+        fake.write_text(
+            f"#!/bin/sh\necho called >> {marker}\n"
+            'echo \'{"access_token": "oauth-token", "token_type": "Bearer"}\'\n'
+        )
+        fake.chmod(0o755)
+        path = os.environ.get("PATH", "")
+        env = {**os.environ, "PATH": f"{tmp_path}{os.pathsep}{path}"}
+        env.pop("DATABRICKS_BEARER", None)
+        env.pop("DATABRICKS_BEARER_COMMAND", None)
+        if command is not None:
+            env["DATABRICKS_BEARER_COMMAND"] = command
+        monkeypatch.setattr("os.environ", env)
+        return marker
+
+    def _broker(self, tmp_path, body: str) -> str:
+        script = tmp_path / "broker.sh"
+        script.write_text(f"#!/bin/sh\n{body}\n")
+        script.chmod(0o755)
+        return str(script)
+
+    def test_serves_the_command_output_without_touching_the_cli(self, tmp_path, monkeypatch):
+        broker = self._broker(tmp_path, 'echo "brokered-token"')
+        marker = self._env(tmp_path, monkeypatch, broker)
+
+        assert get_databricks_token(WS) == "brokered-token"
+        assert not marker.exists()
+
+    def test_reruns_the_command_on_every_fetch(self, tmp_path, monkeypatch):
+        # The whole point: a static env var cannot be rewritten in a running
+        # process, so an expiring bearer has to be re-minted per fetch.
+        counter = tmp_path / "mints"
+        counter.write_text("0")
+        broker = self._broker(
+            tmp_path,
+            f'n=$(cat {counter})\nn=$((n + 1))\necho $n > {counter}\necho "token-$n"',
+        )
+        self._env(tmp_path, monkeypatch, broker)
+
+        assert get_databricks_token(WS) == "token-1"
+        assert get_databricks_token(WS) == "token-2"
+
+    def test_passes_arguments_without_a_shell(self, tmp_path, monkeypatch):
+        # Argv is shlex-split, not handed to `sh -c`, so this stays cross-platform.
+        seen = tmp_path / "args"
+        broker = self._broker(tmp_path, f'printf "%s" "$1:$2" > {seen}\necho tok')
+        self._env(tmp_path, monkeypatch, f"{broker} --coords 'a path'")
+
+        assert get_databricks_token(WS) == "tok"
+        assert seen.read_text() == "--coords:a path"
+
+    def test_windows_hands_the_command_line_over_verbatim(self, tmp_path, monkeypatch):
+        # CreateProcess splits the string itself. shlex's POSIX rules would turn
+        # `C:\bin\broker.exe` into `C:binbroker.exe`, and posix=False would keep
+        # the quotes around a path containing spaces.
+        self._env(tmp_path, monkeypatch, r"C:\bin\broker.exe --arg")
+        monkeypatch.setattr(db_mod.os, "name", "nt")
+        seen = {}
+
+        def fake_run(args, **kwargs):
+            seen["args"] = args
+            return subprocess.CompletedProcess(args, 0, stdout="win-token\n", stderr="")
+
+        monkeypatch.setattr(db_mod, "run", fake_run)
+
+        assert get_databricks_token(WS) == "win-token"
+        assert seen["args"] == r"C:\bin\broker.exe --arg"
+
+    def test_fails_closed_when_the_command_prints_no_token(self, tmp_path, monkeypatch):
+        # Exit 0 with an empty stdout. Falling through to OAuth would report a
+        # misleading stale-login error: a broker-backed profile carries no OAuth
+        # cache to refresh. Stderr rides along so the error names a cause.
+        broker = self._broker(tmp_path, 'echo "nothing to vend" >&2')
+        marker = self._env(tmp_path, monkeypatch, broker)
+
+        with pytest.raises(RuntimeError, match="printed no token") as excinfo:
+            get_databricks_token(WS)
+        assert "nothing to vend" in str(excinfo.value)
+        assert not marker.exists()
+
+    def test_fails_closed_when_the_command_exits_non_zero(self, tmp_path, monkeypatch):
+        # Stdout on a failing command is a diagnostic, not a bearer. Forwarding it
+        # would only resurface as a 401 far from the real cause.
+        broker = self._broker(tmp_path, 'echo "broker unreachable"\nexit 7')
+        marker = self._env(tmp_path, monkeypatch, broker)
+
+        with pytest.raises(RuntimeError, match="exited 7"):
+            get_databricks_token(WS)
+        assert not marker.exists()
+
+    def test_reports_an_unrunnable_command(self, tmp_path, monkeypatch):
+        self._env(tmp_path, monkeypatch, str(tmp_path / "does-not-exist"))
+
+        with pytest.raises(RuntimeError, match="could not be run"):
+            get_databricks_token(WS)
+
+    def test_static_bearer_still_wins(self, tmp_path, monkeypatch):
+        broker = self._broker(tmp_path, 'echo "brokered-token"')
+        self._env(tmp_path, monkeypatch, broker)
+        os.environ["DATABRICKS_BEARER"] = "ci-bearer"
+
+        assert get_databricks_token(WS) == "ci-bearer"
+
+    def test_has_valid_auth_short_circuits(self, tmp_path, monkeypatch):
+        # Otherwise `ensure_databricks_auth` probes the CLI and can open a browser.
+        marker = self._env(tmp_path, monkeypatch, self._broker(tmp_path, "echo tok"))
+
+        assert db_mod.has_valid_databricks_auth(WS) is True
+        assert not marker.exists()