pretensor 0.1.0__py3-none-any.whl

pretensor/core/ids.py

@@ -0,0 +1,101 @@
+"""Stable identifiers for graph nodes and explicit FK edges."""
+
+from __future__ import annotations
+
+import hashlib
+
+from pretensor.connectors.models import ForeignKey, ViewDependency
+
+__all__ = [
+    "table_node_id",
+    "column_node_id",
+    "fk_edge_id",
+    "lineage_edge_id",
+    "dbt_lineage_edge_id",
+    "entity_node_id",
+    "metric_template_node_id",
+    "external_consumer_node_id",
+    "consumes_edge_id",
+]
+
+
+def table_node_id(connection_name: str, schema_name: str, table_name: str) -> str:
+    """Stable node id for a physical table within one indexed connection."""
+    return f"{connection_name}::{schema_name}::{table_name}"
+
+
+def column_node_id(
+    connection_name: str, schema_name: str, table_name: str, column_name: str
+) -> str:
+    """Stable node id for a physical column."""
+    return f"{connection_name}::{schema_name}::{table_name}::{column_name}"
+
+
+def entity_node_id(connection_name: str, entity_name: str) -> str:
+    """Stable node id for a business entity within one indexed connection."""
+    return f"{connection_name}::entity::{entity_name}"
+
+
+def metric_template_node_id(
+    connection_name: str, database: str, template_name: str
+) -> str:
+    """Stable node id for a metric template scoped to one logical database."""
+    return f"{connection_name}::{database}::metric::{template_name}"
+
+
+def fk_edge_id(fk: ForeignKey) -> str:
+    """Stable edge id from explicit FK column pair (per source table)."""
+    return (
+        f"{fk.source_schema}.{fk.source_table}.{fk.source_column}"
+        f"->{fk.target_schema}.{fk.target_table}.{fk.target_column}"
+    )
+
+
+def lineage_edge_id(connection_name: str, dep: ViewDependency) -> str:
+    """Stable id for a LINEAGE edge (dedupes re-index and multiple evidence paths)."""
+    raw = (
+        f"{connection_name}|{dep.source_schema}.{dep.source_table}"
+        f"->{dep.target_schema}.{dep.target_table}|{dep.lineage_type}"
+    )
+    digest = hashlib.sha256(raw.encode()).hexdigest()[:32]
+    return f"lineage::{digest}"
+
+
+def dbt_lineage_edge_id(
+    connection_name: str, parent_dbt_id: str, child_dbt_id: str
+) -> str:
+    """Stable id for a dbt ``parent_map`` LINEAGE edge (idempotent re-runs)."""
+    raw = f"{connection_name}|dbt|{parent_dbt_id}->{child_dbt_id}|model_dependency"
+    digest = hashlib.sha256(raw.encode()).hexdigest()[:32]
+    return f"lineage::dbt::{digest}"
+
+
+def external_consumer_node_id(
+    connection_name: str,
+    service_name: str,
+    file_path: str,
+    line_start: int,
+    sql_fingerprint: str,
+) -> str:
+    """Stable node id for one code location that issues a specific SQL statement.
+
+    Keyed by ``(connection, service, repo-relative file, line, normalized SQL
+    fingerprint)`` so the same statement in the same place dedupes across
+    re-scans, while the same statement at two call sites — or resolved against
+    two different connections — gets its own node. Line drift between scans
+    produces a fresh id; the stale-consumer sweep reclaims the old one.
+    """
+    raw = f"{connection_name}|{service_name}|{file_path}|{line_start}|{sql_fingerprint}"
+    digest = hashlib.sha256(raw.encode()).hexdigest()[:32]
+    return f"consumer::{digest}"
+
+
+def consumes_edge_id(consumer_node_id: str, table_node_id: str, op: str) -> str:
+    """Stable id for a CONSUMES edge (idempotent by consumer, table, and op).
+
+    ``op`` is part of the key so a statement that both reads and writes the same
+    table yields two distinct edges rather than clobbering one.
+    """
+    raw = f"{consumer_node_id}->{table_node_id}|{op}"
+    digest = hashlib.sha256(raw.encode()).hexdigest()[:32]
+    return f"consumes::{digest}"

