dataform-context-mcp 0.5.0__py3-none-any.whl

dataform_context_mcp/__init__.py

@@ -0,0 +1,3 @@
+"""Deterministic, self-hosted MCP server exposing structured Dataform pipeline context."""
+
+__version__ = "0.1.0"

dataform_context_mcp/cli.py

@@ -0,0 +1,229 @@
+"""Command-line entry point: index | report | serve | validate-golden."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+import time
+from pathlib import Path
+
+from .compile import load_graph
+from .db import (
+    AmbiguousError,
+    NotFoundError,
+    get_action,
+    get_meta,
+    open_db,
+    rebuild,
+    resolve,
+    traverse,
+)
+from .golden import GOLDEN_RELPATH, locate_golden, validate_entries
+from .indexer import default_db_path, ensure_fresh
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="dataform-context",
+        description="Deterministic Dataform pipeline context: indexing CLI and MCP server.",
+    )
+    subparsers = parser.add_subparsers(dest="command")
+
+    p_index = subparsers.add_parser(
+        "index", help="Compile the Dataform repo and (re)build the context index"
+    )
+    p_index.add_argument(
+        "--repo",
+        type=Path,
+        default=Path("."),
+        help="Path to the Dataform repository (default: current directory)",
+    )
+    p_index.add_argument("--db", type=Path, help="Index database path (default: cache dir)")
+    p_index.add_argument(
+        "--from-json",
+        dest="from_json",
+        type=Path,
+        help="Load a captured `dataform compile --json` output instead of compiling (tests)",
+    )
+
+    p_report = subparsers.add_parser("report", help="Print a summary of the indexed graph")
+    p_report.add_argument(
+        "--repo",
+        type=Path,
+        default=Path("."),
+        help="Path to the Dataform repository (default: current directory)",
+    )
+    p_report.add_argument("--db", type=Path, help="Index database path (default: cache dir)")
+    p_report.add_argument("--table", help="Print the full context of one table as JSON")
+
+    p_serve = subparsers.add_parser("serve", help="Run the MCP server (stdio)")
+    p_serve.add_argument(
+        "--repo",
+        type=Path,
+        default=Path("."),
+        help="Path to the Dataform repository (default: current directory — MCP clients "
+        "launch project servers from the project root)",
+    )
+    p_serve.add_argument("--db", type=Path)
+
+    p_golden = subparsers.add_parser(
+        "validate-golden",
+        help="Validate column lineage against a golden file "
+        "(JSON list of {table, column, direction, depth, expected_edges, expect_complete})",
+    )
+    p_golden.add_argument("--repo", type=Path, default=Path("."))
+    p_golden.add_argument("--db", type=Path)
+    p_golden.add_argument(
+        "--golden",
+        type=Path,
+        help="Golden file (default: <repo>/.dataform-context/golden_columns.json)",
+    )
+    return parser
+
+
+def _resolve_db(args: argparse.Namespace, parser: argparse.ArgumentParser) -> Path:
+    if args.db:
+        return args.db
+    if args.repo:
+        return default_db_path(args.repo)
+    parser.error(f"{args.command}: --repo or --db is required")
+    raise AssertionError  # unreachable
+
+
+def cmd_index(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
+    db_path = _resolve_db(args, parser)
+    started = time.perf_counter()
+    if args.from_json:
+        graph = load_graph(json.loads(args.from_json.read_text()))
+        conn = open_db(db_path)
+        rebuild(conn, graph, source_hash="manual")
+        meta = get_meta(conn)
+        counts = meta["counts"]
+    else:
+        if not args.repo:
+            parser.error("index: --repo is required unless --from-json is given")
+        conn, index_meta = ensure_fresh(args.repo, db_path)
+        if index_meta["compile_status"] == "error":
+            print(f"compile error:\n{index_meta['compile_error']}", file=sys.stderr)
+            return 1
+        counts = index_meta["counts"]
+    elapsed = time.perf_counter() - started
+    print(f"db: {db_path}")
+    print(f"actions: {counts['actions']}")
+    print(f"table_edges: {counts['table_edges']}")
+    print(f"compile_seconds: {elapsed:.1f}")
+    return 0
+
+
+def cmd_report(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
+    db_path = _resolve_db(args, parser)
+    if not Path(db_path).exists():
+        print(f"no index at {db_path} — run `dataform-context index` first", file=sys.stderr)
+        return 2
+    conn = open_db(db_path)
+    if args.table:
+        try:
+            action_id = resolve(conn, args.table)
+        except NotFoundError as err:
+            print(f"'{args.table}' not found. Suggestions: {err.suggestions}", file=sys.stderr)
+            return 2
+        except AmbiguousError as err:
+            print(f"'{args.table}' is ambiguous. Candidates: {err.candidates}", file=sys.stderr)
+            return 2
+        action = get_action(conn, action_id)
+        detail = {
+            "canonical": action["canonical"],
+            "type": action["action_type"],
+            "layer": action["layer"],
+            "file": action["file_name"],
+            "description": action["description"],
+            "tags": action["tags"],
+            "columns": action["columns"],
+            "upstream": [
+                n["canonical"] for lvl in traverse(conn, action_id, "up", 1) for n in lvl["nodes"]
+            ],
+            "downstream": [
+                n["canonical"] for lvl in traverse(conn, action_id, "down", 1) for n in lvl["nodes"]
+            ],
+        }
+        print(json.dumps(detail, indent=2))
+        return 0
+    meta = get_meta(conn)
+    print(f"indexed_at: {meta.get('indexed_at', '?')}")
+    print(f"actions: {meta['counts']['actions']}")
+    print(f"table_edges: {meta['counts']['table_edges']}")
+    print("by layer:")
+    for layer, count in conn.execute(
+        "SELECT COALESCE(layer, '(none)'), COUNT(*) FROM actions GROUP BY layer ORDER BY layer"
+    ):
+        print(f"  {layer}: {count}")
+    extraction = meta.get("column_extraction", {})
+    if extraction:
+        analyzable = sum(count for status, count in extraction.items() if status != "source")
+        pct_ok = 100 * extraction.get("ok", 0) / analyzable if analyzable else 0.0
+        print("column extraction:")
+        for status in sorted(extraction):
+            print(f"  {status}: {extraction[status]}")
+        print(f"  pct_ok (hors source): {pct_ok:.0f}%")
+    return 0
+
+
+def cmd_validate_golden(args: argparse.Namespace, parser: argparse.ArgumentParser) -> int:
+    db_path = _resolve_db(args, parser)
+    if not Path(db_path).exists():
+        print(f"no index at {db_path} — run `dataform-context index` first", file=sys.stderr)
+        return 2
+    misplaced = None
+    if args.golden:
+        golden_path = args.golden
+    elif args.repo:
+        golden_path, misplaced = locate_golden(args.repo)
+    else:
+        golden_path = None
+    if golden_path is None or not Path(golden_path).exists():
+        print(f"golden file not found: {golden_path}", file=sys.stderr)
+        if misplaced:
+            print(
+                f"found one at {misplaced} instead — move it to "
+                f"{GOLDEN_RELPATH.as_posix()}, or pass --golden {misplaced}",
+                file=sys.stderr,
+            )
+        return 2
+    entries = json.loads(Path(golden_path).read_text())
+    conn = open_db(db_path)
+    summary = validate_entries(conn, entries)
+    for result in summary["results"]:
+        if result["ok"]:
+            print(
+                f"PASS {result['label']} ({result['edges']} edges, complete={result['complete']})"
+            )
+        else:
+            print(f"FAIL {result['label']}: " + " | ".join(result["problems"]))
+    print(f"{summary['passed']}/{summary['total']} golden entries passed")
+    return 0 if summary["passed"] == summary["total"] else 1
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    if args.command is None:
+        parser.print_help()
+        return 1
+    if args.command == "index":
+        return cmd_index(args, parser)
+    if args.command == "report":
+        return cmd_report(args, parser)
+    if args.command == "serve":
+        from .server import create_server
+
+        create_server(args.repo, args.db).run("stdio")
+        return 0
+    if args.command == "validate-golden":
+        return cmd_validate_golden(args, parser)
+    print(f"{args.command}: not implemented", file=sys.stderr)
+    return 1
+
+
+if __name__ == "__main__":
+    sys.exit(main())

dataform_context_mcp/compile.py

@@ -0,0 +1,103 @@
+"""Run `dataform compile --json` and load the CompiledGraph into the typed model."""
+
+from __future__ import annotations
+
+import json
+import subprocess
+from pathlib import Path
+
+from .model import Action, ColumnDoc, CompiledGraphData, Target
+
+_STDERR_CAP = 4000
+
+# enumType of entries in the `tables` array; anything unexpected degrades to "table".
+_ENUM_TYPES = {"TABLE": "table", "VIEW": "view", "INCREMENTAL": "incremental"}
+
+
+class CompileError(Exception):
+    def __init__(self, message: str, exit_code: int | None = None, stderr: str = ""):
+        super().__init__(message)
+        self.exit_code = exit_code
+        self.stderr = stderr
+
+
+def run_dataform_compile(repo: Path, timeout: float = 120.0) -> dict:
+    repo = Path(repo)
+    if not repo.is_dir():
+        raise CompileError(f"repo not found: {repo}", stderr=f"not a directory: {repo}")
+    try:
+        proc = subprocess.run(
+            ["dataform", "compile", "--json"],
+            cwd=repo,
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+        )
+    except FileNotFoundError as err:
+        raise CompileError("dataform CLI not found on PATH", stderr=str(err)) from err
+    except subprocess.TimeoutExpired as err:
+        raise CompileError(f"dataform compile timed out after {timeout}s", stderr=str(err)) from err
+    if proc.returncode != 0:
+        raise CompileError(
+            f"dataform compile failed (exit {proc.returncode})",
+            exit_code=proc.returncode,
+            stderr=(proc.stderr or proc.stdout)[-_STDERR_CAP:],
+        )
+    try:
+        raw = json.loads(proc.stdout)
+    except json.JSONDecodeError as err:
+        raise CompileError("dataform compile produced invalid JSON", stderr=str(err)) from err
+    compilation_errors = (raw.get("graphErrors") or {}).get("compilationErrors") or []
+    if compilation_errors:
+        raise CompileError(
+            f"dataform compilation errors ({len(compilation_errors)})",
+            stderr=json.dumps(compilation_errors)[:_STDERR_CAP],
+        )
+    return raw
+
+
+def load_graph(raw: dict) -> CompiledGraphData:
+    actions: list[Action] = []
+    for entry in raw.get("tables") or []:
+        actions.append(_load_action(entry, _ENUM_TYPES.get(entry.get("enumType", ""), "table")))
+    for entry in raw.get("declarations") or []:
+        actions.append(_load_action(entry, "declaration"))
+    for entry in raw.get("operations") or []:
+        actions.append(_load_action(entry, "operations"))
+    for entry in raw.get("assertions") or []:
+        actions.append(_load_action(entry, "assertion"))
+    return CompiledGraphData(actions=actions, dataform_version=raw.get("dataformCoreVersion"))
+
+
+def _load_target(entry: dict) -> Target:
+    return Target(
+        database=entry.get("database", ""),
+        schema=entry.get("schema", ""),
+        name=entry.get("name", ""),
+    )
+
+
+def _load_action(entry: dict, action_type: str) -> Action:
+    descriptor = entry.get("actionDescriptor") or {}
+    columns = [
+        ColumnDoc(name=".".join(col["path"]), description=col.get("description"))
+        for col in descriptor.get("columns") or []
+        if col.get("path")
+    ]
+    # operations carry a `queries` array instead of `query`
+    query = entry.get("query")
+    if not query and entry.get("queries"):
+        query = ";\n".join(q for q in entry["queries"] if q)
+    return Action(
+        target=_load_target(entry.get("target") or {}),
+        action_type=action_type,
+        file_name=entry.get("fileName"),
+        tags=list(entry.get("tags") or []),
+        description=descriptor.get("description"),
+        columns=columns,
+        dependency_targets=[_load_target(t) for t in entry.get("dependencyTargets") or []],
+        query=query,
+        incremental_query=entry.get("incrementalQuery"),
+        unique_key=list(entry.get("uniqueKey") or []),
+        disabled=bool(entry.get("disabled", False)),
+    )