pretensor 0.1.0__py3-none-any.whl

pretensor/enrichment/analyze/writers.py

@@ -0,0 +1,214 @@
+"""Write ExternalConsumer nodes and CONSUMES edges into the Kuzu graph."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+from pretensor.core.ids import (
+    consumes_edge_id,
+    external_consumer_node_id,
+    table_node_id,
+)
+from pretensor.core.store import KuzuStore
+from pretensor.enrichment.analyze.classify import CONFIDENCE_FLOAT
+from pretensor.enrichment.analyze.extract_python import SqlCandidate
+from pretensor.enrichment.analyze.parse import ParsedSql
+from pretensor.enrichment.analyze.summary import ConsumerRow
+from pretensor.graph_models.consumer import ConsumesEdge, ExternalConsumerNode
+
+__all__ = ["ParsedCandidate", "WriteSummary", "write_consumers_for_run"]
+
+logger = logging.getLogger(__name__)
+
+_CONSUMES_SOURCE = "analyze"
+
+
+@dataclass
+class ParsedCandidate:
+    """A SQL candidate paired with its parsed table refs and per-file provenance."""
+
+    candidate: SqlCandidate
+    parsed: ParsedSql
+    rel_path: str  # repo-relative path (never absolute)
+    language: str
+
+
+@dataclass
+class WriteSummary:
+    """Counts and planned/written rows from ``write_consumers_for_run``."""
+
+    consumers_written: int
+    edges_written: int
+    cross_connection_dropped: int
+    rows: tuple[ConsumerRow, ...] = ()
+
+
+def write_consumers_for_run(
+    store: KuzuStore,
+    connection_name: str,
+    candidates: list[ParsedCandidate],
+    scan_run_id: str,
+    *,
+    service_name: str,
+    min_confidence: float = 0.6,
+    dry_run: bool = False,
+) -> WriteSummary:
+    """Upsert ``ExternalConsumer`` nodes and ``CONSUMES`` edges for one scan run.
+
+    For each candidate:
+    - Splits parsed refs into reads (``table_refs`` minus write targets) and
+      writes (``write_targets``), so each CONSUMES edge carries a truthful ``op``.
+    - Resolves each ``(schema, table)`` ref to an existing ``SchemaTable`` node id;
+      unresolved (cross-connection or unknown) refs are dropped with a counter.
+    - Upserts one ``ExternalConsumer`` node (no raw SQL stored, only the fingerprint).
+    - Upserts one ``CONSUMES`` edge per resolved ``(table, op)``.
+    - Marks ``has_external_consumers = true`` on every touched ``SchemaTable``.
+
+    After all candidates, sweeps ``ExternalConsumer`` rows (and stale
+    ``CONSUMES`` edges) for this ``service_name``/``connection_name`` carrying a
+    different ``scan_run_id``, and clears ``has_external_consumers`` on tables
+    that lost their last consumer.
+    """
+    known_tables = _known_table_ids(store, connection_name)
+    folded_known = _casefold_index(known_tables)
+    written_consumers: set[str] = set()
+    written_edges: set[str] = set()
+    cross_connection_dropped = 0
+    touched_tables: set[str] = set()
+    rows: list[ConsumerRow] = []
+
+    for pc in candidates:
+        cand = pc.candidate
+        parsed = pc.parsed
+
+        confidence = CONFIDENCE_FLOAT.get(cand.confidence_bucket, 0.0)
+        if confidence < min_confidence:
+            continue
+
+        write_set = _dedup_refs(parsed.write_targets)
+        read_set = [r for r in _dedup_refs(parsed.table_refs) if r not in write_set]
+
+        resolved: list[tuple[str, str, str]] = []  # (table_node_id, op, "schema.table")
+        for refs, op in ((read_set, "read"), (write_set, "write")):
+            for schema, table in refs:
+                nid = table_node_id(connection_name, schema, table)
+                if nid not in known_tables:
+                    nid = folded_known.get(nid.casefold(), "")
+                if not nid:
+                    logger.debug(
+                        "analyze writers: unresolved ref %s.%s for connection %s (dropped)",
+                        schema,
+                        table,
+                        connection_name,
+                    )
+                    cross_connection_dropped += 1
+                    continue
+                _, canon_schema, canon_table = nid.split("::", 2)
+                resolved.append((nid, op, f"{canon_schema}.{canon_table}"))
+
+        if not resolved:
+            continue
+
+        consumer_nid = external_consumer_node_id(
+            connection_name,
+            service_name,
+            pc.rel_path,
+            cand.line_start,
+            parsed.fingerprint,
+        )
+        if not dry_run:
+            store.upsert_external_consumer(
+                ExternalConsumerNode(
+                    node_id=consumer_nid,
+                    connection_name=connection_name,
+                    service_name=service_name,
+                    file_path=pc.rel_path,
+                    language=pc.language,
+                    symbol=cand.symbol,
+                    kind=cand.kind,
+                    line_start=cand.line_start,
+                    line_end=cand.line_end,
+                    sql_fingerprint=parsed.fingerprint,
+                    confidence=confidence,
+                    dialect_used=parsed.dialect_used,
+                    scan_run_id=scan_run_id,
+                )
+            )
+        written_consumers.add(consumer_nid)
+
+        for table_nid, op, table_label in resolved:
+            eid = consumes_edge_id(consumer_nid, table_nid, op)
+            if eid in written_edges:
+                continue
+            if not dry_run:
+                store.upsert_consumes_edge(
+                    ConsumesEdge(
+                        edge_id=eid,
+                        source_node_id=consumer_nid,
+                        target_node_id=table_nid,
+                        op=op,
+                        source=_CONSUMES_SOURCE,
+                        confidence=confidence,
+                        scan_run_id=scan_run_id,
+                    )
+                )
+            touched_tables.add(table_nid)
+            written_edges.add(eid)
+            rows.append(
+                ConsumerRow(
+                    service_name=service_name,
+                    table=table_label,
+                    op=op,
+                    file_path=pc.rel_path,
+                    line_start=cand.line_start,
+                    line_end=cand.line_end,
+                    confidence=confidence,
+                )
+            )
+
+    if not dry_run:
+        store.mark_tables_have_external_consumers(sorted(touched_tables))
+        store.sweep_stale_consumers(service_name, connection_name, scan_run_id)
+
+    return WriteSummary(
+        consumers_written=len(written_consumers),
+        edges_written=len(written_edges),
+        cross_connection_dropped=cross_connection_dropped,
+        rows=tuple(rows),
+    )
+
+
+def _known_table_ids(store: KuzuStore, connection_name: str) -> set[str]:
+    """All ``SchemaTable`` node ids indexed under ``connection_name``, in one query."""
+    rows = store.query_all_rows(
+        "MATCH (t:SchemaTable {connection_name: $cn}) RETURN t.node_id",
+        {"cn": connection_name},
+    )
+    return {str(r[0]) for r in rows}
+
+
+def _casefold_index(known: set[str]) -> dict[str, str]:
+    """Casefolded node id → node id, for unique case-insensitive fallback.
+
+    Snowflake stores unquoted identifiers uppercase while application SQL is
+    conventionally lowercase, so an exact-match-only lookup drops every ref.
+    A ref that misses exactly resolves case-insensitively when precisely one
+    stored table matches; ambiguous folds (case-variant twins) stay dropped.
+    """
+    index: dict[str, str | None] = {}
+    for nid in known:
+        key = nid.casefold()
+        index[key] = None if key in index else nid
+    return {k: v for k, v in index.items() if v is not None}
+
+
+def _dedup_refs(refs: list[tuple[str, str]]) -> list[tuple[str, str]]:
+    """Deduplicate (schema, table) pairs while preserving order."""
+    seen: set[tuple[str, str]] = set()
+    out: list[tuple[str, str]] = []
+    for ref in refs:
+        if ref not in seen:
+            seen.add(ref)
+            out.append(ref)
+    return out

pretensor/enrichment/dbt/__init__.py

@@ -0,0 +1,30 @@
+"""dbt manifest parsing for graph enrichment."""
+
+from pretensor.enrichment.dbt.lineage import write_dbt_lineage
+from pretensor.enrichment.dbt.manifest import (
+    DbtExposure,
+    DbtManifest,
+    DbtManifestError,
+    DbtModel,
+    DbtSource,
+    DbtTest,
+)
+from pretensor.enrichment.dbt.metadata import DbtMetadataWriteStats, write_dbt_metadata
+from pretensor.enrichment.dbt.pipeline import DbtEnrichmentSummary, run_dbt_enrichment
+from pretensor.enrichment.dbt.signals import DbtSignalsWriteStats, write_dbt_signals
+
+__all__ = [
+    "DbtExposure",
+    "DbtManifest",
+    "DbtManifestError",
+    "DbtModel",
+    "DbtSource",
+    "DbtEnrichmentSummary",
+    "DbtMetadataWriteStats",
+    "DbtSignalsWriteStats",
+    "DbtTest",
+    "run_dbt_enrichment",
+    "write_dbt_lineage",
+    "write_dbt_metadata",
+    "write_dbt_signals",
+]

pretensor/enrichment/dbt/lineage.py

@@ -0,0 +1,100 @@
+"""Write table-level ``LINEAGE`` edges from dbt ``manifest.json`` ``parent_map``."""
+
+from __future__ import annotations
+
+import logging
+
+from pretensor.core.ids import dbt_lineage_edge_id
+from pretensor.core.store import KuzuStore
+from pretensor.enrichment.dbt.manifest import DbtManifest
+from pretensor.enrichment.dbt.resolution import (
+    resolve_dbt_child_node_id,
+    resolve_dbt_parent_node_id,
+)
+from pretensor.graph_models.edge import LineageEdge
+
+__all__ = ["write_dbt_lineage"]
+
+logger = logging.getLogger(__name__)
+
+_DBT_LINEAGE_SOURCE = "dbt"
+_DBT_LINEAGE_TYPE = "model_dependency"
+_DBT_LINEAGE_CONFIDENCE = 1.0
+
+
+def _schema_table_exists(store: KuzuStore, node_id: str) -> bool:
+    rows = store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {node_id: $nid})
+        RETURN 1
+        LIMIT 1
+        """,
+        {"nid": node_id},
+    )
+    return len(rows) > 0
+
+
+def write_dbt_lineage(
+    manifest: DbtManifest, store: KuzuStore, connection_name: str
+) -> int:
+    """Emit ``LINEAGE`` edges from dbt ``parent_map`` for one indexed connection.
+
+    For each ``model.*`` entry, creates one edge per direct parent that is a
+    ``source.*`` or ``model.*`` node, when both endpoints exist as
+    ``SchemaTable`` rows (matched by ``connection_name``, schema, and physical
+    table name). Upstream is the parent; downstream is the child model.
+
+    Idempotent per connection: existing dbt ``model_dependency`` edges for this
+    connection are removed first, then rewritten. Note that this scopes the
+    clear to edges where **both** endpoints are on ``connection_name`` — in
+    ``--unified`` mode, a dbt project that crosses two Pretensor connections
+    would need per-connection re-runs to fully refresh. Cross-connection dbt
+    lineage is not supported in Phase 1 (out-of-scope).
+
+    Args:
+        manifest: Parsed dbt manifest (includes ``parent_map``).
+        store: Open Kuzu store with schema ensured.
+        connection_name: Pretensor connection name for ``SchemaTable`` nodes.
+
+    Returns:
+        Number of ``LINEAGE`` edges written.
+    """
+    store.clear_dbt_model_dependency_lineage(connection_name)
+    written = 0
+    for child_id, parents in manifest.parent_map.items():
+        child_nid = resolve_dbt_child_node_id(manifest, connection_name, child_id)
+        if child_nid is None:
+            continue
+        if not _schema_table_exists(store, child_nid):
+            logger.debug(
+                "dbt lineage: child SchemaTable missing for %s -> %s",
+                child_id,
+                child_nid,
+            )
+            continue
+        for parent_id in parents:
+            parent_nid = resolve_dbt_parent_node_id(
+                manifest, connection_name, parent_id
+            )
+            if parent_nid is None:
+                continue
+            if not _schema_table_exists(store, parent_nid):
+                logger.debug(
+                    "dbt lineage: parent SchemaTable missing for %s -> %s",
+                    parent_id,
+                    parent_nid,
+                )
+                continue
+            eid = dbt_lineage_edge_id(connection_name, parent_id, child_id)
+            store.upsert_lineage_edge(
+                LineageEdge(
+                    edge_id=eid,
+                    source_node_id=parent_nid,
+                    target_node_id=child_nid,
+                    source=_DBT_LINEAGE_SOURCE,
+                    lineage_type=_DBT_LINEAGE_TYPE,
+                    confidence=_DBT_LINEAGE_CONFIDENCE,
+                )
+            )
+            written += 1
+    return written

