diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py index a0accb9be..8b1b59c6e 100644 --- a/src/ucode/databricks.py +++ b/src/ucode/databricks.py @@ -786,15 +786,30 @@ def perms_ok(path: str) -> bool: return found +# The first `--version` probe's budget, then one longer retry for a slow cold start. +_DATABRICKS_CLI_VERSION_TIMEOUTS_SECONDS = (10, 60) + + def _read_databricks_cli_version(path: str) -> tuple[int, int, int] | None: """Run `` --version`` and parse it. None on any subprocess or parse failure. Used only by discovery, which is what establishes ``path`` in the first place — every other caller reads a version through ``databricks_cli_path()``. + A first run can be slow (e.g. macOS verifying a freshly upgraded binary), so a + timed-out probe is retried once with a longer budget before giving up. """ - try: - result = run([path, "--version"], check=False, capture_output=True, text=True, timeout=10) - except (OSError, subprocess.TimeoutExpired): + result = None + for timeout in _DATABRICKS_CLI_VERSION_TIMEOUTS_SECONDS: + try: + result = run( + [path, "--version"], check=False, capture_output=True, text=True, timeout=timeout + ) + break + except subprocess.TimeoutExpired: + continue + except OSError: + return None + if result is None: return None raw = result.stdout or result.stderr or "" output = (raw if isinstance(raw, str) else raw.decode(errors="replace")).strip() diff --git a/tests/test_databricks.py b/tests/test_databricks.py index bcf26da70..daac11103 100644 --- a/tests/test_databricks.py +++ b/tests/test_databricks.py @@ -4402,11 +4402,29 @@ def boom(*a, **kw): assert db_mod._read_databricks_cli_version("/usr/bin/databricks") is None def test_read_version_swallows_timeout(self, monkeypatch): + timeouts = [] + def boom(*a, **kw): - raise subprocess.TimeoutExpired("databricks", 10) + timeouts.append(kw["timeout"]) + raise subprocess.TimeoutExpired("databricks", kw["timeout"]) monkeypatch.setattr(db_mod, "run", boom) assert db_mod._read_databricks_cli_version("/usr/bin/databricks") is None + assert timeouts == [10, 60] + + def test_read_version_retries_a_slow_first_run(self, monkeypatch): + """A cold `--version` that outlasts the first budget is read on the longer retry.""" + timeouts = [] + + def slow_then_ok(*a, **kw): + timeouts.append(kw["timeout"]) + if len(timeouts) == 1: + raise subprocess.TimeoutExpired("databricks", kw["timeout"]) + return subprocess.CompletedProcess(a, 0, "Databricks CLI v1.20.0", "") + + monkeypatch.setattr(db_mod, "run", slow_then_ok) + assert db_mod._read_databricks_cli_version("/usr/bin/databricks") == (1, 20, 0) + assert timeouts == [10, 60] def test_read_version_parses_output(self, monkeypatch): monkeypatch.setattr(