Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@
## Requirements

- Python 3.12+ — install with `uv` ([uv.astral.sh](https://docs.astral.sh/uv/getting-started/installation/))
- The Databricks CLI on your `PATH` ([install guide](https://docs.databricks.com/aws/en/dev-tools/cli/install)) — `ucode` never installs or upgrades it for you. Keep it current (`brew upgrade databricks/tap/databricks`); features like `ucode configure` rely on recent subcommands such as `databricks aitools`.
- `npm` if tool CLIs need to be installed automatically

## Installation
Expand Down
4 changes: 2 additions & 2 deletions src/ucode/agents/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,9 +19,9 @@
from ucode.config_io import ToolSpec
from ucode.databricks import (
BEDROCK_PROVIDER_TYPES,
ensure_databricks_cli,
get_databricks_token,
install_ai_tools,
install_databricks_cli,
map_bedrock_claude_models,
resolve_provider_service,
)
Expand Down Expand Up @@ -254,7 +254,7 @@ def ensure_bootstrap_dependencies(
update_existing: bool = False,
prompt_optional_updates: bool = True,
) -> None:
install_databricks_cli()
ensure_databricks_cli()
install_tool_binary(
tool,
strict=True,
Expand Down
8 changes: 4 additions & 4 deletions src/ucode/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,11 +43,11 @@
discover_model_services,
ensure_ai_gateway_v2,
ensure_databricks_auth,
ensure_databricks_cli,
ensure_pat_bearer,
find_profile_name_for_host,
get_databricks_profiles,
get_databricks_token,
install_databricks_cli,
is_model_provider_feature_unavailable,
list_profile_entries,
list_tool_provider_services,
Expand Down Expand Up @@ -1648,7 +1648,7 @@ def configure(
set_verbosity(verbose)
prompt_optional_updates = not skip_upgrade
try:
install_databricks_cli()
ensure_databricks_cli()
if agent is not None and agents is not None:
raise RuntimeError("Use either --agent or --agents, not both.")
if workspaces is not None and profiles is not None:
Expand Down Expand Up @@ -1927,7 +1927,7 @@ def configure_tracing(
) -> None:
"""Send coding-session traces to an MLflow experiment in your workspace."""
try:
install_databricks_cli()
ensure_databricks_cli()
configure_tracing_command(disable=disable)
except RuntimeError as exc:
print_err(str(exc))
Expand Down Expand Up @@ -1961,7 +1961,7 @@ def revert_cmd() -> None:
def usage_cmd() -> None:
"""Show Databricks AI Gateway usage summary (last 7 days)."""
try:
install_databricks_cli()
ensure_databricks_cli()
usage_report()
except RuntimeError as exc:
print_err(str(exc))
Expand Down
90 changes: 17 additions & 73 deletions src/ucode/databricks.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,15 +43,8 @@
spinner,
)

UNIX_DATABRICKS_INSTALL_URL = (
"https://raw.githubusercontent.com/databricks/setup-cli/main/install.sh"
)
WINDOWS_DATABRICKS_INSTALL_URL = (
"https://raw.githubusercontent.com/databricks/setup-cli/main/install.ps1"
)
AI_GATEWAY_V2_DOCS_URL = "https://docs.databricks.com/aws/en/ai-gateway/overview-beta"
# v1.0.0 is the release that ships `databricks aitools`.
MIN_DATABRICKS_CLI_VERSION = (1, 0, 0)
DATABRICKS_CLI_INSTALL_DOCS_URL = "https://docs.databricks.com/aws/en/dev-tools/cli/install"
TOKEN_REFRESH_INTERVAL_SECONDS = 1800


Expand Down Expand Up @@ -517,75 +510,26 @@ def workspace_hostname(workspace: str) -> str:
return parsed.hostname


def _parse_databricks_cli_version(output: str) -> tuple[int, int, int] | None:
# Example output: "Databricks CLI v0.299.2"
match = re.search(r"v?(\d+)\.(\d+)\.(\d+)", output)
if not match:
return None
return (int(match.group(1)), int(match.group(2)), int(match.group(3)))


def _run_databricks_cli_installer(brew_subcommand: str = "install") -> None:
system = platform.system()
try:
if system == "Windows":
run(
["powershell", "-Command", f"irm {WINDOWS_DATABRICKS_INSTALL_URL} | iex"],
timeout=240,
)
elif system == "Darwin" and shutil.which("brew"):
run(["brew", brew_subcommand, "databricks/tap/databricks"], timeout=240)
elif shutil.which("curl"):
run(["sh", "-c", f"curl -fsSL {UNIX_DATABRICKS_INSTALL_URL} | sudo sh"], timeout=240)
elif shutil.which("wget"):
run(["sh", "-c", f"wget -qO- {UNIX_DATABRICKS_INSTALL_URL} | sudo sh"], timeout=240)
else:
raise RuntimeError("Neither curl nor wget is available.")
except (subprocess.CalledProcessError, subprocess.TimeoutExpired, RuntimeError) as exc:
raise RuntimeError("Failed to install/upgrade Databricks CLI automatically.") from exc


def ensure_databricks_cli_version() -> None:
try:
result = run(
["databricks", "--version"],
check=False,
capture_output=True,
text=True,
timeout=10,
)
except (OSError, subprocess.TimeoutExpired) as exc:
raise RuntimeError("Failed to read Databricks CLI version.") from exc

raw = result.stdout or result.stderr or ""
output = (raw if isinstance(raw, str) else raw.decode(errors="replace")).strip()
version = _parse_databricks_cli_version(output)
if version is None:
raise RuntimeError(
f"Could not parse Databricks CLI version from `databricks --version` output: {output!r}"
)
if version < MIN_DATABRICKS_CLI_VERSION:
current = ".".join(str(n) for n in version)
required = ".".join(str(n) for n in MIN_DATABRICKS_CLI_VERSION)
print_warning(
f"Databricks CLI v{current} is too old (need v{required} or newer). Upgrading..."
)
_run_databricks_cli_installer(brew_subcommand="upgrade")
ensure_databricks_cli_version()
"""Intentionally a no-op: ucode no longer checks or upgrades the installed
Databricks CLI version.

Kept as a seam so callers don't have to care whether a version policy
exists. Managing the CLI is the user's (or their package manager's) job --
silently replacing a working install, including a locally built one, is
more disruptive than a feature failing with the CLI's own error."""
return

def install_databricks_cli() -> None:
if shutil.which("databricks"):
ensure_databricks_cli_version()
return

print_section("Bootstrap")
print_warning("`databricks` was not found. Installing Databricks CLI...")
_run_databricks_cli_installer(brew_subcommand="install")
def ensure_databricks_cli() -> None:
"""Verify `databricks` is on PATH, raising with install instructions if not.

ucode never installs or upgrades the CLI on the user's behalf."""
if not shutil.which("databricks"):
raise RuntimeError(
"Databricks CLI install completed, but `databricks` is still not on PATH."
"Databricks CLI was not found on PATH. Install it, then re-run this command: "
f"see {DATABRICKS_CLI_INSTALL_DOCS_URL} "
"(on macOS: `brew install databricks/tap/databricks`)."
)
ensure_databricks_cli_version()

Expand All @@ -612,9 +556,9 @@ def install_ai_tools(agent_tokens: list[str], profile: str | None = None) -> Non
timeout=300,
)
except (subprocess.CalledProcessError, subprocess.TimeoutExpired, OSError) as exc:
# The CLI version is already guaranteed by ensure_databricks_cli_version,
# so any failure here is something else (e.g. an agent binary missing
# from PATH). Surface the CLI's own error rather than guessing a cause.
# ucode doesn't police the CLI version, so this may be a CLI too old to
# know `aitools` just as easily as something else (e.g. an agent binary
# missing from PATH). Surface the CLI's own error rather than guessing.
detail = getattr(exc, "stderr", None) or ""
if isinstance(detail, bytes): # TimeoutExpired.stderr is bytes even with text=True
detail = detail.decode(errors="replace")
Expand Down
20 changes: 14 additions & 6 deletions src/ucode/skills_download.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,11 @@

SKILL_NAME_PATTERN = re.compile(r"^[a-z0-9-]+$")

# Skill bundles are served from the Files API's `Skills/` root, not `Volumes/`.
# Both the directory walk and the per-file fetch must use it, since relative
# paths are derived by stripping this prefix off the listing's absolute paths.
SKILL_FILES_ROOT = "Skills"

# Parallel skill fetches per schema; writes stay sequential (they prompt).
_MAX_FETCH_WORKERS = 8

Expand Down Expand Up @@ -86,15 +91,15 @@ def list_skill_files(
) -> tuple[list[str], str | None]:
"""List a skill bundle's files, as paths relative to the skill directory.

Recursively walks the skill's UC Volume directory (including ``SKILL.md``).
Recursively walks the skill's bundle directory (including ``SKILL.md``).
A non-None reason indicates the listing call itself failed.
"""
hostname = workspace_hostname(workspace)
dirs_base = f"https://{hostname}/api/2.0/fs/directories"
volume_prefix = f"/Volumes/{catalog}/{schema}/{leaf}/"
bundle_prefix = f"/{SKILL_FILES_ROOT}/{catalog}/{schema}/{leaf}/"

relative_paths: list[str] = []
pending = [f"Volumes/{catalog}/{schema}/{leaf}"]
pending = [f"{SKILL_FILES_ROOT}/{catalog}/{schema}/{leaf}"]
while pending:
directory = pending.pop()
page_token: str | None = None
Expand All @@ -113,7 +118,7 @@ def list_skill_files(
if entry.get("is_directory"):
pending.append(path.strip("/"))
else:
relative_paths.append(path.removeprefix(volume_prefix))
relative_paths.append(path.removeprefix(bundle_prefix))
page_token = data.get("next_page_token")
if not page_token:
break
Expand All @@ -123,9 +128,12 @@ def list_skill_files(
def fetch_skill_file(
workspace: str, token: str, catalog: str, schema: str, leaf: str, relative_path: str
) -> tuple[bytes | None, str | None]:
"""Fetch one skill bundle file's raw bytes from its UC Volume."""
"""Fetch one skill bundle file's raw bytes from the Files API."""
hostname = workspace_hostname(workspace)
url = f"https://{hostname}/api/2.0/fs/files/Volumes/{catalog}/{schema}/{leaf}/{relative_path}"
url = (
f"https://{hostname}/api/2.0/fs/files/"
f"{SKILL_FILES_ROOT}/{catalog}/{schema}/{leaf}/{relative_path}"
)
return _http_get_bytes(url, token, timeout=30)


Expand Down
Loading
Loading