databricks-mcp-toolkit 0.2.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,10 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """MCP server for Databricks."""
7
+
8
+ __version__ = "0.2.0"
9
+
10
+ __all__ = ["__version__"]
@@ -0,0 +1,10 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ import sys
7
+
8
+ from databricks_mcp.cli import main
9
+
10
+ sys.exit(main())
@@ -0,0 +1,30 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """Thin wrappers over Databricks REST endpoints. Each function takes a DatabricksClient."""
7
+
8
+ from databricks_mcp.api import (
9
+ clusters,
10
+ command_execution,
11
+ dbfs,
12
+ files,
13
+ jobs,
14
+ sql,
15
+ unity_catalog,
16
+ warehouses,
17
+ workspace,
18
+ )
19
+
20
+ __all__ = [
21
+ "clusters",
22
+ "command_execution",
23
+ "dbfs",
24
+ "files",
25
+ "jobs",
26
+ "sql",
27
+ "unity_catalog",
28
+ "warehouses",
29
+ "workspace",
30
+ ]
@@ -0,0 +1,73 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """Clusters API 2.1."""
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import Any
11
+
12
+ from databricks_mcp.client import DatabricksClient
13
+
14
+ BASE = "/api/2.1/clusters"
15
+
16
+
17
+ async def list_clusters(
18
+ client: DatabricksClient,
19
+ page_size: int | None = None,
20
+ page_token: str | None = None,
21
+ ) -> dict[str, Any]:
22
+ return await client.get(f"{BASE}/list", {"page_size": page_size, "page_token": page_token})
23
+
24
+
25
+ async def get_cluster(client: DatabricksClient, cluster_id: str) -> dict[str, Any]:
26
+ return await client.get(f"{BASE}/get", {"cluster_id": cluster_id})
27
+
28
+
29
+ async def create_cluster(client: DatabricksClient, spec: dict[str, Any]) -> dict[str, Any]:
30
+ return await client.post(f"{BASE}/create", spec)
31
+
32
+
33
+ async def start_cluster(client: DatabricksClient, cluster_id: str) -> dict[str, Any]:
34
+ return await client.post(f"{BASE}/start", {"cluster_id": cluster_id})
35
+
36
+
37
+ async def restart_cluster(client: DatabricksClient, cluster_id: str) -> dict[str, Any]:
38
+ return await client.post(f"{BASE}/restart", {"cluster_id": cluster_id})
39
+
40
+
41
+ async def resize_cluster(
42
+ client: DatabricksClient,
43
+ cluster_id: str,
44
+ num_workers: int | None = None,
45
+ autoscale_min_workers: int | None = None,
46
+ autoscale_max_workers: int | None = None,
47
+ ) -> dict[str, Any]:
48
+ body: dict[str, Any] = {"cluster_id": cluster_id}
49
+ if autoscale_min_workers is not None or autoscale_max_workers is not None:
50
+ if autoscale_min_workers is None or autoscale_max_workers is None:
51
+ raise ValueError("Set both autoscale_min_workers and autoscale_max_workers.")
52
+ body["autoscale"] = {
53
+ "min_workers": autoscale_min_workers,
54
+ "max_workers": autoscale_max_workers,
55
+ }
56
+ elif num_workers is not None:
57
+ body["num_workers"] = num_workers
58
+ else:
59
+ raise ValueError("Pass num_workers, or both autoscale_min_workers and autoscale_max_workers.")
60
+ return await client.post(f"{BASE}/resize", body)
61
+
62
+
63
+ async def terminate_cluster(client: DatabricksClient, cluster_id: str) -> dict[str, Any]:
64
+ # "delete" in the API terminates; the config is kept and the cluster can be restarted.
65
+ return await client.post(f"{BASE}/delete", {"cluster_id": cluster_id})
66
+
67
+
68
+ async def list_spark_versions(client: DatabricksClient) -> dict[str, Any]:
69
+ return await client.get(f"{BASE}/spark-versions")
70
+
71
+
72
+ async def list_node_types(client: DatabricksClient) -> dict[str, Any]:
73
+ return await client.get(f"{BASE}/list-node-types")
@@ -0,0 +1,92 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """Command Execution API 1.2. Still the current (and only) version for running code
7
+ on an all-purpose cluster with state kept between commands."""
8
+
9
+ from __future__ import annotations
10
+
11
+ import asyncio
12
+ import time
13
+ from typing import Any
14
+
15
+ from databricks_mcp.client import DatabricksClient
16
+
17
+ BASE = "/api/1.2"
18
+ _RUNNING = {"Queued", "Running", "Cancelling"}
19
+
20
+
21
+ async def create_context(client: DatabricksClient, cluster_id: str, language: str) -> dict[str, Any]:
22
+ return await client.post(f"{BASE}/contexts/create", {"clusterId": cluster_id, "language": language})
23
+
24
+
25
+ async def context_status(client: DatabricksClient, cluster_id: str, context_id: str) -> dict[str, Any]:
26
+ return await client.get(f"{BASE}/contexts/status", {"clusterId": cluster_id, "contextId": context_id})
27
+
28
+
29
+ async def destroy_context(client: DatabricksClient, cluster_id: str, context_id: str) -> dict[str, Any]:
30
+ return await client.post(f"{BASE}/contexts/destroy", {"clusterId": cluster_id, "contextId": context_id})
31
+
32
+
33
+ async def execute(
34
+ client: DatabricksClient, cluster_id: str, context_id: str, command: str, language: str
35
+ ) -> dict[str, Any]:
36
+ return await client.post(
37
+ f"{BASE}/commands/execute",
38
+ {"clusterId": cluster_id, "contextId": context_id, "command": command, "language": language},
39
+ )
40
+
41
+
42
+ async def command_status(client: DatabricksClient, cluster_id: str, context_id: str, command_id: str) -> dict[str, Any]:
43
+ return await client.get(
44
+ f"{BASE}/commands/status",
45
+ {"clusterId": cluster_id, "contextId": context_id, "commandId": command_id},
46
+ )
47
+
48
+
49
+ async def cancel(client: DatabricksClient, cluster_id: str, context_id: str, command_id: str) -> dict[str, Any]:
50
+ return await client.post(
51
+ f"{BASE}/commands/cancel",
52
+ {"clusterId": cluster_id, "contextId": context_id, "commandId": command_id},
53
+ )
54
+
55
+
56
+ async def wait_for_command(
57
+ client: DatabricksClient,
58
+ cluster_id: str,
59
+ context_id: str,
60
+ command_id: str,
61
+ timeout_seconds: float = 120,
62
+ poll_interval_seconds: float = 1,
63
+ ) -> dict[str, Any]:
64
+ deadline = time.monotonic() + timeout_seconds
65
+ while True:
66
+ status = await command_status(client, cluster_id, context_id, command_id)
67
+ if status.get("status") not in _RUNNING or time.monotonic() >= deadline:
68
+ return status
69
+ await asyncio.sleep(poll_interval_seconds)
70
+
71
+
72
+ async def run_once(
73
+ client: DatabricksClient,
74
+ cluster_id: str,
75
+ command: str,
76
+ language: str,
77
+ timeout_seconds: float = 120,
78
+ poll_interval_seconds: float = 1,
79
+ ) -> dict[str, Any]:
80
+ """Create a context, run one command, wait, and always clean the context up."""
81
+ context = await create_context(client, cluster_id, language)
82
+ context_id = context["id"]
83
+ try:
84
+ started = await execute(client, cluster_id, context_id, command, language)
85
+ return await wait_for_command(
86
+ client, cluster_id, context_id, started["id"], timeout_seconds, poll_interval_seconds
87
+ )
88
+ finally:
89
+ try:
90
+ await destroy_context(client, cluster_id, context_id)
91
+ except Exception: # cleanup failure shouldn't hide the command's result
92
+ pass
@@ -0,0 +1,37 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """DBFS API 2.0. Legacy: Databricks recommends Unity Catalog volumes (see files.py).
7
+ I only keep read access here for workspaces that still have data in DBFS."""
8
+
9
+ from __future__ import annotations
10
+
11
+ import base64
12
+ from typing import Any
13
+
14
+ from databricks_mcp.client import DatabricksClient
15
+
16
+ BASE = "/api/2.0/dbfs"
17
+ _MAX_READ = 1_048_576 # the API's own per-call ceiling
18
+
19
+
20
+ async def list_files(client: DatabricksClient, path: str) -> dict[str, Any]:
21
+ return await client.get(f"{BASE}/list", {"path": path})
22
+
23
+
24
+ async def read_file(client: DatabricksClient, path: str, max_bytes: int = 100_000) -> dict[str, Any]:
25
+ length = max(1, min(max_bytes, _MAX_READ))
26
+ result = await client.get(f"{BASE}/read", {"path": path, "offset": 0, "length": length})
27
+ raw = base64.b64decode(result.get("data") or "")
28
+ out: dict[str, Any] = {"path": path, "bytes_read": result.get("bytes_read", len(raw))}
29
+ try:
30
+ out["content"] = raw.decode("utf-8")
31
+ except UnicodeDecodeError:
32
+ out["content"] = None
33
+ out["note"] = "Binary file; content not shown."
34
+ status = await client.get(f"{BASE}/get-status", {"path": path})
35
+ out["total_bytes"] = status.get("file_size")
36
+ out["truncated"] = (status.get("file_size") or 0) > len(raw)
37
+ return out
@@ -0,0 +1,94 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """Files API 2.0 for Unity Catalog volumes (/Volumes/<catalog>/<schema>/<volume>/...).
7
+
8
+ This is what Databricks recommends over DBFS now. Files API 3.0 exists but is still in
9
+ beta, needs an admin toggle and a regional hostname, so I'm staying on 2.0.
10
+ """
11
+
12
+ from __future__ import annotations
13
+
14
+ from typing import Any
15
+ from urllib.parse import quote
16
+
17
+ from databricks_mcp.client import DatabricksClient
18
+
19
+ BASE = "/api/2.0/fs"
20
+
21
+
22
+ def _path(path: str) -> str:
23
+ if not path.startswith("/Volumes/"):
24
+ raise ValueError("Files API paths must be absolute volume paths, e.g. /Volumes/main/default/my_vol/a.csv")
25
+ return quote(path, safe="/")
26
+
27
+
28
+ async def list_directory(
29
+ client: DatabricksClient,
30
+ path: str,
31
+ page_size: int | None = None,
32
+ page_token: str | None = None,
33
+ ) -> dict[str, Any]:
34
+ directory = _path(path.rstrip("/") + "/")
35
+ return await client.get(f"{BASE}/directories{directory}", {"page_size": page_size, "page_token": page_token})
36
+
37
+
38
+ async def read_file(client: DatabricksClient, path: str, max_bytes: int = 100_000) -> dict[str, Any]:
39
+ # Range keeps me from pulling a 5 GB file into memory just to show the first page.
40
+ response = await client.request(
41
+ "GET",
42
+ f"{BASE}/files{_path(path)}",
43
+ headers={"Range": f"bytes=0-{max_bytes - 1}"},
44
+ raw=True,
45
+ )
46
+ data = response.content[:max_bytes]
47
+ total = _total_size(response) or len(data)
48
+ result: dict[str, Any] = {
49
+ "path": path,
50
+ "content_type": response.headers.get("content-type"),
51
+ "bytes_returned": len(data),
52
+ "total_bytes": total,
53
+ "truncated": total > len(data),
54
+ }
55
+ try:
56
+ result["content"] = data.decode("utf-8")
57
+ except UnicodeDecodeError:
58
+ result["content"] = None
59
+ result["note"] = "Binary file; content not shown."
60
+ return result
61
+
62
+
63
+ def _total_size(response: Any) -> int | None:
64
+ content_range = response.headers.get("content-range") # "bytes 0-99/12345"
65
+ if content_range and "/" in content_range:
66
+ tail = content_range.rsplit("/", 1)[1]
67
+ if tail.isdigit():
68
+ return int(tail)
69
+ length = response.headers.get("content-length")
70
+ if length and length.isdigit() and response.status_code == 200:
71
+ return int(length)
72
+ return None
73
+
74
+
75
+ async def write_file(client: DatabricksClient, path: str, content: str, overwrite: bool = False) -> dict[str, Any]:
76
+ await client.request(
77
+ "PUT",
78
+ f"{BASE}/files{_path(path)}",
79
+ params={"overwrite": "true" if overwrite else "false"},
80
+ content=content.encode("utf-8"),
81
+ headers={"Content-Type": "application/octet-stream"},
82
+ raw=True,
83
+ )
84
+ return {"path": path, "written_bytes": len(content.encode("utf-8"))}
85
+
86
+
87
+ async def delete_file(client: DatabricksClient, path: str) -> dict[str, Any]:
88
+ await client.request("DELETE", f"{BASE}/files{_path(path)}", raw=True)
89
+ return {"path": path, "deleted": True}
90
+
91
+
92
+ async def create_directory(client: DatabricksClient, path: str) -> dict[str, Any]:
93
+ await client.request("PUT", f"{BASE}/directories{_path(path.rstrip('/') + '/')}", raw=True)
94
+ return {"path": path, "created": True}
@@ -0,0 +1,144 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """Jobs API 2.2 (current). 2.2 pages tasks/job_clusters past 100 entries and queues jobs by default."""
7
+
8
+ from __future__ import annotations
9
+
10
+ import asyncio
11
+ import time
12
+ from typing import Any
13
+
14
+ from databricks_mcp.client import DatabricksAPIError, DatabricksClient
15
+
16
+ BASE = "/api/2.2/jobs"
17
+ _TERMINAL_LIFE_CYCLE = {"TERMINATED", "SKIPPED", "INTERNAL_ERROR"}
18
+ _MAX_TASK_OUTPUTS = 20
19
+
20
+
21
+ async def list_jobs(
22
+ client: DatabricksClient,
23
+ name: str | None = None,
24
+ limit: int = 25,
25
+ page_token: str | None = None,
26
+ expand_tasks: bool = False,
27
+ ) -> dict[str, Any]:
28
+ return await client.get(
29
+ f"{BASE}/list",
30
+ {
31
+ "name": name,
32
+ "limit": max(1, min(limit, 100)),
33
+ "page_token": page_token,
34
+ "expand_tasks": expand_tasks or None,
35
+ },
36
+ )
37
+
38
+
39
+ async def get_job(client: DatabricksClient, job_id: int, page_token: str | None = None) -> dict[str, Any]:
40
+ return await client.get(f"{BASE}/get", {"job_id": job_id, "page_token": page_token})
41
+
42
+
43
+ async def run_job(
44
+ client: DatabricksClient,
45
+ job_id: int,
46
+ job_parameters: dict[str, str] | None = None,
47
+ notebook_params: dict[str, str] | None = None,
48
+ ) -> dict[str, Any]:
49
+ body: dict[str, Any] = {"job_id": job_id}
50
+ if job_parameters:
51
+ body["job_parameters"] = job_parameters
52
+ if notebook_params:
53
+ body["notebook_params"] = notebook_params
54
+ return await client.post(f"{BASE}/run-now", body)
55
+
56
+
57
+ async def submit_run(client: DatabricksClient, run_spec: dict[str, Any]) -> dict[str, Any]:
58
+ if not run_spec.get("tasks"):
59
+ raise ValueError("run_spec needs a non-empty 'tasks' list.")
60
+ return await client.post(f"{BASE}/runs/submit", run_spec)
61
+
62
+
63
+ async def list_runs(
64
+ client: DatabricksClient,
65
+ job_id: int | None = None,
66
+ active_only: bool = False,
67
+ completed_only: bool = False,
68
+ limit: int = 25,
69
+ page_token: str | None = None,
70
+ ) -> dict[str, Any]:
71
+ return await client.get(
72
+ f"{BASE}/runs/list",
73
+ {
74
+ "job_id": job_id,
75
+ "active_only": active_only or None,
76
+ "completed_only": completed_only or None,
77
+ "limit": max(1, min(limit, 25)),
78
+ "page_token": page_token,
79
+ },
80
+ )
81
+
82
+
83
+ async def get_run(client: DatabricksClient, run_id: int, page_token: str | None = None) -> dict[str, Any]:
84
+ return await client.get(f"{BASE}/runs/get", {"run_id": run_id, "page_token": page_token})
85
+
86
+
87
+ async def cancel_run(client: DatabricksClient, run_id: int) -> dict[str, Any]:
88
+ return await client.post(f"{BASE}/runs/cancel", {"run_id": run_id})
89
+
90
+
91
+ async def get_run_output(client: DatabricksClient, run_id: int) -> dict[str, Any]:
92
+ """runs/get-output only works on task runs, not on a multi-task parent run.
93
+ If I'm handed a parent run, I fetch each task's output instead."""
94
+ run = await get_run(client, run_id)
95
+ tasks = run.get("tasks") or []
96
+ task_runs = [t for t in tasks if t.get("run_id") and t["run_id"] != run_id]
97
+ if not task_runs:
98
+ return await client.get(f"{BASE}/runs/get-output", {"run_id": run_id})
99
+
100
+ outputs = []
101
+ for task in task_runs[:_MAX_TASK_OUTPUTS]:
102
+ task_run_id = task["run_id"]
103
+ try:
104
+ out = await client.get(f"{BASE}/runs/get-output", {"run_id": task_run_id})
105
+ except DatabricksAPIError as exc:
106
+ out = {"error": str(exc)}
107
+ out.pop("metadata", None) # duplicates what get_run already returns
108
+ outputs.append({"task_key": task.get("task_key"), "run_id": task_run_id, **out})
109
+ result: dict[str, Any] = {"run_id": run_id, "state": _state(run), "tasks": outputs}
110
+ if len(task_runs) > _MAX_TASK_OUTPUTS:
111
+ result["note"] = f"Showing {_MAX_TASK_OUTPUTS} of {len(task_runs)} task outputs."
112
+ return result
113
+
114
+
115
+ def _state(run: dict[str, Any]) -> dict[str, Any]:
116
+ return {"status": run.get("status"), "state": run.get("state")}
117
+
118
+
119
+ def is_terminal(run: dict[str, Any]) -> bool:
120
+ status = (run.get("status") or {}).get("state")
121
+ life_cycle = (run.get("state") or {}).get("life_cycle_state")
122
+ return status == "TERMINATED" or life_cycle in _TERMINAL_LIFE_CYCLE
123
+
124
+
125
+ async def wait_for_run(
126
+ client: DatabricksClient,
127
+ run_id: int,
128
+ poll_interval_seconds: float = 10,
129
+ timeout_seconds: float = 600,
130
+ ) -> dict[str, Any]:
131
+ deadline = time.monotonic() + timeout_seconds
132
+ while True:
133
+ run = await get_run(client, run_id)
134
+ if is_terminal(run):
135
+ return {"run_id": run_id, "finished": True, **_state(run), "run_page_url": run.get("run_page_url")}
136
+ if time.monotonic() >= deadline:
137
+ return {
138
+ "run_id": run_id,
139
+ "finished": False,
140
+ **_state(run),
141
+ "run_page_url": run.get("run_page_url"),
142
+ "note": f"Still running after {timeout_seconds}s. Call again to keep waiting.",
143
+ }
144
+ await asyncio.sleep(poll_interval_seconds)
@@ -0,0 +1,99 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """SQL Statement Execution API 2.0."""
7
+
8
+ from __future__ import annotations
9
+
10
+ import asyncio
11
+ import time
12
+ from typing import Any
13
+
14
+ from databricks_mcp.client import DatabricksClient
15
+
16
+ BASE = "/api/2.0/sql/statements"
17
+ _TERMINAL = {"SUCCEEDED", "FAILED", "CANCELED", "CLOSED"}
18
+ # The API lets a request block for at most 50s; after that I poll.
19
+ _MAX_SERVER_WAIT = 50
20
+
21
+
22
+ async def execute(
23
+ client: DatabricksClient,
24
+ statement: str,
25
+ warehouse_id: str,
26
+ catalog: str | None = None,
27
+ schema: str | None = None,
28
+ parameters: dict[str, str] | None = None,
29
+ row_limit: int = 1000,
30
+ timeout_seconds: int = 120,
31
+ poll_interval_seconds: float = 2,
32
+ ) -> dict[str, Any]:
33
+ body: dict[str, Any] = {
34
+ "statement": statement,
35
+ "warehouse_id": warehouse_id,
36
+ "disposition": "INLINE",
37
+ "format": "JSON_ARRAY",
38
+ "row_limit": row_limit,
39
+ "wait_timeout": f"{max(5, min(timeout_seconds, _MAX_SERVER_WAIT))}s",
40
+ "on_wait_timeout": "CONTINUE",
41
+ }
42
+ if catalog:
43
+ body["catalog"] = catalog
44
+ if schema:
45
+ body["schema"] = schema
46
+ if parameters:
47
+ body["parameters"] = [{"name": k, "value": v} for k, v in parameters.items()]
48
+
49
+ started = time.monotonic()
50
+ response = await client.post(BASE, body)
51
+ statement_id = response.get("statement_id")
52
+ while _state(response) not in _TERMINAL and statement_id:
53
+ if time.monotonic() - started >= timeout_seconds:
54
+ break
55
+ await asyncio.sleep(poll_interval_seconds)
56
+ response = await client.get(f"{BASE}/{statement_id}")
57
+ return summarize(response)
58
+
59
+
60
+ async def get_statement(client: DatabricksClient, statement_id: str) -> dict[str, Any]:
61
+ return summarize(await client.get(f"{BASE}/{statement_id}"))
62
+
63
+
64
+ async def cancel_statement(client: DatabricksClient, statement_id: str) -> dict[str, Any]:
65
+ await client.post(f"{BASE}/{statement_id}/cancel")
66
+ return {"statement_id": statement_id, "cancel_requested": True}
67
+
68
+
69
+ def _state(response: dict[str, Any]) -> str | None:
70
+ return (response.get("status") or {}).get("state")
71
+
72
+
73
+ def summarize(response: dict[str, Any]) -> dict[str, Any]:
74
+ """Flatten the API's manifest/result split into columns + rows the model can read."""
75
+ status = response.get("status") or {}
76
+ state = status.get("state")
77
+ out: dict[str, Any] = {"statement_id": response.get("statement_id"), "state": state}
78
+ if status.get("error"):
79
+ out["error"] = status["error"]
80
+ if state not in _TERMINAL:
81
+ out["note"] = "Still running. Call get_sql_statement with this statement_id to check again."
82
+ return out
83
+
84
+ manifest = response.get("manifest") or {}
85
+ columns = (manifest.get("schema") or {}).get("columns") or []
86
+ result = response.get("result") or {}
87
+ rows = result.get("data_array") or []
88
+ out["columns"] = [{"name": c.get("name"), "type": c.get("type_text") or c.get("type_name")} for c in columns]
89
+ out["rows"] = rows
90
+ out["row_count"] = len(rows)
91
+ total = manifest.get("total_row_count")
92
+ if total is not None:
93
+ out["total_row_count"] = total
94
+ out["truncated"] = bool(
95
+ manifest.get("truncated")
96
+ or result.get("next_chunk_index") is not None
97
+ or (total is not None and total > len(rows))
98
+ )
99
+ return out
@@ -0,0 +1,76 @@
1
+ # Copyright (c) 2025-2026 Kush Patel
2
+ # Dual-licensed: AGPL-3.0-only OR commercial license.
3
+ # If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
4
+ # For commercial terms, see the project README.
5
+
6
+ """Unity Catalog API 2.1 (read-only browsing: catalogs, schemas, tables, volumes)."""
7
+
8
+ from __future__ import annotations
9
+
10
+ from typing import Any
11
+ from urllib.parse import quote
12
+
13
+ from databricks_mcp.client import DatabricksClient
14
+
15
+ BASE = "/api/2.1/unity-catalog"
16
+
17
+
18
+ async def list_catalogs(
19
+ client: DatabricksClient, max_results: int = 100, page_token: str | None = None
20
+ ) -> dict[str, Any]:
21
+ return await client.get(f"{BASE}/catalogs", {"max_results": max_results, "page_token": page_token})
22
+
23
+
24
+ async def list_schemas(
25
+ client: DatabricksClient,
26
+ catalog_name: str,
27
+ max_results: int = 100,
28
+ page_token: str | None = None,
29
+ ) -> dict[str, Any]:
30
+ return await client.get(
31
+ f"{BASE}/schemas",
32
+ {"catalog_name": catalog_name, "max_results": max_results, "page_token": page_token},
33
+ )
34
+
35
+
36
+ async def list_tables(
37
+ client: DatabricksClient,
38
+ catalog_name: str,
39
+ schema_name: str,
40
+ max_results: int = 100,
41
+ page_token: str | None = None,
42
+ ) -> dict[str, Any]:
43
+ # omit_columns keeps the listing small; get_table has the full column list.
44
+ return await client.get(
45
+ f"{BASE}/tables",
46
+ {
47
+ "catalog_name": catalog_name,
48
+ "schema_name": schema_name,
49
+ "max_results": max_results,
50
+ "page_token": page_token,
51
+ "omit_columns": True,
52
+ "omit_properties": True,
53
+ },
54
+ )
55
+
56
+
57
+ async def get_table(client: DatabricksClient, full_name: str) -> dict[str, Any]:
58
+ return await client.get(f"{BASE}/tables/{quote(full_name, safe='')}")
59
+
60
+
61
+ async def list_volumes(
62
+ client: DatabricksClient,
63
+ catalog_name: str,
64
+ schema_name: str,
65
+ max_results: int = 100,
66
+ page_token: str | None = None,
67
+ ) -> dict[str, Any]:
68
+ return await client.get(
69
+ f"{BASE}/volumes",
70
+ {
71
+ "catalog_name": catalog_name,
72
+ "schema_name": schema_name,
73
+ "max_results": max_results,
74
+ "page_token": page_token,
75
+ },
76
+ )