"""Define dbt schema types and merging logic.
We generate dbt schema.yml files by translating our metadata into schema.yml
format, then applying human-sourced patches to the auto-generated schemas.
"""
import re
import textwrap
from collections.abc import Callable
from pathlib import Path
from typing import Annotated, Any
import yaml
from pydantic import BaseModel, BeforeValidator, ConfigDict
from pudl.metadata.classes import PUDL_PACKAGE, Resource
[docs]
_DESCRIPTION_WRAP_WIDTH = 88
[docs]
def _normalize_whitespace(text: str) -> str:
"""Collapse all whitespace (including blank lines) to single spaces."""
return re.sub(r"\s+", " ", text).strip()
[docs]
def _normalize_descriptions(obj: Any) -> Any:
"""Recursively collapse whitespace in every ``description`` field.
Normalizing whitespace at parse time reduces spurious diffs and round trip errors,
treating all whitespace as semantically identical.
The ``description`` fields show up in two places: the typed ``description`` field on
``DbtColumn``/``DbtTable``/``DbtSource``, and nested inside arbitrary ``data_tests``
entries which are untyped ``list`` content that pydantic doesn't otherwise inspect.
This function is used as a validator for both.
"""
if isinstance(obj, dict):
return {
key: (
_normalize_whitespace(value)
if key == "description" and isinstance(value, str)
else _normalize_descriptions(value)
)
for key, value in obj.items()
}
if isinstance(obj, list):
return [_normalize_descriptions(item) for item in obj]
return obj
[docs]
def _normalize_description_field(value: str | None) -> str | None:
return value if value is None else _normalize_whitespace(value)
[docs]
def _normalize_data_tests_field(value: list | None) -> list | None:
return value if value is None else _normalize_descriptions(value)
[docs]
_NormalizedDescription = Annotated[
str | None, BeforeValidator(_normalize_description_field)
]
[docs]
_NormalizedDataTests = Annotated[
list | None, BeforeValidator(_normalize_data_tests_field)
]
[docs]
class _LiteralStr(str):
"""Marker subclass telling the dumper to use YAML literal block style (``|``).
Only used for ``description`` fields that we've re-wrapped ourselves, so
they read as human-friendly paragraphs on disk instead of one giant line.
Everything else (regexes, SQL snippets, argument lists) is left completely
alone, since forcing a global line width in the dumper risks reflowing
content where whitespace is significant.
"""
[docs]
def _wrap_description(text: str) -> str | _LiteralStr:
"""Re-wrap a description string to short lines, for readability on disk.
A description may already contain embedded newlines (e.g. from a blank
line in a folded YAML block scalar, or from a previous pass of this same
function). We preserve that line structure and only wrap the text
*within* each line, so we don't invent new paragraph breaks the human
didn't write. We use block style whenever the result spans multiple
lines -- including when wrapping made no change to an already-wrapped
multi-line value -- since a bare multi-line ``str`` would otherwise fall
back to an ugly quoted flow scalar. Single-line results are left as a
plain string so short descriptions keep their current compact
``description: ...`` formatting.
"""
text = text.strip()
wrapped = "\n".join(
textwrap.fill(
line.strip(), width=_DESCRIPTION_WRAP_WIDTH, break_on_hyphens=False
)
if line.strip()
else ""
for line in text.split("\n")
)
if "\n" not in wrapped:
return wrapped
return _LiteralStr(wrapped)
[docs]
def _wrap_descriptions(obj: Any) -> Any:
"""Recursively re-wrap every ``description`` field in a dumped schema dict."""
if isinstance(obj, dict):
return {
key: (
_wrap_description(value)
if key == "description" and isinstance(value, str)
else _wrap_descriptions(value)
)
for key, value in obj.items()
}
if isinstance(obj, list):
return [_wrap_descriptions(item) for item in obj]
return obj
[docs]
def _prettier_yaml_dumps(yaml_contents: dict[str, Any]) -> str:
"""Dump YAML to string that Prettier likes."""
class PrettierCompatibleDumper(yaml.Dumper):
"""Custom Dumper that indents lists like prettier does.
Default dumper behavior::
foo:
- name: a
- name: b
prettier behavior::
foo:
- name: a
- name: b
NOTE (2025-07-10): this code was generated by LLM and Does The
Right Thing. It appears to do so by forcing the ``indentless`` parameter to
always be ``False``. If at any point this stops working we should convert
to using ``ruamel`` which has much more customizable formatting with actual
documentation.
"""
def increase_indent(self, flow=False, indentless=False):
return super().increase_indent(flow, False)
def choose_scalar_style(self):
"""Prefer double quotes over single quotes when quoting is required.
PyYAML's default emitter picks single-quoted style (with ``''`` escaping for
embedded apostrophes) whenever a scalar can't be written plain. Prettier
always uses double-quoted style in that case, so without this override every
such string gets rewritten by ``prek run prettier`` right after we generate
it. This only changes which *quote character* gets used -- it doesn't affect
whether a scalar is quoted at all, so plain-safe strings are still emitted
unquoted exactly as before.
"""
style = super().choose_scalar_style()
return '"' if style == "'" else style
PrettierCompatibleDumper.add_representer(
_LiteralStr,
lambda dumper, data: dumper.represent_scalar(
"tag:yaml.org,2002:str", data, style="|"
),
)
return yaml.dump(
_wrap_descriptions(yaml_contents),
default_flow_style=False,
Dumper=PrettierCompatibleDumper,
indent=2,
sort_keys=False,
width=float("inf"),
)
[docs]
def _foreign_key_data_tests(resource: Resource) -> list[dict] | None:
"""Build ``foreign_key`` data test entries for a resource's outgoing FKs.
One entry per foreign key relationship declared on the resource, in declaration
order, so regenerating a table's schema.yml produces a stable diff.
"""
data_tests = [
{
"foreign_key": {
"arguments": {
"fk_column_names": list(fk.fields),
"pk_table_name": f"source('pudl', '{fk.reference.resource}')",
"pk_column_names": list(fk.reference.fields),
}
}
}
for fk in resource.schema.foreign_keys
]
return data_tests or None
[docs]
class DbtColumn(BaseModel):
"""Define yaml structure of a dbt column."""
# Reject unrecognized keys (e.g. a stray `tests:` instead of
# `data_tests:`) at parse time instead of silently dropping them --
# pydantic's default `extra="ignore"` would otherwise make such content
# vanish from schema.human.yml overrides with no error at all.
[docs]
model_config = ConfigDict(extra="forbid")
[docs]
description: _NormalizedDescription = None
[docs]
data_tests: _NormalizedDataTests = None
[docs]
class DbtTable(BaseModel):
"""Define yaml structure of a dbt table."""
[docs]
model_config = ConfigDict(extra="forbid")
[docs]
description: _NormalizedDescription = None
[docs]
data_tests: _NormalizedDataTests = None
[docs]
columns: list[DbtColumn] | None = None
[docs]
config: dict | None = None # only for models
@classmethod
[docs]
def from_table_name(cls, table_name: str) -> "DbtTable":
"""Construct configuration defining table from PUDL metadata."""
resource = PUDL_PACKAGE.get_resource(table_name)
return cls(
name=table_name,
data_tests=_foreign_key_data_tests(resource),
columns=[DbtColumn(name=f.name) for f in resource.schema.fields],
)
[docs]
class DbtSource(BaseModel):
"""Define basic dbt yml structure to add a pudl table as a dbt source."""
[docs]
model_config = ConfigDict(extra="forbid")
[docs]
tables: list[DbtTable] | None = None
[docs]
description: _NormalizedDescription = None
[docs]
class DbtSchema(BaseModel):
"""Define basic structure of a dbt models yaml file."""
[docs]
model_config = ConfigDict(extra="forbid")
[docs]
sources: list[DbtSource] | None = None
[docs]
models: list[DbtTable] | None = None
@classmethod
[docs]
def from_table_name(cls, table_name: str) -> "DbtSchema":
"""Construct configuration defining table from PUDL metadata."""
return cls(
sources=[
DbtSource(
tables=[DbtTable.from_table_name(table_name)],
)
],
)
@classmethod
[docs]
def from_yaml(cls, schema_path: Path) -> "DbtSchema":
"""Load a DbtSchema object from a YAML file."""
with schema_path.open("r") as schema_yaml:
return cls.model_validate(yaml.safe_load(schema_yaml))
[docs]
def to_yaml(self, schema_path: Path):
"""Write DbtSchema object to YAML file."""
with schema_path.open("w") as schema_file:
yaml_output = _prettier_yaml_dumps(self.model_dump(exclude_none=True))
schema_file.write(yaml_output)
[docs]
def validate_humanity(self):
"""Make sure the human schema matches expectations.
We expect that all human overrides on source tables are data tests or column-level data tests.
We allow the 'name' field so we can match human tables/columns with machine ones.
We do not have any expectations about model definitions since those are human-only.
"""
def enforce_allowlist(
model: DbtSource | DbtTable | DbtColumn,
allowlist: set[str],
model_name: str,
) -> None:
"""Assert that all keys defined on this model are expressly allowed."""
existing_keys = set(model.model_dump(exclude_defaults=True).keys())
invalid_keys = existing_keys - allowlist
assert len(invalid_keys) == 0, (
f"Found {invalid_keys=} in human {model_name}"
)
for source in self.sources or []:
enforce_allowlist(
source, allowlist={"name", "tables"}, model_name=f"source:{source.name}"
)
for table in source.tables or []:
enforce_allowlist(
table,
allowlist={"name", "data_tests", "columns"},
model_name=f"source:{source.name}.{table.name}",
)
for column in table.columns or []:
enforce_allowlist(
column,
allowlist={"name", "data_tests"},
model_name=f"source:{source.name}.{table.name}.{column.name}",
)
[docs]
def merge_schema(machine_schema: DbtSchema, human_schema: DbtSchema) -> DbtSchema:
"""Merge two DbtSchemas by applying human-schema as a patch on top of machine-schema.
Empty merged sources will be stored in the DbtSchema model as None to avoid serializing them.
"""
human_schema.validate_humanity()
merged_sources = merge_sources_by_name(
machine_schema.sources or [], human_schema.sources or []
)
# NOTE 2026-05-11: all models are human-generated.
return DbtSchema(sources=merged_sources or None, models=human_schema.models)
[docs]
def merge_by_name(
machine_elements: list,
human_elements: list,
merger: Callable,
element_factory: Callable,
) -> list:
"""Perform a generic merge of two lists of dbt elements, matching by name.
Args:
machine_elements: can be empty list.
human_elements: can be empty list.
merger: callable that takes two elements of the same dbt type (source, table,
column) and returns a new element that is the merged version.
element_factory: callable that takes the element name and returns an empty instance - used if e.g. the human element doesn't exist.
"""
human_elements_by_name = {element.name: element for element in human_elements}
machine_names = {element.name for element in machine_elements}
if extras := (set(human_elements_by_name.keys()) - machine_names):
raise KeyError(
f"Unmatched {element_factory.__name__} keys not found in machine schema: {extras}"
)
return [
merger(
machine_element,
human_elements_by_name.get(
machine_element.name, element_factory(name=machine_element.name)
),
)
for machine_element in machine_elements
]
[docs]
def merge_sources_by_name(
machine_sources: list[DbtSource], human_sources: list[DbtSource]
) -> list[DbtSource]:
"""Match machine/human sources by name, then merge them."""
return merge_by_name(machine_sources, human_sources, merge_source, DbtSource)
[docs]
def merge_source(machine_source: DbtSource, human_source: DbtSource) -> DbtSource:
"""Merge two DbtSources by applying human-source as a patch on top of machine-source.
Returns a deep copy of the machine source to avoid aliasing,
updating with tables as the merge of the tables of the machine and human sources.
"""
return machine_source.model_copy(
deep=True,
update={
"tables": merge_tables_by_name(
machine_source.tables or [], human_source.tables or []
)
},
)
[docs]
def merge_tables_by_name(
machine_tables: list[DbtTable], human_tables: list[DbtTable]
) -> list[DbtTable]:
"""Match machine/human tables by name, then merge them."""
return merge_by_name(machine_tables, human_tables, merge_table, DbtTable)
[docs]
def merge_table(machine_table: DbtTable, human_table: DbtTable) -> DbtTable:
"""Merge two DbtTables by applying human-table as a patch on top of machine-table.
Returns a deep copy of the machine table to avoid aliasing,
updating with columns and table-level data tests as the merge of the respective machine and human data.
"""
merged_data_tests = (machine_table.data_tests or []) + (
human_table.data_tests or []
)
return machine_table.model_copy(
deep=True,
update={
"data_tests": merged_data_tests or None,
"columns": merge_columns_by_name(
machine_table.columns or [], human_table.columns or []
),
},
)
[docs]
def merge_columns_by_name(
machine_columns: list[DbtColumn], human_columns: list[DbtColumn]
) -> list[DbtColumn]:
"""Match machine/human columns by name, then merge them."""
return merge_by_name(machine_columns, human_columns, merge_column, DbtColumn)
[docs]
def merge_column(machine_column: DbtColumn, human_column: DbtColumn) -> DbtColumn:
"""Merge two DbtColumns by applying human-column as a patch on top of machine-column.
Returns a deep copy of the machine column to avoid aliasing,
updating with data tests as the merge of the data tests of the machine and human columns.
Does **not** update any other attributes (descriptions, etc.).
"""
merged_data_tests = (machine_column.data_tests or []) + (
human_column.data_tests or []
)
return machine_column.model_copy(
deep=True,
update={"data_tests": merged_data_tests or None},
)
