dbt-scribe 0.1.0__py3-none-any.whl

dbt_scribe/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

dbt_scribe/analyzer.py

@@ -0,0 +1,348 @@
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+
+import sqlglot
+from sqlglot import exp
+
+from dbt_scribe.config import ConventionsConfig, ScribeConfig, TestsConfig
+from dbt_scribe.parsers.manifest_parser import ManifestNode
+from dbt_scribe.parsers.yaml_parser import YamlModel, is_description_set
+
+
+class Layer(Enum):
+    STAGING = "staging"
+    INTERMEDIATE = "intermediate"
+    MARTS = "marts"
+    UNKNOWN = "unknown"
+
+
+class ColumnType(Enum):
+    PRIMARY_KEY = "primary_key"
+    FOREIGN_KEY = "foreign_key"
+    ENUM = "enum"
+    SHARED = "shared"
+    BOOLEAN = "boolean"
+    TIMESTAMP = "timestamp"
+    METRIC = "metric"
+    CALCULATED = "calculated"
+    TEXT = "text"
+
+
+@dataclass(frozen=True)
+class EnrichedColumn:
+    name: str
+    data_type: str | None
+    description: str | None
+    tests: list[Any]
+    column_type: ColumnType
+    sql_expression: str | None
+    needs_doc: bool
+    needs_tests: bool
+
+
+@dataclass(frozen=True)
+class EnrichedModel:
+    unique_id: str
+    name: str
+    fqn: list[str]
+    layer: Layer
+    description: str | None
+    compiled_code: str
+    columns: dict[str, EnrichedColumn]
+    depends_on_nodes: list[str]
+    tags: list[str]
+    materialized: str | None
+    path: str
+    adapter_type: str | None
+
+
+AGGREGATION_PATTERN = re.compile(
+    r"\b(count|sum|avg|min|max|median)\s*\(", re.IGNORECASE
+)
+ARITHMETIC_PATTERN = re.compile(r"\s[+\-*/]\s")
+
+
+def detect_layer(fqn: list[str], config: ConventionsConfig) -> Layer:
+    """Detect the dbt layer from a node's fully-qualified name.
+
+    Matches fqn[1] against configured layer prefixes. Also accepts common
+    alternate spellings (e.g. 'mart' for 'marts') to avoid silent failures
+    when the project folder name differs slightly from the config value.
+
+    Args:
+        fqn: Fully-qualified node name list from the manifest.
+        config: Conventions config containing layer folder names.
+
+    Returns:
+        The detected Layer enum value, or Layer.UNKNOWN if unrecognised.
+    """
+    if len(fqn) < 2:
+        return Layer.UNKNOWN
+
+    layer_name = fqn[1]
+
+    # Exact match first
+    if layer_name == config.staging_prefix:
+        return Layer.STAGING
+    if layer_name == config.intermediate_prefix:
+        return Layer.INTERMEDIATE
+    if layer_name == config.marts_prefix:
+        return Layer.MARTS
+
+    # Common alternate spellings — mart vs marts
+    if layer_name in ("mart", "marts", "mrt"):
+        return Layer.MARTS
+    if layer_name in ("staging", "stage", "stg"):
+        return Layer.STAGING
+    if layer_name in ("intermediate", "int"):
+        return Layer.INTERMEDIATE
+
+    return Layer.UNKNOWN
+
+
+def _find_primary_key_column(
+    column_names: list[str],
+    config: TestsConfig,
+) -> str | None:
+    """Return the name of the most likely primary key column, or None.
+
+    A column is considered the primary key only when it is the sole column
+    matching pk_patterns. When multiple columns match (e.g. order_id and
+    user_id in an orders table), none is a true PK — they are foreign keys
+    and should be typed accordingly.
+
+    Args:
+        column_names: Ordered list of column names from the model.
+        config: Tests config containing pk_patterns.
+
+    Returns:
+        The PK column name if exactly one column matches, otherwise None.
+    """
+    pk_candidates = [
+        name for name in column_names
+        if _matches_any(name, config.pk_patterns)
+    ]
+    # Only treat as PK when there is exactly one candidate — unambiguous
+    return pk_candidates[0] if len(pk_candidates) == 1 else None
+
+
+def infer_column_type(
+    column_name: str,
+    sql_expression: str | None,
+    config: TestsConfig,
+    *,
+    shared_columns: list[str] | None = None,
+    pk_column: str | None = None,
+) -> ColumnType:
+    """Infer the semantic type of a column from its name, SQL expression, and config.
+
+    Applies a priority cascade — the first matching rule wins:
+    1. Exact PK match (only when pk_column is set — i.e. unambiguous single PK)
+    2. FK patterns from config
+    3. PK patterns not designated as PK → FK (e.g. order_id in a multi-key model)
+    4. Enum patterns from config
+    5. Shared columns list from docs config
+    6. Boolean name heuristics (is_, has_, did_)
+    7. Timestamp name heuristics (_at, _date)
+    8. Metric (aggregation in SQL expression)
+    9. Calculated (arithmetic in SQL expression)
+    10. Default: text
+
+    Args:
+        column_name: The column name to classify.
+        sql_expression: The SQL expression for this column if computed, else None.
+        config: Tests config with pk/fk/enum patterns.
+        shared_columns: List of shared column names from docs config.
+        pk_column: The designated PK column name for this model, or None.
+
+    Returns:
+        The inferred ColumnType.
+    """
+    # Rule 1 — designated primary key (unambiguous single PK)
+    if pk_column and column_name == pk_column:
+        return ColumnType.PRIMARY_KEY
+
+    # Rule 2 — explicit FK patterns (e.g. customer_fk)
+    if _matches_any(column_name, config.fk_patterns):
+        return ColumnType.FOREIGN_KEY
+
+    # Rule 3 — matches pk_patterns but NOT the designated PK → treat as FK
+    # Example: user_id in an orders table where order_id is the PK
+    if _matches_any(column_name, config.pk_patterns):
+        return ColumnType.FOREIGN_KEY
+
+    # Rule 4 — enum patterns
+    if _matches_any(column_name, config.enum_patterns):
+        return ColumnType.ENUM
+
+    # Rule 5 — shared columns
+    if shared_columns and column_name in shared_columns:
+        return ColumnType.SHARED
+
+    # Rule 6 — boolean heuristics
+    if column_name.startswith(("is_", "has_", "did_")):
+        return ColumnType.BOOLEAN
+
+    # Rule 7 — timestamp heuristics
+    if column_name.endswith(("_at", "_date")):
+        return ColumnType.TIMESTAMP
+
+    # Rule 8 — metric (aggregation function in SQL expression)
+    if sql_expression and AGGREGATION_PATTERN.search(sql_expression):
+        return ColumnType.METRIC
+
+    # Rule 9 — calculated (arithmetic in SQL expression)
+    if sql_expression and ARITHMETIC_PATTERN.search(sql_expression):
+        return ColumnType.CALCULATED
+
+    # Rule 10 — default
+    return ColumnType.TEXT
+
+
+def build_enriched_model(
+    node: ManifestNode,
+    yaml_model: YamlModel | None,
+    config: ScribeConfig,
+    *,
+    overwrite_existing: bool = False,
+) -> EnrichedModel:
+    """Build an EnrichedModel from a manifest node and existing YAML state.
+
+    Combines manifest metadata, existing YAML documentation and tests, compiled
+    SQL expressions, and generation flags (needs_doc, needs_tests) into a single
+    object consumed by the generators and writers.
+
+    Args:
+        node: Parsed manifest node for this model.
+        yaml_model: Existing YAML state for this model, or None if no YAML exists.
+        config: Project-level dbt-scribe configuration.
+        overwrite_existing: When True, sets needs_doc and needs_tests to True
+            even for columns that already have descriptions or tests.
+
+    Returns:
+        A fully populated EnrichedModel instance.
+    """
+    expressions = _extract_select_expressions(node.compiled_code, node.adapter_type)
+    yaml_columns = yaml_model.columns if yaml_model else {}
+
+    # Determine the unambiguous PK column for this model before typing columns
+    pk_column = _find_primary_key_column(
+        list(node.columns.keys()), config.tests
+    )
+
+    columns = {}
+    for column_name, manifest_column in node.columns.items():
+        yaml_column = yaml_columns.get(column_name)
+        description = (
+            yaml_column.description
+            if yaml_column and yaml_column.description is not None
+            else manifest_column.description
+        )
+        tests = list(yaml_column.tests) if yaml_column else []
+        sql_expression = expressions.get(column_name)
+
+        columns[column_name] = EnrichedColumn(
+            name=column_name,
+            data_type=manifest_column.data_type,
+            description=description,
+            tests=tests,
+            column_type=infer_column_type(
+                column_name,
+                sql_expression,
+                config.tests,
+                shared_columns=config.docs.shared_columns,
+                pk_column=pk_column,
+            ),
+            sql_expression=sql_expression,
+            needs_doc=overwrite_existing or not is_description_set(description),
+            needs_tests=overwrite_existing or not tests,
+        )
+
+    return EnrichedModel(
+        unique_id=node.unique_id,
+        name=node.name,
+        fqn=node.fqn,
+        layer=detect_layer(node.fqn, config.conventions),
+        description=yaml_model.description if yaml_model else None,
+        compiled_code=node.compiled_code,
+        columns=columns,
+        depends_on_nodes=node.depends_on_nodes,
+        tags=node.tags,
+        materialized=node.materialized,
+        path=node.path,
+        adapter_type=node.adapter_type,
+    )
+
+
+def _matches_any(value: str, patterns: list[str]) -> bool:
+    """Return True if value matches any of the given regex patterns.
+
+    Args:
+        value: String to test.
+        patterns: List of regex pattern strings.
+
+    Returns:
+        True if any pattern matches.
+    """
+    return any(re.match(pattern, value) for pattern in patterns)
+
+
+def _extract_select_expressions(
+    compiled_code: str,
+    adapter_type: str | None,
+) -> dict[str, str]:
+    """Extract column name → SQL expression mappings from compiled SQL.
+
+    Uses sqlglot to parse the compiled SQL and extract the expression for
+    each selected column. Falls back to a regex-based approach if sqlglot
+    fails to parse the SQL.
+
+    Args:
+        compiled_code: Fully compiled SQL string (Jinja2 already resolved).
+        adapter_type: dbt adapter name used as the sqlglot read dialect.
+
+    Returns:
+        Dict mapping column alias/name to its SQL expression string.
+    """
+    try:
+        tree = sqlglot.parse_one(compiled_code, read=adapter_type or None)
+    except sqlglot.errors.SqlglotError:
+        return _extract_select_expressions_fallback(compiled_code)
+
+    expressions: dict[str, str] = {}
+    for select in tree.find_all(exp.Select):
+        for select_expression in select.expressions:
+            alias = select_expression.alias_or_name
+            if alias:
+                expressions[alias] = (
+                    select_expression.this.sql()
+                    if isinstance(select_expression, exp.Alias)
+                    else select_expression.sql()
+                )
+    return expressions
+
+
+def _extract_select_expressions_fallback(compiled_code: str) -> dict[str, str]:
+    """Regex-based fallback for SQL expression extraction when sqlglot fails.
+
+    Args:
+        compiled_code: Fully compiled SQL string.
+
+    Returns:
+        Dict mapping column alias/name to its SQL expression string.
+    """
+    expressions: dict[str, str] = {}
+    for expression, alias in re.findall(
+        r"(?im)^\s*(.+?)\s+as\s+([a-zA-Z_][a-zA-Z0-9_]*)\s*,?$",
+        compiled_code,
+    ):
+        expressions[alias] = expression.strip()
+    for column in re.findall(
+        r"(?im)^\s*([a-zA-Z_][a-zA-Z0-9_\.]*)\s*,?$", compiled_code
+    ):
+        expressions[column.split(".")[-1]] = column.strip()
+    return expressions

