dbt-costgate 0.7.1__py3-none-any.whl

dbt_costgate/__init__.py

@@ -0,0 +1,4 @@
+# SPDX-License-Identifier: Apache-2.0
+"""dbt-costgate — BigQuery cost gate for dbt pull requests."""
+
+__version__ = "0.7.1"

dbt_costgate/against.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Auto-compile a base ref as the diff baseline, in an isolated git worktree.
+
+The opt-in local convenience carved out by ADR-0006: unlike the pure core (which
+only consumes compiled artifacts) and gitdiff.py (git-only), this edge needs
+*both* git and dbt. It checks ``<ref>`` out into a throwaway worktree, runs
+``dbt compile`` there, and hands the resulting manifest back as the baseline — so
+a local before/after diff is one command with no manual ``--baseline``.
+
+The dbt project may sit at the git repo root (the common case) or in a
+subdirectory (monorepos): the repo root is detected via ``git rev-parse
+--show-toplevel``, so the worktree compile targets the project's actual location.
+Every failure degrades to an ``AgainstError`` with an actionable message; the
+worktree is always removed, even on failure or Ctrl+C.
+"""
+
+from __future__ import annotations
+
+import shutil
+import subprocess
+import sys
+import tempfile
+from collections.abc import Callable
+from pathlib import Path
+
+from dbt_costgate import artifacts
+from dbt_costgate.models import ModelNode
+
+# Third-party packages live here and are gitignored, so a fresh worktree checkout
+# omits them. dbt_modules is the pre-1.0 name; both are handled for old projects.
+_PACKAGE_DIRS = ("dbt_packages", "dbt_modules")
+
+
+class AgainstError(Exception):
+    """Producing the baseline via a worktree compile was not possible."""
+
+
+def _git(project_dir: Path, *args: str, check: bool = True) -> subprocess.CompletedProcess:
+    proc = subprocess.run(["git", *args], cwd=project_dir, capture_output=True, text=True)
+    if check and proc.returncode != 0:
+        raise AgainstError(proc.stderr.strip() or f"git {' '.join(args)} failed")
+    return proc
+
+
+def _resolve_dbt() -> str:
+    """Find dbt without trusting a bare name on $PATH (shell aliases are invisible
+    to subprocess; venvs may not export bin/ to child shells)."""
+    found = shutil.which("dbt")
+    if found:
+        return found
+    candidate = Path(sys.executable).parent / "dbt"
+    if candidate.is_file():
+        return str(candidate)
+    raise AgainstError(
+        "dbt executable not found on PATH or active virtual environment. "
+        "Install dbt, or use --baseline with a pre-compiled manifest."
+    )
+
+
+def _dbt_compile(worktree_dir: Path) -> None:
+    """Default compiler: run `dbt compile` in the worktree. Injectable so tests
+    exercise the whole flow without dbt, a warehouse, or credentials."""
+    proc = subprocess.run(
+        [_resolve_dbt(), "compile"], cwd=worktree_dir, capture_output=True, text=True
+    )
+    if proc.returncode != 0:
+        tail = "\n".join((proc.stderr or proc.stdout).strip().splitlines()[-8:])
+        raise AgainstError("dbt compile failed for the baseline ref:\n" + tail)
+
+
+def _link_packages(src_dir: Path, worktree_dir: Path) -> None:
+    """Make host-installed dbt packages visible to the worktree compile.
+
+    A fresh checkout omits gitignored dbt_packages/, so a project using dbt_utils
+    et al. would crash at compile on the first package macro. Symlink the host's
+    copy in (fast, offline); copy where symlinks aren't available. These are the
+    *current branch's* installed versions — good enough for a cost baseline.
+    """
+    for name in _PACKAGE_DIRS:
+        src = (src_dir / name).resolve()
+        if not src.is_dir():
+            continue
+        dest = worktree_dir / name
+        if dest.exists() or dest.is_symlink():
+            continue
+        try:
+            dest.symlink_to(src, target_is_directory=True)
+        except OSError:
+            shutil.copytree(src, dest)
+
+
+def compiled_baseline(
+    ref: str,
+    project_dir: Path,
+    *,
+    compile_fn: Callable[[Path], None] | None = None,
+) -> dict[str, ModelNode]:
+    """Compile ``ref`` in a throwaway worktree; return its model nodes as the
+    baseline. Raises AgainstError (never a stack trace) on any failure."""
+    # Resolved at call time (not bound as a default) so the real compiler can be
+    # monkeypatched in tests that drive the full worktree flow through the CLI.
+    if compile_fn is None:
+        compile_fn = _dbt_compile
+    project_dir = project_dir.resolve()
+
+    # Self-heal any worktree orphaned by a prior SIGKILL/power-loss — the one case
+    # the finally-block below genuinely can't cover.
+    _git(project_dir, "worktree", "prune", check=False)
+
+    # `git worktree add` checks out the *whole* repo, so a dbt project living in a
+    # subdir lands at <worktree>/<subdir>. Locate that subdir relative to the repo
+    # root (both resolved so the macOS /tmp -> /private/tmp symlink doesn't break
+    # relative_to).
+    repo_root = Path(_git(project_dir, "rev-parse", "--show-toplevel").stdout.strip()).resolve()
+    try:
+        subdir = project_dir.relative_to(repo_root)
+    except ValueError:
+        raise AgainstError(
+            f"project dir {project_dir} is not inside the git repo at {repo_root}."
+        ) from None
+
+    verify = _git(project_dir, "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}", check=False)
+    if verify.returncode != 0:
+        raise AgainstError(
+            f"couldn't resolve ref {ref!r} to compile as the baseline. Use an "
+            f"existing branch/tag/SHA, or --baseline with a pre-compiled manifest."
+        )
+
+    # mkdtemp creates the parent; the worktree path itself must not exist yet, as
+    # `git worktree add` creates it.
+    tmp_parent = Path(tempfile.mkdtemp(prefix="dbt-costgate-against-"))
+    worktree = tmp_parent / "wt"
+    try:
+        _git(project_dir, "worktree", "add", "--detach", "--quiet", str(worktree), ref)
+        project_in_wt = worktree / subdir
+        if not project_in_wt.is_dir():
+            raise AgainstError(
+                f"the dbt project dir {str(subdir)!r} doesn't exist in ref {ref!r} "
+                f"— it may have been added or moved after that ref."
+            )
+        _link_packages(project_dir, project_in_wt)
+        compile_fn(project_in_wt)
+        manifest = artifacts.load_manifest(project_in_wt / "target")
+        return artifacts.model_nodes(manifest)
+    except artifacts.ArtifactError as exc:
+        raise AgainstError(f"the baseline compile produced no usable manifest: {exc}") from exc
+    finally:
+        _git(project_dir, "worktree", "remove", "--force", str(worktree), check=False)
+        _git(project_dir, "worktree", "prune", check=False)
+        shutil.rmtree(tmp_parent, ignore_errors=True)

dbt_costgate/artifacts.py

@@ -0,0 +1,279 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Read dbt artifacts and select the cost-bearing models to estimate.
+
+Pure functions over manifest JSON — no dbt or git invocation. The only I/O is
+reading files the user already produced with ``dbt compile``.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+from dataclasses import dataclass
+from pathlib import Path
+
+from dbt_costgate.models import EstimateBasis, ModelNode
+
+# Statically-unknowable filters make a BigQuery dry-run fall back to a full-table
+# scan (see docs/architecture.md). We can't correct the number, only flag it.
+_DYNAMIC_FILTER = re.compile(
+    r"\b(current_date|current_timestamp|current_datetime)\b", re.IGNORECASE
+)
+_SUBQUERY_IN_PREDICATE = re.compile(r"\b(where|and|or)\b[^;]*\(\s*select\b", re.IGNORECASE)
+_MULTI_STATEMENT = re.compile(r"\b(declare|begin)\b", re.IGNORECASE)
+
+
+class ArtifactError(Exception):
+    """A manifest could not be read or is missing information dbt-costgate needs."""
+
+
+def manifest_path(path: Path) -> Path:
+    """Accept either a manifest.json or a dbt ``target/`` directory."""
+    if path.is_dir():
+        return path / "manifest.json"
+    return path
+
+
+def load_manifest(path: Path) -> dict:
+    mpath = manifest_path(path)
+    if not mpath.is_file():
+        raise ArtifactError(
+            f"No manifest.json at {mpath}. Run `dbt compile` first, then point "
+            f"--current at the target/ directory (or pass the manifest path)."
+        )
+    try:
+        return json.loads(mpath.read_text("utf-8"))
+    except json.JSONDecodeError as exc:  # pragma: no cover - defensive
+        raise ArtifactError(f"{mpath} is not valid JSON: {exc}") from exc
+
+
+def model_nodes(manifest: dict) -> dict[str, ModelNode]:
+    """Cost-bearing model nodes only, keyed by unique_id.
+
+    Filters to SQL models: skips Python models (compiled_code is Python, not SQL),
+    ephemeral models (no queryable relation), and every non-model resource type.
+    """
+    out: dict[str, ModelNode] = {}
+    for unique_id, node in (manifest.get("nodes") or {}).items():
+        if node.get("resource_type") != "model":
+            continue
+        if node.get("language", "sql") != "sql":
+            continue
+        config = node.get("config") or {}
+        if config.get("materialized") == "ephemeral":
+            continue
+        checksum = (node.get("checksum") or {}).get("checksum")
+        out[unique_id] = ModelNode(
+            unique_id=unique_id,
+            name=node.get("name", unique_id),
+            materialized=config.get("materialized", "view"),
+            language=node.get("language", "sql"),
+            database=node.get("database"),
+            schema=node.get("schema"),
+            relation_name=node.get("relation_name"),
+            checksum=checksum,
+            compiled_path=node.get("compiled_path"),
+            compiled_code=node.get("compiled_code"),
+            original_file_path=node.get("original_file_path"),
+            patch_path=node.get("patch_path"),
+            depends_on_macros=list((node.get("depends_on") or {}).get("macros") or []),
+        )
+    return out
+
+
+def has_any_compiled_code(nodes: dict[str, ModelNode]) -> bool:
+    return any(n.compiled_code for n in nodes.values())
+
+
+def resolve_compiled_sql(node: ModelNode, project_dir: Path | None) -> str | None:
+    """Prefer the on-disk compiled file (always written by ``dbt compile``);
+    fall back to the manifest's inlined ``compiled_code``."""
+    if node.compiled_path and project_dir is not None:
+        candidate = project_dir / node.compiled_path
+        if candidate.is_file():
+            return candidate.read_text("utf-8")
+    return node.compiled_code
+
+
+def _resolve_ref(ref: str, nodes: dict[str, ModelNode], side: str) -> str:
+    """Resolve a user-supplied model reference to a unique_id. Accepts a full
+    unique_id (used verbatim) or a bare model name (looked up by ``.name``)."""
+    if ref in nodes:
+        return ref
+    matches = [uid for uid, n in nodes.items() if n.name == ref]
+    if len(matches) == 1:
+        return matches[0]
+    if not matches:
+        raise ArtifactError(
+            f"renames: {side} model {ref!r} not found. Use its dbt model name or its "
+            f"full unique_id (e.g. model.<package>.<name>)."
+        )
+    raise ArtifactError(
+        f"renames: {side} model name {ref!r} is ambiguous across packages "
+        f"({', '.join(sorted(matches))}). Disambiguate with the full unique_id."
+    )
+
+
+def resolve_renames(
+    renames: dict[str, str],
+    current: dict[str, ModelNode],
+    baseline: dict[str, ModelNode],
+) -> dict[str, str]:
+    """Turn a user ``renames`` map (current -> baseline, by name or unique_id) into
+    a ``current_uid -> baseline_uid`` map against the loaded manifests. Fails loudly
+    on an unresolvable, ambiguous, or many-to-one mapping rather than mis-diffing."""
+    resolved: dict[str, str] = {}
+    for current_ref, baseline_ref in renames.items():
+        cur_uid = _resolve_ref(current_ref, current, "current")
+        base_uid = _resolve_ref(baseline_ref, baseline, "baseline")
+        resolved[cur_uid] = base_uid
+    if len(set(resolved.values())) < len(resolved):
+        raise ArtifactError(
+            "renames: two or more current models map to the same baseline model; "
+            "each baseline may be paired with only one current model."
+        )
+    return resolved
+
+
+def _compiled_differs(base: ModelNode, node: ModelNode) -> bool:
+    """Whether the two versions compile to different SQL. A missing side is
+    *unknown*, not different — a parse-only manifest must never select everything."""
+    if base.compiled_code is None or node.compiled_code is None:
+        return False
+    return base.compiled_code != node.compiled_code
+
+
+def select_changed(
+    baseline: dict[str, ModelNode],
+    current: dict[str, ModelNode],
+    renames: dict[str, str] | None = None,
+) -> list[str]:
+    """Models that were added, whose body changed (checksum differs), or whose
+    compiled SQL differs — the last catching config-only and macro-only changes,
+    which never touch the model's own file. A declared rename (``current_uid`` in
+    ``renames``) is always selected."""
+    renames = renames or {}
+    changed = []
+    for uid, node in current.items():
+        if uid in renames:
+            changed.append(uid)
+            continue
+        base = baseline.get(uid)
+        if base is None:
+            changed.append(uid)
+        elif node.checksum is None or base.checksum is None or node.checksum != base.checksum:
+            changed.append(uid)
+        elif _compiled_differs(base, node):
+            changed.append(uid)
+    return changed
+
+
+def indirect_changes(baseline: dict[str, ModelNode], current: dict[str, ModelNode]) -> set[str]:
+    """The models ``select_changed`` picks up *only* from their compiled SQL. Their
+    own file is untouched, so the cause is upstream — surfaced as a warning so a
+    model never appears in a report without an explanation."""
+    out: set[str] = set()
+    for uid, node in current.items():
+        base = baseline.get(uid)
+        if base is None or node.checksum is None or base.checksum is None:
+            continue
+        if node.checksum == base.checksum and _compiled_differs(base, node):
+            out.add(uid)
+    return out
+
+
+@dataclass(frozen=True)
+class MacroIndex:
+    """The macro graph as the manifest records it: the file each macro is defined
+    in, and the macros each macro calls."""
+
+    files: dict[str, str]  # macro unique_id -> original_file_path
+    depends_on: dict[str, list[str]]  # macro unique_id -> macro unique_ids it calls
+
+
+def macro_index(manifest: dict) -> MacroIndex:
+    files: dict[str, str] = {}
+    depends_on: dict[str, list[str]] = {}
+    for uid, macro in (manifest.get("macros") or {}).items():
+        path = macro.get("original_file_path")
+        if path:
+            files[uid] = path
+        depends_on[uid] = list((macro.get("depends_on") or {}).get("macros") or [])
+    return MacroIndex(files=files, depends_on=depends_on)
+
+
+def _changed_macros(paths: set[str], index: MacroIndex) -> set[str]:
+    """Macros defined in a changed file, plus every macro that calls one — dbt
+    treats the caller of a changed macro as changed too, and so must we."""
+    callers: dict[str, set[str]] = {}
+    for uid, deps in index.depends_on.items():
+        for dep in deps:
+            callers.setdefault(dep, set()).add(uid)
+    stack = [uid for uid, path in index.files.items() if path in paths]
+    seen: set[str] = set()
+    while stack:
+        uid = stack.pop()
+        if uid in seen:
+            continue
+        seen.add(uid)
+        stack.extend(callers.get(uid, ()))
+    return seen
+
+
+def _patch_file(patch_path: str) -> str:
+    """``patch_path`` is ``<package>://<path>``; the path half is project-relative,
+    the same basis as ``original_file_path``."""
+    return patch_path.split("://", 1)[-1]
+
+
+def select_by_paths(
+    nodes: dict[str, ModelNode],
+    paths: list[str],
+    macros: MacroIndex | None = None,
+) -> list[str]:
+    """Models a set of changed project-relative paths touches: the model's own file,
+    the YAML that patches it, or any macro in its macro closure. Needs no baseline
+    manifest, so it drives the zero-setup local run."""
+    changed = set(paths)
+    changed_macros = _changed_macros(changed, macros) if macros else set()
+    selected = []
+    for uid, node in nodes.items():
+        if node.original_file_path and node.original_file_path in changed:
+            selected.append(uid)
+        elif node.patch_path and _patch_file(node.patch_path) in changed:
+            selected.append(uid)
+        elif changed_macros.intersection(node.depends_on_macros):
+            selected.append(uid)
+    return selected
+
+
+def touches_project_config(paths: list[str]) -> bool:
+    """A ``dbt_project.yml`` change is project-wide config — real, but not
+    attributable to individual models from paths alone. Paths are project-relative,
+    so only this project's file matches, never a sibling project's."""
+    return any(p == "dbt_project.yml" for p in paths)
+
+
+def detect_basis(node: ModelNode, sql: str | None) -> EstimateBasis:
+    """A self-reference in an incremental model's compiled SQL means it was
+    compiled against an existing table (incremental form); its absence means the
+    full-refresh form. Non-incrementals are always direct."""
+    if not node.is_incremental:
+        return EstimateBasis.DIRECT
+    if sql and node.relation_name and node.relation_name in sql:
+        return EstimateBasis.INCREMENTAL_FORM
+    return EstimateBasis.FULL_REFRESH
+
+
+def sql_warnings(node: ModelNode, sql: str | None) -> list[str]:
+    """Non-fatal caveats about how trustworthy a dry-run of this SQL will be."""
+    warnings: list[str] = []
+    if not sql:
+        return warnings
+    if _DYNAMIC_FILTER.search(sql) or _SUBQUERY_IN_PREDICATE.search(sql):
+        warnings.append("dynamic filter — dry-run may be worst-case (overestimate)")
+    if _MULTI_STATEMENT.search(sql):
+        warnings.append("multi-statement — dry-run bytes may be partial")
+    if node.is_incremental:
+        warnings.append("incremental — figure is the full-refresh scan")
+    return warnings

