diff --git a/.claude-plugin/plugin.json b/.claude-plugin/plugin.json index 31cc1f6..b4eb15c 100644 --- a/.claude-plugin/plugin.json +++ b/.claude-plugin/plugin.json @@ -22,5 +22,6 @@ "vector-search", "data-engineering" ], - "skills": "./skills/" + "skills": "./skills/", + "commands": "./commands/" } diff --git a/.github/workflows/validate-manifest.yml b/.github/workflows/validate-manifest.yml index 4ac26a0..1a5ec3b 100644 --- a/.github/workflows/validate-manifest.yml +++ b/.github/workflows/validate-manifest.yml @@ -9,6 +9,9 @@ on: - 'scripts/skills.py' - 'manifest.json' - '.claude-plugin/**' + - 'hooks/**' + - 'commands/**' + - 'tests/**' push: branches: - main @@ -28,3 +31,6 @@ jobs: - name: Validate manifest is up to date run: python3 scripts/skills.py validate + + - name: Test plugin hooks + run: python3 -m unittest discover -s tests -p '*_test.py' -v diff --git a/CLAUDE.md b/CLAUDE.md index a5f3f08..cb10cbb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -47,6 +47,41 @@ python3 scripts/skills.py sync # sync Codex metadata + icons only python3 scripts/skills.py validate # check Codex metadata + icons + manifest are up to date (CI) ``` +## Plugin components (hooks + commands) + +Beyond skills, the Claude Code plugin ships two component dirs at the repo root. +`commands/` is declared via `"commands"` in `.claude-plugin/plugin.json`, but +**`hooks/hooks.json` is auto-loaded by Claude Code and must NOT be declared** +there. Declaring the standard path double-loads it and fails the plugin with a +"Duplicate hooks file" error. + +- `hooks/`: a UserPromptSubmit prompt router (`databricks-router.py`) that + steers Databricks-related prompts into the skills, a SessionStart context + primer (`databricks-context.py`), and a PostToolUse auth-failure hinter + (`databricks-auth-helper.py`), wired via `hooks/hooks.json`. All + stdlib-only and fail-open. See [hooks/README.md](./hooks/README.md). +- `commands/`: friction-only slash commands (`/databricks:setup`, + `/databricks:doctor`). Product workflows stay in the skills, not commands, to + avoid shadowing a skill of the same name. + +`python3 scripts/skills.py validate` checks these (hooks.json is valid and +references existing scripts, plugin.json does not double-declare hooks, every +command has frontmatter). After changing hook behavior, run the hook test +suite: `python3 -m unittest discover -s tests -p '*_test.py'`. +These ship via the plugin marketplace +(whole-repo source); `databricks aitools install` currently installs skills only. + +**Marketplace entries are load-bearing for installed plugins.** Never remove a +shipped plugin's entry from `.claude-plugin/marketplace.json` (and never rename +the plugin or the marketplace). Claude Code re-resolves installed plugins +against the marketplace catalog at load time, so removing the entry does not +just stop updates: every existing install immediately fails to load ("Plugin +databricks not found in marketplace databricks-agent-skills") and those users +lose all skills, hooks, and commands until they manually uninstall and +reinstall from another source. Verified empirically (2026-06). Listing the +plugin on an additional marketplace, such as Anthropic's official directory, +is additive and never replaces the entry here. + ## Security When documenting examples, obfuscate sensitive info: diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 048b0be..dcb37fe 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -52,6 +52,34 @@ python3 scripts/skills.py validate If validation fails the error tells you which file is missing or stale; the fix is always `python3 scripts/skills.py generate` and committing the result. +## Plugin components (hooks + commands) + +The Claude Code plugin ships more than skills: + +- `hooks/`: `hooks.json` wires a UserPromptSubmit prompt router + (`databricks-router.py`) that steers Databricks-related prompts into the + skills, a SessionStart context primer (`databricks-context.py`), and a + PostToolUse auth-failure hinter (`databricks-auth-helper.py`). All + stdlib-only and fail-open. See [`hooks/README.md`](./hooks/README.md). Each + hook's behavior is pinned by its matching `tests/*_test.py` file; run the + suite with `python3 -m unittest discover -s tests -p '*_test.py'`. + **`hooks/hooks.json` is auto-loaded by Claude Code, so do NOT add a `"hooks"` + key to `.claude-plugin/plugin.json`, or the plugin fails to load with a + "Duplicate hooks file" error.** +- `commands/`: one `*.md` per slash command (`/databricks:`), declared via + `"commands"` in `.claude-plugin/plugin.json`. Each needs frontmatter + (`description`, optional `argument-hint`, `allowed-tools`). + +`scripts/skills.py validate` (run in CI) checks that `hooks/hooks.json` is valid +JSON referencing scripts that exist, that plugin.json does not double-declare the +standard hooks file, and that every command carries a `description` (quoted if it +contains a `:`, since strict YAML rejects unquoted colons). The validate +workflow also runs all hook test files. + +These components ship via the plugin marketplace (the whole repo is the plugin). +`databricks aitools install` packages `skills/` only today; extending it to +hooks/commands is CLI-side follow-up work. + ## Security Please see [SECURITY](./SECURITY) for vulnerability reporting guidelines. diff --git a/README.md b/README.md index 331d959..febbaab 100644 --- a/README.md +++ b/README.md @@ -51,6 +51,7 @@ skill under [`./skills/`](./skills/)): | Stable skills | ✅ (default) | ✅ | | Experimental skills | ✅ (with `--experimental` or by name) | ❌ | | Per-skill selection | ✅ (`databricks aitools install `) | ❌ (all-or-nothing) | +| Commands & hooks | ❌ (skills only today, see below) | ✅ | | Updates | `databricks aitools update` | Plugin marketplace update flow | | Required outside the agent | Databricks CLI v1.0.0+ | None | @@ -89,6 +90,48 @@ originally imported from - See [`experimental/README.md`](./experimental/README.md) for the full list and caveats. +## Commands and hooks (Claude Code) + +When installed as a Claude Code plugin, the `databricks` plugin adds slash +commands and three hooks (prompt routing, session context, auth-failure hints) +on top of the skills. +(These are Claude-Code-specific and ship via the plugin marketplace; the CLI +`databricks aitools install` path installs skills only today; see the note at +the end.) + +**Slash commands**: friction-only entry points; everyday work stays with the +auto-invoked skills. + +- `/databricks:setup [workspace-url]`: auth/onboarding. Install check, then an + OAuth / PAT / service-principal profile, then verify. +- `/databricks:doctor [profile]`: read-only health check (CLI version, auth, + workspace reachability, compute, recent job failures). + +(Product workflows such as apps, jobs, pipelines, DABs, etc. are handled by the +skills, not commands, so they aren't duplicated here.) + +**Hooks** (`hooks/`, all fail-open): + +- **Prompt router** (UserPromptSubmit): a fast keyword regex (sub-50ms, no LLM, + no network) over each prompt. When the prompt is Databricks-related, it injects + a note steering Claude to load `databricks-core` plus the matching product + skill before answering. The full note fires once per session; later Databricks + prompts get a one-line reminder. Unrelated prompts are untouched. No + permission gating, no cost warnings. +- **Context primer** (SessionStart, skipped on resume): injects the routing + rule, CLI version, configured profile names and any + `[__settings__].default_profile` (read locally, no network call, no token + values), and env/in-platform auth state. +- **Auth-failure hint** (PostToolUse on Bash): when a `databricks` command fails + with an auth-shaped error, adds one line suggesting `/databricks:doctor` or + `databricks auth login` before retrying. Never blocks or rewrites commands. + +> **Distribution parity (follow-up).** The plugin marketplace ships the whole +> repo (`marketplace.json` `source: "./"`), so commands and hooks come with it. +> `databricks aitools install` currently packages only `skills/`, so CLI-install +> users don't yet get commands/hooks. Closing that gap is tracked as CLI-side +> work. + ## Structure Each skill follows the [Agent Skills Specification](https://agentskills.io/specification): diff --git a/commands/doctor.md b/commands/doctor.md new file mode 100644 index 0000000..956da99 --- /dev/null +++ b/commands/doctor.md @@ -0,0 +1,41 @@ +--- +description: "Read-only Databricks health check: CLI, profiles, auth validity via one API call. Pass `full` to also check compute and recent job failures." +argument-hint: "[profile] [full]" +allowed-tools: Bash(databricks:*), Read +--- + +# Databricks Doctor + +Run a **read-only** health check and report a short status table. Make no +changes; every step below only reads. If a subcommand or flag is unfamiliar, +check `databricks --help` first rather than guessing. + +Run these in order. Don't stop on the first failure; collect what you can and +report the rest as unknown. + +1. **CLI**: `databricks --version`. Flag only if it's missing; don't gate on a + specific version (the CLI surfaces its own update notice). +2. **Profiles**: `databricks auth profiles`. List configured profiles and + validity. If `$1` is given, use that profile for the rest. Otherwise, if more + than one profile exists, ask the user which to use (**never auto-select**). +3. **Auth method**: `databricks auth describe --profile ` shows the + effective host, user, and credential source (never pass `--sensitive`). +4. **Auth validity**: `databricks current-user me --profile `. This + single API call proves the credentials work end to end (token valid, + workspace reachable, expected identity); don't probe other APIs for it. + For account-level profiles (an `accounts.*` host), `current-user me` does + not exist; report what `auth describe` resolved instead. + +Stop here by default. Run the extended checks below only when the user passed +`full` or asked about compute or jobs: + +5. **Compute**: `databricks warehouses list` and `databricks clusters list` for + the profile. Note what's running. +6. **Recent job failures**: list recent job runs (e.g. + `databricks jobs list-runs --limit 20 --profile `) and surface any + recent failures. + +Then print a compact table: **check | status (✅/⚠️/❌) | detail**. End with the +single most useful next action (e.g. "run `/databricks:setup` to add a profile"). + +This is a status check; it only reads, so don't run anything that changes state. diff --git a/commands/setup.md b/commands/setup.md new file mode 100644 index 0000000..d9edfd4 --- /dev/null +++ b/commands/setup.md @@ -0,0 +1,53 @@ +--- +description: "Set up Databricks CLI auth: install check, then an OAuth / PAT / service-principal profile (workspace or account-level), then verify." +argument-hint: "[workspace-or-account-url]" +allowed-tools: Bash(databricks:*), Read +--- + +# Databricks Setup + +Guide the user through Databricks CLI authentication. Use the **databricks-core** +skill for the authoritative auth details; this command is the step-by-step +wrapper around it. + +1. **CLI present?** `databricks --version`. If it's missing, + follow the install steps in the databricks-core skill + (`databricks-cli-install.md`). In sandboxed environments (Cursor, containers), + print the install command and ask the user to run it in their own terminal. + Don't try to install into the sandbox. +2. **Existing profiles?** `databricks auth profiles`. Show what's already + configured. If a working profile exists, ask whether to reuse it or add a new + one. +3. **Pick an auth method** (ask the user; `$1` may be a workspace or account + console URL): + - **OAuth U2M** (default, interactive): + `databricks auth login --host --profile `. Opens a + browser. Best for laptops. If the user doesn't know their workspace URL, + plain `databricks auth login --profile ` opens login.databricks.com + to sign in and pick a workspace. URLs copied from the browser may carry + `?w=` or `account_id=` query params; the CLI accepts them, + but quote the URL so the shell doesn't interpret the `?`. + - **Account-level**: when the host is an account console URL + (`accounts.cloud.databricks.com`, `accounts.azuredatabricks.net`, + `accounts.gcp.databricks.com`), also pass the account ID: + `databricks auth login --host --account-id --profile `. + Ask for the account ID if it isn't in the URL (it's the UUID shown in the + account console address bar). + - **PAT**: `databricks configure --token --profile `; the user pastes + a personal access token. This command prompts on stdin, so don't run it + yourself (it hangs without a TTY): ask the user to run it in their own + terminal, then continue once it's done. The same applies to + `databricks auth login` when no browser can open (headless or sandboxed + sessions). + - **Service principal (M2M)**: client id/secret via profile or env. Use for + CI/automation; never a personal PAT in CI. + - **In-platform** (notebook/cluster): `DATABRICKS_HOST`/`DATABRICKS_TOKEN` + are already injected, so no setup is needed. +4. **Confirm before writing** any profile; auth writes to `~/.databrickscfg`. +5. **Verify**: `databricks current-user me --profile ` returns the + expected user. For account-level profiles, `current-user me` doesn't exist; + use `databricks auth describe --profile ` and check the resolved host + and account ID. + +Never echo tokens or secrets back. Never auto-select a profile. When done, +suggest `/databricks:doctor` for a full health check. diff --git a/hooks/README.md b/hooks/README.md new file mode 100644 index 0000000..1860dae --- /dev/null +++ b/hooks/README.md @@ -0,0 +1,81 @@ +# Plugin hooks + +Three hooks make sure Databricks work flows through the skills. All are +stdlib-only Python and **fail open** (any error prints `{}` / no output and +exits 0, so a broken hook never blocks a prompt, session start, or tool call). +`hooks.json` wires them in; Claude Code expands `${CLAUDE_PLUGIN_ROOT}`. Claude +Code auto-loads `hooks/hooks.json`, so it is **not** declared in `plugin.json` +(declaring the standard path double-loads it and fails the plugin). + +Each hook is pinned by a test file in `tests/` at the repo root; run the whole +suite with `python3 -m unittest discover -s tests -p '*_test.py'`. + +## `databricks-router.py`: prompt router (UserPromptSubmit) + +Runs a fast keyword regex (sub-50ms, no LLM, no network) over each user prompt. +When the prompt is Databricks-related, it injects an `additionalContext` +instruction telling Claude to load `databricks-core` plus the matching product +skill before answering. When it isn't, it prints `{}` and stays out of the way. + +The full instruction is injected **once per session** (tracked by a marker file +in the temp dir keyed on the payload's `session_id`); later Databricks prompts +in the same session get a one-line reminder instead, so long sessions don't pay +the full routing block on every turn. + +There's no second agent to delegate to. Claude itself drives the `databricks` +CLI through the skills, so "routing" just means "make sure the Databricks skills +are loaded." There is **no permission gating and no cost warning** here. + +Precision is tuned to avoid over-routing: + +- **STRONG** terms (`databricks`, `unity catalog`, `lakeflow`, `dbfs`, + `databricks.yml`, `spark declarative pipelines`, `delta live tables` (the + legacy name still routes), ...) always route, even alongside an + alternative-platform mention, so "migrate from redshift to databricks" routes. +- **AMBIGUOUS** terms (`declarative pipelines`, `model serving`, `vector + search`, `mlflow`, `pyspark`, `genie`, ...) route only when no **SUPPRESS** + term is present. +- **SUPPRESS** terms (alternative data platforms, Jenkins, and plainly-local + dev work like `git commit`, `read the file`, `unit test`, `npm`) hold back an + ambiguous match. +- **URLs**: code-hosting URLs are blanked before matching, so `databricks` + appearing only as a GitHub/GitLab org or repo name + (`github.com/databricks/...`) does not route. URLs whose hostname contains + `databricks` (workspace and docs hosts) still do. + +Edit those three lists when the product surface changes. Behavior is pinned by +`tests/databricks_router_test.py`. + +## `databricks-context.py`: context primer (SessionStart) + +Injects a compact banner at session start: the routing rule (load +`databricks-core` + the product skill), CLI presence + version, configured +profile names plus any `[__settings__].default_profile` (parsed from +`~/.databrickscfg` locally, **no network call**, token values never printed), +and whether env/in-platform auth is set. If the CLI isn't installed it points +at `/databricks:setup`. +Covered by `tests/databricks_context_test.py`. + +Its `hooks.json` entry uses `"matcher": "startup|clear|compact"`: the banner +fires for new sessions, `/clear`, and after compaction, but **not on resume**, +where the prior context already contains it. + +## `databricks-auth-helper.py`: auth-failure hint (PostToolUse) + +Watches Bash tool results (matcher: `Bash`). When a `databricks` command's +output matches a phrase-shaped auth-failure signal (missing default +credentials, `invalid_grant`, `401 unauthorized`, invalid/expired token), it +injects one line suggesting `/databricks:doctor` or `databricks auth login` +before any retry. It never blocks or rewrites tool calls; bare status codes in +ordinary output do not trigger it. Only commands that actually **invoke** the +`databricks` executable count: `databricks` appearing as a repo path, URL, or +argument (`gh pr view --repo databricks/cli`) does not, since such output can +legitimately quote auth-failure phrases without any auth problem. +Covered by `tests/databricks_auth_helper_test.py`. + +## Distribution note + +These ship with the Claude Code plugin (the whole repo is the plugin via +`marketplace.json` `source: "./"`). The Databricks CLI install path +(`databricks aitools install`) currently packages **skills only**. See the repo +README for the parity follow-up. diff --git a/hooks/databricks-auth-helper.py b/hooks/databricks-auth-helper.py new file mode 100644 index 0000000..4c098bf --- /dev/null +++ b/hooks/databricks-auth-helper.py @@ -0,0 +1,127 @@ +#!/usr/bin/env python3 +"""PostToolUse hook: suggest an auth fix when a `databricks` command fails auth. + +Watches Bash tool results. When the command actually invoked the `databricks` +CLI (as a segment executable, not merely "databricks" appearing in a repo +path, URL, or argument) and the output looks like an authentication failure +(missing credentials, expired or invalid token, OAuth refresh failure), it +injects one line of `additionalContext` pointing at `/databricks:doctor` and +`databricks auth login`. Everything else passes through silently. + +No gating: this never blocks or rewrites a tool call, it only adds context +after the fact. + +Contract (Claude Code PostToolUse hook, matcher: Bash): + stdin : JSON with tool_name, tool_input.command, tool_response + stdout: JSON -> hookSpecificOutput.additionalContext, or "{}". + Fail-open: on ANY error print "{}" and exit 0. +""" +import json +import re +import sys + +# `databricks` must be the executable of one of the command's shell segments, +# not a substring anywhere in the command line: `gh pr view --repo +# databricks/cli`, URLs, and file paths mention databricks without invoking +# the CLI, and their output can legitimately quote auth-failure phrases (a PR +# body, this hook's own source). Segments are split on shell connectors and +# command-substitution openers; each segment's executable is its first token +# after env assignments, common wrappers, and wrapper flags. Path-prefixed +# invocations (`/usr/local/bin/databricks`) count; `databricks-test` does not. +_SEGMENT_SPLIT_RE = re.compile(r"&&|\|\||\$\(|[;|&\n`(]") +_ENV_ASSIGNMENT_RE = re.compile(r"[A-Za-z_][A-Za-z0-9_]*=") +_WRAPPERS = frozenset({"sudo", "env", "command", "exec", "time", "nohup", "xargs"}) + + +def _segment_executable(tokens): + """First token that is not an env assignment, wrapper, or wrapper flag.""" + after_wrapper = False + for token in tokens: + if _ENV_ASSIGNMENT_RE.match(token): + continue + if token in _WRAPPERS: + after_wrapper = True + continue + if after_wrapper and token.startswith("-"): + continue + return token + return "" + + +def _invokes_databricks_cli(command): + """True when any segment of the command runs the `databricks` executable.""" + for segment in _SEGMENT_SPLIT_RE.split(command): + executable = _segment_executable(segment.split()) + if executable.rsplit("/", 1)[-1] == "databricks": + return True + return False + +# Phrase-shaped auth-failure signals as emitted by the CLI / Go SDK error +# paths. Deliberately not bare status codes, so ordinary data in stdout +# (e.g. a row containing 401) cannot trip them. +AUTH_ERROR_PATTERNS = [ + r"cannot configure default credentials", + r"\binvalid_grant\b", + r"\b401 unauthorized\b", + r"\binvalid access token\b", + r"\btoken (?:is |has |was )?expired\b", + r"\brefresh token (?:is |was )?(?:invalid|expired|revoked)\b", +] +_AUTH_ERRORS = [re.compile(p, re.IGNORECASE) for p in AUTH_ERROR_PATTERNS] + +AUTH_HINT = ( + "[DATABRICKS] The `databricks` command above failed with what looks like " + "an authentication error. Before retrying, fix auth: run " + "`/databricks:doctor` for a read-only diagnosis, or re-authenticate with " + "`databricks auth login --host --profile ` " + "(`/databricks:setup` walks through it). Never auto-select a profile for " + "the user." +) + + +def check(tool_name, command, response_text): + """Return the auth hint when a databricks command hit an auth error, else None.""" + if tool_name != "Bash": + return None + if not command or not _invokes_databricks_cli(command): + return None + if not response_text: + return None + if any(p.search(response_text) for p in _AUTH_ERRORS): + return AUTH_HINT + return None + + +def main(): + # One outer try so the fail-open guarantee covers the entire main block, + # including JSON serialization; the final print gets its own guard (a + # closed stdout must not surface as a hook failure either). + output = "{}" + try: + data = json.load(sys.stdin) + if not isinstance(data, dict): + raise TypeError("payload is not an object") + tool_input = data.get("tool_input") + command = tool_input.get("command", "") if isinstance(tool_input, dict) else "" + # Serialize the whole response instead of assuming its shape; auth + # errors can land in stdout, stderr, or a combined error field. + response_text = json.dumps(data.get("tool_response", ""), default=str) + result = check(data.get("tool_name", ""), command, response_text) + if result: + output = json.dumps({ + "hookSpecificOutput": { + "hookEventName": "PostToolUse", + "additionalContext": result, + } + }) + except Exception: + output = "{}" + try: + print(output) + except Exception: + pass + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/hooks/databricks-context.py b/hooks/databricks-context.py new file mode 100644 index 0000000..abdf61e --- /dev/null +++ b/hooks/databricks-context.py @@ -0,0 +1,168 @@ +#!/usr/bin/env python3 +"""SessionStart hook: inject a compact Databricks context banner. + +Local-only and fail-open by design. It never makes a network call (so it can't +hang, hit an MCP-style timeout, or trigger an auth prompt at session start) and +any error exits 0 with no output. It surfaces, when available: + + - databricks CLI presence + version + - configured profile names, parsed straight from the config file (no network) + - the `[__settings__].default_profile` the CLI resolves when --profile is omitted + - env-based / in-platform auth (DATABRICKS_HOST, DATABRICKS_CONFIG_PROFILE) + +Token values are never printed, only their presence. + +Contract (Claude Code SessionStart hook): + stdin : JSON (drained, content unused) + stdout: JSON -> hookSpecificOutput.additionalContext, a string injected into + the session context. +""" +import json +import os +import re +import shutil +import subprocess +import sys +from pathlib import Path + +VERSION_RE = re.compile(r"(\d+)\.(\d+)\.(\d+)") +SECTION_RE = re.compile(r"^\[([^\]]+)\]", re.MULTILINE) +# default_profile inside the [__settings__] section ([^\[]*? keeps the search +# from crossing into the next section header). +SETTINGS_DEFAULT_RE = re.compile( + r"^\[__settings__\][^\[]*?^[ \t]*default_profile[ \t]*=[ \t]*(\S+)", + re.MULTILINE, +) +MAX_PROFILES = 12 + + +def cli_version(databricks): + """(major, minor, patch) from `databricks --version`, or None. 3s timeout.""" + try: + out = subprocess.run( + [databricks, "--version"], + capture_output=True, text=True, timeout=3, + ) + except Exception: + return None + m = VERSION_RE.search((out.stdout or "") + (out.stderr or "")) + return tuple(int(x) for x in m.groups()) if m else None + + +def config_profiles(): + """(config_path, [profile names], default_profile) read locally from the config file. + + Parsed directly rather than via `databricks auth profiles` on purpose: this + runs at SessionStart, which must stay offline and fast (no network, no + auth-validation round-trips). Skips CLI-internal sections like + `[__settings__]`, which are not auth profiles, but does surface + `[__settings__].default_profile` since the CLI resolves it when --profile + is omitted. + """ + cfg = os.environ.get("DATABRICKS_CONFIG_FILE") or str(Path.home() / ".databrickscfg") + try: + p = Path(cfg) + # Only read a regular file under a sane size cap, so a FIFO/device or a + # huge file pointed at by DATABRICKS_CONFIG_FILE can never hang or do + # unbounded work at session start. + if not p.is_file() or p.stat().st_size > 1_000_000: + return cfg, [], None + text = p.read_text(errors="replace") + except Exception: + return cfg, [], None + names = [ + n for n in SECTION_RE.findall(text) + if not (n.startswith("__") and n.endswith("__")) + ] + m = SETTINGS_DEFAULT_RE.search(text) + return cfg, names, (m.group(1) if m else None) + + +def _sanitize(value, limit=64): + """Make a config-derived string safe to inject as one context list item. + + Strips control chars / newlines (so a crafted profile name or env value + cannot inject extra bullets or instructions) and caps the length. + """ + s = re.sub(r"[\x00-\x1f\x7f]", " ", str(value)) + s = re.sub(r"\s+", " ", s).strip() + return s[: limit - 1].rstrip() + "…" if len(s) > limit else s + + +def build_context(): + """Return the context banner string, or '' to inject nothing.""" + databricks = shutil.which("databricks") + if not databricks: + return ( + "Databricks CLI (`databricks`) is not on PATH. The Databricks skills " + "and `/databricks:*` commands need it. Run `/databricks:setup` or see " + "the databricks-core skill to install it." + ) + + lines = [] + ver = cli_version(databricks) + if ver: + lines.append(f"CLI v{'.'.join(map(str, ver))}.") + else: + lines.append("CLI present (version unknown).") + + cfg, profiles, default_profile = config_profiles() + if profiles: + shown = [_sanitize(n) for n in profiles[:MAX_PROFILES]] + more = f" (+{len(profiles) - len(shown)} more)" if len(profiles) > len(shown) else "" + lines.append(f"Profiles in {_sanitize(Path(cfg).name)}: {', '.join(shown)}{more}.") + if default_profile: + lines.append( + f"Default profile (from [__settings__]): `{_sanitize(default_profile)}`; " + "the CLI uses it when `--profile` is omitted." + ) + lines.append("Never auto-select a profile. Pass `--profile ` and let the user choose.") + else: + # Basename only, matching the branch above: the full path (possibly a + # custom DATABRICKS_CONFIG_FILE) stays out of the injected context. + lines.append(f"No profiles found in {_sanitize(Path(cfg).name)}.") + + env_profile = os.environ.get("DATABRICKS_CONFIG_PROFILE") + if env_profile: + lines.append(f"DATABRICKS_CONFIG_PROFILE is set to `{_sanitize(env_profile)}`.") + if os.environ.get("DATABRICKS_HOST"): + authed = " with DATABRICKS_TOKEN set" if os.environ.get("DATABRICKS_TOKEN") else "" + lines.append(f"DATABRICKS_HOST is set{authed} (env / in-platform auth).") + + lines.append( + "Route Databricks-related work through the skills: load `databricks-core` " + "(the parent) plus the matching product skill." + ) + return "Databricks context:\n- " + "\n- ".join(lines) + + +def main(): + # One outer try so the fail-open guarantee covers the entire main block, + # including JSON serialization; the final print gets its own guard (a + # closed stdout must never break session startup either). + output = None + try: + try: + json.load(sys.stdin) # drain stdin; content unused + except Exception: + pass + ctx = build_context() + if ctx: + output = json.dumps({ + "hookSpecificOutput": { + "hookEventName": "SessionStart", + "additionalContext": ctx, + } + }) + except Exception: + output = None + try: + if output: + print(output) + except Exception: + pass + + +if __name__ == "__main__": + main() + sys.exit(0) diff --git a/hooks/databricks-router.py b/hooks/databricks-router.py new file mode 100644 index 0000000..3867137 --- /dev/null +++ b/hooks/databricks-router.py @@ -0,0 +1,217 @@ +#!/usr/bin/env python3 +"""UserPromptSubmit hook: route Databricks-related prompts into the skills. + +Reads the user prompt from stdin, runs a fast keyword regex (sub-50ms, no LLM, +no network), and if the prompt is Databricks-related, injects an +`additionalContext` instruction telling Claude to load the `databricks-core` +skill (the parent/router) plus the matching product skill before answering. + +There is no second agent to delegate to: Claude itself drives the `databricks` +CLI through the skills, so "routing" just means "make sure the Databricks skills +are loaded." No permission gating, no cost warnings. + +The full routing instruction is injected once per session (keyed by the +payload's session_id via a marker file in the temp dir); later Databricks +prompts in the same session get a one-line reminder instead, keeping repeat +token cost low. + +Contract (Claude Code UserPromptSubmit hook): + stdin : JSON, e.g. {"prompt": "...", "session_id": "..."} or {"message": "..."} + stdout: JSON -> hookSpecificOutput.additionalContext (injected before the turn), + or "{}" to stay out of the way. + Fail-open: on ANY error print "{}" and exit 0, so a broken hook never blocks a + prompt. +""" +import json +import re +import sys +import tempfile +from pathlib import Path + +# Unambiguously Databricks -> always route, even alongside a mention of an +# alternative platform (e.g. "migrate from redshift to databricks"). +STRONG = [ + r"\bdatabricks\b", + r"\bunity\s+catalog\b", + r"\blakeflow\b", + r"\blakebase\b", + r"\bdbfs\b", + r"\bdbutils\b", + r"\bdbsql\b", + r"\bdatabricks\.yml\b", + r"\basset\s+bundle\b", + r"\bdabs\b", + r"\b(?:lakeflow|spark)\s+declarative\s+pipelines?\b", + r"\bdelta\s+live\s+tables?\b", # legacy name for Declarative Pipelines + r"\bmosaic\s+ai\b", + r"\bdelta\s+sharing\b", + r"\bcloudfiles\b", +] + +# Databricks-likely but also used elsewhere -> route only when no +# alternative-platform / local-dev signal is present. +AMBIGUOUS = [ + r"\bgenie\b", + r"\bdelta\s+(lake|tables?)\b", + r"\bdeclarative\s+pipelines?\b", # bare form collides with Jenkins pipelines + r"\bmodel\s+serving\b", + r"\bvector\s+search\b", + r"\bmlflow\b", + r"\bpyspark\b", + r"\bspark\s*\.\s*(sql|read|write|table)\b", + r"\bserverless\s+(compute|warehouse|migration)\b", + r"\bmedallion\s+(architecture|tables?)\b", + r"\bsql\s+warehouse\b", + r"\bauto\s+loader\b", +] + +# Alternative data platforms + plainly-local dev work -> suppress an AMBIGUOUS +# match. (STRONG matches ignore this list.) +SUPPRESS = [ + r"\bbigquery\b", + r"\bredshift\b", + r"\bsynapse\b", + r"\bsnowflake\b", + r"\bgit\s+(commit|push|pull|status|log|diff|branch|rebase|merge|clone|stash)\b", + r"\b(read|edit|open|write|create|delete)\s+(the\s+|this\s+|a\s+|that\s+)?file\b", + r"\bunit\s+tests?\b", + r"\bnpm\b", + r"\bpip\s+install\b", + r"\bdocker\b", + r"\bkubernetes\b", + r"\bjenkins(?:file)?\b", +] + +_STRONG = [re.compile(p, re.IGNORECASE) for p in STRONG] +_AMBIGUOUS = [re.compile(p, re.IGNORECASE) for p in AMBIGUOUS] +_SUPPRESS = [re.compile(p, re.IGNORECASE) for p in SUPPRESS] + +# "databricks" inside a code-hosting URL (github.com/databricks/...) is an +# org/repo name, not product intent, so URLs are blanked before matching unless +# the hostname itself contains "databricks" (workspace and docs hosts), which +# keeps "why is https://myco.cloud.databricks.com/jobs/123 failing?" routing. +_URL_RE = re.compile( + r"(?:https?://|git@)(?P[\w.-]+)[/:]?\S*" + r"|\b(?:www\.)?(?P(?:github|gitlab|bitbucket)\.(?:com|org))[/:]\S*", + re.IGNORECASE, +) + + +def _strip_non_databricks_urls(text): + def _keep_or_blank(match): + host = match.group("host") or match.group("bare") or "" + return match.group(0) if "databricks" in host.lower() else " " + + return _URL_RE.sub(_keep_or_blank, text) + +ROUTING_INSTRUCTION = ( + "[DATABRICKS] This request is Databricks-related. Handle it through the " + "Databricks skills rather than ad hoc commands. Use the Skill tool to load " + "`databricks-core` first (the parent skill: CLI, auth, profile selection, " + "data exploration), then load the product skill that matches the request:\n" + "- Jobs / Lakeflow / workflows -> databricks-jobs\n" + "- Pipelines / Lakeflow Spark Declarative Pipelines (formerly DLT) -> " + "databricks-pipelines\n" + "- Apps / AppKit -> databricks-apps\n" + "- Asset Bundles / DABs / databricks.yml -> databricks-dabs\n" + "- Model Serving / endpoints -> databricks-model-serving\n" + "- Lakebase / Postgres -> databricks-lakebase\n" + "- Vector Search / RAG -> databricks-vector-search\n" + "- Classic-to-serverless migration -> databricks-serverless-migration\n" + "- Genie / natural-language data Q&A -> databricks-core (Genie CLI support " + "is experimental)\n" + "Then follow the skill's guidance (it drives the `databricks` CLI). If no " + "product skill fits, databricks-core alone is enough." +) + +# After the first routed prompt the skills are loaded (or being loaded), so the +# rest of the session gets this one-liner instead of the full block above. +ROUTING_REMINDER = ( + "[DATABRICKS] Databricks-related prompt: keep routing through the " + "Databricks skills (databricks-core plus the matching product skill); load " + "any that are not already loaded." +) + +_SESSION_ID_SAFE_RE = re.compile(r"[^A-Za-z0-9._-]") + + +def _marker_path(session_id): + """Temp-dir marker recording that this session already got the full instruction.""" + sid = _SESSION_ID_SAFE_RE.sub("", str(session_id or ""))[:64] + if not sid: + return None + return Path(tempfile.gettempdir()) / f"databricks-router-{sid}" + + +def check_prompt(prompt): + """Return the routing instruction if the prompt is Databricks-related, else None.""" + if not prompt or len(prompt.strip()) < 4: + return None + prompt = _strip_non_databricks_urls(prompt) + if any(p.search(prompt) for p in _STRONG): + return ROUTING_INSTRUCTION + if any(p.search(prompt) for p in _SUPPRESS): + return None + if any(p.search(prompt) for p in _AMBIGUOUS): + return ROUTING_INSTRUCTION + return None + + +def routing_context(prompt, session_id): + """Full instruction on the session's first Databricks prompt, reminder after.""" + if check_prompt(prompt) is None: + return None + marker = _marker_path(session_id) + if marker is None: + return ROUTING_INSTRUCTION + try: + if marker.exists(): + return ROUTING_REMINDER + marker.touch() + except Exception: + # Marker bookkeeping must never break routing itself. + pass + return ROUTING_INSTRUCTION + + +def extract_prompt(data): + """Pull the prompt text out of the hook payload (Claude or Codex shapes).""" + if not isinstance(data, dict): + return "" + prompt = data.get("prompt", data.get("message", "")) + if isinstance(prompt, dict): + prompt = prompt.get("content", "") + if isinstance(prompt, list): + prompt = " ".join( + block.get("text", "") for block in prompt if isinstance(block, dict) + ) + return str(prompt) + + +def main(): + # One outer try so the fail-open guarantee covers the entire main block, + # including JSON serialization; the final print gets its own guard (a + # closed stdout must not surface as a hook failure either). + output = "{}" + try: + data = json.load(sys.stdin) + session_id = data.get("session_id", "") if isinstance(data, dict) else "" + result = routing_context(extract_prompt(data), session_id) + if result: + output = json.dumps({ + "hookSpecificOutput": { + "hookEventName": "UserPromptSubmit", + "additionalContext": result, + } + }) + except Exception: + output = "{}" + try: + print(output) + except Exception: + pass + sys.exit(0) + + +if __name__ == "__main__": + main() diff --git a/hooks/hooks.json b/hooks/hooks.json new file mode 100644 index 0000000..22235f1 --- /dev/null +++ b/hooks/hooks.json @@ -0,0 +1,37 @@ +{ + "description": "Databricks plugin hooks: route Databricks prompts into the skills, prime session context, and hint on auth failures.", + "hooks": { + "UserPromptSubmit": [ + { + "hooks": [ + { + "type": "command", + "command": "python3 \"${CLAUDE_PLUGIN_ROOT}/hooks/databricks-router.py\" || python \"${CLAUDE_PLUGIN_ROOT}/hooks/databricks-router.py\" || true" + } + ] + } + ], + "SessionStart": [ + { + "matcher": "startup|clear|compact", + "hooks": [ + { + "type": "command", + "command": "python3 \"${CLAUDE_PLUGIN_ROOT}/hooks/databricks-context.py\" || python \"${CLAUDE_PLUGIN_ROOT}/hooks/databricks-context.py\" || true" + } + ] + } + ], + "PostToolUse": [ + { + "matcher": "Bash", + "hooks": [ + { + "type": "command", + "command": "python3 \"${CLAUDE_PLUGIN_ROOT}/hooks/databricks-auth-helper.py\" || python \"${CLAUDE_PLUGIN_ROOT}/hooks/databricks-auth-helper.py\" || true" + } + ] + } + ] + } +} diff --git a/scripts/skills.py b/scripts/skills.py index 0c8fc80..ce3f7d8 100644 --- a/scripts/skills.py +++ b/scripts/skills.py @@ -420,8 +420,14 @@ def validate_plugin_manifests(repo_root: Path) -> list[str]: if errors: return errors - plugin = json.loads(plugin_path.read_text()) - marketplace = json.loads(marketplace_path.read_text()) + try: + plugin = json.loads(plugin_path.read_text()) + except json.JSONDecodeError as exc: + return [f"{plugin_path.relative_to(repo_root)} is not valid JSON: {exc}"] + try: + marketplace = json.loads(marketplace_path.read_text()) + except json.JSONDecodeError as exc: + return [f"{marketplace_path.relative_to(repo_root)} is not valid JSON: {exc}"] keywords = {k.lower() for k in plugin.get("keywords", [])} @@ -460,6 +466,125 @@ def validate_plugin_manifests(repo_root: Path) -> list[str]: return errors +# --------------------------------------------------------------------------- +# Plugin components (hooks + commands) +# --------------------------------------------------------------------------- +# +# hooks/ and commands/ ship with the Claude Code plugin (the whole repo is the +# plugin via .claude-plugin/marketplace.json `source: "./"`), but they are NOT +# skills, so they live outside the manifest's skills map. These checks keep them +# honest without pulling them into the skill model. Stdlib-only, like the rest +# of this file, so the protected CI runner (no pypi) can run them. + +_HOOK_SCRIPT_RE = re.compile(r"\$\{CLAUDE_PLUGIN_ROOT\}/(\S+?\.py)") + + +def _norm_rel_path(path: str) -> str: + """Normalize a manifest-declared path for comparison ('./x' -> 'x').""" + path = path.strip() + while path.startswith("./"): + path = path[2:] + return path + + +def _read_frontmatter(md_path: Path) -> str | None: + """Return the YAML frontmatter block of a markdown file, or None if absent.""" + text = md_path.read_text() + if not text.startswith("---"): + return None + end = text.find("\n---", 3) + if end == -1: + return None + return text[3:end] + + +def check_plugin_components(repo_root: Path) -> list[str]: + """Validate the non-skill plugin components: hooks/ and commands/. + + - plugin.json must NOT declare "hooks": the standard hooks/hooks.json is + auto-loaded by Claude Code, so declaring it double-loads (a load error) + - plugin.json MUST declare "commands" when commands/ exists (this repo ships + commands via that manifest declaration) + - hooks/hooks.json must be valid JSON, and every ${CLAUDE_PLUGIN_ROOT}/*.py + script it references must exist + - every commands/*.md must have frontmatter carrying a `description`, and + the description must not contain an unquoted ':' (strict YAML parsers + reject it even though some frontmatter readers tolerate it) + + Returns a list of error strings (empty means all good). + """ + errors: list[str] = [] + + plugin_path = repo_root / ".claude-plugin" / "plugin.json" + try: + plugin = json.loads(plugin_path.read_text()) if plugin_path.exists() else {} + except json.JSONDecodeError as exc: + # validate_plugin_manifests reports the broken manifest itself; never + # crash here, just skip the manifest-dependent checks. + return [f".claude-plugin/plugin.json is not valid JSON: {exc}"] + + commands_dir = repo_root / "commands" + if commands_dir.is_dir(): + if "commands" not in plugin: + errors.append( + 'commands/ exists but .claude-plugin/plugin.json does not declare ' + '"commands": "./commands/". Add it, or the commands silently stop ' + "shipping." + ) + md_files = sorted(commands_dir.glob("*.md")) + if not md_files: + errors.append("commands/ exists but contains no *.md command files.") + for md in md_files: + frontmatter = _read_frontmatter(md) + if frontmatter is None: + errors.append( + f"Command 'commands/{md.name}' is missing YAML frontmatter." + ) + elif not re.search(r"^description:\s*\S", frontmatter, re.MULTILINE): + errors.append( + f"Command 'commands/{md.name}' frontmatter is missing a 'description'." + ) + elif re.search(r"^description:[ \t]*[^\s\"'>|].*:(?:\s|$)", frontmatter, re.MULTILINE): + errors.append( + f"Command 'commands/{md.name}' has an unquoted ':' in its " + "description, which strict YAML parsers reject. Quote the " + "whole description string." + ) + + hooks_json = repo_root / "hooks" / "hooks.json" + if hooks_json.exists(): + # Claude Code auto-loads the standard hooks/hooks.json. Declaring that + # same path in plugin.json double-loads it and fails the plugin with a + # "Duplicate hooks file" error, so the manifest must NOT reference it. + declared = plugin.get("hooks", []) + declared = [declared] if isinstance(declared, str) else declared + if isinstance(declared, list) and any( + _norm_rel_path(d) == "hooks/hooks.json" + for d in declared + if isinstance(d, str) + ): + errors.append( + 'plugin.json must not declare "hooks": "./hooks/hooks.json". The ' + "standard hooks/hooks.json is auto-loaded, so declaring it again " + 'double-loads it. Remove the "hooks" key (reserve manifest.hooks ' + "for additional, non-standard hook files)." + ) + try: + hooks_cfg = json.loads(hooks_json.read_text()) + except json.JSONDecodeError as exc: + errors.append(f"hooks/hooks.json is not valid JSON: {exc}") + hooks_cfg = None + if hooks_cfg is not None: + blob = json.dumps(hooks_cfg) + for rel in sorted(set(_HOOK_SCRIPT_RE.findall(blob))): + if not (repo_root / rel).exists(): + errors.append( + f"hooks/hooks.json references '{rel}' which does not exist." + ) + + return errors + + # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- @@ -527,6 +652,16 @@ def main() -> None: print(f" - {err}", file=sys.stderr) ok = False + component_errors = check_plugin_components(repo_root) + if component_errors: + print( + "ERROR: plugin components (hooks/ + commands/) are misconfigured:", + file=sys.stderr, + ) + for err in component_errors: + print(f" - {err}", file=sys.stderr) + ok = False + if not ok: print( "\nRun `python3 scripts/skills.py generate` to fix the " diff --git a/tests/databricks_auth_helper_test.py b/tests/databricks_auth_helper_test.py new file mode 100644 index 0000000..31a7d13 --- /dev/null +++ b/tests/databricks_auth_helper_test.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +"""Unit tests for the databricks-auth-helper PostToolUse hook. + +The helper should fire only for `databricks` Bash commands whose output looks +like an auth failure, and stay silent for everything else. Stdlib-only; run +the suite with: python3 -m unittest discover -s tests -p "*_test.py" +""" +import importlib.util +import unittest +from pathlib import Path + +_HOOKS_DIR = Path(__file__).resolve().parent.parent / "hooks" +_spec = importlib.util.spec_from_file_location( + "databricks_auth_helper", _HOOKS_DIR / "databricks-auth-helper.py" +) +helper = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(helper) + + +class CheckTest(unittest.TestCase): + def test_databricks_auth_failures_hint(self): + for text in [ + "Error: default auth: cannot configure default credentials", + 'oauth2: "invalid_grant" "Token was not recognized"', + "Error: 401 Unauthorized", + "Error: Invalid access token.", + "token is expired, please log in again", + "the refresh token was revoked by the server", + ]: + self.assertIsNotNone( + helper.check("Bash", "databricks jobs list", text), + f"should hint: {text!r}", + ) + + def test_clean_output_no_hint(self): + self.assertIsNone(helper.check("Bash", "databricks jobs list", '{"jobs": []}')) + + def test_non_databricks_command_no_hint(self): + self.assertIsNone(helper.check("Bash", "curl https://example.com", "401 Unauthorized")) + + def test_non_bash_tool_no_hint(self): + self.assertIsNone(helper.check("Read", "databricks", "401 Unauthorized")) + + def test_bare_status_code_in_data_no_hint(self): + # A bare 401 inside ordinary output is not an auth-failure signal. + self.assertIsNone(helper.check("Bash", "databricks jobs list", '{"row_id": 401}')) + + def test_empty_inputs_no_hint(self): + self.assertIsNone(helper.check("Bash", "", "401 unauthorized")) + self.assertIsNone(helper.check("Bash", "databricks auth env", "")) + + +class CommandDetectionTest(unittest.TestCase): + """`databricks` must be a segment executable, not a substring anywhere.""" + + AUTH_ERROR = "Error: 401 Unauthorized" + + def test_databricks_mentioned_but_not_invoked_no_hint(self): + # Observed false positives: gh commands against the databricks GitHub + # org whose output quoted auth-failure phrases (a PR body describing + # this hook, and this hook's own source fetched via the contents API). + for command in [ + "gh pr view 128 --repo databricks/databricks-agent-skills --json body", + 'gh api "repos/databricks/databricks-agent-skills/contents/hooks/databricks-auth-helper.py"', + "git clone https://github.com/databricks/cli", + "curl https://docs.databricks.com/api/auth.html", + "echo databricks", + "cat notes/databricks.md", + "/tmp/databricks-test clusters list", + ]: + self.assertIsNone( + helper.check("Bash", command, self.AUTH_ERROR), + f"should not hint: {command!r}", + ) + + def test_databricks_invoked_hint(self): + for command in [ + "databricks clusters list", + "cd repos/cli && databricks auth describe", + "databricks jobs list | head -5", + "/usr/local/bin/databricks --version", + "./databricks auth env", + "DATABRICKS_CONFIG_PROFILE=dev databricks current-user me", + "sudo -E databricks auth login", + "token=$(databricks auth token --host https://example.com)", + ]: + self.assertIsNotNone( + helper.check("Bash", command, self.AUTH_ERROR), + f"should hint: {command!r}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/databricks_context_test.py b/tests/databricks_context_test.py new file mode 100644 index 0000000..bccb6ee --- /dev/null +++ b/tests/databricks_context_test.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +"""Unit tests for the databricks-context SessionStart hook. + +Covers the pure pieces (config parsing, sanitization) and the build_context +wiring around them. Stdlib-only; run the suite with: +python3 -m unittest discover -s tests -p "*_test.py" +""" +import importlib.util +import os +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +_HOOKS_DIR = Path(__file__).resolve().parent.parent / "hooks" +_spec = importlib.util.spec_from_file_location( + "databricks_context", _HOOKS_DIR / "databricks-context.py" +) +ctx = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(ctx) + + +class ConfigProfilesTest(unittest.TestCase): + def _write_cfg(self, text): + f = tempfile.NamedTemporaryFile("w", suffix=".cfg", delete=False) + f.write(text) + f.close() + self.addCleanup(os.unlink, f.name) + return f.name + + def test_parses_profiles_and_skips_internal_sections(self): + cfg = self._write_cfg( + "[DEFAULT]\nhost = https://x\n\n[__settings__]\nfoo = 1\n\n[prod]\nhost = https://y\n" + ) + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": cfg}): + path, names, default = ctx.config_profiles() + self.assertEqual(path, cfg) + self.assertEqual(names, ["DEFAULT", "prod"]) + self.assertIsNone(default) + + def test_default_profile_from_settings(self): + cfg = self._write_cfg( + "[DEFAULT]\nhost = https://x\n\n[__settings__]\ndefault_profile = prod\n\n" + "[prod]\nhost = https://y\n" + ) + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": cfg}): + _, names, default = ctx.config_profiles() + self.assertEqual(names, ["DEFAULT", "prod"]) + self.assertEqual(default, "prod") + + def test_default_profile_key_outside_settings_ignored(self): + cfg = self._write_cfg( + "[__settings__]\nfoo = 1\n\n[prod]\ndefault_profile = nope\nhost = https://y\n" + ) + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": cfg}): + _, _, default = ctx.config_profiles() + self.assertIsNone(default) + + def test_missing_file_gives_no_profiles(self): + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": "/nonexistent/nope.cfg"}): + _, names, default = ctx.config_profiles() + self.assertEqual(names, []) + self.assertIsNone(default) + + def test_oversized_file_is_skipped(self): + cfg = self._write_cfg("[DEFAULT]\n" + "x = y\n" * 200_000) + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": cfg}): + _, names, _ = ctx.config_profiles() + self.assertEqual(names, []) + + +class SanitizeTest(unittest.TestCase): + def test_strips_control_chars_and_newlines(self): + self.assertEqual(ctx._sanitize("a\nb\x00c"), "a b c") + + def test_truncates_long_values(self): + out = ctx._sanitize("x" * 200, limit=10) + self.assertTrue(out.endswith("…")) + self.assertLessEqual(len(out), 10) + + def test_plain_value_unchanged(self): + self.assertEqual(ctx._sanitize("e2-dogfood"), "e2-dogfood") + + +class BuildContextTest(unittest.TestCase): + def test_cli_missing_points_at_setup(self): + with mock.patch.object(ctx.shutil, "which", return_value=None): + out = ctx.build_context() + self.assertIn("/databricks:setup", out) + + def test_no_profiles_message_uses_basename_only(self): + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": "/secret/dir/custom.cfg"}), \ + mock.patch.object(ctx.shutil, "which", return_value="/usr/bin/databricks"), \ + mock.patch.object(ctx, "cli_version", return_value=(1, 2, 3)): + out = ctx.build_context() + self.assertIn("custom.cfg", out) + self.assertNotIn("/secret/dir", out) + + def test_default_profile_shown_in_banner(self): + f = tempfile.NamedTemporaryFile("w", suffix=".cfg", delete=False) + f.write("[__settings__]\ndefault_profile = prod\n\n[prod]\nhost = https://y\n") + f.close() + self.addCleanup(os.unlink, f.name) + with mock.patch.dict(os.environ, {"DATABRICKS_CONFIG_FILE": f.name}), \ + mock.patch.object(ctx.shutil, "which", return_value="/usr/bin/databricks"), \ + mock.patch.object(ctx, "cli_version", return_value=(1, 2, 3)): + out = ctx.build_context() + self.assertIn("Default profile", out) + self.assertIn("`prod`", out) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/databricks_router_test.py b/tests/databricks_router_test.py new file mode 100644 index 0000000..1ffdbab --- /dev/null +++ b/tests/databricks_router_test.py @@ -0,0 +1,155 @@ +#!/usr/bin/env python3 +"""Unit tests for the databricks-router UserPromptSubmit hook. + +The router decides which prompts get steered into the Databricks skills, so its +precision is pinned here (over-routing is annoying; under-routing misses work). +Stdlib-only; run the suite with: python3 -m unittest discover -s tests -p "*_test.py" +""" +import importlib.util +import tempfile +import unittest +from pathlib import Path +from unittest import mock + +_HOOKS_DIR = Path(__file__).resolve().parent.parent / "hooks" +_spec = importlib.util.spec_from_file_location( + "databricks_router", _HOOKS_DIR / "databricks-router.py" +) +router = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(router) + + +class CheckPromptTest(unittest.TestCase): + def assertRoutes(self, prompt): + self.assertIsNotNone(router.check_prompt(prompt), f"should route: {prompt!r}") + + def assertSkips(self, prompt): + self.assertIsNone(router.check_prompt(prompt), f"should skip: {prompt!r}") + + def test_strong_routes(self): + for p in [ + "How do I deploy a Databricks app?", + "create a unity catalog grant on the sales schema", + "set up a lakeflow job", + "write this dataframe to dbfs", + "validate my databricks.yml asset bundle", + "deploy my dabs to the dev target", + "build a delta live tables pipeline", + ]: + self.assertRoutes(p) + + def test_strong_routes_even_with_alternative_platform(self): + # "databricks" present -> route despite the alternative-platform mention. + self.assertRoutes("migrate my tables from redshift to databricks") + self.assertRoutes("migrate from snowflake to databricks") + + def test_new_strong_terms_route(self): + self.assertRoutes("share this table via delta sharing") + self.assertRoutes("ingest with the cloudFiles format") + + def test_declarative_pipelines_routes(self): + # Current branding for DLT. The bare phrase is AMBIGUOUS (Jenkins also + # has "declarative pipelines"), so Jenkins-flavored prompts stay out. + self.assertRoutes("rewrite this notebook as a spark declarative pipeline") + self.assertRoutes("create a declarative pipeline for the bronze layer") + self.assertSkips("convert this Jenkinsfile to a declarative pipeline") + self.assertSkips("write a jenkins declarative pipeline for the CI build") + + def test_new_ambiguous_terms(self): + self.assertRoutes("create a serverless sql warehouse") + self.assertRoutes("set up auto loader for streaming ingestion") + self.assertSkips("set up a sql warehouse in snowflake") + # One-word "autoloader" (PHP/composer style) must not match. + self.assertSkips("fix the php autoloader config") + + def test_ambiguous_routes_without_alternative_platform(self): + for p in [ + "set up a model serving endpoint", + "create a vector search index for RAG", + "build a medallion architecture for my tables", + "ask Genie about revenue", + ]: + self.assertRoutes(p) + + def test_ambiguous_suppressed_by_alternative_platform(self): + self.assertSkips("set up a model serving endpoint in sagemaker for redshift data") + self.assertSkips("use bigquery for vector search") + + def test_local_dev_skips(self): + for p in [ + "git commit -m 'fix'", + "read the file src/main.py", + "write a unit test for this function", + "npm install react", + "pip install requests", + "build a docker image", + ]: + self.assertSkips(p) + + def test_unrelated_skips(self): + for p in ["hello", "what's the weather", "refactor this react component", "ok", + "explain photon energy in physics"]: + self.assertSkips(p) + + def test_too_short_skips(self): + self.assertSkips("db") + self.assertSkips("") + + def test_code_host_urls_do_not_route(self): + # "databricks" as a GitHub org/repo name is not product intent. + for p in [ + "review https://github.com/databricks/databricks-agent-skills/pull/128 please", + "what changed in github.com/databricks/cli recently?", + "clone git@github.com:databricks/terraform-provider-databricks.git", + ]: + self.assertSkips(p) + + def test_workspace_urls_still_route(self): + # Hostname contains "databricks" -> real product signal. + self.assertRoutes("why is https://myco.cloud.databricks.com/jobs/123 failing?") + + def test_url_plus_real_intent_routes(self): + # Only the URL is blanked; intent outside it still routes. + self.assertRoutes( + "review https://github.com/databricks/cli/pull/5 and then deploy the databricks job" + ) + + def test_extract_prompt_shapes(self): + self.assertEqual(router.extract_prompt({"prompt": "hi"}), "hi") + self.assertEqual(router.extract_prompt({"message": "yo"}), "yo") + self.assertEqual(router.extract_prompt({"prompt": {"content": "x"}}), "x") + self.assertEqual( + router.extract_prompt({"prompt": [{"text": "a"}, {"text": "b"}]}), "a b" + ) + + +class SessionMemoTest(unittest.TestCase): + def test_first_route_full_then_reminder(self): + with tempfile.TemporaryDirectory() as tmp, \ + mock.patch.object(router.tempfile, "gettempdir", return_value=tmp): + first = router.routing_context("deploy a databricks job", "sess-1") + second = router.routing_context("update that databricks job", "sess-1") + other = router.routing_context("deploy a databricks job", "sess-2") + self.assertEqual(first, router.ROUTING_INSTRUCTION) + self.assertEqual(second, router.ROUTING_REMINDER) + self.assertEqual(other, router.ROUTING_INSTRUCTION) + + def test_non_databricks_prompt_does_not_mark_session(self): + with tempfile.TemporaryDirectory() as tmp, \ + mock.patch.object(router.tempfile, "gettempdir", return_value=tmp): + self.assertIsNone(router.routing_context("hello there friend", "sess-3")) + self.assertEqual( + router.routing_context("deploy a databricks job", "sess-3"), + router.ROUTING_INSTRUCTION, + ) + + def test_missing_session_id_always_full_instruction(self): + for sid in (None, "", "!!!"): + self.assertEqual( + router.routing_context("deploy a databricks job", sid), + router.ROUTING_INSTRUCTION, + ) + + +if __name__ == "__main__": + unittest.main()