dbt_scribe/cli.py

@@ -0,0 +1,254 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import click
+
+from dbt_scribe import __version__
+from dbt_scribe.analyzer import build_enriched_model
+from dbt_scribe.config import ConfigError, ScribeConfig, load_config, resolve_provider
+from dbt_scribe.generators.docs_generator import DocsResult, generate_docs
+from dbt_scribe.generators.tests_generator import TestsResult, generate_tests
+from dbt_scribe.parsers.manifest_parser import ManifestNode, parse_manifest
+from dbt_scribe.parsers.yaml_parser import is_description_set, parse_yaml
+from dbt_scribe.resolver import resolve_target
+from dbt_scribe.writers.docs_writer import write_docs_block
+from dbt_scribe.writers.yaml_writer import write_yaml
+
+_REQUIRED_FILES = ["dbt_project.yml", "target/manifest.json", "dbt-scribe.yml"]
+
+_DEFAULT_CONFIG = """\
+version: 1
+
+llm:
+  provider: anthropic
+  # model: claude-sonnet-4-20250514  # defaults to latest Sonnet
+  temperature: 0.2
+
+docs:
+  two_tier: true
+  shared_columns:
+    - created_at
+    - updated_at
+    - _fivetran_synced
+  default_owner: "Data Team"
+
+tests:
+  pk_patterns:
+    - "^.*_id$"
+    - "^id$"
+  fk_patterns:
+    - "^.*_fk$"
+  enum_patterns:
+    - "^.*_type$"
+    - "^.*_status$"
+    - "^.*_category$"
+
+conventions:
+  staging_prefix: staging
+  intermediate_prefix: intermediate
+  marts_prefix: marts
+
+coverage:
+  min_doc_coverage: 80
+  min_test_coverage: 70
+"""
+
+
+class BootstrapError(Exception):
+    pass
+
+
+def _bootstrap(cwd: Path | None = None) -> None:
+    root = cwd or Path.cwd()
+    missing = [f for f in _REQUIRED_FILES if not (root / f).exists()]
+    if not missing:
+        return
+    hints = {
+        "dbt_project.yml": "Are you in the dbt project root directory?",
+        "target/manifest.json": "Run 'dbt compile' to generate the manifest.",
+        "dbt-scribe.yml": "Run 'dbt-scribe init' to create the configuration file.",
+    }
+    lines = ["Bootstrap check failed. Missing files:"]
+    for f in missing:
+        lines.append(f"  • {f}  →  {hints[f]}")
+    raise BootstrapError("\n".join(lines))
+
+
+@click.group()
+@click.version_option(__version__, prog_name="dbt-scribe")
+def cli() -> None:
+    """LLM-powered documentation and test generation for dbt Core projects."""
+
+
+@cli.command()
+def init() -> None:
+    """Generate a dbt-scribe.yml config in the current dbt project."""
+    root = Path.cwd()
+    if not (root / "dbt_project.yml").exists():
+        raise click.ClickException(
+            "dbt_project.yml not found. Are you in the dbt project root directory?"
+        )
+    config_path = root / "dbt-scribe.yml"
+    if config_path.exists():
+        raise click.ClickException(
+            "dbt-scribe.yml already exists. Remove it first or edit it directly."
+        )
+    config_path.write_text(_DEFAULT_CONFIG)
+    click.echo(f"Created {config_path}")
+    click.echo("Edit dbt-scribe.yml to set your LLM provider and project conventions.")
+
+
+@cli.command()
+@click.option("--target", default="models/", help="File, directory, or project root to process.")
+@click.option("--dry-run", is_flag=True, help="Preview changes without writing files.")
+@click.option("--force", is_flag=True, help="Overwrite existing descriptions.")
+def docs(target: str, dry_run: bool, force: bool) -> None:
+    """Generate model and column documentation."""
+    config, nodes = _load_project(target)
+    provider = resolve_provider(config)
+    for node in nodes:
+        model = _build_model(node, config, overwrite_existing=force)
+        docs_result = generate_docs(model, provider, config)
+        yaml_result = write_yaml(
+            model,
+            docs_result,
+            TestsResult(columns={}),
+            config,
+            dry_run,
+            force=force,
+        )
+        docs_block_result = write_docs_block(model, docs_result, config, dry_run)
+        _echo_status("docs", model.name, dry_run, yaml_result.changed or docs_block_result.changed)
+    _echo_summary("docs", nodes, dry_run)
+
+
+@cli.command()
+@click.option("--target", default="models/", help="File, directory, or project root to process.")
+@click.option("--dry-run", is_flag=True, help="Preview changes without writing files.")
+@click.option("--force", is_flag=True, help="Overwrite existing tests.")
+def tests(target: str, dry_run: bool, force: bool) -> None:
+    """Generate generic YAML tests."""
+    config, nodes = _load_project(target)
+    provider = resolve_provider(config)
+    for node in nodes:
+        model = _build_model(node, config, overwrite_existing=force)
+        tests_result = generate_tests(model, provider, config)
+        yaml_result = write_yaml(
+            model,
+            DocsResult(model_description=model.description or "", docs_block_content="", columns={}),
+            tests_result,
+            config,
+            dry_run,
+            force=force,
+        )
+        _echo_status("tests", model.name, dry_run, yaml_result.changed)
+    _echo_summary("tests", nodes, dry_run)
+
+
+@cli.command()
+@click.option("--target", default="models/", help="File, directory, or project root to process.")
+@click.option("--dry-run", is_flag=True, help="Preview changes without writing files.")
+@click.option("--force", is_flag=True, help="Overwrite existing documentation and tests.")
+def generate(target: str, dry_run: bool, force: bool) -> None:
+    """Generate documentation and tests in one pass."""
+    config, nodes = _load_project(target)
+    provider = resolve_provider(config)
+    for node in nodes:
+        model = _build_model(node, config, overwrite_existing=force)
+        docs_result = generate_docs(model, provider, config)
+        tests_result = generate_tests(model, provider, config)
+        yaml_result = write_yaml(
+            model,
+            docs_result,
+            tests_result,
+            config,
+            dry_run,
+            force=force,
+        )
+        docs_block_result = write_docs_block(model, docs_result, config, dry_run)
+        _echo_status("generate", model.name, dry_run, yaml_result.changed or docs_block_result.changed)
+    _echo_summary("generate", nodes, dry_run)
+
+
+@cli.command()
+@click.option("--target", default="models/", help="File, directory, or project root to audit.")
+def audit(target: str) -> None:
+    """Show documentation and test coverage report."""
+    config, nodes = _load_project(target, check_api_key=False)
+    click.echo("Audit summary")
+    for node in nodes:
+        model = _build_model(node, config)
+        total_columns = len(model.columns)
+        documented_columns = sum(
+            1 for column in model.columns.values() if is_description_set(column.description)
+        )
+        tested_columns = sum(1 for column in model.columns.values() if column.tests)
+        doc_coverage = _percentage(documented_columns, total_columns)
+        test_coverage = _percentage(tested_columns, total_columns)
+        click.echo(
+            f"- {model.name}: doc coverage {doc_coverage}% "
+            f"({documented_columns}/{total_columns}), test coverage {test_coverage}% "
+            f"({tested_columns}/{total_columns})"
+        )
+
+
+def _load_project(
+    target: str,
+    *,
+    check_api_key: bool = True,
+) -> tuple[ScribeConfig, list[ManifestNode]]:
+    try:
+        _bootstrap()
+        model_root = _read_model_root()
+        config = load_config("dbt-scribe.yml", check_api_key=check_api_key, model_root=model_root)
+        nodes = resolve_target(target, parse_manifest("target/manifest.json"), model_root=model_root)
+    except (BootstrapError, ValueError, ConfigError) as exc:
+        raise click.ClickException(str(exc)) from exc
+    return config, nodes
+
+
+def _read_model_root() -> str:
+    import yaml as _yaml
+
+    try:
+        raw = _yaml.safe_load(Path("dbt_project.yml").read_text()) or {}
+    except Exception:
+        return "models"
+    paths = raw.get("model-paths") or raw.get("source-paths")
+    if isinstance(paths, list) and paths:
+        return paths[0]
+    return "models"
+
+
+def _build_model(
+    node: ManifestNode,
+    config: ScribeConfig,
+    *,
+    overwrite_existing: bool = False,
+):
+    yaml_path = Path(config.model_root) / Path(node.path).with_suffix(".yml")
+    yaml_model = parse_yaml(yaml_path)
+    return build_enriched_model(
+        node,
+        yaml_model,
+        config,
+        overwrite_existing=overwrite_existing,
+    )
+
+
+def _echo_status(command: str, model_name: str, dry_run: bool, changed: bool) -> None:
+    mode = "dry-run" if dry_run else "write"
+    status = "changed" if changed else "unchanged"
+    click.echo(f"{command}: {model_name} [{mode}, {status}]")
+
+
+def _echo_summary(command: str, nodes: list[ManifestNode], dry_run: bool) -> None:
+    suffix = " (dry-run)" if dry_run else ""
+    click.echo(f"Summary: {command} processed {len(nodes)} model(s){suffix}.")
+
+
+def _percentage(numerator: int, denominator: int) -> int:
+    if denominator == 0:
+        return 100
+    return round((numerator / denominator) * 100)