wraith-modelgen 0.4.0__py3-none-any.whl

modelgen/__init__.py

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

modelgen/cli.py

@@ -0,0 +1,373 @@
+from __future__ import annotations
+
+import logging
+from pathlib import Path
+
+import typer
+from rich.console import Console
+from rich.logging import RichHandler
+from rich.panel import Panel
+from rich.table import Table
+
+from modelgen import __version__
+from modelgen.contract import ContractError, EventContract
+from modelgen.generator import GenerationError, Layer, generate
+from modelgen.introspect import IntrospectionError
+
+app = typer.Typer(
+    name="modelgen",
+    help=(
+        "BigQuery dbt model scaffolder. Generates one layer"
+        " at a time from a YAML data contract."
+    ),
+    no_args_is_help=True,
+    rich_markup_mode="rich",
+)
+console = Console()
+
+
+def _version_callback(value: bool) -> None:
+    if value:
+        console.print(f"modelgen {__version__}")
+        raise typer.Exit()
+
+
+@app.callback()
+def _main(
+    version: bool = typer.Option(
+        False,
+        "--version",
+        "-V",
+        callback=_version_callback,
+        is_eager=True,
+        help="Show version and exit.",
+    ),
+) -> None:
+    pass
+
+
+def _configure_logging(verbose: bool) -> None:
+    level = logging.DEBUG if verbose else logging.INFO
+    logging.basicConfig(
+        level=level,
+        format="%(message)s",
+        datefmt="[%X]",
+        handlers=[
+            RichHandler(
+                console=console, rich_tracebacks=False, show_path=False
+            )
+        ],
+    )
+
+
+@app.command()
+def gen(
+    contract: Path = typer.Argument(
+        ...,
+        help="Path to the YAML data contract file.",
+        exists=True,
+        readable=True,
+    ),
+    raw: bool = typer.Option(
+        False, "--raw", help="Generate the raw (Origin Data Product) layer."
+    ),
+    staging: bool = typer.Option(
+        False,
+        "--staging",
+        help="Generate the staging (Consumption Data Product) layer.",
+    ),
+    output: Path = typer.Option(
+        Path("./models"),
+        "--output",
+        "-o",
+        help="Directory to write generated files into.",
+    ),
+    dry_run: bool = typer.Option(
+        False,
+        "--dry-run",
+        help="Print what would be generated without writing.",
+    ),
+    verbose: bool = typer.Option(
+        False, "--verbose", "-v", help="Enable debug logging."
+    ),
+):
+    """
+    Generate a dbt model layer from a YAML data contract.
+
+    \b
+    Examples:
+        modelgen gen contract.yml --raw -o ./models/raw
+        modelgen gen contract.yml --staging -o ./models/staging
+    """
+    _configure_logging(verbose)
+
+    if raw == staging:
+        console.print("[red]Specify exactly one of --raw or --staging.[/red]")
+        raise typer.Exit(2)
+
+    layer = Layer.RAW if raw else Layer.STAGING
+
+    try:
+        event = EventContract.from_file(contract)
+    except ContractError as e:
+        console.print(f"[red]Invalid contract:[/red] {e}")
+        raise typer.Exit(1)
+
+    from modelgen.introspect import BigQueryIntrospector
+
+    introspector = BigQueryIntrospector()
+
+    target_dataset = (
+        event.raw.dataset if layer == Layer.RAW else event.staging.dataset
+    )
+    target_model = (
+        event.raw_model_name
+        if layer == Layer.RAW
+        else event.staging_model_name
+    )
+
+    console.print(
+        Panel(
+            f"[bold]Event:[/bold] {event.name}\n"
+            f"[bold]Layer:[/bold] {layer.value}\n"
+            f"[bold]Source:[/bold] {event.source.fqn}\n"
+            f"[bold]Target:[/bold] {target_dataset}.{target_model}\n"
+            f"[bold]Unique key:[/bold] {', '.join(event.unique_key)} | "
+            f"[bold]Loaded at:[/bold] {event.loaded_at_field}",
+            title="[bold green]modelgen[/bold green]",
+            expand=False,
+        )
+    )
+
+    try:
+        result = generate(
+            contract=event,
+            introspector=introspector,
+            output_dir=output,
+            layer=layer,
+            contract_path=str(contract),
+            dry_run=dry_run,
+        )
+    except (GenerationError, IntrospectionError) as e:
+        console.print(f"[red]Generation failed:[/red] {e}")
+        raise typer.Exit(1)
+    except Exception as e:
+        console.print(f"[red]Unexpected error:[/red] {type(e).__name__}: {e}")
+        if verbose:
+            console.print_exception()
+        raise typer.Exit(1)
+
+    table = Table(show_header=True, header_style="bold cyan")
+    table.add_column("File")
+    table.add_column("Status")
+    for path in result.written:
+        status = (
+            "[yellow]dry-run[/yellow]" if dry_run else "[green]written[/green]"
+        )
+        table.add_row(str(path), status)
+    console.print(table)
+
+    if result.new_source_columns:
+        console.print(
+            Panel(
+                f"Source has {len(result.new_source_columns)} column(s)"
+                " not yet in the staging contract:\n"
+                f"  {', '.join(result.new_source_columns)}\n\n"
+                "[dim]These flow through raw automatically. "
+                "Add to the staging contract when "
+                "downstream consumers need them.[/dim]",
+                title="[yellow]Schema evolution detected[/yellow]",
+                border_style="yellow",
+            )
+        )
+
+
+@app.command()
+def validate(
+    contract: Path = typer.Argument(
+        ...,
+        help="Path to the YAML data contract file.",
+        exists=True,
+        readable=True,
+    ),
+):
+    """Validate a contract's YAML structure (does not introspect BigQuery)."""
+    try:
+        event = EventContract.from_file(contract)
+        console.print(
+            f"[green]OK[/green] Contract valid: [bold]{event.name}[/bold]\n"
+            f"  Source: {event.source.fqn}\n"
+            f"  Unique key: {', '.join(event.unique_key)}\n"
+            f"  Staging columns declared: {len(event.staging.columns)}"
+        )
+    except ContractError as e:
+        console.print(f"[red]Invalid contract:[/red] {e}")
+        raise typer.Exit(1)
+
+
+@app.command()
+def run(
+    config: Path | None = typer.Option(
+        None,
+        "--config",
+        "-c",
+        help=(
+            "Path to .modelgen.yml. Auto-discovered from cwd"
+            " upward if not given."
+        ),
+    ),
+    dry_run: bool = typer.Option(
+        False,
+        "--dry-run",
+        help="Print what would be generated without writing.",
+    ),
+    verbose: bool = typer.Option(
+        False, "--verbose", "-v", help="Enable debug logging."
+    ),
+):
+    """
+    Discover all contracts in the repo and generate dbt models for every layer.
+
+    \b
+    Reads .modelgen.yml (or uses opinionated defaults) to find contracts
+    and write to:
+      models/raw/     — raw (Origin Data Product) layer
+      models/staging/ — staging layer (skipped if no columns declared)
+
+    \b
+    Default layout expected in the consuming repo:
+      contracts/          <- YAML data contracts
+      models/raw/         <- generated raw models (output)
+      models/staging/     <- generated staging models (output)
+    """
+    _configure_logging(verbose)
+
+    from modelgen.config import ConfigError, ModelgenConfig
+    from modelgen.introspect import BigQueryIntrospector
+
+    try:
+        if config is not None:
+            cfg = ModelgenConfig.from_file(config)
+        else:
+            cfg = ModelgenConfig.discover(Path.cwd())
+    except ConfigError as e:
+        console.print(f"[red]Config error:[/red] {e}")
+        raise typer.Exit(1)
+
+    contract_files = cfg.contract_files()
+    if not contract_files:
+        console.print(
+            f"[yellow]No contracts found.[/yellow] "
+            f"Pattern [dim]{cfg.contracts_glob!r}[/dim]"
+            f" matched nothing in {cfg.root}"
+        )
+        raise typer.Exit(0)
+
+    console.print(
+        Panel(
+            f"[bold]Contracts:[/bold] {len(contract_files)}\n"
+            f"[bold]Root:[/bold] {cfg.root}\n"
+            f"[bold]Raw output:[/bold] {cfg.raw_output}\n"
+            f"[bold]Staging output:[/bold] {cfg.staging_output}",
+            title="[bold green]modelgen run[/bold green]",
+            expand=False,
+        )
+    )
+
+    introspector = BigQueryIntrospector()
+
+    table = Table(show_header=True, header_style="bold cyan")
+    table.add_column("Contract")
+    table.add_column("Layer")
+    table.add_column("File")
+    table.add_column("Status")
+
+    errors: list[tuple[str, str, str]] = []
+    new_cols_notices: list[tuple[str, list[str]]] = []
+
+    for contract_path in contract_files:
+        try:
+            event = EventContract.from_file(contract_path)
+        except ContractError as e:
+            errors.append((contract_path.name, "parse", str(e)))
+            continue
+
+        try:
+            rel = contract_path.relative_to(cfg.root)
+        except ValueError:
+            rel = contract_path
+
+        # Raw — always
+        try:
+            raw_result = generate(
+                contract=event,
+                introspector=introspector,
+                output_dir=cfg.raw_output,
+                layer=Layer.RAW,
+                contract_path=str(rel),
+                dry_run=dry_run,
+            )
+            status = (
+                "[yellow]dry-run[/yellow]"
+                if dry_run
+                else "[green]written[/green]"
+            )
+            for p in raw_result.written:
+                table.add_row(contract_path.name, "raw", p.name, status)
+        except (GenerationError, IntrospectionError) as e:
+            errors.append((contract_path.name, "raw", str(e)))
+
+        # Staging — skip if no columns declared in the contract
+        if not event.staging.columns:
+            continue
+
+        try:
+            stg_result = generate(
+                contract=event,
+                introspector=introspector,
+                output_dir=cfg.staging_output,
+                layer=Layer.STAGING,
+                contract_path=str(rel),
+                dry_run=dry_run,
+            )
+            status = (
+                "[yellow]dry-run[/yellow]"
+                if dry_run
+                else "[green]written[/green]"
+            )
+            for p in stg_result.written:
+                table.add_row(contract_path.name, "staging", p.name, status)
+            if stg_result.new_source_columns:
+                new_cols_notices.append(
+                    (event.name, stg_result.new_source_columns)
+                )
+        except (GenerationError, IntrospectionError) as e:
+            errors.append((contract_path.name, "staging", str(e)))
+
+    console.print(table)
+
+    for event_name, cols in new_cols_notices:
+        console.print(
+            Panel(
+                f"[bold]{event_name}[/bold]: {len(cols)} source"
+                " column(s) not yet in staging contract:\n"
+                f"  {', '.join(cols)}\n\n"
+                "[dim]These flow through raw automatically. "
+                "Add to staging when consumers need them.[/dim]",
+                title="[yellow]Schema evolution detected[/yellow]",
+                border_style="yellow",
+            )
+        )
+
+    if errors:
+        for name, layer_name, msg in errors:
+            console.print(f"[red]Error[/red] [{layer_name}] {name}: {msg}")
+        raise typer.Exit(1)
+
+
+def main():
+    app()
+
+
+if __name__ == "__main__":
+    main()