pretensor/core/portable_export.py

@@ -0,0 +1,276 @@
+"""Portable JSON export helpers for Kuzu graphs."""
+
+from __future__ import annotations
+
+from datetime import date, datetime, timezone
+from pathlib import Path
+from typing import TypeAlias, TypedDict
+
+from pretensor.core.store import KuzuStore
+
+JsonPrimitive: TypeAlias = str | int | float | bool | None
+JsonValue: TypeAlias = JsonPrimitive | list["JsonValue"] | dict[str, "JsonValue"]
+EdgeRow: TypeAlias = dict[str, JsonValue]
+
+_EDGE_ENDPOINT_COLUMNS = frozenset({"from", "to", "src", "dst"})
+
+
+class NodeTypeExport(TypedDict):
+    """Portable representation of one node label."""
+
+    type: str
+    properties: list[str]
+    rows: list[dict[str, JsonValue]]
+
+
+class EdgeTypeExport(TypedDict):
+    """Portable representation of one edge label."""
+
+    type: str
+    properties: list[str]
+    rows: list[EdgeRow]
+
+
+class PortableGraphExport(TypedDict):
+    """Top-level export payload shape."""
+
+    connection_name: str
+    database: str
+    graph_path: str
+    exported_at: str
+    node_types: list[NodeTypeExport]
+    edge_types: list[EdgeTypeExport]
+    stats: dict[str, int]
+
+
+class _NodeExportResult(TypedDict):
+    payload: NodeTypeExport
+    node_ids: set[str]
+
+
+def export_graph_payload(
+    store: KuzuStore,
+    *,
+    connection_name: str,
+    database_name: str,
+    graph_path: Path,
+) -> PortableGraphExport:
+    """Materialize a portable JSON payload for all node and edge labels.
+
+    Args:
+        store: Open graph store.
+        connection_name: Logical registry connection name being exported.
+        database_name: Database key from the registry entry.
+        graph_path: Path to the Kuzu graph file.
+
+    Returns:
+        Portable export payload containing all label rows and properties.
+    """
+    rows = store.query_all_rows("CALL show_tables() RETURN name, type ORDER BY name")
+    node_types: list[NodeTypeExport] = []
+    edge_types: list[EdgeTypeExport] = []
+    included_node_ids: set[str] = set()
+    node_count = 0
+    edge_count = 0
+    node_table_names: list[str] = []
+    edge_table_names: list[str] = []
+    for table_name_raw, table_type_raw in rows:
+        table_name = str(table_name_raw)
+        table_type = str(table_type_raw).upper()
+        if "REL" in table_type:
+            edge_table_names.append(table_name)
+        else:
+            node_table_names.append(table_name)
+
+    for table_name in node_table_names:
+        columns = _table_columns(store, table_name)
+        node_result = _export_node_type(
+            store,
+            table_name=table_name,
+            columns=columns,
+            connection_name=connection_name,
+            database_name=database_name,
+        )
+        node_types.append(node_result["payload"])
+        included_node_ids.update(node_result["node_ids"])
+        node_count += len(node_result["payload"]["rows"])
+
+    for table_name in edge_table_names:
+        columns = _table_columns(store, table_name)
+        edge_payload = _export_edge_type(
+            store,
+            table_name=table_name,
+            columns=columns,
+            included_node_ids=included_node_ids,
+        )
+        edge_types.append(edge_payload)
+        edge_count += len(edge_payload["rows"])
+    return {
+        "connection_name": connection_name,
+        "database": database_name,
+        "graph_path": str(graph_path),
+        "exported_at": datetime.now(timezone.utc).isoformat(),
+        "node_types": node_types,
+        "edge_types": edge_types,
+        "stats": {
+            "node_types": len(node_types),
+            "edge_types": len(edge_types),
+            "node_count": node_count,
+            "edge_count": edge_count,
+        },
+    }
+
+
+def _table_columns(store: KuzuStore, table_name: str) -> list[str]:
+    quoted = table_name.replace("'", "''")
+    rows = store.query_all_rows(f"CALL table_info('{quoted}') RETURN name, type")
+    return [str(row[0]) for row in rows if row and row[0] is not None]
+
+
+def _export_node_type(
+    store: KuzuStore,
+    *,
+    table_name: str,
+    columns: list[str],
+    connection_name: str,
+    database_name: str,
+) -> _NodeExportResult:
+    projections = ", ".join(
+        f"n.{_quote_identifier(column)} AS {_quote_identifier(column)}"
+        for column in columns
+    )
+    order_clause = (
+        f" ORDER BY n.{_quote_identifier('node_id')}" if "node_id" in columns else ""
+    )
+    scope = _scope_predicate(columns, "n")
+    if scope is None:
+        where_clause = ""
+        query_params: dict[str, str] = {}
+    else:
+        scope_predicate, param_name = scope
+        where_clause = f" WHERE {scope_predicate}"
+        query_params = {
+            param_name: connection_name
+            if param_name == "connection_name"
+            else database_name
+        }
+    cypher = (
+        f"MATCH (n:{_quote_identifier(table_name)})"
+        f"{where_clause} "
+        f"RETURN {projections}{order_clause}"
+    )
+    rows = store.query_all_rows(cypher, query_params)
+    materialized_rows = [_materialize_row(columns, row) for row in rows]
+    node_ids = {
+        str(row_map["node_id"])
+        for row_map in materialized_rows
+        if "node_id" in row_map and row_map["node_id"] is not None
+    }
+    return {
+        "payload": {
+            "type": table_name,
+            "properties": columns,
+            "rows": materialized_rows,
+        },
+        "node_ids": node_ids,
+    }
+
+
+def _export_edge_type(
+    store: KuzuStore,
+    *,
+    table_name: str,
+    columns: list[str],
+    included_node_ids: set[str],
+) -> EdgeTypeExport:
+    property_columns = [
+        column for column in columns if column.lower() not in _EDGE_ENDPOINT_COLUMNS
+    ]
+    projections = [
+        f"src.{_quote_identifier('node_id')} AS {_quote_identifier('__from_node_id')}",
+        f"dst.{_quote_identifier('node_id')} AS {_quote_identifier('__to_node_id')}",
+    ]
+    projections.extend(
+        f"r.{_quote_identifier(column)} AS {_quote_identifier(column)}"
+        for column in property_columns
+    )
+    field_names = ["__from_node_id", "__to_node_id", *property_columns]
+    cypher = (
+        f"MATCH (src)-[r:{_quote_identifier(table_name)}]->(dst) "
+        f"RETURN {', '.join(projections)} "
+        f"ORDER BY {_quote_identifier('__from_node_id')}, {_quote_identifier('__to_node_id')}"
+    )
+    rows = store.query_all_rows(cypher)
+    filtered_rows = [
+        row
+        for row in rows
+        if str(row[0]) in included_node_ids and str(row[1]) in included_node_ids
+    ]
+    return {
+        "type": table_name,
+        "properties": property_columns,
+        "rows": [_materialize_edge_row(field_names, row) for row in filtered_rows],
+    }
+
+
+def _scope_predicate(columns: list[str], alias: str) -> tuple[str, str] | None:
+    """Return (predicate, query param name) for the first matching scope column.
+
+    Only one of ``connection_name``/``database_name`` is ever bound per query,
+    so the predicate text and the params dict built from it must stay in sync
+    (Kuzu prepared statements reject params the query text never references).
+    """
+    if "connection_name" in columns:
+        return (
+            f"{alias}.{_quote_identifier('connection_name')} = $connection_name",
+            "connection_name",
+        )
+    if "database" in columns:
+        return (
+            f"{alias}.{_quote_identifier('database')} = $database_name",
+            "database_name",
+        )
+    if "database_key" in columns:
+        return (
+            f"{alias}.{_quote_identifier('database_key')} = $database_name",
+            "database_name",
+        )
+    return None
+
+
+def _materialize_row(
+    field_names: list[str], row: tuple[object, ...]
+) -> dict[str, JsonValue]:
+    return {
+        field_names[idx]: _to_json_value(value)
+        for idx, value in enumerate(row[: len(field_names)])
+    }
+
+
+def _materialize_edge_row(field_names: list[str], row: tuple[object, ...]) -> EdgeRow:
+    row_map = _materialize_row(field_names, row)
+    return {
+        "__from_node_id": row_map["__from_node_id"],
+        "__to_node_id": row_map["__to_node_id"],
+        **{
+            key: value
+            for key, value in row_map.items()
+            if key not in {"__from_node_id", "__to_node_id"}
+        },
+    }
+
+
+def _to_json_value(value: object) -> JsonValue:
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    if isinstance(value, (datetime, date)):
+        return value.isoformat()
+    if isinstance(value, (list, tuple)):
+        return [_to_json_value(item) for item in value]
+    if isinstance(value, dict):
+        return {str(key): _to_json_value(item) for key, item in value.items()}
+    return str(value)
+
+
+def _quote_identifier(identifier: str) -> str:
+    return "`" + identifier.replace("`", "``") + "`"

