pretensor 0.1.0__py3-none-any.whl

pretensor/enrichment/dbt/metadata.py

@@ -0,0 +1,263 @@
+"""Enrich ``SchemaTable`` / ``SchemaColumn`` nodes from dbt ``manifest.json`` metadata."""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+
+from pretensor.core.ids import column_node_id
+from pretensor.core.store import KuzuStore
+from pretensor.enrichment.dbt.manifest import DbtManifest, DbtModel, DbtSource
+from pretensor.enrichment.dbt.resolution import (
+    physical_table_name_for_model,
+    physical_table_name_for_source,
+)
+
+__all__ = ["DbtMetadataWriteStats", "write_dbt_metadata"]
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True, slots=True)
+class DbtMetadataWriteStats:
+    """Counts from ``write_dbt_metadata`` for CLI summaries.
+
+    ``tags_set`` counts tags *actually added* to the graph (delta against any
+    pre-existing tags on the node), not tags offered from the manifest.
+    """
+
+    tables_enriched: int
+    tags_set: int
+
+
+def _empty_stats() -> DbtMetadataWriteStats:
+    return DbtMetadataWriteStats(tables_enriched=0, tags_set=0)
+
+
+def _non_empty_str(value: str | None) -> str | None:
+    if value is None:
+        return None
+    s = value.strip()
+    return s if s else None
+
+
+def _existing_table_tags(
+    store: KuzuStore,
+    connection_name: str,
+    schema_name: str,
+    table_name: str,
+) -> set[str]:
+    rows = store.query_all_rows(
+        """
+        MATCH (t:SchemaTable {connection_name: $cn, schema_name: $s, table_name: $n})
+        RETURN COALESCE(t.tags, CAST([] AS STRING[])) AS tags
+        """,
+        {"cn": connection_name, "s": schema_name, "n": table_name},
+    )
+    if not rows:
+        return set()
+    raw = rows[0][0]
+    if not isinstance(raw, (list, tuple)):
+        return set()
+    return {str(x) for x in raw if str(x).strip()}
+
+
+def _apply_table_metadata(
+    store: KuzuStore,
+    connection_name: str,
+    schema_name: str,
+    table_name: str,
+    *,
+    description: str | None,
+    tags: tuple[str, ...],
+    dbt_ref: str,
+) -> DbtMetadataWriteStats:
+    desc = _non_empty_str(description)
+    tag_list = [t.strip() for t in tags if t and str(t).strip()]
+    if desc is None and not tag_list:
+        return _empty_stats()
+
+    pre_existing = (
+        _existing_table_tags(store, connection_name, schema_name, table_name)
+        if tag_list
+        else set()
+    )
+
+    params: dict[str, object] = {
+        "cn": connection_name,
+        "s": schema_name,
+        "n": table_name,
+    }
+    set_parts: list[str] = []
+    if desc is not None:
+        set_parts.append(
+            "t.description = CASE WHEN t.description IS NULL OR t.description = '' "
+            "THEN $tbl_desc ELSE t.description END"
+        )
+        params["tbl_desc"] = desc
+    if tag_list:
+        set_parts.append(
+            "t.tags = list_distinct(list_concat(COALESCE(t.tags, CAST([] AS STRING[])), $new_tags))"
+        )
+        params["new_tags"] = tag_list
+    if not set_parts:
+        return _empty_stats()
+    result = store.query_all_rows(
+        f"""
+        MATCH (t:SchemaTable {{connection_name: $cn, schema_name: $s, table_name: $n}})
+        SET {", ".join(set_parts)}
+        RETURN t.node_id
+        """,
+        params,
+    )
+    if not result:
+        logger.debug(
+            "dbt metadata: SchemaTable missing for %s (schema=%r table=%r)",
+            dbt_ref,
+            schema_name,
+            table_name,
+        )
+        return _empty_stats()
+
+    new_tags_added = len({t for t in tag_list if t not in pre_existing})
+    return DbtMetadataWriteStats(
+        tables_enriched=1,
+        tags_set=new_tags_added,
+    )
+
+
+def _apply_column_description(
+    store: KuzuStore,
+    connection_name: str,
+    schema_name: str,
+    table_name: str,
+    column_name: str,
+    description: str,
+    *,
+    dbt_ref: str,
+) -> None:
+    desc = _non_empty_str(description)
+    if desc is None:
+        return
+    cid = column_node_id(connection_name, schema_name, table_name, column_name)
+    result = store.query_all_rows(
+        """
+        MATCH (c:SchemaColumn {node_id: $cid})
+        SET c.description = CASE WHEN c.description IS NULL OR c.description = ''
+            THEN $col_desc ELSE c.description END
+        RETURN c.node_id
+        """,
+        {"cid": cid, "col_desc": desc},
+    )
+    if not result:
+        logger.debug(
+            "dbt metadata: SchemaColumn missing for %s column %r (node_id=%s)",
+            dbt_ref,
+            column_name,
+            cid,
+        )
+
+
+def _enrich_model(
+    store: KuzuStore, connection_name: str, model: DbtModel
+) -> DbtMetadataWriteStats:
+    schema = model.schema_name
+    phys = physical_table_name_for_model(model)
+    if not schema or not phys:
+        logger.debug(
+            "dbt metadata: incomplete model metadata for %s (schema=%r table=%r)",
+            model.unique_id,
+            schema,
+            phys,
+        )
+        return _empty_stats()
+    stats = _apply_table_metadata(
+        store,
+        connection_name,
+        schema,
+        phys,
+        description=model.description,
+        tags=model.tags,
+        dbt_ref=model.unique_id,
+    )
+    for col_name, col_desc in model.column_descriptions.items():
+        _apply_column_description(
+            store,
+            connection_name,
+            schema,
+            phys,
+            col_name,
+            col_desc,
+            dbt_ref=model.unique_id,
+        )
+    return stats
+
+
+def _enrich_source(
+    store: KuzuStore, connection_name: str, source: DbtSource
+) -> DbtMetadataWriteStats:
+    schema = source.schema_name
+    phys = physical_table_name_for_source(source)
+    if not schema or not phys:
+        logger.debug(
+            "dbt metadata: incomplete source metadata for %s (schema=%r table=%r)",
+            source.unique_id,
+            schema,
+            phys,
+        )
+        return _empty_stats()
+    stats = _apply_table_metadata(
+        store,
+        connection_name,
+        schema,
+        phys,
+        description=source.description,
+        tags=source.tags,
+        dbt_ref=source.unique_id,
+    )
+    for col_name, col_desc in source.column_descriptions.items():
+        _apply_column_description(
+            store,
+            connection_name,
+            schema,
+            phys,
+            col_name,
+            col_desc,
+            dbt_ref=source.unique_id,
+        )
+    return stats
+
+
+def write_dbt_metadata(
+    manifest: DbtManifest, store: KuzuStore, connection_name: str
+) -> DbtMetadataWriteStats:
+    """Merge dbt model/source descriptions, column descriptions, and tags into the graph.
+
+    Table-level: sets ``SchemaTable.description`` only when the current value is null or
+    empty; merges dbt ``tags`` into ``SchemaTable.tags`` (union, order not preserved).
+
+    Column-level: sets ``SchemaColumn.description`` when null or empty, matching by
+    ``connection_name``, schema, physical table name, and column name.
+
+    Skips nodes that are not present in the store (debug log only). Does not resolve
+    ``docs()`` blocks — only inline string descriptions in the manifest.
+
+    Args:
+        manifest: Parsed dbt manifest.
+        store: Open Kuzu store (``ensure_schema`` recommended).
+        connection_name: Pretensor connection name used when tables were indexed.
+
+    Returns:
+        Counts of tables updated and number of distinct new tags actually added.
+    """
+    total_tables = 0
+    total_tags = 0
+    for model in manifest.nodes.values():
+        st = _enrich_model(store, connection_name, model)
+        total_tables += st.tables_enriched
+        total_tags += st.tags_set
+    for source in manifest.sources.values():
+        st = _enrich_source(store, connection_name, source)
+        total_tables += st.tables_enriched
+        total_tags += st.tags_set
+    return DbtMetadataWriteStats(tables_enriched=total_tables, tags_set=total_tags)