pretensor/enrichment/dbt/manifest.py

@@ -0,0 +1,300 @@
+"""Load and parse dbt ``manifest.json`` artifacts for graph enrichment."""
+
+from __future__ import annotations
+
+import json
+import logging
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Mapping
+
+from pretensor.errors import PretensorError
+
+logger = logging.getLogger(__name__)
+
+_SUPPORTED_SCHEMA_MARKERS = ("manifest/v9", "manifest/v10", "/v9.json", "/v10.json")
+
+
+class DbtManifestError(PretensorError):
+    """Raised when a dbt manifest file is missing or cannot be parsed."""
+
+
+@dataclass(frozen=True, slots=True)
+class DbtModel:
+    """A dbt model node from ``manifest.json`` ``nodes``."""
+
+    unique_id: str
+    name: str
+    description: str | None = None
+    tags: tuple[str, ...] = ()
+    package_name: str | None = None
+    database: str | None = None
+    schema_name: str | None = None
+    alias: str | None = None
+    column_descriptions: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class DbtSource:
+    """A dbt source from ``manifest.json`` ``sources``."""
+
+    unique_id: str
+    name: str
+    source_name: str | None = None
+    description: str | None = None
+    tags: tuple[str, ...] = ()
+    package_name: str | None = None
+    database: str | None = None
+    schema_name: str | None = None
+    identifier: str | None = None
+    column_descriptions: dict[str, str] = field(default_factory=dict)
+
+
+@dataclass(frozen=True, slots=True)
+class DbtExposure:
+    """A dbt exposure from ``manifest.json`` ``exposures``."""
+
+    unique_id: str
+    name: str
+    description: str | None = None
+    tags: tuple[str, ...] = ()
+    package_name: str | None = None
+    depends_on_nodes: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class DbtTest:
+    """A dbt test node (typically under ``nodes`` with ``resource_type`` test)."""
+
+    unique_id: str
+    name: str | None = None
+    package_name: str | None = None
+    attached_node: str | None = None
+
+
+def _is_supported_schema_version(version: str | None) -> bool:
+    if version is None:
+        return True
+    lowered = version.lower()
+    return any(marker in lowered for marker in _SUPPORTED_SCHEMA_MARKERS)
+
+
+def _warn_if_unsupported_schema(version: str | None) -> None:
+    if version is not None and not _is_supported_schema_version(version):
+        logger.warning(
+            "dbt manifest schema version may be unsupported: %s "
+            "(expected v9 or v10 URL in metadata.dbt_schema_version)",
+            version,
+        )
+
+
+def _as_str(value: Any) -> str | None:
+    if value is None:
+        return None
+    if isinstance(value, str):
+        return value
+    return str(value)
+
+
+def _as_str_tuple(value: Any) -> tuple[str, ...]:
+    if value is None:
+        return ()
+    if isinstance(value, str):
+        return (value,)
+    if isinstance(value, (list, tuple)):
+        out: list[str] = []
+        for item in value:
+            s = _as_str(item)
+            if s is not None:
+                out.append(s)
+        return tuple(out)
+    return ()
+
+
+def _depends_on_nodes(raw: Mapping[str, Any]) -> tuple[str, ...]:
+    depends_on = raw.get("depends_on")
+    if not isinstance(depends_on, dict):
+        return ()
+    nodes = depends_on.get("nodes")
+    if not isinstance(nodes, list):
+        return ()
+    return tuple(n for n in (_as_str(x) for x in nodes) if n is not None)
+
+
+def _column_descriptions_from_raw(raw: Mapping[str, Any]) -> dict[str, str]:
+    """Map column name -> non-empty inline ``description`` from manifest ``columns``."""
+    cols = raw.get("columns")
+    if not isinstance(cols, dict):
+        return {}
+    out: dict[str, str] = {}
+    for key, col_raw in cols.items():
+        col_name = _as_str(key)
+        if col_name is None or not isinstance(col_raw, dict):
+            continue
+        desc = _as_str(col_raw.get("description"))
+        if desc is None:
+            continue
+        stripped = desc.strip()
+        if stripped:
+            out[col_name] = stripped
+    return out
+
+
+def _parse_model(unique_id: str, raw: Mapping[str, Any]) -> DbtModel:
+    return DbtModel(
+        unique_id=unique_id,
+        name=_as_str(raw.get("name")) or unique_id,
+        description=_as_str(raw.get("description")),
+        tags=_as_str_tuple(raw.get("tags")),
+        package_name=_as_str(raw.get("package_name")),
+        database=_as_str(raw.get("database")),
+        schema_name=_as_str(raw.get("schema")),
+        alias=_as_str(raw.get("alias")),
+        column_descriptions=_column_descriptions_from_raw(raw),
+    )
+
+
+def _parse_source(unique_id: str, raw: Mapping[str, Any]) -> DbtSource:
+    return DbtSource(
+        unique_id=unique_id,
+        name=_as_str(raw.get("name")) or unique_id,
+        source_name=_as_str(raw.get("source_name")),
+        description=_as_str(raw.get("description")),
+        tags=_as_str_tuple(raw.get("tags")),
+        package_name=_as_str(raw.get("package_name")),
+        database=_as_str(raw.get("database")),
+        schema_name=_as_str(raw.get("schema")),
+        identifier=_as_str(raw.get("identifier")),
+        column_descriptions=_column_descriptions_from_raw(raw),
+    )
+
+
+def _parse_exposure(unique_id: str, raw: Mapping[str, Any]) -> DbtExposure:
+    return DbtExposure(
+        unique_id=unique_id,
+        name=_as_str(raw.get("name")) or unique_id,
+        description=_as_str(raw.get("description")),
+        tags=_as_str_tuple(raw.get("tags")),
+        package_name=_as_str(raw.get("package_name")),
+        depends_on_nodes=_depends_on_nodes(raw),
+    )
+
+
+def _parse_test(unique_id: str, raw: Mapping[str, Any]) -> DbtTest:
+    return DbtTest(
+        unique_id=unique_id,
+        name=_as_str(raw.get("name")),
+        package_name=_as_str(raw.get("package_name")),
+        attached_node=_as_str(raw.get("attached_node")),
+    )
+
+
+def _parse_parent_map(raw: Any) -> dict[str, list[str]]:
+    if not isinstance(raw, dict):
+        return {}
+    result: dict[str, list[str]] = {}
+    for key, value in raw.items():
+        kid = _as_str(key)
+        if kid is None:
+            continue
+        if not isinstance(value, list):
+            result[kid] = []
+            continue
+        parents: list[str] = []
+        for item in value:
+            sid = _as_str(item)
+            if sid is not None:
+                parents.append(sid)
+        result[kid] = parents
+    return result
+
+
+@dataclass(frozen=True, slots=True)
+class DbtManifest:
+    """Parsed dbt ``manifest.json`` (v9/v10-oriented; unknown fields ignored)."""
+
+    schema_version: str | None
+    nodes: dict[str, DbtModel]
+    sources: dict[str, DbtSource]
+    exposures: dict[str, DbtExposure]
+    tests: dict[str, DbtTest]
+    parent_map: dict[str, list[str]]
+
+    @classmethod
+    def load(cls, path: Path) -> DbtManifest:
+        """Load a manifest from ``path``.
+
+        Args:
+            path: Filesystem path to ``manifest.json``.
+
+        Returns:
+            Parsed manifest.
+
+        Raises:
+            DbtManifestError: If the file is missing or JSON is invalid.
+        """
+        if not path.is_file():
+            msg = f"dbt manifest not found or not a file: {path}"
+            raise DbtManifestError(msg)
+        try:
+            text = path.read_text(encoding="utf-8")
+        except OSError as exc:
+            msg = f"cannot read dbt manifest: {path}"
+            raise DbtManifestError(msg) from exc
+        try:
+            payload = json.loads(text)
+        except json.JSONDecodeError as exc:
+            msg = f"dbt manifest is not valid JSON: {path}"
+            raise DbtManifestError(msg) from exc
+        if not isinstance(payload, dict):
+            msg = f"dbt manifest root must be a JSON object: {path}"
+            raise DbtManifestError(msg)
+
+        metadata = payload.get("metadata")
+        schema_version: str | None = None
+        if isinstance(metadata, dict):
+            schema_version = _as_str(metadata.get("dbt_schema_version"))
+        _warn_if_unsupported_schema(schema_version)
+
+        nodes_out: dict[str, DbtModel] = {}
+        tests_out: dict[str, DbtTest] = {}
+        raw_nodes = payload.get("nodes")
+        if isinstance(raw_nodes, dict):
+            for node_id, node_raw in raw_nodes.items():
+                nid = _as_str(node_id)
+                if nid is None or not isinstance(node_raw, dict):
+                    continue
+                rtype = _as_str(node_raw.get("resource_type"))
+                if rtype == "model":
+                    nodes_out[nid] = _parse_model(nid, node_raw)
+                elif rtype == "test":
+                    tests_out[nid] = _parse_test(nid, node_raw)
+
+        sources_out: dict[str, DbtSource] = {}
+        raw_sources = payload.get("sources")
+        if isinstance(raw_sources, dict):
+            for source_id, source_raw in raw_sources.items():
+                sid = _as_str(source_id)
+                if sid is None or not isinstance(source_raw, dict):
+                    continue
+                sources_out[sid] = _parse_source(sid, source_raw)
+
+        exposures_out: dict[str, DbtExposure] = {}
+        raw_exposures = payload.get("exposures")
+        if isinstance(raw_exposures, dict):
+            for exposure_id, exposure_raw in raw_exposures.items():
+                eid = _as_str(exposure_id)
+                if eid is None or not isinstance(exposure_raw, dict):
+                    continue
+                exposures_out[eid] = _parse_exposure(eid, exposure_raw)
+
+        parent_map = _parse_parent_map(payload.get("parent_map"))
+
+        return cls(
+            schema_version=schema_version,
+            nodes=nodes_out,
+            sources=sources_out,
+            exposures=exposures_out,
+            tests=tests_out,
+            parent_map=parent_map,
+        )