modelgen/config.py

@@ -0,0 +1,62 @@
+"""
+Project-level configuration for modelgen.
+
+Reads .modelgen.yml from the repo root (or walks up from cwd to find it).
+All paths in the config file are relative to the .modelgen.yml location.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import yaml
+
+
+class ConfigError(ValueError):
+    """Raised when .modelgen.yml is malformed."""
+
+
+_DEFAULT_CONTRACTS = "contracts/*.yml"
+_DEFAULT_RAW_OUTPUT = "models/raw"
+_DEFAULT_STAGING_OUTPUT = "models/staging"
+
+
+@dataclass
+class ModelgenConfig:
+    root: Path
+    contracts_glob: str
+    raw_output: Path
+    staging_output: Path
+
+    def contract_files(self) -> list[Path]:
+        return sorted(self.root.glob(self.contracts_glob))
+
+    @classmethod
+    def _build(cls, root: Path, raw: dict) -> ModelgenConfig:
+        return cls(
+            root=root,
+            contracts_glob=raw.get("contracts", _DEFAULT_CONTRACTS),
+            raw_output=root / raw.get("raw_output", _DEFAULT_RAW_OUTPUT),
+            staging_output=root
+            / raw.get("staging_output", _DEFAULT_STAGING_OUTPUT),
+        )
+
+    @classmethod
+    def from_file(cls, path: Path) -> ModelgenConfig:
+        try:
+            raw = yaml.safe_load(path.read_text()) or {}
+        except yaml.YAMLError as e:
+            raise ConfigError(f"Failed to parse {path}: {e}") from e
+        if not isinstance(raw, dict):
+            raise ConfigError(f"{path} must be a YAML mapping.")
+        return cls._build(path.parent, raw)
+
+    @classmethod
+    def discover(cls, start: Path) -> ModelgenConfig:
+        """Find .modelgen.yml walking up from start; use defaults if absent."""
+        for directory in [start, *start.parents]:
+            candidate = directory / ".modelgen.yml"
+            if candidate.exists():
+                return cls.from_file(candidate)
+        return cls._build(start, {})

modelgen/contract.py

@@ -0,0 +1,265 @@
+"""
+Data contract schema and parsing.
+
+The contract is the source of truth. Source teams maintain it as part of
+their release process. modelgen propagates contract changes to dbt models;
+downstream FDPs see a stable interface because staging absorbs the drift.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+SUPPORTED_CONTRACT_VERSIONS = {"1"}
+
+
+class ContractError(ValueError):
+    """Raised when a contract is malformed or references invalid state."""
+
+
+@dataclass(frozen=True)
+class ColumnTest:
+    name: str
+    args: dict[str, Any] = field(default_factory=dict)
+
+    @classmethod
+    def from_raw(cls, raw: Any) -> ColumnTest:
+        if isinstance(raw, str):
+            return cls(name=raw)
+        if isinstance(raw, dict):
+            if len(raw) != 1:
+                raise ContractError(
+                    "Test entry must have exactly one key, "
+                    f"got: {list(raw.keys())}"
+                )
+            name = str(next(iter(raw)))
+            args = raw[name] or {}
+            if not isinstance(args, dict):
+                raise ContractError(
+                    f"Test args for {name!r} must be a mapping, "
+                    f"got {type(args).__name__}"
+                )
+            return cls(name=name, args=args)
+        raise ContractError(f"Unrecognised test format: {raw!r}")
+
+
+@dataclass
+class StagingColumn:
+    """One column in staging: source ref + rename + cast + tests."""
+
+    source: str  # column name in raw (== source column name)
+    name: str = ""  # column name in staging (default: source.lower())
+    type: str = ""  # BigQuery type to cast to (default: no cast)
+    description: str = ""
+    tests: list[ColumnTest] = field(default_factory=list)
+
+    def __post_init__(self):
+        if not self.source:
+            raise ContractError("Staging column must declare a 'source'.")
+        if not self.name:
+            object.__setattr__(self, "name", self.source.lower())
+
+    @property
+    def needs_cast(self) -> bool:
+        return bool(self.type)
+
+    @property
+    def needs_rename(self) -> bool:
+        return self.source != self.name
+
+    @classmethod
+    def from_raw(cls, raw: dict) -> StagingColumn:
+        if "source" not in raw:
+            raise ContractError(f"Staging column missing 'source': {raw}")
+        return cls(
+            source=raw["source"],
+            name=raw.get("name", ""),
+            type=raw.get("type", ""),
+            description=raw.get("description", ""),
+            tests=[ColumnTest.from_raw(t) for t in raw.get("tests", [])],
+        )
+
+
+@dataclass(frozen=True)
+class SourceRef:
+    """Fully-qualified BigQuery table reference."""
+
+    project: str
+    dataset: str
+    table: str
+
+    def __post_init__(self):
+        for label, value in [
+            ("project", self.project),
+            ("dataset", self.dataset),
+            ("table", self.table),
+        ]:
+            if not value or not isinstance(value, str):
+                raise ContractError(
+                    f"source.{label} must be a non-empty string."
+                )
+
+    @property
+    def fqn(self) -> str:
+        return f"{self.project}.{self.dataset}.{self.table}"
+
+    @classmethod
+    def from_raw(cls, raw: dict) -> SourceRef:
+        if not raw:
+            raise ContractError("Contract missing 'source' block.")
+        return cls(
+            project=raw["project"],
+            dataset=raw["dataset"],
+            table=raw["table"],
+        )
+
+
+@dataclass(frozen=True)
+class PartitionConfig:
+    """BigQuery partition spec. Maps directly to dbt's partition_by config."""
+
+    field: str
+    data_type: str = "timestamp"
+    granularity: str = "day"
+
+    @classmethod
+    def from_raw(cls, raw: dict | None) -> PartitionConfig | None:
+        if raw is None:
+            return None
+        if "field" not in raw:
+            raise ContractError("partition_by must specify a 'field'.")
+        return cls(
+            field=raw["field"],
+            data_type=raw.get("data_type", "timestamp"),
+            granularity=raw.get("granularity", "day"),
+        )
+
+
+@dataclass
+class RawConfig:
+    dataset: str = "raw"
+    incremental_strategy: str = "merge"
+    dedup: bool = False
+    partition_by: PartitionConfig | None = None
+    cluster_by: list[str] = field(default_factory=list)
+
+    @classmethod
+    def from_raw(cls, raw: dict | None) -> RawConfig:
+        raw = raw or {}
+        cluster_by = raw.get("cluster_by", [])
+        if isinstance(cluster_by, str):
+            cluster_by = [cluster_by]
+        return cls(
+            dataset=raw.get("dataset", "raw"),
+            incremental_strategy=raw.get("incremental_strategy", "merge"),
+            dedup=raw.get("dedup", False),
+            partition_by=PartitionConfig.from_raw(raw.get("partition_by")),
+            cluster_by=cluster_by,
+        )
+
+
+@dataclass
+class StagingConfig:
+    dataset: str = "staging"
+    incremental_strategy: str = "merge"
+    columns: list[StagingColumn] = field(default_factory=list)
+    partition_by: PartitionConfig | None = None
+    cluster_by: list[str] = field(default_factory=list)
+
+    @classmethod
+    def from_raw(cls, raw: dict | None) -> StagingConfig:
+        raw = raw or {}
+        cluster_by = raw.get("cluster_by", [])
+        if isinstance(cluster_by, str):
+            cluster_by = [cluster_by]
+        return cls(
+            dataset=raw.get("dataset", "staging"),
+            incremental_strategy=raw.get("incremental_strategy", "merge"),
+            columns=[
+                StagingColumn.from_raw(c) for c in raw.get("columns", [])
+            ],
+            partition_by=PartitionConfig.from_raw(raw.get("partition_by")),
+            cluster_by=cluster_by,
+        )
+
+
+def _normalise_unique_key(raw: Any) -> list[str]:
+    if isinstance(raw, str):
+        return [raw]
+    if isinstance(raw, list) and all(isinstance(k, str) for k in raw):
+        if not raw:
+            raise ContractError("unique_key list cannot be empty.")
+        return [str(k) for k in raw]
+    raise ContractError(
+        f"unique_key must be a string or list of strings, got {raw!r}"
+    )
+
+
+@dataclass
+class EventContract:
+    name: str
+    description: str
+    unique_key: list[str]
+    loaded_at_field: str
+    source: SourceRef
+    raw: RawConfig
+    staging: StagingConfig
+
+    @property
+    def raw_model_name(self) -> str:
+        return f"raw__{self.name}"
+
+    @property
+    def staging_model_name(self) -> str:
+        return f"stg__{self.name}"
+
+    @property
+    def unique_key_dbt(self) -> str:
+        """Render unique_key in the form dbt expects: a string for single key,
+        a list literal for composite."""
+        if len(self.unique_key) == 1:
+            return f"'{self.unique_key[0]}'"
+        return "[" + ", ".join(f"'{k}'" for k in self.unique_key) + "]"
+
+    @classmethod
+    def from_file(cls, path: Path) -> EventContract:
+        try:
+            raw_yaml = yaml.safe_load(path.read_text())
+        except yaml.YAMLError as e:
+            raise ContractError(f"Failed to parse YAML at {path}: {e}") from e
+
+        if not isinstance(raw_yaml, dict):
+            raise ContractError(
+                "Contract root must be a mapping, "
+                f"got {type(raw_yaml).__name__}"
+            )
+
+        version = str(raw_yaml.get("version", ""))
+        if version not in SUPPORTED_CONTRACT_VERSIONS:
+            raise ContractError(
+                f"Unsupported contract version {version!r}. "
+                f"Supported: {sorted(SUPPORTED_CONTRACT_VERSIONS)}"
+            )
+
+        event = raw_yaml.get("event")
+        if not isinstance(event, dict):
+            raise ContractError("Contract missing 'event' block.")
+
+        required = ["name", "unique_key", "loaded_at_field", "source"]
+        missing = [k for k in required if k not in event]
+        if missing:
+            raise ContractError(f"event missing required fields: {missing}")
+
+        return cls(
+            name=event["name"],
+            description=event.get("description", ""),
+            unique_key=_normalise_unique_key(event["unique_key"]),
+            loaded_at_field=event["loaded_at_field"],
+            source=SourceRef.from_raw(event["source"]),
+            raw=RawConfig.from_raw(event.get("raw")),
+            staging=StagingConfig.from_raw(event.get("staging")),
+        )