pretensor/core/query_runner.py

@@ -0,0 +1,67 @@
+"""Raw Cypher execution wrapper extracted from KuzuStore."""
+
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import kuzu
+
+from pretensor.observability import log_timed_operation
+
+# Use the same logger name as store.py so existing log pipelines see identical
+# source fields regardless of which layer emits the record.
+logger = logging.getLogger("pretensor.core.store")
+
+
+class QueryRunner:
+    """Thin wrapper around a kuzu.Connection that handles all Cypher execution.
+
+    Owns the three execution methods that were previously on KuzuStore:
+    ``execute``, ``execute_write``, and ``query_all_rows``.
+    """
+
+    def __init__(self, conn: kuzu.Connection, path: Path) -> None:
+        self._conn = conn
+        # Stored so execute_write can include graph_path in the log record,
+        # matching the field emitted by the original KuzuStore.execute_write.
+        self._path = path
+
+    def execute_write(self, query: str, params: dict[str, Any] | None = None) -> None:
+        """Execute a write (mutation) Cypher query. Use this for DELETE / SET operations."""
+        with log_timed_operation(
+            logger,
+            event="graph.execute_write",
+            level=logging.DEBUG,
+            has_params=bool(params),
+            query_preview=query.strip().splitlines()[0][:120] if query.strip() else "",
+            graph_path=str(self._path),
+        ):
+            if params:
+                self._conn.execute(query, params)
+            else:
+                self._conn.execute(query)
+
+    def execute(
+        self, cypher: str, parameters: dict[str, Any] | None = None
+    ) -> kuzu.QueryResult | list[kuzu.QueryResult]:
+        """Run a read query (or arbitrary Cypher) with optional parameters."""
+        return self._conn.execute(cypher, parameters)
+
+    def query_all_rows(
+        self, cypher: str, parameters: dict[str, Any] | None = None
+    ) -> list[tuple[Any, ...]]:
+        """Execute Cypher and materialize all result rows as tuples.
+
+        Keeps callers (MCP, search index) off raw ``kuzu.QueryResult`` handles
+        while still routing reads through :class:`QueryRunner`.
+        """
+        result = self.execute(cypher, parameters)
+        if isinstance(result, list):
+            raise TypeError("Expected a single QueryResult from query_all_rows")
+        rows: list[tuple[Any, ...]] = []
+        while result.has_next():
+            nxt = result.get_next()
+            rows.append(tuple(nxt) if not isinstance(nxt, tuple) else nxt)
+        return rows

