152 lines
6.0 KiB
Python
152 lines
6.0 KiB
Python
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:
|
|
if physical_file is None and cfg.paths.catalog_metadata_snapshot is not None:
|
|
from tht.mschema.catalog_snapshot import (
|
|
CatalogSnapshotError,
|
|
load_catalog_metadata_snapshot,
|
|
)
|
|
|
|
try:
|
|
snapshot = load_catalog_metadata_snapshot(
|
|
cfg.paths.catalog_metadata_snapshot, cfg._workspace_id
|
|
)
|
|
physical, annotations, relationships = snapshot.to_schema_inputs()
|
|
except CatalogSnapshotError as exc:
|
|
raise SchemaContextError(str(exc)) from exc
|
|
return SchemaContext(
|
|
physical=physical,
|
|
annotations=annotations,
|
|
effective_relationships=relationships,
|
|
)
|
|
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),
|
|
)
|