datalex-cli 0.1.1__py3-none-any.whl

datalex_core/dbt/emit.py

@@ -0,0 +1,344 @@
+"""dbt YAML emitter.
+
+Given a loaded DataLexProject, emits:
+  * sources/<source_name>.yml         — one file per `kind: source`
+  * models/_schema.yml                — schema.yml for every `kind: model`
+
+Output is dbt v2 format and includes:
+  * contracts (`config.contract.enforced: true` with `data_type` per column)
+  * column-level constraints (primary_key / unique / not_null / foreign_key / check)
+  * tests (unique / not_null / accepted_values / relationships / custom)
+  * freshness (at source level and per-table)
+  * meta round-trip via `meta.datalex.*` so reimports never clobber user intent
+
+The dict payloads returned by `build_sources_yaml` / `build_models_yaml` are plain,
+serialization-ready dicts — callers choose how to write them (single file, per-file,
+etc.) via `write_dbt_yaml`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import yaml
+
+from datalex_core.datalex.project import DataLexProject
+
+
+# ------------------------ build payloads ------------------------
+
+
+def build_sources_yaml(project: DataLexProject) -> Dict[str, Dict[str, Any]]:
+    """Return { relative_path: source_doc } for every `kind: source` file.
+
+    We split by source name so dbt's `source(name, table)` reference stays stable
+    and edit-friendly: sources/<name>.yml.
+    """
+    out: Dict[str, Dict[str, Any]] = {}
+    for src in project.sources.values():
+        doc = _source_to_dict(src)
+        out[f"sources/{src['name']}.yml"] = {"version": 2, "sources": [doc]}
+    return out
+
+
+def build_models_yaml(project: DataLexProject) -> Dict[str, Dict[str, Any]]:
+    """Return { relative_path: models_doc } for every `kind: model` file.
+
+    We collect all models into a single `models/_schema.yml` so dbt can `dbt parse`
+    them in one read. Power users can split later; keeping everything in one file
+    is the dbt community's default and prevents discovery surprises.
+    """
+    models = [_model_to_dict(m) for m in project.models.values()]
+    if not models:
+        return {}
+    return {"models/_schema.yml": {"version": 2, "models": models}}
+
+
+# ------------------------ writing ------------------------
+
+
+@dataclass
+class EmitReport:
+    files: List[str] = field(default_factory=list)
+    sources: int = 0
+    models: int = 0
+
+    def summary(self) -> str:
+        lines = ["dbt emission complete:"]
+        lines.append(f"  source files: {self.sources}")
+        lines.append(f"  model files:  {self.models}")
+        for f in self.files:
+            lines.append(f"    - {f}")
+        return "\n".join(lines)
+
+
+def emit_dbt(
+    project: DataLexProject,
+    out_dir: str,
+    include_sources: bool = True,
+    include_models: bool = True,
+) -> EmitReport:
+    """Render a DataLexProject into a dbt-parseable YAML tree under out_dir."""
+    report = EmitReport()
+    out = Path(out_dir)
+    out.mkdir(parents=True, exist_ok=True)
+
+    if include_sources:
+        for rel, payload in build_sources_yaml(project).items():
+            target = out / rel
+            target.parent.mkdir(parents=True, exist_ok=True)
+            _write_yaml(target, payload)
+            report.files.append(str(target))
+            report.sources += 1
+
+    if include_models:
+        for rel, payload in build_models_yaml(project).items():
+            target = out / rel
+            target.parent.mkdir(parents=True, exist_ok=True)
+            _write_yaml(target, payload)
+            report.files.append(str(target))
+            report.models += 1
+
+    return report
+
+
+def _write_yaml(path: Path, doc: Dict[str, Any]) -> None:
+    with path.open("w", encoding="utf-8") as f:
+        yaml.safe_dump(
+            doc,
+            f,
+            sort_keys=False,
+            default_flow_style=False,
+            allow_unicode=True,
+            width=120,
+        )
+
+
+# ------------------------ translators ------------------------
+
+
+def _source_to_dict(src: Dict[str, Any]) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": src["name"]}
+    _copy_if_set(doc, src, ("description", "database", "schema", "loader", "loaded_at_field"))
+    if src.get("freshness"):
+        doc["freshness"] = _translate_freshness(src["freshness"])
+
+    tables_out: List[Dict[str, Any]] = []
+    for tbl in src.get("tables", []) or []:
+        tables_out.append(_source_table_to_dict(tbl))
+    doc["tables"] = tables_out
+
+    meta = _build_meta(src, extra={"kind": "source"})
+    if meta:
+        doc["meta"] = meta
+
+    return doc
+
+
+def _source_table_to_dict(tbl: Dict[str, Any]) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": tbl["name"]}
+    _copy_if_set(doc, tbl, ("description", "identifier", "loaded_at_field"))
+    if tbl.get("freshness"):
+        doc["freshness"] = _translate_freshness(tbl["freshness"])
+    cols_out: List[Dict[str, Any]] = []
+    for c in tbl.get("columns", []) or []:
+        cols_out.append(_source_column_to_dict(c))
+    if cols_out:
+        doc["columns"] = cols_out
+    meta = _build_meta(tbl)
+    if meta:
+        doc["meta"] = meta
+    return doc
+
+
+def _source_column_to_dict(col: Dict[str, Any]) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": col["name"]}
+    _copy_if_set(doc, col, ("description",))
+    # sources pass `type` through dbt-side as data_type so contract works downstream
+    if col.get("type"):
+        doc["data_type"] = col["type"]
+    if col.get("tests"):
+        doc["tests"] = list(col["tests"])
+    meta = _build_meta(col, extra=_sensitivity_meta(col))
+    if meta:
+        doc["meta"] = meta
+    return doc
+
+
+def _model_to_dict(m: Dict[str, Any]) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": m["name"]}
+    _copy_if_set(doc, m, ("description",))
+
+    config = _model_config(m)
+    if config:
+        doc["config"] = config
+
+    cols_out: List[Dict[str, Any]] = []
+    contract_enforced = bool((m.get("contract") or {}).get("enforced"))
+    for c in m.get("columns", []) or []:
+        cols_out.append(_model_column_to_dict(c, contract_enforced=contract_enforced))
+    if cols_out:
+        doc["columns"] = cols_out
+
+    meta = _build_meta(
+        m,
+        extra={
+            "kind": "model",
+            **(
+                {"depends_on": [_ref_to_string(r) for r in m["depends_on"]]}
+                if m.get("depends_on")
+                else {}
+            ),
+            **({"derived_sql": m["derived_sql"]} if m.get("derived_sql") else {}),
+            **({"sql_path": m["sql_path"]} if m.get("sql_path") else {}),
+            **({"owner": m["owner"]} if m.get("owner") else {}),
+            **({"domain": m["domain"]} if m.get("domain") else {}),
+        },
+    )
+    if meta:
+        doc["meta"] = meta
+
+    if m.get("tags"):
+        doc.setdefault("config", {})["tags"] = list(m["tags"])
+
+    return doc
+
+
+def _model_config(m: Dict[str, Any]) -> Dict[str, Any]:
+    cfg: Dict[str, Any] = {}
+    if m.get("materialization"):
+        cfg["materialized"] = m["materialization"]
+    if m.get("database"):
+        cfg["database"] = m["database"]
+    if m.get("schema"):
+        cfg["schema"] = m["schema"]
+    contract = m.get("contract") or {}
+    if contract.get("enforced") is not None:
+        cfg["contract"] = {"enforced": bool(contract["enforced"])}
+    return cfg
+
+
+def _model_column_to_dict(col: Dict[str, Any], contract_enforced: bool) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": col["name"]}
+    _copy_if_set(doc, col, ("description",))
+
+    # Contract enforcement requires data_type; always emit it if present.
+    if col.get("type"):
+        doc["data_type"] = col["type"]
+    elif contract_enforced:
+        # dbt parse will fail without data_type on contract-enforced models.
+        # Surface this as a YAML-visible TODO rather than silently dropping it.
+        doc["data_type"] = "UNSPECIFIED"
+
+    constraints = _translate_constraints(col)
+    if constraints:
+        doc["constraints"] = constraints
+
+    if col.get("tests"):
+        doc["tests"] = list(col["tests"])
+
+    meta = _build_meta(col, extra=_sensitivity_meta(col))
+    if meta:
+        doc["meta"] = meta
+    return doc
+
+
+def _translate_constraints(col: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """Convert DataLex column constraint rules to dbt constraint entries.
+
+    Pulls from both shorthand fields (primary_key / unique / nullable) and the
+    explicit `constraints:` array. Deduplicates by (type, expression) so the
+    same intent declared both ways doesn't produce duplicate entries.
+    """
+    out: List[Dict[str, Any]] = []
+    seen: set = set()
+
+    def _add(entry: Dict[str, Any]) -> None:
+        key = (entry.get("type"), entry.get("expression"))
+        if key not in seen:
+            seen.add(key)
+            out.append(entry)
+
+    if col.get("primary_key"):
+        _add({"type": "primary_key"})
+    if col.get("unique"):
+        _add({"type": "unique"})
+    if col.get("nullable") is False and not col.get("primary_key"):
+        _add({"type": "not_null"})
+
+    ref = col.get("references")
+    if ref and ref.get("entity") and ref.get("column"):
+        _add({"type": "foreign_key", "expression": f"{ref['entity']}({ref['column']})"})
+
+    for c in col.get("constraints", []) or []:
+        ctype = c.get("type")
+        if ctype in ("primary_key", "unique", "not_null"):
+            _add({"type": ctype})
+        elif ctype == "check" and c.get("expression"):
+            _add({"type": "check", "expression": c["expression"]})
+        elif ctype == "foreign_key" and c.get("expression"):
+            _add({"type": "foreign_key", "expression": c["expression"]})
+    return out
+
+
+def _translate_freshness(f: Dict[str, Any]) -> Dict[str, Any]:
+    out: Dict[str, Any] = {}
+    for k in ("warn_after", "error_after"):
+        if f.get(k):
+            out[k] = {"count": f[k]["count"], "period": f[k]["period"]}
+    if f.get("filter"):
+        out["filter"] = f["filter"]
+    return out
+
+
+def _build_meta(
+    obj: Dict[str, Any],
+    extra: Optional[Dict[str, Any]] = None,
+) -> Dict[str, Any]:
+    """Merge the object's declared `meta` with governance round-trip keys under
+    `meta.datalex.*`. User-declared keys win — we never overwrite."""
+    out: Dict[str, Any] = {}
+    # start with any user-declared meta
+    user_meta = obj.get("meta") or {}
+    if isinstance(user_meta, dict):
+        out.update({k: v for k, v in user_meta.items() if k != "datalex"})
+
+    datalex_meta: Dict[str, Any] = {}
+    # Preserve existing meta.datalex if present (idempotent re-emits).
+    if isinstance(user_meta.get("datalex"), dict):
+        datalex_meta.update(user_meta["datalex"])
+    if extra:
+        for k, v in extra.items():
+            datalex_meta.setdefault(k, v)
+    if datalex_meta:
+        out["datalex"] = datalex_meta
+    return out
+
+
+def _sensitivity_meta(obj: Dict[str, Any]) -> Dict[str, Any]:
+    extra: Dict[str, Any] = {}
+    if obj.get("sensitivity"):
+        extra["sensitivity"] = obj["sensitivity"]
+    if obj.get("tags"):
+        extra["tags"] = list(obj["tags"])
+    if obj.get("terms"):
+        extra["terms"] = list(obj["terms"])
+    return extra
+
+
+def _copy_if_set(dst: Dict[str, Any], src: Dict[str, Any], keys: Tuple[str, ...]) -> None:
+    for k in keys:
+        v = src.get(k)
+        if v is not None and v != "":
+            dst[k] = v
+
+
+def _ref_to_string(dep: Dict[str, Any]) -> str:
+    if "ref" in dep:
+        return f"ref:{dep['ref']}"
+    if "source" in dep:
+        s = dep["source"]
+        return f"source:{s.get('source')}.{s.get('name')}"
+    return str(dep)

datalex_core/dbt/manifest.py

@@ -0,0 +1,329 @@
+"""dbt manifest.json -> DataLex source/model importer with idempotent round-trip.
+
+Design:
+  * Stable key = `unique_id` from the manifest (e.g. `source.my_project.raw.orders`,
+    `model.my_project.stg_orders`). Stored under `meta.datalex.dbt.unique_id`.
+  * On re-import, existing DataLex files are *merged*, not overwritten. User-authored
+    fields (description, tests, sensitivity, owner, etc.) are preserved; only fields
+    the manifest owns (database/schema/columns' data_type) get refreshed.
+  * The importer emits ready-to-write dicts; callers choose where to write them
+    (typically under `sources/` and `models/`).
+
+What we do NOT do here: write files. A thin wrapper does that — `write_import_result`
+in this module — but users can choose to merge into an existing project tree manually
+via their own logic.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any, Dict, List, Optional
+
+import yaml
+
+
+# ------------------------ public API ------------------------
+
+
+@dataclass
+class ImportResult:
+    sources: Dict[str, Dict[str, Any]] = field(default_factory=dict)  # name -> doc
+    models: Dict[str, Dict[str, Any]] = field(default_factory=dict)
+    warnings: List[str] = field(default_factory=list)
+
+
+def import_manifest(
+    manifest_path: str,
+    existing_project_root: Optional[str] = None,
+) -> ImportResult:
+    """Parse a dbt manifest.json and return merged DataLex source/model docs.
+
+    When `existing_project_root` is provided, documents with matching
+    `meta.datalex.dbt.unique_id` are merged (user-authored fields preserved).
+    """
+    with open(manifest_path, "r", encoding="utf-8") as f:
+        manifest = json.load(f)
+
+    existing = _load_existing_by_unique_id(existing_project_root) if existing_project_root else {}
+
+    result = ImportResult()
+
+    nodes = manifest.get("nodes") or {}
+    sources = manifest.get("sources") or {}
+
+    # Sources are keyed by source_name in dbt; group per source_name so we emit one file per source.
+    sources_grouped: Dict[str, List[Dict[str, Any]]] = {}
+    for uid, node in sources.items():
+        source_name = node.get("source_name") or node.get("name")
+        sources_grouped.setdefault(source_name, []).append(node)
+
+    for source_name, tables in sources_grouped.items():
+        doc = _build_source_doc(source_name, tables, existing)
+        result.sources[doc["name"]] = doc
+
+    for uid, node in nodes.items():
+        if node.get("resource_type") != "model":
+            continue
+        doc = _build_model_doc(node, existing)
+        result.models[doc["name"]] = doc
+
+    return result
+
+
+def write_import_result(result: ImportResult, out_root: str) -> List[str]:
+    """Persist an ImportResult into a DataLex-style tree under out_root.
+
+    Writes:
+      <out_root>/sources/<name>.yaml
+      <out_root>/models/dbt/<name>.yaml
+    """
+    out = Path(out_root)
+    written: List[str] = []
+
+    for doc in result.sources.values():
+        path = out / "sources" / f"{doc['name']}.yaml"
+        _write_yaml(path, doc)
+        written.append(str(path))
+
+    for doc in result.models.values():
+        path = out / "models" / "dbt" / f"{doc['name']}.yaml"
+        _write_yaml(path, doc)
+        written.append(str(path))
+
+    return written
+
+
+# ------------------------ builders ------------------------
+
+
+def _build_source_doc(
+    source_name: str,
+    tables: List[Dict[str, Any]],
+    existing: Dict[str, Dict[str, Any]],
+) -> Dict[str, Any]:
+    # Source-level attributes come from the first table (dbt stores them per-node; pick any).
+    first = tables[0]
+    database = first.get("database")
+    schema = first.get("schema")
+
+    # Look for an existing source doc matching any of these tables' unique_ids so we
+    # preserve cross-table user fields (e.g., source-level owner).
+    existing_doc: Optional[Dict[str, Any]] = None
+    for t in tables:
+        uid = t.get("unique_id")
+        existing_doc = existing.get(uid) or existing_doc
+
+    doc: Dict[str, Any] = {
+        "kind": "source",
+        "name": _safe_name(source_name),
+    }
+    if existing_doc:
+        _merge_preserving_user_fields(doc, existing_doc, keys=("description", "owner", "tags", "loader", "loaded_at_field", "freshness"))
+
+    if database:
+        doc["database"] = database
+    if schema:
+        doc["schema"] = schema
+
+    table_docs: List[Dict[str, Any]] = []
+    for t in tables:
+        table_docs.append(_build_source_table_doc(t, existing_doc))
+    doc["tables"] = table_docs
+
+    # meta.datalex.dbt.unique_id list, so re-import can find this doc even if one table is renamed.
+    doc.setdefault("meta", {}).setdefault("datalex", {})["dbt"] = {
+        "unique_ids": sorted(t.get("unique_id") for t in tables if t.get("unique_id")),
+    }
+    return doc
+
+
+def _build_source_table_doc(
+    t: Dict[str, Any],
+    existing_source: Optional[Dict[str, Any]],
+) -> Dict[str, Any]:
+    name = _safe_name(t.get("name", ""))
+    table_doc: Dict[str, Any] = {"name": name}
+
+    # Locate prior table body if present
+    prior_table: Dict[str, Any] = {}
+    if existing_source:
+        for candidate in existing_source.get("tables", []) or []:
+            if candidate.get("name") == name:
+                prior_table = candidate
+                break
+
+    if t.get("description"):
+        table_doc["description"] = t["description"]
+    elif prior_table.get("description"):
+        table_doc["description"] = prior_table["description"]
+
+    if t.get("identifier") and t["identifier"] != name:
+        table_doc["identifier"] = t["identifier"]
+    if t.get("loaded_at_field"):
+        table_doc["loaded_at_field"] = t["loaded_at_field"]
+    if t.get("freshness"):
+        table_doc["freshness"] = t["freshness"]
+
+    # columns
+    cols_out: List[Dict[str, Any]] = []
+    prior_cols = {c.get("name"): c for c in (prior_table.get("columns") or []) if c.get("name")}
+    for c in t.get("columns", {}).values() if isinstance(t.get("columns"), dict) else (t.get("columns") or []):
+        cols_out.append(_build_source_column_doc(c, prior_cols.get(c.get("name"), {})))
+    if cols_out:
+        table_doc["columns"] = cols_out
+
+    # unique_id preserved at table level too
+    if t.get("unique_id"):
+        table_doc.setdefault("meta", {}).setdefault("datalex", {}).setdefault("dbt", {})[
+            "unique_id"
+        ] = t["unique_id"]
+
+    return table_doc
+
+
+def _build_source_column_doc(c: Dict[str, Any], prior: Dict[str, Any]) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": c.get("name")}
+    # type: manifest owns it; prefer manifest value if present
+    if c.get("data_type"):
+        doc["type"] = c["data_type"]
+    elif prior.get("type"):
+        doc["type"] = prior["type"]
+
+    # user-authored: preserve
+    for k in ("description", "sensitivity", "tags"):
+        if prior.get(k):
+            doc[k] = prior[k]
+    # manifest description wins only if user has no override
+    if c.get("description") and "description" not in doc:
+        doc["description"] = c["description"]
+
+    return doc
+
+
+def _build_model_doc(
+    node: Dict[str, Any],
+    existing: Dict[str, Dict[str, Any]],
+) -> Dict[str, Any]:
+    name = _safe_name(node.get("name", ""))
+    uid = node.get("unique_id")
+    prior = existing.get(uid, {}) if uid else {}
+
+    doc: Dict[str, Any] = {
+        "kind": "model",
+        "name": name,
+    }
+    # user-owned fields preserved
+    _merge_preserving_user_fields(
+        doc, prior, keys=("description", "owner", "domain", "tags", "materialization", "contract"),
+    )
+
+    # manifest-owned fields
+    config = node.get("config") or {}
+    if config.get("materialized") and "materialization" not in doc:
+        doc["materialization"] = config["materialized"]
+    if node.get("database"):
+        doc["database"] = node["database"]
+    if node.get("schema"):
+        doc["schema"] = node["schema"]
+    if node.get("description") and "description" not in doc:
+        doc["description"] = node["description"]
+
+    # depends_on — from manifest; represent both refs and sources
+    depends: List[Dict[str, Any]] = []
+    for parent_uid in (node.get("depends_on", {}) or {}).get("nodes", []) or []:
+        if parent_uid.startswith("model."):
+            depends.append({"ref": _safe_name(parent_uid.rsplit(".", 1)[-1])})
+        elif parent_uid.startswith("source."):
+            parts = parent_uid.split(".")
+            if len(parts) >= 4:
+                depends.append({"source": {"source": _safe_name(parts[-2]), "name": _safe_name(parts[-1])}})
+    if depends:
+        doc["depends_on"] = depends
+
+    # columns
+    prior_cols = {c.get("name"): c for c in (prior.get("columns") or []) if c.get("name")}
+    cols_out: List[Dict[str, Any]] = []
+    columns_raw = node.get("columns") or {}
+    column_iter = columns_raw.values() if isinstance(columns_raw, dict) else columns_raw
+    for c in column_iter:
+        cols_out.append(_build_model_column_doc(c, prior_cols.get(c.get("name"), {})))
+    if cols_out:
+        doc["columns"] = cols_out
+
+    if uid:
+        doc.setdefault("meta", {}).setdefault("datalex", {})["dbt"] = {"unique_id": uid}
+
+    return doc
+
+
+def _build_model_column_doc(c: Dict[str, Any], prior: Dict[str, Any]) -> Dict[str, Any]:
+    doc: Dict[str, Any] = {"name": c.get("name")}
+    if c.get("data_type"):
+        doc["type"] = c["data_type"]
+    elif prior.get("type"):
+        doc["type"] = prior["type"]
+
+    for k in ("description", "sensitivity", "tags", "terms", "tests", "constraints"):
+        if prior.get(k):
+            doc[k] = prior[k]
+    if c.get("description") and "description" not in doc:
+        doc["description"] = c["description"]
+    return doc
+
+
+# ------------------------ helpers ------------------------
+
+
+def _load_existing_by_unique_id(project_root: str) -> Dict[str, Dict[str, Any]]:
+    """Walk the project tree and index every doc by its `meta.datalex.dbt.unique_id(s)`."""
+    out: Dict[str, Dict[str, Any]] = {}
+    root = Path(project_root)
+    if not root.exists():
+        return out
+    for path in root.rglob("*.yaml"):
+        try:
+            with path.open("r", encoding="utf-8") as f:
+                doc = yaml.safe_load(f)
+        except Exception:
+            continue
+        if not isinstance(doc, dict):
+            continue
+        meta = (doc.get("meta") or {}).get("datalex") or {}
+        dbt_meta = meta.get("dbt") or {}
+        uid = dbt_meta.get("unique_id")
+        uids = dbt_meta.get("unique_ids") or ([uid] if uid else [])
+        for u in uids:
+            if u:
+                out[u] = doc
+    return out
+
+
+def _merge_preserving_user_fields(
+    dst: Dict[str, Any],
+    src: Dict[str, Any],
+    keys: tuple,
+) -> None:
+    for k in keys:
+        if src.get(k) not in (None, "", [], {}):
+            dst[k] = src[k]
+
+
+def _safe_name(name: str) -> str:
+    """DataLex names must match ^[a-z][a-z0-9_]*$ — coerce dbt names that drift."""
+    import re
+
+    if not name:
+        return "unnamed"
+    s = name.strip().lower()
+    s = re.sub(r"[^a-z0-9_]", "_", s)
+    if not re.match(r"^[a-z]", s):
+        s = "n_" + s
+    return s
+
+
+def _write_yaml(path: Path, doc: Dict[str, Any]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", encoding="utf-8") as f:
+        yaml.safe_dump(doc, f, sort_keys=False, default_flow_style=False, allow_unicode=True)