pretensor/enrichment/dbt/pipeline.py

@@ -0,0 +1,77 @@
+"""Orchestrate dbt manifest enrichment passes (lineage, metadata, signals)."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+from pretensor.core.store import KuzuStore
+from pretensor.enrichment.dbt.lineage import write_dbt_lineage
+from pretensor.enrichment.dbt.manifest import DbtManifest
+from pretensor.enrichment.dbt.metadata import write_dbt_metadata
+from pretensor.enrichment.dbt.signals import write_dbt_signals
+
+__all__ = ["DbtEnrichmentSummary", "run_dbt_enrichment", "run_dbt_enrichment_from_manifest"]
+
+
+@dataclass(frozen=True, slots=True)
+class DbtEnrichmentSummary:
+    """Aggregate counts after running all dbt enrichment writers."""
+
+    lineage_edges: int
+    tables_enriched: int
+    tags_set: int
+    exposures_marked: int
+    freshness_rows_applied: int
+    tests_counted: int
+
+
+def run_dbt_enrichment_from_manifest(
+    manifest: DbtManifest,
+    sources_path: Path | None,
+    store: KuzuStore,
+    connection_name: str,
+) -> DbtEnrichmentSummary:
+    """Run lineage → metadata → signals on an already-parsed dbt manifest.
+
+    Exists so CLI commands can preload/validate the manifest *before* mutating
+    the graph and still share the same orchestration as ``run_dbt_enrichment``.
+    """
+    lineage_edges = write_dbt_lineage(manifest, store, connection_name)
+    meta_stats = write_dbt_metadata(manifest, store, connection_name)
+    signal_stats = write_dbt_signals(
+        manifest, store, connection_name, sources_path=sources_path
+    )
+    return DbtEnrichmentSummary(
+        lineage_edges=lineage_edges,
+        tables_enriched=meta_stats.tables_enriched,
+        tags_set=meta_stats.tags_set,
+        exposures_marked=signal_stats.exposures_marked,
+        freshness_rows_applied=signal_stats.freshness_rows_applied,
+        tests_counted=signal_stats.tests_counted,
+    )
+
+
+def run_dbt_enrichment(
+    manifest_path: Path,
+    sources_path: Path | None,
+    store: KuzuStore,
+    connection_name: str,
+) -> DbtEnrichmentSummary:
+    """Load a dbt manifest and run lineage, metadata, then exposure/freshness signals.
+
+    Order matches the dbt enrichment pipeline sequence.
+
+    Args:
+        manifest_path: Path to dbt ``target/manifest.json`` (after ``dbt compile`` / ``dbt run``).
+        sources_path: Optional ``sources.json`` from ``dbt source freshness`` (or ``None``).
+        store: Open Kuzu graph store for the indexed connection.
+        connection_name: Logical Pretensor connection name for this database.
+
+    Returns:
+        Counts for CLI summary lines.
+    """
+    manifest = DbtManifest.load(manifest_path)
+    return run_dbt_enrichment_from_manifest(
+        manifest, sources_path, store, connection_name
+    )

pretensor/enrichment/dbt/resolution.py

@@ -0,0 +1,101 @@
+"""Name and node-id resolution helpers shared across dbt enrichment writers.
+
+These map dbt ``unique_id`` values (``model.<pkg>.<name>`` / ``source.<pkg>.<src>.<name>``)
+to physical ``SchemaTable`` node ids for an indexed Pretensor connection.
+"""
+
+from __future__ import annotations
+
+import logging
+
+from pretensor.core.ids import table_node_id
+from pretensor.enrichment.dbt.manifest import DbtManifest, DbtModel, DbtSource
+
+__all__ = [
+    "physical_table_name_for_model",
+    "physical_table_name_for_source",
+    "resolve_dbt_parent_node_id",
+    "resolve_dbt_child_node_id",
+]
+
+logger = logging.getLogger(__name__)
+
+
+def physical_table_name_for_model(model: DbtModel) -> str | None:
+    """Return the physical table name dbt writes the model to (``alias`` or ``name``)."""
+    return model.alias or model.name
+
+
+def physical_table_name_for_source(source: DbtSource) -> str | None:
+    """Return the physical table name dbt reads the source from (``identifier`` or ``name``)."""
+    return source.identifier or source.name
+
+
+def resolve_dbt_parent_node_id(
+    manifest: DbtManifest,
+    connection_name: str,
+    parent_id: str,
+) -> str | None:
+    """Resolve a dbt ``unique_id`` (model or source) to its ``SchemaTable`` node id."""
+    if parent_id.startswith("source."):
+        src = manifest.sources.get(parent_id)
+        if src is None:
+            logger.debug("dbt resolution: unknown source id %s", parent_id)
+            return None
+        schema = src.schema_name
+        phys = physical_table_name_for_source(src)
+        if not schema or not phys:
+            logger.debug(
+                "dbt resolution: incomplete source metadata for %s (schema=%r table=%r)",
+                parent_id,
+                schema,
+                phys,
+            )
+            return None
+        return table_node_id(connection_name, schema, phys)
+    if parent_id.startswith("model."):
+        mdl = manifest.nodes.get(parent_id)
+        if mdl is None:
+            logger.debug("dbt resolution: unknown model id %s", parent_id)
+            return None
+        schema = mdl.schema_name
+        phys = physical_table_name_for_model(mdl)
+        if not schema or not phys:
+            logger.debug(
+                "dbt resolution: incomplete model metadata for %s (schema=%r table=%r)",
+                parent_id,
+                schema,
+                phys,
+            )
+            return None
+        return table_node_id(connection_name, schema, phys)
+    logger.debug("dbt resolution: skipping non-model non-source id %s", parent_id)
+    return None
+
+
+def resolve_dbt_child_node_id(
+    manifest: DbtManifest,
+    connection_name: str,
+    child_id: str,
+) -> str | None:
+    """Resolve a dbt ``model.*`` ``unique_id`` to its ``SchemaTable`` node id.
+
+    Only models can be lineage children (sources have no upstream in dbt).
+    """
+    if not child_id.startswith("model."):
+        return None
+    mdl = manifest.nodes.get(child_id)
+    if mdl is None:
+        logger.debug("dbt resolution: unknown child model id %s", child_id)
+        return None
+    schema = mdl.schema_name
+    phys = physical_table_name_for_model(mdl)
+    if not schema or not phys:
+        logger.debug(
+            "dbt resolution: incomplete child model metadata for %s (schema=%r table=%r)",
+            child_id,
+            schema,
+            phys,
+        )
+        return None
+    return table_node_id(connection_name, schema, phys)