ThothII/harness/tht/mschema/context.py

import json
from dataclasses import dataclass
from pathlib import Path
from typing import Annotated, Literal

from pydantic import BaseModel, Field, ValidationError, model_validator

from tht.mschema.models import Annotations, ForeignKey, PhysicalSchema


class SchemaContextError(ValueError):
    """The effective schema inputs cannot be used safely."""


RelationshipName = Annotated[str, Field(min_length=1)]


class EffectiveRelationship(BaseModel):
    source_table: RelationshipName = Field(alias="sourceTable")
    source_columns: list[RelationshipName] = Field(alias="sourceColumns", min_length=1)
    target_table: RelationshipName = Field(alias="targetTable")
    target_columns: list[RelationshipName] = Field(alias="targetColumns", min_length=1)
    origin: Literal["physical", "generated", "manual"]

    model_config = {"extra": "forbid", "populate_by_name": True}

    @model_validator(mode="after")
    def columns_are_paired(self):
        if len(self.source_columns) != len(self.target_columns):
            raise ValueError("sourceColumns and targetColumns must have the same length")
        return self


class EffectiveRelationshipSnapshot(BaseModel):
    schema_version: Literal[1] = Field(alias="schemaVersion")
    workspace_id: RelationshipName = Field(alias="workspaceId")
    relationships: list[EffectiveRelationship]

    model_config = {"extra": "forbid", "populate_by_name": True}


@dataclass(frozen=True)
class SchemaContext:
    physical: PhysicalSchema
    annotations: Annotations
    # None means legacy physical + annotation FK merging. A dict, including an
    # empty one, means the catalog snapshot is the exclusive relationship source.
    effective_relationships: dict[str, list[ForeignKey]] | None


def physical_path(cfg) -> Path:
    from tht.jobs.dwh_pipeline import resolve_dwh_snapshot

    if not (cfg.paths.artifacts.parent / ".tht-dwh").exists():
        return cfg.paths.artifacts / "mschema" / "physical.yaml"
    return resolve_dwh_snapshot(cfg).physical


def annotations_path(cfg) -> Path:
    if cfg.paths.annotations_root is not None:
        return cfg.paths.annotations_root / "mschema" / "annotations.yaml"
    return cfg.paths.artifacts / "mschema" / "annotations.yaml"


def _load_effective_relationships(cfg, physical: PhysicalSchema) -> dict[str, list[ForeignKey]] | None:
    path = cfg.paths.effective_relationships
    if path is None:
        return None
    if not path.is_file():
        raise SchemaContextError(f"effective relationship snapshot is missing: {path}")
    try:
        raw = json.loads(path.read_text())
        snapshot = EffectiveRelationshipSnapshot.model_validate(raw)
    except (OSError, json.JSONDecodeError, ValidationError) as exc:
        raise SchemaContextError(f"effective relationship snapshot is invalid: {path}") from exc
    if snapshot.workspace_id != cfg._workspace_id:
        raise SchemaContextError(
            "effective relationship snapshot workspace does not match runtime workspace"
        )

    by_table: dict[str, list[ForeignKey]] = {}
    seen: set[tuple[str, tuple[str, ...], str, tuple[str, ...]]] = set()
    for relationship in snapshot.relationships:
        source = physical.tables.get(relationship.source_table)
        target = physical.tables.get(relationship.target_table)
        if source is None or target is None:
            raise SchemaContextError(
                "effective relationship endpoint table is absent from physical schema: "
                f"{relationship.source_table}->{relationship.target_table}"
            )
        missing_source = [name for name in relationship.source_columns if name not in source.columns]
        missing_target = [name for name in relationship.target_columns if name not in target.columns]
        if missing_source or missing_target:
            missing = ", ".join(
                [f"{relationship.source_table}.{name}" for name in missing_source]
                + [f"{relationship.target_table}.{name}" for name in missing_target]
            )
            raise SchemaContextError(
                f"effective relationship endpoint column is absent from physical schema: {missing}"
            )
        key = (
            relationship.source_table,
            tuple(relationship.source_columns),
            relationship.target_table,
            tuple(relationship.target_columns),
        )
        if key in seen:
            continue
        seen.add(key)
        by_table.setdefault(relationship.source_table, []).append(
            ForeignKey(
                columns=relationship.source_columns,
                ref_table=relationship.target_table,
                ref_columns=relationship.target_columns,
            )
        )
    return by_table


def load_schema_context(cfg, *, physical_file: Path | None = None) -> SchemaContext:
    physical_file = physical_file or physical_path(cfg)
    if not physical_file.is_file():
        raise SchemaContextError(f"physical schema is missing: {physical_file}")
    try:
        physical = PhysicalSchema.from_yaml(physical_file)
        annotations = Annotations.from_yaml(annotations_path(cfg))
    except (OSError, ValidationError, ValueError) as exc:
        raise SchemaContextError("physical schema or annotations are invalid") from exc
    return SchemaContext(
        physical=physical,
        annotations=annotations,
        effective_relationships=_load_effective_relationships(cfg, physical),
    )