From cffcd146ef8f240b2d25c23f2bdf232b3ed966b3 Mon Sep 17 00:00:00 2001
From: Harry Yao <harry.yao@databricks.com>
Date: Thu, 22 Oct 2026 21:25:52 +0000
Subject: [PATCH] Retry a slow first `databricks --version` before giving up

Discovery probes each databricks binary with a 10 s timeout and records a timed-out probe as unparseable. A cold first run (e.g. right after an upgrade) can take longer, so the launch failed with "Could not parse Databricks CLI version ... 'Databricks CLI v1.20.0'" even though the fallback read in the same call succeeded. Retry a timed-out probe once with a 60 s budget.

Signed-off-by: Harry Yao <harry.yao@databricks.com>
Co-authored-by: Isaac <no-reply@databricks.com>
---
 src/ucode/databricks.py  | 21 ++++++++++++++++++---
 tests/test_databricks.py | 20 +++++++++++++++++++-
 2 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/src/ucode/databricks.py b/src/ucode/databricks.py
index f32b2c52a..02d43e492 100644
--- a/src/ucode/databricks.py
+++ b/src/ucode/databricks.py
@@ -770,15 +770,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 ``<path> --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 9910d1a1f..afad31ee3 100644
--- a/tests/test_databricks.py
+++ b/tests/test_databricks.py
@@ -4088,11 +4088,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(