dbt_costgate/bigquery.py

@@ -0,0 +1,110 @@
+# SPDX-License-Identifier: Apache-2.0
+"""The one network edge: BigQuery dry-run, behind a ``DryRunner`` protocol.
+
+dbt-costgate's only warehouse interaction. `dryRun=true` executes nothing, reads no
+table data, and is never billed. Credentials are never handled here — the client
+resolves them through Application Default Credentials, exactly like dbt-bigquery.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+from dbt_costgate.models import ErrorKind
+
+
+@dataclass
+class DryRunResult:
+    total_bytes: int | None = None
+    location: str | None = None
+    error_kind: ErrorKind | None = None
+    error_detail: str | None = None
+
+    @property
+    def ok(self) -> bool:
+        return self.error_kind is None and self.total_bytes is not None
+
+
+class DryRunner(Protocol):
+    def dry_run(self, sql: str, self_relation: str | None = None) -> DryRunResult: ...
+
+
+def categorize(exc: Exception, self_relation: str | None) -> tuple[ErrorKind, str]:
+    """Map a BigQuery client exception to a dbt-costgate error kind.
+
+    Matches on exception type name and message rather than importing google
+    exception classes, so the pure test suite can exercise it with stand-ins.
+    """
+    name = type(exc).__name__
+    msg = str(exc)
+    low = msg.lower()
+    if "notfound" in name.lower() or "not found" in low:
+        if self_relation and _relation_in_message(self_relation, msg):
+            return ErrorKind.DESTINATION_MISSING, msg
+        return ErrorKind.UPSTREAM_MISSING, msg
+    if "forbidden" in name.lower() or "permissiondenied" in name.lower() or "permission" in low:
+        return ErrorKind.PERMISSION, msg
+    if (
+        "serviceunavailable" in name.lower()
+        or "internalserver" in name.lower()
+        or "toomanyrequests" in name.lower()
+        or any(code in msg for code in ("429", "500", "503"))
+    ):
+        return ErrorKind.TRANSIENT, msg
+    if "badrequest" in name.lower() or "invalid" in low or "syntax" in low:
+        return ErrorKind.INVALID_SQL, msg
+    return ErrorKind.OTHER, msg
+
+
+def _relation_in_message(relation: str, msg: str) -> bool:
+    """A dbt relation_name looks like `project`.`dataset`.`table`; BigQuery's
+    404 text uses `project:dataset.table` or `dataset.table`. Compare on the bare
+    table token so backticks/separators don't matter."""
+    table = relation.replace("`", "").split(".")[-1]
+    return bool(table) and table in msg
+
+
+class BigQueryDryRunner:
+    """Real dry-runner. Imports google-cloud-bigquery lazily so the module (and
+    the pure test suite) load without credentials or the client installed."""
+
+    def __init__(self, project: str | None = None, deadline_seconds: float = 60.0):
+        self._project = project
+        self._deadline = deadline_seconds
+        self._client = None
+        self._retry = None
+
+    def _ensure_client(self):
+        if self._client is not None:
+            return
+        from google.api_core import exceptions as gexc
+        from google.api_core.retry import Retry
+        from google.cloud import bigquery
+
+        self._bigquery = bigquery
+        self._client = bigquery.Client(project=self._project)
+        self._retry = Retry(
+            predicate=lambda e: isinstance(
+                e,
+                (
+                    gexc.TooManyRequests,
+                    gexc.ServiceUnavailable,
+                    gexc.InternalServerError,
+                ),
+            ),
+            deadline=self._deadline,
+        )
+
+    def dry_run(self, sql: str, self_relation: str | None = None) -> DryRunResult:
+        try:
+            self._ensure_client()
+        except Exception as exc:  # ADC / import failure — operational, surfaced by caller
+            return DryRunResult(error_kind=ErrorKind.OTHER, error_detail=str(exc))
+        job_config = self._bigquery.QueryJobConfig(dry_run=True, use_query_cache=False)
+        try:
+            job = self._client.query(sql, job_config=job_config, retry=self._retry)
+            return DryRunResult(total_bytes=int(job.total_bytes_processed), location=job.location)
+        except Exception as exc:
+            kind, detail = categorize(exc, self_relation)
+            return DryRunResult(error_kind=kind, error_detail=detail)