databricks-explore 0.5.0__py3-none-any.whl

databricks_explore/demo_history.py

@@ -0,0 +1,157 @@
+"""Demo mode's invented query history, table lineage and column lineage, as the warehouse's system tables return them.
+
+Everything is invented and deterministic (the statements, users and their order never change; the timestamps are
+anchored to the start of the current UTC day). `system_query` answers the SQL that mining.py and
+Explorer.column_lineage send, filtering on the same parameters the real tables would.
+
+Set EXPLORER_DEMO_REDACTED=1 (databricks-explore --demo --redacted) to make every statement_text read <REDACTED>, as it does
+for a reader who is neither an account admin nor in databricks_pii_access. Without it, three rows are redacted, so
+a run reports them as skipped.
+"""
+import os
+import random
+import re
+from datetime import datetime, timedelta, timezone
+
+USERS = ["ana.k@example.invalid", "ben.o@example.invalid", "chloe.m@example.invalid", "dev.p@example.invalid",
+         "elif.s@example.invalid", "farid.n@example.invalid"]
+REDACTED = "<REDACTED>"
+LINEAGE_SEEN = "2026-11-10T08:00:00.000Z"
+
+O, C, I, P = "demo.sales.orders", "demo.sales.customers", "demo.sales.order_items", "demo.sales.products"
+R, S, E, D = "demo.sales.returns", "demo.sales.shipments", "demo.hr.employees", "demo.hr.departments"
+T, V = "samples.nyctaxi.trips", "demo.sales.v_daily_revenue"
+SYS = ["system.access.table_lineage", "system.query.history"]
+
+# source table, source column, target table, target column
+COLUMN_LINEAGE = [
+    (O, "order_date", V, "order_date"), (O, "amount", V, "revenue"), (O, "order_id", V, "orders"),
+    (T, "trip_id", "demo.metrics.wide_metrics", "id"), (T, "pickup_ts", "demo.metrics.wide_metrics", "event_ts"),
+    (T, "distance_mi", "demo.metrics.wide_metrics", "metric_01_throughput_per_sec"),
+    (T, "fare", "demo.metrics.wide_metrics", "metric_02_error_rate"),
+]
+
+# (count, user indexes, statement type, status, lineage tables or None for no lineage row, SQL with {date}/{email})
+_SPECS = [
+    (70, [0, 1, 2, 3, 4, 5], "SELECT", "FINISHED", [O, C],
+     "SELECT o.order_id, c.name, o.amount FROM demo.sales.orders o JOIN demo.sales.customers c "
+     "ON o.customer_id = c.customer_id WHERE c.email = '{email}' AND o.order_date >= '{date}'"),
+    (60, [0, 1, 3], "SELECT", "FINISHED", [I, P],
+     "SELECT i.order_id, p.category, SUM(i.qty * i.unit_price) AS revenue FROM demo.sales.order_items i "
+     "INNER JOIN demo.sales.products p ON i.sku = p.sku WHERE p.list_price > 10 GROUP BY 1, 2"),
+    (35, [1, 2, 4], "SELECT", "FINISHED", [O, R],
+     "SELECT o.order_id, r.reason FROM demo.sales.orders o, demo.sales.returns r "
+     "WHERE o.order_id = r.order_id AND r.returned_at > '{date}'"),
+    (30, [0, 2, 5], "SELECT", "FINISHED", [O, C, S],
+     "WITH recent AS (SELECT order_id AS oid, customer_id AS cid, amount FROM demo.sales.orders "
+     "WHERE order_date >= '{date}'), shipped AS (SELECT * FROM demo.sales.shipments) "
+     "SELECT c.name, recent.amount, s.carrier FROM recent JOIN demo.sales.customers c ON c.customer_id = recent.cid "
+     "LEFT JOIN shipped s ON s.order_id = recent.oid"),
+    (25, [1, 3], "SELECT", "FINISHED", [I, R],
+     "SELECT * FROM demo.sales.order_items i LEFT JOIN demo.sales.returns r USING (order_id, line_no)"),
+    (20, [0, 4], "SELECT", "FINISHED", [I, S],
+     "SELECT i.order_id, i.units, s.carrier FROM (SELECT order_id, SUM(qty) AS units FROM demo.sales.order_items "
+     "GROUP BY order_id) i JOIN demo.sales.shipments s ON s.order_id = i.order_id"),
+    (22, [5, 2], "SELECT", "FINISHED", [E, D],
+     "SELECT e.name, d.budget FROM demo.hr.employees e JOIN demo.hr.departments d ON e.dept = d.dept"),
+    (8, [5], "SELECT", "FINISHED", [E, D],
+     "SELECT d.dept, e.name AS manager FROM demo.hr.departments d JOIN demo.hr.employees e ON e.emp_id = d.manager_id"),
+    (25, [1, 2, 3, 4], "SELECT", "FINISHED", [O, S],  # short names: the lineage rows say which tables they are
+     "SELECT * FROM orders o JOIN shipments s ON s.order_id = o.order_id WHERE s.shipped_at >= '{date}'"),
+    (6, [2], "SELECT", "FINISHED", None,  # a notebook run: no lineage row, short names cannot be resolved
+     "SELECT * FROM orders o JOIN shipments s ON s.order_id = o.order_id"),
+    (5, [3], "MERGE", "FINISHED", [I],
+     "MERGE INTO demo.sales.returns t USING (SELECT order_id, line_no FROM demo.sales.order_items) s "
+     "ON t.order_id = s.order_id AND t.line_no = s.line_no WHEN MATCHED THEN UPDATE SET t.reason = 'merged'"),
+    (10, [0, 1], "SELECT", "FINISHED", [O, R],
+     "SELECT * FROM demo.sales.orders o WHERE EXISTS (SELECT 1 FROM demo.sales.returns r WHERE r.order_id = o.order_id)"),
+    (60, [0, 1, 2, 3, 4, 5], "SELECT", "FINISHED", [O],
+     "SELECT count(*) FROM demo.sales.orders WHERE status = 'shipped' AND order_date >= '{date}'"),
+    (25, [4, 5], "SELECT", "FINISHED", [T], "SELECT zone, avg(fare) FROM samples.nyctaxi.trips GROUP BY zone"),
+    (20, [0, 2], "SELECT", "FINISHED", [V], "SELECT * FROM demo.sales.v_daily_revenue ORDER BY order_date DESC"),
+    (4, [1], "SELECT", "FINISHED", [O, C],  # does not parse
+     "SELECT o.order_id FROM demo.sales.orders o JOIN demo.sales.customers c ON (o.customer_id = c.customer_id"),
+    (3, [3], "SELECT", "FINISHED", [O], "SELECT * FROM (SELECT * FROM demo.sales.orders"),
+    (3, [2], "SELECT", "FINISHED", [O, C], "SELECT o.order_id FROM demo.sales.orders o LEFT JOIN demo.sales.customers c ON c.customer_id = o.customer_id WHERE"),
+    (5, [4], "SELECT", "FINISHED", [O, C],  # unqualified columns, two tables: no way to know which side each is on
+     "SELECT * FROM demo.sales.orders JOIN demo.sales.customers ON customer_id = lifetime_value"),
+    (6, [1], "SELECT", "FAILED", [O, S],  # failed runs are not read
+     "SELECT * FROM demo.sales.orders o JOIN demo.sales.shipments s ON s.shipped_at = o.order_ts"),
+    (8, [0, 2], "SHOW", "FINISHED", None, "SHOW TABLES IN demo.sales"),
+    (5, [0], "SELECT", "FINISHED", SYS,  # Databricks' own tables are not the user's data model
+     "SELECT * FROM system.access.table_lineage l JOIN system.query.history h ON l.statement_id = h.statement_id"),
+]
+_REDACTED_ROWS = 3  # in an ordinary demo, this many rows read <REDACTED>
+_DATES = ["2026-09-01", "2026-09-15", "2026-10-01", "2026-10-20"]
+_EMAILS = ["jane.doe@example.invalid", "li.wei@example.invalid", "sam.o@example.invalid"]
+
+_cache = {}
+
+
+def redacted_mode():
+    """True when every statement_text should read <REDACTED> (EXPLORER_DEMO_REDACTED=1)."""
+    return os.environ.get("EXPLORER_DEMO_REDACTED") == "1"
+
+
+def _base():
+    return datetime.now(timezone.utc).replace(hour=0, minute=0, second=0, microsecond=0)
+
+
+def _fmt(dt):
+    return dt.strftime("%Y-%m-%dT%H:%M:%S.") + f"{dt.microsecond // 1000:03d}Z"
+
+
+def history():
+    """Every invented statement, oldest first: dicts with the columns of system.query.history and a `tables` list."""
+    base = _base()
+    if base in _cache:
+        return _cache[base]
+    r = random.Random(20261113)
+    rows = []
+    for count, users, kind, status, tables, sql in _SPECS:
+        for _ in range(count):
+            when = base - timedelta(hours=r.uniform(2, 38 * 24))
+            rows.append({"executed_by": USERS[r.choice(users)], "start_time": when, "statement_type": kind,
+                         "execution_status": status, "tables": tables,
+                         "statement_text": sql.format(date=r.choice(_DATES), email=r.choice(_EMAILS))})
+    rows.sort(key=lambda x: x["start_time"])
+    for n, row in enumerate(rows, 1):
+        row["statement_id"] = f"demo-hist-{n:04d}"
+    for k in range(1, _REDACTED_ROWS + 1):  # a few, spread out
+        rows[len(rows) * k // (_REDACTED_ROWS + 1)]["redacted"] = True
+    _cache.clear()
+    _cache[base] = rows
+    return rows
+
+
+def _param_time(params, name):
+    return datetime.fromisoformat(str(params[name]).replace("Z", "+00:00"))
+
+
+def _selected(low, params, cap):
+    """The statements the history query's WHERE and LIMIT would return."""
+    since, until = _param_time(params, "since"), _param_time(params, "until")
+    types = re.search(r"statement_type in \(([^)]*)\)", low)
+    wanted = set(re.findall(r"'(\w+)'", types.group(1))) if types else None
+    rows = [x for x in history()
+            if since <= x["start_time"] < until and x["execution_status"] == "FINISHED"
+            and (wanted is None or x["statement_type"].lower() in {w.lower() for w in wanted})]
+    return rows[:cap]
+
+
+def system_query(low, params, cap):
+    """(columns, rows) for a statement on system.query.history, system.access.table_lineage or column_lineage."""
+    strs = lambda *names: [{"name": n, "type": "STRING"} for n in names]  # noqa: E731
+    if "column_lineage" in low:
+        t = (params or {}).get("t")
+        hits = [c for c in COLUMN_LINEAGE if t in (c[0], c[2])]
+        return (strs("source_table_full_name", "source_column_name", "target_table_full_name", "target_column_name",
+                     "last_seen"), [[a, b, c, d, LINEAGE_SEEN] for a, b, c, d in hits])
+    rows = _selected(low, params, cap)
+    if "table_lineage" in low:
+        return (strs("statement_id", "source_table_full_name"),
+                [[x["statement_id"], t] for x in rows if x["tables"] for t in x["tables"]])
+    redact_all = redacted_mode()
+    return (strs("statement_id", "executed_by", "start_time", "statement_text"),
+            [[x["statement_id"], x["executed_by"], _fmt(x["start_time"]),
+              REDACTED if redact_all or x.get("redacted") else x["statement_text"]] for x in rows])

databricks_explore/genie.py

@@ -0,0 +1,281 @@
+"""Genie One over MCP: a small JSON-RPC client for ask, poll, query result and cancel, plus the response shaping."""
+import json
+import re
+import threading
+import time
+from urllib.parse import unquote
+
+import requests
+
+PROTOCOL = "2025-06-18"
+SERVER_PATH = "/ai-gateway/mcp-services/system.ai.genie_one_mcp"
+CLIENT_INFO = {"name": "databricks-explore", "version": "0"}
+TIMEOUT = 60                      # seconds for one HTTP call; every tool call returns quickly (the browser polls)
+MAX_QUESTION = 4000
+POLL_RETRY_PAUSE = 0.5            # seconds before the one retry of a poll that failed in transit
+MAX_ERROR = 300                   # characters of an upstream error body kept in a message
+
+_PERMISSION = re.compile(r"permission|EXECUTE|denied|forbidden|not authorized|unauthorized", re.I)
+_SESSION = re.compile(r"session", re.I)
+# https://<host>/explore/data/<catalog>/<schema>/<table>[?query][#fragment], as Genie writes them in its answers
+_TABLE_URL = re.compile(r"https?://[^/\s)\]>\"']+/explore/data/([^/\s?#)\]>\"']+)/([^/\s?#)\]>\"']+)/([^/\s?#)\]>\"']+)[^\s)\]>\"']*")
+_AGENT_COMMENT = re.compile(r"<!--.*?-->", re.S)
+
+
+class GenieError(Exception):
+    """Genie could not be reached or refused. `status` is the HTTP status the route returns (403, 429 or 502).
+    `transient` marks a failure in transit (HTTP 5xx, no connection, timeout) that may work on a second try."""
+
+    def __init__(self, message, status=502, transient=False):
+        super().__init__(message)
+        self.status = status
+        self.transient = transient
+
+
+class GenieInvalid(ValueError):
+    """A request body or query the ask routes cannot use (route returns HTTP 400)."""
+
+
+# ---------------------------------------------------------------- validation
+
+def _text(value, name, limit=None):
+    if not isinstance(value, str) or not value.strip():
+        raise GenieInvalid(f"{name}: must be a non-empty string")
+    if limit and len(value) > limit:
+        raise GenieInvalid(f"{name}: at most {limit} characters")
+    return value.strip()
+
+
+def parse_ask(body):
+    """(question, conversation_id or None) from a POST /api/ask body."""
+    if not isinstance(body, dict):
+        raise GenieInvalid("(root): body must be a JSON object")
+    cid = body.get("conversation_id")
+    return _text(body.get("question"), "question", MAX_QUESTION), (_text(cid, "conversation_id") if cid else None)
+
+
+def parse_ids(source, *names):
+    """The named non-empty string ids from a query string or body."""
+    if not hasattr(source, "get"):
+        raise GenieInvalid("(root): body must be a JSON object")
+    return tuple(_text(source.get(n), n) for n in names)
+
+
+# ---------------------------------------------------------------- shaping
+
+def tables_from_markdown(markdown):
+    """Tables an answer links to, as [{"name": "cat.sch.tbl", "url": ...}], first mention first, no repeats."""
+    seen, out = set(), []
+    for m in _TABLE_URL.finditer(markdown or ""):
+        parts = [unquote(g).strip("`") for g in m.groups()]
+        name = ".".join(parts)
+        if all(parts) and name.lower() not in seen:
+            seen.add(name.lower())
+            out.append({"name": name, "url": m.group(0)})
+    return out
+
+
+def shape_ask(sc):
+    """The ask response: ids and the first status."""
+    return {"conversation_id": sc.get("conversation_id"), "response_id": sc.get("response_id"),
+            "status": sc.get("status") or "in_progress"}
+
+
+def shape_poll(sc):
+    """One poll, in the shape the page reads. `tables` is parsed from the answer's links."""
+    status = sc.get("status") or "in_progress"
+    answer = sc.get("final_answer")
+    items = [{"item_id": q.get("item_id"), "sql": q.get("sql") or ""} for q in (sc.get("query_items") or []) if q.get("item_id")]
+    return {"conversation_id": sc.get("conversation_id"), "response_id": sc.get("response_id"), "status": status,
+            "deep_link": sc.get("deep_link"), "final_answer": answer,
+            "progress_steps": [str(s) for s in (sc.get("progress_steps") or [])],
+            "query_items": items, "tables": tables_from_markdown(answer),
+            "error": str(sc.get("error") or sc.get("message")) if status in ("failed", "incomplete") and (sc.get("error") or sc.get("message")) else None}
+
+
+def shape_result(sc):
+    """A query result in the Result shape of POST /api/query (contracts/api.md section 3), so the results grid reads it."""
+    cols = []
+    for c in sc.get("columns") or []:
+        text = c.get("type_text")
+        base = re.split(r"[(<\s]", (text or "").strip())[0].upper() or None
+        cols.append({"name": c.get("name"), "type": c.get("type") or base, "type_text": text})
+    rows = [list(r) for r in (sc.get("rows") or [])]
+    return {"statement_id": sc.get("item_id"), "state": sc.get("statement_state"), "ready": bool(sc.get("ready")),
+            "columns": cols, "rows": rows, "row_count": len(rows), "total_row_count": sc.get("total_row_count"),
+            "truncated": bool(sc.get("truncated")), "ms": None}
+
+
+# ---------------------------------------------------------------- transport
+
+def default_post(url, headers, body, timeout):
+    """POST a JSON body. Returns (status, response headers with lower-case keys, body text). Network failures are GenieError."""
+    try:
+        r = requests.post(url, headers=headers, data=json.dumps(body), timeout=timeout)
+    except requests.Timeout:
+        raise GenieError("The Genie One MCP server did not answer in time. Try again.", 502, transient=True) from None
+    except requests.RequestException as e:
+        raise GenieError(f"Could not reach the Genie One MCP server: {type(e).__name__}. Check the connection and the profile.", 502, transient=True) from None
+    return r.status_code, {k.lower(): v for k, v in r.headers.items()}, r.text
+
+
+def _snippet(text):
+    return " ".join(str(text).split())[:MAX_ERROR]
+
+
+def _permission_message(detail=""):
+    base = ("Genie One refused the request. Your user needs EXECUTE on system.ai.genie_one_mcp, "
+            "and the workspace needs Genie One turned on.")
+    return f"{base} ({_snippet(detail)})" if detail else base
+
+
+def _classify(message):
+    """GenieError for an error message from the server: 403 when it reads like a permission problem, else 502."""
+    return GenieError(_permission_message(message), 403) if _PERMISSION.search(message) else GenieError(_snippet(message) or "Genie failed.", 502)
+
+
+def _parse_body(text, rpc_id):
+    """The JSON-RPC message in a response body: plain JSON, or an event stream with the message in a data line."""
+    text = (text or "").strip()
+    if not text:
+        return None
+    if text.startswith(("{", "[")):
+        data = json.loads(text)
+        return data if isinstance(data, dict) else next((m for m in data if m.get("id") == rpc_id), None)
+    found = None
+    for line in text.splitlines():
+        if line.startswith("data:"):
+            try:
+                msg = json.loads(line[5:].strip())
+            except ValueError:
+                continue
+            if isinstance(msg, dict) and (msg.get("id") == rpc_id or found is None):
+                found = msg
+    return found
+
+
+class GenieClient:
+    """Talks to the Genie One MCP server as the connected user. One MCP session is reused until the server drops it."""
+
+    def __init__(self, w, warehouse_id=None, post=default_post, timeout=TIMEOUT):
+        self.w = w
+        self.warehouse_id = warehouse_id    # sent in _meta with every tool call; the page's selected warehouse
+        self._post = post
+        self._timeout = timeout
+        self._session = None
+        self._next_id = 0
+        self._lock = threading.Lock()
+
+    # ---- plumbing
+    def url(self):
+        """The MCP endpoint of the connected workspace."""
+        try:
+            wid = self.w.get_workspace_id()
+        except Exception as e:
+            raise GenieError(f"Could not read the workspace id: {_snippet(e)}", 502) from None
+        return f"{self.w.config.host.rstrip('/')}{SERVER_PATH}?o={wid}"
+
+    def _headers(self):
+        h = dict(self.w.config.authenticate())  # a fresh bearer token each call; never logged or returned
+        h.update({"Content-Type": "application/json", "Accept": "application/json, text/event-stream",
+                  "MCP-Protocol-Version": PROTOCOL})
+        if self._session:
+            h["Mcp-Session-Id"] = self._session
+        return h
+
+    def _send(self, body):
+        """(status, headers, text) of one POST; an HTTP error is turned into a GenieError except 404, which the caller handles."""
+        status, headers, text = self._post(self.url(), self._headers(), body, self._timeout)
+        if status in (401, 403):
+            raise GenieError(_permission_message(), 403)
+        if status == 429:
+            raise GenieError("Genie is rate limited right now. Wait a moment and try again.", 429)
+        if status >= 400 and status != 404:
+            raise GenieError(f"The Genie One MCP server answered HTTP {status}: {_snippet(text)}", 502, transient=status >= 500)
+        return status, headers, text
+
+    def _rpc_id(self):
+        self._next_id += 1
+        return self._next_id
+
+    def _initialize(self):
+        self._session = None
+        rid = self._rpc_id()
+        status, headers, text = self._send({"jsonrpc": "2.0", "id": rid, "method": "initialize",
+                                            "params": {"protocolVersion": PROTOCOL, "capabilities": {}, "clientInfo": CLIENT_INFO}})
+        if status == 404:
+            raise GenieError("This workspace has no Genie One MCP server (system.ai.genie_one_mcp). It may not be enabled here.", 502)
+        msg = _parse_body(text, rid)
+        if msg and msg.get("error"):
+            raise _classify(str(msg["error"].get("message", msg["error"])))
+        self._session = headers.get("mcp-session-id")
+        self._send({"jsonrpc": "2.0", "method": "notifications/initialized"})
+
+    def _call(self, name, args, retry=True):
+        with self._lock:
+            if not self._session:
+                self._initialize()
+        rid = self._rpc_id()
+        params = {"name": name, "arguments": args}
+        if self.warehouse_id:
+            params["_meta"] = {"warehouse_id": self.warehouse_id}
+        status, _, text = self._send({"jsonrpc": "2.0", "id": rid, "method": "tools/call", "params": params})
+        msg = None
+        try:
+            msg = _parse_body(text, rid) if status != 404 else None
+        except ValueError:
+            raise GenieError("The Genie One MCP server sent a reply that is not JSON.", 502) from None
+        expired = status == 404 or (msg and msg.get("error") and _SESSION.search(str(msg["error"].get("message", ""))))
+        if expired and retry:  # the server forgot the session: start a new one and send the call once more
+            with self._lock:
+                self._session = None
+            return self._call(name, args, retry=False)
+        if status == 404:
+            raise GenieError("The Genie One MCP server lost the session. Ask again.", 502)
+        if msg is None:
+            raise GenieError("The Genie One MCP server sent an empty reply.", 502)
+        if msg.get("error"):
+            raise _classify(str(msg["error"].get("message", msg["error"])))
+        result = msg.get("result") or {}
+        sc = result.get("structuredContent")
+        if sc is None:
+            content = _AGENT_COMMENT.sub("", "".join(c.get("text", "") for c in result.get("content") or []))
+            try:
+                sc = json.loads(content)
+            except ValueError:
+                sc = {"message": content.strip()}
+        if result.get("isError"):
+            raise _classify(str(sc.get("message") or sc.get("error") or json.dumps(sc)))
+        return sc
+
+    # ---- the four tools
+    def ask(self, question, conversation_id=None):
+        """Start a turn (or a follow-up in `conversation_id`). Returns ids and the status."""
+        args = {"question": question}
+        if conversation_id:
+            args["conversation_id"] = conversation_id
+        return shape_ask(self._call("genie_ask", args))
+
+    def poll(self, conversation_id, response_id):
+        """One upstream poll, shaped by shape_poll. A transient failure (HTTP 5xx, no connection, timeout) is tried
+        once more after a short pause; only poll does this, because the page polls again anyway and a repeated ask,
+        result or cancel is not harmless."""
+        args = {"conversation_id": conversation_id, "response_id": response_id}
+        try:
+            sc = self._call("genie_poll_response", args)
+        except GenieError as e:
+            if not e.transient:
+                raise
+            time.sleep(POLL_RETRY_PAUSE)
+            sc = self._call("genie_poll_response", args)
+        return shape_poll(sc)
+
+    def query_result(self, conversation_id, response_id, item_id):
+        """The rows of one SQL item, in the Result shape."""
+        return shape_result(self._call("genie_get_query_result", {"conversation_id": conversation_id,
+                                                                  "response_id": response_id, "item_id": item_id}))
+
+    def cancel(self, conversation_id):
+        """Stop the in-flight turn of a conversation."""
+        self._call("genie_cancel_response", {"conversation_id": conversation_id})
+        return {"conversation_id": conversation_id, "cancel_requested": True}