ripple-sql 0.1.0__py3-none-any.whl

ripple/mcp_server.py

@@ -0,0 +1,600 @@
+"""MCP server over stdio.
+
+Deliberately small: nine deterministic tools, no prompts, no resources.
+The MCP client (Claude Code, Cursor, Codex, ...) does the reasoning; Ripple
+answers lineage questions exactly and admits what it can't resolve.
+
+Hand-rolled JSON-RPC rather than an SDK dependency: the protocol surface
+we need is tiny and the tool must stay a one-dependency install.
+"""
+
+from __future__ import annotations
+
+import difflib
+import json
+import sys
+import threading
+from pathlib import Path
+from typing import Any
+
+from ripple import answer_page
+from ripple.answer import attach, breaks_lines, noun_for, plain_text
+
+PROTOCOL_VERSION = "2025-06-18"
+# Versions whose stdio framing and tools/* surface match what this server
+# implements. Anything else (including future revisions) gets PROTOCOL_VERSION
+# back, per spec, instead of a blind echo that claims semantics we don't have.
+SUPPORTED_PROTOCOL_VERSIONS = {"2024-11-05", "2025-03-26", PROTOCOL_VERSION}
+
+DEPTH_ARG = {
+    "type": "integer",
+    "description": (
+        "how many hops to follow (default 25). Raise it when the result says truncated_at_depth."
+    ),
+}
+PAGE_ARG = {
+    "type": "boolean",
+    "description": (
+        "also write a self-contained HTML answer page under .ripple/answers "
+        "and return its path as html_path, for a person to open in a browser. "
+        "Nothing is sent anywhere."
+    ),
+}
+
+TOOLS = [
+    {
+        "name": "lineage_summary",
+        "description": (
+            "Overview of the SQL project's column lineage: model count, parse rate, "
+            "edge count, and how many edges are uncertain (review_required). "
+            "Call this first."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {"refresh": {"type": "boolean", "description": "Re-scan the project"}},
+        },
+    },
+    {
+        "name": "breaks",
+        "description": (
+            "What breaks if this column changes: every downstream column, model, filter "
+            "and join that depends on it, each with a trust label. Edges marked "
+            "review_required could not be fully verified; treat them as 'check by hand'."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "target": {
+                    "type": "string",
+                    "description": "model.column, e.g. stg_payments.amount",
+                },
+                "depth": DEPTH_ARG,
+                "page": PAGE_ARG,
+            },
+            "required": ["target"],
+        },
+    },
+    {
+        "name": "trace",
+        "description": "Where a column's value comes from, hop by hop, back to sources.",
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "target": {"type": "string", "description": "model.column"},
+                "depth": DEPTH_ARG,
+                "page": PAGE_ARG,
+            },
+            "required": ["target"],
+        },
+    },
+    {
+        "name": "model_columns",
+        "description": "List a model's output columns (use to find the exact column name before calling breaks/trace).",
+        "inputSchema": {
+            "type": "object",
+            "properties": {"model": {"type": "string", "description": "model name"}},
+            "required": ["model"],
+        },
+    },
+    {
+        "name": "unresolved_tables",
+        "description": (
+            "External tables this project reads but does not define, sorted by how much "
+            "coverage they block: each entry has the table name, how many models reference "
+            "it, and how many of those are stuck (star_only or failed) because its columns "
+            "are unknown. Follow-up: Ripple never connects to a warehouse, but you likely "
+            "can. Fetch these tables' columns yourself and hand them back via ingest_schema. "
+            "Canonical query: SELECT table_name, column_name FROM information_schema.columns "
+            "WHERE table_name IN ('t1', 't2') ORDER BY table_name, ordinal_position. "
+            "Dialect notes: Snowflake stores unquoted identifiers uppercase, so match with "
+            "UPPER(table_name) or uppercase literals; BigQuery scopes it per dataset as "
+            "`project.dataset.INFORMATION_SCHEMA.COLUMNS`; Postgres and Redshift use "
+            "information_schema.columns as written (add table_schema = '...' if names "
+            "repeat across schemas). Then call ingest_schema with the results."
+        ),
+        "inputSchema": {"type": "object", "properties": {}},
+    },
+    {
+        "name": "ingest_schema",
+        "description": (
+            "Save warehouse column lists for external tables so Ripple can resolve them. "
+            'Accepts {"tables": {"name": {"columns": ["a", "b"]}}} or, straight '
+            'from a tabular query result, {"rows": [{"table": "name", "column": '
+            '"a"}]}. Table names may be bare or schema qualified; case does not matter. '
+            "Merges into .ripple/schemas.json at the project root (committable, share it "
+            "with your team) and returns what changed. Rebuild happens automatically on "
+            "next query."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "tables": {
+                    "type": "object",
+                    "description": 'table name -> {"columns": [...]} or a bare column list',
+                },
+                "rows": {
+                    "type": "array",
+                    "description": 'alternative tabular form: [{"table": ..., "column": ...}]',
+                    "items": {
+                        "type": "object",
+                        "properties": {
+                            "table": {"type": "string"},
+                            "column": {"type": "string"},
+                        },
+                    },
+                },
+            },
+        },
+    },
+    {
+        "name": "ingest_usage",
+        "description": (
+            "Ingest a query-history export: what actually ran in the warehouse. "
+            "Takes a FILE PATH only, by design: never fetch or paste query rows "
+            "yourself, query text must not pass through a model. Have the user run "
+            'the export with a local command, for example: snow sql -q "SELECT '
+            "query_id, query_text, database_name, schema_name, user_name, "
+            "start_time FROM TABLE(INFORMATION_SCHEMA.QUERY_HISTORY("
+            'RESULT_LIMIT => 10000))" --format json > history.json, then pass '
+            "that path here. JSONL or a JSON array with a query_text field works; "
+            "Snowflake QUERY_HISTORY column names are understood as-is. Ripple "
+            "stores per-table counts only, never the query text."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "path": {
+                    "type": "string",
+                    "description": "path to the export file on this machine",
+                }
+            },
+            "required": ["path"],
+        },
+    },
+    {
+        "name": "collect_usage",
+        "description": (
+            "Collect query history by running the user's already-authenticated "
+            "warehouse CLI (snow, bq, or databricks) locally. Ripple's process "
+            "spawns the CLI, spools the rows to a temp file, aggregates, and "
+            "deletes the file inside one call: the query text never reaches you "
+            "and is never stored, and no credential is read or transmitted. "
+            "Call with no arguments first: you get back which CLIs are installed "
+            "and which connections exist. When several connections are "
+            "configured, NEVER pick one for the user; show the options and ask. "
+            "scope 'mine' (default) reads only the user's own queries and needs "
+            "no permission. scope 'account' covers everyone's queries and may "
+            "need a grant; a permission error includes the exact statement to "
+            "hand an admin. Databricks has no mine/account split: its history "
+            "API returns whatever this login can see, and the result says "
+            "'visible-to-this-login'. If the call times out because a browser sign-in "
+            "window opened, tell the user to finish signing in and call again."
+        ),
+        "inputSchema": {
+            "type": "object",
+            "properties": {
+                "platform": {
+                    "type": "string",
+                    "enum": ["snowflake", "bigquery", "databricks"],
+                    "description": "omit to get the capability table",
+                },
+                "connection": {
+                    "type": "string",
+                    "description": "connection or profile name (bigquery: project id)",
+                },
+                "days": {"type": "integer", "description": "window size, default 7"},
+                "scope": {"type": "string", "enum": ["mine", "account"]},
+                "region": {
+                    "type": "string",
+                    "description": "BigQuery INFORMATION_SCHEMA region, default us",
+                },
+            },
+        },
+    },
+    {
+        "name": "usage_summary",
+        "description": (
+            "Aggregated usage from the ingested query history: which models were "
+            "seen running in the window, which were not, attribution counts, and "
+            "the tables queries read outside this project. Aggregates only, no "
+            "query text. Never report a model as 'unused' from this; the honest "
+            "claim is 'not seen in this window', and the window is in the manifest."
+        ),
+        "inputSchema": {"type": "object", "properties": {}},
+    },
+]
+
+
+INGEST_NEXT = (
+    "For each table, get its columns from the warehouse (DESCRIBE TABLE or "
+    "INFORMATION_SCHEMA.COLUMNS through your warehouse MCP) and call ingest_schema "
+    'with {"tables": {"<table>": {"columns": ["a", "b"]}}}. Tables with '
+    "blocked_models > 0 unlock the most."
+)
+
+
+def _summary_next(graph, summary: dict) -> str:
+    """One sentence an agent can act on, the way the CLI ends on 'try:'.
+
+    Counted before it blames: on balboa only 9 of 17 review links pointed at
+    external tables; the rest came from ambiguous names and an unparsed CTE,
+    which no ingest fixes."""
+    review = summary.get("review_required_edges", 0)
+    if not review:
+        return "Ask breaks(model.column) about any column; model_columns lists a model's columns."
+    review_edges = [e for e in graph.edges if e.trust == "review_required"]
+    external = sum(1 for e in review_edges if "not found" in e.reason)
+    incomplete = [e for e in review_edges if "declared columns of" in e.reason]
+    other = review - external - len(incomplete)
+    parts = []
+    if incomplete:
+        tables = sorted({e.src_model for e in incomplete})[:5]
+        parts.append(
+            f"{len(incomplete)} of the {review} review links read columns that the "
+            f"declared schema lacks ({', '.join(tables)}): re-ingest complete column "
+            "lists for those tables."
+        )
+    if external:
+        parts.append(
+            f"{external} of the {review} review links point at tables outside the "
+            "project: call unresolved_tables, then ingest_schema with their column "
+            "lists, then ask breaks(model.column)."
+        )
+    if other > 0:
+        parts.append(
+            f"{other}{' other' if external or incomplete else ''} review links come from ambiguous "
+            "names, disabled models, or statements that did not fully parse; no ingest "
+            "fixes those, and each breaks hit carries its reason."
+        )
+    return " ".join(parts)
+
+
+def _hits_next(result: dict) -> str:
+    """Counted like the summary: only hits whose reason is an external table
+    are sent to unresolved_tables; the rest are named for what they are."""
+    hits = [h for hits in result.get("by_model", {}).values() for h in hits]
+    review = [h for h in hits if h.get("trust") == "review_required"]
+    external = sum(1 for h in review if "not found" in (h.get("reason") or ""))
+    incomplete = sum(1 for h in review if "declared columns of" in (h.get("reason") or ""))
+    other = len(review) - external - incomplete
+    parts = []
+    if external:
+        parts.append(
+            f"{external} of the {len(review)} review hits trace through tables outside "
+            "the project: unresolved_tables lists them; ingest_schema their columns, "
+            "then ask again."
+        )
+    if incomplete:
+        parts.append(
+            f"{incomplete} read columns the declared schema lacks: re-ingest complete column lists."
+        )
+    if other > 0:
+        parts.append(
+            f"{other}{' other' if parts else ''} review hits come from ambiguous names, "
+            "disabled models, or statements that did not fully parse; each hit carries "
+            "its reason."
+        )
+    return " ".join(parts) or f"{result['review_required']} hits need review."
+
+
+def _depth(arguments: dict) -> int:
+    """The caller's depth if it is a sane integer; the default otherwise.
+
+    Agents send strings, floats and worse; a bad value must degrade to the
+    default, never crash the tool call."""
+    depth = arguments.get("depth")
+    if isinstance(depth, bool) or not isinstance(depth, int) or not 1 <= depth <= 10000:
+        return 25
+    return depth
+
+
+class RippleService:
+    """Builds and caches the lineage graph, rebuilding when the SQL, YAML, or
+    LookML files change.
+
+    The agent's loop is edit SQL then ask breaks; answering from a stale
+    graph would be a quiet lie, so staleness is checked on every call.
+    """
+
+    def __init__(self, path: str = ".", dialect: str | None = None):
+        self.path = path
+        self.dialect = dialect
+        self._graph = None
+        self._fingerprint = None
+        self._lock = threading.RLock()
+
+    def _root(self) -> Path:
+        from ripple.project import find_project_root
+
+        return find_project_root(Path(self.path))
+
+    def _current_fingerprint(self):
+        from ripple.sourcefiles import stamp
+
+        return stamp(self._root())
+
+    def graph(self, refresh: bool = False):
+        # the walk stays outside the lock so requests do not queue behind each other's stats
+        fingerprint = self._current_fingerprint()
+        # serve() calls this from many threads at once; one build, not one per caller
+        with self._lock:
+            if self._graph is None or refresh or fingerprint != self._fingerprint:
+                from ripple import cache
+                from ripple.graph import LineageGraph
+                from ripple.project import load_project
+
+                # Same disk cache as the CLI: an MCP session dies with the client,
+                # and without this every client restart re-pays the full build,
+                # which on a big repo is exactly the tool-call the client times out.
+                root = self._root()
+                graph = None if refresh else cache.load(root, self.dialect)
+                if graph is None:
+                    project = load_project(self.path, dialect=self.dialect)
+                    graph = LineageGraph.build(project)
+                    cache.store(graph, root, self.dialect)
+                self._graph = graph
+                self._fingerprint = fingerprint
+            return self._graph
+
+    def call(self, name: str, arguments: dict[str, Any]) -> dict:
+        if name == "lineage_summary":
+            graph = self.graph(refresh=bool(arguments.get("refresh")))
+            summary = graph.stats()
+            summary["project_warnings"] = graph.project.warnings
+            summary["next"] = _summary_next(graph, summary)
+            return summary
+        if name == "unresolved_tables":
+            unresolved = self.graph().unresolved_tables()
+            return {
+                "unresolved": unresolved,
+                "total": len(unresolved),
+                **(
+                    {"next": INGEST_NEXT}
+                    if unresolved
+                    else {"note": "every referenced table resolves; nothing to ingest"}
+                ),
+            }
+        if name == "ingest_schema":
+            from ripple.schemas import merge_schemas, normalize_tables
+
+            tables = normalize_tables(arguments)
+            before = self.graph().stats()
+            delta = merge_schemas(self._root(), tables)
+            after = self.graph(refresh=True).stats()
+            # the promise of this work is a smaller review count; show whether
+            # it moved, including an honest "still N" (the CLI already does)
+            delta["review_required_before"] = before["review_required_edges"]
+            delta["review_required_after"] = after["review_required_edges"]
+            delta["star_only_before"] = before["star_only"]
+            delta["star_only_after"] = after["star_only"]
+            delta["edges_before"] = before["edges"]
+            delta["edges_after"] = after["edges"]
+            delta["note"] = (
+                "graph rebuilt"
+                if (before["review_required_edges"], before["star_only"])
+                != (after["review_required_edges"], after["star_only"])
+                else "graph rebuilt; nothing moved, so these tables were not what blocked the review links"
+            )
+            delta["next"] = (
+                "Re-ask breaks(model.column), or call lineage_summary to see how many "
+                "links still need review."
+            )
+            return delta
+        if name == "ingest_usage":
+            from ripple.usage import ingest_file
+            from ripple.usage.ingest import write_store
+
+            project = self.graph().project
+            store = ingest_file(
+                arguments["path"],
+                [m.name for m in project.models],
+                project.dialect,
+                model_aliases={m.name: set(m.aliases) for m in project.models},
+            )
+            write_store(self._root(), store)
+            return store
+        if name == "collect_usage":
+            from ripple.usage import collect as collect_mod
+            from ripple.usage.discover import capabilities
+            from ripple.usage.ingest import write_store
+
+            platform = arguments.get("platform")
+            if not platform:
+                return capabilities()
+            project = self.graph().project
+            # 55s: finish (or fail loudly) inside a typical 60s client tool budget
+            store = collect_mod.collect(
+                platform,
+                [m.name for m in project.models],
+                connection=arguments.get("connection"),
+                days=int(arguments.get("days") or 7),
+                scope=arguments.get("scope") or "mine",
+                region=arguments.get("region") or "us",
+                timeout=55.0,
+                model_aliases={m.name: set(m.aliases) for m in project.models},
+            )
+            write_store(self._root(), store)
+            return store
+        if name == "usage_summary":
+            from ripple.usage.ingest import read_store
+
+            store = read_store(self._root())
+            if store is None:
+                return {
+                    "note": (
+                        "no usage ingested yet: call collect_usage to run the "
+                        "user's own warehouse CLI locally, or have the user "
+                        "export query history and call ingest_usage with the path"
+                    )
+                }
+            recorded = store.get("models_fingerprint")
+            if recorded:
+                from ripple.usage.ingest import models_fingerprint
+
+                current = models_fingerprint([m.name for m in self.graph().project.models])
+                if current != recorded:
+                    store = {
+                        **store,
+                        "warning": (
+                            "stale: the project's models changed since this usage "
+                            "was ingested; per-model counts reflect the old model "
+                            "set. Re-run collect_usage or ingest_usage."
+                        ),
+                    }
+            return store
+        from ripple.graph import UnknownTarget
+
+        # the page path is for a person with a screen; a remote agent keeps the text
+        if name == "breaks":
+            graph = self.graph()
+            model, column = self._split(arguments["target"])
+            try:
+                result = graph.breaks(model, column, max_depth=_depth(arguments))
+            except UnknownTarget as e:
+                return {"error": str(e), "did_you_mean": e.suggestions}
+            if result.get("review_required"):
+                result["next"] = _hits_next(result)
+            result = attach(result, "breaks", noun_for(graph.project.mode))
+            if arguments.get("page"):
+                full = plain_text(breaks_lines(result["answer"], full=True))
+                page = answer_page.write_for(graph, self._root(), result["answer"], full)
+                result["html_path"] = str(page)
+            return result
+        if name == "trace":
+            graph = self.graph()
+            model, column = self._split(arguments["target"])
+            try:
+                result = graph.trace(model, column, max_depth=_depth(arguments))
+            except UnknownTarget as e:
+                return {"error": str(e), "did_you_mean": e.suggestions}
+            if any(h.get("trust") == "review_required" for h in result.get("upstream", [])):
+                result["next"] = (
+                    "Some hops need review. unresolved_tables lists the external tables "
+                    "whose columns would settle them; ingest_schema them, then ask again."
+                )
+            result = attach(result, "trace", noun_for(graph.project.mode))
+            if arguments.get("page"):
+                page = answer_page.write_for(graph, self._root(), result["answer"], result["text"])
+                result["html_path"] = str(page)
+            return result
+        if name == "model_columns":
+            graph = self.graph()
+            wanted = arguments["model"].lower()
+            # a dbt-disabled model is analyzed for its own schema but never
+            # binds references; the summary counts it, so it must be nameable
+            # (dbt-ga4 keeps 30 such models behind vars)
+            candidates = graph._candidates.get(wanted) or graph._all_candidates.get(wanted, [])
+            # a name shared by a model and a source is the model's here;
+            # two models sharing it is a real ambiguity, refused with the
+            # spellings that would settle it
+            analyzed = [c for c in candidates if c in graph.reports]
+            exact = [c for c in analyzed if c.lower() == wanted]
+            if len(exact) == 1:
+                analyzed = exact  # the spelling is this model's own name
+            if len(analyzed) > 1:
+                return {
+                    "error": f"'{arguments['model']}' names {len(analyzed)} models; ask by full name",
+                    "did_you_mean": sorted(graph._askable_name(c) for c in analyzed),
+                }
+            report = graph.reports.get(analyzed[0]) if analyzed else None
+            if report is None:
+                known = sorted(graph.reports.keys())
+                close = difflib.get_close_matches(arguments["model"], known, n=5)
+                return {"error": f"unknown model '{arguments['model']}'", "did_you_mean": close}
+            model = next((m for m in graph.project.models if m.name == report.name), None)
+            return {
+                "model": report.name,
+                "columns": report.columns,
+                "status": report.status,
+                "analyzed": report.parsed,
+                **({"enabled": False} if model is not None and not model.enabled else {}),
+            }
+        raise ValueError(f"unknown tool: {name}")
+
+    @staticmethod
+    def _split(target: str) -> tuple[str, str]:
+        model, _, column = target.rpartition(".")
+        if not model:
+            raise ValueError(f"expected model.column, got '{target}'")
+        return model, column
+
+
+def run_stdio(path: str = ".", dialect: str | None = None) -> None:
+    service = RippleService(path=path, dialect=dialect)
+    for line in sys.stdin:
+        line = line.strip()
+        if not line:
+            continue
+        try:
+            message = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        response = _handle(message, service)
+        if response is not None:
+            sys.stdout.write(json.dumps(response) + "\n")
+            sys.stdout.flush()
+
+
+def _handle(message: dict, service: RippleService) -> dict | None:
+    method = message.get("method")
+    msg_id = message.get("id")
+
+    if method == "initialize":
+        requested = message.get("params", {}).get("protocolVersion")
+        version = requested if requested in SUPPORTED_PROTOCOL_VERSIONS else PROTOCOL_VERSION
+        info = {"name": "ripple", "version": _version()}
+        payload = {"protocolVersion": version, "capabilities": {"tools": {}}, "serverInfo": info}
+        return _result(msg_id, payload)
+    if msg_id is None:  # notification
+        return None
+    if method == "tools/list":
+        return _result(msg_id, {"tools": TOOLS})
+    if method == "tools/call":
+        params = message.get("params", {})
+        try:
+            payload = service.call(params.get("name", ""), params.get("arguments", {}) or {})
+            text = json.dumps(payload, indent=2, default=str)
+            return _result(msg_id, {"content": [{"type": "text", "text": text}]})
+        except Exception as e:
+            content = [{"type": "text", "text": f"error: {e}"}]
+            return _result(msg_id, {"content": content, "isError": True})
+    if method == "ping":
+        return _result(msg_id, {})
+    return {
+        "jsonrpc": "2.0",
+        "id": msg_id,
+        "error": {"code": -32601, "message": f"method not found: {method}"},
+    }
+
+
+def _result(msg_id, payload: dict) -> dict:
+    return {"jsonrpc": "2.0", "id": msg_id, "result": payload}
+
+
+def _version() -> str:
+    try:
+        from importlib.metadata import version
+
+        return version("ripple-sql")
+    except Exception:
+        return "0.1.0"

