diff --git a/changelog.d/update-policyengine-bundle-6.12.0.changed.md b/changelog.d/update-policyengine-bundle-6.12.0.changed.md new file mode 100644 index 000000000..cc85c185c --- /dev/null +++ b/changelog.d/update-policyengine-bundle-6.12.0.changed.md @@ -0,0 +1 @@ +Update the PolicyEngine bundle to 6.12.0. diff --git a/changelog.d/v2plan-1f-115-catalog-metadata-fields.added.md b/changelog.d/v2plan-1f-115-catalog-metadata-fields.added.md new file mode 100644 index 000000000..d4355c067 --- /dev/null +++ b/changelog.d/v2plan-1f-115-catalog-metadata-fields.added.md @@ -0,0 +1 @@ +Catalog format 2 stores the metadata the app's v3 workspace reads, taken only from PolicyEngine.py's public records: each model version's entities and basic inputs; each variable's input and hidden-input flags, module, the parameter path its adds or subtracts names, its enum values' labels, its unit and its definition period; each parameter's economy and household flags. Bracket schedules are stored as parameter nodes with their stored label. A generated revision adds the columns as nullable, so catalog versions published earlier keep them empty and are compared without them when republished. An opt-in parity test, which CI does not run, compares every new field for the US and the UK with /{country}/metadata. Each country's default dataset now comes from the installed PolicyEngine.py release's manifest; the Public API no longer fixes its own (populace_us_2024 and microcosm_uk_2024_25), and publication stops if the release names no default or gives the national region another dataset. diff --git a/docker/Dockerfile b/docker/Dockerfile index 303243995..09513f8cc 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,4 +1,4 @@ FROM python:3.12 # Match the API bundle in pyproject.toml; the bundle updater changes both pins. # The models extra pins the complete certified runtime package set. -RUN pip install "policyengine[models]==6.11.1" ipython +RUN pip install "policyengine[models]==6.12.0" ipython diff --git a/docs/migration/stage-12a-user-login.md b/docs/migration/stage-12a-user-login.md index c01657bc3..3a34626ab 100644 --- a/docs/migration/stage-12a-user-login.md +++ b/docs/migration/stage-12a-user-login.md @@ -314,7 +314,7 @@ from the final SQLModel metadata. It has no post-generation corrections. "Rollback"). The models drop the `user_associations` relationships from v2 simulations and -reports. `V2_ALEMBIC_HEAD_REVISION` is `fbb7488ebe8a`. +reports. `V2_ALEMBIC_HEAD_REVISION` was `fbb7488ebe8a` when Stage 12A shipped. ### Reshape safety diff --git a/docs/migration/stage-9-v2-metadata.md b/docs/migration/stage-9-v2-metadata.md index 243826d98..aa035d6d0 100644 --- a/docs/migration/stage-9-v2-metadata.md +++ b/docs/migration/stage-9-v2-metadata.md @@ -93,6 +93,36 @@ metadata time-period options are stored on the model-version snapshot. PolicyEngine Core and country-package versions are derived diagnostic evidence and are not separately persisted. +Catalog format 2 adds the metadata the app's v3 workspace reads, taken only +from PolicyEngine.py's public records, which carry it from the release that +publishes entities, basic inputs and these flags: + +| Record | Fields | +| --- | --- | +| Model version | `catalog_format` (2), `entities` (`key`, `plural`, `is_person`, in the package's order), `basic_inputs` (the household builder's default fields) | +| Variable | `is_input_variable`, `hidden_input`, `module_name`, `adds_parameter_path` and `subtracts_parameter_path` (the parameter whose dated value lists the variables summed; `adds` and `subtracts` keep the list PolicyEngine.py resolved), `enum_values` (`value` and declared `label`, in order), `unit` (for example `currency-USD`) and `definition_period` (`year`, `month` or `eternity`), which the household builder and breakdown read | +| Parameter | `economy`, `household` | + +Bracket schedules are parameter nodes with their stored label, or none when the +package stores none. An empty basic-input list is stored as empty. Extraction +stops before reading any record if the installed PolicyEngine.py does not +publish these fields, and names each missing field: 6.7.5 and 6.11.1 +publish only the variables' unit and definition period. It also stops if the +model has other than one person entity, names a basic input or a parameter path +that is not published, gives enum labels that do not match the variable's +possible values or are not JSON scalars, or gives a flag that is not a boolean. +An opt-in test, `tests/integration/test_v2_catalog_metadata_parity.py`, +compares every format 2 field for the US and the UK with what +`/{country}/metadata` builds. CI does not run it, because it extracts every +variable and parameter of the installed models, with their values, and builds +each country's v1 metadata; set `RUN_V2_CATALOG_METADATA_PARITY=1` to run it. + +Publication never rewrites a stored version. A version published before +format 2 has a null `catalog_format` and null format 2 fields, and later +publications of the same PolicyEngine.py version compare it without those +fields. The first PolicyEngine.py version published after format 2 deploys +carries them. + Parameter extraction requires the PolicyEngine.py public history contract from 5.0.4 and later: values are ordered from oldest to newest, bounded `end_date` values are inclusive and fall one day before the following `start_date`, and @@ -101,20 +131,24 @@ does not reorder, deduplicate, or repair package output. A value superseded on the immediately following day therefore has equal start and end dates and is a valid one-day interval. -The US national region uses `populace_us_2024`. A US state, +Each country's default dataset is the one the installed PolicyEngine.py +release names in its manifest (`default_dataset`); the Public API sets none of +its own. Publication stops if the manifest names no default, or if the release +gives the national region another dataset. The US national region uses the +release's default (`populace_us_2024` in 6.11.1). A US state, congressional-district, or place region uses its PolicyEngine.py regional alternative when one is available. Until a later reviewed change removes the -fallback, a subnational region without such an alternative uses -`populace_us_2024`. The command emits one summary warning containing only the -affected region types and counts. The UK national and country regions use -`microcosm_uk_2024_25`, the default certified by PolicyEngine.py 6.10.0; a -country filters that national file. UK constituency and local-authority -regions use the dataset PolicyEngine.py routes them to, -`microcosm_uk_2024_25_dense`, because the national file has no area codes. +fallback, a subnational region without such an alternative uses the default. +The command emits one summary warning containing only the affected region +types and counts. The UK national and country regions use the release's +default (`microcosm_uk_2024_25` in 6.11.1); a country filters that national +file. UK constituency and local-authority regions use the dataset +PolicyEngine.py routes them to, `microcosm_uk_2024_25_dense`, because the +national file has no area codes. Review the fallback summary on every new PolicyEngine.py release. A changed -count may identify an upstream catalog change that needs a reviewed dataset -selection even when publication otherwise succeeds. +count may identify an upstream change to the release's datasets even when +publication otherwise succeeds. ## Validation evidence and retry diff --git a/migrations/v2/versions/f5c1d9b4685b_store_catalog_format_2_metadata.py b/migrations/v2/versions/f5c1d9b4685b_store_catalog_format_2_metadata.py new file mode 100644 index 000000000..48356a9a2 --- /dev/null +++ b/migrations/v2/versions/f5c1d9b4685b_store_catalog_format_2_metadata.py @@ -0,0 +1,94 @@ +"""store catalog format 2 metadata + +Revision ID: f5c1d9b4685b +Revises: fb1a3e40d3d6 +Create Date: 2026-10-26 20:13:40.733924 +Generation: uv run alembic -c alembic-v2.ini revision --autogenerate +""" + +from typing import Sequence, Union + +from alembic import op +import sqlalchemy as sa +import sqlmodel + + +revision: str = "f5c1d9b4685b" +down_revision: Union[str, None] = "fb1a3e40d3d6" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.add_column("parameters", sa.Column("economy", sa.Boolean(), nullable=True)) + op.add_column("parameters", sa.Column("household", sa.Boolean(), nullable=True)) + op.add_column( + "tax_benefit_model_versions", + sa.Column("catalog_format", sa.Integer(), nullable=True), + ) + op.add_column( + "tax_benefit_model_versions", sa.Column("entities", sa.JSON(), nullable=True) + ) + op.add_column( + "tax_benefit_model_versions", + sa.Column("basic_inputs", sa.JSON(), nullable=True), + ) + op.add_column( + "variables", sa.Column("is_input_variable", sa.Boolean(), nullable=True) + ) + op.add_column("variables", sa.Column("hidden_input", sa.Boolean(), nullable=True)) + op.add_column( + "variables", + sa.Column( + "module_name", sqlmodel.sql.sqltypes.AutoString(length=512), nullable=True + ), + ) + op.add_column( + "variables", + sa.Column( + "adds_parameter_path", + sqlmodel.sql.sqltypes.AutoString(length=512), + nullable=True, + ), + ) + op.add_column( + "variables", + sa.Column( + "subtracts_parameter_path", + sqlmodel.sql.sqltypes.AutoString(length=512), + nullable=True, + ), + ) + op.add_column("variables", sa.Column("enum_values", sa.JSON(), nullable=True)) + op.add_column( + "variables", + sa.Column("unit", sqlmodel.sql.sqltypes.AutoString(length=128), nullable=True), + ) + op.add_column( + "variables", + sa.Column( + "definition_period", + sqlmodel.sql.sqltypes.AutoString(length=32), + nullable=True, + ), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_column("variables", "definition_period") + op.drop_column("variables", "unit") + op.drop_column("variables", "enum_values") + op.drop_column("variables", "subtracts_parameter_path") + op.drop_column("variables", "adds_parameter_path") + op.drop_column("variables", "module_name") + op.drop_column("variables", "hidden_input") + op.drop_column("variables", "is_input_variable") + op.drop_column("tax_benefit_model_versions", "basic_inputs") + op.drop_column("tax_benefit_model_versions", "entities") + op.drop_column("tax_benefit_model_versions", "catalog_format") + op.drop_column("parameters", "household") + op.drop_column("parameters", "economy") + # ### end Alembic commands ### diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py index 9aa18b61c..4971d4c03 100644 --- a/policyengine_api/data/v2/catalog/extraction.py +++ b/policyengine_api/data/v2/catalog/extraction.py @@ -2,19 +2,19 @@ from __future__ import annotations -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from importlib import metadata as importlib_metadata -from typing import Any +from typing import TYPE_CHECKING, Any from policyengine_api.data.v2.catalog.extraction_regions import ( DATASET_YEAR_PATTERN, MAX_REGION_FILTER_VALUE_LENGTH, - REVIEWED_DEFAULT_DATASETS, SUPPORTED_REGION_TYPES, dataset_year, normalize_region, + release_default_dataset, select_region_datasets, - verified_default_dataset, verified_source_regions, ) from policyengine_api.data.v2.catalog.extraction_values import ( @@ -23,17 +23,21 @@ catalog_identifier, normalize_json_value, optional_text, + required_bool, required_text, ) from policyengine_api.data.v2.catalog.extraction_variables_and_parameters import ( + VARIABLE_TEXT_LIMITS, normalize_parameter, normalize_parameter_nodes, normalize_variable, public_named_records, ) from policyengine_api.data.v2.catalog.records import ( + CATALOG_FORMAT, CountryCatalog, DatasetRecord, + EntityRecord, FallbackSummary, ModelRecord, ModelVersionRecord, @@ -41,6 +45,15 @@ ) from policyengine_api.supported_countries import SUPPORTED_COUNTRY_IDS +if TYPE_CHECKING: + from uuid import UUID + + from policyengine_api.data.v2.catalog.records import ( + ParameterNodeRecord, + ParameterRecord, + VariableRecord, + ) + __all__ = [ "CURRENT_LAW_IDS", "DATASET_YEAR_PATTERN", @@ -49,13 +62,15 @@ "METADATA_TIME_PERIODS", "PLACEHOLDER_VERSIONS", "REQUIRED_DEPENDENCIES", - "REVIEWED_DEFAULT_DATASETS", "SUPPORTED_COUNTRIES", "SUPPORTED_REGION_TYPES", + "VARIABLE_TEXT_LIMITS", "CatalogExtractionError", + "ModelMetadata", "extract_catalog", "extract_installed_catalog", "normalize_json_value", + "normalize_model_metadata", ] @@ -176,6 +191,196 @@ def _verified_release_manifest( return release_manifest +# The public fields catalog format 2 reads, by the record that carries them. +FORMAT_2_FIELDS: tuple[tuple[str, str, tuple[str, ...]], ...] = ( + ("model version", "", ("entities", "basic_inputs")), + ( + "variable", + "variables_by_name", + ( + "is_input_variable", + "hidden_input", + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "enum_values", + "unit", + "definition_period", + ), + ), + ("parameter", "parameters_by_name", ("economy", "household")), +) + + +def _require_format_2_fields(country_id: str, source: object) -> None: + """Fail first, naming each field, when PolicyEngine.py does not publish them. + + Older releases, such as 6.7.5 and 6.11.1, publish only the + variables' unit and definition period of catalog format 2's fields. + """ + missing: list[str] = [] + for record_name, mapping_attribute, field_names in FORMAT_2_FIELDS: + if mapping_attribute: + records = getattr(source, mapping_attribute, None) + holders = list(records.values()) if isinstance(records, Mapping) else [] + else: + holders = [source] + absent = [ + field_name + for field_name in field_names + if any(not hasattr(holder, field_name) for holder in holders) + ] + if absent: + missing.append(f"{record_name} {', '.join(absent)}") + if missing: + msg = ( + "the installed PolicyEngine.py does not publish catalog format 2's " + f"fields for {country_id}: {'; '.join(missing)}" + ) + raise CatalogExtractionError(msg) + + +def _source_sequence(country_id: str, source: object, attribute: str) -> Sequence: + values = getattr(source, attribute, None) + if not isinstance(values, Sequence) or isinstance(values, (str, bytes)): + msg = f"{country_id} public model {attribute} must be a sequence" + raise CatalogExtractionError(msg) + return values + + +def _normalize_entities(country_id: str, source: object) -> tuple[EntityRecord, ...]: + """Normalize the public model's entities, in its order, with one person entity.""" + source_entities = _source_sequence(country_id, source, "entities") + entities: list[EntityRecord] = [] + for source_entity in source_entities: + key = required_text( + getattr(source_entity, "key", None), + field_name=f"{country_id} entity key", + maximum=128, + ) + if any(entity.key == key for entity in entities): + msg = f"{country_id} entity key {key!r} is repeated" + raise CatalogExtractionError(msg) + entities.append( + EntityRecord( + key=key, + plural=required_text( + getattr(source_entity, "plural", None), + field_name=f"{country_id} entity {key} plural", + maximum=128, + ), + is_person=required_bool( + getattr(source_entity, "is_person", None), + field_name=f"{country_id} entity {key} is_person", + ), + ) + ) + if sum(entity.is_person for entity in entities) != 1: + msg = f"{country_id} public model must have exactly one person entity" + raise CatalogExtractionError(msg) + return tuple(entities) + + +def _normalize_basic_inputs( + country_id: str, + source: object, + *, + variable_names: frozenset[str], +) -> tuple[str, ...]: + """Normalize the country's default household fields; each is a variable. + + An empty list, which PolicyEngine.py gives a package that declares none, is + stored as empty. + """ + source_inputs = _source_sequence(country_id, source, "basic_inputs") + basic_inputs = tuple( + required_text(value, field_name=f"{country_id} basic input", maximum=512) + for value in source_inputs + ) + for name in basic_inputs: + if name not in variable_names: + msg = f"{country_id} basic input {name!r} is not a published variable" + raise CatalogExtractionError(msg) + return basic_inputs + + +def _verify_parameter_paths( + variables: tuple[VariableRecord, ...], + *, + parameter_names: frozenset[str], +) -> None: + """Require each parameter named by ``adds`` or ``subtracts`` to be published.""" + for variable in variables: + for field_name, path in ( + ("adds_parameter_path", variable.adds_parameter_path), + ("subtracts_parameter_path", variable.subtracts_parameter_path), + ): + if path is not None and path not in parameter_names: + msg = ( + f"variable {variable.name} {field_name} {path!r} is not a " + "published parameter" + ) + raise CatalogExtractionError(msg) + + +@dataclass(frozen=True, slots=True) +class ModelMetadata: + """The model metadata of one country's catalog, without regions or datasets.""" + + entities: tuple[EntityRecord, ...] + basic_inputs: tuple[str, ...] + variables: tuple[VariableRecord, ...] + parameter_nodes: tuple[ParameterNodeRecord, ...] + parameters: tuple[ParameterRecord, ...] + + +def normalize_model_metadata( + country_id: str, + source: object, + *, + model_version_id: UUID, +) -> ModelMetadata: + """Normalize a public model's entities, variables and parameter tree. + + Everything comes from PolicyEngine.py's public records. Separate from the + certified data release, so the metadata can be checked on its own. + """ + _require_format_2_fields(country_id, source) + variables = tuple( + normalize_variable(variable, model_version_id=model_version_id) + for variable in public_named_records( + source, + attribute="variables_by_name", + entity_name="variable", + ) + ) + parameters = tuple( + normalize_parameter(parameter, model_version_id=model_version_id) + for parameter in public_named_records( + source, + attribute="parameters_by_name", + entity_name="parameter", + ) + ) + _verify_parameter_paths( + variables, + parameter_names=frozenset(parameter.name for parameter in parameters), + ) + return ModelMetadata( + entities=_normalize_entities(country_id, source), + basic_inputs=_normalize_basic_inputs( + country_id, + source, + variable_names=frozenset(variable.name for variable in variables), + ), + variables=variables, + parameter_nodes=normalize_parameter_nodes( + source, model_version_id=model_version_id + ), + parameters=parameters, + ) + + def _normalize_country( *, country_id: str, @@ -196,34 +401,18 @@ def _normalize_country( expected_country_package_version=expected_country_package_version, ) - variables = tuple( - normalize_variable(variable, model_version_id=model_version_id) - for variable in public_named_records( - source, - attribute="variables_by_name", - entity_name="variable", - ) - ) - parameter_nodes = normalize_parameter_nodes( - source, model_version_id=model_version_id - ) - parameters = tuple( - normalize_parameter(parameter, model_version_id=model_version_id) - for parameter in public_named_records( - source, - attribute="parameters_by_name", - entity_name="parameter", - ) + metadata = normalize_model_metadata( + country_id, source, model_version_id=model_version_id ) - reviewed_default, default_dataset_uri = verified_default_dataset( + default_dataset, default_dataset_uri = release_default_dataset( country_id, release_manifest ) source_regions = verified_source_regions(country_id, source) dataset_selection = select_region_datasets( country_id, source_regions, - reviewed_default=reviewed_default, + default_dataset=default_dataset, default_dataset_uri=default_dataset_uri, ) @@ -272,10 +461,13 @@ def _normalize_country( ), current_law_id=CURRENT_LAW_IDS[country_id], metadata_time_periods=METADATA_TIME_PERIODS[country_id], + catalog_format=CATALOG_FORMAT, + entities=metadata.entities, + basic_inputs=metadata.basic_inputs, ), - variables=variables, - parameter_nodes=parameter_nodes, - parameters=parameters, + variables=metadata.variables, + parameter_nodes=metadata.parameter_nodes, + parameters=metadata.parameters, datasets=datasets, regions=tuple(sorted(regions, key=lambda record: record.code)), fallback_summaries=tuple( diff --git a/policyengine_api/data/v2/catalog/extraction_regions.py b/policyengine_api/data/v2/catalog/extraction_regions.py index 986976451..474712453 100644 --- a/policyengine_api/data/v2/catalog/extraction_regions.py +++ b/policyengine_api/data/v2/catalog/extraction_regions.py @@ -20,10 +20,6 @@ from collections.abc import Mapping from uuid import UUID -REVIEWED_DEFAULT_DATASETS = { - "us": "populace_us_2024", - "uk": "microcosm_uk_2024_25", -} SUPPORTED_REGION_TYPES = frozenset( { "national", @@ -102,7 +98,7 @@ class _RegionDatasetSelection(NamedTuple): """The logical input datasets a country's regions select. Names every selected dataset, each region code's dataset, and how many - regions of each type fall back to the reviewed default dataset. + regions of each type fall back to the release's default dataset. """ dataset_names: set[str] @@ -110,20 +106,20 @@ class _RegionDatasetSelection(NamedTuple): fallback_counts: Counter[str] -def verified_default_dataset( +def release_default_dataset( country_id: str, release_manifest: object, ) -> tuple[str, object]: - """Return the reviewed default dataset and the manifest's URI for it. + """Return the default dataset the installed PolicyEngine.py release names. - The manifest's default dataset must be the reviewed one. + Also returns the manifest's URI for it. The Public API sets no default of + its own. """ - reviewed_default = REVIEWED_DEFAULT_DATASETS[country_id] - manifest_default = getattr(release_manifest, "default_dataset", None) - if manifest_default != reviewed_default: - msg = f"{country_id} default dataset does not match the reviewed selection" + default_dataset = getattr(release_manifest, "default_dataset", None) + if not isinstance(default_dataset, str) or not default_dataset: + msg = f"{country_id} release manifest names no default dataset" raise CatalogExtractionError(msg) - return reviewed_default, getattr(release_manifest, "default_dataset_uri", None) + return default_dataset, getattr(release_manifest, "default_dataset_uri", None) def verified_source_regions(country_id: str, source: object) -> tuple[Any, ...]: @@ -155,7 +151,7 @@ def select_region_datasets( country_id: str, source_regions: tuple[Any, ...], *, - reviewed_default: str, + default_dataset: str, default_dataset_uri: object, ) -> _RegionDatasetSelection: """Choose each region's logical input dataset. @@ -163,11 +159,11 @@ def select_region_datasets( A subnational region with a dataset path uses that dataset: US states and districts, and the UK constituencies and local authorities that .py routes to its local-area release. Every other region, such as a UK country that - filters the national file, uses the reviewed default, and the national + filters the national file, uses the release's default, and the national region may only name it. A US subnational region without a path is counted as a fallback. """ - selected_dataset_names: set[str] = {reviewed_default} + selected_dataset_names: set[str] = {default_dataset} region_dataset_names: dict[str, str] = {} fallback_counts: Counter[str] = Counter() for region in source_regions: @@ -189,14 +185,17 @@ def select_region_datasets( dataset_name = _dataset_name_from_path(dataset_path) selected_dataset_names.add(dataset_name) else: - dataset_name = reviewed_default + dataset_name = default_dataset if country_id == "us" and region_type != "national" and not dataset_path: fallback_counts[region_type] += 1 if region_type == "national" and dataset_path not in { None, default_dataset_uri, }: - msg = f"{country_id} national region does not use the reviewed dataset" + msg = ( + f"{country_id} national region does not use the release's default " + f"dataset {default_dataset}" + ) raise CatalogExtractionError(msg) region_dataset_names[code] = dataset_name return _RegionDatasetSelection( diff --git a/policyengine_api/data/v2/catalog/extraction_values.py b/policyengine_api/data/v2/catalog/extraction_values.py index 3fb5de8c4..3bf3ea8c2 100644 --- a/policyengine_api/data/v2/catalog/extraction_values.py +++ b/policyengine_api/data/v2/catalog/extraction_values.py @@ -53,6 +53,22 @@ def optional_text( return value +def optional_bool(value: object, *, field_name: str) -> bool | None: + """Return a boolean or ``None``; raise for any other value.""" + if value is not None and not isinstance(value, bool): + msg = f"{field_name} must be a boolean or null" + raise CatalogExtractionError(msg) + return value + + +def required_bool(value: object, *, field_name: str) -> bool: + """Return a boolean; raise for any other value.""" + if not isinstance(value, bool): + msg = f"{field_name} must be a boolean" + raise CatalogExtractionError(msg) + return value + + def _normalize_json_float(value: float) -> JSONValue: """Keep a finite float, write an infinity as a string and reject NaN.""" if math.isnan(value): diff --git a/policyengine_api/data/v2/catalog/extraction_variables_and_parameters.py b/policyengine_api/data/v2/catalog/extraction_variables_and_parameters.py index f4f06d11c..1c6c0e8d8 100644 --- a/policyengine_api/data/v2/catalog/extraction_variables_and_parameters.py +++ b/policyengine_api/data/v2/catalog/extraction_variables_and_parameters.py @@ -10,20 +10,47 @@ CatalogExtractionError, catalog_identifier, normalize_json_value, + optional_bool, optional_text, + required_bool, required_text, ) from policyengine_api.data.v2.catalog.records import ( + EnumValueRecord, ParameterNodeRecord, ParameterRecord, ParameterValueRecord, VariableRecord, ) +from policyengine_api.data.v2.models import Variable if TYPE_CHECKING: from uuid import UUID +def _variable_column_length(column_name: str) -> int: + """Return the length limit of one text column of the v2 ``variables`` table.""" + length = getattr(Variable.__table__.c[column_name].type, "length", None) + if not isinstance(length, int): + msg = f"variables.{column_name} has no length limit" + raise TypeError(msg) + return length + + +# Catalog format 2's variable text fields, checked against the columns that +# store them, so the two limits cannot drift apart. +VARIABLE_TEXT_LIMITS: Mapping[str, int] = { + column_name: _variable_column_length(column_name) + for column_name in ( + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "unit", + "definition_period", + ) +} + + def _type_name(value: object) -> str | None: if value is None: return None @@ -78,6 +105,12 @@ def normalize_variable( else: normalized_possible_values = None + enum_values = _normalize_enum_values( + name, + getattr(source, "enum_values", None), + possible_values=normalized_possible_values, + ) + def optional_names(attribute: str) -> list[str] | None: values = getattr(source, attribute, None) if values is None: @@ -117,9 +150,87 @@ def optional_names(attribute: str) -> list[str] | None: default_value=normalize_json_value(getattr(source, "default_value", None)), adds=optional_names("adds"), subtracts=optional_names("subtracts"), + is_input_variable=optional_bool( + getattr(source, "is_input_variable", None), + field_name=f"variable {name} is_input_variable", + ), + hidden_input=optional_bool( + getattr(source, "hidden_input", None), + field_name=f"variable {name} hidden_input", + ), + module_name=optional_text( + getattr(source, "module_name", None), + field_name=f"variable {name} module_name", + maximum=VARIABLE_TEXT_LIMITS["module_name"], + ), + adds_parameter_path=optional_text( + getattr(source, "adds_parameter_path", None), + field_name=f"variable {name} adds_parameter_path", + maximum=VARIABLE_TEXT_LIMITS["adds_parameter_path"], + ), + subtracts_parameter_path=optional_text( + getattr(source, "subtracts_parameter_path", None), + field_name=f"variable {name} subtracts_parameter_path", + maximum=VARIABLE_TEXT_LIMITS["subtracts_parameter_path"], + ), + enum_values=enum_values, + unit=optional_text( + getattr(source, "unit", None), + field_name=f"variable {name} unit", + maximum=VARIABLE_TEXT_LIMITS["unit"], + ), + definition_period=optional_text( + getattr(source, "definition_period", None), + field_name=f"variable {name} definition_period", + maximum=VARIABLE_TEXT_LIMITS["definition_period"], + ), ) +def _normalize_enum_values( + name: str, + source_values: object, + *, + possible_values: list[str] | None, +) -> tuple[EnumValueRecord, ...] | None: + """Normalize an enum variable's value labels, which must match its values. + + Each label is kept as declared when it is a JSON scalar. + """ + if source_values is None and possible_values is None: + return None + if ( + source_values is None + or not isinstance(source_values, Sequence) + or isinstance(source_values, (str, bytes)) + ): + msg = f"variable {name} enum values do not match its possible values" + raise CatalogExtractionError(msg) + records: list[EnumValueRecord] = [] + for source_value in source_values: + value = required_text( + getattr(source_value, "value", None), + field_name=f"variable {name} enum value", + maximum=255, + ) + try: + label = normalize_json_value(getattr(source_value, "label", None)) + except CatalogExtractionError as error: + msg = ( + f"variable {name} enum label for {value} is not JSON-compatible: " + f"{error}" + ) + raise CatalogExtractionError(msg) from error + if isinstance(label, (dict, list)): + msg = f"variable {name} enum label for {value} must be a JSON scalar" + raise CatalogExtractionError(msg) + records.append(EnumValueRecord(value=value, label=label)) + if [record.value for record in records] != possible_values: + msg = f"variable {name} enum values do not match its possible values" + raise CatalogExtractionError(msg) + return tuple(records) + + def normalize_parameter( source: object, *, @@ -213,6 +324,14 @@ def normalize_parameter( maximum=128, ), values=tuple(normalized_values), + economy=required_bool( + getattr(source, "economy", None), + field_name=f"parameter {name} economy", + ), + household=required_bool( + getattr(source, "household", None), + field_name=f"parameter {name} household", + ), ) diff --git a/policyengine_api/data/v2/catalog/publication_reconciliation.py b/policyengine_api/data/v2/catalog/publication_reconciliation.py index 0977676c6..f292e4bf3 100644 --- a/policyengine_api/data/v2/catalog/publication_reconciliation.py +++ b/policyengine_api/data/v2/catalog/publication_reconciliation.py @@ -76,9 +76,45 @@ def _versioned_table_from(table: sa.Table) -> sa.Join: ) +def stored_catalog_format( + connection: Connection, country: CountryCatalog +) -> int | None: + """Return the catalog format this country's stored version was published with. + + None means it was published before catalog format 2, so its format 2 + fields are empty. + """ + statement = ( + sa.select(MODEL_VERSIONS.c.catalog_format) + .select_from( + MODEL_VERSIONS.join( + MODELS, + MODELS.c.id == MODEL_VERSIONS.c.model_id, + ) + ) + .where( + MODELS.c.name == country.model.name, + MODEL_VERSIONS.c.version == country.model_version.version, + ) + ) + return connection.execute(statement).scalar_one_or_none() + + def _comparison_pairs( country: CountryCatalog, + *, + with_format_2: bool, ) -> tuple[tuple[Select, Select], ...]: + """Return the staged and stored rows of each compared table. + + Without ``with_format_2``, catalog format 2's fields are left out: a + version published before format 2 keeps them empty, and publication never + rewrites a stored version. + """ + + def format_2(*columns: sa.ColumnElement) -> tuple[sa.ColumnElement, ...]: + return columns if with_format_2 else () + staged_country = STAGE_CATALOG_MODELS.c.country_id == country.country_id selected_version = sa.and_( MODELS.c.name == country.model.name, @@ -89,12 +125,22 @@ def _comparison_pairs( STAGE_CATALOG_MODELS.c.version_description, STAGE_CATALOG_MODELS.c.current_law_id, STAGE_CATALOG_MODELS.c.metadata_time_periods, + *format_2( + STAGE_CATALOG_MODELS.c.catalog_format, + STAGE_CATALOG_MODELS.c.entities, + STAGE_CATALOG_MODELS.c.basic_inputs, + ), ).where(staged_country) actual_model_version = ( sa.select( MODEL_VERSIONS.c.description, MODEL_VERSIONS.c.current_law_id, sa.cast(MODEL_VERSIONS.c.metadata_time_periods, JSONB), + *format_2( + MODEL_VERSIONS.c.catalog_format, + sa.cast(MODEL_VERSIONS.c.entities, JSONB), + sa.cast(MODEL_VERSIONS.c.basic_inputs, JSONB), + ), ) .select_from( MODEL_VERSIONS.join( @@ -115,6 +161,16 @@ def _comparison_pairs( STAGE_CATALOG_VARIABLES.c.default_value, STAGE_CATALOG_VARIABLES.c.adds, STAGE_CATALOG_VARIABLES.c.subtracts, + *format_2( + STAGE_CATALOG_VARIABLES.c.is_input_variable, + STAGE_CATALOG_VARIABLES.c.hidden_input, + STAGE_CATALOG_VARIABLES.c.module_name, + STAGE_CATALOG_VARIABLES.c.adds_parameter_path, + STAGE_CATALOG_VARIABLES.c.subtracts_parameter_path, + STAGE_CATALOG_VARIABLES.c.enum_values, + STAGE_CATALOG_VARIABLES.c.unit, + STAGE_CATALOG_VARIABLES.c.definition_period, + ), ).where(STAGE_CATALOG_VARIABLES.c.country_id == country.country_id) actual_variables = ( sa.select( @@ -127,6 +183,16 @@ def _comparison_pairs( sa.cast(VARIABLES.c.default_value, JSONB), sa.cast(VARIABLES.c.adds, JSONB), sa.cast(VARIABLES.c.subtracts, JSONB), + *format_2( + VARIABLES.c.is_input_variable, + VARIABLES.c.hidden_input, + VARIABLES.c.module_name, + VARIABLES.c.adds_parameter_path, + VARIABLES.c.subtracts_parameter_path, + sa.cast(VARIABLES.c.enum_values, JSONB), + VARIABLES.c.unit, + VARIABLES.c.definition_period, + ), ) .select_from(_versioned_table_from(VARIABLES)) .where(selected_version) @@ -153,6 +219,10 @@ def _comparison_pairs( STAGE_CATALOG_PARAMETERS.c.description, STAGE_CATALOG_PARAMETERS.c.data_type, STAGE_CATALOG_PARAMETERS.c.unit, + *format_2( + STAGE_CATALOG_PARAMETERS.c.economy, + STAGE_CATALOG_PARAMETERS.c.household, + ), ).where(STAGE_CATALOG_PARAMETERS.c.country_id == country.country_id) actual_parameters = ( sa.select( @@ -161,6 +231,7 @@ def _comparison_pairs( PARAMETERS.c.description, PARAMETERS.c.data_type, PARAMETERS.c.unit, + *format_2(PARAMETERS.c.economy, PARAMETERS.c.household), ) .select_from(_versioned_table_from(PARAMETERS)) .where(selected_version) @@ -288,9 +359,12 @@ def assert_country_matches( """Raise CatalogPublicationError if the stored version differs from the staged one. Rows are compared as sets in both directions, so a missing row and an extra - row both fail. + row both fail. A version published before catalog format 2 is compared + without format 2's fields. """ - for staged, actual in _comparison_pairs(country): + stored_format = stored_catalog_format(connection, country) + pairs = _comparison_pairs(country, with_format_2=stored_format is not None) + for staged, actual in pairs: if _sets_differ(connection, staged, actual): msg = f"persisted {country.country_id} catalog differs from PolicyEngine.py" raise CatalogPublicationError(msg) @@ -320,6 +394,9 @@ def assert_country_matches( "description", "current_law_id", "metadata_time_periods", + "catalog_format", + "entities", + "basic_inputs", ], sa.select( STAGE_CATALOG_MODELS.c.version_id, @@ -332,6 +409,12 @@ def assert_country_matches( STAGE_CATALOG_MODELS.c.metadata_time_periods, MODEL_VERSIONS.c.metadata_time_periods.type, ), + STAGE_CATALOG_MODELS.c.catalog_format, + sa.cast(STAGE_CATALOG_MODELS.c.entities, MODEL_VERSIONS.c.entities.type), + sa.cast( + STAGE_CATALOG_MODELS.c.basic_inputs, + MODEL_VERSIONS.c.basic_inputs.type, + ), ) .select_from( STAGE_CATALOG_MODELS.join( @@ -395,6 +478,14 @@ def assert_country_matches( "default_value", "adds", "subtracts", + "is_input_variable", + "hidden_input", + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "enum_values", + "unit", + "definition_period", "tax_benefit_model_version_id", ], sa.select( @@ -418,6 +509,17 @@ def assert_country_matches( STAGE_CATALOG_VARIABLES.c.subtracts, VARIABLES.c.subtracts.type, ), + STAGE_CATALOG_VARIABLES.c.is_input_variable, + STAGE_CATALOG_VARIABLES.c.hidden_input, + STAGE_CATALOG_VARIABLES.c.module_name, + STAGE_CATALOG_VARIABLES.c.adds_parameter_path, + STAGE_CATALOG_VARIABLES.c.subtracts_parameter_path, + sa.cast( + STAGE_CATALOG_VARIABLES.c.enum_values, + VARIABLES.c.enum_values.type, + ), + STAGE_CATALOG_VARIABLES.c.unit, + STAGE_CATALOG_VARIABLES.c.definition_period, MODEL_VERSIONS.c.id, ) .select_from( @@ -481,6 +583,8 @@ def assert_country_matches( "description", "data_type", "unit", + "economy", + "household", "tax_benefit_model_version_id", ], sa.select( @@ -491,6 +595,8 @@ def assert_country_matches( STAGE_CATALOG_PARAMETERS.c.description, STAGE_CATALOG_PARAMETERS.c.data_type, STAGE_CATALOG_PARAMETERS.c.unit, + STAGE_CATALOG_PARAMETERS.c.economy, + STAGE_CATALOG_PARAMETERS.c.household, MODEL_VERSIONS.c.id, ) .select_from( diff --git a/policyengine_api/data/v2/catalog/publication_staging.py b/policyengine_api/data/v2/catalog/publication_staging.py index 76820b27b..f4243fa64 100644 --- a/policyengine_api/data/v2/catalog/publication_staging.py +++ b/policyengine_api/data/v2/catalog/publication_staging.py @@ -16,8 +16,10 @@ from policyengine_api.data.v2.catalog.records import ( CountryCatalog, NormalizedCatalog, + VariableRecord, iter_batches, ) +from policyengine_api.data.v2.models import CatalogEntity, CatalogEnumValue if TYPE_CHECKING: from collections.abc import Callable, Iterable, Iterator, Sequence @@ -39,6 +41,9 @@ sa.Column("version_description", sa.Text), sa.Column("current_law_id", sa.Integer, nullable=False), sa.Column("metadata_time_periods", JSONB, nullable=False), + sa.Column("catalog_format", sa.Integer, nullable=False), + sa.Column("entities", JSONB, nullable=False), + sa.Column("basic_inputs", JSONB, nullable=False), prefixes=["TEMPORARY"], postgresql_on_commit="DROP", ) @@ -57,6 +62,14 @@ sa.Column("default_value", JSONB, nullable=False), sa.Column("adds", JSONB), sa.Column("subtracts", JSONB), + sa.Column("is_input_variable", sa.Boolean), + sa.Column("hidden_input", sa.Boolean), + sa.Column("module_name", sa.Text), + sa.Column("adds_parameter_path", sa.Text), + sa.Column("subtracts_parameter_path", sa.Text), + sa.Column("enum_values", JSONB), + sa.Column("unit", sa.Text), + sa.Column("definition_period", sa.Text), prefixes=["TEMPORARY"], postgresql_on_commit="DROP", ) @@ -83,6 +96,8 @@ sa.Column("description", sa.Text), sa.Column("data_type", sa.Text), sa.Column("unit", sa.Text), + sa.Column("economy", sa.Boolean, nullable=False), + sa.Column("household", sa.Boolean, nullable=False), prefixes=["TEMPORARY"], postgresql_on_commit="DROP", ) @@ -163,6 +178,30 @@ def _model_rows(country: CountryCatalog) -> Iterator[tuple[object, ...]]: country.model_version.description, country.model_version.current_law_id, Jsonb(country.model_version.metadata_time_periods), + country.model_version.catalog_format, + Jsonb( + [ + CatalogEntity( + key=entity.key, + plural=entity.plural, + is_person=entity.is_person, + ) + for entity in country.model_version.entities + ] + ), + Jsonb(list(country.model_version.basic_inputs)), + ) + + +def _enum_values_json(record: VariableRecord) -> Jsonb | None: + """Return an enum variable's values as ordered value and label objects.""" + if record.enum_values is None: + return None + return Jsonb( + [ + CatalogEnumValue(value=item.value, label=item.label) + for item in record.enum_values + ] ) @@ -182,6 +221,14 @@ def _variable_rows(country: CountryCatalog) -> Iterator[tuple[object, ...]]: Jsonb(record.default_value), _optional_json(record.adds), _optional_json(record.subtracts), + record.is_input_variable, + record.hidden_input, + record.module_name, + record.adds_parameter_path, + record.subtracts_parameter_path, + _enum_values_json(record), + record.unit, + record.definition_period, ) @@ -213,6 +260,8 @@ def _parameter_rows(country: CountryCatalog) -> Iterator[tuple[object, ...]]: record.description, record.data_type, record.unit, + record.economy, + record.household, ) diff --git a/policyengine_api/data/v2/catalog/records.py b/policyengine_api/data/v2/catalog/records.py index afccc4add..4c026cacd 100644 --- a/policyengine_api/data/v2/catalog/records.py +++ b/policyengine_api/data/v2/catalog/records.py @@ -11,6 +11,37 @@ from uuid import UUID +CATALOG_FORMAT = 2 +"""The catalog content format this extractor publishes. + +Format 2 adds the metadata the app's v3 workspace reads, from PolicyEngine.py's +public records: entities and basic inputs; each variable's input and +hidden-input flags, module, the parameter its ``adds`` or ``subtracts`` names, +and enum value labels; each parameter's economy and household flags. Versions +published before it store null for these fields. +""" + + +@dataclass(frozen=True, slots=True) +class EntityRecord: + """A model entity: its key, plural name, and whether it is the person entity.""" + + key: str + plural: str + is_person: bool + + +@dataclass(frozen=True, slots=True) +class EnumValueRecord: + """One value of an enum variable: its stored name and its declared label. + + The label is kept as the country package declares it, usually a string. + """ + + value: str + label: str | int | float | bool | None + + @dataclass(frozen=True, slots=True) class ModelRecord: """One country's tax-benefit model as named in the catalog.""" @@ -35,11 +66,19 @@ class ModelVersionRecord: description: str | None current_law_id: int metadata_time_periods: tuple[int, ...] + catalog_format: int + entities: tuple[EntityRecord, ...] + basic_inputs: tuple[str, ...] @dataclass(frozen=True, slots=True) class VariableRecord: - """One model variable, with its entity, type and the variables it sums.""" + """One model variable, with its entity, type and the variables it sums. + + ``adds_parameter_path`` and ``subtracts_parameter_path`` name the parameter + whose dated value lists the variables summed; ``adds`` and ``subtracts`` + keep the list PolicyEngine.py resolved from it. + """ id: UUID model_version_id: UUID @@ -52,11 +91,22 @@ class VariableRecord: default_value: Any adds: list[str] | None subtracts: list[str] | None + is_input_variable: bool | None + hidden_input: bool | None + module_name: str | None + adds_parameter_path: str | None + subtracts_parameter_path: str | None + enum_values: tuple[EnumValueRecord, ...] | None + unit: str | None + definition_period: str | None @dataclass(frozen=True, slots=True) class ParameterNodeRecord: - """A non-leaf node of the parameter tree; it groups parameters and has no values.""" + """A non-leaf node of the parameter tree; it groups parameters and has no values. + + Bracket schedules are nodes too, with their stored label or none. + """ id: UUID model_version_id: UUID @@ -88,6 +138,8 @@ class ParameterRecord: data_type: str | None unit: str | None values: tuple[ParameterValueRecord, ...] + economy: bool + household: bool @dataclass(frozen=True, slots=True) diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index 2c7b6986d..d463fad16 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -24,7 +24,7 @@ from collections.abc import Mapping V2_ALEMBIC_DISPOSABLE_TEST = "V2_ALEMBIC_DISPOSABLE_TEST" -V2_ALEMBIC_HEAD_REVISION = "fb1a3e40d3d6" +V2_ALEMBIC_HEAD_REVISION = "f5c1d9b4685b" DISPOSABLE_DATABASE_NAME = "policyengine_v2_alembic_test" DISPOSABLE_HOSTS = frozenset({"localhost", "127.0.0.1", "::1", "postgres"}) MIGRATION_ROLE = "policyengine_v2_migrator" diff --git a/policyengine_api/data/v2/models/__init__.py b/policyengine_api/data/v2/models/__init__.py index a3aedae09..8089fc90a 100644 --- a/policyengine_api/data/v2/models/__init__.py +++ b/policyengine_api/data/v2/models/__init__.py @@ -41,6 +41,8 @@ HouseholdJobStatus, ) from policyengine_api.data.v2.models.metadata import ( # noqa: E402 + CatalogEntity, + CatalogEnumValue, Dataset, DatasetVersion, Parameter, @@ -100,6 +102,8 @@ "AggregateOutput", "AggregateType", "BudgetSummary", + "CatalogEntity", + "CatalogEnumValue", "ChangeAggregate", "CongressionalDistrictImpact", "ConstituencyImpact", diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index 8d91864df..61c210338 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -7,6 +7,7 @@ import sqlalchemy as sa from sqlmodel import Field, Relationship +from typing_extensions import TypedDict from policyengine_api.data.v2.models.base import ( IdentifiedModel, @@ -36,6 +37,28 @@ class RegionType(StrEnum): PLACE = "place" +class CatalogEntity(TypedDict): + """One entity of a catalog version's ``entities`` column. + + ``is_person`` marks the person entity; each version has exactly one. + """ + + key: str + plural: str + is_person: bool + + +class CatalogEnumValue(TypedDict): + """One value of a variable's ``enum_values`` column. + + ``value`` is the stored name and ``label`` the label the country package + declares, kept as declared (usually a string). + """ + + value: str + label: str | int | float | bool | None + + class TaxBenefitModel(TimestampedModel, table=True): """A country's tax-benefit model, unique by name. @@ -85,6 +108,11 @@ class TaxBenefitModelVersion(IdentifiedModel, table=True): description: str | None = None current_law_id: int metadata_time_periods: list[int] = Field(sa_type=sa.JSON) + # Catalog format 2 fields. Versions published before format 2 keep them + # null. + catalog_format: int | None = None + entities: list[CatalogEntity] | None = Field(default=None, sa_type=sa.JSON) + basic_inputs: list[str] | None = Field(default=None, sa_type=sa.JSON) model: TaxBenefitModel = Relationship(back_populates="versions") variables: list["Variable"] = Relationship( @@ -276,6 +304,19 @@ class Variable(IdentifiedModel, table=True): default_value: Any = Field(default=None, sa_type=sa.JSON) adds: list[str] | None = Field(default=None, sa_type=sa.JSON) subtracts: list[str] | None = Field(default=None, sa_type=sa.JSON) + # Catalog format 2 fields; null in versions published before it. A + # parameter path names the parameter whose dated value lists the variables + # this one adds or subtracts. + is_input_variable: bool | None = None + hidden_input: bool | None = None + module_name: str | None = Field(default=None, max_length=512) + adds_parameter_path: str | None = Field(default=None, max_length=512) + subtracts_parameter_path: str | None = Field(default=None, max_length=512) + enum_values: list[CatalogEnumValue] | None = Field(default=None, sa_type=sa.JSON) + # The unit (for example ``currency-USD``) and the period a value covers + # (``year``, ``month`` or ``eternity``), as the country package declares. + unit: str | None = Field(default=None, max_length=128) + definition_period: str | None = Field(default=None, max_length=32) tax_benefit_model_version_id: UUID = Field( foreign_key="tax_benefit_model_versions.id", ondelete="CASCADE", @@ -336,6 +377,9 @@ class Parameter(IdentifiedModel, table=True): description: str | None = None data_type: str | None = Field(default=None, max_length=128) unit: str | None = Field(default=None, max_length=128) + # Catalog format 2 fields; null in versions published before it. + economy: bool | None = None + household: bool | None = None tax_benefit_model_version_id: UUID = Field( foreign_key="tax_benefit_model_versions.id", ondelete="CASCADE", diff --git a/pyproject.toml b/pyproject.toml index 08c441d99..7513ff086 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -44,7 +44,7 @@ dependencies = [ # versions certified by this PolicyEngine bundle. Cloud Run installs the # complete frozen environment from uv.lock; the generic image installs # this requirement directly. - "policyengine[models]==6.11.1", + "policyengine[models]==6.12.0", "pydantic", "pymysql", "python-dotenv", diff --git a/tests/fixtures/v2_catalog.py b/tests/fixtures/v2_catalog.py index 77f5cbd20..e21372173 100644 --- a/tests/fixtures/v2_catalog.py +++ b/tests/fixtures/v2_catalog.py @@ -107,6 +107,14 @@ def _model_source( default_value=0.0, adds=None, subtracts=None, + is_input_variable=True, + hidden_input=False, + module_name="household.income.employment_income", + adds_parameter_path=None, + subtracts_parameter_path=None, + enum_values=None, + unit="currency-USD", + definition_period="year", ) parameter_node = SimpleNamespace( name="gov.example", @@ -119,6 +127,8 @@ def _model_source( description="An example rate", data_type=float, unit="/1", + economy=True, + household=False, parameter_values=[ SimpleNamespace( value=0.1, @@ -148,6 +158,11 @@ def _model_source( default_dataset_uri=default_uri, ) version_holder.region_registry = RegionRegistry(country_id, regions) + version_holder.entities = [ + SimpleNamespace(key="person", plural="people", is_person=True), + SimpleNamespace(key="household", plural="households", is_person=False), + ] + version_holder.basic_inputs = ["employment_income"] version_holder.variables = [variable] version_holder.parameter_nodes = [parameter_node] version_holder.parameters = [parameter] @@ -235,3 +250,78 @@ def normalized_catalog( models=source_models(policyengine_version=policyengine_version), installed_version=installed_version, ) + + +def add_format_2_variable_records(model: SimpleNamespace) -> None: + """Add two variables that between them fill every format 2 variable field. + + ``filing_status`` is an enum whose labels are a string and an integer, with + no unit and a yearly period. ``countable_income`` names the parameters that + hold the lists it adds and subtracts, with a currency unit and a monthly + period. + """ + for name in ("gov.example.countable_income", "gov.example.deductions"): + parameter = SimpleNamespace( + name=name, + label=None, + description=None, + data_type=list, + unit=None, + economy=True, + household=True, + parameter_values=[ + SimpleNamespace( + value=["employment_income"], + start_date=datetime(2025, 1, 1, tzinfo=UTC), + end_date=None, + ) + ], + ) + model.parameters.append(parameter) + model.parameters_by_name[name] = parameter + variables = ( + SimpleNamespace( + name="filing_status", + label="Filing status", + entity="person", + description=None, + data_type="Enum", + possible_values=["SINGLE", "JOINT"], + default_value="SINGLE", + adds=None, + subtracts=None, + is_input_variable=True, + hidden_input=False, + module_name="household.demographic.filing_status", + adds_parameter_path=None, + subtracts_parameter_path=None, + enum_values=[ + SimpleNamespace(value="SINGLE", label="Single"), + SimpleNamespace(value="JOINT", label=2), + ], + unit=None, + definition_period="year", + ), + SimpleNamespace( + name="countable_income", + label="Countable income", + entity="person", + description=None, + data_type=float, + possible_values=None, + default_value=0.0, + adds=["employment_income"], + subtracts=["employment_income"], + is_input_variable=False, + hidden_input=False, + module_name="household.income.countable_income", + adds_parameter_path="gov.example.countable_income", + subtracts_parameter_path="gov.example.deductions", + enum_values=None, + unit="currency-USD", + definition_period="month", + ), + ) + for variable in variables: + model.variables.append(variable) + model.variables_by_name[variable.name] = variable diff --git a/tests/fixtures/v2_catalog_value_shapes.py b/tests/fixtures/v2_catalog_value_shapes.py index 09dd8990e..3cadf1da5 100644 --- a/tests/fixtures/v2_catalog_value_shapes.py +++ b/tests/fixtures/v2_catalog_value_shapes.py @@ -73,6 +73,8 @@ def source(self) -> SimpleNamespace: description=None, data_type=self.data_type, unit=None, + economy=True, + household=False, parameter_values=[ SimpleNamespace( value=value.value, start_date=value.start, end_date=value.end @@ -339,6 +341,17 @@ def _variables() -> list[SimpleNamespace]: default_value=_FilingStatus.SINGLE, adds=None, subtracts=None, + is_input_variable=True, + hidden_input=False, + module_name="household.demographic.filing_status", + adds_parameter_path=None, + subtracts_parameter_path=None, + enum_values=[ + SimpleNamespace(value="SINGLE", label="Single"), + SimpleNamespace(value="JOINT", label="Joint"), + ], + unit=None, + definition_period="year", ), SimpleNamespace( name="is_blind", @@ -350,6 +363,14 @@ def _variables() -> list[SimpleNamespace]: default_value=False, adds=None, subtracts=None, + is_input_variable=True, + hidden_input=False, + module_name="household.demographic.is_blind", + adds_parameter_path=None, + subtracts_parameter_path=None, + enum_values=None, + unit=None, + definition_period="year", ), SimpleNamespace( name="household_net_income", @@ -361,6 +382,14 @@ def _variables() -> list[SimpleNamespace]: default_value=0.0, adds=["employment_income"], subtracts=["income_tax"], + is_input_variable=False, + hidden_input=False, + module_name="household.income.household_net_income", + adds_parameter_path=None, + subtracts_parameter_path=None, + enum_values=None, + unit="currency-USD", + definition_period="year", ), ] diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index 2df169511..a5bd7d20c 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -1248,3 +1248,136 @@ def test_a_household_only_country_row_lets_records_name_that_country() -> None: transaction.rollback() finally: engine.dispose() + + +CATALOG_FORMAT_2_PARENT_REVISION = "fb1a3e40d3d6" +CATALOG_FORMAT_2_COLUMNS = { + "tax_benefit_model_versions": ("catalog_format", "entities", "basic_inputs"), + "variables": ( + "is_input_variable", + "hidden_input", + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "enum_values", + "unit", + "definition_period", + ), + "parameters": ("economy", "household"), +} + + +def _catalog_version_published_before_format_2(connection: sa.Connection) -> UUID: + """Store a model version with one variable and one parameter; return its id.""" + model_id = uuid4() + version_id = uuid4() + connection.execute( + text( + "INSERT INTO tax_benefit_models (id, name) " + "VALUES (:model_id, 'catalog-format-probe')" + ), + {"model_id": model_id}, + ) + connection.execute( + text( + "INSERT INTO tax_benefit_model_versions " + "(id, model_id, version, current_law_id, metadata_time_periods) " + "VALUES (:version_id, :model_id, '0.0.1-probe', 1, '[2026]')" + ), + {"version_id": version_id, "model_id": model_id}, + ) + connection.execute( + text( + "INSERT INTO variables " + "(id, name, entity, default_value, tax_benefit_model_version_id) " + "VALUES (:variable_id, 'probe_income', 'person', '0', :version_id)" + ), + {"variable_id": uuid4(), "version_id": version_id}, + ) + connection.execute( + text( + "INSERT INTO parameters (id, name, tax_benefit_model_version_id) " + "VALUES (:parameter_id, 'gov.probe.rate', :version_id)" + ), + {"parameter_id": uuid4(), "version_id": version_id}, + ) + return version_id + + +def _format_2_values(connection: sa.Connection, version_id: UUID) -> list[tuple]: + """Return the format 2 columns of the version's model, variable and parameter.""" + rows = [] + for table_name, columns in CATALOG_FORMAT_2_COLUMNS.items(): + key = ( + "id" + if table_name == "tax_benefit_model_versions" + else "tax_benefit_model_version_id" + ) + rows.append( + tuple( + connection.execute( + text( + f"SELECT {', '.join(columns)} FROM {table_name} " + f"WHERE {key} = :version_id" + ), + {"version_id": version_id}, + ).one() + ) + ) + return rows + + +def test_catalog_format_2_columns_leave_earlier_versions_empty() -> None: + """A version stored before catalog format 2 keeps its format 2 fields null. + + Upgrading adds the columns as nullable without writing to stored rows, and + downgrading drops them and keeps every row. + """ + config = _config() + engine = create_engine(_disposable_url()) + version_id = None + try: + command.downgrade(config, CATALOG_FORMAT_2_PARENT_REVISION) + for table_name, columns in CATALOG_FORMAT_2_COLUMNS.items(): + assert set(columns).isdisjoint(_columns(engine, table_name)) + with engine.begin() as connection: + version_id = _catalog_version_published_before_format_2(connection) + + command.upgrade(config, "head") + + for table_name, columns in CATALOG_FORMAT_2_COLUMNS.items(): + stored_columns = _columns(engine, table_name) + for column_name in columns: + assert stored_columns[column_name]["nullable"] + assert stored_columns[column_name]["default"] is None + with engine.connect() as connection: + assert _format_2_values(connection, version_id) == [ + (None, None, None), + (None, None, None, None, None, None, None, None), + (None, None), + ] + + command.downgrade(config, CATALOG_FORMAT_2_PARENT_REVISION) + + for table_name, columns in CATALOG_FORMAT_2_COLUMNS.items(): + assert set(columns).isdisjoint(_columns(engine, table_name)) + with engine.connect() as connection: + counts = connection.execute( + text( + "SELECT (SELECT count(*) FROM variables " + "WHERE tax_benefit_model_version_id = :version_id), " + "(SELECT count(*) FROM parameters " + "WHERE tax_benefit_model_version_id = :version_id)" + ), + {"version_id": version_id}, + ).one() + assert tuple(counts) == (1, 1) + finally: + with engine.begin() as connection: + connection.execute( + text( + "DELETE FROM tax_benefit_models WHERE name = 'catalog-format-probe'" + ) + ) + command.upgrade(config, "head") + engine.dispose() diff --git a/tests/integration/test_v2_catalog_installed_folders.py b/tests/integration/test_v2_catalog_installed_folders.py index 2b4de803d..7f178773d 100644 --- a/tests/integration/test_v2_catalog_installed_folders.py +++ b/tests/integration/test_v2_catalog_installed_folders.py @@ -5,8 +5,8 @@ publication instead, each reduced to the parameters under the folders in SELECTED_FOLDERS, and checks that every extracted value comes back from PostgreSQL unchanged. Extraction gets only what it needs besides those -parameters: each model's identity, package and release manifest, its national -region, and no variables. +parameters: each model's identity, package, release manifest and entities, its +national region, and no variables or basic inputs. tests/fixtures/v2_catalog_value_shapes.py covers the value shapes extraction accepts that the installed models don't hold, such as text and mappings. @@ -115,6 +115,9 @@ def selected_model(country_id: str, model: object, folders: Iterable[str]) -> ob country_id, [region for region in model.region_registry if region.code == country_id], ), + entities=model.entities, + # Basic inputs name variables, and the reduced model has none. + basic_inputs=[], variables_by_name={}, parameter_nodes_by_name={ name: node diff --git a/tests/integration/test_v2_catalog_metadata_parity.py b/tests/integration/test_v2_catalog_metadata_parity.py new file mode 100644 index 000000000..b8aa6485d --- /dev/null +++ b/tests/integration/test_v2_catalog_metadata_parity.py @@ -0,0 +1,208 @@ +"""Parity of catalog format 2's fields with the v1 ``/{country}/metadata`` document. + +The catalog reads these fields from PolicyEngine.py's public records; the v1 +document builds them from the country package's tax-benefit system. Each new +field must agree for the US and the UK. The check normalizes the installed +public models' metadata only, so it does not depend on the certified data +release that region publication reads. +""" + +from __future__ import annotations + +import os +from typing import TYPE_CHECKING, Any +from uuid import uuid4 + +import pytest + +from policyengine_api.data.v2.catalog.extraction import normalize_model_metadata + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping + + from policyengine_api.data.v2.catalog.extraction import ModelMetadata + +# CI never runs this test: it extracts every variable and parameter of the +# installed models, with their values, and builds each country's v1 metadata. +pytestmark = pytest.mark.skipif( + os.environ.get("RUN_V2_CATALOG_METADATA_PARITY") != "1", + reason=( + "extracts the complete installed catalog; set RUN_V2_CATALOG_METADATA_PARITY=1" + ), +) + +COUNTRY_PACKAGES = {"us": "policyengine_us", "uk": "policyengine_uk"} + + +@pytest.fixture(scope="module", params=("us", "uk")) +def country_pair(request: pytest.FixtureRequest) -> tuple[ModelMetadata, Any]: + """Return one country's normalized metadata and its v1 metadata builder.""" + import policyengine + + from policyengine_api.country import PolicyEngineCountry + + country_id = request.param + metadata = normalize_model_metadata( + country_id, + getattr(policyengine, country_id).model, + model_version_id=uuid4(), + ) + return metadata, PolicyEngineCountry(COUNTRY_PACKAGES[country_id], country_id) + + +def _report(mismatches: list[str]) -> str: + shown = "\n".join(mismatches[:20]) + return f"{len(mismatches)} mismatches, first {min(len(mismatches), 20)}:\n{shown}" + + +def _typed_enum_values( + values: Iterable[Mapping[str, Any]], +) -> list[tuple[str, Any, str]]: + """Return each enum value with its label and the label's type name. + + Comparing the type name as well keeps 1, 1.0 and True apart, which ``==`` + treats as equal. + """ + return [ + (value["value"], value["label"], type(value["label"]).__name__) + for value in values + ] + + +def test_entities_and_basic_inputs_match_v1( + country_pair: tuple[ModelMetadata, Any], +) -> None: + metadata, v1_country = country_pair + v1_entities = v1_country.metadata["entities"] + + assert [ + (entity.key, entity.plural, entity.is_person) for entity in metadata.entities + ] == [ + (key, entity["plural"], entity["is_person"]) + for key, entity in v1_entities.items() + ] + assert list(metadata.basic_inputs) == (v1_country.metadata["basicInputs"]) + + +def test_variable_flags_modules_enums_parts_units_and_periods_match_v1( + country_pair: tuple[ModelMetadata, Any], +) -> None: + metadata, v1_country = country_pair + v1_variables = v1_country.metadata["variables"] + mismatches: list[str] = [] + + assert {variable.name for variable in metadata.variables} == set(v1_variables) + for variable in metadata.variables: + v1_variable = v1_variables[variable.name] + enum_values = ( + _typed_enum_values( + {"value": item.value, "label": item.label} + for item in variable.enum_values + ) + if variable.enum_values is not None + else None + ) + v1_enum_values = ( + _typed_enum_values(v1_variable["possibleValues"]) + if "possibleValues" in v1_variable + else None + ) + comparisons = { + "is_input_variable": ( + variable.is_input_variable, + v1_variable["isInputVariable"], + ), + "hidden_input": (variable.hidden_input, v1_variable["hidden_input"]), + "module_name": (variable.module_name, v1_variable["moduleName"]), + "enum_values": (enum_values, v1_enum_values), + "unit": (variable.unit, v1_variable["unit"]), + "definition_period": ( + variable.definition_period, + v1_variable["definitionPeriod"], + ), + "adds_parameter_path": ( + variable.adds_parameter_path, + v1_variable["adds"] if isinstance(v1_variable["adds"], str) else None, + ), + "subtracts_parameter_path": ( + variable.subtracts_parameter_path, + v1_variable["subtracts"] + if isinstance(v1_variable["subtracts"], str) + else None, + ), + } + mismatches.extend( + f"{variable.name}.{field}: catalog {stored!r}, v1 {expected!r}" + for field, (stored, expected) in comparisons.items() + if stored != expected or type(stored) is not type(expected) + ) + + assert not mismatches, _report(mismatches) + + +def test_parameter_flags_match_v1(country_pair: tuple[ModelMetadata, Any]) -> None: + metadata, v1_country = country_pair + v1_parameters = { + name: entry + for name, entry in v1_country.metadata["parameters"].items() + if entry["type"] == "parameter" + } + catalog_parameters = { + parameter.name: parameter + for parameter in metadata.parameters + if parameter.name.startswith("gov") + } + mismatches: list[str] = [] + + assert set(catalog_parameters) == set(v1_parameters) + for name, parameter in catalog_parameters.items(): + expected = (v1_parameters[name]["economy"], v1_parameters[name]["household"]) + if (parameter.economy, parameter.household) != expected: + mismatches.append( + f"{name}: catalog {(parameter.economy, parameter.household)}, " + f"v1 {expected}" + ) + + assert not mismatches, _report(mismatches) + + +def test_bracket_schedules_are_nodes_with_the_stored_label_v1_shows( + country_pair: tuple[ModelMetadata, Any], +) -> None: + """Every schedule v1 lists is a catalog node; a stored label matches v1's. + + v1 falls back to the name's last part when a schedule has no stored label; + the catalog stores no label, and the routes generate one. + """ + from policyengine_core.parameters import ParameterScale + + metadata, v1_country = country_pair + nodes = {node.name: node for node in metadata.parameter_nodes} + v1_parameters = v1_country.metadata["parameters"] + schedules = [ + parameter + for parameter in v1_country.tax_benefit_system.parameters.get_descendants() + if isinstance(parameter, ParameterScale) and parameter.name.startswith("gov") + ] + mismatches: list[str] = [] + + assert schedules + for schedule in schedules: + node = nodes.get(schedule.name) + stored_label = schedule.metadata.get("label") + if node is None: + mismatches.append(f"{schedule.name}: not a catalog node") + elif node.label != stored_label: + mismatches.append( + f"{schedule.name}: catalog label {node.label!r}, " + f"stored label {stored_label!r}" + ) + elif stored_label is not None and ( + v1_parameters[schedule.name]["label"] != stored_label + ): + mismatches.append( + f"{schedule.name}: catalog label {node.label!r}, " + f"v1 label {v1_parameters[schedule.name]['label']!r}" + ) + + assert not mismatches, _report(mismatches) diff --git a/tests/integration/test_v2_catalog_publication.py b/tests/integration/test_v2_catalog_publication.py index aeab2e5b6..6fb506732 100644 --- a/tests/integration/test_v2_catalog_publication.py +++ b/tests/integration/test_v2_catalog_publication.py @@ -16,10 +16,12 @@ from sqlalchemy.pool import NullPool from sqlmodel import Session, select +from policyengine_api.data.v2.catalog.extraction import extract_catalog from policyengine_api.data.v2.catalog.publication import ( CatalogPublicationError, publish_catalog, ) +from policyengine_api.data.v2.catalog.records import EnumValueRecord from policyengine_api.data.v2.models import ( Dataset, DatasetVersion, @@ -31,7 +33,14 @@ TaxBenefitModelVersion, ) from policyengine_api.data.v2.settings import V2_MIGRATION_DATABASE_URL -from tests.fixtures.v2_catalog import normalized_catalog +from tests.fixtures.v2_catalog import ( + POLICYENGINE_VERSION, + add_format_2_variable_records, + bundle, + installed_version, + normalized_catalog, + source_models, +) from tests.fixtures.v2_catalog_value_shapes import ( PARAMETER_CASES, REGION_FILTER_VALUES, @@ -384,3 +393,243 @@ def test_publication_stores_every_value_shape_unchanged( model_name = catalog.country(country_id).model.name stored_columns = json.loads(stored_regions[(model_name, code)]) assert stored_columns[filter_value_index] == filter_value, code + + +def _format_2_values(engine: Engine) -> dict[str, object]: + """Return the US version's stored catalog format 2 fields.""" + with engine.connect() as connection: + version = connection.execute( + text( + "SELECT version.catalog_format, version.entities, version.basic_inputs " + "FROM tax_benefit_model_versions AS version " + "JOIN tax_benefit_models AS model ON model.id = version.model_id " + "WHERE model.name = 'policyengine-us'" + ) + ).one() + variable = connection.execute( + text( + "SELECT is_input_variable, hidden_input, module_name, " + "adds_parameter_path, subtracts_parameter_path, enum_values, " + "unit, definition_period " + "FROM variables WHERE name = 'employment_income' " + "AND tax_benefit_model_version_id = (" + "SELECT version.id FROM tax_benefit_model_versions AS version " + "JOIN tax_benefit_models AS model ON model.id = version.model_id " + "WHERE model.name = 'policyengine-us')" + ) + ).one() + parameter = connection.execute( + text( + "SELECT economy, household FROM parameters " + "WHERE name = 'gov.example.rate' " + "AND tax_benefit_model_version_id = (" + "SELECT version.id FROM tax_benefit_model_versions AS version " + "JOIN tax_benefit_models AS model ON model.id = version.model_id " + "WHERE model.name = 'policyengine-us')" + ) + ).one() + return { + "version": tuple(version), + "variable": tuple(variable), + "parameter": tuple(parameter), + } + + +def _clear_format_2_fields(engine: Engine) -> None: + """Make every stored version look as if published before catalog format 2.""" + with engine.begin() as connection: + connection.execute( + text( + "UPDATE tax_benefit_model_versions SET catalog_format = NULL, " + "entities = NULL, basic_inputs = NULL" + ) + ) + connection.execute( + text( + "UPDATE variables SET is_input_variable = NULL, hidden_input = NULL, " + "module_name = NULL, adds_parameter_path = NULL, " + "subtracts_parameter_path = NULL, enum_values = NULL, " + "unit = NULL, definition_period = NULL" + ) + ) + connection.execute( + text("UPDATE parameters SET economy = NULL, household = NULL") + ) + + +def test_publication_stores_catalog_format_2_fields( + publication_engine: Engine, +) -> None: + publish_catalog(publication_engine, normalized_catalog()) + + assert _format_2_values(publication_engine) == { + "version": ( + 2, + [ + {"key": "person", "plural": "people", "is_person": True}, + {"key": "household", "plural": "households", "is_person": False}, + ], + ["employment_income"], + ), + "variable": ( + True, + False, + "household.income.employment_income", + None, + None, + None, + "currency-USD", + "year", + ), + "parameter": (True, False), + } + + +def test_version_stored_before_format_2_keeps_its_fields_empty_on_republish( + publication_engine: Engine, +) -> None: + catalog = normalized_catalog() + publish_catalog(publication_engine, catalog) + _clear_format_2_fields(publication_engine) + before_ids = _identifiers(publication_engine) + + result = publish_catalog(publication_engine, catalog) + + assert result.as_dict()["outcome"] == "ok" + assert _identifiers(publication_engine) == before_ids + assert _format_2_values(publication_engine) == { + "version": (None, None, None), + "variable": (None, None, None, None, None, None, None, None), + "parameter": (None, None), + } + + +def test_version_stored_with_format_2_fails_when_a_format_2_field_changes( + publication_engine: Engine, +) -> None: + catalog = normalized_catalog() + publish_catalog(publication_engine, catalog) + before_counts = _counts(publication_engine) + us = catalog.country("us") + changed_parameter = replace(us.parameters[0], household=True) + changed = replace( + catalog, + countries=( + replace(us, parameters=(changed_parameter,)), + catalog.country("uk"), + ), + ) + + with pytest.raises(CatalogPublicationError, match="differs"): + publish_catalog(publication_engine, changed) + + assert _counts(publication_engine) == before_counts + assert _format_2_values(publication_engine)["parameter"] == (True, False) + + +def _catalog_with_format_2_variable_records(): + models = source_models() + add_format_2_variable_records(models["us"]) + return extract_catalog( + bundle=bundle(), + policyengine_version=POLICYENGINE_VERSION, + models=models, + installed_version=installed_version, + ) + + +def _stored_us_variables(engine: Engine) -> dict[str, tuple[object, ...]]: + with engine.connect() as connection: + rows = connection.execute( + text( + "SELECT variable.name, variable.enum_values, " + "variable.adds_parameter_path, variable.subtracts_parameter_path, " + "variable.unit, variable.definition_period " + "FROM variables AS variable " + "JOIN tax_benefit_model_versions AS version " + "ON version.id = variable.tax_benefit_model_version_id " + "JOIN tax_benefit_models AS model ON model.id = version.model_id " + "WHERE model.name = 'policyengine-us'" + ) + ).all() + return {row[0]: tuple(row[1:]) for row in rows} + + +def _with_changed_us_variable(catalog, name: str, **changes: object): + us = catalog.country("us") + variables = tuple( + replace(variable, **changes) if variable.name == name else variable + for variable in us.variables + ) + return replace( + catalog, + countries=(replace(us, variables=variables), catalog.country("uk")), + ) + + +def test_enum_labels_parameter_paths_units_and_periods_round_trip_exactly( + publication_engine: Engine, +) -> None: + catalog = _catalog_with_format_2_variable_records() + + publish_catalog(publication_engine, catalog) + before_ids = _identifiers(publication_engine) + stored = _stored_us_variables(publication_engine) + republished = publish_catalog(publication_engine, catalog) + + assert stored["filing_status"] == ( + [{"value": "SINGLE", "label": "Single"}, {"value": "JOINT", "label": 2}], + None, + None, + None, + "year", + ) + assert type(stored["filing_status"][0][1]["label"]) is int + assert stored["countable_income"] == ( + None, + "gov.example.countable_income", + "gov.example.deductions", + "currency-USD", + "month", + ) + assert republished.as_dict()["outcome"] == "ok" + assert _identifiers(publication_engine) == before_ids + + +@pytest.mark.parametrize( + ("variable_name", "changes"), + [ + ( + "filing_status", + { + "enum_values": ( + EnumValueRecord(value="SINGLE", label="Single filer"), + EnumValueRecord(value="JOINT", label=2), + ) + }, + ), + ("countable_income", {"adds_parameter_path": "gov.example.deductions"}), + ("countable_income", {"subtracts_parameter_path": None}), + ("filing_status", {"unit": "currency-USD"}), + ("countable_income", {"definition_period": "year"}), + ], +) +def test_republication_detects_a_changed_enum_label_parameter_path_unit_or_period( + publication_engine: Engine, + variable_name: str, + changes: dict[str, object], +) -> None: + catalog = _catalog_with_format_2_variable_records() + publish_catalog(publication_engine, catalog) + stored = _stored_us_variables(publication_engine) + + with pytest.raises( + CatalogPublicationError, + match=r"^persisted us catalog differs from PolicyEngine.py$", + ): + publish_catalog( + publication_engine, + _with_changed_us_variable(catalog, variable_name, **changes), + ) + + assert _stored_us_variables(publication_engine) == stored diff --git a/tests/unit/test_supported_countries.py b/tests/unit/test_supported_countries.py index db9a42740..68fe331fc 100644 --- a/tests/unit/test_supported_countries.py +++ b/tests/unit/test_supported_countries.py @@ -146,10 +146,9 @@ def test_every_country_list_follows_a_country_added_to_the_one_list() -> None: def test_catalog_settings_cover_exactly_the_listed_countries() -> None: - # Each listed country needs reviewed catalog settings, so adding a country + # Each listed country needs catalog settings, so adding a country # without them fails here rather than during publication. listed_countries = set(SUPPORTED_COUNTRY_IDS) - assert set(extraction.REVIEWED_DEFAULT_DATASETS) == listed_countries assert set(extraction.CURRENT_LAW_IDS) == listed_countries assert set(extraction.METADATA_TIME_PERIODS) == listed_countries diff --git a/tests/unit/v2/test_alembic_v2.py b/tests/unit/v2/test_alembic_v2.py index 56e31410b..9eca8f5b3 100644 --- a/tests/unit/v2/test_alembic_v2.py +++ b/tests/unit/v2/test_alembic_v2.py @@ -55,6 +55,8 @@ f"{PROJECT_REF}:test-password@aws-0-us-east-2.pooler.supabase.com:5432/" "postgres?sslmode=require" ) +# The catalog format 2 revision's parent, the countries revision. +CATALOG_FORMAT_2_PARENT_REVISION = "fb1a3e40d3d6" DISPOSABLE_URL = ( "postgresql+psycopg://postgres:test-password@127.0.0.1:5432/" f"{DISPOSABLE_DATABASE_NAME}" @@ -236,6 +238,7 @@ def test_v2_revision_chain_has_generated_policy_and_user_identity_changes() -> N assert script.get_heads() == [V2_ALEMBIC_HEAD_REVISION] assert [revision.revision for revision in script.walk_revisions()] == [ V2_ALEMBIC_HEAD_REVISION, + CATALOG_FORMAT_2_PARENT_REVISION, "fbb7488ebe8a", "91af4fc8d8c1", "60d6518b6a98", @@ -605,7 +608,7 @@ def _countries_revision_source() -> str: for revision in script.walk_revisions() if revision.down_revision == "fbb7488ebe8a" ) - assert countries_revision.revision == V2_ALEMBIC_HEAD_REVISION + assert countries_revision.revision == CATALOG_FORMAT_2_PARENT_REVISION assert countries_revision.path is not None return Path(countries_revision.path).read_text(encoding="utf-8") @@ -828,3 +831,54 @@ def test_v2_environment_commits_after_persistent_target_qualification() -> None: assert "current_heads != previous_heads" in env_source assert "current_heads == script_heads" in env_source assert "validate_v2_head_schema(connection, target_metadata)" in env_source + + +def test_catalog_format_2_revision_adds_only_nullable_columns() -> None: + """The revision after the countries revision adds catalog format 2's columns. + + It is generated, adds each column as nullable so versions published earlier + keep them empty, drops them on downgrade, and writes no data. + """ + script = ScriptDirectory.from_config(Config(str(REPO / "alembic-v2.ini"))) + children = [ + revision + for revision in script.walk_revisions() + if revision.down_revision == CATALOG_FORMAT_2_PARENT_REVISION + ] + assert len(children) == 1 + revision = Path(children[0].path).read_text(encoding="utf-8") + upgrade = revision[ + revision.index("def upgrade()") : revision.index("def downgrade()") + ] + downgrade = revision[revision.index("def downgrade()") :] + + assert ( + "Generation: uv run alembic -c alembic-v2.ini revision --autogenerate" + in revision + ) + added_columns = { + "tax_benefit_model_versions": ("catalog_format", "entities", "basic_inputs"), + "variables": ( + "is_input_variable", + "hidden_input", + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "enum_values", + "unit", + "definition_period", + ), + "parameters": ("economy", "household"), + } + column_count = sum(len(columns) for columns in added_columns.values()) + assert upgrade.count("op.add_column(") == column_count + assert upgrade.count("nullable=True") == column_count + assert downgrade.count("op.drop_column(") == column_count + for table_name, columns in added_columns.items(): + for column_name in columns: + assert f'"{column_name}"' in upgrade + assert f'op.drop_column("{table_name}", "{column_name}")' in downgrade + for body in (upgrade, downgrade): + assert "op.execute(" not in body + assert "op.bulk_insert(" not in body + assert "server_default" not in body diff --git a/tests/unit/v2/test_catalog_extraction.py b/tests/unit/v2/test_catalog_extraction.py index 302fee0c5..e1738d24f 100644 --- a/tests/unit/v2/test_catalog_extraction.py +++ b/tests/unit/v2/test_catalog_extraction.py @@ -300,9 +300,61 @@ def test_rejects_incomplete_public_model_catalogs() -> None: installed_version=installed_version, ) + +def _use_another_default(model: SimpleNamespace, dataset_name: str) -> None: + """Make the release default to another dataset, which its national region uses.""" + manifest = model.release_manifest + manifest.default_dataset = dataset_name + manifest.default_dataset_uri = f"hf://policyengine/{dataset_name}.h5@fixture" + for region in model.region_registry: + if region.region_type == "national": + region.dataset_path = manifest.default_dataset_uri + + +def test_each_country_uses_the_default_dataset_its_release_names() -> None: + """The release's own default is used; the Public API fixes none of its own.""" + models = source_models() + _use_another_default(models["us"], "populace_us_2025") + _use_another_default(models["uk"], "microcosm_uk_2025_26") + + catalog = extract_catalog( + bundle=bundle(), + policyengine_version=POLICYENGINE_VERSION, + models=models, + installed_version=installed_version, + ) + + for country_id, default_dataset in ( + ("us", "populace_us_2025"), + ("uk", "microcosm_uk_2025_26"), + ): + country = catalog.country(country_id) + dataset_names = {dataset.id: dataset.name for dataset in country.datasets} + national = next( + region for region in country.regions if region.region_type == "national" + ) + assert dataset_names[national.default_dataset_id] == default_dataset + assert default_dataset in set(dataset_names.values()) + assert not {"populace_us_2024", "microcosm_uk_2024_25"} & set( + dataset_names.values() + ) + + +@pytest.mark.parametrize( + ("default_dataset", "message"), + [ + (None, "^us release manifest names no default dataset$"), + ("", "^us release manifest names no default dataset$"), + ], +) +def test_a_missing_release_default_fails( + default_dataset: str | None, + message: str, +) -> None: models = source_models() - models["us"].release_manifest.default_dataset = "unreviewed_us_2025" - with pytest.raises(CatalogExtractionError, match="reviewed selection"): + models["us"].release_manifest.default_dataset = default_dataset + + with pytest.raises(CatalogExtractionError, match=message): extract_catalog( bundle=bundle(), policyengine_version=POLICYENGINE_VERSION, diff --git a/tests/unit/v2/test_catalog_extraction_failures.py b/tests/unit/v2/test_catalog_extraction_failures.py index fa46c3f80..c72b2f039 100644 --- a/tests/unit/v2/test_catalog_extraction_failures.py +++ b/tests/unit/v2/test_catalog_extraction_failures.py @@ -103,7 +103,8 @@ def _give_filter_additional_filters(models: dict[str, SimpleNamespace]) -> None: (_remove_national_region, "^us national region is absent$"), ( _give_national_region_another_dataset, - "^uk national region does not use the reviewed dataset$", + "^uk national region does not use the release's default dataset " + "microcosm_uk_2024_25$", ), ( _give_filter_an_unsupported_value, @@ -126,7 +127,7 @@ def test_inconsistent_public_models_fail_with_a_specific_message( _extract(models) -def test_national_region_without_a_dataset_path_uses_the_reviewed_dataset() -> None: +def test_national_region_without_a_dataset_path_uses_the_release_default() -> None: models = source_models() models["us"].region_registry[0].dataset_path = None diff --git a/tests/unit/v2/test_catalog_metadata_fields.py b/tests/unit/v2/test_catalog_metadata_fields.py new file mode 100644 index 000000000..5e6d01559 --- /dev/null +++ b/tests/unit/v2/test_catalog_metadata_fields.py @@ -0,0 +1,488 @@ +"""Catalog format 2: the metadata v3 reads, taken from .py's public records.""" + +from __future__ import annotations + +from datetime import UTC, datetime +from types import SimpleNamespace +from typing import TYPE_CHECKING + +import pytest + +from policyengine_api.data.v2.catalog.extraction import ( + CatalogExtractionError, + extract_catalog, +) +from policyengine_api.data.v2.catalog.records import ( + CATALOG_FORMAT, + EntityRecord, + EnumValueRecord, + NormalizedCatalog, +) +from tests.fixtures.v2_catalog import ( + POLICYENGINE_VERSION, + bundle, + installed_version, + normalized_catalog, + source_models, +) + +if TYPE_CHECKING: + from collections.abc import Callable + + +def _extract(models: dict[str, SimpleNamespace]) -> NormalizedCatalog: + return extract_catalog( + bundle=bundle(), + policyengine_version=POLICYENGINE_VERSION, + models=models, + installed_version=installed_version, + ) + + +def _add_variable(model: SimpleNamespace, **fields: object) -> None: + variable = SimpleNamespace( + label=None, + entity="person", + description=None, + data_type=float, + possible_values=None, + default_value=0.0, + adds=None, + subtracts=None, + is_input_variable=False, + hidden_input=False, + module_name=None, + adds_parameter_path=None, + subtracts_parameter_path=None, + enum_values=None, + unit=None, + definition_period="year", + ) + for name, value in fields.items(): + setattr(variable, name, value) + model.variables.append(variable) + model.variables_by_name[variable.name] = variable + + +def _add_parameter(model: SimpleNamespace, name: str, *, value: object) -> None: + parameter = SimpleNamespace( + name=name, + label=None, + description=None, + data_type=type(value), + unit=None, + economy=True, + household=True, + parameter_values=[ + SimpleNamespace( + value=value, + start_date=datetime(2025, 1, 1, tzinfo=UTC), + end_date=None, + ) + ], + ) + model.parameters.append(parameter) + model.parameters_by_name[name] = parameter + + +def _add_parameter_node(model: SimpleNamespace, name: str, label: str | None) -> None: + node = SimpleNamespace(name=name, label=label, description=None) + model.parameter_nodes.append(node) + model.parameter_nodes_by_name[name] = node + + +def test_catalog_format_2_publishes_entities_and_basic_inputs() -> None: + catalog = normalized_catalog() + + for country_id in ("us", "uk"): + model_version = catalog.country(country_id).model_version + assert model_version.catalog_format == CATALOG_FORMAT == 2 + assert model_version.entities == ( + EntityRecord(key="person", plural="people", is_person=True), + EntityRecord(key="household", plural="households", is_person=False), + ) + assert model_version.basic_inputs == ("employment_income",) + + +def test_variables_carry_flags_module_unit_period_and_unset_parts() -> None: + variable = normalized_catalog().country("us").variables[0] + + assert variable.name == "employment_income" + assert variable.is_input_variable is True + assert variable.hidden_input is False + assert variable.module_name == "household.income.employment_income" + assert variable.adds_parameter_path is None + assert variable.subtracts_parameter_path is None + assert variable.enum_values is None + assert variable.unit == "currency-USD" + assert variable.definition_period == "year" + + +def test_enum_values_keep_their_declared_labels_in_order() -> None: + models = source_models() + _add_variable( + models["us"], + name="filing_status", + data_type="Enum", + possible_values=["SINGLE", "JOINT", "EXEMPT"], + default_value="SINGLE", + is_input_variable=True, + enum_values=[ + SimpleNamespace(value="SINGLE", label="Single"), + SimpleNamespace(value="JOINT", label=2), + SimpleNamespace(value="EXEMPT", label=None), + ], + ) + + variables = { + variable.name: variable for variable in _extract(models).country("us").variables + } + + assert variables["filing_status"].enum_values == ( + EnumValueRecord(value="SINGLE", label="Single"), + EnumValueRecord(value="JOINT", label=2), + EnumValueRecord(value="EXEMPT", label=None), + ) + assert isinstance(variables["filing_status"].enum_values[1].label, int) + + +def test_parameter_named_parts_record_the_parameter_path_beside_the_list() -> None: + models = source_models() + _add_parameter( + models["us"], + "gov.example.countable_income", + value=["employment_income"], + ) + _add_parameter(models["us"], "gov.example.deductions", value=["employment_income"]) + _add_variable( + models["us"], + name="countable_income", + adds=["employment_income"], + subtracts=["employment_income"], + adds_parameter_path="gov.example.countable_income", + subtracts_parameter_path="gov.example.deductions", + ) + + variables = { + variable.name: variable for variable in _extract(models).country("us").variables + } + + countable_income = variables["countable_income"] + assert countable_income.adds_parameter_path == "gov.example.countable_income" + assert countable_income.subtracts_parameter_path == "gov.example.deductions" + assert countable_income.adds == ["employment_income"] + assert countable_income.subtracts == ["employment_income"] + + +def test_parameters_carry_their_economy_and_household_flags() -> None: + parameter = normalized_catalog().country("us").parameters[0] + + assert parameter.name == "gov.example.rate" + assert (parameter.economy, parameter.household) == (True, False) + + +def test_bracket_schedules_are_stored_as_nodes_with_their_stored_label() -> None: + models = source_models() + _add_parameter_node(models["uk"], "gov.example.rates", "Example rates") + _add_parameter_node(models["uk"], "gov.example.rates[0]", None) + _add_parameter_node(models["uk"], "gov.example.thresholds", None) + _add_parameter_node(models["uk"], "gov.example.thresholds[0]", None) + _add_parameter(models["uk"], "gov.example.rates[0].rate", value=0.2) + _add_parameter(models["uk"], "gov.example.thresholds[0].threshold", value=0) + + nodes = {node.name: node for node in _extract(models).country("uk").parameter_nodes} + + assert nodes["gov.example.rates"].label == "Example rates" + assert nodes["gov.example.thresholds"].label is None + assert "gov.example.rates[0]" in nodes + + +def _remove_entities(models: dict[str, SimpleNamespace]) -> None: + del models["us"].entities + + +def _make_entity_flag_text(models: dict[str, SimpleNamespace]) -> None: + models["us"].entities[0].is_person = "yes" + + +def _remove_person_entity(models: dict[str, SimpleNamespace]) -> None: + models["us"].entities = models["us"].entities[1:] + + +def _repeat_entity_key(models: dict[str, SimpleNamespace]) -> None: + models["us"].entities.append( + SimpleNamespace(key="household", plural="homes", is_person=False) + ) + + +def _remove_basic_inputs(models: dict[str, SimpleNamespace]) -> None: + del models["us"].basic_inputs + + +def _make_entities_null(models: dict[str, SimpleNamespace]) -> None: + models["us"].entities = None + + +def _make_basic_inputs_text(models: dict[str, SimpleNamespace]) -> None: + models["us"].basic_inputs = "employment_income" + + +def _remove_variable_enum_values(models: dict[str, SimpleNamespace]) -> None: + del models["us"].variables[0].enum_values + + +def _make_enum_label_nan(models: dict[str, SimpleNamespace]) -> None: + variable = models["us"].variables[0] + variable.possible_values = ["SINGLE"] + variable.enum_values = [SimpleNamespace(value="SINGLE", label=float("nan"))] + + +def _name_unknown_basic_input(models: dict[str, SimpleNamespace]) -> None: + models["us"].basic_inputs = ["employment_income", "unpublished_input"] + + +def _remove_variable_unit(models: dict[str, SimpleNamespace]) -> None: + del models["us"].variables[0].unit + + +def _make_unit_a_number(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].unit = 1 + + +def _make_definition_period_a_list(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].definition_period = ["year"] + + +def _make_input_flag_text(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].is_input_variable = "true" + + +def _make_hidden_flag_number(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].hidden_input = 0 + + +def _name_unknown_adds_parameter(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].adds_parameter_path = "gov.example.missing" + + +def _name_unknown_subtracts_parameter(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].subtracts_parameter_path = "gov.example.missing" + + +def _mismatch_enum_values(models: dict[str, SimpleNamespace]) -> None: + variable = models["us"].variables[0] + variable.possible_values = ["SINGLE", "JOINT"] + variable.enum_values = [SimpleNamespace(value="SINGLE", label="Single")] + + +def _give_enum_labels_without_values(models: dict[str, SimpleNamespace]) -> None: + models["us"].variables[0].enum_values = [ + SimpleNamespace(value="SINGLE", label="Single") + ] + + +def _make_enum_label_an_object(models: dict[str, SimpleNamespace]) -> None: + variable = models["us"].variables[0] + variable.possible_values = ["SINGLE"] + variable.enum_values = [SimpleNamespace(value="SINGLE", label={"text": "Single"})] + + +def _make_economy_flag_missing(models: dict[str, SimpleNamespace]) -> None: + del models["us"].parameters[0].economy + + +def _make_economy_flag_null(models: dict[str, SimpleNamespace]) -> None: + models["us"].parameters[0].economy = None + + +def _make_household_flag_text(models: dict[str, SimpleNamespace]) -> None: + models["us"].parameters[0].household = "false" + + +@pytest.mark.parametrize( + ("break_models", "message"), + [ + ( + _remove_entities, + "^the installed PolicyEngine.py does not publish catalog format 2's " + "fields for us: model version entities$", + ), + (_make_entities_null, "^us public model entities must be a sequence$"), + (_make_entity_flag_text, "^us entity person is_person must be a boolean$"), + ( + _remove_person_entity, + "^us public model must have exactly one person entity$", + ), + (_repeat_entity_key, "^us entity key 'household' is repeated$"), + ( + _remove_basic_inputs, + "^the installed PolicyEngine.py does not publish catalog format 2's " + "fields for us: model version basic_inputs$", + ), + ( + _make_basic_inputs_text, + "^us public model basic_inputs must be a sequence$", + ), + ( + _remove_variable_enum_values, + "^the installed PolicyEngine.py does not publish catalog format 2's " + "fields for us: variable enum_values$", + ), + ( + _make_enum_label_nan, + "^variable employment_income enum label for SINGLE is not " + "JSON-compatible: NaN is not JSON-compatible$", + ), + ( + _name_unknown_basic_input, + "^us basic input 'unpublished_input' is not a published variable$", + ), + ( + _remove_variable_unit, + "^the installed PolicyEngine.py does not publish catalog format 2's " + "fields for us: variable unit$", + ), + ( + _make_unit_a_number, + "^variable employment_income unit must be a string or null$", + ), + ( + _make_definition_period_a_list, + "^variable employment_income definition_period must be a string or null$", + ), + ( + _make_input_flag_text, + "^variable employment_income is_input_variable must be a boolean or null$", + ), + ( + _make_hidden_flag_number, + "^variable employment_income hidden_input must be a boolean or null$", + ), + ( + _name_unknown_adds_parameter, + "^variable employment_income adds_parameter_path 'gov.example.missing' " + "is not a published parameter$", + ), + ( + _name_unknown_subtracts_parameter, + "^variable employment_income subtracts_parameter_path " + "'gov.example.missing' is not a published parameter$", + ), + ( + _mismatch_enum_values, + "^variable employment_income enum values do not match its possible values$", + ), + ( + _give_enum_labels_without_values, + "^variable employment_income enum values do not match its possible values$", + ), + ( + _make_enum_label_an_object, + "^variable employment_income enum label for SINGLE must be a JSON scalar$", + ), + ( + _make_economy_flag_missing, + "^the installed PolicyEngine.py does not publish catalog format 2's " + "fields for us: parameter economy$", + ), + ( + _make_economy_flag_null, + "^parameter gov.example.rate economy must be a boolean$", + ), + ( + _make_household_flag_text, + "^parameter gov.example.rate household must be a boolean$", + ), + ], +) +def test_inconsistent_format_2_metadata_fails_with_a_specific_message( + break_models: Callable[[dict[str, SimpleNamespace]], None], + message: str, +) -> None: + models = source_models() + break_models(models) + + with pytest.raises(CatalogExtractionError, match=message): + _extract(models) + + +def test_a_release_without_format_2_fields_fails_naming_each_missing_field() -> None: + """A release without the fields fails before any record is checked, naming each. + + PolicyEngine.py 6.7.5 and 6.11.1, for example, publish only the + variables' unit and definition period of these fields. + """ + models = source_models() + model = models["us"] + del model.entities, model.basic_inputs + for variable in model.variables: + for field_name in ( + "is_input_variable", + "hidden_input", + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "enum_values", + "unit", + "definition_period", + ): + delattr(variable, field_name) + for parameter in model.parameters: + del parameter.economy, parameter.household + + with pytest.raises( + CatalogExtractionError, + match=r"^the installed PolicyEngine.py does not publish catalog format 2's " + r"fields for us: model version entities, basic_inputs; variable " + r"is_input_variable, hidden_input, module_name, adds_parameter_path, " + r"subtracts_parameter_path, enum_values, unit, definition_period; parameter " + r"economy, household$", + ): + _extract(models) + + +def test_an_empty_basic_input_list_is_stored_empty() -> None: + models = source_models() + models["us"].basic_inputs = [] + + assert _extract(models).country("us").model_version.basic_inputs == () + + +FORMAT_2_VARIABLE_TEXT_FIELDS = ( + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + "unit", + "definition_period", +) + + +def _variable_column_length(field_name: str) -> int: + from policyengine_api.data.v2.models import Variable + + return Variable.__table__.c[field_name].type.length + + +def test_variable_text_limits_are_the_variable_column_lengths() -> None: + """Extraction checks each format 2 text field against its column's length.""" + from policyengine_api.data.v2.catalog.extraction import VARIABLE_TEXT_LIMITS + + column_lengths = { + field_name: _variable_column_length(field_name) + for field_name in FORMAT_2_VARIABLE_TEXT_FIELDS + } + + assert column_lengths == VARIABLE_TEXT_LIMITS + + +@pytest.mark.parametrize("field_name", FORMAT_2_VARIABLE_TEXT_FIELDS) +def test_an_over_long_format_2_text_field_stops_extraction(field_name: str) -> None: + length = _variable_column_length(field_name) + models = source_models() + setattr(models["us"].variables[0], field_name, "x" * (length + 1)) + + with pytest.raises( + CatalogExtractionError, + match=rf"^variable employment_income {field_name} exceeds {length} characters$", + ): + _extract(models) diff --git a/tests/unit/v2/test_catalog_publication.py b/tests/unit/v2/test_catalog_publication.py index a3cbe0d65..ec7ef503b 100644 --- a/tests/unit/v2/test_catalog_publication.py +++ b/tests/unit/v2/test_catalog_publication.py @@ -213,12 +213,32 @@ def test_publication_uses_mapped_tables_and_sqlalchemy_statements() -> None: assert all("VALUES" not in statement for statement in compiled_inserts) catalog = normalized_catalog() - comparisons = publication_reconciliation._comparison_pairs(catalog.country("us")) - assert all( - isinstance(statement, sa.sql.Select) and not isinstance(statement, TextClause) - for pair in comparisons - for statement in pair + format_1_comparisons, format_2_comparisons = ( + publication_reconciliation._comparison_pairs( + catalog.country("us"), with_format_2=with_format_2 + ) + for with_format_2 in (False, True) ) + for comparisons in (format_1_comparisons, format_2_comparisons): + assert all( + isinstance(statement, sa.sql.Select) + and not isinstance(statement, TextClause) + for pair in comparisons + for statement in pair + ) + assert all( + len(staged.selected_columns) == len(stored.selected_columns) + for staged, stored in comparisons + ) + # Format 2 compares model versions, variables and parameters on more fields; + # the pairs are model versions, variables, parameter nodes, parameters, + # parameter values, datasets and regions. + assert [ + len(format_2[0].selected_columns) - len(format_1[0].selected_columns) + for format_1, format_2 in zip( + format_1_comparisons, format_2_comparisons, strict=True + ) + ] == [3, 8, 0, 2, 0, 0, 0] class Result: def scalar_one(self): diff --git a/tests/unit/v2/test_catalog_staging_rows.py b/tests/unit/v2/test_catalog_staging_rows.py index f43f2772f..e8b6f0bd4 100644 --- a/tests/unit/v2/test_catalog_staging_rows.py +++ b/tests/unit/v2/test_catalog_staging_rows.py @@ -3,13 +3,16 @@ from __future__ import annotations import inspect +from dataclasses import replace from psycopg.types.json import Jsonb from policyengine_api.data.v2.catalog.publication_staging import ( + STAGE_CATALOG_VARIABLES, STAGING_TABLES, _catalog_rows, ) +from policyengine_api.data.v2.catalog.records import EnumValueRecord from tests.fixtures.v2_catalog import normalized_catalog @@ -49,6 +52,15 @@ def test_catalog_rows_project_every_record_field_in_column_order() -> None: country.model_version.description, country.model_version.current_law_id, ("jsonb", country.model_version.metadata_time_periods), + 2, + ( + "jsonb", + [ + {"key": "person", "plural": "people", "is_person": True}, + {"key": "household", "plural": "households", "is_person": False}, + ], + ), + ("jsonb", ["employment_income"]), ) ] assert rows["stage_catalog_variables"] == [ @@ -66,6 +78,14 @@ def test_catalog_rows_project_every_record_field_in_column_order() -> None: ("jsonb", record.default_value), None if record.adds is None else ("jsonb", record.adds), None if record.subtracts is None else ("jsonb", record.subtracts), + record.is_input_variable, + record.hidden_input, + record.module_name, + record.adds_parameter_path, + record.subtracts_parameter_path, + None, + record.unit, + record.definition_period, ) for record in country.variables ] @@ -82,6 +102,8 @@ def test_catalog_rows_project_every_record_field_in_column_order() -> None: record.description, record.data_type, record.unit, + record.economy, + record.household, ) for record in country.parameters ] @@ -121,3 +143,24 @@ def test_catalog_rows_project_every_record_field_in_column_order() -> None: ] assert len(rows["stage_catalog_parameter_values"]) == 2 assert len(rows["stage_catalog_regions"]) == 3 + + +def test_enum_values_are_staged_as_value_and_label_objects_in_order() -> None: + country = normalized_catalog().country("us") + variable = replace( + country.variables[0], + possible_values=["SINGLE", "JOINT"], + enum_values=( + EnumValueRecord(value="SINGLE", label="Single"), + EnumValueRecord(value="JOINT", label=2), + ), + ) + country = replace(country, variables=(variable,)) + + (row,) = (_plain(row) for row in _catalog_rows(country)[STAGE_CATALOG_VARIABLES]) + + staged_column_names = [column.name for column in STAGE_CATALOG_VARIABLES.columns] + assert row[staged_column_names.index("enum_values")] == ( + "jsonb", + [{"value": "SINGLE", "label": "Single"}, {"value": "JOINT", "label": 2}], + ) diff --git a/tests/unit/v2/test_models.py b/tests/unit/v2/test_models.py index baa0e2590..72a986c00 100644 --- a/tests/unit/v2/test_models.py +++ b/tests/unit/v2/test_models.py @@ -4,6 +4,7 @@ from sqlalchemy.dialects import postgresql from sqlalchemy.orm import configure_mappers from sqlalchemy.schema import CreateTable +from sqlmodel.sql.sqltypes import AutoString from policyengine_api.data.v1_models import V1Base from policyengine_api.data.v2.models import ( @@ -719,3 +720,78 @@ def test_v1_copy_mappings_keep_their_us_and_uk_checks() -> None: assert "countries.id" not in { foreign_key.target_fullname for foreign_key in table.foreign_keys } + + +def test_catalog_format_2_columns_are_nullable_for_earlier_versions() -> None: + """Catalog format 2 stores v3's metadata; versions published before it keep null. + + The model version records its catalog format, entities and basic inputs; each + variable its input flags, module, parameter-named parts and enum labels; each + parameter its economy and household flags. + """ + expected_columns = { + "tax_benefit_model_versions": { + "catalog_format": sa.Integer, + "entities": sa.JSON, + "basic_inputs": sa.JSON, + }, + "variables": { + "is_input_variable": sa.Boolean, + "hidden_input": sa.Boolean, + "module_name": AutoString, + "adds_parameter_path": AutoString, + "subtracts_parameter_path": AutoString, + "enum_values": sa.JSON, + "unit": AutoString, + "definition_period": AutoString, + }, + "parameters": { + "economy": sa.Boolean, + "household": sa.Boolean, + }, + } + + for table_name, columns in expected_columns.items(): + table = V2_METADATA.tables[table_name] + for column_name, column_type in columns.items(): + column = table.c[column_name] + assert column.nullable, f"{table_name}.{column_name}" + assert column.server_default is None, f"{table_name}.{column_name}" + assert isinstance(column.type, column_type), f"{table_name}.{column_name}" + + variables = V2_METADATA.tables["variables"] + for column_name in ( + "module_name", + "adds_parameter_path", + "subtracts_parameter_path", + ): + assert variables.c[column_name].type.length == 512 + assert variables.c["unit"].type.length == 128 + assert variables.c["definition_period"].type.length == 32 + + +def test_catalog_format_2_json_columns_have_typed_shapes() -> None: + """The entity and enum-value columns name typed shapes the catalog records match. + + Readers of these columns get the keys and value types without re-checking. + """ + from typing import get_type_hints + + from policyengine_api.data.v2.catalog.records import EntityRecord, EnumValueRecord + from policyengine_api.data.v2.models import ( + CatalogEntity, + CatalogEnumValue, + Variable, + ) + + assert TaxBenefitModelVersion.model_fields["entities"].annotation == ( + list[CatalogEntity] | None + ) + assert Variable.model_fields["enum_values"].annotation == ( + list[CatalogEnumValue] | None + ) + for typed_shape, record in ( + (CatalogEntity, EntityRecord), + (CatalogEnumValue, EnumValueRecord), + ): + assert get_type_hints(typed_shape) == get_type_hints(record) diff --git a/uv.lock b/uv.lock index f415f60f0..321628b5f 100644 --- a/uv.lock +++ b/uv.lock @@ -2712,7 +2712,7 @@ wheels = [ [[package]] name = "policyengine" -version = "6.11.1" +version = "6.12.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "h5py" }, @@ -2724,9 +2724,9 @@ dependencies = [ { name = "pydantic" }, { name = "requests" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/34/5c/b0900e634ea1073426877e929850816914a182eb1cecacd2e99d705be63a/policyengine-6.11.1.tar.gz", hash = "sha256:af277f572638377522834f7335e6668905e3c414859c61e8ea4f201b522a8625", size = 1549029, upload-time = "2026-10-27T01:31:08.885Z" } +sdist = { url = "https://files.pythonhosted.org/packages/c4/ed/6f1b85109092d740a8121c98933c4848af6e6355241012688b13347a7d97/policyengine-6.12.0.tar.gz", hash = "sha256:b8d4a6346206015c508f416f2e48e150a37df3d3103966e0b294d4f15e269cf6", size = 1590873, upload-time = "2026-11-02T14:23:39.916Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/79/8788bb21fbcdd5404864656daa9d92b06d9fb5314bfc9b598a780dbcb1ee/policyengine-6.11.1-py3-none-any.whl", hash = "sha256:9ada4083d27ee4f4ca637335371186e3b3fd1607e00980c03a14d293e385080b", size = 539750, upload-time = "2026-10-27T01:31:07.018Z" }, + { url = "https://files.pythonhosted.org/packages/3a/8e/cdae9cd07eea17b33ee682249c758326fdfb4f9c837bec6832484c7428c2/policyengine-6.12.0-py3-none-any.whl", hash = "sha256:de08e9eb4e2de9e12010d1bff0a224e10658cce0c7dafee4783eec46978edb59", size = 557145, upload-time = "2026-11-02T14:23:37.915Z" }, ] [package.optional-dependencies] @@ -2818,7 +2818,7 @@ requires-dist = [ { name = "mypy", marker = "extra == 'dev'", specifier = ">=1.15,<2" }, { name = "openai" }, { name = "packaging", specifier = ">=24,<27" }, - { name = "policyengine", extras = ["models"], specifier = "==6.11.1" }, + { name = "policyengine", extras = ["models"], specifier = "==6.12.0" }, { name = "policyengine-observability", extras = ["flask", "google", "httpx", "otlp-grpc"], specifier = ">=3.0.3,<4" }, { name = "psycopg", extras = ["binary"], specifier = ">=3.3,<4" }, { name = "pydantic" },