databricks-explore 0.5.0__py3-none-any.whl

databricks_explore/mcp_server/__init__.py

@@ -0,0 +1,93 @@
+"""An MCP server inside the Flask app: streamable HTTP at /mcp, local only, bearer token, JSON-RPC handled directly.
+
+`register(app, explorer, store, miner)` adds /mcp and the two /api/agent-access routes the page uses. Read access,
+local edits and read-only SQL only: no tool writes to the workspace or connects a profile. See contracts/mcp.md.
+"""
+import json
+from typing import Callable
+
+from flask import Flask, Response, jsonify, request
+
+from . import auth, rpc
+from .tooling import REGISTRY, Context
+
+SESSION_HEADER = "Mcp-Session-Id"
+
+
+def access_info(port: str) -> dict:
+    """The URL, the token and the ready-to-paste `claude mcp add` command for the Agent access panel."""
+    url = f"http://127.0.0.1:{port}/mcp"
+    token = auth.get_token()
+    return {"url": url, "token": token,
+            "command": f'claude mcp add --transport http databricks-explore {url} --header "Authorization: Bearer {token}"',
+            "tools": len(REGISTRY)}
+
+
+def register(app: Flask, explorer: Callable, store: Callable, miner: Callable) -> None:
+    """Add the MCP endpoint and the agent-access routes to the app."""
+    ctx = Context(explorer=explorer, store=store, miner=miner)
+    auth.get_token()  # create the token file at startup, so it exists before the first request
+
+    def refuse_non_local():
+        """403 for a Host or Origin that is not this machine (DNS rebinding); None when fine."""
+        if not auth.host_is_local(request.host) or not auth.origin_is_local(request.headers.get("Origin")):
+            return jsonify({"error": "forbidden: this endpoint only answers local requests"}), 403
+        return None
+
+    def json_rpc(body: dict, status: int = 200, headers: dict | None = None) -> Response:
+        return Response(json.dumps(body), status=status, mimetype="application/json", headers=headers or {})
+
+    @app.route("/mcp", methods=["GET", "POST", "DELETE"])
+    def mcp_endpoint():
+        refused = refuse_non_local()
+        if refused:
+            return refused
+        if not auth.bearer_ok(request.headers.get("Authorization")):
+            return Response(json.dumps({"error": "unauthorized: send Authorization: Bearer <token> (see Agent access in the app)"}),
+                            status=401, mimetype="application/json", headers={"WWW-Authenticate": 'Bearer realm="databricks-explore"'})
+        if request.method == "GET":  # no server-to-client stream is offered
+            return Response("", status=405, headers={"Allow": "POST, DELETE"})
+        session = request.headers.get(SESSION_HEADER)
+        if request.method == "DELETE":
+            return Response("", status=204 if auth.session_ok(session) else 404)
+        version = request.headers.get("MCP-Protocol-Version")
+        if version is not None and version not in rpc.PROTOCOL_VERSIONS:
+            return json_rpc(rpc.error(rpc.INVALID_REQUEST, f"unsupported MCP-Protocol-Version {version}; use one of {', '.join(rpc.PROTOCOL_VERSIONS)}"), 400)
+        try:
+            body = json.loads(request.get_data(as_text=True) or "")
+        except ValueError:
+            return json_rpc(rpc.error(rpc.PARSE_ERROR, "the body is not valid JSON"), 400)
+        batch = body if isinstance(body, list) else [body]
+        if not batch:
+            return json_rpc(rpc.error(rpc.INVALID_REQUEST, "empty batch"), 400)
+        if all(isinstance(m, dict) and m.get("method") == "server/discover" for m in batch):
+            # The 2026-07-28 handshake-free probe. We speak the initialize handshake, so say "no such method" and let the client fall back.
+            return json_rpc(rpc.error(rpc.METHOD_NOT_FOUND, "method not found: server/discover; use initialize", batch[0].get("id")))
+        is_init = any(isinstance(m, dict) and m.get("method") == "initialize" for m in batch)
+        if not is_init:
+            if not session:
+                return json_rpc(rpc.error(rpc.INVALID_REQUEST, f"missing {SESSION_HEADER} header: send initialize first"), 400)
+            if not auth.session_ok(session):
+                return json_rpc(rpc.error(rpc.INVALID_REQUEST, "unknown session: send initialize again"), 404)
+        replies = [r for r in (rpc.handle(m, ctx) for m in batch) if r is not None]
+        headers = {SESSION_HEADER: auth.new_session_id()} if is_init else {}
+        if not replies:
+            return Response("", status=202, headers=headers)
+        if isinstance(body, list):
+            return Response(json.dumps(replies), mimetype="application/json", headers=headers)
+        status = 400 if replies[0].get("error", {}).get("code") in (rpc.INVALID_REQUEST, rpc.PARSE_ERROR) else 200
+        return json_rpc(replies[0], status, headers)
+
+    @app.get("/api/agent-access")
+    def agent_access():
+        """The MCP URL, token and ready-to-paste command for the Agent access panel. Local requests only."""
+        return refuse_non_local() or jsonify(access_info(request.environ.get("SERVER_PORT", "8765")))
+
+    @app.post("/api/agent-access/token")
+    def agent_access_new_token():
+        """Replace the token. Agents using the old one get a 401 until they are given the new one."""
+        refused = refuse_non_local()
+        if refused:
+            return refused
+        auth.new_token()
+        return jsonify(access_info(request.environ.get("SERVER_PORT", "8765")))

