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

databricks_mcp/api/warehouses.py

@@ -0,0 +1,32 @@
+# Copyright (c) 2025-2026 Kush Patel
+# Dual-licensed: AGPL-3.0-only OR commercial license.
+# If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
+# For commercial terms, see the project README.
+
+"""SQL Warehouses API 2.0."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from databricks_mcp.client import DatabricksClient
+
+BASE = "/api/2.0/sql/warehouses"
+
+
+async def list_warehouses(client: DatabricksClient) -> dict[str, Any]:
+    return await client.get(BASE)
+
+
+async def get_warehouse(client: DatabricksClient, warehouse_id: str) -> dict[str, Any]:
+    return await client.get(f"{BASE}/{warehouse_id}")
+
+
+async def start_warehouse(client: DatabricksClient, warehouse_id: str) -> dict[str, Any]:
+    await client.post(f"{BASE}/{warehouse_id}/start")
+    return {"warehouse_id": warehouse_id, "start_requested": True}
+
+
+async def stop_warehouse(client: DatabricksClient, warehouse_id: str) -> dict[str, Any]:
+    await client.post(f"{BASE}/{warehouse_id}/stop")
+    return {"warehouse_id": warehouse_id, "stop_requested": True}

databricks_mcp/api/workspace.py

@@ -0,0 +1,67 @@
+# Copyright (c) 2025-2026 Kush Patel
+# Dual-licensed: AGPL-3.0-only OR commercial license.
+# If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
+# For commercial terms, see the project README.
+
+"""Workspace API 2.0 (notebooks, folders, workspace files)."""
+
+from __future__ import annotations
+
+import base64
+from typing import Any
+
+from databricks_mcp.client import DatabricksClient
+
+BASE = "/api/2.0/workspace"
+
+
+async def list_objects(client: DatabricksClient, path: str) -> dict[str, Any]:
+    return await client.get(f"{BASE}/list", {"path": path})
+
+
+async def get_status(client: DatabricksClient, path: str) -> dict[str, Any]:
+    return await client.get(f"{BASE}/get-status", {"path": path})
+
+
+async def export_notebook(client: DatabricksClient, path: str, format: str = "SOURCE") -> dict[str, Any]:
+    """Export and decode. SOURCE/JUPYTER/HTML come back as text; DBC stays base64."""
+    result = await client.get(f"{BASE}/export", {"path": path, "format": format})
+    encoded = result.get("content")
+    if encoded and format != "DBC":
+        try:
+            result["content"] = base64.b64decode(encoded).decode("utf-8")
+            result["encoding"] = "utf-8"
+        except (ValueError, UnicodeDecodeError):
+            result["encoding"] = "base64"
+    elif encoded:
+        result["encoding"] = "base64"
+    return result
+
+
+async def import_notebook(
+    client: DatabricksClient,
+    path: str,
+    content: str,
+    language: str | None = None,
+    format: str = "SOURCE",
+    overwrite: bool = False,
+) -> dict[str, Any]:
+    body: dict[str, Any] = {
+        "path": path,
+        "format": format,
+        "content": base64.b64encode(content.encode("utf-8")).decode("ascii"),
+        "overwrite": overwrite,
+    }
+    if language:
+        body["language"] = language
+    elif format == "SOURCE":
+        raise ValueError("language is required when format is SOURCE.")
+    return await client.post(f"{BASE}/import", body)
+
+
+async def mkdirs(client: DatabricksClient, path: str) -> dict[str, Any]:
+    return await client.post(f"{BASE}/mkdirs", {"path": path})
+
+
+async def delete(client: DatabricksClient, path: str, recursive: bool = False) -> dict[str, Any]:
+    return await client.post(f"{BASE}/delete", {"path": path, "recursive": recursive})

databricks_mcp/cli.py

@@ -0,0 +1,100 @@
+# Copyright (c) 2025-2026 Kush Patel
+# Dual-licensed: AGPL-3.0-only OR commercial license.
+# If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
+# For commercial terms, see the project README.
+
+"""`databricks-mcp` command line entry point."""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import logging
+import os
+import sys
+from dataclasses import replace
+
+from databricks_mcp import __version__
+from databricks_mcp.config import Settings
+from databricks_mcp.server import TOOLSETS, build_server
+
+
+def _parse_toolsets(raw: str | None) -> list[str] | None:
+    if not raw:
+        return None
+    return [t.strip() for t in raw.split(",") if t.strip()]
+
+
+def _parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="databricks-mcp", description="MCP server for Databricks.")
+    parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}")
+    sub = parser.add_subparsers(dest="command")
+
+    run = sub.add_parser("run", help="Run the MCP server (the default command).")
+    run.add_argument(
+        "--transport",
+        choices=["stdio", "streamable-http"],
+        default=os.environ.get("DATABRICKS_MCP_TRANSPORT", "stdio"),
+        help="stdio for desktop clients (default); streamable-http to serve over HTTP.",
+    )
+    run.add_argument("--host", default="127.0.0.1", help="Bind address for streamable-http.")
+    run.add_argument("--port", type=int, default=8000, help="Port for streamable-http.")
+    for p in (run, sub.add_parser("list-tools", help="Print the tools the server would expose.")):
+        p.add_argument("--read-only", action="store_true", help="Only expose tools that don't change anything.")
+        p.add_argument(
+            "--toolsets",
+            default=os.environ.get("DATABRICKS_MCP_TOOLSETS"),
+            help=f"Comma-separated subset of: {','.join(TOOLSETS)}. Default: all.",
+        )
+    run.add_argument(
+        "--log-level",
+        type=str.upper,
+        choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"],
+        default=os.environ.get("DATABRICKS_MCP_LOG_LEVEL", "WARNING").upper(),
+    )
+    return parser
+
+
+def _settings(args: argparse.Namespace) -> Settings:
+    settings = Settings.from_env()
+    if args.read_only:
+        settings = replace(settings, read_only=True)
+    return settings
+
+
+def main(argv: list[str] | None = None) -> int:
+    argv = list(sys.argv[1:] if argv is None else argv)
+    # `run` is the default, so `databricks-mcp --read-only` works in client configs.
+    if not argv or argv[0] not in {"run", "list-tools", "-h", "--help", "--version"}:
+        argv = ["run", *argv]
+    args = _parser().parse_args(argv)
+
+    try:
+        server = build_server(
+            settings=_settings(args),
+            toolsets=_parse_toolsets(args.toolsets),
+            log_level=getattr(args, "log_level", "WARNING"),
+        )
+    except ValueError as exc:
+        print(f"databricks-mcp: {exc}", file=sys.stderr)
+        return 2
+
+    if args.command == "list-tools":
+        for tool in asyncio.run(server.list_tools()):
+            ro = " [read-only]" if tool.annotations and tool.annotations.read_only_hint else ""
+            print(f"{tool.name}{ro}\n    {(tool.description or '').splitlines()[0]}")
+        return 0
+
+    # stdout is the MCP channel under stdio, so logs must go to stderr.
+    logging.basicConfig(stream=sys.stderr, level=args.log_level, format="%(levelname)s %(name)s: %(message)s")
+    # httpx logs every request URL at INFO; keep that out unless asked for.
+    logging.getLogger("httpx").setLevel(max(logging.WARNING, logging.getLevelName(args.log_level)))
+    if args.transport == "stdio":
+        server.run("stdio")
+    else:
+        server.run("streamable-http", host=args.host, port=args.port)
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

databricks_mcp/client.py

@@ -0,0 +1,208 @@
+# Copyright (c) 2025-2026 Kush Patel
+# Dual-licensed: AGPL-3.0-only OR commercial license.
+# If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
+# For commercial terms, see the project README.
+
+"""Async HTTP client for the Databricks REST API."""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from collections.abc import Callable, Mapping
+from typing import Any
+
+import httpx
+
+from databricks_mcp import __version__
+
+logger = logging.getLogger(__name__)
+
+# Returns (workspace host, auth headers). Called per request so OAuth tokens refresh.
+AuthProvider = Callable[[], tuple[str, dict[str, str]]]
+
+# Safe to resend after a 5xx or timeout. POSTs are not: a retried clusters/create
+# can leave you with two clusters.
+_IDEMPOTENT = {"GET", "HEAD", "PUT", "DELETE"}
+_MAX_BACKOFF_SECONDS = 30.0
+_USER_AGENT = f"databricks-mcp-toolkit/{__version__}"
+
+
+class DatabricksAPIError(Exception):
+    def __init__(
+        self,
+        message: str,
+        status_code: int | None = None,
+        error_code: str | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.message = message
+        self.status_code = status_code
+        self.error_code = error_code
+
+    def __str__(self) -> str:
+        parts = []
+        if self.status_code:
+            parts.append(f"HTTP {self.status_code}")
+        if self.error_code:
+            parts.append(self.error_code)
+        prefix = " ".join(parts)
+        return f"{prefix}: {self.message}" if prefix else self.message
+
+
+def sdk_auth_provider() -> AuthProvider:
+    """Auth via databricks-sdk, so every auth type the SDK supports works here too.
+
+    I build the Config lazily: `databricks-mcp list-tools` shouldn't need credentials.
+    """
+    config = None
+
+    def provide() -> tuple[str, dict[str, str]]:
+        nonlocal config
+        if config is None:
+            from databricks.sdk.config import Config
+
+            try:
+                config = Config()
+            except Exception as exc:  # the SDK raises ValueError with a decent message
+                raise DatabricksAPIError(
+                    "Databricks auth is not configured. Set DATABRICKS_HOST plus "
+                    "DATABRICKS_TOKEN (or OAuth / a config profile). "
+                    f"SDK said: {exc}"
+                ) from exc
+        return config.host.rstrip("/"), config.authenticate()
+
+    return provide
+
+
+def _clean_params(params: Mapping[str, Any] | None) -> dict[str, Any] | None:
+    if not params:
+        return None
+    return {k: v for k, v in params.items() if v is not None}
+
+
+def _error_from_response(response: httpx.Response) -> DatabricksAPIError:
+    message = ""
+    error_code = None
+    try:
+        body = response.json()
+        if isinstance(body, dict):
+            error_code = body.get("error_code")
+            message = body.get("message") or body.get("error") or ""
+    except ValueError:
+        pass
+    if not message:
+        message = response.text[:500] or response.reason_phrase
+    return DatabricksAPIError(message, status_code=response.status_code, error_code=error_code)
+
+
+class DatabricksClient:
+    def __init__(
+        self,
+        auth: AuthProvider | None = None,
+        *,
+        timeout_seconds: float = 60.0,
+        max_retries: int = 3,
+        backoff_seconds: float = 1.0,
+        transport: httpx.AsyncBaseTransport | None = None,
+    ) -> None:
+        self._auth = auth or sdk_auth_provider()
+        self._timeout = httpx.Timeout(timeout_seconds, connect=10.0)
+        self._max_retries = max_retries
+        self._backoff = backoff_seconds
+        self._transport = transport
+        self._http: httpx.AsyncClient | None = None
+
+    def _client(self) -> httpx.AsyncClient:
+        if self._http is None:
+            self._http = httpx.AsyncClient(
+                timeout=self._timeout,
+                transport=self._transport,
+                headers={"User-Agent": _USER_AGENT},
+            )
+        return self._http
+
+    async def aclose(self) -> None:
+        if self._http is not None:
+            await self._http.aclose()
+            self._http = None
+
+    def _delay(self, attempt: int, response: httpx.Response | None = None) -> float:
+        if response is not None:
+            retry_after = response.headers.get("Retry-After")
+            if retry_after:
+                try:
+                    return min(float(retry_after), _MAX_BACKOFF_SECONDS)
+                except ValueError:
+                    pass
+        return min(self._backoff * (2**attempt), _MAX_BACKOFF_SECONDS)
+
+    async def request(
+        self,
+        method: str,
+        path: str,
+        *,
+        params: Mapping[str, Any] | None = None,
+        json: Any = None,
+        content: bytes | None = None,
+        headers: Mapping[str, str] | None = None,
+        raw: bool = False,
+    ) -> Any:
+        """Send a request. Returns parsed JSON ({} for an empty body), or the
+        httpx.Response itself when raw=True (used by the Files API)."""
+        method = method.upper()
+        host, auth_headers = self._auth()
+        url = f"{host}{path if path.startswith('/') else '/' + path}"
+        merged_headers = {**auth_headers, **(headers or {})}
+        clean = _clean_params(params)
+
+        attempt = 0
+        while True:
+            try:
+                response = await self._client().request(
+                    method,
+                    url,
+                    params=clean,
+                    json=json,
+                    content=content,
+                    headers=merged_headers,
+                )
+            except httpx.ConnectError as exc:
+                # Never reached the server, so any method is safe to resend.
+                if attempt < self._max_retries:
+                    await asyncio.sleep(self._delay(attempt))
+                    attempt += 1
+                    continue
+                raise DatabricksAPIError(f"Could not connect to {host}: {exc}") from exc
+            except (httpx.TimeoutException, httpx.TransportError) as exc:
+                if method in _IDEMPOTENT and attempt < self._max_retries:
+                    await asyncio.sleep(self._delay(attempt))
+                    attempt += 1
+                    continue
+                raise DatabricksAPIError(f"{method} {path} failed: {exc!r}") from exc
+
+            status = response.status_code
+            retryable = status == 429 or (status >= 500 and method in _IDEMPOTENT)
+            if retryable and attempt < self._max_retries:
+                delay = self._delay(attempt, response)
+                logger.warning("%s %s -> %s, retrying in %.1fs", method, path, status, delay)
+                await asyncio.sleep(delay)
+                attempt += 1
+                continue
+
+            if status >= 400:
+                raise _error_from_response(response)
+            if raw:
+                return response
+            if not response.content:
+                return {}
+            try:
+                return response.json()
+            except ValueError:
+                return {"raw": response.text}
+
+    async def get(self, path: str, params: Mapping[str, Any] | None = None) -> Any:
+        return await self.request("GET", path, params=params)
+
+    async def post(self, path: str, json: Any = None) -> Any:
+        return await self.request("POST", path, json=json if json is not None else {})

databricks_mcp/config.py

@@ -0,0 +1,56 @@
+# Copyright (c) 2025-2026 Kush Patel
+# Dual-licensed: AGPL-3.0-only OR commercial license.
+# If you obtained this from GitHub or PyPI, you may use it under the AGPL-3.0 terms.
+# For commercial terms, see the project README.
+
+"""Server settings.
+
+Databricks auth is NOT configured here. I hand that to databricks-sdk's unified auth,
+so anything the Databricks CLI accepts works: DATABRICKS_HOST + DATABRICKS_TOKEN,
+OAuth M2M (DATABRICKS_CLIENT_ID / DATABRICKS_CLIENT_SECRET), a ~/.databrickscfg
+profile via DATABRICKS_CONFIG_PROFILE, Azure CLI, and so on.
+"""
+
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+_TRUE = {"1", "true", "yes", "on"}
+
+
+def _env_bool(name: str, default: bool = False) -> bool:
+    raw = os.environ.get(name)
+    if raw is None or raw.strip() == "":
+        return default
+    return raw.strip().lower() in _TRUE
+
+
+def _env_int(name: str, default: int) -> int:
+    raw = os.environ.get(name)
+    if raw is None or raw.strip() == "":
+        return default
+    try:
+        return int(raw)
+    except ValueError as exc:
+        raise ValueError(f"{name} must be an integer, got {raw!r}") from exc
+
+
+@dataclass(frozen=True)
+class Settings:
+    warehouse_id: str | None = None
+    read_only: bool = False
+    # Hard cap on what a single tool returns to the model. Big JSON blobs eat context fast.
+    max_response_chars: int = 50_000
+    http_timeout_seconds: int = 60
+    max_retries: int = 3
+
+    @classmethod
+    def from_env(cls) -> Settings:
+        return cls(
+            warehouse_id=os.environ.get("DATABRICKS_WAREHOUSE_ID") or None,
+            read_only=_env_bool("DATABRICKS_MCP_READ_ONLY"),
+            max_response_chars=_env_int("DATABRICKS_MCP_MAX_RESPONSE_CHARS", 50_000),
+            http_timeout_seconds=_env_int("DATABRICKS_MCP_HTTP_TIMEOUT", 60),
+            max_retries=_env_int("DATABRICKS_MCP_MAX_RETRIES", 3),
+        )