ripple/names.py

@@ -0,0 +1,40 @@
+"""Which spellings name which model.
+
+`LineageGraph._require_target` accepts more than a model's canonical name:
+an alias, a different case, a shorter dotted path. A surface with no graph
+behind it (the answer page, answering a follow-up from an export) has to
+resolve a typed name the same way, so the map it needs is built here rather
+than growing graph.py.
+"""
+
+from __future__ import annotations
+
+
+def askable_names(graph) -> dict[str, list[str]]:
+    """Every spelling the graph accepts for a model, mapped to the canonical
+    name or names it means.
+
+    Only the spellings that carry information: one that already is the
+    canonical name it resolves to is left out, which keeps the map small
+    enough to embed in a page."""
+    merged = {**graph._all_candidates, **graph._candidates}
+    return {
+        spelling: list(canonical)
+        for spelling, canonical in merged.items()
+        if len(canonical) > 1 or canonical[0] != spelling
+    }
+
+
+def unlinked_columns(graph) -> list[str]:
+    """Columns no link touches. `select 1 as x` is a real answer (nothing
+    downstream reads it), so the page has to know the column exists."""
+    linked = set()
+    for edge in graph.edges:
+        linked.add(f"{edge.src_model}.{edge.src_column}")
+        linked.add(f"{edge.dst_model}.{edge.dst_column}")
+    return sorted(
+        column_id
+        for name, report in graph.reports.items()
+        for column in report.columns
+        if column != "*" and (column_id := f"{name}.{column}") not in linked
+    )