databricks_explore/mcp_server/auth.py

@@ -0,0 +1,95 @@
+"""The MCP bearer token (a 0600 file under the data directory), the local-only Host and Origin checks, and session ids."""
+import hashlib
+import hmac
+import os
+import secrets
+import uuid
+from typing import Optional
+from urllib.parse import urlsplit
+
+from .. import dbx
+
+LOCAL_HOSTS = ("127.0.0.1", "localhost", "::1")
+
+
+def token_path():
+    """data/mcp_token, under whatever dbx.DATA is when called."""
+    return dbx.DATA / "mcp_token"
+
+
+def _write(path, token: str) -> None:
+    """Write the token with mode 0600 from the start (never briefly readable by others), replacing any older file."""
+    tmp = path.parent / (path.name + ".tmp")
+    tmp.unlink(missing_ok=True)
+    fd = os.open(tmp, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
+    with os.fdopen(fd, "w", encoding="utf-8") as f:
+        f.write(token + "\n")
+    os.chmod(tmp, 0o600)
+    os.replace(tmp, path)
+
+
+def get_token() -> str:
+    """The token. Created on first use. A file with looser permissions is tightened to 0600."""
+    path = token_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        token = path.read_text(encoding="utf-8").strip()
+    except OSError:
+        token = ""
+    if not token:
+        token = secrets.token_urlsafe(32)
+        _write(path, token)
+    elif path.stat().st_mode & 0o077:
+        os.chmod(path, 0o600)
+    return token
+
+
+def new_token() -> str:
+    """Replace the token. Every agent configured with the old one gets a 401 until it is given the new one."""
+    path = token_path()
+    path.parent.mkdir(parents=True, exist_ok=True)
+    token = secrets.token_urlsafe(32)
+    _write(path, token)
+    return token
+
+
+def bearer_ok(header: Optional[str]) -> bool:
+    """True when the Authorization header is `Bearer <the token>`. Compared in constant time."""
+    if not header or not header.lower().startswith("bearer "):
+        return False
+    return hmac.compare_digest(header[7:].strip().encode(), get_token().encode())
+
+
+def _hostname(netloc_or_url: str) -> str:
+    parts = urlsplit(netloc_or_url if "//" in netloc_or_url else "//" + netloc_or_url)
+    return (parts.hostname or "").lower()
+
+
+def host_is_local(host_header: Optional[str]) -> bool:
+    """The Host header names this machine. A DNS-rebinding page reaches us under its own name, which fails this."""
+    return _hostname(host_header or "") in LOCAL_HOSTS
+
+
+def origin_is_local(origin: Optional[str]) -> bool:
+    """No Origin (a non-browser client) is fine. A browser's Origin must be a local address; "null" is not."""
+    if origin is None:
+        return True
+    return _hostname(origin) in LOCAL_HOSTS
+
+
+def new_session_id() -> str:
+    """A session id the server can recognize later without keeping a table: a random part and its HMAC under the token."""
+    sid = uuid.uuid4().hex
+    return f"{sid}.{_mac(sid)}"
+
+
+def session_ok(session_id: Optional[str]) -> bool:
+    """True for an id new_session_id() made under the current token."""
+    if not session_id or "." not in session_id:
+        return False
+    sid, mac = session_id.rsplit(".", 1)
+    return hmac.compare_digest(mac, _mac(sid))
+
+
+def _mac(sid: str) -> str:
+    return hmac.new(get_token().encode(), sid.encode(), hashlib.sha256).hexdigest()[:24]

databricks_explore/mcp_server/rpc.py

@@ -0,0 +1,88 @@
+"""JSON-RPC 2.0 for MCP: initialize, ping, tools/list and tools/call, with protocol version negotiation and proper errors."""
+import time
+from typing import Any, Optional
+
+from .. import dbx
+
+from . import tools_canvas, tools_data  # noqa: F401  (importing registers the tools)
+from .schema import validate
+from .tooling import REGISTRY, Context, ToolError
+
+PROTOCOL_VERSIONS = ("2025-11-25", "2025-06-18", "2025-03-26", "2024-11-05")  # newest first
+SERVER_INFO = {"name": "databricks-explore", "title": "Databricks Explorer", "version": "0.5.0"}
+INSTRUCTIONS = (
+    "Local Unity Catalog explorer. Read the catalog, relationships and concepts, run read-only SQL, and edit canvases next to the person using the browser. "
+    "Start with status. Tools that need a connection say so: the person connects in the app, you cannot. "
+    "No tool writes to the workspace. Canvas edits are saved at once and undoable with undo_canvas. "
+    "On canvases prefer structural placement (place: right_of/below/inside an anchor id) over absolute coordinates, and read the doctor score each edit returns.")
+
+PARSE_ERROR, INVALID_REQUEST, METHOD_NOT_FOUND, INVALID_PARAMS, INTERNAL_ERROR = -32700, -32600, -32601, -32602, -32603
+
+
+def error(code: int, message: str, rid: Any = None) -> dict:
+    """A JSON-RPC error response."""
+    return {"jsonrpc": "2.0", "id": rid, "error": {"code": code, "message": message}}
+
+
+def negotiate(requested: Optional[str]) -> str:
+    """The protocol version to use: the client's when we speak it, else our newest."""
+    return requested if requested in PROTOCOL_VERSIONS else PROTOCOL_VERSIONS[0]
+
+
+def handle(message: Any, ctx: Context) -> Optional[dict]:
+    """Answer one JSON-RPC message. None for a notification or a response from the client (nothing to send back)."""
+    if not isinstance(message, dict) or message.get("jsonrpc") != "2.0":
+        return error(INVALID_REQUEST, "not a JSON-RPC 2.0 message")
+    method, rid = message.get("method"), message.get("id")
+    if method is None:
+        return None if ("result" in message or "error" in message) else error(INVALID_REQUEST, "missing method", rid)
+    if not isinstance(method, str):
+        return error(INVALID_REQUEST, "method must be a string", rid)
+    if "id" not in message:
+        return None  # notifications/initialized, notifications/cancelled, ...: accepted, no reply
+    if isinstance(rid, bool) or not isinstance(rid, (int, str)):
+        return error(INVALID_REQUEST, "id must be a string or an integer")
+    params = message.get("params", {})
+    if not isinstance(params, dict):
+        return error(INVALID_PARAMS, "params must be an object", rid)
+    if method == "initialize":
+        return {"jsonrpc": "2.0", "id": rid, "result": {
+            "protocolVersion": negotiate(params.get("protocolVersion")), "capabilities": {"tools": {"listChanged": False}},
+            "serverInfo": SERVER_INFO, "instructions": INSTRUCTIONS}}
+    if method == "ping":
+        return {"jsonrpc": "2.0", "id": rid, "result": {}}
+    if method == "tools/list":
+        return {"jsonrpc": "2.0", "id": rid, "result": {"tools": [t.listing() for t in REGISTRY.values()]}}
+    if method == "tools/call":
+        return _call(params, rid, ctx)
+    return error(METHOD_NOT_FOUND, f"method not found: {method}", rid)
+
+
+def _call(params: dict, rid: Any, ctx: Context) -> dict:
+    name = params.get("name")
+    tool = REGISTRY.get(name) if isinstance(name, str) else None
+    if tool is None:
+        return error(INVALID_PARAMS, f"unknown tool: {name!r}", rid)
+    args = params.get("arguments", {})
+    problems = validate(args, tool.schema) if isinstance(args, dict) else ["arguments: must be an object"]
+    if problems:  # a tool execution error, so the model can read it and correct the call
+        return _result(rid, [{"type": "text", "text": "Invalid arguments: " + "; ".join(problems[:6])}], True)
+    started = time.time()
+    try:
+        content, failed = tool.handler(ctx, args), False
+    except Exception as e:  # ToolError and every domain error (validation, not found, warehouse): message back to the agent
+        content, failed = [{"type": "text", "text": _message(e)}], True
+    # The name and timing only: never arguments, SQL or results.
+    dbx.log({"op": f"mcp.{tool.name}", "ms": int((time.time() - started) * 1000), "ok": not failed})
+    return _result(rid, content, failed)
+
+
+def _message(e: Exception) -> str:
+    if isinstance(e, ToolError):
+        return str(e)
+    msg = dbx.clean_error(e)
+    return f"{type(e).__name__}: {msg}" if msg and not isinstance(e, (ValueError, LookupError)) else (msg or type(e).__name__)
+
+
+def _result(rid: Any, content: list, is_error: bool) -> dict:
+    return {"jsonrpc": "2.0", "id": rid, "result": {"content": content, "isError": is_error}}

databricks_explore/mcp_server/schema.py

@@ -0,0 +1,68 @@
+"""A small JSON Schema checker for tool arguments: type, enum, const, bounds, items, properties, required, no extra keys."""
+from typing import Any
+
+_TYPES = {
+    "string": lambda v: isinstance(v, str),
+    "integer": lambda v: isinstance(v, int) and not isinstance(v, bool),
+    "number": lambda v: isinstance(v, (int, float)) and not isinstance(v, bool),
+    "boolean": lambda v: isinstance(v, bool),
+    "array": lambda v: isinstance(v, list),
+    "object": lambda v: isinstance(v, dict),
+    "null": lambda v: v is None,
+}
+
+
+def validate(value: Any, schema: dict, path: str = "arguments") -> list[str]:
+    """Every way `value` breaks `schema`, as messages that name the path. Empty when it is fine."""
+    errors: list[str] = []
+    if "anyOf" in schema and not any(not validate(value, alt, path) for alt in schema["anyOf"]):
+        return [f"{path}: does not match any allowed form ({'; '.join(_describe(a) for a in schema['anyOf'])})"]
+    want = schema.get("type")
+    if want is not None:
+        kinds = want if isinstance(want, list) else [want]
+        if not any(_TYPES[k](value) for k in kinds):
+            return [f"{path}: expected {' or '.join(kinds)}, got {_name(value)}"]
+    if "const" in schema and value != schema["const"]:
+        errors.append(f"{path}: must be {schema['const']!r}")
+    if "enum" in schema and value not in schema["enum"]:
+        errors.append(f"{path}: must be one of {', '.join(map(str, schema['enum']))}")
+    if isinstance(value, (int, float)) and not isinstance(value, bool):
+        if "minimum" in schema and value < schema["minimum"]:
+            errors.append(f"{path}: must be at least {schema['minimum']}")
+        if "exclusiveMinimum" in schema and value <= schema["exclusiveMinimum"]:
+            errors.append(f"{path}: must be greater than {schema['exclusiveMinimum']}")
+        if "maximum" in schema and value > schema["maximum"]:
+            errors.append(f"{path}: must be at most {schema['maximum']}")
+    if isinstance(value, str):
+        if "minLength" in schema and len(value) < schema["minLength"]:
+            errors.append(f"{path}: must not be empty" if schema["minLength"] == 1 else f"{path}: too short")
+        if "maxLength" in schema and len(value) > schema["maxLength"]:
+            errors.append(f"{path}: longer than {schema['maxLength']} characters")
+    if isinstance(value, list):
+        if "minItems" in schema and len(value) < schema["minItems"]:
+            errors.append(f"{path}: needs at least {schema['minItems']} item(s)")
+        if "maxItems" in schema and len(value) > schema["maxItems"]:
+            errors.append(f"{path}: at most {schema['maxItems']} items")
+        for i, item in enumerate(value):
+            if "items" in schema:
+                errors += validate(item, schema["items"], f"{path}[{i}]")
+    if isinstance(value, dict):
+        props = schema.get("properties", {})
+        for key in schema.get("required", []):
+            if key not in value:
+                errors.append(f"{path}: missing required field {key!r}")
+        for key, item in value.items():
+            if key in props:
+                errors += validate(item, props[key], f"{path}.{key}")
+            elif schema.get("additionalProperties") is False:
+                errors.append(f"{path}: unknown field {key!r} (known: {', '.join(props) or 'none'})")
+    return errors
+
+
+def _name(v: Any) -> str:
+    return "null" if v is None else type(v).__name__
+
+
+def _describe(schema: dict) -> str:
+    kind = schema.get("type", "value")
+    return " or ".join(kind) if isinstance(kind, list) else kind

databricks_explore/mcp_server/sqlguard.py

@@ -0,0 +1,59 @@
+"""The read-only gate for run_sql: exactly one SELECT, DESCRIBE, SHOW or EXPLAIN statement, checked with sqlglot."""
+import logging
+
+import sqlglot
+from sqlglot import exp
+from sqlglot.errors import SqlglotError
+
+logging.getLogger("sqlglot").setLevel(logging.ERROR)  # the Command fallback logs a warning for SHOW and EXPLAIN
+
+QUERY_ROOTS = (exp.Select, exp.SetOperation)  # WITH ... SELECT is a Select that carries a `with`
+COMMAND_WORDS = ("SHOW", "EXPLAIN", "DESCRIBE", "DESC")  # sqlglot reads these as a Command when it cannot model them
+# Anything of these kinds anywhere in the tree is refused, so DML or DDL cannot hide in a CTE or a subquery.
+_FORBIDDEN_NAMES = ("Insert", "Update", "Delete", "Merge", "Create", "Drop", "Alter", "AlterTable", "TruncateTable", "Copy",
+                    "Command", "Use", "Set", "SetItem", "Into", "Grant", "Revoke", "Analyze", "Attach", "Detach", "Refresh",
+                    "Cache", "Uncache", "LoadData", "Kill", "Transaction", "Commit", "Rollback", "Pragma")
+FORBIDDEN = tuple(getattr(exp, n) for n in _FORBIDDEN_NAMES if hasattr(exp, n))
+# Functions that run code or call out of the warehouse, so a SELECT is not enough to be read-only.
+FORBIDDEN_FUNCTIONS = {"reflect", "java_method", "try_reflect", "http_request"}
+
+
+class NotReadOnly(ValueError):
+    """The statement is not one run_sql accepts. str(e) says why."""
+
+
+def check_read_only(sql: str) -> str:
+    """Return the one statement to run (as the caller wrote it, without a trailing `;`), or raise NotReadOnly."""
+    if not isinstance(sql, str) or not sql.strip():
+        raise NotReadOnly("statement is empty")
+    try:
+        parsed = sqlglot.parse(sql, read="databricks")
+    except SqlglotError as e:
+        raise NotReadOnly(f"could not parse the statement ({_first_line(e)}); only a plain SELECT, DESCRIBE, SHOW or EXPLAIN is accepted") from e
+    statements = [s for s in parsed if s is not None]
+    if not statements:
+        raise NotReadOnly("statement is empty")
+    if len(statements) > 1:
+        raise NotReadOnly(f"{len(statements)} statements found; send exactly one")
+    root = statements[0]
+    while isinstance(root, (exp.Subquery, exp.Paren)) and root.this is not None:
+        root = root.this  # (SELECT 1) is still a query
+    if isinstance(root, exp.Command):
+        word = str(root.this).upper()
+        if word not in COMMAND_WORDS:
+            raise NotReadOnly(f"{word} is not allowed; run_sql accepts SELECT, DESCRIBE, SHOW and EXPLAIN")
+        return sql.strip().rstrip(";").rstrip()
+    if not isinstance(root, QUERY_ROOTS + (exp.Describe,)):
+        raise NotReadOnly(f"{type(root).__name__.upper()} statements are not allowed; run_sql accepts SELECT, DESCRIBE, SHOW and EXPLAIN")
+    for node in root.walk():
+        if isinstance(node, FORBIDDEN):
+            raise NotReadOnly(f"{type(node).__name__.upper()} is not allowed inside a query; run_sql is read-only")
+        if isinstance(node, (exp.Anonymous, exp.Func)):
+            name = node.name.lower() if isinstance(node, exp.Anonymous) else type(node).__name__.lower()
+            if name in FORBIDDEN_FUNCTIONS:
+                raise NotReadOnly(f"{name}() is not allowed; it can run code or call outside the warehouse")
+    return sql.strip().rstrip(";").rstrip()
+
+
+def _first_line(e: Exception) -> str:
+    return str(e).strip().splitlines()[0][:120] if str(e).strip() else type(e).__name__

databricks_explore/mcp_server/tooling.py

@@ -0,0 +1,74 @@
+"""What a tool is: the registry, the decorator that adds one, the context handlers get, and the error they raise."""
+from dataclasses import dataclass, field
+from typing import Any, Callable, Optional
+
+REGISTRY: dict[str, "Tool"] = {}
+
+
+class ToolError(Exception):
+    """A tool could not do what was asked. The message goes back to the agent as an error result."""
+
+
+@dataclass
+class Context:
+    """The app's connection state, as the tools see it. The tools never connect, and never choose a profile."""
+    explorer: Callable[[], Any]                      # the connected Explorer, or None
+    store: Callable[[], Any]                         # the relationship store of the connected profile (a context manager)
+    miner: Callable[[], Any]                         # the Miner of the connected Explorer
+    extra: dict = field(default_factory=dict)
+
+    def need_explorer(self):
+        """The connected Explorer, or a ToolError that tells the agent to ask the person to connect."""
+        e = self.explorer()
+        if e is None:
+            raise ToolError("Not connected: connect a profile in the app first. The agent cannot connect or choose a profile.")
+        return e
+
+
+@dataclass
+class Tool:
+    name: str
+    description: str
+    schema: dict
+    handler: Callable[[Context, dict], Any]
+    read_only: bool = True
+    destructive: bool = False
+
+    def listing(self) -> dict:
+        """The tools/list entry."""
+        return {"name": self.name, "description": self.description, "inputSchema": self.schema,
+                "annotations": {"readOnlyHint": self.read_only, "destructiveHint": self.destructive,
+                                "idempotentHint": False, "openWorldHint": False}}
+
+
+def obj(properties: Optional[dict] = None, required: tuple = ()) -> dict:
+    """An input schema: an object with exactly these properties."""
+    return {"type": "object", "properties": properties or {}, "required": list(required), "additionalProperties": False}
+
+
+S = {"type": "string"}
+B = {"type": "boolean"}
+FORMAT = {"type": "string", "enum": ["text", "json"], "description": "text (default, compact) or json (everything)"}
+
+
+def integer(minimum: int, maximum: int, description: str = "") -> dict:
+    """A bounded integer property."""
+    return {"type": "integer", "minimum": minimum, "maximum": maximum, **({"description": description} if description else {})}
+
+
+def strings(description: str = "") -> dict:
+    """A list of strings property."""
+    return {"type": "array", "items": S, **({"description": description} if description else {})}
+
+
+def tool(name: str, description: str, schema: dict, read_only: bool = True, destructive: bool = False):
+    """Register a handler `fn(ctx, args) -> str | list of MCP content blocks` as a tool."""
+    def add(fn):
+        REGISTRY[name] = Tool(name, description, schema, fn, read_only, destructive)
+        return fn
+    return add
+
+
+def text(s: str) -> list[dict]:
+    """One MCP text block."""
+    return [{"type": "text", "text": s}]