pretensor/core/registry.py

@@ -0,0 +1,209 @@
+"""JSON registry of indexed databases under ``.pretensor/``."""
+
+from __future__ import annotations
+
+import json
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from pretensor.core.dsn_crypto import DSNEncryptor
+from pretensor.core.secure_io import atomic_write_text
+from pretensor.errors import PretensorError
+
+__all__ = [
+    "GraphRegistry",
+    "MultiDatabaseRegistry",
+    "RegistryEntry",
+    "DatabaseNotFoundError",
+]
+
+_REGISTRY_VERSION = 1
+
+
+class DatabaseNotFoundError(PretensorError, KeyError):
+    """Raised when a connection name is missing from the registry."""
+
+
+class RegistryEntry(BaseModel):
+    """One indexed connection."""
+
+    model_config = ConfigDict(frozen=True)
+
+    connection_name: str
+    database: str
+    dsn: str = ""
+    dsn_encrypted: str | None = None
+    graph_path: str
+    unified_graph_path: str | None = None
+    last_indexed_at: datetime
+    # Persisted by index/add so reindex/detect_changes can parse stored DSNs.
+    dialect: Literal["postgres", "mysql", "snowflake", "bigquery"] = "postgres"
+    table_count: int | None = None
+    description: str | None = None
+    # Capability provenance for ``list_databases``: distinguishes "user didn't
+    # run this" ("not_attempted") from "ran but found nothing" ("empty").
+    dbt_manifest_path: str | None = None
+    llm_enrichment_ran: bool = False
+
+    def plaintext_dsn(self, encryptor: DSNEncryptor | None) -> str:
+        """Return the connection URL, decrypting when ``dsn_encrypted`` is set."""
+        if self.dsn_encrypted:
+            if encryptor is None:
+                raise RuntimeError(
+                    "DSNEncryptor required for encrypted registry entries"
+                )
+            return encryptor.decrypt(self.dsn_encrypted)
+        return self.dsn
+
+
+class _RegistryFile(BaseModel):
+    version: int = _REGISTRY_VERSION
+    entries: dict[str, RegistryEntry] = Field(default_factory=dict)
+
+
+class GraphRegistry:
+    """Load and save ``registry.json`` tracking graph files per connection."""
+
+    def __init__(self, registry_path: Path) -> None:
+        self._path = registry_path
+        self._data = _RegistryFile()
+
+    @property
+    def path(self) -> Path:
+        return self._path
+
+    def load(self) -> GraphRegistry:
+        """Load from disk if the file exists; otherwise start empty."""
+        if not self._path.exists():
+            return self
+        raw = json.loads(self._path.read_text(encoding="utf-8"))
+        if isinstance(raw, dict) and "entries" not in raw and "version" not in raw:
+            raw = {"version": _REGISTRY_VERSION, "entries": raw}
+        self._data = _RegistryFile.model_validate(raw)
+        return self
+
+    def get(self, connection_name: str) -> RegistryEntry | None:
+        return self._data.entries.get(connection_name)
+
+    def require(self, connection_name: str) -> RegistryEntry:
+        """Return an entry or raise :class:`DatabaseNotFoundError`."""
+        entry = self.get(connection_name)
+        if entry is None:
+            raise DatabaseNotFoundError(
+                f"No registry entry for connection {connection_name!r}. "
+                "Use `pretensor list` or `pretensor add`."
+            )
+        return entry
+
+    def list_entries(self) -> list[RegistryEntry]:
+        """Return all entries sorted by connection name."""
+        return sorted(self._data.entries.values(), key=lambda e: e.connection_name)
+
+    def remove(self, connection_name: str) -> None:
+        """Drop a connection from the in-memory registry (call :meth:`save`)."""
+        entries = dict(self._data.entries)
+        entries.pop(connection_name, None)
+        self._data = _RegistryFile(version=self._data.version, entries=entries)
+
+    def upsert(
+        self,
+        *,
+        connection_name: str,
+        database: str,
+        dsn: str,
+        graph_path: Path,
+        indexed_at: datetime | None = None,
+        encrypt_dsn: bool = False,
+        encryptor: DSNEncryptor | None = None,
+        unified_graph_path: Path | None = None,
+        dialect: Literal["postgres", "mysql", "snowflake", "bigquery"] = "postgres",
+        table_count: int | None = None,
+        description: str | None = None,
+        dbt_manifest_path: str | None = None,
+        llm_enrichment_ran: bool = False,
+    ) -> None:
+        """Record or replace an entry for ``connection_name``."""
+        when = indexed_at or datetime.now(timezone.utc)
+        dsn_plain = ""
+        dsn_enc: str | None = None
+        if encrypt_dsn:
+            if encryptor is None:
+                raise ValueError("encryptor is required when encrypt_dsn is True")
+            dsn_enc = encryptor.encrypt(dsn)
+        else:
+            dsn_plain = dsn
+        uni = str(unified_graph_path.resolve()) if unified_graph_path else None
+        entry = RegistryEntry(
+            connection_name=connection_name,
+            database=database,
+            dsn=dsn_plain,
+            dsn_encrypted=dsn_enc,
+            graph_path=str(graph_path.resolve()),
+            unified_graph_path=uni,
+            last_indexed_at=when,
+            dialect=dialect,
+            table_count=table_count,
+            description=description,
+            dbt_manifest_path=dbt_manifest_path,
+            llm_enrichment_ran=llm_enrichment_ran,
+        )
+        entries = dict(self._data.entries)
+        entries[connection_name] = entry
+        self._data = _RegistryFile(version=_REGISTRY_VERSION, entries=entries)
+
+    def update_indexed_metadata(
+        self,
+        connection_name: str,
+        *,
+        table_count: int,
+        indexed_at: datetime | None = None,
+    ) -> None:
+        """Refresh ``table_count`` and optional timestamp after indexing."""
+        cur = self.require(connection_name)
+        when = indexed_at or datetime.now(timezone.utc)
+        entries = dict(self._data.entries)
+        entries[connection_name] = RegistryEntry(
+            connection_name=cur.connection_name,
+            database=cur.database,
+            dsn=cur.dsn,
+            dsn_encrypted=cur.dsn_encrypted,
+            graph_path=cur.graph_path,
+            unified_graph_path=cur.unified_graph_path,
+            last_indexed_at=when,
+            dialect=cur.dialect,
+            table_count=table_count,
+            description=cur.description,
+            dbt_manifest_path=cur.dbt_manifest_path,
+            llm_enrichment_ran=cur.llm_enrichment_ran,
+        )
+        self._data = _RegistryFile(version=self._data.version, entries=entries)
+
+    def unified_graph_file(self) -> Path | None:
+        """Return the shared Kuzu path when all entries agree, else ``None``."""
+        entries = self.list_entries()
+        if not entries:
+            return None
+        paths = {e.unified_graph_path or e.graph_path for e in entries}
+        if len(paths) != 1:
+            return None
+        return Path(next(iter(paths)))
+
+    def save(self) -> None:
+        """Write the registry atomically with owner-only (0o600) permissions.
+
+        The registry can hold a cleartext DSN for entries that predate
+        default-on encryption, so it is written 0o600 to keep credentials off
+        a shared host's other users.
+        """
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        payload: dict[str, Any] = self._data.model_dump(mode="json")
+        atomic_write_text(self._path, json.dumps(payload, indent=2), mode=0o600)
+
+
+class MultiDatabaseRegistry(GraphRegistry):
+    """Alias for the multi-connection registry (Phase 3 unified graph)."""
+
+    pass