pretensor 0.1.0__py3-none-any.whl

pretensor/skills/generator.py

@@ -0,0 +1,235 @@
+"""Generate a compact Claude/Cursor skill markdown from a Kuzu graph."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from pathlib import Path
+
+from pretensor.core.registry import RegistryEntry
+from pretensor.core.store import KuzuStore
+from pretensor.mcp.server import _build_oss_registry
+
+__all__ = ["SkillGenerator", "write_skill_files_for_index"]
+
+
+def _render_tool_lines() -> list[str]:
+    # Source the canonical tool list from the MCP registry so the skill body
+    # never drifts from `pretensor serve`. Handlers are closures that this code
+    # never invokes; the placeholder Path is only captured, never read.
+    registry = _build_oss_registry(Path("."))
+    out: list[str] = []
+    for tool in registry.list_tools():
+        desc = (tool.description or "").strip().split("\n", 1)[0]
+        first_sentence = desc.split(". ", 1)[0].rstrip(".")
+        if first_sentence:
+            out.append(f"- `{tool.name}` — {first_sentence}")
+        else:
+            out.append(f"- `{tool.name}`")
+    return out
+
+
+def write_skill_files_for_index(
+    *,
+    body: str,
+    connection_name: str,
+    cwd: Path | None = None,
+    skills_target: str = "claude",
+) -> list[Path]:
+    """Write skill markdown to one or more paths (``--skills-target``)."""
+    root = (cwd or Path.cwd()).resolve()
+    safe = connection_name.replace("/", "_")
+    key = skills_target.strip().lower() if skills_target.strip() else "claude"
+    out: list[Path] = []
+
+    def _write_claude() -> Path:
+        out_dir = root / ".claude" / "skills" / f"pretensor-{safe}"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        path = out_dir / "SKILL.md"
+        path.write_text(body, encoding="utf-8")
+        return path
+
+    def _write_cursor() -> Path:
+        out_dir = root / ".cursor" / "rules"
+        out_dir.mkdir(parents=True, exist_ok=True)
+        path = out_dir / f"pretensor-{safe}.mdc"
+        path.write_text(body, encoding="utf-8")
+        return path
+
+    if key in ("claude", ""):
+        out.append(_write_claude())
+    elif key == "cursor":
+        out.append(_write_cursor())
+    elif key == "all":
+        out.append(_write_claude())
+        out.append(_write_cursor())
+    else:
+        custom = Path(skills_target).expanduser()
+        if not custom.is_absolute():
+            custom = root / custom
+        custom.parent.mkdir(parents=True, exist_ok=True)
+        custom.write_text(body, encoding="utf-8")
+        out.append(custom)
+    return out
+
+
+@dataclass(frozen=True, slots=True)
+class SkillGenerator:
+    """Build `.claude/skills/pretensor-{db_name}/SKILL.md` after indexing."""
+
+    store: KuzuStore
+    connection_name: str
+    database: str
+    last_indexed_at: datetime
+
+    def render(self) -> str:
+        """Return markdown body for the skill file."""
+        table_count, row_sum = self._table_stats()
+        entities = self._entity_map_lines()
+        top_joins = self._top_joins()
+
+        safe = self.connection_name.replace("/", "_")
+        skill_name = f"pretensor-{safe}"
+        skill_description = (
+            f"Pretensor graph for the `{self.database}` database "
+            f"(connection `{self.connection_name}`, {table_count} tables). "
+            "Use to look up tables, joins, and MCP query tools for this database."
+        )
+
+        lines = [
+            "---",
+            f"name: {skill_name}",
+            f"description: {skill_description}",
+            "---",
+            "",
+            f"# Pretensor graph: {self.database}",
+            "",
+            "## Database overview",
+            "",
+            f"- **Connection name:** `{self.connection_name}`",
+            f"- **Logical database:** `{self.database}`",
+            f"- **Tables:** {table_count}",
+            f"- **Approx. row count (sum over tables):** {row_sum}",
+            f"- **Last indexed:** {self.last_indexed_at.isoformat()}",
+            "",
+            "## Key entities",
+            "",
+        ]
+        if entities:
+            lines.extend(entities)
+        else:
+            lines.append(
+                "_No Entity nodes yet — run indexing with entity extraction enabled._"
+            )
+        lines.extend(
+            [
+                "",
+                "## MCP tools (Pretensor)",
+                "",
+                "Use the **pretensor** MCP server (`pretensor serve`):",
+                "",
+                *_render_tool_lines(),
+                "",
+                "Resources: `pretensor://databases`, `pretensor://db/{name}/overview`.",
+                "",
+                "## Common patterns",
+                "",
+                "_Placeholder — join-path and query patterns land in a later epic._",
+                "",
+                "## Top relationships (sample)",
+                "",
+            ]
+        )
+        if top_joins:
+            lines.extend(top_joins)
+        else:
+            lines.append("_No join edges recorded._")
+        lines.append("")
+        return "\n".join(lines)
+
+    def _table_stats(self) -> tuple[int, int]:
+        rows = self.store.query_all_rows(
+            "MATCH (t:SchemaTable) RETURN count(*), sum(t.row_count)"
+        )
+        if not rows:
+            return 0, 0
+        c, s = rows[0]
+        tc = int(c) if c is not None else 0
+        rs = int(s) if s is not None else 0
+        return tc, rs
+
+    def _entity_map_lines(self) -> list[str]:
+        pairs = self.store.query_all_rows(
+            """
+            MATCH (e:Entity)-[:REPRESENTS]->(t:SchemaTable)
+            RETURN e.name, t.schema_name, t.table_name
+            ORDER BY e.name, t.schema_name, t.table_name
+            """
+        )
+        by_entity: dict[str, list[str]] = {}
+        for name, schema_name, table_name in pairs:
+            key = str(name)
+            qn = f"{schema_name}.{table_name}"
+            by_entity.setdefault(key, []).append(qn)
+        names_only = self.store.query_all_rows(
+            "MATCH (e:Entity) RETURN e.name ORDER BY e.name"
+        )
+        out: list[str] = []
+        for (n,) in names_only:
+            key = str(n)
+            tlist = sorted(set(by_entity.get(key, [])))
+            if tlist:
+                joined = ", ".join(tlist)
+                out.append(f"- **{key}** — {joined}")
+            else:
+                out.append(f"- **{key}**")
+        return out
+
+    def _top_joins(self) -> list[str]:
+        fk = self.store.query_all_rows(
+            """
+            MATCH (a:SchemaTable)-[:FK_REFERENCES]->(b:SchemaTable)
+            RETURN concat(a.schema_name, '.', a.table_name),
+                   concat(b.schema_name, '.', b.table_name)
+            LIMIT 8
+            """
+        )
+        inf = self.store.query_all_rows(
+            """
+            MATCH (a:SchemaTable)-[r:INFERRED_JOIN]->(b:SchemaTable)
+            RETURN concat(a.schema_name, '.', a.table_name),
+                   concat(b.schema_name, '.', b.table_name),
+                   r.confidence
+            LIMIT 8
+            """
+        )
+        lines: list[str] = []
+        for src, dst in fk:
+            lines.append(f"- `{src}` —[FK_REFERENCES]→ `{dst}`")
+        for src, dst, conf in inf:
+            conf_s = f" (confidence {float(conf):.2f})" if conf is not None else ""
+            lines.append(f"- `{src}` —[INFERRED_JOIN]→ `{dst}`{conf_s}")
+        return lines[:12]
+
+    @classmethod
+    def write_for_index(
+        cls,
+        *,
+        store: KuzuStore,
+        entry: RegistryEntry,
+        cwd: Path | None = None,
+        skills_target: str = "claude",
+    ) -> list[Path]:
+        """Write skill file(s) per ``skills_target`` (see ``write_skill_files_for_index``)."""
+        body = cls(
+            store=store,
+            connection_name=entry.connection_name,
+            database=entry.database,
+            last_indexed_at=entry.last_indexed_at,
+        ).render()
+        return write_skill_files_for_index(
+            body=body,
+            connection_name=entry.connection_name,
+            cwd=cwd,
+            skills_target=skills_target,
+        )

pretensor/staleness/__init__.py

@@ -0,0 +1,15 @@
+"""Schema drift: snapshots, graph patching, and impact analysis."""
+
+from __future__ import annotations
+
+from pretensor.staleness.graph_patcher import GraphPatcher, PatchResult
+from pretensor.staleness.impact_analyzer import ImpactAnalyzer, ImpactReport
+from pretensor.staleness.snapshot_store import SnapshotStore
+
+__all__ = [
+    "GraphPatcher",
+    "ImpactAnalyzer",
+    "ImpactReport",
+    "PatchResult",
+    "SnapshotStore",
+]

pretensor/staleness/graph_patcher.py

@@ -0,0 +1,355 @@
+"""Apply :class:`SchemaChange` list to the Kuzu graph (incremental reindex)."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+
+from pretensor.connectors.models import Column, SchemaSnapshot, Table
+from pretensor.connectors.snapshot import ChangeTarget, ChangeType, SchemaChange
+from pretensor.core.ids import column_node_id, fk_edge_id, table_node_id
+from pretensor.core.store import KuzuStore
+from pretensor.graph_models.edge import GraphEdge
+from pretensor.graph_models.node import GraphNode
+
+__all__ = ["GraphPatcher", "PatchResult"]
+
+
+@dataclass
+class PatchResult:
+    """Counts of mutations applied (or that would be applied in dry-run)."""
+
+    tables_added: int = 0
+    tables_removed: int = 0
+    columns_added: int = 0
+    columns_removed: int = 0
+    columns_updated: int = 0
+    fk_edges_removed: int = 0
+    inferred_joins_removed: int = 0
+    join_paths_marked_stale: int = 0
+    cluster_memberships_removed: int = 0
+    clusters_marked_stale: int = 0
+    foreign_keys_synced: int = 0
+    notes: list[str] = field(default_factory=list)
+
+
+class GraphPatcher:
+    """Translate pipeline :class:`SchemaChange` rows into Kuzu writes."""
+
+    def __init__(self, store: KuzuStore) -> None:
+        self._store = store
+
+    def apply(
+        self,
+        changes: list[SchemaChange],
+        new_snapshot: SchemaSnapshot,
+        *,
+        dry_run: bool = False,
+    ) -> PatchResult:
+        """Apply changes using ``new_snapshot`` for full row data.
+
+        Args:
+            changes: Output of :func:`pretensor.connectors.snapshot.diff_snapshots`.
+            new_snapshot: Fresh introspection result (source of truth for adds/updates).
+            dry_run: If True, compute counts only (no graph writes).
+
+        Returns:
+            Summary of operations.
+        """
+        self._store.ensure_schema()
+        result = PatchResult()
+        tables_by_key = {(t.schema_name, t.name): t for t in new_snapshot.tables}
+        connection_name = new_snapshot.connection_name
+        database_key = new_snapshot.database
+
+        def invalidate_table(schema_name: str, table_name: str) -> None:
+            tid = table_node_id(connection_name, schema_name, table_name)
+            if dry_run:
+                return
+            self._store.mark_join_paths_stale_for_table(tid, database_key)
+            self._store.remove_table_from_clusters(tid)
+            self._store.mark_clusters_stale_for_table(tid, database_key)
+
+        for ch in changes:
+            key = (ch.schema_name, ch.table_name)
+            if ch.change_type == ChangeType.ADDED and ch.target == ChangeTarget.TABLE:
+                tbl = tables_by_key.get(key)
+                if tbl is None:
+                    result.notes.append(f"skip ADDED table missing in snapshot: {key}")
+                    continue
+                if dry_run:
+                    result.tables_added += 1
+                    result.columns_added += len(tbl.columns)
+                    continue
+                invalidate_table(ch.schema_name, ch.table_name)
+                grants_json = (
+                    json.dumps(tbl.grants, ensure_ascii=False) if tbl.grants else None
+                )
+                node = GraphNode(
+                    node_id=table_node_id(connection_name, tbl.schema_name, tbl.name),
+                    connection_name=connection_name,
+                    database=new_snapshot.database,
+                    schema_name=tbl.schema_name,
+                    table_name=tbl.name,
+                    row_count=tbl.row_count,
+                    comment=tbl.comment,
+                    table_type=tbl.table_type,
+                    seq_scan_count=tbl.seq_scan_count,
+                    idx_scan_count=tbl.idx_scan_count,
+                    insert_count=tbl.insert_count,
+                    update_count=tbl.update_count,
+                    delete_count=tbl.delete_count,
+                    is_partitioned=tbl.is_partitioned,
+                    partition_key=tbl.partition_key,
+                    grants_json=grants_json,
+                    access_read_count=tbl.access_read_count,
+                    access_write_count=tbl.access_write_count,
+                    days_since_last_access=tbl.days_since_last_access,
+                    potentially_unused=tbl.potentially_unused,
+                    table_bytes=tbl.table_bytes,
+                    clustering_key=tbl.clustering_key,
+                )
+                self._store.upsert_table(node)
+                for col in tbl.columns:
+                    self._upsert_column(
+                        connection_name, new_snapshot.database, tbl, col
+                    )
+                result.tables_added += 1
+                result.columns_added += len(tbl.columns)
+
+            elif (
+                ch.change_type == ChangeType.REMOVED and ch.target == ChangeTarget.TABLE
+            ):
+                tid = table_node_id(connection_name, ch.schema_name, ch.table_name)
+                if dry_run:
+                    result.tables_removed += 1
+                    continue
+                self._store.delete_join_paths_for_table(tid, database_key)
+                self._store.remove_table_from_clusters(tid)
+                self._store.mark_clusters_stale_for_table(tid, database_key)
+                self._store.delete_table_node_cascade(tid)
+                result.tables_removed += 1
+
+            elif (
+                ch.change_type == ChangeType.ADDED and ch.target == ChangeTarget.COLUMN
+            ):
+                tbl = tables_by_key.get(key)
+                if tbl is None or not ch.column_name:
+                    continue
+                col = next((c for c in tbl.columns if c.name == ch.column_name), None)
+                if col is None:
+                    continue
+                if dry_run:
+                    result.columns_added += 1
+                    continue
+                invalidate_table(ch.schema_name, ch.table_name)
+                self._upsert_column(connection_name, new_snapshot.database, tbl, col)
+                result.columns_added += 1
+
+            elif (
+                ch.change_type == ChangeType.REMOVED
+                and ch.target == ChangeTarget.COLUMN
+            ):
+                if not ch.column_name:
+                    continue
+                tid = table_node_id(connection_name, ch.schema_name, ch.table_name)
+                cid = column_node_id(
+                    connection_name, ch.schema_name, ch.table_name, ch.column_name
+                )
+                if dry_run:
+                    result.columns_removed += 1
+                    result.fk_edges_removed += self._count_fk_touching(
+                        tid, ch.column_name
+                    )
+                    result.inferred_joins_removed += self._count_inferred_touching(
+                        tid, ch.column_name
+                    )
+                    continue
+                invalidate_table(ch.schema_name, ch.table_name)
+                result.fk_edges_removed += self._store.delete_fk_edges_touching_column(
+                    tid, ch.column_name
+                )
+                result.inferred_joins_removed += (
+                    self._store.delete_inferred_joins_touching_column(
+                        tid, ch.column_name
+                    )
+                )
+                self._store.delete_column_node(cid)
+                result.columns_removed += 1
+
+            elif (
+                ch.change_type == ChangeType.MODIFIED
+                and ch.target == ChangeTarget.COLUMN
+            ):
+                tbl = tables_by_key.get(key)
+                if tbl is None or not ch.column_name:
+                    continue
+                col = next((c for c in tbl.columns if c.name == ch.column_name), None)
+                if col is None:
+                    continue
+                if dry_run:
+                    result.columns_updated += 1
+                    continue
+                invalidate_table(ch.schema_name, ch.table_name)
+                self._upsert_column(connection_name, new_snapshot.database, tbl, col)
+                result.columns_updated += 1
+
+            elif (
+                ch.change_type == ChangeType.MODIFIED
+                and ch.target == ChangeTarget.TABLE
+            ):
+                tbl = tables_by_key.get(key)
+                if tbl is None:
+                    continue
+                if dry_run:
+                    result.columns_updated += len(tbl.columns)
+                    continue
+                invalidate_table(ch.schema_name, ch.table_name)
+                grants_json = (
+                    json.dumps(tbl.grants, ensure_ascii=False) if tbl.grants else None
+                )
+                node = GraphNode(
+                    node_id=table_node_id(connection_name, tbl.schema_name, tbl.name),
+                    connection_name=connection_name,
+                    database=new_snapshot.database,
+                    schema_name=tbl.schema_name,
+                    table_name=tbl.name,
+                    row_count=tbl.row_count,
+                    comment=tbl.comment,
+                    table_type=tbl.table_type,
+                    seq_scan_count=tbl.seq_scan_count,
+                    idx_scan_count=tbl.idx_scan_count,
+                    insert_count=tbl.insert_count,
+                    update_count=tbl.update_count,
+                    delete_count=tbl.delete_count,
+                    is_partitioned=tbl.is_partitioned,
+                    partition_key=tbl.partition_key,
+                    grants_json=grants_json,
+                    access_read_count=tbl.access_read_count,
+                    access_write_count=tbl.access_write_count,
+                    days_since_last_access=tbl.days_since_last_access,
+                    potentially_unused=tbl.potentially_unused,
+                    table_bytes=tbl.table_bytes,
+                    clustering_key=tbl.clustering_key,
+                )
+                self._store.upsert_table(node)
+                for col in tbl.columns:
+                    self._upsert_column(
+                        connection_name, new_snapshot.database, tbl, col
+                    )
+                result.columns_updated += len(tbl.columns)
+
+        if not dry_run:
+            result.foreign_keys_synced = self._sync_foreign_keys(new_snapshot)
+
+        return result
+
+    def _upsert_column(
+        self,
+        connection_name: str,
+        database: str,
+        table: Table,
+        col: Column,
+    ) -> None:
+        tid = table_node_id(connection_name, table.schema_name, table.name)
+        cid = column_node_id(connection_name, table.schema_name, table.name, col.name)
+        mcv_json = (
+            json.dumps(col.most_common_values, ensure_ascii=False)
+            if col.most_common_values
+            else None
+        )
+        hb_json = (
+            json.dumps(col.histogram_bounds, ensure_ascii=False)
+            if col.histogram_bounds
+            else None
+        )
+        self._store.upsert_column_for_table(
+            column_node_id=cid,
+            connection_name=connection_name,
+            database=database,
+            schema_name=table.schema_name,
+            table_name=table.name,
+            column_name=col.name,
+            data_type=col.data_type,
+            nullable=col.nullable,
+            is_primary_key=col.is_primary_key,
+            is_foreign_key=col.is_foreign_key,
+            table_node_id=tid,
+            comment=col.comment,
+            default_value=col.default_value,
+            is_indexed=col.is_indexed,
+            check_constraints=col.check_constraints,
+            ordinal_position=col.ordinal_position,
+            most_common_values_json=mcv_json,
+            histogram_bounds_json=hb_json,
+            stats_correlation=col.stats_correlation,
+        )
+
+    def _sync_foreign_keys(self, snapshot: SchemaSnapshot) -> int:
+        """Replace FK edges with the set implied by ``snapshot``."""
+        known = {
+            table_node_id(snapshot.connection_name, t.schema_name, t.name)
+            for t in snapshot.tables
+        }
+        new_edge_ids: set[str] = set()
+        edges: list[GraphEdge] = []
+        for t in snapshot.tables:
+            src_id = table_node_id(snapshot.connection_name, t.schema_name, t.name)
+            for fk in t.foreign_keys:
+                dst_id = table_node_id(
+                    snapshot.connection_name, fk.target_schema, fk.target_table
+                )
+                if src_id not in known or dst_id not in known:
+                    continue
+                edge = GraphEdge(
+                    edge_id=fk_edge_id(fk),
+                    source_node_id=src_id,
+                    target_node_id=dst_id,
+                    source_column=fk.source_column,
+                    target_column=fk.target_column,
+                    constraint_name=fk.constraint_name,
+                )
+                new_edge_ids.add(edge.edge_id)
+                edges.append(edge)
+
+        rows = self._store.query_all_rows(
+            """
+            MATCH (src:SchemaTable)-[r:FK_REFERENCES]->(dst:SchemaTable)
+            WHERE src.connection_name = $cn
+            RETURN r.edge_id
+            """,
+            {"cn": snapshot.connection_name},
+        )
+        old_ids = {str(r[0]) for r in rows if r[0] is not None}
+        for eid in old_ids - new_edge_ids:
+            self._store.execute(
+                "MATCH ()-[r:FK_REFERENCES {edge_id: $eid}]->() DELETE r",
+                {"eid": eid},
+            )
+        for edge in edges:
+            self._store.merge_fk_edge(edge)
+        return len(edges)
+
+    def _count_fk_touching(self, table_node_id: str, column_name: str) -> int:
+        rows = self._store.query_all_rows(
+            """
+            MATCH (src:SchemaTable)-[r:FK_REFERENCES]->(dst:SchemaTable)
+            WHERE (src.node_id = $tid AND r.source_column = $col)
+               OR (dst.node_id = $tid AND r.target_column = $col)
+            RETURN count(*)
+            """,
+            {"tid": table_node_id, "col": column_name},
+        )
+        return int(rows[0][0]) if rows and rows[0][0] is not None else 0
+
+    def _count_inferred_touching(self, table_node_id: str, column_name: str) -> int:
+        rows = self._store.query_all_rows(
+            """
+            MATCH (src:SchemaTable)-[r:INFERRED_JOIN]->(dst:SchemaTable)
+            WHERE (src.node_id = $tid AND r.source_column = $col)
+               OR (dst.node_id = $tid AND r.target_column = $col)
+            RETURN count(*)
+            """,
+            {"tid": table_node_id, "col": column_name},
+        )
+        return int(rows[0][0]) if rows and rows[0][0] is not None else 0