diff --git a/.github/bump_version.py b/.github/bump_version.py index 1e7b49e18..e29c6bf8c 100644 --- a/.github/bump_version.py +++ b/.github/bump_version.py @@ -6,6 +6,7 @@ def get_current_version(pyproject_path: Path) -> str: + """Return the version in pyproject.toml, exiting with status 1 if it is absent.""" text = pyproject_path.read_text() match = re.search(r'^version\s*=\s*"(\d+\.\d+\.\d+)"', text, re.MULTILINE) if not match: @@ -18,6 +19,11 @@ def get_current_version(pyproject_path: Path) -> str: def infer_bump(changelog_dir: Path) -> str: + """Return "major", "minor" or "patch" from the pending changelog fragment types. + + Any breaking fragment makes a major bump and any added or removed fragment a + minor one; exit with status 1 when there are no fragments. + """ fragments = [ f for f in changelog_dir.iterdir() if f.is_file() and f.name != ".gitkeep" ] @@ -39,6 +45,7 @@ def infer_bump(changelog_dir: Path) -> str: def bump_version(version: str, bump: str) -> str: + """Return ``version`` raised by one ``bump`` level, resetting the lower parts.""" major, minor, patch = (int(x) for x in version.split(".")) if bump == "major": return f"{major + 1}.0.0" @@ -48,6 +55,7 @@ def bump_version(version: str, bump: str) -> str: def update_file(path: Path, old_version: str, new_version: str) -> None: + """Rewrite the ``version = "..."`` line in ``path``; leave other files untouched.""" text = path.read_text() updated = text.replace( f'version = "{old_version}"', @@ -59,6 +67,7 @@ def update_file(path: Path, old_version: str, new_version: str) -> None: def main() -> None: + """Bump the version in pyproject.toml according to the changelog.d fragments.""" root = Path(__file__).resolve().parent.parent pyproject = root / "pyproject.toml" changelog_dir = root / "changelog.d" diff --git a/.github/find-api-model-versions.py b/.github/find-api-model-versions.py index 0c0096db2..6f80d3f57 100644 --- a/.github/find-api-model-versions.py +++ b/.github/find-api-model-versions.py @@ -1,3 +1,9 @@ +"""CI helper that reads the versions pinned by the installed PolicyEngine bundle. + +It checks the package and data versions against the API's version constants, +then writes them to GITHUB_ENV or, with --shell, prints shell assignments. +""" + import argparse import os import shlex @@ -48,9 +54,7 @@ def _data_release_version(bundle: dict, country_id: str) -> str: def find_api_model_versions() -> dict[str, str]: - """ - Find and validate the package and data versions in the installed bundle. - """ + """Find and validate the package and data versions in the installed bundle.""" bundle = get_current_bundle() if not isinstance(bundle, dict): raise RuntimeError("Installed PolicyEngine bundle manifest is not an object.") @@ -89,9 +93,7 @@ def find_api_model_versions() -> dict[str, str]: def find_api_model_versions_and_output_to_github() -> None: - """ - Find the API model versions and output them to a file for GitHub. - """ + """Find the API model versions and output them to a file for GitHub.""" versions = find_api_model_versions() with open(os.environ["GITHUB_ENV"], "a") as f: for key, value in versions.items(): diff --git a/.github/scripts/check_updates.py b/.github/scripts/check_updates.py index eccbc982e..e3c232e9f 100644 --- a/.github/scripts/check_updates.py +++ b/.github/scripts/check_updates.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 -""" -Fetch and format changelog entries between two versions of a package. +"""Fetch and format changelog entries between two versions of a package. Fetches CHANGELOG.md from the package's GitHub repo and extracts entries between the old and new version, formatted as markdown. @@ -25,6 +24,8 @@ class ChangelogEntry(TypedDict): + """One released version's changelog items, grouped by lowercase category.""" + version: str changes: dict[str, list[str]] @@ -119,6 +120,10 @@ def format_changes(entries: list[ChangelogEntry]) -> str: def main() -> int: + """Print the markdown changes between two versions of a package. + + Always return 0, so an unreachable changelog never fails the update job. + """ parser = argparse.ArgumentParser( description="Fetch changelog entries between two package versions." ) diff --git a/.github/scripts/verify_observability_collector.py b/.github/scripts/verify_observability_collector.py index 8156a70bc..633c7213f 100644 --- a/.github/scripts/verify_observability_collector.py +++ b/.github/scripts/verify_observability_collector.py @@ -64,6 +64,8 @@ class VerificationProbe(NamedTuple): + """Identifiers of one verification export, for finding its stored telemetry.""" + verification_id: str trace_id: str started_at: float @@ -315,7 +317,6 @@ def verify( delivery_timeout_seconds: float = 600.0, ) -> VerificationProbe: """Export real telemetry, verify storage, and require log rejection.""" - started_at = time.time() verification_id = f"collector-verify-{uuid.uuid4().hex}" trace_request, metric_request, log_request, trace_id = _probe_requests( @@ -370,6 +371,11 @@ def verify( def main() -> None: + """Verify the collector, reading both tokens from the environment. + + COLLECTOR_ID_TOKEN authenticates to the collector and GOOGLE_OAUTH_ACCESS_TOKEN + reads Cloud Trace and Cloud Monitoring. + """ parser = argparse.ArgumentParser() parser.add_argument("--endpoint", required=True) parser.add_argument("--project-id", required=True) diff --git a/AGENTS.md b/AGENTS.md index 6c2a17c8a..cb75d95ec 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -123,7 +123,7 @@ To add or change a rule, read `docs/engineering/skills/rules-registry.md`. - `api.quality.environment-variables-load`: Required environment variables and credentials load in a deployed-like environment. (blocking: CI job `test_env_vars`) - `api.quality.mypy-v2`: Shared query models, FastAPI dependencies, v2 routes and services, and the user sign-in and account-deletion libraries pass mypy with untyped definitions disallowed. (blocking: mypy) - `api.quality.ruff-format`: All Python files are formatted with ruff format. (blocking: ruff format) -- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)) +- `api.quality.ruff-lint`: All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. (blocking: ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)) - `api.quality.runtime-evaluated-models-resolve`: Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. (blocking: pytest `tests/unit/test_runtime_evaluated_annotations.py`) - `api.quality.tests-pass`: The unit, contract and to-refactor test suites pass. (blocking: CI job `test`) - `api.quality.tests-ruff-exemptions`: Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. (blocking: ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**`) diff --git a/changelog.d/ruff-docstrings.changed.md b/changelog.d/ruff-docstrings.changed.md new file mode 100644 index 000000000..cbdc5aa29 --- /dev/null +++ b/changelog.d/ruff-docstrings.changed.md @@ -0,0 +1 @@ +Enabled ruff's docstring rules (D) and documented every module, class and function outside tests; published models and routes now carry these descriptions in their JSON schema and OpenAPI, with no structural schema change. diff --git a/dashboard/experiments/gpt4_api_user.py b/dashboard/experiments/gpt4_api_user.py index 27b21f890..1fd37047b 100644 --- a/dashboard/experiments/gpt4_api_user.py +++ b/dashboard/experiments/gpt4_api_user.py @@ -1,3 +1,10 @@ +"""Streamlit experiment that answers a UK policy question with GPT-4 and the API. + +It finds related parameters by embedding similarity, asks GPT-4 to write the +reform as JSON, runs it on the public PolicyEngine API, and asks GPT-4 to answer +from the resulting economic impact. +""" + import json import os import time @@ -14,6 +21,7 @@ @st.cache_data def get_embeddings_df() -> pd.DataFrame: + """Load parameter_embeddings.csv.gz, evaluating each embedding string to a list.""" embeddings_df = pd.read_csv("parameter_embeddings.csv.gz", compression="gzip") embeddings_df.parameter_embedding = embeddings_df.parameter_embedding.apply( lambda x: eval(x) @@ -23,6 +31,7 @@ def get_embeddings_df() -> pd.DataFrame: @st.cache_data def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]: + """Return the OpenAI embedding of ``prompt``, cached across Streamlit reruns.""" return get_embedding(prompt, engine=engine) @@ -59,6 +68,7 @@ def embed(prompt: str, engine: str = "text-embedding-ada-002") -> list[float]: def write_prompt(question: str, relevant_parameters: object) -> str: + """Build the GPT-4 prompt that rewrites ``question`` as a reform JSON object.""" return f"""Question: {question} Here's some metadata about the parameters that are relevant to your question: Relevant parameters: {relevant_parameters} @@ -79,6 +89,7 @@ def write_prompt(question: str, relevant_parameters: object) -> str: @st.cache_data def get_gpt4_response(prompt: str) -> str: + """Return GPT-4's reply to ``prompt`` as one user message, cached across reruns.""" return openai.ChatCompletion.create( model="gpt-4", messages=[{"role": "user", "content": prompt}], @@ -94,6 +105,10 @@ def get_gpt4_response(prompt: str) -> str: @st.cache_data def get_policyengine_impact(reform: object) -> object: + """Create ``reform`` on the public UK API and return its 2023 UK economic impact. + + Poll every five seconds, with no time limit, while the impact is computing. + """ payload = {"data": reform} url = "https://api.policyengine.org/uk/policy" diff --git a/docs/generated/rules.md b/docs/generated/rules.md index 22059e6d1..b28af3c56 100644 --- a/docs/generated/rules.md +++ b/docs/generated/rules.md @@ -272,7 +272,7 @@ flowchart LR | `api.quality.environment-variables-load` | Required environment variables and credentials load in a deployed-like environment. | mechanical | enforced | yes | CI job `test_env_vars` | `pr.yml:test_env_vars` | — | | `api.quality.mypy-v2` | Shared query models, FastAPI dependencies, v2 routes and services, and the user sign-in and account-deletion libraries pass mypy with untyped definitions disallowed. | mechanical | enforced | yes | mypy | `pr.yml:typecheck-v2` | — | | `api.quality.ruff-format` | All Python files are formatted with ruff format. | mechanical | enforced | yes | ruff format | `pr.yml:lint` | — | -| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — | +| `api.quality.ruff-lint` | All Python code passes ruff check with every rule family selected in pyproject.toml, with no ratchet; families are added one themed group at a time until the selection is ALL. | mechanical | enforced | yes | ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W) | `pr.yml:lint` | — | | `api.quality.runtime-evaluated-models-resolve` | Every package model is one ruff treats as runtime-evaluated (a base resolves, through the import in its own module, to a name in runtime-evaluated-base-classes, or is a same-module class that does), and every Pydantic model and SQLAlchemy mapper resolves its annotations once the package is imported. | mechanical | enforced | yes | pytest `tests/unit/test_runtime_evaluated_annotations.py` | `pr.yml:test` | **unproven** | | `api.quality.tests-pass` | The unit, contract and to-refactor test suites pass. | mechanical | enforced | yes | CI job `test` | `pr.yml:test` | — | | `api.quality.tests-ruff-exemptions` | Test files may use pytest's bare assert and need no docstrings, annotations or named constants for magic numbers (S101, D1, ANN, PLR2004); every other ruff rule applies to tests. | mechanical | enforced | yes | ruff per-file ignore of ANN, D1, PLR2004, S101 for `**/tests/**` | `pr.yml:lint` | — | @@ -292,7 +292,7 @@ flowchart LR e_ruff_format["ruff format"] j_CI_pr_yml_lint["CI pr.yml:lint"] r_api_quality_ruff_lint("api.quality.ruff-lint (enforced)") - e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"] + e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_["ruff check (ANN, D, E4, E7, E9, F, FLY, FURB, I, PIE, PYI, RET, RSE, TC, UP, W)"] r_api_quality_runtime_evaluated_models_resolve("api.quality.runtime-evaluated-models-resolve (enforced)") e_pytest_tests_unit_test_runtime_evaluated_annotations_py["pytest tests/unit/test_runtime_evaluated_annotations.py"] j_CI_pr_yml_test["CI pr.yml:test"] @@ -308,8 +308,8 @@ flowchart LR e_mypy --> j_CI_pr_yml_typecheck_v2 r_api_quality_ruff_format --> e_ruff_format e_ruff_format --> j_CI_pr_yml_lint - r_api_quality_ruff_lint --> e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ - e_ruff_check__ANN__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint + r_api_quality_ruff_lint --> e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ + e_ruff_check__ANN__D__E4__E7__E9__F__FLY__FURB__I__PIE__PYI__RET__RSE__TC__UP__W_ --> j_CI_pr_yml_lint r_api_quality_runtime_evaluated_models_resolve --> e_pytest_tests_unit_test_runtime_evaluated_annotations_py e_pytest_tests_unit_test_runtime_evaluated_annotations_py --> j_CI_pr_yml_test r_api_quality_tests_pass --> e_CI_job_test diff --git a/migrations/v1/env.py b/migrations/v1/env.py index ab4b6983e..c2ce10a5c 100644 --- a/migrations/v1/env.py +++ b/migrations/v1/env.py @@ -40,6 +40,7 @@ def run_migrations_offline() -> None: + """Emit the migration SQL for the configured MySQL URL without connecting.""" context.configure( url=config.get_main_option("sqlalchemy.url"), target_metadata=target_metadata, @@ -52,6 +53,7 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: + """Migrate the connection the caller supplied, or else a new configured engine.""" if provided_connection is not None: _run_migrations_with_connection(provided_connection) return diff --git a/migrations/v2/env.py b/migrations/v2/env.py index 17b8e7d6d..302ca39c5 100644 --- a/migrations/v2/env.py +++ b/migrations/v2/env.py @@ -38,7 +38,6 @@ def _include_application_object( _compare_to: SchemaItem | None, ) -> bool: """Keep Alembic's own version table outside application drift.""" - return not (type_ == "table" and name == "alembic_version") @@ -65,6 +64,7 @@ def _configure(connection: Connection) -> None: def run_migrations_offline() -> None: + """Raise, because v2 migrations need a live connection to qualify the target.""" raise RuntimeError( "v2 migrations require an online connection so target identity and " "generated application-data before/after states can be verified" @@ -72,6 +72,11 @@ def run_migrations_offline() -> None: def run_migrations_online() -> None: + """Migrate the caller's connection, or else a new engine in one transaction. + + The new engine's transaction commits, so the target qualification's reads + cannot leave the completed migration to be rolled back. + """ provided_connection = config.attributes.get("connection") if provided_connection is not None: _configure(provided_connection) diff --git a/policyengine_api/api.py b/policyengine_api/api.py index 81fc92106..df4fea71c 100644 --- a/policyengine_api/api.py +++ b/policyengine_api/api.py @@ -1,7 +1,5 @@ # ruff: noqa: E402 -""" -This is the main Flask app for the PolicyEngine API. -""" +"""The PolicyEngine API Flask application, with its blueprints and extensions.""" import sys import time @@ -10,6 +8,7 @@ def log_timing(message: str) -> None: + """Print a startup milestone to stderr with the seconds since import began.""" elapsed = time.time() - start_time print(f"[{elapsed:.2f}s] {message}", file=sys.stderr, flush=True) diff --git a/policyengine_api/asgi_factory.py b/policyengine_api/asgi_factory.py index 774b0f268..7663d3968 100644 --- a/policyengine_api/asgi_factory.py +++ b/policyengine_api/asgi_factory.py @@ -99,7 +99,6 @@ def _matched_route_template(routes: Iterable[BaseRoute], scope: Scope) -> str | def _native_route_template(app: FastAPI, scope: Scope) -> str | None: """Return the matched FastAPI route template, excluding the Flask mount.""" - return _matched_route_template(app.router.routes, scope) @@ -112,7 +111,6 @@ def create_asgi_app( observability_runtime: ObservabilityRuntime | None = None, ) -> ASGIApp: """Create the Stage 2 FastAPI shell around the existing Flask app.""" - if route_settings is None: route_settings = RouteImplementationSettings.from_environment() if dependencies is None: diff --git a/policyengine_api/constants.py b/policyengine_api/constants.py index cea6ecd7d..b17a41b9f 100644 --- a/policyengine_api/constants.py +++ b/policyengine_api/constants.py @@ -1,3 +1,5 @@ +"""API version, country package versions and bundle metadata, resolved at import.""" + import hashlib import json from datetime import datetime @@ -49,6 +51,7 @@ def _resolve_distribution_version( def get_py_manifest() -> Any: # noqa: ANN401 # SimplePath, which open() rejects; a Path at runtime + """Return the path of the bundle manifest in the installed policyengine wheel.""" return distribution("policyengine").locate_file( "policyengine/data/bundle/manifest.json" ) @@ -89,6 +92,7 @@ def _resolve_bundle_package_versions(bundle: dict | None) -> dict[str, str]: def get_bundle_data_release(country_id: str) -> dict: + """Return the bundle's data release entry for a country, or {} if it has none.""" if not isinstance(_policyengine_bundle, dict): return {} @@ -104,6 +108,7 @@ def get_bundle_data_release(country_id: str) -> dict: def get_bundle_default_dataset(country_id: str) -> str | None: + """Return the bundle's default dataset name for a country, or None if unset.""" release = get_bundle_data_release(country_id) default_dataset = release.get("default_dataset") if default_dataset is None: @@ -116,6 +121,11 @@ def _bundle_dataset_display_label(default_dataset: object) -> str: def get_bundle_default_dataset_option(country_id: str) -> dict: + """Build the metadata dataset option for the bundle's certified default dataset. + + The option carries ``dataset_uri`` and ``data_version`` only when the bundle + records them. + """ release = get_bundle_data_release(country_id) default_dataset = release.get("default_dataset") label = _bundle_dataset_display_label(default_dataset) @@ -185,10 +195,10 @@ def get_bundle_default_dataset_option(country_id: str) -> dict: def _build_runtime_cache_version( scope: str, country_id: str, caller_version: str | None = None ) -> str: - """ - Build a compact version token for cache keys stored in legacy VARCHAR(10) - columns. The token changes whenever the relevant runtime or payload schema - changes, even if the country package version is unchanged. + """Build a compact cache version token that fits legacy VARCHAR(10) columns. + + The token changes whenever the relevant runtime or payload schema changes, + even if the country package version is unchanged. """ schema_version = str(RUNTIME_CACHE_SCHEMA_VERSIONS[scope]) prefix = "e" if scope == "economy_impact" else "r" @@ -216,10 +226,16 @@ def _build_runtime_cache_version( def get_economy_impact_cache_version( country_id: str, caller_version: str | None = None ) -> str: + """Return the cache version stored with economy impacts for this country. + + ``caller_version`` is the country model version the request asked for; it + defaults to the installed country package version. + """ return _build_runtime_cache_version("economy_impact", country_id, caller_version) def get_report_output_cache_version(country_id: str) -> str: + """Return the cache version that marks a report output current for this runtime.""" return _build_runtime_cache_version("report_output", country_id) diff --git a/policyengine_api/country.py b/policyengine_api/country.py index 332871fae..ec0499113 100644 --- a/policyengine_api/country.py +++ b/policyengine_api/country.py @@ -1,3 +1,5 @@ +"""Loaded country tax-benefit systems, their metadata and household calculation.""" + import importlib import inspect import json @@ -121,7 +123,10 @@ def _record_calculation_failure( class PolicyEngineCountry: + """A country package's tax-benefit system with its prebuilt metadata document.""" + def __init__(self, country_package_name: str, country_id: str) -> None: + """Import the country package, build its tax-benefit system and its metadata.""" self.country_package_name = country_package_name self.country_id = country_id self.country_package = importlib.import_module(country_package_name) @@ -131,6 +136,7 @@ def __init__(self, country_package_name: str, country_id: str) -> None: self.build_metadata() def build_metadata(self) -> None: + """Build the JSON-safe metadata document served by the metadata endpoints.""" self.metadata = self._json_safe( dict( variables=self.build_variables(), @@ -172,6 +178,7 @@ def _json_safe( return value def build_microsimulation_options(self) -> dict: + """Build the regions, time periods and datasets offered for economy runs.""" # { region: [{ name: "uk", label: "the UK" }], time_period: [{ name: 2022, label: "2022", ... }] } options = dict() if self.country_id == "uk": @@ -334,6 +341,7 @@ def build_microsimulation_options(self) -> dict: return options def build_variables(self) -> dict: + """Describe every variable by name; enum variables list their values.""" variables = self.tax_benefit_system.variables variable_data = {} for variable_name, variable in variables.items(): @@ -369,6 +377,7 @@ def build_variables(self) -> dict: return variable_data def build_parameters(self) -> dict: + """Describe every parameter under ``gov``, with leaf parameters' values.""" parameters = self.tax_benefit_system.parameters parameter_data = {} for parameter in parameters.get_descendants(): @@ -428,6 +437,7 @@ def build_parameters(self) -> dict: return parameter_data def build_entities(self) -> dict: + """Describe every entity by key, with its roles (empty when it has none).""" data = {} for entity in self.tax_benefit_system.entities: entity_data = { @@ -557,7 +567,6 @@ def prepare_calculation( spm_requested: bool = False, ) -> PreparedCountryCalculation: """Parse a situation before the API accepts it as a calculation.""" - simulation, system = self._create_simulation(household, reform, spm=spm) household = json.loads(json.dumps(household)) has_axes = "axes" in household @@ -648,14 +657,14 @@ def _apply_reform_to_system( def create_policy_reform(policy_data: dict) -> dict: - """ - Create a policy reform. + """Create a policy reform. Args: policy_data (dict): The policy data. Returns: dict: The reform. + """ def modify_parameters(parameters: ParameterNode) -> ParameterNode: @@ -696,6 +705,12 @@ def get_requested_computations( include_provided_values: bool = False, variable_names: set[str] | None = None, ) -> list[tuple[str, str, str, str]]: + """List the (entity plural, entity id, variable, period) cells to compute. + + A cell is requested when its value is None, or, with + ``include_provided_values``, whatever its value. ``variable_names`` keeps only + cells for those variables. + """ requested_computations = dpath.util.search( household, "*/*/*/*", diff --git a/policyengine_api/country_validation.py b/policyengine_api/country_validation.py index 28b08a280..a716ab60f 100644 --- a/policyengine_api/country_validation.py +++ b/policyengine_api/country_validation.py @@ -8,6 +8,8 @@ class CountryErrorPayload(TypedDict): + """The JSON error body returned for an unsupported country ID.""" + status: Literal["error"] message: str @@ -20,6 +22,7 @@ def __init__( country_id: str, available_country_ids: tuple[str, ...] = COUNTRIES, ) -> None: + """Record the rejected ID and the supported IDs listed in the message.""" self.country_id = country_id self.available_country_ids = available_country_ids super().__init__( @@ -28,6 +31,7 @@ def __init__( ) def to_payload(self) -> CountryErrorPayload: + """Return the ``{"status": "error", "message": ...}`` response body.""" return {"status": "error", "message": str(self)} diff --git a/policyengine_api/data/congressional_districts.py b/policyengine_api/data/congressional_districts.py index c78752304..5c6603260 100644 --- a/policyengine_api/data/congressional_districts.py +++ b/policyengine_api/data/congressional_districts.py @@ -1,5 +1,4 @@ -""" -US Congressional District Metadata +"""Metadata for every US congressional district and helpers that use it. This module defines the metadata for all 435 US Congressional districts based on the 2020 Census apportionment, effective for the 118th Congress @@ -68,8 +67,7 @@ class CongressionalDistrictMetadataItem(BaseModel): - """ - Metadata for a single US Congressional district. + """Metadata for a single US Congressional district. Uses Pydantic BaseModel for: - Runtime validation of data integrity @@ -588,8 +586,7 @@ class CongressionalDistrictMetadataItem(BaseModel): def _get_ordinal_suffix(number: int) -> str: - """ - Get the ordinal suffix for a number (st, nd, rd, th). + """Get the ordinal suffix for a number (st, nd, rd, th). Examples: 1 -> "st" @@ -599,6 +596,7 @@ def _get_ordinal_suffix(number: int) -> str: 11 -> "th" 21 -> "st" 22 -> "nd" + """ if 10 <= number % 100 <= 20: # Special case for 11th, 12th, 13th, etc. @@ -609,34 +607,35 @@ def _get_ordinal_suffix(number: int) -> str: def _format_district_number(number: int) -> str: - """ - Format district number with leading zero for single digits. + """Format district number with leading zero for single digits. Examples: 1 -> "01" 9 -> "09" 10 -> "10" 38 -> "38" + """ return f"{number:02d}" def _build_district_name(state_code: str, number: int) -> str: - """ - Build the district name in the format: congressional_district/- + """Build a district's region name, ``congressional_district/-``. Examples: ("CA", 5) -> "congressional_district/CA-05" ("TX", 38) -> "congressional_district/TX-38" ("DC", 1) -> "congressional_district/DC-01" + """ return f"congressional_district/{state_code}-{_format_district_number(number)}" def _build_district_label(state_code: str, number: int) -> str: - """ - Build the district label in the format: 's th congressional district - For at-large districts (states with only 1 district), use: 's at-large congressional district + """Build a district's display label, such as "Texas's 3rd congressional district". + + At-large districts (states with only one district) are labelled + "'s at-large congressional district" instead of numbered. Examples: ("CA", 1) -> "California's 1st congressional district" @@ -645,6 +644,7 @@ def _build_district_label(state_code: str, number: int) -> str: ("FL", 21) -> "Florida's 21st congressional district" ("AK", 1) -> "Alaska's at-large congressional district" ("WY", 1) -> "Wyoming's at-large congressional district" + """ state_name = STATE_CODE_TO_NAME[state_code] if state_code in AT_LARGE_STATES: @@ -654,8 +654,7 @@ def _build_district_label(state_code: str, number: int) -> str: def build_congressional_district_metadata() -> list[dict]: - """ - Build the complete congressional district metadata structure for use in country.py. + """Build the congressional district region entries that country.py publishes. Returns a list of dictionaries with 'name', 'label', 'type', 'state_abbreviation', and 'state_name' keys, formatted as: @@ -679,6 +678,7 @@ def build_congressional_district_metadata() -> list[dict]: Returns: List of 436 dictionaries (435 voting districts + DC) + """ return [ { @@ -693,23 +693,23 @@ def build_congressional_district_metadata() -> list[dict]: def get_valid_state_codes() -> set[str]: - """ - Get the set of valid US state codes (lowercase for case-insensitive matching). + """Get the set of valid US state codes (lowercase for case-insensitive matching). Returns: Set of 51 lowercase state codes (50 states + DC) + """ return {code.lower() for code in STATE_CODE_TO_NAME.keys()} def get_valid_congressional_districts() -> set[str]: - """ - Get the set of valid congressional district identifiers (lowercase for case-insensitive matching). + """Return every valid district identifier, lowercased for case-insensitive matching. Format: "-" (e.g., "ca-37", "tx-01") Returns: Set of 436 lowercase district identifiers + """ return { f"{district.state_code.lower()}-{_format_district_number(district.number)}" @@ -718,8 +718,7 @@ def get_valid_congressional_districts() -> set[str]: def normalize_us_region(region: str) -> str: - """ - Normalize a US region string to the standard prefixed format. + """Normalize a US region string to the standard prefixed format. This function handles legacy region formats (bare state codes like "ca") and converts them to the standard format ("state/ca"). It should be called @@ -739,6 +738,7 @@ def normalize_us_region(region: str) -> str: Note: This function does NOT validate that the region is valid - it only normalizes the format. Use _validate_us_region for validation. + """ # Already has a valid prefix - return as-is if region.startswith(("state/", "place/", "congressional_district/")): diff --git a/policyengine_api/data/orm.py b/policyengine_api/data/orm.py index 1a1ffaa67..12e9a5ec2 100644 --- a/policyengine_api/data/orm.py +++ b/policyengine_api/data/orm.py @@ -29,6 +29,11 @@ def get_remote_database_config() -> dict[str, str]: + """Read the Cloud SQL instance, user and database name from the environment. + + Raise KeyError when POLICYENGINE_DB_INSTANCE_CONNECTION_NAME is unset; the + user and database name default to ``policyengine``. + """ return { "instance_connection_name": os.environ[ "POLICYENGINE_DB_INSTANCE_CONNECTION_NAME" @@ -40,7 +45,6 @@ def get_remote_database_config() -> dict[str, str]: def build_session_factory(engine: Engine) -> sessionmaker[Session]: """Return SQLAlchemy's standard configurable Session factory.""" - return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) @@ -77,7 +81,6 @@ def get_connection() -> DBAPIConnection: def get_v1_engine() -> Engine: """Return the process-owned Cloud SQL engine; never select SQLite.""" - global _v1_engine if _v1_engine is None: _v1_engine = _build_remote_engine() @@ -86,7 +89,6 @@ def get_v1_engine() -> Engine: def get_v1_session_factory() -> sessionmaker[Session]: """Return one configured Session factory for the Cloud SQL engine.""" - global _v1_session_factory if _v1_session_factory is None: _v1_session_factory = build_session_factory(get_v1_engine()) @@ -95,14 +97,12 @@ def get_v1_session_factory() -> sessionmaker[Session]: def clear_v1_session_factories() -> None: """Forget cached factories, primarily after replacing an Engine in tests.""" - global _v1_session_factory _v1_session_factory = None def close_v1_engines() -> None: """Release process-owned SQLAlchemy pools and Cloud SQL connectors.""" - global _cloud_sql_connector, _v1_engine clear_v1_session_factories() if _v1_engine is not None: diff --git a/policyengine_api/data/places.py b/policyengine_api/data/places.py index f24467f83..3a73b950b 100644 --- a/policyengine_api/data/places.py +++ b/policyengine_api/data/places.py @@ -1,5 +1,4 @@ -""" -US Census place code parsing and validation utilities. +"""US Census place code parsing and validation utilities. Place codes follow the format: STATE_ABBREV-PLACE_FIPS Example: NJ-57000 for Newark, NJ @@ -9,8 +8,7 @@ def parse_place_code(place_code: str) -> tuple[str, str]: - """ - Parse a place code into its state abbreviation and FIPS components. + """Parse a place code into its state abbreviation and FIPS components. Args: place_code: Place code in format STATE_ABBREV-PLACE_FIPS (e.g., "NJ-57000") @@ -20,6 +18,7 @@ def parse_place_code(place_code: str) -> tuple[str, str]: Raises: ValueError: If the place code format is invalid + """ if "-" not in place_code: raise ValueError( @@ -30,14 +29,14 @@ def parse_place_code(place_code: str) -> tuple[str, str]: def validate_place_code(place_code: str) -> None: - """ - Validate a place code has valid state abbreviation and FIPS format. + """Validate a place code has valid state abbreviation and FIPS format. Args: place_code: Place code in format STATE_ABBREV-PLACE_FIPS (e.g., "NJ-57000") Raises: ValueError: If the state abbreviation or FIPS code is invalid + """ state_abbrev, place_fips = parse_place_code(place_code) diff --git a/policyengine_api/data/v1_models.py b/policyengine_api/data/v1_models.py index 79f83a173..71578d8aa 100644 --- a/policyengine_api/data/v1_models.py +++ b/policyengine_api/data/v1_models.py @@ -28,10 +28,12 @@ class V1Base(DeclarativeBase): - pass + """Declarative base whose metadata is the v1 schema Alembic compares against.""" class Household(V1Base): + """A saved v1 household input, with its JSON hash and country package version.""" + __tablename__ = "household" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -95,6 +97,11 @@ class HouseholdMirrorEvent(V1Base): class ComputedHousehold(V1Base): + """A legacy cached household calculation, keyed by household, policy and country. + + Mapped so Alembic sees the table; the API no longer reads or writes it. + """ + __tablename__ = "computed_household" household_id: Mapped[int] = mapped_column(Integer, primary_key=True) policy_id: Mapped[int] = mapped_column(Integer, primary_key=True) @@ -105,6 +112,8 @@ class ComputedHousehold(V1Base): class Policy(V1Base): + """A saved v1 policy reform, identified within a country by its id and JSON hash.""" + __tablename__ = "policy" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3), primary_key=True) @@ -115,6 +124,11 @@ class Policy(V1Base): class Economy(V1Base): + """A legacy cached economy-wide result for one policy and options hash. + + Mapped so Alembic sees the table; the API no longer reads or writes it. + """ + __tablename__ = "economy" economy_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True @@ -132,6 +146,12 @@ class Economy(V1Base): class ReformImpact(V1Base): + """A cached economy comparison of a reform against a baseline, with its job status. + + Rows are keyed by policies, region, dataset, time period, options hash and + API version, and record the simulation execution that computes them. + """ + __tablename__ = "reform_impact" reform_impact_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True @@ -154,6 +174,11 @@ class ReformImpact(V1Base): class Analysis(V1Base): + """A legacy AI-analysis prompt and its generated text. + + Mapped so Alembic sees the table; the API no longer reads or writes it. + """ + __tablename__ = "analysis" prompt_id: Mapped[int] = mapped_column( Integer, primary_key=True, autoincrement=True @@ -166,6 +191,12 @@ class Analysis(V1Base): class UserPolicy(V1Base): + """A user's saved reform/baseline pair and how they set it up. + + ``mirror_revision`` increases on every change so v2 mirroring applies + events in order. + """ + __tablename__ = "user_policies" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -253,6 +284,8 @@ class UserPolicyMirrorEvent(V1Base): class UserProfile(V1Base): + """A user's profile, unique by Auth0 id and by username when one is set.""" + __tablename__ = "user_profiles" user_id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) auth0_id: Mapped[str] = mapped_column(String(255), unique=True) @@ -262,6 +295,12 @@ class UserProfile(V1Base): class Simulation(V1Base): + """A v1 simulation of one policy on a household or geography population. + + Holds the latest output inline, plus pointers to its active and latest + successful runs in ``simulation_runs``. + """ + __tablename__ = "simulations" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -279,6 +318,12 @@ class Simulation(V1Base): class ReportOutput(V1Base): + """A v1 report comparing one or two simulations, with its latest output inline. + + Also stores the report spec and pointers to its active and latest + successful runs in ``report_output_runs``. + """ + __tablename__ = "report_outputs" id: Mapped[int] = mapped_column(Integer, primary_key=True, autoincrement=True) country_id: Mapped[str] = mapped_column(String(3)) @@ -298,6 +343,8 @@ class ReportOutput(V1Base): class ReportOutputRun(V1Base): + """One numbered execution of a report output, with the versions it ran against.""" + __tablename__ = "report_output_runs" __table_args__ = ( UniqueConstraint( @@ -332,6 +379,8 @@ class ReportOutputRun(V1Base): class SimulationRun(V1Base): + """One numbered execution of a simulation, possibly as an input to a report run.""" + __tablename__ = "simulation_runs" __table_args__ = ( UniqueConstraint( @@ -366,6 +415,8 @@ class SimulationRun(V1Base): class LegacyReportOutputAlias(V1Base): + """Redirect from a duplicate legacy report output id to its canonical one.""" + __tablename__ = "legacy_report_output_aliases" legacy_report_output_id: Mapped[int] = mapped_column( Integer, diff --git a/policyengine_api/data/v2/catalog/catalog_selection.py b/policyengine_api/data/v2/catalog/catalog_selection.py index f998cf224..f72a43192 100644 --- a/policyengine_api/data/v2/catalog/catalog_selection.py +++ b/policyengine_api/data/v2/catalog/catalog_selection.py @@ -44,7 +44,6 @@ class SelectedCatalog: def validate_policyengine_version(value: str) -> str: """Return one bounded canonical PEP 440 version string.""" - if not isinstance(value, str) or not value or value != value.strip(): raise InvalidPolicyEngineVersionError( "policyengine_version must be a non-empty canonical version" @@ -74,7 +73,6 @@ def select_catalog( policyengine_version: str | None = None, ) -> SelectedCatalog: """Select one country catalog using the requested or running package version.""" - if country_id not in SUPPORTED_V2_COUNTRY_IDS: raise UnsupportedPreviewCountryError(country_id) explicit_version = policyengine_version is not None diff --git a/policyengine_api/data/v2/catalog/extraction.py b/policyengine_api/data/v2/catalog/extraction.py index 421d4633b..4d8a24fc8 100644 --- a/policyengine_api/data/v2/catalog/extraction.py +++ b/policyengine_api/data/v2/catalog/extraction.py @@ -142,7 +142,6 @@ def _type_name(value: object) -> str | None: def normalize_json_value(value: object) -> JSONValue: """Convert supported public values into strict JSON-compatible values.""" - if value is None or isinstance(value, (str, bool, int)): return value if isinstance(value, float): @@ -718,7 +717,6 @@ def extract_catalog( installed_version: Callable[[str], str] = importlib_metadata.version, ) -> NormalizedCatalog: """Validate and normalize supplied PolicyEngine.py public objects.""" - dependencies = _verify_bundle( bundle=bundle, policyengine_version=policyengine_version, @@ -751,7 +749,6 @@ def extract_catalog( def extract_installed_catalog() -> NormalizedCatalog: """Load and normalize the installed PolicyEngine.py certified catalog.""" - import policyengine from policyengine.bundle import get_current_bundle diff --git a/policyengine_api/data/v2/catalog/initialization.py b/policyengine_api/data/v2/catalog/initialization.py index 556519694..9fcf31c54 100644 --- a/policyengine_api/data/v2/catalog/initialization.py +++ b/policyengine_api/data/v2/catalog/initialization.py @@ -32,7 +32,6 @@ def build_data_write_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated, unpooled engine for one initialization attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -46,7 +45,6 @@ def initialize_catalog( ), ) -> PublicationEvidence: """Load only row-write settings, extract fully, then publish explicitly.""" - settings = load_v2_data_write_database_settings(environ) catalog = extractor() engine = engine_builder(settings) @@ -78,7 +76,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run explicit initialization and return a shell-compatible status.""" - try: evidence = initialize_catalog() except Exception as error: # noqa: BLE001 - command must return safe evidence diff --git a/policyengine_api/data/v2/catalog/publication.py b/policyengine_api/data/v2/catalog/publication.py index 8f33fc664..cb97ed763 100644 --- a/policyengine_api/data/v2/catalog/publication.py +++ b/policyengine_api/data/v2/catalog/publication.py @@ -93,7 +93,6 @@ def publish_catalog( checkpoint: Callable[[str, Connection], None] | None = None, ) -> PublicationEvidence: """Publish one complete catalog atomically and return non-secret evidence.""" - started = time.monotonic() with engine.begin() as connection: _verify_expected_revision(connection) diff --git a/policyengine_api/data/v2/catalog/publication_reconciliation.py b/policyengine_api/data/v2/catalog/publication_reconciliation.py index 75d6ede2a..3296bf172 100644 --- a/policyengine_api/data/v2/catalog/publication_reconciliation.py +++ b/policyengine_api/data/v2/catalog/publication_reconciliation.py @@ -48,6 +48,7 @@ def version_exists(connection: Connection, country: CountryCatalog) -> bool: + """Return whether the model name and version of this country are already stored.""" statement = ( sa.select(MODEL_VERSIONS.c.id) .select_from( @@ -284,6 +285,11 @@ def assert_country_matches( connection: Connection, country: CountryCatalog, ) -> None: + """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. + """ for staged, actual in _comparison_pairs(country): if _sets_differ(connection, staged, actual): raise CatalogPublicationError( @@ -626,11 +632,20 @@ def assert_country_matches( def publish_new_country(connection: Connection, country: CountryCatalog) -> None: + """Insert one country's staged catalog into the v2 tables, parents first. + + The model row is reused if one with the same name already exists. + """ for statement in SET_BASED_INSERT_STATEMENTS: connection.execute(statement, {"country_id": country.country_id}) def assert_canonical_value_uniqueness(connection: Connection) -> None: + """Raise CatalogPublicationError if a parameter has duplicate current-law values. + + Current-law values are those with neither a policy nor a dynamic; each + parameter may have only one per start date. + """ duplicates = ( sa.select(PARAMETER_VALUES.c.parameter_id) .where( diff --git a/policyengine_api/data/v2/catalog/publication_staging.py b/policyengine_api/data/v2/catalog/publication_staging.py index d45dcb5f6..8656b73f8 100644 --- a/policyengine_api/data/v2/catalog/publication_staging.py +++ b/policyengine_api/data/v2/catalog/publication_staging.py @@ -275,7 +275,6 @@ def copy_rows( rows: Iterable[Sequence[object]], ) -> int: """Write one bounded source stream through Psycopg COPY.""" - raw_connection = connection.connection.driver_connection statement = sql.SQL("COPY {} ({}) FROM STDIN").format( sql.Identifier(table.name), @@ -291,6 +290,7 @@ def copy_rows( def create_staging_tables(connection: Connection) -> None: + """Create the temporary staging tables, which Postgres drops at commit.""" STAGING_METADATA.create_all(connection, checkfirst=False) @@ -300,6 +300,11 @@ def stage_catalog( *, checkpoint: Callable[[str, Connection], None] | None = None, ) -> dict[str, int]: + """COPY every country's records into the staging tables and return rows per table. + + Raise CatalogPublicationError if the copied counts differ from the + catalog's own counts. + """ observed = {table.name: 0 for table in STAGING_TABLES} for country in catalog.countries: for table, rows in _catalog_rows(country).items(): diff --git a/policyengine_api/data/v2/catalog/publication_types.py b/policyengine_api/data/v2/catalog/publication_types.py index b9e388fc0..884bcef07 100644 --- a/policyengine_api/data/v2/catalog/publication_types.py +++ b/policyengine_api/data/v2/catalog/publication_types.py @@ -20,6 +20,7 @@ class PublicationEvidence: elapsed_seconds: float def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object printed after a publication.""" return { "outcome": "ok", "policyengine_version": self.policyengine_version, diff --git a/policyengine_api/data/v2/catalog/records.py b/policyengine_api/data/v2/catalog/records.py index 6a0adf232..555311294 100644 --- a/policyengine_api/data/v2/catalog/records.py +++ b/policyengine_api/data/v2/catalog/records.py @@ -13,6 +13,8 @@ @dataclass(frozen=True, slots=True) class ModelRecord: + """One country's tax-benefit model as named in the catalog.""" + id: UUID country_id: str name: str @@ -21,6 +23,12 @@ class ModelRecord: @dataclass(frozen=True, slots=True) class ModelVersionRecord: + """The model version a catalog was extracted from. + + Also carries the country's current-law policy id and the years its metadata + covers. + """ + id: UUID model_id: UUID version: str @@ -31,6 +39,8 @@ class ModelVersionRecord: @dataclass(frozen=True, slots=True) class VariableRecord: + """One model variable, with its entity, type and the variables it sums.""" + id: UUID model_version_id: UUID name: str @@ -46,6 +56,8 @@ class VariableRecord: @dataclass(frozen=True, slots=True) class ParameterNodeRecord: + """A non-leaf node of the parameter tree; it groups parameters and has no values.""" + id: UUID model_version_id: UUID name: str @@ -55,6 +67,8 @@ class ParameterNodeRecord: @dataclass(frozen=True, slots=True) class ParameterValueRecord: + """A parameter's value from ``start_date`` until ``end_date``, or indefinitely.""" + id: UUID parameter_id: UUID value_json: Any @@ -64,6 +78,8 @@ class ParameterValueRecord: @dataclass(frozen=True, slots=True) class ParameterRecord: + """A leaf parameter with its unit and every value it has held over time.""" + id: UUID model_version_id: UUID name: str @@ -76,6 +92,8 @@ class ParameterRecord: @dataclass(frozen=True, slots=True) class DatasetRecord: + """A logical input dataset, referenced by name; ``storage_path`` is always None.""" + id: UUID model_version_id: UUID name: str @@ -87,6 +105,12 @@ class DatasetRecord: @dataclass(frozen=True, slots=True) class RegionRecord: + """A region that can be simulated, with the dataset it uses by default. + + Regions that need a subset of that dataset (``requires_filter``) also name + the variable, value and strategy that select their households. + """ + id: UUID model_version_id: UUID default_dataset_id: UUID @@ -104,6 +128,11 @@ class RegionRecord: @dataclass(frozen=True, slots=True) class FallbackSummary: + """How many US sub-national regions of one type fell back to the default dataset. + + These regions had no dataset path of their own. + """ + region_type: str count: int @@ -117,7 +146,6 @@ def iter_batches( batch_size: int, ) -> Iterator[tuple[RecordT, ...]]: """Yield bounded immutable slices without constructing a flattened copy.""" - if batch_size < 1: raise ValueError("batch_size must be positive") for start in range(0, len(records), batch_size): @@ -126,6 +154,8 @@ def iter_batches( @dataclass(frozen=True, slots=True) class CountryCatalog: + """Every record the catalog publishes for one country's model version.""" + country_id: str model: ModelRecord model_version: ModelVersionRecord @@ -142,7 +172,6 @@ def parameter_value_batches( batch_size: int, ) -> Iterator[tuple[ParameterValueRecord, ...]]: """Yield canonical parameter values in bounded batches.""" - if batch_size < 1: raise ValueError("batch_size must be positive") batch: list[ParameterValueRecord] = [] @@ -157,7 +186,6 @@ def parameter_value_batches( def entity_counts(self) -> dict[str, int]: """Return non-secret record counts for validation and deployment evidence.""" - return { "models": 1, "model_versions": 1, @@ -174,13 +202,14 @@ def entity_counts(self) -> dict[str, int]: @dataclass(frozen=True, slots=True) class NormalizedCatalog: + """A complete extracted catalog for one PolicyEngine.py version, all countries.""" + policyengine_version: str dependency_versions: tuple[tuple[str, str], ...] countries: tuple[CountryCatalog, ...] def country(self, country_id: str) -> CountryCatalog: """Return one supported country catalog.""" - for catalog in self.countries: if catalog.country_id == country_id: return catalog @@ -188,7 +217,6 @@ def country(self, country_id: str) -> CountryCatalog: def entity_counts(self) -> dict[str, int]: """Return aggregate counts without catalog content.""" - counts: dict[str, int] = {} for country in self.countries: for name, count in country.entity_counts().items(): @@ -198,5 +226,4 @@ def entity_counts(self) -> dict[str, int]: def record_content(record: object) -> tuple[Any, ...]: """Return a deterministic field-ordered representation for comparisons.""" - return tuple(getattr(record, item.name) for item in fields(record)) diff --git a/policyengine_api/data/v2/database.py b/policyengine_api/data/v2/database.py index b742177bd..0f93f9b16 100644 --- a/policyengine_api/data/v2/database.py +++ b/policyengine_api/data/v2/database.py @@ -33,7 +33,6 @@ def _url_fingerprint(settings: V2DatabaseSettings) -> tuple[object, ...]: """Build a private comparison value that is never rendered or logged.""" - url = settings.connection.url return ( url.drivername, @@ -50,7 +49,6 @@ def _url_fingerprint(settings: V2DatabaseSettings) -> tuple[object, ...]: def build_v2_engine(settings: V2DatabaseSettings) -> Engine: """Construct an unconnected SQLModel engine for an explicit v2 target.""" - return create_engine( settings.connection.url, pool_pre_ping=True, @@ -67,13 +65,11 @@ def build_v2_engine(settings: V2DatabaseSettings) -> Engine: def build_v2_session_factory(engine: Engine) -> sessionmaker[Session]: """Build a factory whose ordinary sessions are SQLModel Sessions.""" - return sessionmaker(bind=engine, class_=Session, expire_on_commit=False) def _discard_inherited_state(current_pid: int) -> None: """Drop pool state inherited through a process fork without touching parent FDs.""" - global _engine, _engine_pid, _engine_url_fingerprint, _session_factory if _engine is not None: @@ -86,7 +82,6 @@ def _discard_inherited_state(current_pid: int) -> None: def get_v2_engine(settings: V2DatabaseSettings | None = None) -> Engine: """Return the one lazy v2 runtime engine owned by this process.""" - global _engine, _engine_pid, _engine_url_fingerprint current_pid = os.getpid() @@ -114,7 +109,6 @@ def get_v2_session_factory( settings: V2DatabaseSettings | None = None, ) -> sessionmaker[Session]: """Return the process-owned SQLModel Session factory without opening a session.""" - global _session_factory engine = get_v2_engine(settings) @@ -126,7 +120,6 @@ def get_v2_session_factory( def close_v2_database() -> None: """Dispose this process's v2 pool and forget its lazy configuration.""" - global _engine, _engine_pid, _engine_url_fingerprint, _session_factory with _state_lock: diff --git a/policyengine_api/data/v2/household_migration_qualification.py b/policyengine_api/data/v2/household_migration_qualification.py index 11ca1616e..0297a7d85 100644 --- a/policyengine_api/data/v2/household_migration_qualification.py +++ b/policyengine_api/data/v2/household_migration_qualification.py @@ -54,9 +54,11 @@ class HouseholdDataCounts: @property def total(self) -> int: + """The number of affected rows; anything but zero stops the migration.""" return sum(self.as_dict().values()) def as_dict(self) -> dict[str, int]: + """Return the counts keyed by the names used in evidence and errors.""" return { "households": self.households, "user_household_associations": self.user_household_associations, @@ -76,6 +78,10 @@ class HouseholdMigrationQualification: required: bool = True def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object the qualification command prints. + + ``qualification`` is "not-required" when the revision was already applied. + """ return { "outcome": "ok", "qualification": "performed" if self.required else "not-required", @@ -89,6 +95,7 @@ class RetainedHouseholdDataError(RuntimeError): """Raised when a target contains household data requiring preservation.""" def __init__(self, counts: HouseholdDataCounts) -> None: + """Keep the counts and say in the message how to proceed safely.""" self.counts = counts rendered_counts = ", ".join( f"{name}={count}" for name, count in counts.as_dict().items() @@ -142,14 +149,12 @@ def count(statement: object, table_name: str) -> int: def require_no_retained_household_data(counts: HouseholdDataCounts) -> None: """Stop the migration when any affected row requires a decision.""" - if counts.total: raise RetainedHouseholdDataError(counts) def build_qualification_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated connection pool for one qualification attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -165,7 +170,6 @@ def _read_and_require_counts(connection: Connection) -> HouseholdDataCounts: def household_migration_is_pending(connection: Connection) -> bool: """Return whether the Stage 11 household revision is in the upgrade path.""" - current_heads = MigrationContext.configure(connection).get_current_heads() revisions = ScriptDirectory.from_config( Config(V2_ALEMBIC_CONFIG) @@ -183,7 +187,6 @@ def qualify_household_migration_target( ), ) -> HouseholdMigrationQualification: """Qualify the configured Supabase target without changing it.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -212,7 +215,6 @@ def qualify_household_migration_if_pending( pending_checker: Callable[[Connection], bool] = household_migration_is_pending, ) -> HouseholdMigrationQualification: """Qualify affected rows only while the Stage 11 revision is pending.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -257,7 +259,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run qualification and return a shell-compatible status.""" - try: evidence = qualify_household_migration_if_pending() except Exception as error: # noqa: BLE001 - command must emit safe evidence diff --git a/policyengine_api/data/v2/metadata_validation.py b/policyengine_api/data/v2/metadata_validation.py index a37e92aaf..786908eb4 100644 --- a/policyengine_api/data/v2/metadata_validation.py +++ b/policyengine_api/data/v2/metadata_validation.py @@ -41,7 +41,6 @@ class V2MetadataValidationError(RuntimeError): def validate_v2_metadata_table_names(table_names: Iterable[str]) -> None: """Reject known v1-only and explicitly prohibited v2 table names.""" - actual = frozenset(table_names) prohibited = sorted(actual & PROHIBITED_V2_TABLES) v1_only = sorted(actual & V1_ONLY_TABLES) diff --git a/policyengine_api/data/v2/migration_target.py b/policyengine_api/data/v2/migration_target.py index ac78e6a43..a8373e4b9 100644 --- a/policyengine_api/data/v2/migration_target.py +++ b/policyengine_api/data/v2/migration_target.py @@ -36,6 +36,11 @@ class V2MigrationTargetError(V2ConfigurationError): @dataclass(frozen=True) class ConfiguredSupabaseTarget: + """The persistent Supabase database v2 Alembic may migrate, and its required role. + + The freshness audit fields gate the first migration of an unstamped database. + """ + environment: str project_ref: str database_name: str @@ -46,6 +51,11 @@ class ConfiguredSupabaseTarget: @dataclass(frozen=True) class V2AlembicSettings: + """The URL v2 Alembic connects to, and whether it is the disposable test database. + + ``target`` is None only in disposable test mode. + """ + url: URL disposable_test: bool target: ConfiguredSupabaseTarget | None @@ -88,7 +98,6 @@ def load_v2_alembic_settings( environ: Mapping[str, str] | None = None, ) -> V2AlembicSettings: """Select either the configured persistent target or isolated test Postgres.""" - values = os.environ if environ is None else environ raw_url = _required(values, V2_MIGRATION_DATABASE_URL) url = _parse_url(raw_url) @@ -134,7 +143,6 @@ def qualify_v2_connection( settings: V2AlembicSettings, ) -> None: """Verify live target identity and first-use state before Alembic runs.""" - if connection.dialect.name != "postgresql": raise V2MigrationTargetError("v2 Alembic requires a Postgres connection") if settings.disposable_test: @@ -172,7 +180,6 @@ def qualify_v2_connection( def validate_v2_head_schema(connection: Connection, metadata: MetaData) -> None: """Require live application tables to match authoritative ORM metadata.""" - public_tables = set(inspect(connection).get_table_names(schema="public")) application_tables = public_tables - {"alembic_version"} metadata_tables = {table.name for table in metadata.tables.values()} diff --git a/policyengine_api/data/v2/models/associations.py b/policyengine_api/data/v2/models/associations.py index 4b24bb4b0..2c5cd21b1 100644 --- a/policyengine_api/data/v2/models/associations.py +++ b/policyengine_api/data/v2/models/associations.py @@ -39,6 +39,11 @@ def _jsonb_type() -> sa.types.TypeEngine[Any]: class UserHouseholdAssociation(TimestampedModel, table=True): + """A user's saved, named link to a v2 or v1 household in the same country. + + Deleted with the user; a household that users have saved cannot be deleted. + """ + __tablename__ = "user_household_associations" __table_args__ = ( sa.UniqueConstraint( @@ -95,6 +100,12 @@ class UserHouseholdAssociation(TimestampedModel, table=True): class UserPolicy(TimestampedModel, table=True): + """A user's saved, named link to a v2 or v1 policy in the same country. + + Deleted with the user, together with any v1 mapping it was mirrored from; a + policy that users have saved cannot be deleted. + """ + __tablename__ = "user_policies" __table_args__ = ( sa.UniqueConstraint( diff --git a/policyengine_api/data/v2/models/base.py b/policyengine_api/data/v2/models/base.py index 23c3809a9..fae3a25e2 100644 --- a/policyengine_api/data/v2/models/base.py +++ b/policyengine_api/data/v2/models/base.py @@ -45,7 +45,6 @@ def utc_now() -> datetime: """Return an aware UTC timestamp for application-created rows.""" - return datetime.now(UTC) @@ -78,7 +77,6 @@ class TimestampedModel(IdentifiedModel): def enum_type(enum_class: type, name: str) -> sa.Enum: """Build a stable lowercase-value Postgres enum for a string Enum.""" - # Native named enums and value serialization are SQLAlchemy features that # SQLModel intentionally exposes through Field(sa_type=...). return sa.Enum( diff --git a/policyengine_api/data/v2/models/comparison_runs.py b/policyengine_api/data/v2/models/comparison_runs.py index da24b6886..e45e992af 100644 --- a/policyengine_api/data/v2/models/comparison_runs.py +++ b/policyengine_api/data/v2/models/comparison_runs.py @@ -11,6 +11,11 @@ class Stage12RunStatus(StrEnum): + """Lifecycle state of a Stage 12 run and of each of its simulations. + + Failed and skipped rows always carry an error code. + """ + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -20,6 +25,11 @@ class Stage12RunStatus(StrEnum): class Stage12AggregationStatus(StrEnum): + """Progress of combining a Stage 12 run's simulations into its aggregate. + + A run can only succeed once its aggregation has succeeded. + """ + NOT_STARTED = "not_started" RUNNING = "running" SUCCEEDED = "succeeded" @@ -27,6 +37,11 @@ class Stage12AggregationStatus(StrEnum): class Stage12ResultComparisonStatus(StrEnum): + """Outcome of comparing a Stage 12 aggregate exactly with the production result. + + Direct submissions have no production result, so they stay ``not_requested``. + """ + NOT_REQUESTED = "not_requested" PENDING = "pending" RUNNING = "running" @@ -36,6 +51,8 @@ class Stage12ResultComparisonStatus(StrEnum): class Stage12SimulationRole(StrEnum): + """Role of a Stage 12 simulation in its run: baseline, reform or standalone.""" + BASELINE = "baseline" REFORM = "reform" STANDALONE = "standalone" diff --git a/policyengine_api/data/v2/models/households.py b/policyengine_api/data/v2/models/households.py index a6831d659..78e76b2ab 100644 --- a/policyengine_api/data/v2/models/households.py +++ b/policyengine_api/data/v2/models/households.py @@ -25,6 +25,8 @@ class HouseholdJobStatus(StrEnum): + """Lifecycle state of a queued household calculation job.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -32,6 +34,13 @@ class HouseholdJobStatus(StrEnum): class Household(TimestampedModel, table=True): + """Immutable household input, stored once per canonicalization version and hash. + + ``content_hash`` is the SHA-256 of the canonical document, so saving the same + household again finds the existing row. Households that simulations, reports + or saved links use cannot be deleted. + """ + __tablename__ = "households" __table_args__ = ( sa.UniqueConstraint( @@ -88,6 +97,11 @@ class Household(TimestampedModel, table=True): class HouseholdJob(IdentifiedModel, table=True): + """A queued household calculation request, with its status and result or error. + + Deleting the job's policy or dynamic keeps the job but clears the reference. + """ + __tablename__ = "household_jobs" __table_args__ = ( sa.Index("ix_household_jobs_status_created_at", "status", "created_at"), diff --git a/policyengine_api/data/v2/models/metadata.py b/policyengine_api/data/v2/models/metadata.py index 1b7b9edb7..8d91864df 100644 --- a/policyengine_api/data/v2/models/metadata.py +++ b/policyengine_api/data/v2/models/metadata.py @@ -21,6 +21,11 @@ class RegionType(StrEnum): + """Kind of geography a region covers. + + ``national`` is the whole model country; ``country`` is a UK nation. + """ + NATIONAL = "national" COUNTRY = "country" STATE = "state" @@ -32,6 +37,11 @@ class RegionType(StrEnum): class TaxBenefitModel(TimestampedModel, table=True): + """A country's tax-benefit model, unique by name. + + Its versions own the metadata catalog and are deleted with it. + """ + __tablename__ = "tax_benefit_models" __table_args__ = (sa.UniqueConstraint("name", name="uq_tax_benefit_models_name"),) @@ -50,6 +60,13 @@ class TaxBenefitModel(TimestampedModel, table=True): class TaxBenefitModelVersion(IdentifiedModel, table=True): + """One published catalog version of a model, unique per model and version. + + Owns its variables, parameters and parameter nodes, which are deleted with it, + and records the country's current-law policy id and the years its metadata + covers. + """ + __tablename__ = "tax_benefit_model_versions" __table_args__ = ( sa.UniqueConstraint( @@ -91,6 +108,13 @@ class TaxBenefitModelVersion(IdentifiedModel, table=True): class Region(TimestampedModel, table=True): + """A region in one model version, unique by code within that version. + + Its default dataset must belong to the same version. Filtered regions + (``requires_filter``) also name the variable and value that select their + households. + """ + __tablename__ = "regions" __table_args__ = ( sa.UniqueConstraint( @@ -142,6 +166,11 @@ class Region(TimestampedModel, table=True): class Dataset(TimestampedModel, table=True): + """A dataset of one model version, unique by name within that version. + + Output datasets (``is_output_dataset``) must record a ``storage_path``. + """ + __tablename__ = "datasets" __table_args__ = ( sa.UniqueConstraint( @@ -198,6 +227,8 @@ class Dataset(TimestampedModel, table=True): class DatasetVersion(IdentifiedModel, table=True): + """A named version of a dataset, unique per dataset and deleted with it.""" + __tablename__ = "dataset_versions" __table_args__ = ( sa.UniqueConstraint( @@ -225,6 +256,8 @@ class DatasetVersion(IdentifiedModel, table=True): class Variable(IdentifiedModel, table=True): + """A model variable in one model version, unique by name and deleted with it.""" + __tablename__ = "variables" __table_args__ = ( sa.UniqueConstraint( @@ -255,6 +288,11 @@ class Variable(IdentifiedModel, table=True): class ParameterNode(IdentifiedModel, table=True): + """A non-leaf node of a version's parameter tree, which groups parameters. + + Unique by name within the version, and has no values of its own. + """ + __tablename__ = "parameter_nodes" __table_args__ = ( sa.UniqueConstraint( @@ -279,6 +317,11 @@ class ParameterNode(IdentifiedModel, table=True): class Parameter(IdentifiedModel, table=True): + """A leaf parameter of a model version, unique by name within it. + + Its values include current law and the values each policy or dynamic sets. + """ + __tablename__ = "parameters" __table_args__ = ( sa.UniqueConstraint( @@ -309,6 +352,12 @@ class Parameter(IdentifiedModel, table=True): class ParameterValue(IdentifiedModel, table=True): + """A parameter's value from ``start_date`` until ``end_date``, or indefinitely. + + A value owned by a policy or dynamic is deleted with its owner. A value with + neither is current law, and a parameter has at most one per start date. + """ + __tablename__ = "parameter_values" __table_args__ = ( sa.CheckConstraint( diff --git a/policyengine_api/data/v2/models/policies.py b/policyengine_api/data/v2/models/policies.py index 9b9957b30..97e4a619e 100644 --- a/policyengine_api/data/v2/models/policies.py +++ b/policyengine_api/data/v2/models/policies.py @@ -22,6 +22,12 @@ class Policy(TimestampedModel, table=True): + """Immutable policy content, stored once per canonicalization version and hash. + + Its parameter values are deleted with it. Policies that saved links use + cannot be deleted. + """ + __tablename__ = "policies" __table_args__ = ( sa.UniqueConstraint( @@ -90,6 +96,11 @@ class Policy(TimestampedModel, table=True): class Dynamic(TimestampedModel, table=True): + """Named parameter values, separate from any policy, that a simulation can apply. + + Its parameter values are deleted with it. + """ + __tablename__ = "dynamics" name: str = Field(max_length=255) diff --git a/policyengine_api/data/v2/models/reports.py b/policyengine_api/data/v2/models/reports.py index 170248a82..d10c36947 100644 --- a/policyengine_api/data/v2/models/reports.py +++ b/policyengine_api/data/v2/models/reports.py @@ -25,6 +25,8 @@ class ReportRunStatus(StrEnum): + """Lifecycle state of a report run; finished runs record a completion time.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -32,12 +34,16 @@ class ReportRunStatus(StrEnum): class ReportRunTrigger(StrEnum): + """Why a report run started: the report's first run, a manual or a system rerun.""" + INITIAL = "initial" MANUAL = "manual" SYSTEM = "system" class OutputStatus(StrEnum): + """Lifecycle state of one computed output row of a report run.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -45,12 +51,16 @@ class OutputStatus(StrEnum): class AggregateType(StrEnum): + """How an aggregate combines a variable's values.""" + SUM = "sum" MEAN = "mean" COUNT = "count" class DecileType(StrEnum): + """Whether deciles rank by income or by wealth.""" + INCOME = "income" WEALTH = "wealth" @@ -232,6 +242,11 @@ class ReportRun(TimestampedModel, table=True): class AggregateOutput(IdentifiedModel, table=True): + """A sum, mean or count of one variable in one simulation of a report run. + + ``filter_config`` optionally restricts which entities are included. + """ + __tablename__ = "aggregates" __table_args__ = ( sa.Index("ix_aggregates_report_run_status", "report_run_id", "status"), @@ -262,6 +277,12 @@ class AggregateOutput(IdentifiedModel, table=True): class ChangeAggregate(IdentifiedModel, table=True): + """An aggregate of a variable's change from baseline to reform in a report run. + + ``change_geq`` and ``change_leq`` optionally restrict it to entities whose + change falls within those bounds. + """ + __tablename__ = "change_aggregates" __table_args__ = ( sa.Index( @@ -302,6 +323,11 @@ class ChangeAggregate(IdentifiedModel, table=True): class BudgetSummary(IdentifiedModel, table=True): + """Baseline and reform totals of one budget variable for a report run. + + One row per run, variable and entity. + """ + __tablename__ = "budget_summary" __table_args__ = ( sa.UniqueConstraint( @@ -334,6 +360,12 @@ class BudgetSummary(IdentifiedModel, table=True): class DecileImpact(IdentifiedModel, table=True): + """Mean income change and better/worse-off counts in one income quantile. + + ``decile`` runs from 1 to ``quantiles``; there is one row per run, income + variable, entity, quantile count and decile. + """ + __tablename__ = "decile_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -379,6 +411,11 @@ class DecileImpact(IdentifiedModel, table=True): class IntraDecileImpact(IdentifiedModel, table=True): + """Shares of one decile whose income rises or falls by more or less than 5%. + + One row per run, decile type and decile, where ``decile`` is 0 to 10. + """ + __tablename__ = "intra_decile_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -420,6 +457,8 @@ class IntraDecileImpact(IdentifiedModel, table=True): class Inequality(IdentifiedModel, table=True): + """Gini and top/bottom income shares for one simulation of a report run.""" + __tablename__ = "inequality" __table_args__ = ( sa.UniqueConstraint( @@ -450,6 +489,8 @@ class Inequality(IdentifiedModel, table=True): class Poverty(IdentifiedModel, table=True): + """Poverty headcount and rate of one type for one simulation of a report run.""" + __tablename__ = "poverty" __table_args__ = ( sa.Index( @@ -479,6 +520,11 @@ class Poverty(IdentifiedModel, table=True): class ProgramStatistics(IdentifiedModel, table=True): + """Baseline and reform totals, counts, winners and losers for one program. + + One row per run, program and entity; ``is_tax`` marks taxes. + """ + __tablename__ = "program_statistics" __table_args__ = ( sa.UniqueConstraint( @@ -536,6 +582,11 @@ class GeographicImpactBase(IdentifiedModel): class CongressionalDistrictImpact(GeographicImpactBase, table=True): + """Household income change and population in one US congressional district. + + One row per run and district GEOID. + """ + __tablename__ = "congressional_district_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -555,6 +606,11 @@ class CongressionalDistrictImpact(GeographicImpactBase, table=True): class ConstituencyImpact(GeographicImpactBase, table=True): + """Household income change in one UK constituency, with its map coordinates. + + One row per run and constituency code. + """ + __tablename__ = "constituency_impacts" __table_args__ = ( sa.UniqueConstraint( @@ -573,6 +629,11 @@ class ConstituencyImpact(GeographicImpactBase, table=True): class LocalAuthorityImpact(GeographicImpactBase, table=True): + """Household income change in one UK local authority, with its map coordinates. + + One row per run and local authority code. + """ + __tablename__ = "local_authority_impacts" __table_args__ = ( sa.UniqueConstraint( diff --git a/policyengine_api/data/v2/models/simulations.py b/policyengine_api/data/v2/models/simulations.py index 613c4e1c1..a6821eadc 100644 --- a/policyengine_api/data/v2/models/simulations.py +++ b/policyengine_api/data/v2/models/simulations.py @@ -22,6 +22,8 @@ class SimulationStatus(StrEnum): + """Lifecycle state of a v2 simulation.""" + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -29,11 +31,19 @@ class SimulationStatus(StrEnum): class SimulationType(StrEnum): + """Whether a simulation runs on one household or on a population dataset.""" + HOUSEHOLD = "household" ECONOMY = "economy" class Simulation(TimestampedModel, table=True): + """A simulation of exactly one household or one dataset, matching its type. + + Deleting its policy, dynamic, region or output dataset clears that reference; + its household and input dataset cannot be deleted while it exists. + """ + __tablename__ = "simulations" __table_args__ = ( sa.CheckConstraint( diff --git a/policyengine_api/data/v2/models/users.py b/policyengine_api/data/v2/models/users.py index 14a294f30..85c0ea19a 100644 --- a/policyengine_api/data/v2/models/users.py +++ b/policyengine_api/data/v2/models/users.py @@ -27,11 +27,20 @@ class UserType(StrEnum): + """Whether a user is backed by sign-in identities or by a signed-out browser.""" + UNAUTHENTICATED = "unauthenticated" STANDARD = "standard" class User(IdentifiedModel, table=True): + """A v2 user account, unique by email when one is given. + + A ``standard`` user always has a primary country. Deleting a user deletes + their sign-in identities and saved associations; reports they own are kept + without an owner. + """ + __tablename__ = "users" __table_args__ = ( sa.UniqueConstraint("email", name="uq_users_email"), diff --git a/policyengine_api/data/v2/policy_migration_qualification.py b/policyengine_api/data/v2/policy_migration_qualification.py index b88c237e7..597539aff 100644 --- a/policyengine_api/data/v2/policy_migration_qualification.py +++ b/policyengine_api/data/v2/policy_migration_qualification.py @@ -46,9 +46,11 @@ class PolicyDataCounts: @property def total(self) -> int: + """The number of affected rows; anything but zero stops the migration.""" return self.policies + self.policy_parameter_values + self.user_policies def as_dict(self) -> dict[str, int]: + """Return the counts keyed by the names used in evidence and errors.""" return { "policies": self.policies, "policy_parameter_values": self.policy_parameter_values, @@ -66,6 +68,10 @@ class PolicyMigrationQualification: required: bool = True def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object the qualification command prints. + + ``qualification`` is "not-required" when the revision was already applied. + """ return { "outcome": "ok", "qualification": "performed" if self.required else "not-required", @@ -79,6 +85,7 @@ class RetainedPolicyDataError(RuntimeError): """Raised when a target contains policy data requiring preservation.""" def __init__(self, counts: PolicyDataCounts) -> None: + """Keep the counts and say in the message how to proceed safely.""" self.counts = counts super().__init__( "the configured Supabase target contains retained v2 policy data " @@ -125,14 +132,12 @@ def table_exists(name: str) -> bool: def require_no_retained_policy_data(counts: PolicyDataCounts) -> None: """Stop the migration when any predecessor policy row requires a decision.""" - if counts.total: raise RetainedPolicyDataError(counts) def build_qualification_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated connection pool for one qualification attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -157,7 +162,6 @@ def _qualify_connection(connection: Connection) -> PolicyDataCounts: def policy_migration_is_pending(connection: Connection) -> bool: """Return whether the Phase 10 policy revision is in the upgrade path.""" - current_heads = MigrationContext.configure(connection).get_current_heads() revisions = ScriptDirectory.from_config( Config(V2_ALEMBIC_CONFIG) @@ -173,7 +177,6 @@ def qualify_policy_migration_target( ), ) -> PolicyMigrationQualification: """Qualify the configured Supabase target without changing it.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -197,7 +200,6 @@ def qualify_policy_migration_if_pending( pending_checker: Callable[[Connection], bool] = policy_migration_is_pending, ) -> PolicyMigrationQualification: """Qualify predecessor rows only while the Phase 10 revision is pending.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -242,7 +244,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run qualification and return a shell-compatible status.""" - try: evidence = qualify_policy_migration_if_pending() except Exception as error: # noqa: BLE001 - command must emit safe evidence diff --git a/policyengine_api/data/v2/report_runs.py b/policyengine_api/data/v2/report_runs.py index e4da22634..6aa32f955 100644 --- a/policyengine_api/data/v2/report_runs.py +++ b/policyengine_api/data/v2/report_runs.py @@ -55,7 +55,6 @@ def set_report_type( report_type: str | None, ) -> Report: """Set a report's optional type only while it has no execution history.""" - report = _locked_report(session, report_id) if report.type == report_type: return report @@ -84,7 +83,6 @@ def create_report_run( idempotency_key: UUID | None = None, ) -> ReportRun: """Create one run, or return the run for a retried idempotent request.""" - _locked_report(session, report_id) if trigger is ReportRunTrigger.MANUAL and idempotency_key is None: raise ValueError("manual report reruns require an idempotency key") @@ -130,7 +128,6 @@ def begin_report_run( started_at: datetime | None = None, ) -> ReportRun: """Start or resume the same pending/running durable run after worker retry.""" - run = session.exec( select(ReportRun).where(ReportRun.id == report_run_id).with_for_update() ).one_or_none() @@ -154,7 +151,6 @@ def complete_report_run( markdown: str | None = None, ) -> ReportRun: """Mark the selected durable run successful without replacing history.""" - run = begin_report_run(session, report_run_id=report_run_id) run.status = ReportRunStatus.SUCCEEDED run.completed_at = completed_at or utc_now() @@ -173,7 +169,6 @@ def fail_report_run( completed_at: datetime | None = None, ) -> ReportRun: """Mark the selected durable run failed without deleting older success.""" - run = begin_report_run(session, report_run_id=report_run_id) run.status = ReportRunStatus.FAILED run.completed_at = completed_at or utc_now() @@ -191,7 +186,6 @@ def select_current_report_run( policyengine_version: str = POLICYENGINE_VERSION, ) -> ReportRun | None: """Select the deterministic current successful run for deployed versions.""" - report = session.get(Report, report_id) if report is None: raise ReportNotFoundError(f"report {report_id} does not exist") diff --git a/policyengine_api/data/v2/settings.py b/policyengine_api/data/v2/settings.py index 4ef27980f..5aa3833db 100644 --- a/policyengine_api/data/v2/settings.py +++ b/policyengine_api/data/v2/settings.py @@ -49,16 +49,15 @@ class PostgresConnectionSettings: @property def url(self) -> URL: """Return SQLAlchemy's immutable URL for engine construction.""" - return self._url @property def redacted_url(self) -> str: """Return a log-safe URL with the password hidden.""" - return self._url.render_as_string(hide_password=True) def __str__(self) -> str: + """Return the URL with its password hidden, so printing never leaks it.""" return self.redacted_url @@ -92,7 +91,6 @@ def _required(environ: Mapping[str, str], name: str) -> str: @cache def _load_secret_from_secret_manager(resource_name: str) -> str: """Resolve a v2 runtime URL only when a preview request selects it.""" - from google.cloud import secretmanager client = secretmanager.SecretManagerServiceClient() @@ -140,7 +138,6 @@ def load_supabase_target_settings( environ: Mapping[str, str] | None = None, ) -> SupabaseTargetSettings: """Load the externally configured non-secret Supabase target identity.""" - values = _environment(environ) project_ref = _required(values, V2_SUPABASE_PROJECT_REF) environment = _required(values, V2_SUPABASE_ENVIRONMENT) @@ -165,7 +162,6 @@ def parse_persistent_postgres_url( setting_name: str, ) -> PostgresConnectionSettings: """Validate an explicit persistent Postgres URL without echoing it.""" - try: url = make_url(raw_url) except ArgumentError as error: @@ -200,7 +196,6 @@ def validate_supabase_database_identity( setting_name: str, ) -> None: """Require a persistent URL to identify the configured Supabase project.""" - url = connection.url direct_host = f"db.{target.project_ref}.supabase.co" is_direct = url.host == direct_host @@ -245,7 +240,6 @@ def load_v2_runtime_database_settings( secret_loader: Callable[[str], str] | None = None, ) -> V2DatabaseSettings: """Load the future ordinary-runtime Postgres identity explicitly.""" - values = _environment(environ) raw_url = _resolve_runtime_database_url( values, @@ -269,7 +263,6 @@ def load_v2_migration_database_settings( environ: Mapping[str, str] | None = None, ) -> V2DatabaseSettings: """Load the schema-migration Postgres identity explicitly.""" - values = _environment(environ) return _database_settings( _required(values, V2_MIGRATION_DATABASE_URL), @@ -282,7 +275,6 @@ def load_v2_data_write_database_settings( environ: Mapping[str, str] | None = None, ) -> V2DatabaseSettings: """Load the one-time catalog row-write Postgres identity explicitly.""" - values = _environment(environ) return _database_settings( _required(values, V2_DATA_WRITE_DATABASE_URL), diff --git a/policyengine_api/data/v2/user_association_migration_qualification.py b/policyengine_api/data/v2/user_association_migration_qualification.py index 9c74f44d0..3629782d7 100644 --- a/policyengine_api/data/v2/user_association_migration_qualification.py +++ b/policyengine_api/data/v2/user_association_migration_qualification.py @@ -66,9 +66,11 @@ class UserAssociationCounts: @property def dormant_total(self) -> int: + """Rows in the two reshaped tables, which must be empty before the reshape.""" return self.user_simulation_associations + self.user_report_associations def as_dict(self) -> dict[str, int]: + """Return the counts keyed by table name, for the evidence and errors.""" return { "user_policies": self.user_policies, "user_household_associations": self.user_household_associations, @@ -87,6 +89,7 @@ class UserAssociationQualification: migration_pending: bool def as_dict(self) -> dict[str, object]: + """Return the evidence as the JSON object the qualification command prints.""" return { "outcome": "ok", "revision": USER_ASSOCIATION_MIGRATION_REVISION, @@ -102,6 +105,7 @@ class RetainedDormantAssociationError(RuntimeError): """Raised when a dormant association table holds rows before the reshape.""" def __init__(self, counts: UserAssociationCounts) -> None: + """Keep the counts and say in the message how to proceed safely.""" self.counts = counts super().__init__( "the configured Supabase target contains user_simulation_associations=" @@ -142,13 +146,13 @@ def require_empty_dormant_associations( *, migration_pending: bool, ) -> None: + """Raise RetainedDormantAssociationError if a pending reshape would affect rows.""" if migration_pending and counts.dormant_total: raise RetainedDormantAssociationError(counts) def user_association_migration_is_pending(connection: Connection) -> bool: """Return whether the Stage 12A revision is in the upgrade path.""" - current_heads = MigrationContext.configure(connection).get_current_heads() revisions = ScriptDirectory.from_config( Config(V2_ALEMBIC_CONFIG) @@ -160,7 +164,6 @@ def user_association_migration_is_pending(connection: Connection) -> bool: def build_qualification_engine(settings: V2DatabaseSettings) -> Engine: """Build an isolated connection pool for one qualification attempt.""" - return create_engine(settings.connection.url, poolclass=NullPool) @@ -182,7 +185,6 @@ def qualify_user_association_migration( ), ) -> UserAssociationQualification: """Report counts and stop when a pending reshape would affect rows.""" - settings = load_v2_migration_database_settings(environ) engine = engine_builder(settings) try: @@ -226,7 +228,6 @@ def _error_payload(error: Exception) -> dict[str, object]: def main() -> int: """Run qualification and return a shell-compatible status.""" - try: evidence = qualify_user_association_migration() except Exception as error: # noqa: BLE001 - command must emit safe evidence diff --git a/policyengine_api/dataset_display.py b/policyengine_api/dataset_display.py index b303ed3bd..d092fa7e8 100644 --- a/policyengine_api/dataset_display.py +++ b/policyengine_api/dataset_display.py @@ -11,7 +11,6 @@ def get_dataset_display_label(dataset_name: object) -> str: """Return the public label for an internal logical dataset name.""" - normalized_name = str(dataset_name or "") for prefix, label in DATASET_FAMILY_DISPLAY_LABELS: if normalized_name.startswith(prefix): diff --git a/policyengine_api/fastapi_routes/dependencies.py b/policyengine_api/fastapi_routes/dependencies.py index 206d1614a..a7d0712a3 100644 --- a/policyengine_api/fastapi_routes/dependencies.py +++ b/policyengine_api/fastapi_routes/dependencies.py @@ -88,13 +88,24 @@ class MetadataReader(Protocol): """Read the already-loaded metadata document for a country.""" - def get_metadata(self, country_id: str) -> JSONObject: ... + def get_metadata(self, country_id: str) -> JSONObject: + """Return the loaded v1 metadata document for a supported country.""" + ... class V2MetadataResourceReader(Protocol): - """Own the database session used by one v2 metadata resource request.""" + """Own the database session used by one v2 metadata resource request. - def close(self) -> None: ... + Every read selects one country catalog: the one published for + ``policyengine_version``, or the running package's catalog when it is None. + Absent resources raise ``MetadataResourceNotFoundError``, and list reads + raise ``InvalidMetadataPageError`` unless ``offset >= 0`` and + ``1 <= limit <= 500``. + """ + + def close(self) -> None: + """Release the request's database session, even after a failed read.""" + ... def list_models( self, @@ -103,20 +114,26 @@ def list_models( *, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataModel]: ... + ) -> MetadataPageResult[MetadataModel]: + """Return a page holding the selected catalog's single tax-benefit model.""" + ... def get_model( self, country_id: str, model_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataModel]: ... + ) -> MetadataDetailResult[MetadataModel]: + """Return the selected catalog's model when ``model_id`` is its UUID.""" + ... def get_model_by_country( self, country_id: str, policyengine_version: str | None = None, - ) -> MetadataModelSelectionResult: ... + ) -> MetadataModelSelectionResult: + """Return the country's model and the selected catalog's model version.""" + ... def list_model_versions( self, @@ -125,14 +142,18 @@ def list_model_versions( *, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataModelVersionDetail]: ... + ) -> MetadataPageResult[MetadataModelVersionDetail]: + """Return a page holding the selected catalog's single model version.""" + ... def get_model_version( self, country_id: str, version_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataModelVersionDetail]: ... + ) -> MetadataDetailResult[MetadataModelVersionDetail]: + """Return the selected model version when ``version_id`` is its UUID.""" + ... def list_variables( self, @@ -142,14 +163,21 @@ def list_variables( offset: int = 0, limit: int = 100, search: str | None = None, - ) -> MetadataPageResult[MetadataVariable]: ... + ) -> MetadataPageResult[MetadataVariable]: + """Return variables in name order, optionally filtered by ``search``. + + ``search`` is a case-insensitive substring of the name, label or description. + """ + ... def get_variable( self, country_id: str, variable_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataVariable]: ... + ) -> MetadataDetailResult[MetadataVariable]: + """Return one variable of the selected catalog by UUID.""" + ... def list_parameters( self, @@ -159,14 +187,21 @@ def list_parameters( offset: int = 0, limit: int = 100, search: str | None = None, - ) -> MetadataPageResult[MetadataParameterSummary]: ... + ) -> MetadataPageResult[MetadataParameterSummary]: + """Return parameters in name order, optionally filtered by ``search``. + + ``search`` is a case-insensitive substring of the name, label or description. + """ + ... def get_parameter( self, country_id: str, parameter_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataParameterSummary]: ... + ) -> MetadataDetailResult[MetadataParameterSummary]: + """Return one parameter of the selected catalog by UUID.""" + ... def list_parameter_children( self, @@ -176,7 +211,12 @@ def list_parameter_children( parent_path: str = "", offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataParameterChild]: ... + ) -> MetadataPageResult[MetadataParameterChild]: + """Return the nodes and parameters one level below ``parent_path``. + + An empty ``parent_path`` lists the top level of the parameter tree. + """ + ... def list_parameter_values( self, @@ -188,14 +228,22 @@ def list_parameter_values( offset: int = 0, limit: int = 100, now: datetime | None = None, - ) -> MetadataPageResult[MetadataCanonicalParameterValue]: ... + ) -> MetadataPageResult[MetadataCanonicalParameterValue]: + """Return canonical values, optionally for one parameter. + + With ``current``, only values in effect on the UTC day of ``now`` (default: + the present) are returned. + """ + ... def get_parameter_value( self, country_id: str, value_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: ... + ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: + """Return one canonical parameter value of the selected catalog by UUID.""" + ... def list_datasets( self, @@ -204,14 +252,18 @@ def list_datasets( *, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataDataset]: ... + ) -> MetadataPageResult[MetadataDataset]: + """Return the logical input datasets of the selected catalog.""" + ... def get_dataset( self, country_id: str, dataset_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataDataset]: ... + ) -> MetadataDetailResult[MetadataDataset]: + """Return one logical input dataset of the selected catalog by UUID.""" + ... def list_regions( self, @@ -221,27 +273,41 @@ def list_regions( region_type: str | None = None, offset: int = 0, limit: int = 100, - ) -> MetadataPageResult[MetadataRegion]: ... + ) -> MetadataPageResult[MetadataRegion]: + """Return the catalog's regions by code, optionally of one ``region_type``.""" + ... def get_region( self, country_id: str, region_id: UUID, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataRegion]: ... + ) -> MetadataDetailResult[MetadataRegion]: + """Return one region of the selected catalog by UUID.""" + ... def get_region_by_code( self, country_id: str, region_code: str, policyengine_version: str | None = None, - ) -> MetadataDetailResult[MetadataRegion]: ... + ) -> MetadataDetailResult[MetadataRegion]: + """Return one region by code; the national region's code is the country id.""" + ... def get_economy_options( self, country_id: str, policyengine_version: str | None = None, - ) -> MetadataEconomyOptionsResult: ... + ) -> MetadataEconomyOptionsResult: + """Return the options a client needs to configure an economy comparison. + + These are the current-law policy id, regions, years and datasets. + + Raise ``MetadataCatalogUnavailableError`` when the national region, its + dataset or the model version's options are missing from the catalog. + """ + ... class V2PolicyResourceService(Protocol): @@ -250,9 +316,20 @@ class V2PolicyResourceService(Protocol): def create_policy( self, command: NativePolicyCreationInput, - ) -> NativePolicyCreation: ... + ) -> NativePolicyCreation: + """Store the policy, or return the identical one already stored. + + Policies are immutable and deduplicated by canonical content hash; + ``created`` says whether this call inserted the policy. + """ + ... + + def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: + """Return one policy of the country with its parameter values. - def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: ... + Raise ``PolicyNotFoundError`` when no such policy exists in that country. + """ + ... def list_policies( self, @@ -261,7 +338,9 @@ def list_policies( tax_benefit_model_id: UUID | None = None, offset: int = 0, limit: int = 100, - ) -> PolicyPage: ... + ) -> PolicyPage: + """Return a page of the country's policies, optionally for one model.""" + ... class V2HouseholdResourceService(Protocol): @@ -270,14 +349,25 @@ class V2HouseholdResourceService(Protocol): def create_household( self, household_input: NativeHouseholdCreationInput, - ) -> NativeHouseholdCreation: ... + ) -> NativeHouseholdCreation: + """Store the household, or return the identical one already stored. + + Households are immutable and deduplicated by canonical content hash; + ``created`` says whether this call inserted the household. + """ + ... def get_household( self, *, country_id: str, household_id: UUID, - ) -> HouseholdRead: ... + ) -> HouseholdRead: + """Return one household of the country. + + Raise ``HouseholdNotFoundError`` when no such household exists there. + """ + ... def list_households( self, @@ -286,7 +376,9 @@ def list_households( default_year: int | None, offset: int, limit: int, - ) -> HouseholdPage: ... + ) -> HouseholdPage: + """Return a page of the country's households, optionally of one default year.""" + ... class V2UserPolicyResourceService(Protocol): @@ -297,7 +389,9 @@ def create_user_policy( *, principal: AccountPrincipal, association_input: UserPolicyCreateFields, - ) -> UserPolicyRead: ... + ) -> UserPolicyRead: + """Save a policy for the caller, provisioning their user if needed.""" + ... def get_user_policy( self, @@ -305,7 +399,9 @@ def get_user_policy( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> UserPolicyRead: ... + ) -> UserPolicyRead: + """Return the caller's association, or raise ``UserPolicyNotFoundError``.""" + ... def list_user_policies( self, @@ -316,7 +412,9 @@ def list_user_policies( legacy_policy_id: int | None = None, offset: int = 0, limit: int = 100, - ) -> UserPolicyPage: ... + ) -> UserPolicyPage: + """Return a page of the caller's associations, optionally for one policy.""" + ... def patch_user_policy( self, @@ -325,7 +423,12 @@ def patch_user_policy( country_id: str, association_id: UUID, association_input: UserPolicyUpdateInput, - ) -> UserPolicyRead: ... + ) -> UserPolicyRead: + """Update the supplied fields of the caller's association and return it. + + A new ``legacy_policy_id`` repoints it to another v1 policy. + """ + ... def delete_user_policy( self, @@ -333,7 +436,9 @@ def delete_user_policy( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> None: ... + ) -> None: + """Delete the caller's association, or raise ``UserPolicyNotFoundError``.""" + ... class V2UserHouseholdResourceService(Protocol): @@ -344,7 +449,9 @@ def create_user_household( *, principal: AccountPrincipal, association_input: UserHouseholdCreateFields, - ) -> UserHouseholdRead: ... + ) -> UserHouseholdRead: + """Save a household for the caller, provisioning their user if needed.""" + ... def get_user_household( self, @@ -352,7 +459,9 @@ def get_user_household( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> UserHouseholdRead: ... + ) -> UserHouseholdRead: + """Return the caller's association, or raise ``UserHouseholdNotFoundError``.""" + ... def list_user_households( self, @@ -363,7 +472,9 @@ def list_user_households( legacy_household_id: int | None = None, offset: int = 0, limit: int = 100, - ) -> UserHouseholdPage: ... + ) -> UserHouseholdPage: + """Return a page of the caller's associations, optionally for one household.""" + ... def patch_user_household( self, @@ -372,7 +483,12 @@ def patch_user_household( country_id: str, association_id: UUID, association_input: UserHouseholdUpdateInput, - ) -> UserHouseholdRead: ... + ) -> UserHouseholdRead: + """Update the supplied fields of the caller's association and return it. + + A new ``legacy_household_id`` repoints it to another v1 household. + """ + ... def delete_user_household( self, @@ -380,7 +496,9 @@ def delete_user_household( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> None: ... + ) -> None: + """Delete the caller's association, or raise ``UserHouseholdNotFoundError``.""" + ... class V2UserSimulationResourceService(Protocol): @@ -391,7 +509,9 @@ def create_user_simulation( *, principal: AccountPrincipal, association_input: UserSimulationCreateFields, - ) -> UserSimulationRead: ... + ) -> UserSimulationRead: + """Save a v1 simulation for the caller, provisioning their user if needed.""" + ... def get_user_simulation( self, @@ -399,7 +519,9 @@ def get_user_simulation( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> UserSimulationRead: ... + ) -> UserSimulationRead: + """Return the caller's association, or raise ``UserSimulationNotFoundError``.""" + ... def list_user_simulations( self, @@ -409,7 +531,9 @@ def list_user_simulations( legacy_simulation_id: int | None = None, offset: int = 0, limit: int = 100, - ) -> UserSimulationPage: ... + ) -> UserSimulationPage: + """Return a page of the caller's associations, optionally for one simulation.""" + ... def patch_user_simulation( self, @@ -418,7 +542,12 @@ def patch_user_simulation( country_id: str, association_id: UUID, association_input: UserSimulationUpdateInput, - ) -> UserSimulationRead: ... + ) -> UserSimulationRead: + """Update the supplied fields of the caller's association and return it. + + A new ``legacy_simulation_id`` repoints it to another v1 simulation. + """ + ... def delete_user_simulation( self, @@ -426,7 +555,9 @@ def delete_user_simulation( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> None: ... + ) -> None: + """Delete the caller's association, or raise ``UserSimulationNotFoundError``.""" + ... class V2UserReportResourceService(Protocol): @@ -437,7 +568,12 @@ def create_user_report( *, principal: AccountPrincipal, association_input: UserReportCreateFields, - ) -> UserReportRead: ... + ) -> UserReportRead: + """Save a v1 report for the caller, provisioning their user if needed. + + Linked saved reforms must be the caller's own. + """ + ... def get_user_report( self, @@ -445,14 +581,21 @@ def get_user_report( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> UserReportRead: ... + ) -> UserReportRead: + """Return the caller's association, or raise ``UserReportNotFoundError``.""" + ... def get_shared_user_report( self, *, country_id: str, association_id: UUID, - ) -> SharedUserReportRead: ... + ) -> SharedUserReportRead: + """Return any user's saved report to a link holder, without its owner. + + Raise ``UserReportNotFoundError`` when the country has no such report. + """ + ... def list_user_reports( self, @@ -462,7 +605,9 @@ def list_user_reports( legacy_report_id: int | None = None, offset: int = 0, limit: int = 100, - ) -> UserReportPage: ... + ) -> UserReportPage: + """Return a page of the caller's associations, optionally for one v1 report.""" + ... def patch_user_report( self, @@ -471,7 +616,9 @@ def patch_user_report( country_id: str, association_id: UUID, association_input: UserReportUpdateInput, - ) -> UserReportRead: ... + ) -> UserReportRead: + """Update the supplied fields of the caller's association and return it.""" + ... def delete_user_report( self, @@ -479,7 +626,9 @@ def delete_user_report( principal: AccountPrincipal, country_id: str, association_id: UUID, - ) -> None: ... + ) -> None: + """Delete the caller's association, or raise ``UserReportNotFoundError``.""" + ... class V2AccountResourceService(Protocol): @@ -487,7 +636,13 @@ class V2AccountResourceService(Protocol): def get_account( self, *, principal: AccountPrincipal, country_id: str - ) -> AccountRead: ... + ) -> AccountRead: + """Return the caller's account, provisioning a signed-in user on first use. + + Raise ``AccountNotFoundError`` for an unauthenticated user that has not + saved anything yet. + """ + ... def update_account( self, @@ -495,11 +650,18 @@ def update_account( principal: AccountPrincipal, country_id: str, account_input: AccountUpdateInput, - ) -> AccountRead: ... + ) -> AccountRead: + """Change the supplied fields of a signed-in user's account and return it. + + Raise ``AccountEditRequiresSignInError`` for an unauthenticated user. + """ + ... def export_account( self, *, principal: AccountPrincipal, country_id: str - ) -> AccountExport: ... + ) -> AccountExport: + """Return the caller's account and all its associations, in every country.""" + ... def upgrade_account( self, @@ -507,15 +669,29 @@ def upgrade_account( principal: UserPrincipal, unauthenticated_user_id: UUID, country_id: str, - ) -> AccountUpgrade: ... + ) -> AccountUpgrade: + """Move a browser's unauthenticated user into the signed-in account. + + The outcome says whether it was converted in place, merged into an + existing account, or absent. + """ + ... + + def delete_account(self, *, principal: AccountPrincipal) -> AccountDeletion: + """Delete the caller's user and associations; an absent user succeeds. - def delete_account(self, *, principal: AccountPrincipal) -> AccountDeletion: ... + A signed-in user's Auth0 user is deleted too when Auth0 Management is + configured. + """ + ... class SimulationGatewayProbe(Protocol): """Minimal simulation-entrypoint health-check interface.""" - def health_check(self) -> bool: ... + def health_check(self) -> bool: + """Return whether the simulation entrypoint's health endpoint returns 200.""" + ... def _default_readiness_probe() -> bool: @@ -607,7 +783,6 @@ def _default_require_user( unauthenticated_user: str | None = Depends(_unauthenticated_user_scheme), ) -> AccountPrincipal: """Authenticate the request's account; fails closed without a credential.""" - from policyengine_api.fastapi_routes.v2.user_auth import ( authenticate_account_request, ) @@ -634,7 +809,6 @@ def _build_v1_legacy_reference_service() -> LegacyReferenceChecker: def _default_legacy_reference_checker() -> LegacyReferenceChecker: """Defer the v1 session factory until a request checks a v1 reference.""" - from policyengine_api.services.v2.legacy_refs.services import ( DeferredLegacyReferenceService, ) diff --git a/policyengine_api/fastapi_routes/health.py b/policyengine_api/fastapi_routes/health.py index b411512a8..6e0da7ec1 100644 --- a/policyengine_api/fastapi_routes/health.py +++ b/policyengine_api/fastapi_routes/health.py @@ -13,10 +13,14 @@ class HealthResponse(BaseModel): + """Health-check result showing that the API is up and serving requests.""" + status: Literal["healthy"] class SimulationGatewayHealthResponse(BaseModel): + """Health-check result showing that the API and its simulation gateway are up.""" + status: Literal["healthy"] simulation_gateway: Literal["healthy"] diff --git a/policyengine_api/fastapi_routes/metadata.py b/policyengine_api/fastapi_routes/metadata.py index 24a214a72..44bbe3f6d 100644 --- a/policyengine_api/fastapi_routes/metadata.py +++ b/policyengine_api/fastapi_routes/metadata.py @@ -26,6 +26,8 @@ class MetadataSuccessPayload(TypedDict): + """The v1 ``{status, message, result}`` envelope around a metadata document.""" + status: Literal["ok"] message: None result: JSONObject diff --git a/policyengine_api/fastapi_routes/responses.py b/policyengine_api/fastapi_routes/responses.py index a58ab27ec..114d51494 100644 --- a/policyengine_api/fastapi_routes/responses.py +++ b/policyengine_api/fastapi_routes/responses.py @@ -13,7 +13,6 @@ def render_legacy_json(content: JSONValue) -> bytes: """Serialize once with the same JSON encoder used by legacy metadata.""" - return json.dumps(content).encode("utf-8") @@ -23,4 +22,5 @@ class LegacyJSONResponse(Response): media_type = "application/json" def render(self, content: JSONValue) -> bytes: + """Encode ``content`` byte-for-byte as the Flask metadata route does.""" return render_legacy_json(content) diff --git a/policyengine_api/fastapi_routes/v2/errors.py b/policyengine_api/fastapi_routes/v2/errors.py index 02ca98597..c5ba09cfd 100644 --- a/policyengine_api/fastapi_routes/v2/errors.py +++ b/policyengine_api/fastapi_routes/v2/errors.py @@ -18,6 +18,11 @@ class V2AuthenticationRequiredError(Exception): """Raised when a user route has no valid bearer access token (401).""" def __init__(self, message: str, *, invalid_token: bool = False) -> None: + """Choose the WWW-Authenticate challenge sent with the 401. + + ``invalid_token`` marks a credential that was sent but rejected, which the + challenge reports as ``error="invalid_token"``. + """ super().__init__(message) self.challenge = ( INVALID_TOKEN_BEARER_CHALLENGE if invalid_token else BEARER_CHALLENGE @@ -48,7 +53,6 @@ def v2_error_response( headers: Mapping[str, str] | None = None, ) -> JSONResponse: """Serialize one API v2 error without resource-path dispatch.""" - error = V2ErrorResponse(message=message) return JSONResponse( status_code=status_code, diff --git a/policyengine_api/fastapi_routes/v2/households/document_models.py b/policyengine_api/fastapi_routes/v2/households/document_models.py index d2b702b52..d64c01e9e 100644 --- a/policyengine_api/fastapi_routes/v2/households/document_models.py +++ b/policyengine_api/fastapi_routes/v2/households/document_models.py @@ -28,16 +28,31 @@ class StrictHouseholdDocumentModel(BaseModel): + """Base for household document records; rejects unknown keys and NaN or infinity.""" + model_config = ConfigDict(extra="forbid", allow_inf_nan=False) class HouseholdEntityRecord(StrictHouseholdDocumentModel): + """One group entity in a household, such as a tax unit or a benefit unit. + + ``id`` is unique within its collection and is what people's ``memberships`` + refer to. ``values`` maps variable names to values, and ``source_name`` + optionally keeps the entity's original name, such as its key in a v1 household. + """ + id: DocumentIdentifier source_name: DocumentIdentifier | None = None values: EntityValues class HouseholdPersonRecord(HouseholdEntityRecord): + """One person in a household. + + ``memberships`` maps each of the country's entity collections to the ``id`` + of the group the person belongs to, and must name every collection. + """ + memberships: dict[str, DocumentIdentifier] @@ -51,6 +66,12 @@ def _omit_spm_default(schema: dict[str, Any]) -> None: class USHouseholdDocument(StrictHouseholdDocumentModel): + """US household content: people and the US group entities they belong to. + + ``spm`` optionally selects an SPM forecast scenario and geography; omit it + rather than sending null. + """ + people: PersonRecords household: EntityRecords family: EntityRecords @@ -64,6 +85,7 @@ class USHouseholdDocument(StrictHouseholdDocumentModel): @field_validator("spm", mode="before") @classmethod def require_spm_object(cls, value: object) -> object: + """Reject an explicit null ``spm``; omit the field to make no selection.""" if value is None: raise ValueError("spm must be an object when supplied") return value @@ -72,6 +94,11 @@ def require_spm_object(cls, value: object) -> object: def serialize_document( self, handler: SerializerFunctionWrapHandler ) -> dict[str, Any]: + """Serialize the document, leaving out ``spm`` when no selection was made. + + Omitting the key keeps such documents identical to those stored before + ``spm`` existed. + """ document: dict[str, Any] = handler(self) if self.spm is None: document.pop("spm", None) @@ -79,6 +106,8 @@ def serialize_document( class UKHouseholdDocument(StrictHouseholdDocumentModel): + """UK household content: people, households and benefit units.""" + people: PersonRecords household: EntityRecords benunit: EntityRecords diff --git a/policyengine_api/fastapi_routes/v2/households/request_models.py b/policyengine_api/fastapi_routes/v2/households/request_models.py index 32619006a..80af28ab9 100644 --- a/policyengine_api/fastapi_routes/v2/households/request_models.py +++ b/policyengine_api/fastapi_routes/v2/households/request_models.py @@ -26,6 +26,7 @@ class HouseholdCreateRequest(BaseModel): @model_validator(mode="after") def require_country_document(self) -> HouseholdCreateRequest: + """Reject ``household_data`` whose collections belong to the other country.""" expected_type = ( USHouseholdDocument if self.country_id == "us" else UKHouseholdDocument ) @@ -34,6 +35,7 @@ def require_country_document(self) -> HouseholdCreateRequest: return self def to_service_input(self) -> NativeHouseholdCreationInput: + """Convert the request to the household service's creation input.""" return NativeHouseholdCreationInput.model_validate( self.model_dump(mode="python") ) diff --git a/policyengine_api/fastapi_routes/v2/households/response_models.py b/policyengine_api/fastapi_routes/v2/households/response_models.py index c51a8e5d6..35c645c9c 100644 --- a/policyengine_api/fastapi_routes/v2/households/response_models.py +++ b/policyengine_api/fastapi_routes/v2/households/response_models.py @@ -22,6 +22,8 @@ class StrictHouseholdAPIModel(BaseModel): + """Base for household responses; rejects unknown fields and NaN or infinity.""" + model_config = ConfigDict( extra="forbid", allow_inf_nan=False, @@ -30,6 +32,12 @@ class StrictHouseholdAPIModel(BaseModel): class HouseholdItem(StrictHouseholdAPIModel): + """A stored household. + + Households are immutable and deduplicated: creating one with identical + content returns the existing household and its ``id``. + """ + id: UUID country_id: CountryId default_year: DefaultYear | None @@ -39,14 +47,23 @@ class HouseholdItem(StrictHouseholdAPIModel): @classmethod def from_read(cls, item: HouseholdRead) -> HouseholdItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class HouseholdDetailResult(StrictHouseholdAPIModel): + """Result holding one household in ``item``.""" + item: HouseholdItem class HouseholdPageResult(StrictHouseholdAPIModel): + """One page of households, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[HouseholdItem] offset: int limit: int @@ -54,6 +71,7 @@ class HouseholdPageResult(StrictHouseholdAPIModel): @classmethod def from_page(cls, page: HouseholdPage) -> HouseholdPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[HouseholdItem.from_read(item) for item in page.items], offset=page.offset, @@ -66,17 +84,19 @@ def from_page(cls, page: HouseholdPage) -> HouseholdPageResult: class HouseholdSuccessResponse(StrictHouseholdAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class HouseholdDetailResponse(HouseholdSuccessResponse[HouseholdDetailResult]): - pass + """Success response holding one household.""" class HouseholdPageResponse(HouseholdSuccessResponse[HouseholdPageResult]): - pass + """Success response holding one page of households.""" HOUSEHOLD_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { diff --git a/policyengine_api/fastapi_routes/v2/households/routes.py b/policyengine_api/fastapi_routes/v2/households/routes.py index cf0fe4006..6e17e322a 100644 --- a/policyengine_api/fastapi_routes/v2/households/routes.py +++ b/policyengine_api/fastapi_routes/v2/households/routes.py @@ -50,7 +50,6 @@ async def enforce_household_request_size(request: Request) -> None: """Reject declared or actual request bodies larger than 1 MiB.""" - content_length = request.headers.get("content-length") if content_length is not None: try: @@ -103,7 +102,6 @@ def build_v2_household_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build native household routes without opening a database connection.""" - router = APIRouter(prefix="/v2", responses=HOUSEHOLD_ERROR_RESPONSES) create_query = query_dependency(HouseholdCreateQuery) detail_query = query_dependency(HouseholdDetailQuery) diff --git a/policyengine_api/fastapi_routes/v2/me/response_models.py b/policyengine_api/fastapi_routes/v2/me/response_models.py index f04e2079a..c0abd72fc 100644 --- a/policyengine_api/fastapi_routes/v2/me/response_models.py +++ b/policyengine_api/fastapi_routes/v2/me/response_models.py @@ -31,10 +31,18 @@ class StrictAccountAPIModel(BaseModel): + """Base for account responses; rejects unknown fields.""" + model_config = ConfigDict(extra="forbid", from_attributes=True) class MeItem(StrictAccountAPIModel): + """The current user: a signed-in account or a browser's unauthenticated user. + + ``user_type`` says which. An unauthenticated user has no primary country, and + only a signed-in user can set its names or primary country. + """ + id: UUID user_type: Literal["unauthenticated", "standard"] # Null only for an unauthenticated user. @@ -46,14 +54,22 @@ class MeItem(StrictAccountAPIModel): @classmethod def from_read(cls, item: AccountRead) -> MeItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class MeDetailResult(StrictAccountAPIModel): + """Result holding the current user in ``item``.""" + item: MeItem class MeExportResult(StrictAccountAPIModel): + """The current user and every policy, household, simulation and report it saved. + + The export covers every country, not only the requested one. + """ + # Null when an unauthenticated user has not saved anything yet. user: MeItem | None policies: list[UserPolicyItem] @@ -63,6 +79,7 @@ class MeExportResult(StrictAccountAPIModel): @classmethod def from_export(cls, export: AccountExport) -> MeExportResult: + """Build the export result from the service export, converting each record.""" return cls( user=None if export.user is None else MeItem.from_read(export.user), policies=[UserPolicyItem.from_read(item) for item in export.policies], @@ -77,11 +94,19 @@ def from_export(cls, export: AccountExport) -> MeExportResult: class MeUpgradeResult(StrictAccountAPIModel): + """The signed-in account after a browser's unauthenticated user moved into it. + + ``outcome`` is ``converted`` when the unauthenticated user became the account + in place, ``merged`` when its records moved to an existing account, and + ``none`` when there was no unauthenticated user to move. + """ + item: MeItem outcome: Literal["converted", "merged", "none"] @classmethod def from_upgrade(cls, upgrade: AccountUpgrade) -> MeUpgradeResult: + """Build the upgrade result from the service's account and outcome.""" return cls( item=MeItem.from_read(upgrade.account), outcome=upgrade.outcome.value, @@ -92,18 +117,20 @@ def from_upgrade(cls, upgrade: AccountUpgrade) -> MeUpgradeResult: class AccountSuccessResponse(StrictAccountAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class MeDetailResponse(AccountSuccessResponse[MeDetailResult]): - pass + """Success response holding the current user.""" class MeExportResponse(AccountSuccessResponse[MeExportResult]): - pass + """Success response holding the current user's export.""" class MeUpgradeResponse(AccountSuccessResponse[MeUpgradeResult]): - pass + """Success response holding the signed-in account after an upgrade.""" diff --git a/policyengine_api/fastapi_routes/v2/me/routes.py b/policyengine_api/fastapi_routes/v2/me/routes.py index b6694d30d..7bb04376b 100644 --- a/policyengine_api/fastapi_routes/v2/me/routes.py +++ b/policyengine_api/fastapi_routes/v2/me/routes.py @@ -68,7 +68,6 @@ def _service_factory( def build_v2_me_router(dependencies: NativeRouteDependencies) -> APIRouter: """Build account routes without opening a connection.""" - require_user = dependencies.require_user router = APIRouter( prefix="/v2", diff --git a/policyengine_api/fastapi_routes/v2/metadata/common.py b/policyengine_api/fastapi_routes/v2/metadata/common.py index bfe6fc90e..899243334 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/common.py +++ b/policyengine_api/fastapi_routes/v2/metadata/common.py @@ -68,6 +68,12 @@ def read_resource( response_type: type[ResponseT], operation: Callable[[V2MetadataResourceReader], object], ) -> ResponseT | JSONResponse: + """Run one metadata read and wrap its result in ``response_type``. + + Version, pagination and catalog failures become typed v2 error responses + (400, 404 or 503, and 500 for anything unexpected), and the reader is + closed whatever happens. + """ reader = None try: factory = dependencies.v2_metadata_reader_factory diff --git a/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py b/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py index cf3653a54..17ff62e27 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/geography_routes.py @@ -31,6 +31,11 @@ def build_v2_metadata_geography_router( dependencies: NativeRouteDependencies, ) -> APIRouter: + """Build the v2 dataset, region and economy-option routes. + + Each request opens and closes its own metadata reader, so building the + router touches no database. + """ router = APIRouter(prefix="/v2", responses=ERROR_RESPONSES) @router.get( diff --git a/policyengine_api/fastapi_routes/v2/metadata/model_routes.py b/policyengine_api/fastapi_routes/v2/metadata/model_routes.py index 68336b1fe..927b52536 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/model_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/model_routes.py @@ -32,6 +32,11 @@ def build_v2_metadata_model_router( dependencies: NativeRouteDependencies, ) -> APIRouter: + """Build the v2 tax-benefit model, model-version and variable routes. + + Each request opens and closes its own metadata reader, so building the + router touches no database. + """ router = APIRouter(prefix="/v2", responses=ERROR_RESPONSES) @router.get( diff --git a/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py b/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py index 5683a3801..dfe11f152 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py +++ b/policyengine_api/fastapi_routes/v2/metadata/parameter_routes.py @@ -30,6 +30,11 @@ def build_v2_metadata_parameter_router( dependencies: NativeRouteDependencies, ) -> APIRouter: + """Build the v2 parameter, parameter-tree and parameter-value routes. + + Each request opens and closes its own metadata reader, so building the + router touches no database. + """ router = APIRouter(prefix="/v2", responses=ERROR_RESPONSES) @router.get( diff --git a/policyengine_api/fastapi_routes/v2/metadata/response_models.py b/policyengine_api/fastapi_routes/v2/metadata/response_models.py index 6637998e9..d53660834 100644 --- a/policyengine_api/fastapi_routes/v2/metadata/response_models.py +++ b/policyengine_api/fastapi_routes/v2/metadata/response_models.py @@ -24,6 +24,8 @@ class MetadataResourceSuccessResponse(StrictResponseModel, Generic[ResourceT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResourceT @@ -32,67 +34,73 @@ class MetadataResourceSuccessResponse(StrictResponseModel, Generic[ResourceT]): class MetadataModelPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataModel]] ): - pass + """Success response holding a page of the selected catalog's tax-benefit models. + + A catalog has one model per country, so only the first page has an item. + """ class MetadataModelDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataModel]] ): - pass + """Success response holding one tax-benefit model.""" class MetadataModelSelectionResponse( MetadataResourceSuccessResponse[MetadataModelSelectionResult] ): - pass + """Success response holding a country's model and its selected model version.""" class MetadataModelVersionPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataModelVersionDetail]] ): - pass + """Success response holding a page of the selected catalog's model versions. + + A catalog has one model version, so only the first page has an item. + """ class MetadataModelVersionDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataModelVersionDetail]] ): - pass + """Success response holding one model version.""" class MetadataVariablePageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataVariable]] ): - pass + """Success response holding one page of variables.""" class MetadataVariableDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataVariable]] ): - pass + """Success response holding one variable.""" class MetadataParameterPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataParameterSummary]] ): - pass + """Success response holding one page of parameters.""" class MetadataParameterDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataParameterSummary]] ): - pass + """Success response holding one parameter.""" class MetadataParameterChildPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataParameterChild]] ): - pass + """Success response holding one page of a parameter path's direct children.""" class MetadataParameterValuePageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataCanonicalParameterValue]] ): - pass + """Success response holding one page of parameter values.""" class MetadataParameterValueDetailResponse( @@ -100,34 +108,34 @@ class MetadataParameterValueDetailResponse( MetadataDetailResult[MetadataCanonicalParameterValue] ] ): - pass + """Success response holding one parameter value.""" class MetadataDatasetPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataDataset]] ): - pass + """Success response holding one page of datasets.""" class MetadataDatasetDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataDataset]] ): - pass + """Success response holding one dataset.""" class MetadataRegionPageResponse( MetadataResourceSuccessResponse[MetadataPageResult[MetadataRegion]] ): - pass + """Success response holding one page of regions.""" class MetadataRegionDetailResponse( MetadataResourceSuccessResponse[MetadataDetailResult[MetadataRegion]] ): - pass + """Success response holding one region.""" class MetadataEconomyOptionsResponse( MetadataResourceSuccessResponse[MetadataEconomyOptionsResult] ): - pass + """Success response holding the options for configuring an economy comparison.""" diff --git a/policyengine_api/fastapi_routes/v2/policies/response_models.py b/policyengine_api/fastapi_routes/v2/policies/response_models.py index 30cb2a242..0ae6b4887 100644 --- a/policyengine_api/fastapi_routes/v2/policies/response_models.py +++ b/policyengine_api/fastapi_routes/v2/policies/response_models.py @@ -25,6 +25,12 @@ class StrictPolicyAPIModel(BaseModel): class PolicyParameterValueItem(StrictPolicyAPIModel): + """One value a policy sets for a parameter from ``start_date``. + + ``parameter_name`` is the parameter's path in the model, and a null + ``end_date`` means the value has no end. + """ + id: UUID parameter_id: UUID parameter_name: str @@ -34,6 +40,12 @@ class PolicyParameterValueItem(StrictPolicyAPIModel): class PolicyItem(StrictPolicyAPIModel): + """A stored policy: parameter values for one tax-benefit model. + + Policies are immutable and deduplicated: creating one with identical + content returns the existing policy and its ``id``. + """ + id: UUID country_id: str tax_benefit_model_id: UUID @@ -43,14 +55,23 @@ class PolicyItem(StrictPolicyAPIModel): @classmethod def from_read(cls, item: PolicyRead) -> PolicyItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class PolicyDetailResult(StrictPolicyAPIModel): + """Result holding one policy in ``item``.""" + item: PolicyItem class PolicyPageResult(StrictPolicyAPIModel): + """One page of policies, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[PolicyItem] offset: int limit: int @@ -58,6 +79,7 @@ class PolicyPageResult(StrictPolicyAPIModel): @classmethod def from_page(cls, page: PolicyPage) -> PolicyPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[PolicyItem.from_read(item) for item in page.items], offset=page.offset, @@ -70,17 +92,19 @@ def from_page(cls, page: PolicyPage) -> PolicyPageResult: class PolicySuccessResponse(StrictPolicyAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class PolicyDetailResponse(PolicySuccessResponse[PolicyDetailResult]): - pass + """Success response holding one policy.""" class PolicyPageResponse(PolicySuccessResponse[PolicyPageResult]): - pass + """Success response holding one page of policies.""" POLICY_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = { diff --git a/policyengine_api/fastapi_routes/v2/policies/routes.py b/policyengine_api/fastapi_routes/v2/policies/routes.py index 3ecd854ec..9b33d467a 100644 --- a/policyengine_api/fastapi_routes/v2/policies/routes.py +++ b/policyengine_api/fastapi_routes/v2/policies/routes.py @@ -55,7 +55,6 @@ async def enforce_policy_request_size(request: Request) -> None: """Bound both declared and actual request bytes before service creation.""" - content_length = request.headers.get("content-length") if content_length is not None: try: @@ -108,7 +107,6 @@ def build_v2_policy_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build native policy routes without opening a database connection.""" - router = APIRouter(prefix="/v2", responses=POLICY_ERROR_RESPONSES) create_query = query_dependency(PolicyCreateQuery) detail_query = query_dependency(PolicyDetailQuery) diff --git a/policyengine_api/fastapi_routes/v2/routes.py b/policyengine_api/fastapi_routes/v2/routes.py index 863280494..0521719be 100644 --- a/policyengine_api/fastapi_routes/v2/routes.py +++ b/policyengine_api/fastapi_routes/v2/routes.py @@ -44,7 +44,6 @@ def build_v2_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build isolated resource routes without loading v2 configuration.""" - router = APIRouter() router.include_router(build_v2_me_router(dependencies)) router.include_router(build_v2_user_report_router(dependencies)) diff --git a/policyengine_api/fastapi_routes/v2/user_auth.py b/policyengine_api/fastapi_routes/v2/user_auth.py index 42b3098c9..66bc9848a 100644 --- a/policyengine_api/fastapi_routes/v2/user_auth.py +++ b/policyengine_api/fastapi_routes/v2/user_auth.py @@ -47,7 +47,6 @@ def authenticate_user_request(authorization: str | None) -> UserPrincipal: """Return the signed-in user principal or raise a typed HTTP failure.""" - try: settings = load_user_auth_settings() except UserAuthConfigurationError as error: @@ -88,7 +87,6 @@ def authenticate_account_request( unauthenticated-user header is then ignored. Without one, the header selects the browser's unauthenticated user. """ - if authorization is None and unauthenticated_user is not None: try: return UnauthenticatedPrincipal( @@ -103,7 +101,6 @@ def authenticate_account_request( def require_signed_in_principal(principal: AccountPrincipal) -> UserPrincipal: """Refuse an unauthenticated principal where only a signed-in user may act.""" - if isinstance(principal, UserPrincipal): return principal raise V2AuthenticationRequiredError("Sign in with a bearer access token") diff --git a/policyengine_api/fastapi_routes/v2/user_households/response_models.py b/policyengine_api/fastapi_routes/v2/user_households/response_models.py index 2802dde5a..df4137499 100644 --- a/policyengine_api/fastapi_routes/v2/user_households/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_households/response_models.py @@ -21,10 +21,18 @@ class StrictUserHouseholdAPIModel(BaseModel): + """Base for user-household responses; rejects unknown fields.""" + model_config = ConfigDict(extra="forbid", from_attributes=True) class UserHouseholdItem(StrictUserHouseholdAPIModel): + """A household saved by a v2 user, with the user's name and description for it. + + Unlike households, associations can be renamed, pointed at another + household of the same country, or deleted. + """ + id: UUID country_id: CountryId user_id: UUID @@ -37,14 +45,23 @@ class UserHouseholdItem(StrictUserHouseholdAPIModel): @classmethod def from_read(cls, item: UserHouseholdRead) -> UserHouseholdItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class UserHouseholdDetailResult(StrictUserHouseholdAPIModel): + """Result holding one user-household association in ``item``.""" + item: UserHouseholdItem class UserHouseholdPageResult(StrictUserHouseholdAPIModel): + """One page of user-household associations, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[UserHouseholdItem] offset: int limit: int @@ -52,6 +69,7 @@ class UserHouseholdPageResult(StrictUserHouseholdAPIModel): @classmethod def from_page(cls, page: UserHouseholdPage) -> UserHouseholdPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[UserHouseholdItem.from_read(item) for item in page.items], offset=page.offset, @@ -64,6 +82,8 @@ def from_page(cls, page: UserHouseholdPage) -> UserHouseholdPageResult: class UserHouseholdSuccessResponse(StrictUserHouseholdAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT @@ -72,11 +92,11 @@ class UserHouseholdSuccessResponse(StrictUserHouseholdAPIModel, Generic[ResultT] class UserHouseholdDetailResponse( UserHouseholdSuccessResponse[UserHouseholdDetailResult] ): - pass + """Success response holding one user-household association.""" class UserHouseholdPageResponse(UserHouseholdSuccessResponse[UserHouseholdPageResult]): - pass + """Success response holding one page of user-household associations.""" USER_HOUSEHOLD_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = ( diff --git a/policyengine_api/fastapi_routes/v2/user_households/routes.py b/policyengine_api/fastapi_routes/v2/user_households/routes.py index 9bfd9e67c..c5239a1cf 100644 --- a/policyengine_api/fastapi_routes/v2/user_households/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_households/routes.py @@ -64,7 +64,6 @@ def build_v2_user_household_router( Authentication is a router-level dependency, so a missing or invalid credential is rejected before query or body validation. """ - require_user = dependencies.require_user router = APIRouter( prefix="/v2", diff --git a/policyengine_api/fastapi_routes/v2/user_policies/response_models.py b/policyengine_api/fastapi_routes/v2/user_policies/response_models.py index 693b53b4d..7886c40c7 100644 --- a/policyengine_api/fastapi_routes/v2/user_policies/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_policies/response_models.py @@ -40,6 +40,11 @@ class PolicyProvenanceItem(BaseModel): class UserPolicyItem(StrictUserPolicyAPIModel): + """A policy saved by a v2 user, with the user's name and description for it. + + Unlike policies, associations can be renamed or deleted. + """ + id: UUID country_id: CountryId user_id: UUID @@ -54,14 +59,23 @@ class UserPolicyItem(StrictUserPolicyAPIModel): @classmethod def from_read(cls, item: UserPolicyRead) -> UserPolicyItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class UserPolicyDetailResult(StrictUserPolicyAPIModel): + """Result holding one user-policy association in ``item``.""" + item: UserPolicyItem class UserPolicyPageResult(StrictUserPolicyAPIModel): + """One page of user-policy associations, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[UserPolicyItem] offset: int limit: int @@ -69,6 +83,7 @@ class UserPolicyPageResult(StrictUserPolicyAPIModel): @classmethod def from_page(cls, page: UserPolicyPage) -> UserPolicyPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[UserPolicyItem.from_read(item) for item in page.items], offset=page.offset, @@ -81,17 +96,19 @@ def from_page(cls, page: UserPolicyPage) -> UserPolicyPageResult: class UserPolicySuccessResponse(StrictUserPolicyAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class UserPolicyDetailResponse(UserPolicySuccessResponse[UserPolicyDetailResult]): - pass + """Success response holding one user-policy association.""" class UserPolicyPageResponse(UserPolicySuccessResponse[UserPolicyPageResult]): - pass + """Success response holding one page of user-policy associations.""" USER_POLICY_ERROR_RESPONSES: dict[int | str, dict[str, Any]] = ( diff --git a/policyengine_api/fastapi_routes/v2/user_policies/routes.py b/policyengine_api/fastapi_routes/v2/user_policies/routes.py index e201c1635..d9c360585 100644 --- a/policyengine_api/fastapi_routes/v2/user_policies/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_policies/routes.py @@ -61,7 +61,6 @@ def build_v2_user_policy_router( Authentication is a router-level dependency, so a missing or invalid credential is rejected before query or body validation. """ - require_user = dependencies.require_user router = APIRouter( prefix="/v2", diff --git a/policyengine_api/fastapi_routes/v2/user_reports/response_models.py b/policyengine_api/fastapi_routes/v2/user_reports/response_models.py index 6099ab5f8..de4be2d19 100644 --- a/policyengine_api/fastapi_routes/v2/user_reports/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_reports/response_models.py @@ -19,6 +19,8 @@ class StrictUserReportAPIModel(BaseModel): + """Base for user-report responses; rejects unknown fields.""" + model_config = ConfigDict(extra="forbid", from_attributes=True) @@ -35,6 +37,8 @@ class StoredSnapshotItem(BaseModel): class ReportProvisionIntervalItem(StoredSnapshotItem): + """A provision's value for the dates from ``start_date`` to ``end_date``.""" + # The contract makes this key optional rather than nullable. parameter: str | None = Field(default=None, exclude_if=lambda value: value is None) start_date: str @@ -43,6 +47,13 @@ class ReportProvisionIntervalItem(StoredSnapshotItem): class ReportProvisionItem(StoredSnapshotItem): + """One parameter the report's reform changes, with its baseline and reform values. + + ``intervals`` gives the reform's dated values when they vary over time. The + ``growth`` keys, present only when set, say how a value set from one year on + moves in later years and how grown values round. + """ + path: str breadcrumb: str unit: str | None = None @@ -66,6 +77,8 @@ class ReportProvisionItem(StoredSnapshotItem): class ReportCalibrationWorstItem(StoredSnapshotItem): + """The calibration target that one match fits worst, with its relative error.""" + name: str source: str geography: str @@ -73,6 +86,8 @@ class ReportCalibrationWorstItem(StoredSnapshotItem): class ReportCalibrationMatchItem(StoredSnapshotItem): + """One variable's calibration match: its targets, mean error and worst target.""" + variable: str depth: int ring: str @@ -82,6 +97,8 @@ class ReportCalibrationMatchItem(StoredSnapshotItem): class ReportCalibrationItem(StoredSnapshotItem): + """Calibration results saved with a report, for one geography and data release.""" + release_id: str | None = None geography: str reached_count: int @@ -89,10 +106,18 @@ class ReportCalibrationItem(StoredSnapshotItem): class ReportScorecardItem(StoredSnapshotItem): + """The programs on the scorecard saved with a report's validation.""" + programs: list[str] class ReportValidationItem(StoredSnapshotItem): + """The validation snapshot saved with a report. + + It records when the report was matched and with which map model version, + and optionally the calibration results and scorecard. + """ + matched_at: str map_model_version: str calibration: ReportCalibrationItem | None = None @@ -100,6 +125,15 @@ class ReportValidationItem(StoredSnapshotItem): class UserReportItem(StrictUserReportAPIModel): + """A report saved by a user, with the user's name, notes and saved presentation. + + ``legacy_report_id`` is the v1 report output it shows. ``user_policy_id`` is + the user's saved reform the report was run from, and + ``baseline_user_policy_id`` the saved reform it was compared with instead of + current law; each is null when not set or once that reform is deleted. + ``budget_window_years`` is null for a report without a budget window. + """ + id: UUID country_id: CountryId user_id: UUID @@ -119,10 +153,13 @@ class UserReportItem(StrictUserReportAPIModel): @classmethod def from_read(cls, item: UserReportRead) -> UserReportItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class SharedBaselineReformItem(StrictUserReportAPIModel): + """The saved reform a shared report was compared with, named but not linked.""" + name: str | None @@ -144,18 +181,29 @@ class SharedUserReportItem(StrictUserReportAPIModel): @classmethod def from_read(cls, item: SharedUserReportRead) -> SharedUserReportItem: + """Build the shared item by validating the service read record's fields.""" return cls.model_validate(item) class SharedUserReportDetailResult(StrictUserReportAPIModel): + """Result holding one shared report in ``item``.""" + item: SharedUserReportItem class UserReportDetailResult(StrictUserReportAPIModel): + """Result holding one user-report association in ``item``.""" + item: UserReportItem class UserReportPageResult(StrictUserReportAPIModel): + """One page of user-report associations, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[UserReportItem] offset: int limit: int @@ -163,6 +211,7 @@ class UserReportPageResult(StrictUserReportAPIModel): @classmethod def from_page(cls, page: UserReportPage) -> UserReportPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[UserReportItem.from_read(item) for item in page.items], offset=page.offset, @@ -175,20 +224,22 @@ def from_page(cls, page: UserReportPage) -> UserReportPageResult: class UserReportSuccessResponse(StrictUserReportAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT class UserReportDetailResponse(UserReportSuccessResponse[UserReportDetailResult]): - pass + """Success response holding one user-report association.""" class UserReportPageResponse(UserReportSuccessResponse[UserReportPageResult]): - pass + """Success response holding one page of user-report associations.""" class SharedUserReportDetailResponse( UserReportSuccessResponse[SharedUserReportDetailResult] ): - pass + """Success response holding one shared report.""" diff --git a/policyengine_api/fastapi_routes/v2/user_reports/routes.py b/policyengine_api/fastapi_routes/v2/user_reports/routes.py index ee58a678c..502ca468b 100644 --- a/policyengine_api/fastapi_routes/v2/user_reports/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_reports/routes.py @@ -84,7 +84,6 @@ def build_v2_user_report_router( Authentication is a router-level dependency, so a missing or invalid credential is rejected before query or body validation. """ - require_user = dependencies.require_user router = APIRouter( prefix="/v2", @@ -226,7 +225,6 @@ def build_v2_shared_report_router( dependencies: NativeRouteDependencies, ) -> APIRouter: """Build the credential-free read of a saved report by its link.""" - router = APIRouter(prefix="/v2", responses=SHARED_REPORT_ERROR_RESPONSES) country_query = query_dependency(CountryQuery) service_factory = _service_factory(dependencies) diff --git a/policyengine_api/fastapi_routes/v2/user_resources.py b/policyengine_api/fastapi_routes/v2/user_resources.py index d5bd13085..f7dc494a6 100644 --- a/policyengine_api/fastapi_routes/v2/user_resources.py +++ b/policyengine_api/fastapi_routes/v2/user_resources.py @@ -147,7 +147,6 @@ def _log_failure(event: str, error: Exception) -> None: the type and the innermost frame locations are logged. Logging never changes the response. """ - try: frames = traceback.extract_tb(error.__traceback__)[-8:] logger.error( @@ -171,7 +170,6 @@ def user_resource_operation( operation: Callable[[], OperationT], ) -> OperationT | JSONResponse: """Run one service call and map typed failures to secret-safe errors.""" - try: return operation() except StandardUserRequiresSignInError as error: diff --git a/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py b/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py index 21cc621b9..70df6f78a 100644 --- a/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py +++ b/policyengine_api/fastapi_routes/v2/user_simulations/response_models.py @@ -18,10 +18,18 @@ class StrictUserSimulationAPIModel(BaseModel): + """Base for user-simulation responses; rejects unknown fields.""" + model_config = ConfigDict(extra="forbid", from_attributes=True) class UserSimulationItem(StrictUserSimulationAPIModel): + """A simulation saved by a user, with the user's name and description for it. + + ``legacy_simulation_id`` is the v1 simulation it saves. Associations can be + renamed, pointed at another v1 simulation of the same country, or deleted. + """ + id: UUID country_id: CountryId user_id: UUID @@ -34,14 +42,23 @@ class UserSimulationItem(StrictUserSimulationAPIModel): @classmethod def from_read(cls, item: UserSimulationRead) -> UserSimulationItem: + """Build the response item by validating the service read record's fields.""" return cls.model_validate(item) class UserSimulationDetailResult(StrictUserSimulationAPIModel): + """Result holding one user-simulation association in ``item``.""" + item: UserSimulationItem class UserSimulationPageResult(StrictUserSimulationAPIModel): + """One page of user-simulation associations, oldest first. + + ``offset`` and ``limit`` echo the request; when ``has_more`` is true, ask + for the next page with ``offset + limit``. + """ + items: list[UserSimulationItem] offset: int limit: int @@ -49,6 +66,7 @@ class UserSimulationPageResult(StrictUserSimulationAPIModel): @classmethod def from_page(cls, page: UserSimulationPage) -> UserSimulationPageResult: + """Build the response page from a service page, converting each item.""" return cls( items=[UserSimulationItem.from_read(item) for item in page.items], offset=page.offset, @@ -61,6 +79,8 @@ def from_page(cls, page: UserSimulationPage) -> UserSimulationPageResult: class UserSimulationSuccessResponse(StrictUserSimulationAPIModel, Generic[ResultT]): + """Success envelope: ``status`` "ok", a null ``message`` and the ``result``.""" + status: Literal["ok"] = "ok" message: None = None result: ResultT @@ -69,10 +89,10 @@ class UserSimulationSuccessResponse(StrictUserSimulationAPIModel, Generic[Result class UserSimulationDetailResponse( UserSimulationSuccessResponse[UserSimulationDetailResult] ): - pass + """Success response holding one user-simulation association.""" class UserSimulationPageResponse( UserSimulationSuccessResponse[UserSimulationPageResult] ): - pass + """Success response holding one page of user-simulation associations.""" diff --git a/policyengine_api/fastapi_routes/v2/user_simulations/routes.py b/policyengine_api/fastapi_routes/v2/user_simulations/routes.py index 39a3961e5..0915ba779 100644 --- a/policyengine_api/fastapi_routes/v2/user_simulations/routes.py +++ b/policyengine_api/fastapi_routes/v2/user_simulations/routes.py @@ -61,7 +61,6 @@ def build_v2_user_simulation_router( Authentication is a router-level dependency, so a missing or invalid credential is rejected before query or body validation. """ - require_user = dependencies.require_user router = APIRouter( prefix="/v2", diff --git a/policyengine_api/gcp_logging.py b/policyengine_api/gcp_logging.py index 6c3b55b7e..8e89aafe2 100644 --- a/policyengine_api/gcp_logging.py +++ b/policyengine_api/gcp_logging.py @@ -25,7 +25,6 @@ def log_struct( labels: Mapping[str, Any] | None = None, ) -> None: """Record an allowlisted structured message without affecting callers.""" - try: message = str(info.get("message") or "API event") attributes = { diff --git a/policyengine_api/libs/auth0_management.py b/policyengine_api/libs/auth0_management.py index e6c165203..275b0a051 100644 --- a/policyengine_api/libs/auth0_management.py +++ b/policyengine_api/libs/auth0_management.py @@ -57,6 +57,13 @@ class Auth0ManagementError(RuntimeError): @dataclass(frozen=True) class Auth0ManagementSettings: + """Validated Management API settings for one Auth0 tenant. + + Exactly one of ``client_secret`` and ``client_secret_resource`` (a Secret + Manager version name) is set. The secret is kept out of ``repr``, and equal + settings share one client, so its token is reused. + """ + domain: str client_id: str client_secret: str = field(default="", repr=False) @@ -67,7 +74,6 @@ def load_auth0_management_settings( environ: Mapping[str, str] | None = None, ) -> Auth0ManagementSettings | None: """Return complete Management API settings, or ``None`` when absent.""" - values = os.environ if environ is None else environ domain = values.get(AUTH0_MANAGEMENT_DOMAIN_ENV, "").strip() client_id = values.get(AUTH0_MANAGEMENT_CLIENT_ID_ENV, "").strip() @@ -123,6 +129,12 @@ def __init__( http_client: httpx.Client | None = None, secret_loader: Callable[[str], str] = _load_secret_from_secret_manager, ) -> None: + """Use ``settings`` for the tenant and client credentials. + + ``http_client`` replaces the per-request client (tests pass a mock + transport), and ``secret_loader`` reads ``client_secret_resource`` when + no direct secret is set. + """ self._settings = settings self._http_client = http_client self._secret_loader = secret_loader @@ -194,7 +206,6 @@ def _access_token(self) -> str: def _discard_token(self, token: str) -> None: """Forget a rejected token unless another thread already replaced it.""" - with self._lock: if self._token == token: self._token = None @@ -206,7 +217,6 @@ def prepare_deletion(self) -> None: Callers run this before deleting their own records, so a credential or Secret Manager failure stops the deletion with nothing removed. """ - self._access_token() def _delete_request(self, subject: str, token: str) -> httpx.Response: @@ -222,7 +232,6 @@ def delete_user(self, subject: str) -> None: A 401 means the cached token was revoked or rotated, so it is dropped and the request is retried once with a fresh token. """ - token = self._access_token() response = self._delete_request(subject, token) if response.status_code == 401: @@ -241,7 +250,6 @@ def delete_user(self, subject: str) -> None: def _shared_client(settings: Auth0ManagementSettings) -> Auth0ManagementClient: """Return one process-wide client per settings so its token is reused.""" - with _shared_clients_lock: client = _shared_clients.get(settings) if client is None: @@ -260,7 +268,6 @@ def build_auth0_user_deleter( expires. Settings are read on each call, so callers invoke this only when they are about to delete an account. """ - settings = load_auth0_management_settings(environ) if settings is None: return None diff --git a/policyengine_api/libs/gateway_auth.py b/policyengine_api/libs/gateway_auth.py index 8c4b28252..dcdc0dde7 100644 --- a/policyengine_api/libs/gateway_auth.py +++ b/policyengine_api/libs/gateway_auth.py @@ -49,7 +49,10 @@ class GatewayAuthError(RuntimeError): def gateway_auth_required() -> bool: - """True iff this runtime requires gateway auth to be configured.""" + """Return whether ``GATEWAY_AUTH_REQUIRED`` makes missing gateway auth an error. + + Only ``1``, ``true``, ``yes`` and ``on`` (any case) enable it. + """ return os.environ.get(GATEWAY_AUTH_REQUIRED_ENV, "").lower() in ( "1", "true", @@ -120,6 +123,7 @@ def __init__( *, http_timeout: float = 10.0, ) -> None: + """Read each Auth0 setting from its argument, else its environment variable.""" self._issuer = ( issuer if issuer is not None @@ -256,11 +260,13 @@ class GatewayBearerAuth(httpx.Auth): """Attach a bearer token and retry once on a 401.""" def __init__(self, token_provider: GatewayAuthTokenProvider) -> None: + """Authenticate requests with tokens from ``token_provider``.""" self._token_provider = token_provider def auth_flow( self, request: httpx.Request ) -> Generator[httpx.Request, httpx.Response, None]: + """Add a bearer token; on a 401, refetch the token and resend once.""" request.headers["Authorization"] = f"Bearer {self._token_provider.get_token()}" response = yield request diff --git a/policyengine_api/libs/simulation_entrypoint.py b/policyengine_api/libs/simulation_entrypoint.py index ab91b8d41..80af9f989 100644 --- a/policyengine_api/libs/simulation_entrypoint.py +++ b/policyengine_api/libs/simulation_entrypoint.py @@ -61,7 +61,6 @@ def resolve_simulation_entrypoint_url(entrypoint: str | None = None) -> str: def _attach_current_request_id(request: httpx.Request) -> None: """Attach the active API request ID to an outbound simulation request.""" - request_id = current_request_id() if request_id is not None: request.headers[REQUEST_ID_HEADER] = request_id @@ -71,6 +70,8 @@ def _attach_current_request_id(request: httpx.Request) -> None: class SimulationRunSegments(BaseModel): + """How many of a running job's segments have finished, out of its total.""" + model_config = ConfigDict(extra="ignore", frozen=True) completed: int = Field(ge=0) @@ -106,14 +107,13 @@ def from_status(cls, data: dict) -> "SimulationRunProgress | None": return run_progress if run_progress.to_dict() else None def to_dict(self) -> dict: + """Return the reported fields as JSON values, leaving out unreported ones.""" return self.model_dump(mode="json", exclude_none=True) @dataclass class ModalSimulationExecution: - """ - Represents a Modal simulation job execution. - """ + """Represents a Modal simulation job execution.""" job_id: str status: str @@ -131,9 +131,7 @@ def name(self) -> str: @dataclass class ModalBudgetWindowBatchExecution: - """ - Represents a budget-window batch execution in the Simulation Entrypoint. - """ + """Represents a budget-window batch execution in the Simulation Entrypoint.""" batch_job_id: str status: str @@ -152,14 +150,18 @@ def name(self) -> str: class SimulationEntrypointClient: - """ - HTTP client for the configured Simulation Entrypoint. + """HTTP client for the configured Simulation Entrypoint. This class provides methods for submitting simulation jobs and polling for their status/results via HTTP endpoints. """ def __init__(self, entrypoint: str | None = None) -> None: + """Connect to ``entrypoint`` (default: the configured one) with gateway auth. + + Raise GatewayAuthError if auth is partly configured, or is required here + but not configured. + """ self.entrypoint = entrypoint or get_sim_entrypoint() self.base_url = resolve_simulation_entrypoint_url(self.entrypoint) self._token_provider = GatewayAuthTokenProvider() @@ -232,8 +234,7 @@ def _normalize_submission_payload(self, payload: dict) -> dict: return modal_payload def run(self, payload: dict) -> ModalSimulationExecution: - """ - Submit a simulation job to the selected Simulation entrypoint. + """Submit a simulation job to the selected Simulation entrypoint. Parameters ---------- @@ -250,6 +251,7 @@ def run(self, payload: dict) -> ModalSimulationExecution: ------ httpx.HTTPStatusError If the API returns an error response. + """ try: modal_payload = self._normalize_submission_payload(payload) @@ -301,9 +303,7 @@ def run(self, payload: dict) -> ModalSimulationExecution: raise def run_budget_window_batch(self, payload: dict) -> ModalBudgetWindowBatchExecution: - """ - Submit a budget-window batch job to the selected Simulation entrypoint. - """ + """Submit a budget-window batch job to the selected Simulation entrypoint.""" try: modal_payload = self._normalize_submission_payload(payload) @@ -356,6 +356,13 @@ def get_spm_capability( *, policyengine_version: str | None = None, ) -> object: + """Return the worker's published SPM capability for the bundle it would run. + + The bundle is the one the worker's ``/versions`` routes resolve for + ``policyengine_version`` or else ``country``/``version``; None means that + bundle publishes no capability. Raise SPMValidationError if the routes + are unreachable or ambiguous. + """ from policyengine_api.spm import SPMValidationError try: @@ -445,8 +452,7 @@ def resolve_app_name( ) from exc def get_execution_id(self, execution: ModalSimulationExecution) -> str: - """ - Get the job ID from an execution. + """Get the job ID from an execution. Parameters ---------- @@ -457,12 +463,12 @@ def get_execution_id(self, execution: ModalSimulationExecution) -> str: ------- str The job ID. + """ return execution.job_id def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: - """ - Poll the selected Simulation entrypoint for the current status of a job. + """Poll the selected Simulation entrypoint for the current status of a job. Parameters ---------- @@ -473,6 +479,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: ------- ModalSimulationExecution Execution object with current status and result if complete. + """ try: response = self.client.get(f"{self.base_url}/jobs/{job_id}") @@ -513,9 +520,7 @@ def get_execution_by_id(self, job_id: str) -> ModalSimulationExecution: def get_budget_window_batch_by_id( self, batch_job_id: str ) -> ModalBudgetWindowBatchExecution: - """ - Poll the selected Simulation entrypoint for a budget-window batch. - """ + """Poll the selected Simulation entrypoint for a budget-window batch.""" try: response = self.client.get( f"{self.base_url}/budget-window-jobs/{batch_job_id}" @@ -557,8 +562,7 @@ def get_budget_window_batch_by_id( raise def get_execution_status(self, execution: ModalSimulationExecution) -> str: - """ - Get the status string from an execution. + """Get the status string from an execution. Parameters ---------- @@ -569,12 +573,12 @@ def get_execution_status(self, execution: ModalSimulationExecution) -> str: ------- str The status string ("submitted", "running", "complete", "failed"). + """ return execution.status def get_execution_result(self, execution: ModalSimulationExecution) -> dict | None: - """ - Get the result from a completed execution. + """Get the result from a completed execution. Parameters ---------- @@ -585,17 +589,18 @@ def get_execution_result(self, execution: ModalSimulationExecution) -> dict | No ------- dict or None The simulation result if complete, None otherwise. + """ return execution.result def health_check(self) -> bool: - """ - Check if the selected Simulation entrypoint is healthy. + """Check if the selected Simulation entrypoint is healthy. Returns ------- bool True if the API is healthy, False otherwise. + """ try: response = self.client.get(f"{self.base_url}/health") diff --git a/policyengine_api/libs/user_auth.py b/policyengine_api/libs/user_auth.py index df218da9a..6f6a3f0ad 100644 --- a/policyengine_api/libs/user_auth.py +++ b/policyengine_api/libs/user_auth.py @@ -81,6 +81,11 @@ class InvalidUnauthenticatedUserError(UserAuthError): class PrincipalKind(StrEnum): + """Whether a token was issued to a person or to a machine-to-machine client. + + User routes refuse client principals. + """ + USER = "user" CLIENT = "client" @@ -95,6 +100,7 @@ class UserAuthSettings: @property def configured(self) -> bool: + """Whether an issuer and audience are set, so bearer tokens can be checked.""" return bool(self.issuer and self.audience) @@ -109,7 +115,11 @@ class VerifiedPrincipal: class TokenDecoder(Protocol): - def decode(self, token: str) -> Mapping[str, Any]: ... + """Verify an access token and return its claims; tests inject a fake.""" + + def decode(self, token: str) -> Mapping[str, Any]: + """Return the claims of a valid token, or raise a PyJWT error.""" + ... def _normalized_issuer(value: str) -> str: @@ -137,7 +147,6 @@ def load_user_auth_settings( environ: Mapping[str, str] | None = None, ) -> UserAuthSettings: """Load and validate the inbound user-authentication configuration.""" - values = os.environ if environ is None else environ required_value = values.get(PUBLIC_API_AUTH_REQUIRED_ENV, "").strip() if required_value not in {"", "0", "1"}: @@ -172,6 +181,11 @@ def __init__( *, jwks_client: jwt.PyJWKClient | None = None, ) -> None: + """Accept tokens from ``issuer`` for ``audience``. + + Without ``jwks_client``, signing keys are fetched from the issuer's + ``.well-known/jwks.json`` and cached by key ID. + """ self.issuer = issuer self.audience = audience self.jwks_client = jwks_client or jwt.PyJWKClient( @@ -181,6 +195,11 @@ def __init__( ) def decode(self, token: str) -> Mapping[str, Any]: + """Verify the signature, issuer, audience and required claims; return claims. + + Raise ``jwt.PyJWKClientConnectionError`` when the signing keys cannot be + fetched, and another ``jwt.PyJWTError`` when the token is invalid. + """ signing_key = self.jwks_client.get_signing_key_from_jwt(token).key return jwt.decode( token, @@ -199,13 +218,11 @@ def _decoder(issuer: str, audience: str) -> JWTDecoder: def reset_decoder_cache() -> None: """Clear cached decoders for tests and signing-key or tenant rotations.""" - _decoder.cache_clear() def bearer_token(authorization: str | None) -> str: """Return the credentials of an ``Authorization: Bearer`` header.""" - if not authorization: raise MissingBearerTokenError("a bearer access token is required") scheme, _, credentials = authorization.strip().partition(" ") @@ -217,7 +234,6 @@ def bearer_token(authorization: str | None) -> str: def unauthenticated_user_id(value: str) -> UUID: """Return the browser UUID carried by the unauthenticated-user header.""" - if re.fullmatch(CANONICAL_UUID_PATTERN, value) is None: raise InvalidUnauthenticatedUserError( f"{UNAUTHENTICATED_USER_HEADER} must be a canonical UUID" @@ -237,7 +253,6 @@ def _verified_email(claims: Mapping[str, Any]) -> str | None: def classify_principal(claims: Mapping[str, Any]) -> VerifiedPrincipal: """Build a principal from already-validated token claims.""" - subject = claims.get("sub") issuer = claims.get("iss") if not isinstance(subject, str) or not subject.strip(): @@ -263,7 +278,6 @@ def authenticate_bearer( decoder: TokenDecoder | None = None, ) -> VerifiedPrincipal: """Validate one Authorization header and classify its principal.""" - token = bearer_token(authorization) if not settings.configured: raise UserAuthUnavailableError("user authentication is not configured") diff --git a/policyengine_api/migration_flags.py b/policyengine_api/migration_flags.py index 62ecb0ae8..142c1cfd1 100644 --- a/policyengine_api/migration_flags.py +++ b/policyengine_api/migration_flags.py @@ -47,6 +47,8 @@ class RouteImplementation(StrEnum): @dataclass(frozen=True) class MigrationContext: + """The migration flag values in force for one request, as recorded in its logs.""" + route_group: str route_impl: RouteImplementation db_entity: str | None @@ -57,6 +59,7 @@ class MigrationContext: sim_compute: str | None def to_log_dict(self) -> dict: + """Return the flags as a flat log record, with ``route_impl`` as its value.""" return { "route_group": self.route_group, "route_impl": self.route_impl.value, @@ -109,6 +112,10 @@ def infer_route_group(path: str) -> str: def get_route_impl(route_group: str) -> RouteImplementation: + """Read ``ROUTE_IMPL_``, defaulting to the Flask fallback. + + Raises ValueError if the variable names an unknown implementation. + """ env_name = f"ROUTE_IMPL_{route_group.upper()}" return RouteImplementation( _read_choice( @@ -138,18 +145,25 @@ def from_environment(cls) -> RouteImplementationSettings: def get_db_write(entity: str) -> str: + """Read the ``DB_WRITE_`` write source, defaulting to Cloud SQL. + + Raises ValueError if the variable names an unknown source. + """ env_name = f"DB_WRITE_{entity.upper()}" return _read_choice(env_name, DEFAULT_DB_SOURCE, DB_WRITE_SOURCES) def get_db_read(entity: str) -> str: + """Read the ``DB_READ_`` read source, defaulting to Cloud SQL. + + Raises ValueError if the variable names an unknown source. + """ env_name = f"DB_READ_{entity.upper()}" return _read_choice(env_name, DEFAULT_DB_SOURCE, DB_READ_SOURCES) def get_v1_policy_write_source() -> str: """Select Cloud SQL alone or immediate Cloud SQL-to-Supabase mirroring.""" - return _read_choice( "DB_WRITE_POLICY", DEFAULT_DB_SOURCE, @@ -159,7 +173,6 @@ def get_v1_policy_write_source() -> str: def get_v1_policy_read_source() -> str: """Require every v1 policy read to remain on Cloud SQL in Phase 10.""" - return _read_choice( "DB_READ_POLICY", DEFAULT_DB_SOURCE, @@ -169,7 +182,6 @@ def get_v1_policy_read_source() -> str: def get_v1_household_write_source() -> str: """Select Cloud SQL alone or immediate copying to Supabase.""" - return _read_choice( "DB_WRITE_HOUSEHOLD", DEFAULT_DB_SOURCE, @@ -179,7 +191,6 @@ def get_v1_household_write_source() -> str: def get_v1_household_read_source() -> str: """Require every v1 household read to remain on Cloud SQL in Phase 11.""" - return _read_choice( "DB_READ_HOUSEHOLD", DEFAULT_DB_SOURCE, @@ -188,6 +199,10 @@ def get_v1_household_read_source() -> str: def get_sim_compute(flow: str) -> str: + """Read the ``SIM_COMPUTE_`` compute backend, defaulting to the old gateway. + + Raises ValueError if the variable names an unknown backend. + """ env_name = f"SIM_COMPUTE_{flow.upper()}" return _read_choice( env_name, diff --git a/policyengine_api/migration_logging.py b/policyengine_api/migration_logging.py index 08a6d0e86..ce32de74e 100644 --- a/policyengine_api/migration_logging.py +++ b/policyengine_api/migration_logging.py @@ -162,7 +162,6 @@ def log_migration_request( route_impl: RouteImplementation | None = None, ) -> None: """Log a migration-aware API request in the shared structured format.""" - elapsed_ms = None if started_at is not None: elapsed_ms = round((time.time() - started_at) * 1000, 2) diff --git a/policyengine_api/migration_registry.py b/policyengine_api/migration_registry.py index 4d0100f7a..e49de6986 100644 --- a/policyengine_api/migration_registry.py +++ b/policyengine_api/migration_registry.py @@ -7,6 +7,12 @@ @dataclass(frozen=True) class RouteGroupConfig: + """A migratable route group and the URL path segments that identify its requests. + + ``db_entity`` and ``sim_flow`` select which database-source and + simulation-compute flags apply to the group's requests. + """ + name: str path_segments: tuple[str, ...] db_entity: str | None = None diff --git a/policyengine_api/observability/identifiers.py b/policyengine_api/observability/identifiers.py index 57be0c593..64d8899d0 100644 --- a/policyengine_api/observability/identifiers.py +++ b/policyengine_api/observability/identifiers.py @@ -9,13 +9,11 @@ def generate_observability_id() -> str: """Create an identifier used only to correlate observability records.""" - return str(uuid4()) def normalize_observability_id(value: object) -> str | None: """Return a canonical UUID string, or ``None`` for malformed input.""" - if not isinstance(value, str): return None try: diff --git a/policyengine_api/observability/runtime.py b/policyengine_api/observability/runtime.py index aea6f10c2..a560b5378 100644 --- a/policyengine_api/observability/runtime.py +++ b/policyengine_api/observability/runtime.py @@ -68,12 +68,12 @@ def _build_runtime() -> ObservabilityRuntime: def get_runtime() -> ObservabilityRuntime: + """Return the process-wide runtime that was configured when this module loaded.""" return runtime def set_runtime_context(**attributes: object) -> None: """Bind local telemetry attributes without affecting application behavior.""" - try: runtime.set_context(**attributes) except Exception: diff --git a/policyengine_api/observability/stages.py b/policyengine_api/observability/stages.py index 2c517542b..698a69475 100644 --- a/policyengine_api/observability/stages.py +++ b/policyengine_api/observability/stages.py @@ -12,12 +12,16 @@ class RunConfiguration(StrEnum): + """A kind of calculation request, each with its own registered stage plan.""" + HOUSEHOLD = "household" ECONOMY_ANNUAL = "economy_annual" ECONOMY_BUDGET_WINDOW = "economy_budget_window" class Stage(StrEnum): + """A named calculation step; each value is the span name telemetry records.""" + HOUSEHOLD_LOAD_INPUTS = "household.load_inputs" HOUSEHOLD_CACHE_LOOKUP = "household.cache_lookup" HOUSEHOLD_INPUT_NORMALIZATION = "household.input_normalization" @@ -43,10 +47,13 @@ class Stage(StrEnum): @dataclass(frozen=True) class StagePlan: + """The stages one run configuration is allowed to record as spans.""" + configuration: RunConfiguration stages: tuple[Stage, ...] def name(self, stage: Stage) -> str: + """Return the span name for ``stage``; raise ValueError if not in this plan.""" if stage not in self.stages: raise ValueError( f"{stage.value!r} is not registered for {self.configuration.value!r}" diff --git a/policyengine_api/query_parameters.py b/policyengine_api/query_parameters.py index 7584c7d68..90535d5a3 100644 --- a/policyengine_api/query_parameters.py +++ b/policyengine_api/query_parameters.py @@ -45,13 +45,11 @@ def normalize_country_id(value: object) -> object: """Lowercase a textual country ID before its supported-value check.""" - return value.lower() if isinstance(value, str) else value def validate_legacy_user_id(value: str) -> str: """Reject an empty or whitespace-only legacy user identifier.""" - if not value.strip(): raise ValueError( "legacy user_id must contain at least one non-whitespace character" @@ -184,6 +182,7 @@ class BudgetWindowEconomyQuery(EconomyQuery): @model_validator(mode="after") def validate_end_year(self) -> BudgetWindowEconomyQuery: + """Reject a window whose last year falls after BUDGET_WINDOW_MAX_END_YEAR.""" if int(self.start_year) + self.window_size - 1 > BUDGET_WINDOW_MAX_END_YEAR: raise ValueError( f"budget-window end_year must be {BUDGET_WINDOW_MAX_END_YEAR} or earlier" @@ -287,6 +286,7 @@ class DuplicateScalarQueryParameterError(ValueError): """Raised when a scalar query field is supplied more than once.""" def __init__(self, parameter_name: str) -> None: + """Record the public name of the repeated query parameter.""" self.parameter_name = parameter_name super().__init__( f"scalar query parameter {parameter_name!r} must not be repeated" @@ -306,7 +306,6 @@ def query_field_multiplicity( model_type: type[StrictQueryParameters], ) -> Mapping[str, bool]: """Return public query names mapped to whether repeated values are valid.""" - return { (field.alias or field_name): _annotation_is_list(field.annotation) for field_name, field in model_type.model_fields.items() @@ -318,7 +317,6 @@ def validate_scalar_query_multiplicity( items: Iterable[tuple[str, str]], ) -> None: """Reject a repeated declared scalar while allowing declared list fields.""" - multiplicity = query_field_multiplicity(model_type) seen: set[str] = set() for name, _value in items: @@ -334,7 +332,6 @@ def parse_query_items( items: Iterable[tuple[str, str]], ) -> QueryParametersT: """Validate ordered query pairs through one canonical Pydantic schema.""" - pairs = list(items) validate_scalar_query_multiplicity(model_type, pairs) multiplicity = query_field_multiplicity(model_type) @@ -355,7 +352,6 @@ def parse_multidict_query( query_input: object, ) -> QueryParametersT: """Adapt a Flask/Werkzeug-style MultiDict to the canonical parser.""" - items_method = getattr(query_input, "items", None) if not callable(items_method): raise TypeError("query input must provide an items method") diff --git a/policyengine_api/readiness.py b/policyengine_api/readiness.py index 106ca424a..7d7eb4bb3 100644 --- a/policyengine_api/readiness.py +++ b/policyengine_api/readiness.py @@ -16,7 +16,6 @@ def validate_policy_runtime_configuration() -> None: """Validate v1 source selectors, v2 persistence, user auth and measurement.""" - from policyengine_api.data.v2.settings import ( load_v2_runtime_database_settings, ) diff --git a/policyengine_api/request_context.py b/policyengine_api/request_context.py index 6e7fd8ef3..8c344594d 100644 --- a/policyengine_api/request_context.py +++ b/policyengine_api/request_context.py @@ -30,13 +30,11 @@ def generate_request_id() -> str: """Generate an API request correlation identifier.""" - return uuid.uuid4().hex def current_request_id() -> str | None: """Return the correlation identifier for the current request, if any.""" - if flask.has_request_context(): return getattr(flask.g, "request_id", None) return _asgi_request_id.get() @@ -44,7 +42,6 @@ def current_request_id() -> str | None: def current_observability_id() -> str | None: """Return the diagnostic correlation identifier for the current request.""" - if flask.has_request_context(): return getattr(flask.g, "observability_id", None) return _asgi_observability_id.get() @@ -52,7 +49,6 @@ def current_observability_id() -> str | None: def incoming_observability_id() -> str | None: """Return the validated identifier candidate supplied by the caller.""" - if flask.has_request_context(): return getattr(flask.g, "incoming_observability_id", None) return _asgi_incoming_observability_id.get() @@ -74,7 +70,6 @@ def _bind_observability_id(observability_id: str) -> str: def start_observability_id(value: object = None) -> str: """Bind the identifier for a newly accepted calculation or report.""" - current = current_observability_id() if current is not None: return current @@ -86,7 +81,6 @@ def start_observability_id(value: object = None) -> str: def restore_observability_id(value: object) -> str | None: """Bind a valid identifier read from durable functional state.""" - observability_id = normalize_observability_id(value) if observability_id is None: return current_observability_id() @@ -95,5 +89,4 @@ def restore_observability_id(value: object) -> str | None: def resolve_observability_id(value: object) -> str: """Use a valid caller value or create a new diagnostic identifier.""" - return normalize_observability_id(value) or generate_observability_id() diff --git a/policyengine_api/routes/economy_routes.py b/policyengine_api/routes/economy_routes.py index 8a2f83aeb..5ab331b4e 100644 --- a/policyengine_api/routes/economy_routes.py +++ b/policyengine_api/routes/economy_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for society-wide economic impacts of a policy against a baseline.""" + import json from http import HTTPStatus @@ -73,6 +75,13 @@ def _bad_request_response(error: str | ValueError) -> Response: def get_economic_impact( country_id: str, policy_id: int, baseline_policy_id: int ) -> Response: + """Return the single-year economic impact, starting it if none is cached. + + A policy ID of 0 means current law. The result stays ``computing`` until the + simulation finishes, and while computing the body also carries the stage + fields the Simulation Entrypoint reports. Invalid query parameters give 400 + and a failed calculation gives 502. + """ policy_id = int(policy_id or get_current_law_policy_id(country_id)) baseline_policy_id = int( baseline_policy_id or get_current_law_policy_id(country_id) @@ -128,6 +137,13 @@ def get_economic_impact( def get_budget_window_economic_impact( country_id: str, policy_id: int, baseline_policy_id: int ) -> Response: + """Return a multi-year budget-window impact with per-year progress. + + The body lists completed, computing and queued years, and the + ``X-PolicyEngine-Budget-Window-Cache`` header reports the cache outcome when + known. Invalid query parameters give 400; a failed calculation is reported + in the body with status 200. + """ policy_id = int(policy_id or get_current_law_policy_id(country_id)) baseline_policy_id = int( baseline_policy_id or get_current_law_policy_id(country_id) diff --git a/policyengine_api/routes/error_routes.py b/policyengine_api/routes/error_routes.py index b5f7ddb5f..ab6005351 100644 --- a/policyengine_api/routes/error_routes.py +++ b/policyengine_api/routes/error_routes.py @@ -1,3 +1,5 @@ +"""App-wide Flask error handlers that return the shared JSON error envelope.""" + from flask import Blueprint, Response from werkzeug.exceptions import ( HTTPException, @@ -10,41 +12,41 @@ @error_bp.app_errorhandler(404) def response_404(error: HTTPException) -> Response: - """Specific handler for 404 Not Found errors""" + """Return 404 Not Found as a JSON error body.""" return _make_error_response(error, 404, result=None) @error_bp.app_errorhandler(400) def response_400(error: HTTPException) -> Response: - """Specific handler for 400 Bad Request errors""" + """Return 400 Bad Request as a JSON error body.""" return _make_error_response(error, 400, result=None) @error_bp.app_errorhandler(401) def response_401(error: HTTPException) -> Response: - """Specific handler for 401 Unauthorized errors""" + """Return 401 Unauthorized as a JSON error body.""" return _make_error_response(error, 401, result=None) @error_bp.app_errorhandler(403) def response_403(error: HTTPException) -> Response: - """Specific handler for 403 Forbidden errors""" + """Return 403 Forbidden as a JSON error body.""" return _make_error_response(error, 403, result=None) @error_bp.app_errorhandler(500) def response_500(error: HTTPException) -> Response: - """Specific handler for 500 Internal Server errors""" + """Return 500 Internal Server Error as a JSON error body.""" return _make_error_response(error, 500, result=None) @error_bp.app_errorhandler(HTTPException) def response_http_exception(error: HTTPException) -> Response: - """Generic handler for HTTPException; should be raised if no specific handler is found""" + """Return any other HTTPException as a JSON error body with its own status.""" return _make_error_response(error, error.code, result=None) @error_bp.app_errorhandler(Exception) def response_generic_error(error: Exception) -> Response: - """Handler for any unhandled exceptions""" + """Return an unhandled exception as a 500 JSON error body carrying its message.""" return _make_error_response(error, 500, result=None) diff --git a/policyengine_api/routes/home_routes.py b/policyengine_api/routes/home_routes.py index ac05ccf44..b550d8ddb 100644 --- a/policyengine_api/routes/home_routes.py +++ b/policyengine_api/routes/home_routes.py @@ -1,3 +1,5 @@ +"""Flask route for the API's HTML landing page.""" + from flask import Blueprint home_bp = Blueprint("home", __name__) diff --git a/policyengine_api/routes/household_routes.py b/policyengine_api/routes/household_routes.py index 7eefd896e..0235fea20 100644 --- a/policyengine_api/routes/household_routes.py +++ b/policyengine_api/routes/household_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for v1 households and calculations, with optional copying to API v2.""" + import json import logging import time @@ -258,12 +260,12 @@ def _calculation_cache_key(*args: object, **kwargs: object) -> str: @household_bp.route("//household/", methods=["GET"]) @validate_country def get_household(country_id: str, household_id: int) -> Response: - """ - Get a household's input data with a given ID. + """Get a household's input data with a given ID. Args: country_id (str): The country ID. household_id (int): The household ID. + """ try: get_v1_household_read_source() @@ -290,13 +292,12 @@ def get_household(country_id: str, household_id: int) -> Response: @household_bp.route("//household", methods=["POST"]) @validate_country def post_household(country_id: str) -> Response: - """ - Set a household's input data. + """Set a household's input data. Args: country_id (str): The country ID. - """ + """ # Validate payload payload = request.json is_payload_valid, message = validate_household_payload(payload) diff --git a/policyengine_api/routes/metadata_routes.py b/policyengine_api/routes/metadata_routes.py index 274585fe4..ad658cbe8 100644 --- a/policyengine_api/routes/metadata_routes.py +++ b/policyengine_api/routes/metadata_routes.py @@ -1,3 +1,5 @@ +"""Flask route serving a country's model metadata document.""" + import json from flask import Blueprint, Response @@ -20,8 +22,8 @@ def get_metadata(country_id: str) -> Response: Args: country_id (str): The country ID. - """ + """ # Retrieve country metadata and add status and message to the response country_metadata = metadata_service.get_metadata(country_id) body = json.dumps( diff --git a/policyengine_api/routes/policy_routes.py b/policyengine_api/routes/policy_routes.py index d27884f2a..5f5e4d61a 100644 --- a/policyengine_api/routes/policy_routes.py +++ b/policyengine_api/routes/policy_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for v1 policies and saved user policies, optionally mirrored to v2.""" + import json import time @@ -41,7 +43,6 @@ def _should_mirror_to_v2(country_id: str, write_source: str) -> bool: """Select immediate mirroring only for countries supported by API v2.""" - return write_source == "dual_write" and country_id in SUPPORTED_V2_COUNTRY_IDS @@ -87,7 +88,6 @@ def _user_policy_persistence_failure( legacy_user_policy_id: object | None = None, ) -> Response: """Return and record a failure without serializing exception details.""" - status_code = 503 if error.retryable else 500 try: logger.log_struct( @@ -144,18 +144,17 @@ def _serialize_policy(policy: Policy) -> dict: @policy_bp.route("//policy/", methods=["GET"]) @validate_country def get_policy(country_id: str, policy_id: int | str) -> Response: - """ - Get policy data for a given country and policy ID. + """Return one stored v1 policy. Args: - country_id (str) - policy_id (int | str) + country_id (str): Country whose policies are searched. + policy_id (int | str): ID of the stored policy; cast to int. Returns: - Response: A Flask response object containing the - policy data in JSON format - """ + Response: The policy as JSON under ``result``. An unknown ID gives 404, + and an invalid policy read-source setting gives 503. + """ try: get_v1_policy_read_source() except ValueError: @@ -179,14 +178,16 @@ def get_policy(country_id: str, policy_id: int | str) -> Response: @policy_bp.route("//policy", methods=["POST"]) @validate_country def set_policy(country_id: str) -> Response: - """ - Set policy data for given country and policy. If policy already exists, - return existing policy and 200. + """Store a policy, or return the ID of an identical one already stored. + + A policy is identical when its country, content hash and label all match; + that case returns 200, a new policy returns 201. An invalid payload gives + 400, and an unavailable write configuration or v2 mirror gives 503. Args: - country_id (str) - """ + country_id (str): Country to store the policy under. + """ payload = request.json is_payload_valid, message = validate_set_policy_payload(payload) diff --git a/policyengine_api/routes/reform_impact_routes.py b/policyengine_api/routes/reform_impact_routes.py index c3a9d2921..21bbf6df9 100644 --- a/policyengine_api/routes/reform_impact_routes.py +++ b/policyengine_api/routes/reform_impact_routes.py @@ -1,3 +1,5 @@ +"""Flask route listing recent cached reform impact calculations.""" + import json from datetime import datetime from typing import Any @@ -24,7 +26,6 @@ def _serialize_v1_datetime(value: datetime | None) -> str | None: def _serialize_v1_reform_impact(impact: CachedReformImpact) -> dict[str, object]: """Project a cached impact onto the explicit historical v1 response shape.""" - return { "reform_impact_id": impact.reform_impact_id, "baseline_policy_id": impact.baseline_policy_id, diff --git a/policyengine_api/routes/report_output_routes.py b/policyengine_api/routes/report_output_routes.py index b105b33ba..b8a919de3 100644 --- a/policyengine_api/routes/report_output_routes.py +++ b/policyengine_api/routes/report_output_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for creating, reading and updating v1 report outputs.""" + import json import jsonschema @@ -19,7 +21,6 @@ def _serialize_v1_report_output(view: ReportOutputView) -> dict: """Project mapped report state onto the historical v1 response shape.""" - report_output = view.report_output result = { column.name: getattr(report_output, column.name) @@ -40,8 +41,7 @@ def _serialize_v1_report_output(view: ReportOutputView) -> dict: @report_output_bp.route("//report", methods=["POST"]) @validate_country def create_report_output(country_id: str) -> Response: - """ - Create a new report output record with pending status. + """Create a new report output record with pending status. Args: country_id (str): The country ID. @@ -50,6 +50,7 @@ def create_report_output(country_id: str) -> Response: - simulation_1_id (int): The first simulation ID (required) - simulation_2_id (int | None): The second simulation ID (optional, for comparisons) - year (str | None): The year for the report (optional) + """ print(f"Creating report output for country {country_id}") payload = request.json @@ -117,8 +118,7 @@ def create_report_output(country_id: str) -> Response: @report_output_bp.route("//report/", methods=["GET"]) @validate_country def get_report_output(country_id: str, report_id: int) -> Response: - """ - Get a report output record by ID. + """Get a report output record by ID. The response result may include requested_at, started_at, and finished_at values projected from the selected report_output_runs row. Those fields are @@ -128,6 +128,7 @@ def get_report_output(country_id: str, report_id: int) -> Response: Args: country_id (str): The country ID. report_id (int): The report output ID. + """ print(f"Getting report output {report_id} for country {country_id}") @@ -152,8 +153,7 @@ def get_report_output(country_id: str, report_id: int) -> Response: @report_output_bp.route("//report", methods=["PATCH"]) @validate_country def update_report_output(country_id: str) -> Response: - """ - Update a report output record with results or error. + """Update a report output record with results or error. Args: country_id (str): The country ID. @@ -164,8 +164,8 @@ def update_report_output(country_id: str) -> Response: - output (dict): The result output (for complete status) - api_version (str): The API version of the report - error_message (str): The error message (for error status) - """ + """ payload = request.json if payload is None: raise BadRequest("Payload missing from request") diff --git a/policyengine_api/routes/simulation_routes.py b/policyengine_api/routes/simulation_routes.py index 5e9b329e6..76ca6e7db 100644 --- a/policyengine_api/routes/simulation_routes.py +++ b/policyengine_api/routes/simulation_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for creating, reading and updating v1 simulation records.""" + import json import jsonschema @@ -28,7 +30,6 @@ def _spm_override_response(payload: dict) -> Response | None: def _serialize_v1_simulation(simulation: Simulation) -> dict: """Project canonical ORM JSON objects onto the legacy v1 response shape.""" - response = { column.name: getattr(simulation, column.name) for column in Simulation.__table__.columns @@ -47,8 +48,7 @@ def _serialize_v1_simulation(simulation: Simulation) -> dict: @simulation_bp.route("//simulation", methods=["POST"]) @validate_country def create_simulation(country_id: str) -> Response: - """ - Create a new simulation record. + """Create a new simulation record. Args: country_id (str): The country ID. @@ -57,6 +57,7 @@ def create_simulation(country_id: str) -> Response: - population_id (str): The population identifier (household or geography ID) - population_type (str): Type of population ('household' or 'geography') - policy_id (int): The policy ID + """ print(f"Creating simulation for country {country_id}") @@ -129,12 +130,12 @@ def create_simulation(country_id: str) -> Response: @simulation_bp.route("//simulation/", methods=["GET"]) @validate_country def get_simulation(country_id: str, simulation_id: int) -> Response: - """ - Get a simulation record by ID. + """Get a simulation record by ID. Args: country_id (str): The country ID. simulation_id (int): The simulation ID. + """ print(f"Getting simulation {simulation_id} for country {country_id}") @@ -165,8 +166,7 @@ def get_simulation(country_id: str, simulation_id: int) -> Response: @simulation_bp.route("//simulation", methods=["PATCH"]) @validate_country def update_simulation(country_id: str) -> Response: - """ - Update a simulation record with results or error. + """Update a simulation record with results or error. Args: country_id (str): The country ID. @@ -177,8 +177,8 @@ def update_simulation(country_id: str) -> Response: - output (dict): The result output (for complete status) - api_version (str): The API version of the simulation - error_message (str): The error message (for error status) - """ + """ payload = request.json if payload is None: raise BadRequest("Payload missing from request") diff --git a/policyengine_api/routes/system_routes.py b/policyengine_api/routes/system_routes.py index 450459083..67b7959d1 100644 --- a/policyengine_api/routes/system_routes.py +++ b/policyengine_api/routes/system_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for the liveness and readiness probes and the legacy OpenAPI spec.""" + from flask import Blueprint, Response, jsonify from policyengine_api.readiness import is_ready @@ -8,11 +10,13 @@ @system_bp.route("/liveness-check", methods=["GET"]) def liveness_check() -> Response: + """Return 200 "OK" whenever the process can serve a request, even during warmup.""" return Response("OK", status=200, headers={"Content-Type": "text/plain"}) @system_bp.route("/readiness-check", methods=["GET"]) def readiness_check() -> Response: + """Return 200 "OK" once startup warmup has finished, and 503 "NOT READY" before.""" # The service is not ready until startup warmup compiles the simulation # machinery. Liveness remains unconditional so the worker is not restarted. if not is_ready(): @@ -26,4 +30,5 @@ def readiness_check() -> Response: @system_bp.route("/specification", methods=["GET"]) def get_specification() -> Response: + """Return the legacy public OpenAPI specification as JSON.""" return jsonify(OPENAPI_SPECIFICATION) diff --git a/policyengine_api/routes/user_profile_routes.py b/policyengine_api/routes/user_profile_routes.py index 8e36f9707..4ffcccbad 100644 --- a/policyengine_api/routes/user_profile_routes.py +++ b/policyengine_api/routes/user_profile_routes.py @@ -1,3 +1,5 @@ +"""Flask routes for creating, reading and updating v1 user profiles.""" + import json from flask import Blueprint, Response, request @@ -30,10 +32,11 @@ def _serialize_user_profile( @user_profile_bp.route("//user-profile", methods=["POST"]) @validate_country def set_user_profile(country_id: str) -> Response: - """ - Creates a new user_profile - """ + """Create a user profile for an Auth0 ID, or return the one it already has. + Returns 201 for a new profile and 200 for an existing one; the Auth0 ID is + never echoed back. + """ payload = request.json if payload is None: raise BadRequest("Payload missing from request") @@ -65,6 +68,11 @@ def set_user_profile(country_id: str) -> Response: @user_profile_bp.route("//user-profile", methods=["GET"]) @validate_country def get_user_profile(country_id: str) -> Response: + """Return a user profile looked up by ``auth0_id`` or, failing that, ``user_id``. + + The Auth0 ID is included only when the caller looked the profile up by it. + Neither identifier gives 400 and an unknown user gives 404. + """ auth0_id = request.args.get("auth0_id") user_id = request.args.get("user_id") @@ -104,12 +112,11 @@ def get_user_profile(country_id: str) -> Response: @user_profile_bp.route("//user-profile", methods=["PUT"]) @validate_country def update_user_profile(country_id: str) -> Response: - """ - Update any part of a user_profile, given a user_id, - except the auth0_id value; any attempt to edit this - will assume malicious intent and 403 - """ + """Update the username, primary country or start date of profile ``user_id``. + The Auth0 ID cannot be changed: an ``auth0_id`` in the payload is ignored. + A null ``user_id`` gives 400 and an unknown one gives 404. + """ payload = request.json if payload is None: diff --git a/policyengine_api/runtime_cache/claims.py b/policyengine_api/runtime_cache/claims.py index 3fcdc58ab..7ce7008f0 100644 --- a/policyengine_api/runtime_cache/claims.py +++ b/policyengine_api/runtime_cache/claims.py @@ -19,16 +19,24 @@ class ExpiringClaimStore: + """Token-owned keys that let one worker claim shared work until a TTL lapses.""" + def __init__( self, client: CacheBackend, *, family: str = "coordination", ) -> None: + """Store claims in ``client``, reporting events under the ``family`` label.""" self.client = client self.family = family def acquire(self, key: str, token: str, *, ttl_seconds: int) -> bool: + """Claim ``key`` for ``token`` if it is free, returning whether it succeeded. + + Raise CacheCoordinationError when the backend fails, so callers never + proceed as if they owned the work. + """ if not token or ttl_seconds <= 0: raise ValueError("claim token and positive TTL are required") started_at = time.perf_counter() @@ -54,6 +62,11 @@ def acquire(self, key: str, token: str, *, ttl_seconds: int) -> bool: return acquired def release(self, key: str, token: str) -> bool: + """Delete the claim only if ``token`` still owns it; return whether it did. + + A claim that expired or passed to another owner is left alone. Raise + CacheCoordinationError when the backend fails. + """ started_at = time.perf_counter() try: released = bool(self.client.eval(COMPARE_AND_DELETE, 1, key, token)) diff --git a/policyengine_api/runtime_cache/client.py b/policyengine_api/runtime_cache/client.py index e6bbf0dae..13ad1cc51 100644 --- a/policyengine_api/runtime_cache/client.py +++ b/policyengine_api/runtime_cache/client.py @@ -18,6 +18,10 @@ def build_runtime_cache_client(settings: RuntimeCacheSettings) -> redis.Redis: + """Build a Redis client with bounded timeouts and no retry on timeout. + + Raise RuntimeError when the cache is disabled or TLS has no CA bundle. + """ if not settings.enabled or settings.url is None: raise RuntimeError("the shared runtime cache is disabled") kwargs: dict[str, Any] = { @@ -40,7 +44,6 @@ def get_runtime_cache_client( settings: RuntimeCacheSettings | None = None, ) -> redis.Redis: """Return one client per process without connecting during import.""" - process_id = os.getpid() if process_id not in _clients: with _client_lock: @@ -52,6 +55,7 @@ def get_runtime_cache_client( def close_runtime_cache_clients() -> None: + """Close and forget every cached client; the ASGI app calls this at shutdown.""" for client in _clients.values(): client.close() _clients.clear() diff --git a/policyengine_api/runtime_cache/core.py b/policyengine_api/runtime_cache/core.py index 7d50bdfda..f6bc712ea 100644 --- a/policyengine_api/runtime_cache/core.py +++ b/policyengine_api/runtime_cache/core.py @@ -21,23 +21,65 @@ class CachePipeline(Protocol): """The queued commands that cache writers issue inside one pipeline.""" - def __enter__(self) -> Self: ... - def __exit__(self, *args: object) -> object: ... - def set(self, key: str, value: str, **kwargs: object) -> object: ... - def delete(self, *keys: str) -> object: ... - def expire(self, key: str, seconds: int) -> object: ... - def zadd(self, key: str, mapping: dict[str, float]) -> object: ... - def zrem(self, key: str, *members: str) -> object: ... - def zremrangebyrank(self, key: str, start: int, end: int) -> object: ... - def execute(self) -> list[object]: ... + def __enter__(self) -> Self: + """Open the pipeline for queuing commands.""" + ... + + def __exit__(self, *args: object) -> object: + """Discard any commands that were queued but not executed.""" + ... + + def set(self, key: str, value: str, **kwargs: object) -> object: + """Queue a Redis SET; ``kwargs`` carries options such as ``ex``.""" + ... + + def delete(self, *keys: str) -> object: + """Queue a Redis DEL of every given key.""" + ... + + def expire(self, key: str, seconds: int) -> object: + """Queue a Redis EXPIRE that sets the key's TTL in seconds.""" + ... + + def zadd(self, key: str, mapping: dict[str, float]) -> object: + """Queue a Redis ZADD of member-to-score pairs into a sorted set.""" + ... + + def zrem(self, key: str, *members: str) -> object: + """Queue a Redis ZREM of the given sorted-set members.""" + ... + + def zremrangebyrank(self, key: str, start: int, end: int) -> object: + """Queue a Redis ZREMRANGEBYRANK over inclusive ranks (``-1`` is last).""" + ... + + def execute(self) -> list[object]: + """Run the queued commands and return their replies in queue order.""" + ... class CacheBackend(Protocol): - def get(self, key: str) -> str | bytes | None: ... - def set(self, key: str, value: str, **kwargs: object) -> object: ... - def delete(self, *keys: str) -> object: ... - def pipeline(self, transaction: bool = True) -> CachePipeline: ... - def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: ... + """The Redis client subset the runtime cache uses; fakes implement it too.""" + + def get(self, key: str) -> str | bytes | None: + """Return the stored value, or None when the key is missing or expired.""" + ... + + def set(self, key: str, value: str, **kwargs: object) -> object: + """Store a value; ``kwargs`` carries Redis options such as ``ex`` and ``nx``.""" + ... + + def delete(self, *keys: str) -> object: + """Delete every given key and return how many existed.""" + ... + + def pipeline(self, transaction: bool = True) -> CachePipeline: + """Return a pipeline whose commands run together on ``execute``.""" + ... + + def eval(self, script: str, numkeys: int, *keys_and_args: str) -> object: + """Run a Lua script atomically on the server against the given keys.""" + ... class CacheCoordinationError(RuntimeError): @@ -46,6 +88,8 @@ class CacheCoordinationError(RuntimeError): @dataclass(frozen=True) class CacheNamespace: + """The environment and service prefix that keeps deployments' keys apart.""" + environment: str service: str @@ -57,6 +101,11 @@ def key( *, suffix: str | None = None, ) -> str: + """Build a key from a SHA-256 of the canonical JSON of ``inputs``. + + Equal inputs give the same key whatever their dict order; values that + are not JSON-serialisable are hashed by their ``str()``. + """ encoded = json.dumps( inputs, sort_keys=True, @@ -71,6 +120,7 @@ def key( return f"{key}:{suffix}" if suffix else key def family_key(self, family: str, schema_version: int, name: str) -> str: + """Build an unhashed key for a fixed, named entry within a cache family.""" return ( f"{CACHE_KEY_ROOT}:{self.environment}:{self.service}:" f"{family}:v{schema_version}:{name}" @@ -82,6 +132,7 @@ def encode_envelope( schema_version: int, payload: object, ) -> str: + """Serialise a payload with the family and schema version it was written under.""" return json.dumps( { "family": family, @@ -100,6 +151,11 @@ def decode_envelope( family: str, schema_version: int, ) -> object: + """Return an envelope's payload, or None if it is absent or unusable. + + Undecodable JSON and envelopes from another family or schema version are + treated as misses rather than errors. + """ if value is None: return None try: @@ -127,7 +183,6 @@ def jittered_ttl( Coordination and claim TTLs are safety bounds and must not use this helper. """ - if ttl_seconds <= 0: raise ValueError("cache TTL must be positive") if not 0 <= jitter_fraction < 1: @@ -148,7 +203,6 @@ def record_cache_event( operation: str | None = None, ) -> None: """Emit one metric-ready event without accepting cache keys or values.""" - payload: dict[str, str | int | float] = { "message": "Runtime cache operation", "metric_name": "runtime_cache_operations", @@ -174,6 +228,7 @@ def __init__( schema_version: int, ttl_seconds: int, ) -> None: + """Bind the cache to one key family; raise ValueError for a non-positive TTL.""" if ttl_seconds <= 0: raise ValueError("cache TTL must be positive") self.client = client @@ -183,9 +238,11 @@ def __init__( self.ttl_seconds = ttl_seconds def key(self, inputs: dict[str, Any]) -> str: + """Return the namespaced key this cache stores ``inputs`` under.""" return self.namespace.key(self.family, self.schema_version, inputs) def get(self, inputs: dict[str, Any]) -> object: + """Return the cached payload, or None on a miss or any backend failure.""" started_at = time.perf_counter() try: value = self.client.get(self.key(inputs)) @@ -210,6 +267,7 @@ def get(self, inputs: dict[str, Any]) -> object: return payload def set(self, inputs: dict[str, Any], payload: object) -> bool: + """Store a payload with a jittered TTL; return False if the write failed.""" started_at = time.perf_counter() try: result = self.client.set( diff --git a/policyengine_api/runtime_cache/dependencies.py b/policyengine_api/runtime_cache/dependencies.py index c5cf2a97d..17324ef22 100644 --- a/policyengine_api/runtime_cache/dependencies.py +++ b/policyengine_api/runtime_cache/dependencies.py @@ -13,6 +13,8 @@ @dataclass(frozen=True) class RuntimeCacheContext: + """The process's cache backend and key namespace, built once from settings.""" + client: CacheBackend namespace: CacheNamespace enabled: bool @@ -20,6 +22,10 @@ class RuntimeCacheContext: @lru_cache(maxsize=1) def get_runtime_cache_context() -> RuntimeCacheContext: + """Return the shared cache context, a no-storage backend when disabled. + + Built on first use, not at import, and memoised for the process. + """ settings = load_runtime_cache_settings() namespace = CacheNamespace(settings.environment, settings.service) client: CacheBackend = ( @@ -35,4 +41,5 @@ def get_runtime_cache_context() -> RuntimeCacheContext: def clear_runtime_cache_context() -> None: + """Forget the memoised context so the next call rereads the settings.""" get_runtime_cache_context.cache_clear() diff --git a/policyengine_api/runtime_cache/fake.py b/policyengine_api/runtime_cache/fake.py index 74fe06318..89b83a0c7 100644 --- a/policyengine_api/runtime_cache/fake.py +++ b/policyengine_api/runtime_cache/fake.py @@ -13,6 +13,7 @@ class InMemoryCacheBackend: """Small Redis-like fake with explicit time advancement and atomic pipelines.""" def __init__(self, *, now: Callable[[], float] | None = None) -> None: + """Start empty at time zero, or follow the ``now`` clock when one is given.""" self._external_now = now self._time = 0.0 self._values: dict[str, str] = {} @@ -24,6 +25,10 @@ def _now(self) -> float: return self._external_now() if self._external_now else self._time def advance(self, seconds: float) -> None: + """Move the fake clock forward so TTLs can expire without sleeping. + + Raises RuntimeError when the fake follows an external clock instead. + """ if self._external_now is not None: raise RuntimeError("cannot advance an externally clocked fake") self._time += seconds @@ -36,6 +41,7 @@ def _purge(self, key: str) -> None: self._expires.pop(key, None) def get(self, key: str) -> str | None: + """Emulate GET: return the string value, or None if missing or expired.""" with self._lock: self._purge(key) return self._values.get(key) @@ -48,6 +54,11 @@ def set( ex: int | None = None, nx: bool = False, ) -> bool | None: + """Emulate SET with the ``ex`` and ``nx`` options only. + + Like redis-py, return True on a write and None when ``nx`` finds the key + already set; a write without ``ex`` clears any earlier TTL. + """ with self._lock: self._purge(key) if nx and key in self._values: @@ -60,6 +71,7 @@ def set( return True def delete(self, *keys: str) -> int: + """Emulate DEL for strings and sorted sets, returning how many existed.""" with self._lock: deleted = 0 for key in keys: @@ -72,6 +84,7 @@ def delete(self, *keys: str) -> int: return deleted def expire(self, key: str, seconds: int) -> bool: + """Emulate EXPIRE: set a TTL and return False if the key does not exist.""" with self._lock: self._purge(key) if key not in self._values and key not in self._sorted_sets: @@ -80,6 +93,7 @@ def expire(self, key: str, seconds: int) -> bool: return True def zadd(self, key: str, mapping: dict[str, float]) -> int: + """Emulate ZADD without flags: upsert scores and count new members only.""" with self._lock: self._purge(key) values = self._sorted_sets.setdefault(key, {}) @@ -88,6 +102,11 @@ def zadd(self, key: str, mapping: dict[str, float]) -> int: return added def zrevrange(self, key: str, start: int, end: int) -> list[str]: + """Emulate ZREVRANGE: members by descending score over inclusive ranks. + + Equal scores fall back to reverse member order, as in Redis. The only + negative ``end`` supported is ``-1`` (the last member). + """ with self._lock: self._purge(key) ordered = sorted( @@ -99,6 +118,7 @@ def zrevrange(self, key: str, start: int, end: int) -> list[str]: return [member for member, _ in ordered[start:stop]] def zrem(self, key: str, *members: str) -> int: + """Emulate ZREM, returning how many of the members were present.""" with self._lock: values = self._sorted_sets.get(key, {}) removed = sum(member in values for member in members) @@ -107,6 +127,10 @@ def zrem(self, key: str, *members: str) -> int: return removed def zremrangebyrank(self, key: str, start: int, end: int) -> int: + """Emulate ZREMRANGEBYRANK over ascending-score ranks, inclusive of ``end``. + + As with ``zrevrange``, the only negative ``end`` supported is ``-1``. + """ with self._lock: self._purge(key) ordered = sorted( @@ -123,6 +147,11 @@ def eval( numkeys: int, *keys_and_args: str, ) -> int: + """Emulate EVAL of the claim store's compare-and-delete script only. + + The script text is ignored: delete the one key if it still holds the + expected token, and raise ValueError for any other key/argument shape. + """ if numkeys != 1 or len(keys_and_args) != 2: raise ValueError("fake supports one-key compare-and-delete only") key, expected = keys_and_args @@ -132,16 +161,22 @@ def eval( return self.delete(key) def pipeline(self, transaction: bool = True) -> InMemoryPipeline: + """Return a pipeline that queues commands until ``execute``.""" return InMemoryPipeline(self, transaction=transaction) class InMemoryPipeline: + """A fake Redis pipeline that records any method call and replays it later.""" + def __init__(self, backend: InMemoryCacheBackend, *, transaction: bool) -> None: + """Queue commands for ``backend``; ``transaction`` makes ``execute`` atomic.""" self.backend = backend self.transaction = transaction self.operations: list[tuple[str, tuple[object, ...], dict[str, object]]] = [] def __getattr__(self, name: str) -> Callable[..., InMemoryPipeline]: + """Return a recorder that queues a call to the backend method ``name``.""" + def queue(*args: object, **kwargs: object) -> InMemoryPipeline: self.operations.append((name, args, kwargs)) return self @@ -149,6 +184,11 @@ def queue(*args: object, **kwargs: object) -> InMemoryPipeline: return queue def execute(self) -> list[object]: + """Replay the queued commands on the backend, then clear the queue. + + With ``transaction`` the replay holds the backend lock, as MULTI/EXEC + would; without it, other callers' commands may interleave. + """ lock = self.backend._lock if self.transaction else RLock() with lock: results = [ @@ -159,9 +199,11 @@ def execute(self) -> list[object]: return results def __enter__(self) -> Self: + """Return the pipeline itself for use in a ``with`` block.""" return self def __exit__(self, *_args: object) -> None: + """Drop queued commands that were never executed.""" self.operations.clear() @@ -169,27 +211,34 @@ class DisabledCacheBackend(InMemoryCacheBackend): """No-storage backend for unselected unit-test/application cache mode.""" def get(self, key: str) -> None: - return None + """Return None: every read is a miss.""" def set(self, key: str, value: str, **kwargs: object) -> bool: + """Discard the value and report the write as not stored.""" return False def delete(self, *keys: str) -> int: + """Return 0: no key ever exists.""" return 0 def zadd(self, key: str, mapping: dict[str, float]) -> int: + """Discard the members and report none added.""" return 0 def zrevrange(self, key: str, start: int, end: int) -> list[str]: + """Return an empty list: no sorted set ever exists.""" return [] def zrem(self, key: str, *members: str) -> int: + """Return 0: there is nothing to remove.""" return 0 def zremrangebyrank(self, key: str, start: int, end: int) -> int: + """Return 0: there is nothing to remove.""" return 0 def expire(self, key: str, seconds: int) -> bool: + """Return False, as Redis does for a key that does not exist.""" return False def eval( @@ -198,4 +247,5 @@ def eval( numkeys: int, *keys_and_args: str, ) -> int: + """Raise RuntimeError, so claim operations fail closed when disabled.""" raise RuntimeError("the shared runtime cache is disabled") diff --git a/policyengine_api/runtime_cache/household_calculations.py b/policyengine_api/runtime_cache/household_calculations.py index f490c3bd2..1082054ef 100644 --- a/policyengine_api/runtime_cache/household_calculations.py +++ b/policyengine_api/runtime_cache/household_calculations.py @@ -21,6 +21,13 @@ @dataclass(frozen=True) class HouseholdCalculationIdentity: + """Every input that determines a calculated household; the cache key hashes it. + + The content hashes and package versions make edited inputs and upgraded + models miss rather than reuse stale results. ``spm`` is the requested + resolved SPM configuration, or None when the request selected no SPM. + """ + country_id: str household_id: int policy_id: int @@ -33,6 +40,8 @@ class HouseholdCalculationIdentity: @dataclass(frozen=True) class CachedHouseholdCalculation: + """A calculated household, its warnings, and the SPM receipt if SPM was used.""" + household: dict[str, Any] warnings: tuple[str, ...] = () spm_provenance: dict[str, Any] | None = None @@ -42,6 +51,7 @@ class HouseholdCalculationCache: """Cache a calculated household, its warnings and its receipts atomically.""" def __init__(self, client: CacheBackend, namespace: CacheNamespace) -> None: + """Store results in ``client`` under ``namespace`` for about a day.""" self._cache = RecoverableJSONCache( client, namespace, @@ -51,12 +61,18 @@ def __init__(self, client: CacheBackend, namespace: CacheNamespace) -> None: ) def cache_key(self, identity: HouseholdCalculationIdentity) -> str: + """Return the backend key a result for ``identity`` is stored under.""" return self._cache.key(asdict(identity)) def get( self, identity: HouseholdCalculationIdentity, ) -> CachedHouseholdCalculation | None: + """Return the cached result, or None on a miss or any malformed entry. + + An entry whose SPM receipt is missing, invalid, or resolves to a + different SPM configuration than ``identity`` asks for is also a miss. + """ payload = self._cache.get(asdict(identity)) if not isinstance(payload, dict): return None @@ -96,6 +112,11 @@ def set( identity: HouseholdCalculationIdentity, value: CachedHouseholdCalculation, ) -> bool: + """Cache a result, returning False if it was not written. + + Results whose SPM receipt does not match the identity's SPM request are + refused rather than stored, so a later read cannot serve them. + """ if identity.spm is None: if value.spm_provenance is not None: return False diff --git a/policyengine_api/runtime_cache/reform_impacts.py b/policyengine_api/runtime_cache/reform_impacts.py index e21bdfe2b..32d7a867d 100644 --- a/policyengine_api/runtime_cache/reform_impacts.py +++ b/policyengine_api/runtime_cache/reform_impacts.py @@ -34,6 +34,7 @@ class ReformImpactStartClaim: observability_id: str def to_payload(self) -> dict[str, str]: + """Return the JSON object stored as the claim's value.""" return { "submission_claim_id": self.submission_claim_id, "observability_id": self.observability_id, @@ -41,6 +42,7 @@ def to_payload(self) -> dict[str, str]: @classmethod def from_payload(cls, payload: object) -> "ReformImpactStartClaim | None": + """Rebuild a claim, or return None if a field is missing or invalid.""" if not isinstance(payload, dict): return None submission_claim_id = payload.get("submission_claim_id") @@ -57,6 +59,11 @@ def from_payload(cls, payload: object) -> "ReformImpactStartClaim | None": @dataclass(frozen=True) class CachedReformImpact: + """An economy-impact job and its result, shaped like a legacy reform_impact row. + + ``status`` is ``computing``, ``ok`` or ``error``; naive datetimes are UTC. + """ + reform_impact_id: int baseline_policy_id: int reform_policy_id: int @@ -162,6 +169,7 @@ class ReformImpactCache: family = "reform-impact" def __init__(self, client: CacheBackend, namespace: CacheNamespace) -> None: + """Keep values, indexes and start claims in ``client`` under ``namespace``.""" self.client = client self.namespace = namespace self._start_claims = ExpiringClaimStore(client, family=self.family) @@ -219,7 +227,6 @@ def claim_start( observability_id: str, ) -> bool: """Atomically claim ownership of one reform-impact submission.""" - claim = ReformImpactStartClaim.from_payload( { "submission_claim_id": claim_token, @@ -260,7 +267,6 @@ def get_start_claim( target: str, ) -> ReformImpactStartClaim | None: """Read the current submission owner and its diagnostic identifier.""" - started_at = time.perf_counter() key = self._start_claim_key( country_id=country_id, @@ -329,7 +335,6 @@ def release_start( observability_id: str, ) -> bool: """Release a start claim only when its complete state still matches.""" - claim = ReformImpactStartClaim.from_payload( { "submission_claim_id": claim_token, @@ -400,6 +405,11 @@ def set( *, record_execution_id: str | None = None, ) -> bool: + """Store an impact and index it by scope and recency in one transaction. + + Each index keeps only the newest 1,000 entries. Return False if the + write failed; raise ValueError when there is no record identifier. + """ # Terminal errors clear the worker handle in the payload while keeping # the original record key and lookup indexes stable. record_execution_id = record_execution_id or impact.execution_id @@ -439,6 +449,7 @@ def set( return True def get_by_execution_id(self, execution_id: str) -> CachedReformImpact | None: + """Return the job's impact, or None if absent, unreadable or unreachable.""" try: value = self.client.get(self._record_key(execution_id)) except Exception: @@ -470,6 +481,7 @@ def _from_index(self, index: str, limit: int) -> list[CachedReformImpact]: return [impact for impact in impacts if impact is not None] def recent(self, limit: int) -> list[CachedReformImpact]: + """Return up to ``limit`` impacts, most recently started first.""" return self._from_index(self._recent_index(), limit) def matching( @@ -485,6 +497,12 @@ def matching( options_hash: str, options_hash_pattern: str | None = None, ) -> list[CachedReformImpact]: + """Return the cached impacts for one policy pair and scope, best match first. + + Exact ``options_hash`` matches sort before those matching + ``options_hash_pattern``, a SQL ``LIKE`` pattern; ties go to the most + recently started. Only the scope index's newest 1,000 entries are read. + """ probe = CachedReformImpact( reform_impact_id=0, baseline_policy_id=baseline_policy_id, @@ -534,6 +552,11 @@ def update( clear_execution_id: bool = False, **changes: Any, # noqa: ANN401 # forwarded unchanged to dataclasses.replace ) -> CachedReformImpact | None: + """Apply field changes to a cached impact and rewrite it under the same key. + + ``clear_execution_id`` drops the worker handle without moving the record. + Return the updated impact, or None if it is missing or the write failed. + """ impact = self.get_by_execution_id(execution_id) if impact is None: return None @@ -553,6 +576,10 @@ def delete_matching_computing( time_period: str, options_hash: str, ) -> None: + """Remove in-flight (``computing``) impacts for this scope and options. + + Cache failures are ignored, so a stale entry may survive until it expires. + """ # Deletion historically omits api_version, so search the bounded recent # index and remove only matching in-flight values. impacts = self.recent(REFORM_IMPACT_INDEX_LIMIT) @@ -584,6 +611,5 @@ def delete_matching_computing( def reform_impact_id(execution_id: str) -> int: """Stable positive cache-local identifier for the historical response field.""" - digest = hashlib.sha256(execution_id.encode("utf-8")).hexdigest()[:15] return int(digest, 16) diff --git a/policyengine_api/runtime_cache/settings.py b/policyengine_api/runtime_cache/settings.py index 6fb2937da..958b0ac43 100644 --- a/policyengine_api/runtime_cache/settings.py +++ b/policyengine_api/runtime_cache/settings.py @@ -33,6 +33,8 @@ class RuntimeCacheConfigurationError(RuntimeError): @dataclass(frozen=True) class RuntimeCacheSettings: + """Validated cache configuration; the URL and CA bundle stay secret in reprs.""" + mode: str environment: str service: str @@ -45,6 +47,7 @@ class RuntimeCacheSettings: @property def enabled(self) -> bool: + """Whether a real Redis backend is configured (local or deployed mode).""" return self.mode != "disabled" @@ -129,7 +132,6 @@ def load_runtime_cache_settings( environ: Mapping[str, str] | None = None, ) -> RuntimeCacheSettings: """Load disabled, explicit-local, or fail-closed deployed cache settings.""" - values = os.environ if environ is None else environ raw_mode = values.get(RUNTIME_CACHE_MODE, "").strip() deployed_runtime = _is_deployed(values) diff --git a/policyengine_api/services/__init__.py b/policyengine_api/services/__init__.py index ca7322caf..5bd08f8f8 100644 --- a/policyengine_api/services/__init__.py +++ b/policyengine_api/services/__init__.py @@ -1 +1,3 @@ +"""Business logic for the v1 API resources; ``v2`` holds the v2 resource services.""" + # from .economy_service import EconomyService, JobService, PolicyService diff --git a/policyengine_api/services/budget_window_cache.py b/policyengine_api/services/budget_window_cache.py index b6190341a..edff6530b 100644 --- a/policyengine_api/services/budget_window_cache.py +++ b/policyengine_api/services/budget_window_cache.py @@ -50,10 +50,12 @@ class BudgetWindowCacheState(BaseModel): error: dict[str, Any] | None = None def to_payload(self) -> dict[str, Any]: + """Return the state as the cache JSON payload, omitting unset fields.""" return self.model_dump(exclude_none=True) @model_validator(mode="after") def require_fields_for_status(self) -> BudgetWindowCacheState: + """Raise ValueError unless every field this status depends on is set.""" missing = [ field for field in self._required_fields[self.status] @@ -65,6 +67,7 @@ def require_fields_for_status(self) -> BudgetWindowCacheState: @classmethod def from_payload(cls, payload: object) -> BudgetWindowCacheState | None: + """Parse a decoded cache payload, or return None if it is not a valid state.""" try: return cls.model_validate(payload) except ValidationError: @@ -79,6 +82,7 @@ def __init__( client: CacheBackend | None = None, namespace: CacheNamespace | None = None, ) -> None: + """Bind to a cache client and namespace, defaulting to the runtime cache.""" if client is None or namespace is None: context = get_runtime_cache_context() client = client or context.client @@ -102,6 +106,11 @@ def build_key( options_hash: str | None, api_version: str, ) -> str: + """Return the cache key for one budget-window report request. + + Every argument is part of the key, so a change to any of them (including the + API version or the options hash) starts a separate report. + """ return self.namespace.key( BUDGET_WINDOW_CACHE_FAMILY, BUDGET_WINDOW_CACHE_SCHEMA_VERSION, @@ -146,7 +155,6 @@ def _handle_cache_error( def get_state(self, cache_key: str) -> BudgetWindowCacheState | None: """Read the complete report state or fail closed on cache outage.""" - started_at = time.perf_counter() state_key = self._state_key(cache_key) try: @@ -203,6 +211,12 @@ def claim_batch_start( claim_token: str, observability_id: str | None, ) -> bool: + """Atomically claim the right to submit this report's batch job. + + Return True only if no state existed for the key, in which case a short-lived + "starting" state now holds the claim. Raise CacheCoordinationError if the cache + is unreachable, so callers never submit duplicate work blindly. + """ state = BudgetWindowCacheState( status="starting", observability_id=observability_id, @@ -220,6 +234,11 @@ def clear_starting_claim( claim_token: str, observability_id: str | None, ) -> None: + """Release this caller's "starting" claim so another request can submit. + + The claim is deleted only if it still holds this token; cache errors are + swallowed because the claim expires on its own. + """ state = BudgetWindowCacheState( status="starting", observability_id=observability_id, @@ -239,6 +258,11 @@ def store_submitted( batch_job_id: str, observability_id: str | None, ) -> None: + """Replace the starting claim with the submitted batch job ID. + + Raise CacheCoordinationError if the state cannot be written, since later polls + would otherwise lose track of the running batch. + """ state = BudgetWindowCacheState( status="submitted", observability_id=observability_id, @@ -282,6 +306,7 @@ def set_completed_result( result: dict[str, Any], observability_id: str | None, ) -> bool: + """Store the completed report result; return False if the cache write failed.""" return self._set_recoverable_state( cache_key, BudgetWindowCacheState( @@ -298,6 +323,10 @@ def set_terminal_error( error: dict[str, str], observability_id: str | None, ) -> bool: + """Store an SPM validation failure so later requests replay it, not resubmit. + + Return False if the cache write failed. + """ return self._set_recoverable_state( cache_key, BudgetWindowCacheState( @@ -315,6 +344,10 @@ def set_execution_failure( result: dict[str, Any], observability_id: str | None, ) -> bool: + """Store a failed batch result so later requests return it without resubmitting. + + Return False if the cache write failed. + """ return self._set_recoverable_state( cache_key, BudgetWindowCacheState( diff --git a/policyengine_api/services/economy_comparison.py b/policyengine_api/services/economy_comparison.py index 08de03777..4e8c5f573 100644 --- a/policyengine_api/services/economy_comparison.py +++ b/policyengine_api/services/economy_comparison.py @@ -13,6 +13,11 @@ def budgetary_impact(baseline: dict, reform: dict) -> dict: + """Return the reform's change in tax revenue, benefit spending and net budget. + + A positive ``budgetary_impact`` means the reform raises net revenue (tax + revenue change minus benefit spending change). + """ tax_revenue_impact = reform["total_tax"] - baseline["total_tax"] state_tax_revenue_impact = reform["total_state_tax"] - baseline["total_state_tax"] benefit_spending_impact = reform["total_benefits"] - baseline["total_benefits"] @@ -28,6 +33,11 @@ def budgetary_impact(baseline: dict, reform: dict) -> dict: def labor_supply_response(baseline: dict, reform: dict) -> dict: + """Return the reform's labour supply responses relative to the baseline. + + This covers the income and substitution effects in total, as shares of + original earnings, by income decile, and in weekly hours. + """ substitution_lsr = reform["substitution_lsr"] - baseline["substitution_lsr"] income_lsr = reform["income_lsr"] - baseline["income_lsr"] total_change = substitution_lsr + income_lsr @@ -103,6 +113,10 @@ def labor_supply_response(baseline: dict, reform: dict) -> dict: def detailed_budgetary_impact(baseline: dict, reform: dict, country_id: str) -> dict: + """Return baseline, reform and difference for each UK programme's budget. + + Only the UK reports programme-level totals; other countries get an empty dict. + """ result = {} if country_id == "uk": for program in baseline["programs"]: @@ -116,8 +130,7 @@ def detailed_budgetary_impact(baseline: dict, reform: dict, country_id: str) -> def decile_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on the deciles of the population. + """Compare the impact of a reform on the deciles of the population. Args: baseline (dict): The baseline economy. @@ -125,6 +138,7 @@ def decile_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on the deciles of the population. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] @@ -157,8 +171,7 @@ def decile_impact(baseline: dict, reform: dict) -> dict: def wealth_decile_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on the deciles of the population. + """Compare the impact of a reform on the deciles of the population. Args: baseline (dict): The baseline economy. @@ -166,6 +179,7 @@ def wealth_decile_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on the deciles of the population. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] @@ -197,8 +211,7 @@ def wealth_decile_impact(baseline: dict, reform: dict) -> dict: def inequality_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on inequality. + """Compare the impact of a reform on inequality. Args: baseline (dict): The baseline economy. @@ -206,8 +219,8 @@ def inequality_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on inequality. - """ + """ return dict( gini=dict( baseline=baseline["gini"], @@ -225,8 +238,7 @@ def inequality_impact(baseline: dict, reform: dict) -> dict: def poverty_impact(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on poverty. + """Compare the impact of a reform on poverty. Args: baseline (dict): The baseline economy. @@ -234,6 +246,7 @@ def poverty_impact(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on poverty. + """ baseline_poverty = MicroSeries( baseline["person_in_poverty"], weights=baseline["person_weight"] @@ -296,14 +309,21 @@ def poverty_impact(baseline: dict, reform: dict) -> dict: def compute_income_change( baseline_values: np.ndarray, reform_values: np.ndarray ) -> np.ndarray: - """Percentage income change with a floor of 1 on the baseline - to avoid division by zero for zero/negative incomes.""" + """Return each household's income change as a fraction of baseline income. + + The baseline is floored at 1 so zero or negative incomes do not divide by zero. + """ absolute_change = reform_values - baseline_values capped_baseline = np.maximum(baseline_values, 1) return absolute_change / capped_baseline def intra_decile_impact(baseline: dict, reform: dict) -> dict: + """Return the share of people in each income decile by size of income change. + + People fall into five buckets (gain or lose more or less than 5%, or no change + within 0.1%). ``all`` is the unweighted mean of the ten decile shares. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] ) @@ -356,6 +376,11 @@ def intra_decile_impact(baseline: dict, reform: dict) -> dict: def intra_wealth_decile_impact(baseline: dict, reform: dict) -> dict: + """Return the share of people in each wealth decile by size of income change. + + This uses the same buckets as ``intra_decile_impact``, grouped by household + wealth decile instead of income decile. + """ baseline_income = MicroSeries( baseline["household_net_income"], weights=baseline["household_weight"] ) @@ -408,8 +433,7 @@ def intra_wealth_decile_impact(baseline: dict, reform: dict) -> dict: def poverty_gender_breakdown(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on poverty. + """Compare the impact of a reform on poverty. Args: baseline (dict): The baseline economy. @@ -417,6 +441,7 @@ def poverty_gender_breakdown(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on poverty. + """ if baseline["is_male"] is None: return {} @@ -463,8 +488,7 @@ def poverty_gender_breakdown(baseline: dict, reform: dict) -> dict: def poverty_racial_breakdown(baseline: dict, reform: dict) -> dict: - """ - Compare the impact of a reform on poverty. + """Compare the impact of a reform on poverty. Args: baseline (dict): The baseline economy. @@ -472,6 +496,7 @@ def poverty_racial_breakdown(baseline: dict, reform: dict) -> dict: Returns: dict: The impact of the reform on poverty. + """ if baseline["race"] is None: return {} @@ -508,6 +533,8 @@ def poverty_racial_breakdown(baseline: dict, reform: dict) -> dict: class UKConstituencyBreakdownByConstituency(BaseModel): + """One constituency's household income change and its map position (x, y).""" + average_household_income_change: float relative_household_income_change: float x: int @@ -515,11 +542,19 @@ class UKConstituencyBreakdownByConstituency(BaseModel): class UKConstituencyBreakdown(BaseModel): + """UK reform results by constituency. + + ``outcomes_by_region`` counts how many constituencies in the UK and each nation + fall into each income-change bucket. + """ + by_constituency: dict[str, UKConstituencyBreakdownByConstituency] outcomes_by_region: dict[str, dict[str, int]] class UKLocalAuthorityBreakdownByLA(BaseModel): + """One local authority's household income change and its map position (x, y).""" + average_household_income_change: float relative_household_income_change: float x: int @@ -527,6 +562,12 @@ class UKLocalAuthorityBreakdownByLA(BaseModel): class UKLocalAuthorityBreakdown(BaseModel): + """UK reform results by local authority. + + ``outcomes_by_region`` counts how many local authorities in the UK and each + nation fall into each income-change bucket. + """ + by_local_authority: dict[str, UKLocalAuthorityBreakdownByLA] outcomes_by_region: dict[str, dict[str, int]] @@ -534,6 +575,12 @@ class UKLocalAuthorityBreakdown(BaseModel): def uk_constituency_breakdown( baseline: dict, reform: dict, country_id: str, region: str | None = None ) -> UKConstituencyBreakdown | None: + """Return each UK constituency's average and relative household income change. + + Return None outside the UK or for a local-authority region. A constituency or + country region limits the output to that area. Constituency weights and + names come from PolicyEngine's Hugging Face datasets. + """ if country_id != "uk": return None @@ -663,6 +710,12 @@ def uk_constituency_breakdown( def uk_local_authority_breakdown( baseline: dict, reform: dict, country_id: str, region: str | None = None ) -> UKLocalAuthorityBreakdown | None: + """Return each UK local authority's average and relative household income change. + + Return None outside the UK or for a constituency region. A local-authority or + country region limits the output to that area. Local-authority weights and + names come from PolicyEngine's Hugging Face datasets. + """ if country_id != "uk": return None @@ -785,8 +838,7 @@ def compare_economic_outputs( country_id: str = None, region: str | None = None, ) -> dict: - """ - Compare the economic outputs of two economies. + """Compare the economic outputs of two economies. Args: baseline (dict): The baseline economy. @@ -797,6 +849,7 @@ def compare_economic_outputs( Returns: dict: The comparison of the two economies. + """ if baseline.get("type") == "general": budgetary_impact_data = budgetary_impact(baseline, reform) diff --git a/policyengine_api/services/economy_service.py b/policyengine_api/services/economy_service.py index 6bac954d6..cdf82dfe2 100644 --- a/policyengine_api/services/economy_service.py +++ b/policyengine_api/services/economy_service.py @@ -1,3 +1,5 @@ +"""Economic impact service backed by cached reform impacts and simulation jobs.""" + import datetime import hashlib import json @@ -72,9 +74,7 @@ class ImpactAction(Enum): - """ - Enum for the action to take based on the status of an economic impact calculation. - """ + """What to do with an economic impact request, given its stored calculation.""" COMPLETED = "completed" FAILED = "failed" @@ -83,9 +83,7 @@ class ImpactAction(Enum): class ImpactStatus(Enum): - """ - Enum for the status of an economic impact calculation. - """ + """Enum for the status of an economic impact calculation.""" COMPUTING = "computing" OK = "ok" @@ -103,6 +101,8 @@ class ImpactStatus(Enum): class SimulationOptions(BaseModel): + """The job payload sent to the Simulation entrypoint for one economy simulation.""" + spm: SPMSelection | None = None country: str scope: Literal["macro", "household"] = "macro" @@ -116,6 +116,12 @@ class SimulationOptions(BaseModel): class EconomicImpactSetupOptions(BaseModel): + """One resolved economic impact request. + + It holds the request's cache identity, the runtime bundle versions and the + submission claim token that guards against duplicate simulation jobs. + """ + submission_claim_id: str observability_id: str | None = None country_id: str @@ -135,9 +141,10 @@ class EconomicImpactSetupOptions(BaseModel): class EconomicImpactResult(BaseModel): - """ - Model for the result of an economic impact calculation. - Contains the status and the reform impact JSON. + """The status of an economic impact calculation and, once complete, its result. + + ``data`` holds the reform impact JSON when the status is ok, and ``message`` + explains an error. """ status: ImpactStatus @@ -150,9 +157,7 @@ class EconomicImpactResult(BaseModel): model_config = {"frozen": True} # Make model immutable def to_dict(self) -> dict[str, str | dict | None]: - """ - Convert the EconomicImpactResult to a dictionary. - """ + """Convert the EconomicImpactResult to a dictionary.""" return { "status": self.status.value, "data": self.data, @@ -161,33 +166,25 @@ def to_dict(self) -> dict[str, str | dict | None]: @classmethod def completed(cls, data: dict) -> "EconomicImpactResult": - """ - Create an EconomicImpactResult for a completed impact calculation. - """ + """Create an EconomicImpactResult for a completed impact calculation.""" return cls(status=ImpactStatus.OK, data=data) @classmethod def computing( cls, run_progress: SimulationRunProgress | None = None ) -> "EconomicImpactResult": - """ - Create an EconomicImpactResult for a computing impact calculation. - """ + """Create an EconomicImpactResult for a computing impact calculation.""" return cls(status=ImpactStatus.COMPUTING, data=None, run_progress=run_progress) @classmethod def error(cls, message: str) -> "EconomicImpactResult": - """ - Create an EconomicImpactResult for an error in the impact calculation. - """ + """Create an EconomicImpactResult for an error in the impact calculation.""" logger.log_struct({"message": message}, severity="ERROR") return cls(status=ImpactStatus.ERROR, data=None, message=message) class BudgetWindowEconomicImpactResult(BaseModel): - """ - Model for a batch budget-window economic impact response. - """ + """Model for a batch budget-window economic impact response.""" status: ImpactStatus data: dict | None = None @@ -202,6 +199,7 @@ class BudgetWindowEconomicImpactResult(BaseModel): model_config = {"frozen": True} def to_dict(self) -> dict[str, Any]: + """Return the response body; ``cache_status`` is internal and is left out.""" return { "status": self.status.value, "data": self.data, @@ -217,6 +215,7 @@ def to_dict(self) -> dict[str, Any]: def completed( cls, data: dict, *, cache_status: str | None = None ) -> "BudgetWindowEconomicImpactResult": + """Build a finished result with progress at 100.""" return cls( status=ImpactStatus.OK, data=data, @@ -235,6 +234,7 @@ def computing( message: str, cache_status: str | None = None, ) -> "BudgetWindowEconomicImpactResult": + """Build an in-progress result listing the done, running and queued years.""" return cls( status=ImpactStatus.COMPUTING, data=None, @@ -256,6 +256,7 @@ def failed( queued_years: list[str] | None = None, cache_status: str | None = None, ) -> "BudgetWindowEconomicImpactResult": + """Build an error result, copying the message into ``error`` and logging it.""" logger.log_struct({"message": message}, severity="ERROR") return cls( status=ImpactStatus.ERROR, @@ -270,10 +271,10 @@ def failed( class EconomyService: - """ - Service for calculating economic impact of policy reforms; this is connected - to the /economy route, which does not have its own table; therefore, it connects - with other services to access their respective tables + """Service that calculates the economic impact of policy reforms. + + It backs the /economy routes, which have no table of their own, so it reads + policies and cached reform impacts through the other services. """ def __init__( @@ -285,6 +286,7 @@ def __init__( budget_window_cache_: BudgetWindowCache | None = None, simulation_entrypoint_: SimulationEntrypointClient | None = None, ) -> None: + """Inject collaborators; omitted ones are built lazily or use the defaults.""" self._primary_session_factory = primary_session_factory self._injected_policy_service = policy_service_ self._injected_reform_impacts_service = reform_impacts_service_ @@ -334,7 +336,6 @@ def _get_policy_jsons( @staticmethod def _parse_json_object(value: JSONValue) -> dict[str, Any]: """Accept ORM-decoded objects and legacy JSON text at a read boundary.""" - parsed = json.loads(value) if isinstance(value, str) else value if not isinstance(parsed, dict): raise TypeError("Expected a JSON object") @@ -353,15 +354,12 @@ def get_economic_impact( api_version: str, target: Literal["general", "cliff"] = "general", ) -> EconomicImpactResult: - """ - Calculate the society-wide economic impact of a policy reform. + """Calculate the society-wide economic impact of a policy reform. - Returns a tuple of: - - status: "ok", "error", or "computing" - - reform_impact_json: The JSON object containing the economic impact data, or None if - the status is "computing" or "error". + Return a result whose status is "ok" with the impact JSON in ``data``, + "computing" while the calculation runs, or "error" with a ``message``. + Unexpected failures are logged and re-raised. """ - set_runtime_context( country_id=country_id, policy_id=policy_id, @@ -429,6 +427,14 @@ def get_budget_window_economic_impact( target: Literal["general", "cliff"] = "general", max_active_years: int = BUDGET_WINDOW_MAX_ACTIVE_YEARS, ) -> BudgetWindowEconomicImpactResult: + """Return the multi-year budget-window impact, submitting a batch job if needed. + + The shared cache records each report's state (starting, submitted, completed or + failed) so concurrent requests share one batch job; later calls poll it and + report progress until it finishes. Raise CacheCoordinationError if the cache is + unavailable or submission ownership keeps changing, and SPMValidationError when + the SPM selection or a stored result cannot be certified. + """ set_runtime_context( country_id=country_id, policy_id=policy_id, @@ -1167,10 +1173,7 @@ def _get_previous_impacts( options_hash: str, api_version: str, ) -> list[CachedReformImpact]: - """ - Fetch any previous simulation runs for the given policy reform. - """ - + """Fetch any previous simulation runs for the given policy reform.""" return self._reform_impacts.get_all_reform_impacts_by_options_hash_prefix( country_id, policy_id, @@ -1187,9 +1190,10 @@ def _get_most_recent_impact( self, setup_options: EconomicImpactSetupOptions, ) -> CachedReformImpact | None: - """ - Get the first impact from the reform_impacts table based on the provided setup options. - Returns the first impact if it exists, otherwise returns None. + """Return the newest cached impact for these options, or None if there is none. + + An impact with exactly this options hash wins; otherwise the newest impact + whose hash extends this one (with extra bundle fields) is used. """ previous_impacts = self._get_previous_impacts( country_id=setup_options.country_id, @@ -1237,8 +1241,7 @@ def _handle_execution_state( reform_impact: ReformImpact, execution: ModalSimulationExecution | None = None, ) -> EconomicImpactResult: - """ - Handle the state of the execution and return the appropriate status and result. + """Map an execution's state to the impact status and result to return. Supports both GCP Workflow statuses (SUCCEEDED, FAILED, ACTIVE) and Modal statuses (complete, failed, running, submitted). @@ -1509,10 +1512,7 @@ def _setup_sim_options( policyengine_version: str | None = None, spm: dict | None = None, ) -> SimulationOptions: - """ - Set up the simulation options for the simulation API job. - """ - + """Set up the simulation options for the simulation API job.""" return SimulationOptions.model_validate( { "country": country_id, @@ -1655,13 +1655,11 @@ def _with_policyengine_bundle( return response def _setup_region(self, country_id: str, region: str) -> str: - """ - Validate the region for the given country. + """Validate the region for the given country. Assumes region has already been normalized (e.g., "ca" -> "state/ca"). Raises ValueError for invalid regions. """ - # For US regions, validate (skip validation for national "us") if country_id == "us" and region != "us": self._validate_us_region(region) @@ -1669,8 +1667,7 @@ def _setup_region(self, country_id: str, region: str) -> str: return region def _validate_us_region(self, region: str) -> None: - """ - Validate a prefixed US region string. + """Validate a prefixed US region string. Raises ValueError if the region is not valid. """ @@ -1748,9 +1745,7 @@ def _set_reform_impact_computing( setup_options: EconomicImpactSetupOptions, execution_id: str, ) -> None: - """ - In the reform_impact table, set the status of the impact to "computing". - """ + """Record the impact as "computing" for this run; log and re-raise errors.""" try: self._reform_impacts.set_reform_impact( country_id=setup_options.country_id, @@ -1786,9 +1781,7 @@ def _set_reform_impact_complete( reform_impact_json: dict[str, Any], execution_id: str, ) -> None: - """ - In the reform_impact table, set the status of the impact to "ok" and store the reform impact JSON. - """ + """Record the impact as "ok" with its result JSON; log and re-raise errors.""" try: self._reform_impacts.set_complete_reform_impact( country_id=setup_options.country_id, @@ -1820,9 +1813,7 @@ def _set_reform_impact_error( execution_id: str, error_code: str | None = None, ) -> None: - """ - In the reform_impact table, set the status of the impact to "error" and store the error message. - """ + """Record the impact as "error" with its message; log and re-raise errors.""" try: self._reform_impacts.set_error_reform_impact( country_id=setup_options.country_id, @@ -1847,5 +1838,4 @@ def _set_reform_impact_error( def _create_submission_claim_id(self) -> str: """Create an opaque token for one submission ownership claim.""" - return str(uuid.uuid4()) diff --git a/policyengine_api/services/household_calculation_service.py b/policyengine_api/services/household_calculation_service.py index 655b7700b..f8ab17a67 100644 --- a/policyengine_api/services/household_calculation_service.py +++ b/policyengine_api/services/household_calculation_service.py @@ -1,3 +1,5 @@ +"""Calculations for stored and request-supplied households, with result caching.""" + from __future__ import annotations import time @@ -39,6 +41,8 @@ @dataclass(frozen=True) class CalculationResult: + """One country-package calculation: the household, warnings and SPM receipt.""" + household: dict warnings: tuple[str, ...] = () spm_provenance: dict | None = None @@ -46,6 +50,8 @@ class CalculationResult: @dataclass(frozen=True) class HouseholdCalculationResult: + """A household calculation as returned to the route, and whether it was cached.""" + household: dict warnings: tuple[str, ...] = () cached: bool = False @@ -66,15 +72,18 @@ class PreparedHouseholdCalculation: class HouseholdNotFoundError(LookupError): - pass + """Raised when the stored household does not exist in the requested country.""" class PolicyNotFoundError(LookupError): - pass + """Raised when the stored policy does not exist in the requested country.""" class InvalidHouseholdInputsError(ValueError): + """Raised when a household or policy has inputs the country model does not know.""" + def __init__(self, invalid_inputs: list[Any]) -> None: + """Keep the unrecognised inputs so the route can report them.""" self.invalid_inputs = invalid_inputs super().__init__("Household or policy contains unrecognized inputs") @@ -95,6 +104,10 @@ def _validated_spm_provenance( def get_household_year(household: dict) -> int | str: + """Return the year of the first age entered for ``you``, or the current year. + + This is the period used when filling in missing yearly variables. + """ household_year: int | str = date.today().year household_age_list = list( household.get("people", {}).get("you", {}).get("age", {}).keys() @@ -109,6 +122,12 @@ def add_yearly_variables( country_id: str, countries: dict | None = None, ) -> dict: + """Add every yearly, monthly and eternal variable the household's entities lack. + + Input variables get their default value and computed variables get None, which + asks the model to calculate them, for the household's year. The household is + changed in place and returned. + """ if countries is None: from policyengine_api.country import COUNTRIES @@ -143,6 +162,7 @@ def __init__( cache: HouseholdCalculationCache | None = None, country_provider: Callable[[], dict] | None = None, ) -> None: + """Inject the session factory, cache and country source, or use the defaults.""" self._injected_primary_session_factory = primary_session_factory if cache is None: context = get_runtime_cache_context() @@ -225,6 +245,13 @@ def calculate_stored_household( *, on_accepted: Callable[[], object] | None = None, ) -> HouseholdCalculationResult: + """Calculate a saved household under a saved policy, reusing a cached result. + + Raise HouseholdNotFoundError or PolicyNotFoundError if either is missing in this + country, and InvalidHouseholdInputsError for unrecognised inputs. + ``on_accepted`` runs once the inputs have loaded, before the cache lookup and + calculation. + """ api_version = COUNTRY_PACKAGE_VERSIONS[country_id] with observability_runtime.span( HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_LOAD_INPUTS) @@ -362,7 +389,6 @@ def prepare_household_calculation( spm_requested: bool = False, ) -> PreparedHouseholdCalculation: """Validate request inputs before accepting a calculation.""" - countries = self._countries() country = countries.get(country_id) spm = normalize_spm_selection(country_id, spm) @@ -398,7 +424,6 @@ def calculate_prepared_household( prepared: PreparedHouseholdCalculation, ) -> HouseholdCalculationResult: """Calculate inputs already accepted by ``prepare_household_calculation``.""" - with observability_runtime.span( HOUSEHOLD_STAGES.name(Stage.HOUSEHOLD_CALCULATION) ): @@ -443,7 +468,6 @@ def parse_prepared_household( on_accepted: Callable[[], object] | None = None, ) -> PreparedHouseholdCalculation: """Parse a prepared situation without performing requested calculations.""" - prepare_calculation = getattr(prepared.country, "prepare_calculation", None) country_calculation = prepared.country_calculation if callable(prepare_calculation): @@ -476,7 +500,6 @@ def calculate_household( spm_requested: bool = False, ) -> HouseholdCalculationResult: """Validate and calculate request-provided household and policy data.""" - prepared = self.prepare_household_calculation( country_id, household_json, diff --git a/policyengine_api/services/household_mirroring.py b/policyengine_api/services/household_mirroring.py index 276e7ca90..bcf370718 100644 --- a/policyengine_api/services/household_mirroring.py +++ b/policyengine_api/services/household_mirroring.py @@ -36,7 +36,9 @@ class LegacyHouseholdMirror(Protocol): def mirror_legacy_household( self, snapshot: LegacyHouseholdSnapshot, - ) -> LegacyHouseholdPersistenceResult: ... + ) -> LegacyHouseholdPersistenceResult: + """Copy one committed v1 household into v2 and map its v1 id to the result.""" + ... class HouseholdMirrorUnavailableError(RuntimeError): @@ -124,7 +126,6 @@ def process_household_event_after_commit( require_pending: bool = False, ) -> LegacyHouseholdPersistenceResult: """Copy one exact retained create event and then record its completion.""" - selected_event_service = event_service or HouseholdService() started_at = time.perf_counter() destination_result: LegacyHouseholdPersistenceResult | None = None diff --git a/policyengine_api/services/household_service.py b/policyengine_api/services/household_service.py index 2957d5bd9..dd3b57010 100644 --- a/policyengine_api/services/household_service.py +++ b/policyengine_api/services/household_service.py @@ -1,3 +1,5 @@ +"""v1 household storage and the retained events that copy new households into v2.""" + from __future__ import annotations import copy @@ -68,6 +70,8 @@ class HouseholdMirrorEventIntegrityError(RuntimeError): class HouseholdPersistenceFailureCategory(StrEnum): + """Safe-to-report reasons a household write failed.""" + TIMEOUT = "timeout" UNAVAILABLE = "unavailable" INTEGRITY = "integrity" @@ -79,11 +83,13 @@ class HouseholdPersistenceError(RuntimeError): """Domain error containing only safe persistence-failure attributes.""" def __init__(self, category: HouseholdPersistenceFailureCategory) -> None: + """Record the failure category; the message never carries database details.""" super().__init__("Household persistence failed") self.category = HouseholdPersistenceFailureCategory(category) @property def retryable(self) -> bool: + """Whether a retry may succeed: true for timeouts and a database outage.""" return self.category in { HouseholdPersistenceFailureCategory.TIMEOUT, HouseholdPersistenceFailureCategory.UNAVAILABLE, @@ -91,6 +97,7 @@ def retryable(self) -> bool: @classmethod def from_exception(cls, error: Exception) -> HouseholdPersistenceError: + """Classify a database or unexpected error as a safe persistence error.""" if isinstance(error, SQLAlchemyTimeoutError): return cls(HouseholdPersistenceFailureCategory.TIMEOUT) if isinstance(error, OperationalError): @@ -104,6 +111,8 @@ def from_exception(cls, error: Exception) -> HouseholdPersistenceError: @dataclass(frozen=True) class PendingHouseholdMirrorEvent: + """A validated retained household copy event, and whether it was processed.""" + event_id: int snapshot: LegacyHouseholdSnapshot source_fingerprint_sha256: str @@ -112,6 +121,8 @@ class PendingHouseholdMirrorEvent: @dataclass(frozen=True) class HouseholdCreateResult: + """A newly stored v1 household, with its v2 copy snapshot and any event id.""" + household: Household snapshot: LegacyHouseholdSnapshot | None mirror_event_id: int | None @@ -124,6 +135,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -135,6 +147,10 @@ def get_household( country_id: str, household_id: int, ) -> Household | None: + """Return the stored household in this country, or None if there is none. + + Raise Exception if ``household_id`` is not a non-negative int. + """ if type(household_id) is not int or household_id < 0: raise Exception( f"Invalid household ID: {household_id}. Must be a positive integer." @@ -164,6 +180,11 @@ def create_household( record_mirror_event: bool = False, spm: object = None, ) -> HouseholdCreateResult: + """Store a v1 household, optionally recording an event to copy it into v2. + + Raise SPMValidationError for an invalid SPM selection before writing, and + HouseholdPersistenceError, carrying a safe category, if the write fails. + """ # Validate before the persistence boundary so typed input failures remain 400. household_json = household_storage_json(country_id, household_json, spm) try: @@ -278,7 +299,6 @@ def process_mirror_event( require_pending: bool = False, ) -> LegacyHouseholdPersistenceResult: """Process one explicitly selected retained create event.""" - with self._sessions.begin() as session: event = session.scalar( select(HouseholdMirrorEvent) diff --git a/policyengine_api/services/metadata_service.py b/policyengine_api/services/metadata_service.py index b71791a23..35bdf86b2 100644 --- a/policyengine_api/services/metadata_service.py +++ b/policyengine_api/services/metadata_service.py @@ -1,9 +1,14 @@ +"""Read access to the metadata document of each loaded country model.""" + from policyengine_api.country import COUNTRIES from policyengine_api.json_types import JSONObject class MetadataService: + """Service that returns a country's metadata from the loaded country models.""" + def get_metadata(self, country_id: str) -> JSONObject: + """Return the country's metadata; raise RuntimeError for an unknown country.""" country = COUNTRIES.get(country_id) if country is None: raise RuntimeError( diff --git a/policyengine_api/services/policy_mirroring.py b/policyengine_api/services/policy_mirroring.py index 257481f7f..1de90ee88 100644 --- a/policyengine_api/services/policy_mirroring.py +++ b/policyengine_api/services/policy_mirroring.py @@ -36,7 +36,9 @@ class LegacyPolicyMirror(Protocol): def mirror_legacy_policy( self, snapshot: LegacyPolicySnapshot, - ) -> LegacyPolicyPersistenceResult: ... + ) -> LegacyPolicyPersistenceResult: + """Copy one committed v1 policy into v2 and map its v1 id to the result.""" + ... class PolicyMirrorUnavailableError(RuntimeError): @@ -125,7 +127,6 @@ def mirror_policy_after_commit( mirror_factory: Callable[[], LegacyPolicyMirror] | None = None, ) -> LegacyPolicyPersistenceResult: """Require the Supabase policy and mapping transaction before v1 success.""" - started_at = time.perf_counter() try: result = (mirror_factory or _default_mirror_factory)().mirror_legacy_policy( diff --git a/policyengine_api/services/policy_service.py b/policyengine_api/services/policy_service.py index 5112d9429..d74f21c06 100644 --- a/policyengine_api/services/policy_service.py +++ b/policyengine_api/services/policy_service.py @@ -1,3 +1,5 @@ +"""v1 policy storage: lookup, label search and deduplicated creation.""" + from __future__ import annotations import copy @@ -31,7 +33,6 @@ class PolicySetResult: def __iter__(self) -> Iterator[int | str | bool]: """Preserve the established three-value internal unpacking interface.""" - yield self.policy_id yield self.message yield self.is_existing_policy @@ -44,6 +45,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -62,6 +64,11 @@ def get_policy( country_id: str, policy_id: int, ) -> Policy | None: + """Return the stored policy in this country, or None if there is none. + + Raise Exception if ``policy_id`` is not a non-negative int, and ValueError for + an empty ``country_id``. + """ self._validate_policy_id(policy_id) if not country_id: raise ValueError("country_id cannot be empty or None") @@ -86,6 +93,7 @@ def get_policy_json( country_id: str, policy_id: int, ) -> JSONValue: + """Return the policy's stored JSON, or None if the policy does not exist.""" policy = self.get_policy(country_id, policy_id) return None if policy is None else policy.policy_json @@ -95,7 +103,6 @@ def get_policy_snapshot( policy_id: int, ) -> LegacyPolicySnapshot | None: """Return detached fields required to mirror one existing v1 policy.""" - policy = self.get_policy(country_id, policy_id) if policy is None: return None @@ -115,6 +122,11 @@ def search_policies( *, unique_only: bool = False, ) -> list[Policy]: + """Return this country's policies whose label contains ``query``. + + With ``unique_only``, keep only the first policy for each label and content + hash. + """ with self._sessions() as session: results = list( session.scalars( @@ -144,6 +156,11 @@ def set_policy( *, prepare_for_mirroring: bool = False, ) -> PolicySetResult: + """Store a policy, or return the existing one with the same content and label. + + Raise ValueError for an unknown country. With ``prepare_for_mirroring`` the + result carries a snapshot of the stored row, for copying the policy into v2. + """ country_id = country_id.lower() if country_id not in COUNTRY_PACKAGE_VERSIONS: raise ValueError(f"Invalid country_id: {country_id}") diff --git a/policyengine_api/services/reform_impacts_service.py b/policyengine_api/services/reform_impacts_service.py index 1658ff655..e199a18b0 100644 --- a/policyengine_api/services/reform_impacts_service.py +++ b/policyengine_api/services/reform_impacts_service.py @@ -21,6 +21,7 @@ class ReformImpactsService: """Preserve historical lookup contracts over an expiring cache repository.""" def __init__(self, cache: ReformImpactCache | None = None) -> None: + """Use this reform-impact cache, or one on the runtime cache if omitted.""" if cache is None: context = get_runtime_cache_context() cache = ReformImpactCache(context.client, context.namespace) @@ -30,6 +31,7 @@ def get_recent_reform_impacts( self, max_results: int, ) -> list[CachedReformImpact]: + """Return up to ``max_results`` cached impacts, most recently started first.""" return self._cache.recent(max_results) def get_all_reform_impacts( @@ -43,6 +45,7 @@ def get_all_reform_impacts( options_hash: str, api_version: str, ) -> list[CachedReformImpact]: + """Return cached impacts with exactly this options hash, newest first.""" return self._cache.matching( country_id=country_id, reform_policy_id=policy_id, @@ -66,6 +69,11 @@ def get_all_reform_impacts_by_options_hash_prefix( options_hash_prefix: str, api_version: str, ) -> list[CachedReformImpact]: + """Return cached impacts matching the options hash or the prefix pattern. + + ``options_hash_prefix`` is a SQL ``LIKE`` pattern. Exact hash matches come + first, then the most recently started. + """ return self._cache.matching( country_id=country_id, reform_policy_id=policy_id, @@ -94,7 +102,6 @@ def claim_reform_impact_start( observability_id: str, ) -> bool: """Fail closed unless this request atomically owns job submission.""" - return self._cache.claim_start( country_id=country_id, reform_policy_id=policy_id, @@ -123,7 +130,6 @@ def get_reform_impact_start_claim( target: str, ) -> ReformImpactStartClaim | None: """Return the current submission owner and diagnostic identifier.""" - return self._cache.get_start_claim( country_id=country_id, reform_policy_id=policy_id, @@ -152,7 +158,6 @@ def release_reform_impact_start( observability_id: str, ) -> None: """Best-effort release; an unavailable cache safely falls back to expiry.""" - try: self._cache.release_start( country_id=country_id, @@ -187,6 +192,11 @@ def set_reform_impact( execution_id: str, observability_id: str | None = None, ) -> CachedReformImpact: + """Cache a newly submitted impact by execution ID so later requests can poll it. + + Raise ReformImpactHandoffError if the write fails, because the accepted job + could not otherwise be found again. + """ impact = CachedReformImpact( reform_impact_id=reform_impact_id(execution_id), country_id=country_id, @@ -222,6 +232,7 @@ def delete_reform_impact( time_period: str, options_hash: str, ) -> None: + """Remove this scope's in-flight (computing) impacts, ignoring cache errors.""" self._cache.delete_matching_computing( country_id=country_id, reform_policy_id=policy_id, @@ -245,6 +256,13 @@ def set_error_reform_impact( execution_id: str, error_code: str | None = None, ) -> CachedReformImpact | None: + """Mark this execution's cached impact as failed. + + Only ``execution_id`` selects the record; the scope arguments are kept for the + historical signature and ignored. An ``error_code`` records a typed, terminal + failure and clears the worker handle. Return the updated impact, or None if + it is missing or the write failed. + """ del ( country_id, policy_id, @@ -275,6 +293,11 @@ def set_complete_reform_impact( reform_impact_json: dict[str, Any], execution_id: str, ) -> CachedReformImpact | None: + """Store the finished result on this execution's cached impact. + + Only ``execution_id`` selects the record; the scope arguments are ignored. + Return the updated impact, or None if it is missing or the write failed. + """ del ( country_id, reform_policy_id, diff --git a/policyengine_api/services/report_output_alias_service.py b/policyengine_api/services/report_output_alias_service.py index 90d23c6bc..14feae677 100644 --- a/policyengine_api/services/report_output_alias_service.py +++ b/policyengine_api/services/report_output_alias_service.py @@ -1,3 +1,5 @@ +"""Aliases that redirect a legacy v1 report output ID to its canonical report output.""" + from sqlalchemy.orm import Session from policyengine_api.data.v1_models import LegacyReportOutputAlias, ReportOutput @@ -10,11 +12,17 @@ class ReportOutputAliasService: def get_alias( session: Session, legacy_report_output_id: int ) -> LegacyReportOutputAlias | None: + """Return the alias stored for this legacy report output ID, or None.""" return session.get(LegacyReportOutputAlias, legacy_report_output_id) def resolve_canonical_report_output_id( self, session: Session, requested_report_output_id: int ) -> int | None: + """Return the report output ID to serve for a requested ID, following any alias. + + Return None if there is neither an alias nor a report output with this ID, and + raise ValueError if the alias points to a missing report output. + """ alias = self.get_alias(session, requested_report_output_id) if alias is not None: canonical_id = alias.canonical_report_output_id @@ -32,6 +40,13 @@ def set_alias( legacy_report_output_id: int, canonical_report_output_id: int, ) -> bool: + """Point a legacy report output at a canonical one describing the same report. + + Repeating an existing alias is a no-op; the caller owns the transaction. Raise + ValueError if either output is missing, both IDs are the same, the legacy output + is already aliased elsewhere, or the two differ in country, simulations or year. + Always return True. + """ legacy = session.get(ReportOutput, legacy_report_output_id) if legacy is None: raise ValueError( diff --git a/policyengine_api/services/report_output_service.py b/policyengine_api/services/report_output_service.py index cbcfc7cf8..96dd213b1 100644 --- a/policyengine_api/services/report_output_service.py +++ b/policyengine_api/services/report_output_service.py @@ -1,3 +1,5 @@ +"""v1 report outputs, with every write mirrored into the report's run history.""" + import json from dataclasses import dataclass from datetime import UTC, datetime @@ -29,7 +31,6 @@ def _canonicalize_annual_spm_output( expected_year: str | None, ) -> JSONObject: """Validate and serialize compact SPM provenance on one economy result.""" - if "spm_config" in output: raise ValueError( "Completed US economy output must not include legacy spm_config" @@ -58,7 +59,6 @@ def _canonicalize_society_wide_spm_output( expected_year: str | None, ) -> object: """Validate compact annual or budget-window society-wide SPM output.""" - if not isinstance(output, dict): return output typed_output = cast("JSONObject", output) @@ -87,6 +87,12 @@ def _canonicalize_society_wide_spm_output( @dataclass(frozen=True) class ReportOutputView: + """A report output as served: the record, the run to display and the ID to report. + + ``response_id`` is the requested ID when an outdated report was redirected to its + current equivalent; None means the record's own ID is reported. + """ + report_output: ReportOutput display_run: ReportOutputRun | None response_id: int | None = None @@ -94,6 +100,8 @@ class ReportOutputView: @dataclass(frozen=True) class ReportCreateResult: + """A report output view and whether this call created the report.""" + view: ReportOutputView created: bool @@ -105,6 +113,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory for reports and simulations, or the v1 one.""" self._injected_session_factory = session_factory self.report_spec_service = ReportSpecService() self.report_run_service = ReportRunService() @@ -120,6 +129,11 @@ def _utc_timestamp() -> datetime: @staticmethod def format_run_timestamp(value: datetime | str | None) -> str | None: + """Return a timestamp as an ISO 8601 UTC string ending in ``Z``, or None. + + Naive values are taken as UTC; an unparseable string is returned with ``Z`` + added. + """ if value is None: return None if isinstance(value, datetime): @@ -609,6 +623,10 @@ def create_or_reuse_report_output( simulation_2_id: int | None = None, year: str = "2025", ) -> ReportCreateResult: + """Return the current report for these simulations and year, creating it if new. + + Raise ValueError if a referenced simulation does not exist in this country. + """ with self._sessions.begin() as session: existing = self._find_existing_report_output( session, @@ -641,6 +659,13 @@ def resolve_report_output( country_id: str, report_output_id: int, ) -> ReportOutputView | None: + """Return the report output to serve for a requested ID, or None if missing. + + A report stored under an outdated cache version is served as the current report + for the same simulations and year (created if needed), with the requested ID + kept as ``response_id``. Raise Exception for an ID that is not a non-negative + int. + """ if type(report_output_id) is not int or report_output_id < 0: raise Exception( f"Invalid report output ID: {report_output_id}. " @@ -674,6 +699,12 @@ def update_report_output( output: dict | list | str | None = None, error_message: str | None = None, ) -> ReportOutputView | None: + """Apply the given status, output or error message and sync the run history. + + Return None when nothing was given. Raise LookupError if the report does not + exist, and ValueError when marking it running without an active run or + completing a US economy report whose SPM provenance is invalid. + """ values = { key: value for key, value in { diff --git a/policyengine_api/services/report_run_service.py b/policyengine_api/services/report_run_service.py index aaf6daf53..7db920901 100644 --- a/policyengine_api/services/report_run_service.py +++ b/policyengine_api/services/report_run_service.py @@ -1,3 +1,5 @@ +"""Run history for v1 report outputs, one numbered row per calculation attempt.""" + import uuid from datetime import UTC, datetime from typing import Any @@ -44,6 +46,12 @@ def create_report_output_run( version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, ) -> ReportOutputRun: + """Append the next numbered run to a report output and return it. + + The parent row is locked while the sequence number is chosen, and timestamps are + set to now according to ``status``. Raise ValueError if the report output does + not exist. The caller owns the transaction. + """ parent = session.scalar( select(ReportOutput) .where(ReportOutput.id == report_output_id) @@ -90,11 +98,13 @@ def create_report_output_run( def get_report_output_run( self, session: Session, run_id: str ) -> ReportOutputRun | None: + """Return the run with this ID, or None.""" return session.get(ReportOutputRun, run_id) def list_report_output_runs( self, session: Session, report_output_id: int ) -> list[ReportOutputRun]: + """Return all runs of a report output, oldest first.""" return list( session.scalars( select(ReportOutputRun) @@ -106,6 +116,7 @@ def list_report_output_runs( def get_newest_report_output_run( self, session: Session, report_output_id: int ) -> ReportOutputRun | None: + """Return the report output's highest-numbered run, or None if it has none.""" return session.scalar( select(ReportOutputRun) .where(ReportOutputRun.report_output_id == report_output_id) @@ -115,6 +126,12 @@ def get_newest_report_output_run( def select_display_run( self, session: Session, report_output: ReportOutput ) -> ReportOutputRun | None: + """Return the run whose result the report output should display, or None. + + Prefer the active run, then (for a failed report) the run holding its error, + then the latest successful run, then any run matching the report's result, and + finally the newest run. + """ runs = list( session.scalars( select(ReportOutputRun) diff --git a/policyengine_api/services/report_spec_service.py b/policyengine_api/services/report_spec_service.py index 6bcb82609..a3aa8444b 100644 --- a/policyengine_api/services/report_spec_service.py +++ b/policyengine_api/services/report_spec_service.py @@ -1,3 +1,5 @@ +"""Explicit specifications of what each v1 report output calculates.""" + import json from typing import Any, Literal @@ -14,12 +16,16 @@ class ReportSimulationInput(BaseModel): + """The population and policy of one simulation in a household report.""" + population_type: Literal["household", "geography"] population_id: str policy_id: int class HouseholdReportSpec(BaseModel): + """A household report: one household under one policy, or under two to compare.""" + country_id: str report_kind: Literal["household_single", "household_comparison"] time_period: str @@ -28,6 +34,11 @@ class HouseholdReportSpec(BaseModel): class EconomyReportSpec(BaseModel): + """An economy report: one region under a baseline and a reform policy. + + A single-policy report uses the same policy for both. + """ + country_id: str report_kind: Literal["economy_single", "economy_comparison"] time_period: str @@ -43,6 +54,8 @@ class EconomyReportSpec(BaseModel): class ReportSpecService: + """Service that builds, checks and stores report specs for v1 report outputs.""" + @staticmethod def _validate_schema_version(schema_version: int | None) -> None: if schema_version != REPORT_SPEC_SCHEMA_VERSION: @@ -120,6 +133,10 @@ def _validate_report_country( def infer_report_kind( simulation_1: Simulation, simulation_2: Simulation | None = None ) -> str: + """Return the report kind implied by the simulations' population type and count. + + Raise ValueError if the two simulations' population types differ or are unknown. + """ population_type = simulation_1.population_type if simulation_2 is not None and simulation_2.population_type != population_type: raise ValueError( @@ -140,6 +157,11 @@ def build_report_spec( target: Literal["general", "cliff"] = "general", options: dict[str, Any] | None = None, ) -> ReportSpec: + """Build the spec implied by a report output and its linked simulations. + + Raise ValueError if the simulations do not match the report's simulation IDs or + country, or if a comparison's simulations cover different households or regions. + """ self._validate_report_simulation_linkage( report_output, simulation_1, simulation_2 ) @@ -243,6 +265,11 @@ def _parse_report_spec(report_kind: str, raw_spec: dict) -> ReportSpec: def get_report_spec( self, session: Session, report_output_id: int ) -> ReportSpec | None: + """Return the stored spec for a report output, or None if it has none. + + Raise ValueError if the stored spec has an unsupported schema version or no + longer matches the report's linked simulations. + """ report_output = session.get(ReportOutput, report_output_id) if report_output is None or report_output.report_spec_json is None: return None @@ -264,6 +291,11 @@ def set_report_spec( report_spec_status: Literal["explicit", "backfilled_assumed"], schema_version: int = REPORT_SPEC_SCHEMA_VERSION, ) -> bool: + """Store a report spec after checking it matches the linked simulations. + + Raise ValueError for an unknown status or schema version, a missing report, or a + mismatched spec. Always return True; the caller owns the transaction. + """ if report_spec_status not in REPORT_SPEC_STATUSES: raise ValueError(f"Unsupported report spec status: {report_spec_status}") self._validate_schema_version(schema_version) diff --git a/policyengine_api/services/simulation_run_service.py b/policyengine_api/services/simulation_run_service.py index 677f2b6d5..91b3929fa 100644 --- a/policyengine_api/services/simulation_run_service.py +++ b/policyengine_api/services/simulation_run_service.py @@ -1,3 +1,5 @@ +"""Run history for v1 simulations, one numbered row per calculation attempt.""" + import uuid from typing import Any @@ -16,6 +18,8 @@ class SimulationRunService: + """Simulation-run operations performed through a caller-owned ORM session.""" + def create_simulation_run( self, session: Session, @@ -31,6 +35,12 @@ def create_simulation_run( version_manifest: dict[str, str | None] | None = None, run_id: str | None = None, ) -> SimulationRun: + """Append the next numbered run to a simulation and return it. + + The parent row is locked while the sequence number is chosen; timestamps start + unset. Raise ValueError if the simulation does not exist. The caller owns the + transaction. + """ parent = session.scalar( select(Simulation).where(Simulation.id == simulation_id).with_for_update() ) @@ -78,6 +88,7 @@ def get_simulation_run( session: Session, run_id: str, ) -> SimulationRun | None: + """Return the run with this ID, or None.""" return session.get(SimulationRun, run_id) def list_simulation_runs( @@ -85,6 +96,7 @@ def list_simulation_runs( session: Session, simulation_id: int, ) -> list[SimulationRun]: + """Return all runs of a simulation, oldest first.""" return list( session.scalars( select(SimulationRun) @@ -98,6 +110,7 @@ def get_newest_simulation_run( session: Session, simulation_id: int, ) -> SimulationRun | None: + """Return the simulation's highest-numbered run, or None if it has none.""" return session.scalar( select(SimulationRun) .where(SimulationRun.simulation_id == simulation_id) @@ -109,6 +122,10 @@ def select_display_run( session: Session, simulation: Simulation, ) -> SimulationRun | None: + """Return the run whose result the simulation should display, or None. + + Prefer the active run, then the latest successful run, then the newest run. + """ if simulation.active_run_id: active_run = self.get_simulation_run(session, simulation.active_run_id) if active_run is not None: diff --git a/policyengine_api/services/simulation_service.py b/policyengine_api/services/simulation_service.py index 3ef805673..da9e0a8b0 100644 --- a/policyengine_api/services/simulation_service.py +++ b/policyengine_api/services/simulation_service.py @@ -1,3 +1,5 @@ +"""v1 simulations, with every write mirrored into the simulation's run history.""" + import json import uuid from dataclasses import dataclass @@ -17,6 +19,8 @@ @dataclass(frozen=True) class SimulationCreateResult: + """A simulation and whether this call created it.""" + simulation: Simulation created: bool @@ -28,6 +32,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -207,6 +212,12 @@ def get_or_create_simulation( population_type: str, policy_id: int, ) -> SimulationCreateResult: + """Return the simulation for this population and policy, creating it if needed. + + A numeric household ID also matches its stored aliases. Stored output that + predates the current SPM receipt format is cleared and the simulation reset to + pending, so it is calculated again. + """ numeric_identity = canonical_numeric_household_id( country_id, population_id, population_type ) @@ -252,6 +263,10 @@ def get_simulation( country_id: str, simulation_id: int, ) -> Simulation | None: + """Return the simulation in this country, or None. + + Raise Exception if ``simulation_id`` is not a non-negative int. + """ if type(simulation_id) is not int or simulation_id < 0: raise Exception( f"Invalid simulation ID: {simulation_id}. Must be a positive integer." @@ -267,6 +282,12 @@ def update_simulation( output: dict | list | str | None = None, error_message: str | None = None, ) -> Simulation | None: + """Apply the given status, output or error message and sync the run history. + + Return None when nothing was given. The simulation is restamped with the current + country package version. Raise ValueError for output with legacy ``spm_config`` + or an invalid SPM receipt, and LookupError if the simulation does not exist. + """ values = { key: value for key, value in { diff --git a/policyengine_api/services/simulation_spec_service.py b/policyengine_api/services/simulation_spec_service.py index f6c5a7373..d75f6a76f 100644 --- a/policyengine_api/services/simulation_spec_service.py +++ b/policyengine_api/services/simulation_spec_service.py @@ -1,3 +1,5 @@ +"""Explicit specifications of what each v1 simulation calculates.""" + import json from typing import Literal @@ -10,6 +12,8 @@ class SimulationSpec(BaseModel): + """The population and policy one simulation runs, as stored on its row.""" + country_id: str population_id: str population_type: Literal["household", "geography"] @@ -17,6 +21,8 @@ class SimulationSpec(BaseModel): class SimulationSpecService: + """Service that builds, checks and stores the spec saved on a v1 simulation.""" + def _validate_schema_version(self, schema_version: int | None) -> None: if schema_version != SIMULATION_SPEC_SCHEMA_VERSION: raise ValueError( @@ -38,6 +44,7 @@ def _validate_simulation_spec_matches_model( raise ValueError("Simulation spec must match the linked simulation row") def build_simulation_spec(self, simulation: Simulation) -> SimulationSpec: + """Return the spec implied by a simulation's country, population and policy.""" return SimulationSpec( country_id=simulation.country_id, population_id=simulation.population_id, @@ -50,6 +57,11 @@ def get_simulation_spec( session: Session, simulation_id: int, ) -> SimulationSpec | None: + """Return the stored spec for a simulation, or None if it has none. + + Raise ValueError if the stored spec has an unsupported schema version or no + longer matches the simulation row. + """ simulation = session.get(Simulation, simulation_id) if simulation is None or simulation.simulation_spec_json is None: return None @@ -71,6 +83,11 @@ def set_simulation_spec( simulation_spec: SimulationSpec, schema_version: int = SIMULATION_SPEC_SCHEMA_VERSION, ) -> bool: + """Store a spec on a simulation after checking it matches the simulation row. + + Raise ValueError for an unsupported schema version, a missing simulation or a + mismatched spec. Always return True; the caller owns the transaction. + """ self._validate_schema_version(schema_version) simulation = session.get(Simulation, simulation_id) if simulation is None: diff --git a/policyengine_api/services/user_policy_mirroring.py b/policyengine_api/services/user_policy_mirroring.py index 7f24d8de9..4b28a7aac 100644 --- a/policyengine_api/services/user_policy_mirroring.py +++ b/policyengine_api/services/user_policy_mirroring.py @@ -40,6 +40,8 @@ class LegacyUserPolicyMirror(Protocol): + """Supabase transaction service used after a v1 saved-policy commit.""" + def mirror_legacy_user_policy( self, snapshot: LegacyUserPolicySnapshot, @@ -47,7 +49,9 @@ def mirror_legacy_user_policy( *, source_revision: int, changed_fields: frozenset[str], - ) -> LegacyUserPolicyPersistenceResult: ... + ) -> LegacyUserPolicyPersistenceResult: + """Copy one v1 saved-policy revision, with its reform policy, into v2.""" + ... class UserPolicyMirrorUnavailableError(RuntimeError): @@ -169,7 +173,6 @@ def mirror_user_policy_after_commit( mirror_factory: Callable[[], LegacyUserPolicyMirror] | None = None, ) -> LegacyUserPolicyPersistenceResult: """Require one complete Supabase association transaction before v1 success.""" - started_at = time.perf_counter() try: result = _run_user_policy_mirror( @@ -217,8 +220,12 @@ def mirror_pending_user_policy_events_after_commit( | None = None, mirror_factory: Callable[[], LegacyUserPolicyMirror] | None = None, ) -> LegacyUserPolicyPersistenceResult: - """Synchronously apply retained source events through one request's revision.""" + """Copy this saved policy's retained events into v2, up to the request's revision. + Runs synchronously after the v1 commit and returns the result for + ``through_revision``. Raise UserPolicyMirrorUnavailableError, after logging the + failure, if an event cannot be copied. + """ selected_event_service = event_service or UserPolicyService() selected_reform_loader = ( reform_snapshot_loader or PolicyService().get_policy_snapshot diff --git a/policyengine_api/services/user_policy_service.py b/policyengine_api/services/user_policy_service.py index 3818791e5..98d26eb12 100644 --- a/policyengine_api/services/user_policy_service.py +++ b/policyengine_api/services/user_policy_service.py @@ -1,3 +1,5 @@ +"""v1 saved policies and the retained revision events that copy them into v2.""" + from __future__ import annotations from dataclasses import dataclass @@ -64,6 +66,8 @@ class UserPolicyMirrorEventIntegrityError(RuntimeError): class UserPolicyPersistenceFailureCategory(StrEnum): + """Safe-to-report reasons a saved-policy write failed.""" + TIMEOUT = "timeout" UNAVAILABLE = "unavailable" INTEGRITY = "integrity" @@ -75,11 +79,13 @@ class UserPolicyPersistenceError(RuntimeError): """Domain error containing only safe persistence-failure attributes.""" def __init__(self, category: UserPolicyPersistenceFailureCategory) -> None: + """Record the failure category; the message never carries database details.""" super().__init__("Saved-policy persistence failed") self.category = UserPolicyPersistenceFailureCategory(category) @property def retryable(self) -> bool: + """Whether a retry may succeed: true for timeouts and a database outage.""" return self.category in { UserPolicyPersistenceFailureCategory.TIMEOUT, UserPolicyPersistenceFailureCategory.UNAVAILABLE, @@ -87,6 +93,7 @@ def retryable(self) -> bool: @classmethod def from_exception(cls, error: Exception) -> UserPolicyPersistenceError: + """Classify a database or unexpected error as a safe persistence error.""" if isinstance(error, SQLAlchemyTimeoutError): return cls(UserPolicyPersistenceFailureCategory.TIMEOUT) if isinstance(error, OperationalError): @@ -100,6 +107,8 @@ def from_exception(cls, error: Exception) -> UserPolicyPersistenceError: @dataclass(frozen=True) class PendingUserPolicyMirrorEvent: + """A validated retained saved-policy event: one revision and what it changed.""" + event_id: int source_revision: int event_type: str @@ -110,6 +119,8 @@ class PendingUserPolicyMirrorEvent: @dataclass(frozen=True) class UserPolicyCreateResult: + """A saved policy, whether this call created it, and any mirror revision.""" + user_policy: UserPolicy created: bool snapshot: LegacyUserPolicySnapshot | None @@ -118,6 +129,8 @@ class UserPolicyCreateResult: @dataclass(frozen=True) class UserPolicyUpdateResult: + """An updated saved policy, the fields set, and its mirror revision if recorded.""" + user_policy: UserPolicy snapshot: LegacyUserPolicySnapshot | None changed_fields: frozenset[str] @@ -131,6 +144,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -277,6 +291,12 @@ def create_or_get_user_policy( *, record_mirror_event: bool = False, ) -> UserPolicyCreateResult: + """Return the saved policy with these identity fields, creating it if missing. + + With ``record_mirror_event`` the row is locked and a create revision event is + recorded, even for an existing policy. Raise UserPolicyPersistenceError, + carrying a safe category, if the write fails. + """ try: with self._sessions.begin() as session: user_policy = self._find_matching_user_policy( @@ -315,6 +335,7 @@ def list_user_policies( country_id: str, user_id: str, ) -> list[UserPolicy]: + """Return every saved policy this user has in this country.""" with self._sessions() as session: return list( session.scalars( @@ -333,6 +354,12 @@ def update_user_policy( *, record_mirror_event: bool = False, ) -> UserPolicyUpdateResult | None: + """Set the given fields on a saved policy; return None if not in this country. + + With ``record_mirror_event`` the row is locked and an update revision event is + recorded. Raise UserPolicyPersistenceError, carrying a safe category, if the + write fails. + """ try: with self._sessions.begin() as session: statement = select(UserPolicy).where( @@ -389,7 +416,6 @@ def process_pending_mirror_events( | None = None, ) -> LegacyUserPolicyPersistenceResult: """Process retained source events in order through one request's revision.""" - latest_result: LegacyUserPolicyPersistenceResult | None = None while True: with self._sessions.begin() as session: diff --git a/policyengine_api/services/user_service.py b/policyengine_api/services/user_service.py index f27fa7d1e..e60a1a11b 100644 --- a/policyengine_api/services/user_service.py +++ b/policyengine_api/services/user_service.py @@ -1,3 +1,5 @@ +"""v1 user profiles, keyed by internal user ID and by Auth0 ID.""" + from __future__ import annotations from typing import TYPE_CHECKING @@ -18,6 +20,7 @@ def __init__( self, session_factory: sessionmaker[Session] | None = None, ) -> None: + """Use this session factory, or the application's v1 one if omitted.""" self._injected_session_factory = session_factory @property @@ -31,6 +34,10 @@ def create_profile( username: str | None, user_since: int, ) -> tuple[bool, UserProfile]: + """Create a profile for this Auth0 ID, or return the one that already exists. + + Return ``(created, profile)``. + """ with self._sessions.begin() as session: return self._create_profile( session, @@ -67,6 +74,10 @@ def get_profile( auth0_id: str | None = None, user_id: int | str | None = None, ) -> UserProfile | None: + """Return the profile with this user ID, or else this Auth0 ID, or None. + + ``user_id`` wins when both are given. Raise ValueError if neither is given. + """ if auth0_id is None and user_id is None: raise ValueError("you must specify either auth0_id or user_id") with self._sessions() as session: @@ -92,6 +103,10 @@ def update_profile( username: str | None, user_since: int | None, ) -> UserProfile | None: + """Set the given non-None fields on a profile; return None if it does not exist. + + Raise ValueError if ``user_id`` is None. + """ if user_id is None: raise ValueError("you must specify either auth0_id or user_id") with self._sessions.begin() as session: diff --git a/policyengine_api/services/v2/associations/services.py b/policyengine_api/services/v2/associations/services.py index 04c02c7b3..a8fbe8130 100644 --- a/policyengine_api/services/v2/associations/services.py +++ b/policyengine_api/services/v2/associations/services.py @@ -21,7 +21,6 @@ def require_legacy_reference( legacy_id: int, ) -> None: """Check one v1 reference before the caller opens its v2 transaction.""" - require_legacy_reference_checker(legacy_references).require_existing( country_id=country_id, requested={kind: (legacy_id,)}, diff --git a/policyengine_api/services/v2/associations/types.py b/policyengine_api/services/v2/associations/types.py index 1fe4cf552..e3da0b36c 100644 --- a/policyengine_api/services/v2/associations/types.py +++ b/policyengine_api/services/v2/associations/types.py @@ -62,11 +62,15 @@ def existing_ids( *, country_id: str, requested: Mapping[LegacyRecordKind, Collection[int]], - ) -> dict[LegacyRecordKind, frozenset[int]]: ... + ) -> dict[LegacyRecordKind, frozenset[int]]: + """Return, for every requested kind, the IDs that exist in this country.""" + ... def require_existing( self, *, country_id: str, requested: Mapping[LegacyRecordKind, Collection[int]], - ) -> None: ... + ) -> None: + """Raise ``LegacyReferenceNotFoundError`` if any requested ID is missing.""" + ... diff --git a/policyengine_api/services/v2/associations/validators.py b/policyengine_api/services/v2/associations/validators.py index 68159d4d7..b87673eb0 100644 --- a/policyengine_api/services/v2/associations/validators.py +++ b/policyengine_api/services/v2/associations/validators.py @@ -37,6 +37,7 @@ class LegacyReferenceCheckUnavailableError(RuntimeError): def require_legacy_reference_checker( legacy_references: LegacyReferenceChecker | None, ) -> LegacyReferenceChecker: + """Return the checker, or raise LegacyReferenceCheckUnavailableError if None.""" if legacy_references is None: raise LegacyReferenceCheckUnavailableError( "v1 reference checks are not configured" @@ -45,6 +46,7 @@ def require_legacy_reference_checker( def require_not_mirror_owned(mirror_owned: bool, *, association_id: UUID) -> None: + """Raise MirrorOwnedAssociationError if the v1 saved-policy mirror owns the row.""" if mirror_owned: raise MirrorOwnedAssociationError( f"association {association_id} is maintained from v1 saved policies " diff --git a/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py b/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py index 32c0ca17c..14a0412d1 100644 --- a/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py +++ b/policyengine_api/services/v2/comparison_runs/database_connectors/creates.py @@ -29,6 +29,11 @@ def create_comparison_report( session: Session, record: ComparisonReportRecord, ) -> UUID | None: + """Insert a report row unless a row with the same identity already exists. + + Return the new evaluation_id, or None when uq_stage12_eval_reports_identity + already holds a row for this identity. + """ statement = ( insert(Stage12ComparisonReport) .values(**report_insert_values(record)) @@ -42,6 +47,11 @@ def create_comparison_simulation( session: Session, record: ComparisonSimulationRecord, ) -> UUID | None: + """Insert a simulation row unless a row with the same identity already exists. + + Return the new simulation_execution_id, or None when + uq_stage12_eval_simulations_identity already holds a row for this identity. + """ statement = ( insert(Stage12ComparisonSimulation) .values(**simulation_insert_values(record)) diff --git a/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py b/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py index c06b6d337..032f3feb2 100644 --- a/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py +++ b/policyengine_api/services/v2/comparison_runs/database_connectors/reads.py @@ -26,6 +26,10 @@ def read_comparison_report( *, lock: bool, ) -> Stage12ComparisonReport | None: + """Return the report row for evaluation_id, or None. + + With lock=True the row is read FOR UPDATE until the transaction ends. + """ statement = select(Stage12ComparisonReport).where( Stage12ComparisonReport.evaluation_id == evaluation_id ) @@ -38,6 +42,11 @@ def read_comparison_report_by_identity( session: Session, record: ComparisonReportRecord, ) -> Stage12ComparisonReport | None: + """Return the report row with the record's identity, or None. + + The identity is the environment, calculation flow, production identity, + worker version and version-manifest digest. + """ return session.exec( select(Stage12ComparisonReport).where( Stage12ComparisonReport.environment == record.environment, @@ -56,6 +65,10 @@ def read_comparison_simulation( *, lock: bool, ) -> Stage12ComparisonSimulation | None: + """Return the simulation row for simulation_execution_id, or None. + + With lock=True the row is read FOR UPDATE until the transaction ends. + """ statement = select(Stage12ComparisonSimulation).where( Stage12ComparisonSimulation.simulation_execution_id == simulation_execution_id ) @@ -68,6 +81,11 @@ def read_comparison_simulation_by_identity( session: Session, record: ComparisonSimulationRecord, ) -> Stage12ComparisonSimulation | None: + """Return the simulation row with the record's identity, or None. + + The identity is the parent evaluation, role, input digest, worker version + and version-manifest digest. + """ return session.exec( select(Stage12ComparisonSimulation).where( Stage12ComparisonSimulation.evaluation_id == record.evaluation_id, diff --git a/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py b/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py index 91f4fa773..edfb0b31f 100644 --- a/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py +++ b/policyengine_api/services/v2/comparison_runs/database_connectors/updates.py @@ -60,6 +60,10 @@ def update_comparison_report( row: Stage12ComparisonReport, record: ComparisonReportRecord, ) -> Stage12ComparisonReport: + """Copy the record's statuses and mutable output fields onto the report row. + + The row is flushed and refreshed; committing is left to the caller. + """ row.status = Stage12RunStatus(record.status.value) row.aggregation_status = Stage12AggregationStatus(record.aggregation_status.value) for field_name in REPORT_MUTABLE_FIELDS: @@ -75,6 +79,10 @@ def update_comparison_simulation( row: Stage12ComparisonSimulation, record: ComparisonSimulationRecord, ) -> Stage12ComparisonSimulation: + """Copy the record's status, row identity and mutable output fields onto the row. + + The row is flushed and refreshed; committing is left to the caller. + """ row.status = Stage12RunStatus(record.status.value) row.row_identity_columns = ( list(record.row_identity_columns) @@ -94,6 +102,10 @@ def update_report_result_comparison( row: Stage12ComparisonReport, record: ComparisonReportRecord, ) -> Stage12ComparisonReport: + """Copy the record's result-comparison status and fields onto the report row. + + The row is flushed and refreshed; committing is left to the caller. + """ row.comparison_status = Stage12ResultComparisonStatus( record.comparison_status.value ) diff --git a/policyengine_api/services/v2/comparison_runs/database_session.py b/policyengine_api/services/v2/comparison_runs/database_session.py index e3d156915..04a1ddc53 100644 --- a/policyengine_api/services/v2/comparison_runs/database_session.py +++ b/policyengine_api/services/v2/comparison_runs/database_session.py @@ -13,15 +13,20 @@ class ComparisonRunDatabaseSession: + """Session lifetimes for comparison-run services: reads and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Open every session from session_factory.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads that is closed, not committed, on exit.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session in a transaction that commits on success, else rolls back.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/comparison_runs/services.py b/policyengine_api/services/v2/comparison_runs/services.py index abb22116a..5d94fad94 100644 --- a/policyengine_api/services/v2/comparison_runs/services.py +++ b/policyengine_api/services/v2/comparison_runs/services.py @@ -61,6 +61,11 @@ def create_or_resolve_report( session: Session, record: ComparisonReportRecord, ) -> ComparisonReportPersistenceResult: + """Insert a report row, or resolve the existing row with the same identity. + + Raise ComparisonRunIdentityError when an existing row disagrees with the + record on its release or data identity, so a retry cannot reinterpret it. + """ created_id = create_comparison_report(session, record) if created_id is not None: row = read_comparison_report(session, created_id, lock=False) @@ -86,6 +91,12 @@ def create_or_resolve_simulation( session: Session, record: ComparisonSimulationRecord, ) -> ComparisonSimulationPersistenceResult: + """Insert a simulation row under its report, or resolve the existing row. + + Raise ComparisonRunNotFoundError when the parent report is missing, and + ComparisonRunIdentityError when the record disagrees with its parent's + release and retention fields or with an existing row of the same identity. + """ parent_row = read_comparison_report(session, record.evaluation_id, lock=False) if parent_row is None: raise ComparisonRunNotFoundError( @@ -117,6 +128,12 @@ def replace_report_lifecycle( session: Session, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's new lifecycle, aggregation and output fields under a row lock. + + Identity and result-comparison fields must be unchanged and both status moves + allowed. A replay against a succeeded report must match it exactly, and then + returns the stored record without writing. + """ row = read_comparison_report(session, record.evaluation_id, lock=True) if row is None: raise ComparisonRunNotFoundError( @@ -140,6 +157,12 @@ def replace_report_result_comparison( session: Session, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's result-comparison outcome under a row lock. + + Only the comparison fields and updated_at may differ from the stored row. Once + the comparison is matched or different, an identical replay returns the + stored record without writing and any change raises. + """ row = read_comparison_report(session, record.evaluation_id, lock=True) if row is None: raise ComparisonRunNotFoundError( @@ -164,6 +187,12 @@ def replace_simulation_lifecycle( session: Session, record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: + """Write a simulation's new lifecycle and output fields under a row lock. + + Identity fields must be unchanged and the status move allowed. A replay + against a succeeded simulation must match it exactly, and then returns the + stored record without writing. + """ row = read_comparison_simulation( session, record.simulation_execution_id, @@ -183,13 +212,17 @@ def replace_simulation_lifecycle( class V2ComparisonRunService: + """Comparison-run operations, each run in its own database transaction.""" + def __init__(self, database_session: ComparisonRunDatabaseSession) -> None: + """Open a new session from database_session for every call.""" self._database_session = database_session def create_or_resolve_report( self, record: ComparisonReportRecord, ) -> ComparisonReportPersistenceResult: + """Insert a report row or resolve its existing identity, in one transaction.""" with self._database_session.transaction() as session: return create_or_resolve_report(session, record) @@ -197,6 +230,7 @@ def create_or_resolve_simulation( self, record: ComparisonSimulationRecord, ) -> ComparisonSimulationPersistenceResult: + """Insert a simulation row or resolve its identity, in one transaction.""" with self._database_session.transaction() as session: return create_or_resolve_simulation(session, record) @@ -204,6 +238,7 @@ def replace_report_lifecycle( self, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's lifecycle update in one transaction, under a row lock.""" with self._database_session.transaction() as session: return replace_report_lifecycle(session, record) @@ -211,6 +246,7 @@ def replace_simulation_lifecycle( self, record: ComparisonSimulationRecord, ) -> ComparisonSimulationRecord: + """Write a simulation lifecycle update in one transaction, under a row lock.""" with self._database_session.transaction() as session: return replace_simulation_lifecycle(session, record) @@ -218,10 +254,12 @@ def replace_report_result_comparison( self, record: ComparisonReportRecord, ) -> ComparisonReportRecord: + """Write a report's result-comparison outcome in one transaction.""" with self._database_session.transaction() as session: return replace_report_result_comparison(session, record) def get_report(self, evaluation_id: UUID) -> ComparisonReportRecord: + """Return the stored report, or raise ComparisonRunNotFoundError.""" with self._database_session.read() as session: row = read_comparison_report(session, evaluation_id, lock=False) if row is None: diff --git a/policyengine_api/services/v2/comparison_runs/transformations.py b/policyengine_api/services/v2/comparison_runs/transformations.py index 5c0c247d8..fc71b24d3 100644 --- a/policyengine_api/services/v2/comparison_runs/transformations.py +++ b/policyengine_api/services/v2/comparison_runs/transformations.py @@ -19,6 +19,7 @@ def report_insert_values(record: ComparisonReportRecord) -> dict[str, Any]: + """Return a report record's column values, with statuses as the table's enums.""" values = record.model_dump(mode="python") values["status"] = Stage12RunStatus(record.status.value) values["aggregation_status"] = Stage12AggregationStatus( @@ -31,6 +32,10 @@ def report_insert_values(record: ComparisonReportRecord) -> dict[str, Any]: def simulation_insert_values(record: ComparisonSimulationRecord) -> dict[str, Any]: + """Return a simulation record's column values in the table's types. + + Status and role become the table's enums, and row_identity_columns a list. + """ values = record.model_dump(mode="python") values["status"] = Stage12RunStatus(record.status.value) values["role"] = Stage12SimulationRole(record.role.value) @@ -43,6 +48,10 @@ def simulation_insert_values(record: ComparisonSimulationRecord) -> dict[str, An def report_record(row: Stage12ComparisonReport) -> ComparisonReportRecord: + """Rebuild the validated contract record from a stored report row. + + Raise pydantic.ValidationError if the row breaks the contract's state rules. + """ return ComparisonReportRecord.model_validate( { field_name: getattr(row, field_name) @@ -54,6 +63,10 @@ def report_record(row: Stage12ComparisonReport) -> ComparisonReportRecord: def simulation_record( row: Stage12ComparisonSimulation, ) -> ComparisonSimulationRecord: + """Rebuild the validated contract record from a stored simulation row. + + Raise pydantic.ValidationError if the row breaks the contract's state rules. + """ return ComparisonSimulationRecord.model_validate( { field_name: getattr(row, field_name) diff --git a/policyengine_api/services/v2/comparison_runs/types.py b/policyengine_api/services/v2/comparison_runs/types.py index 34be2e22f..00a127455 100644 --- a/policyengine_api/services/v2/comparison_runs/types.py +++ b/policyengine_api/services/v2/comparison_runs/types.py @@ -36,6 +36,11 @@ def _validate_retention(created_at: datetime, retention_expires_at: datetime) -> class ComparisonRunLifecycleStatus(StrEnum): + """Lifecycle of a comparison report or simulation. + + Succeeded and skipped are final; failed and incomplete runs may restart. + """ + PENDING = "pending" RUNNING = "running" SUCCEEDED = "succeeded" @@ -45,6 +50,8 @@ class ComparisonRunLifecycleStatus(StrEnum): class ComparisonRunAggregationStatus(StrEnum): + """Progress of a comparison report's aggregation step; succeeded is final.""" + NOT_STARTED = "not_started" RUNNING = "running" SUCCEEDED = "succeeded" @@ -52,6 +59,11 @@ class ComparisonRunAggregationStatus(StrEnum): class ResultComparisonStatus(StrEnum): + """Outcome of comparing a comparison run's aggregate with the production result. + + Not-requested, matched, different and failed are final. + """ + NOT_REQUESTED = "not_requested" PENDING = "pending" RUNNING = "running" @@ -106,6 +118,13 @@ class ComparisonReportRecord(StrictContractModel): @model_validator(mode="after") def validate_state(self) -> ComparisonReportRecord: + """Reject a report whose lifecycle, output and comparison fields disagree. + + The aggregate output fields are all set or all unset; a succeeded report + needs successful aggregation, its output and completed_at; a failed or + skipped one needs an error_code; and the comparison fields must fit + comparison_status. Retention must be timezone-aware and 1-30 days long. + """ _validate_retention(self.created_at, self.retention_expires_at) output_fields = ( self.aggregate_output_uri, @@ -208,6 +227,13 @@ class ComparisonSimulationRecord(StrictContractModel): @model_validator(mode="after") def validate_state(self) -> ComparisonSimulationRecord: + """Reject a simulation whose lifecycle and output fields disagree. + + The output artifact and row-identity fields are all set or all unset, + row_identity_columns is non-empty and unique, a succeeded simulation + needs its output and completed_at, and a failed or skipped one needs an + error_code. Retention must be timezone-aware and 1-30 days long. + """ _validate_retention(self.created_at, self.retention_expires_at) output_fields = ( self.output_uri, @@ -245,11 +271,15 @@ def validate_state(self) -> ComparisonSimulationRecord: @dataclass(frozen=True) class ComparisonReportPersistenceResult: + """A stored report record, and whether this call inserted it or found it.""" + record: ComparisonReportRecord created: bool @dataclass(frozen=True) class ComparisonSimulationPersistenceResult: + """A stored simulation record, and whether this call inserted it or found it.""" + record: ComparisonSimulationRecord created: bool diff --git a/policyengine_api/services/v2/comparison_runs/validators.py b/policyengine_api/services/v2/comparison_runs/validators.py index 56dd42d8b..45c104a9d 100644 --- a/policyengine_api/services/v2/comparison_runs/validators.py +++ b/policyengine_api/services/v2/comparison_runs/validators.py @@ -135,6 +135,7 @@ def require_report_identity( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if an update changes an immutable field.""" _require_equal_fields(existing, candidate, REPORT_IMMUTABLE_FIELDS) @@ -142,6 +143,11 @@ def require_report_conflict_matches( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if a duplicate insert names a different release. + + The stored report and the retried record must agree on every release, + package and dataset field. + """ _require_equal_fields(existing, candidate, REPORT_CONFLICT_FIELDS) @@ -149,6 +155,7 @@ def require_simulation_identity( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ComparisonRunIdentityError if an update changes an immutable field.""" _require_equal_fields(existing, candidate, SIMULATION_IMMUTABLE_FIELDS) @@ -156,6 +163,11 @@ def require_simulation_conflict_matches( existing: ComparisonSimulationRecord, candidate: ComparisonSimulationRecord, ) -> None: + """Raise ComparisonRunIdentityError if a duplicate insert names a different worker. + + The stored simulation and the retried record must agree on their parent, + role, input digest and worker release. + """ _require_equal_fields(existing, candidate, SIMULATION_CONFLICT_FIELDS) @@ -164,7 +176,6 @@ def require_simulation_parent_matches( child: ComparisonSimulationRecord, ) -> None: """Require a child to retain the parent's release and retention identity.""" - if child.evaluation_id != parent.evaluation_id: raise ComparisonRunIdentityError( "comparison child names a different parent evaluation_id" @@ -188,7 +199,6 @@ def require_successful_report_replay( candidate: ComparisonReportRecord, ) -> None: """Reject attempts to alter an already successful report receipt.""" - _require_equal_fields(existing, candidate, REPORT_SUCCEEDED_REPLAY_FIELDS) @@ -197,7 +207,6 @@ def require_successful_simulation_replay( candidate: ComparisonSimulationRecord, ) -> None: """Reject attempts to alter an already successful simulation receipt.""" - _require_equal_fields(existing, candidate, SIMULATION_SUCCEEDED_REPLAY_FIELDS) @@ -205,6 +214,10 @@ def require_comparison_update_only( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if anything but the comparison fields changes. + + updated_at may also differ; every other report field must match the stored row. + """ fields = tuple( field_name for field_name in ComparisonReportRecord.model_fields @@ -218,6 +231,7 @@ def require_result_comparison_unchanged( existing: ComparisonReportRecord, candidate: ComparisonReportRecord, ) -> None: + """Raise ComparisonRunIdentityError if the result-comparison fields differ.""" _require_equal_fields(existing, candidate, REPORT_RESULT_COMPARISON_FIELDS) @@ -319,6 +333,11 @@ def require_lifecycle_transition( current: ComparisonRunLifecycleStatus, candidate: ComparisonRunLifecycleStatus, ) -> None: + """Raise ComparisonRunStateTransitionError if lifecycle cannot move to candidate. + + Every status may repeat; succeeded and skipped are final, and failed or + incomplete runs may return to running. + """ if candidate not in ALLOWED_TRANSITIONS[current]: raise ComparisonRunStateTransitionError( f"comparison-run lifecycle cannot move from {current.value} to {candidate.value}" @@ -329,6 +348,10 @@ def require_aggregation_transition( current: ComparisonRunAggregationStatus, candidate: ComparisonRunAggregationStatus, ) -> None: + """Raise ComparisonRunStateTransitionError if aggregation cannot move to candidate. + + Every status may repeat; succeeded is final and failed may return to running. + """ if candidate not in ALLOWED_AGGREGATION_TRANSITIONS[current]: raise ComparisonRunStateTransitionError( "comparison-run aggregation cannot move " @@ -340,6 +363,10 @@ def require_result_comparison_transition( current: ResultComparisonStatus, candidate: ResultComparisonStatus, ) -> None: + """Raise ComparisonRunStateTransitionError if comparison cannot move to candidate. + + Not-requested, matched, different and failed are all final. + """ if candidate not in ALLOWED_RESULT_COMPARISON_TRANSITIONS[current]: raise ComparisonRunStateTransitionError( f"result comparison cannot move from {current.value} to {candidate.value}" diff --git a/policyengine_api/services/v2/households/database_connectors/creates.py b/policyengine_api/services/v2/households/database_connectors/creates.py index 72624119c..721a0096e 100644 --- a/policyengine_api/services/v2/households/database_connectors/creates.py +++ b/policyengine_api/services/v2/households/database_connectors/creates.py @@ -21,6 +21,11 @@ def create_household( canonicalization_version: int, content_hash: str, ) -> UUID | None: + """Insert a household unless one with the same canonical version and hash exists. + + Return the new id, or None when the unique content-hash constraint made the + insert a no-op; the caller then reads the existing row. + """ household_id = uuid4() statement = ( insert(Household) @@ -50,6 +55,11 @@ def create_legacy_household_mapping( fingerprint_version: int, fingerprint_sha256: str, ) -> UUID | None: + """Insert the mapping from a v1 household id to its v2 household. + + Return the new mapping id, or None when a mapping for this country and v1 id + already exists (for example, a concurrent copy won the race). + """ mapping_id = uuid4() statement = ( insert(LegacyHouseholdMapping) diff --git a/policyengine_api/services/v2/households/database_connectors/reads.py b/policyengine_api/services/v2/households/database_connectors/reads.py index 107a72ca8..8a1e6ad8c 100644 --- a/policyengine_api/services/v2/households/database_connectors/reads.py +++ b/policyengine_api/services/v2/households/database_connectors/reads.py @@ -15,6 +15,7 @@ def read_household_by_content_identity( session: Session, *, canonicalization_version: int, content_hash: str ) -> Household | None: + """Return the household with this canonical version and content hash, or None.""" return session.exec( select(Household).where( Household.canonicalization_version == canonicalization_version, @@ -26,6 +27,7 @@ def read_household_by_content_identity( def read_household_row( session: Session, *, country_id: str, household_id: UUID ) -> Household | None: + """Return the household with this id in this country, or None.""" return session.exec( select(Household).where( Household.id == household_id, @@ -42,6 +44,10 @@ def read_household_rows( offset: int, limit: int, ) -> list[Household]: + """Return up to ``limit + 1`` households in creation order, starting at ``offset``. + + The extra row lets the caller report whether another page exists. + """ statement = select(Household).where(Household.country_id == country_id) if default_year is not None: statement = statement.where(Household.default_year == default_year) @@ -61,6 +67,11 @@ def read_legacy_household_mapping( legacy_household_id: int, lock: bool, ) -> LegacyHouseholdMapping | None: + """Return the v2 mapping for a v1 household id, or None. + + With ``lock``, the row is selected FOR UPDATE so concurrent copies of the same + v1 household serialize on it. + """ statement = select(LegacyHouseholdMapping).where( LegacyHouseholdMapping.country_id == country_id, LegacyHouseholdMapping.legacy_household_id == legacy_household_id, diff --git a/policyengine_api/services/v2/households/database_session.py b/policyengine_api/services/v2/households/database_session.py index c0d580c90..6470744dd 100644 --- a/policyengine_api/services/v2/households/database_session.py +++ b/policyengine_api/services/v2/households/database_session.py @@ -13,15 +13,20 @@ class HouseholdDatabaseSession: + """Session lifetimes for household services: read-only sessions and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/households/services.py b/policyengine_api/services/v2/households/services.py index 0bfe957fa..c698ea951 100644 --- a/policyengine_api/services/v2/households/services.py +++ b/policyengine_api/services/v2/households/services.py @@ -58,6 +58,7 @@ def normalize_creation_input( household_input: HouseholdCreationInput, ) -> HouseholdCreationInput: + """Return the input with its household document validated and normalized.""" document = household_input.household_data assert isinstance(document, dict) return HouseholdCreationInput( @@ -78,6 +79,11 @@ def create_normalized_household( [HouseholdCreationInput], CanonicalHouseholdContent ] = canonicalize_household, ) -> HouseholdCreationResult: + """Store a normalized household, or return the existing one with identical content. + + Raise HouseholdCreationIntegrityError if the hash conflict finds no row, and + HouseholdContentHashCollisionError if the stored row's canonical bytes differ. + """ content = canonicalizer(household_input) created_id = create_household( session, @@ -108,6 +114,7 @@ def create_normalized_household( def read_complete_household( session: Session, *, country_id: str, household_id: UUID ) -> HouseholdRead: + """Return one household of this country, or raise HouseholdNotFoundError.""" household = read_household_row( session, country_id=country_id, @@ -126,6 +133,7 @@ def read_household_page( offset: int, limit: int, ) -> HouseholdPage: + """Return a page of households, optionally filtered by default year.""" return household_page( read_household_rows( session, @@ -143,6 +151,12 @@ def mirror_legacy_household_in_session( session: Session, snapshot: LegacyHouseholdSnapshot, ) -> LegacyHouseholdPersistenceResult: + """Copy one committed v1 household into v2 and map its v1 id to the result. + + Replaying the same v1 household is idempotent. Raise + LegacyHouseholdMappingIntegrityError if the v1 id is already mapped with a + different source fingerprint or to a different v2 household. + """ fingerprint = legacy_household_fingerprint(snapshot) existing = read_legacy_household_mapping( session, @@ -210,12 +224,19 @@ def mirror_legacy_household_in_session( class V2HouseholdService: + """Service for creating, reading and copying immutable v2 households.""" + def __init__(self, database_session: HouseholdDatabaseSession) -> None: + """Open every session through ``database_session``.""" self._database_session = database_session def create_household( self, household_input: NativeHouseholdCreationInput ) -> NativeHouseholdCreation: + """Store native household content, deduplicated by content hash. + + Return the stored household and whether this call created it. + """ normalized = normalize_creation_input( HouseholdCreationInput.model_validate(household_input.model_dump()) ) @@ -229,6 +250,7 @@ def create_household( return NativeHouseholdCreation(item=item, created=result.created) def get_household(self, *, country_id: str, household_id: UUID) -> HouseholdRead: + """Return one household of this country, or raise HouseholdNotFoundError.""" with self._database_session.read() as session: return read_complete_household( session, @@ -244,6 +266,7 @@ def list_households( offset: int, limit: int, ) -> HouseholdPage: + """Return a page of households, optionally filtered by default year.""" with self._database_session.read() as session: return read_household_page( session, @@ -256,5 +279,6 @@ def list_households( def mirror_legacy_household( self, snapshot: LegacyHouseholdSnapshot ) -> LegacyHouseholdPersistenceResult: + """Copy one committed v1 household into v2 in its own transaction.""" with self._database_session.transaction() as session: return mirror_legacy_household_in_session(session, snapshot) diff --git a/policyengine_api/services/v2/households/transformations.py b/policyengine_api/services/v2/households/transformations.py index def489d9d..a8cf8127f 100644 --- a/policyengine_api/services/v2/households/transformations.py +++ b/policyengine_api/services/v2/households/transformations.py @@ -72,6 +72,11 @@ def _canonical_json( def canonical_household_document(household_input: HouseholdCreationInput) -> bytes: + """Serialize a household to the canonical bytes that its content hash covers. + + People and entity records are sorted by id and keys are sorted, so the same + household always yields the same bytes regardless of input order. + """ content = dict(cast("dict[str, Any]", household_input.household_data)) for collection, records in content.items(): if collection in { @@ -96,6 +101,7 @@ def canonical_household_document(household_input: HouseholdCreationInput) -> byt def canonicalize_household( household_input: HouseholdCreationInput, ) -> CanonicalHouseholdContent: + """Return the canonical bytes, SHA-256 hash and canonicalization version.""" document = canonical_household_document(household_input) return CanonicalHouseholdContent( version=HOUSEHOLD_CANONICALIZATION_VERSION, @@ -105,6 +111,7 @@ def canonicalize_household( def stored_household_creation_input(household: Household) -> HouseholdCreationInput: + """Rebuild a stored row's creation input so its canonical bytes can be compared.""" return HouseholdCreationInput.model_validate( { "country_id": household.country_id, @@ -115,6 +122,7 @@ def stored_household_creation_input(household: Household) -> HouseholdCreationIn def household_read(household: Household) -> HouseholdRead: + """Convert a stored household row to the service's read model.""" return HouseholdRead( id=household.id, country_id=household.country_id, @@ -126,6 +134,10 @@ def household_read(household: Household) -> HouseholdRead: def household_page(rows: list[Household], *, offset: int, limit: int) -> HouseholdPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ displayed = rows[:limit] return HouseholdPage( items=tuple(household_read(row) for row in displayed), @@ -176,6 +188,13 @@ def _derive_default_year(values: list[object]) -> int | None: def translate_legacy_household( snapshot: LegacyHouseholdSnapshot, ) -> HouseholdCreationInput: + """Translate a v1 household document into normalized v2 household content. + + Named v1 people and groups get stable ids in name order, and group members + become person memberships. The default year is derived from the value periods. + Raise LegacyHouseholdTranslationError when the v1 document cannot be + translated exactly. + """ source = snapshot.household_json assert isinstance(source, dict) collections = SUPPORTED_ENTITY_COLLECTIONS[snapshot.country_id] @@ -277,6 +296,7 @@ def translate_legacy_household( def legacy_household_fingerprint(snapshot: LegacyHouseholdSnapshot) -> str: + """Return a v1 snapshot's SHA-256 fingerprint, which detects changed v1 content.""" document = { "fingerprint_version": LEGACY_HOUSEHOLD_FINGERPRINT_VERSION, **snapshot.model_dump(mode="json"), diff --git a/policyengine_api/services/v2/households/types.py b/policyengine_api/services/v2/households/types.py index dab71705a..23eefa1cb 100644 --- a/policyengine_api/services/v2/households/types.py +++ b/policyengine_api/services/v2/households/types.py @@ -40,6 +40,7 @@ class HouseholdCreationInput(StrictHouseholdInput): @field_validator("household_data") @classmethod def require_document(cls, value: object) -> object: + """Reject household_data that is not a JSON object.""" if type(value) is not dict: raise ValueError("household_data must be an object") return value @@ -64,6 +65,7 @@ class LegacyHouseholdSnapshot(StrictHouseholdInput): @field_validator("household_json") @classmethod def require_legacy_document(cls, value: object) -> object: + """Reject a v1 household_json that is not a JSON object.""" if type(value) is not dict: raise ValueError("legacy household_json must be an object") return value @@ -71,6 +73,12 @@ def require_legacy_document(cls, value: object) -> object: @dataclass(frozen=True) class CanonicalHouseholdContent: + """Canonical bytes of one household and the SHA-256 digest that deduplicates it. + + ``version`` is the canonicalization version; the same content hashed under + a different version is a different stored household. + """ + version: int document: bytes content_hash: str @@ -78,12 +86,20 @@ class CanonicalHouseholdContent: @dataclass(frozen=True) class HouseholdCreationResult: + """The stored household id, and whether this call inserted it or found it.""" + household_id: UUID created: bool @dataclass(frozen=True) class LegacyHouseholdPersistenceResult: + """Outcome of copying one v1 household into v2. + + ``household_id`` is the v2 household the v1 id maps to; the flags say whether + this call created the household and the v1-to-v2 mapping or found them. + """ + household_id: UUID household_created: bool mapping_created: bool @@ -91,6 +107,8 @@ class LegacyHouseholdPersistenceResult: @dataclass(frozen=True) class HouseholdRead: + """A stored v2 household as the service returns it to routes.""" + id: UUID country_id: str default_year: int | None @@ -101,6 +119,8 @@ class HouseholdRead: @dataclass(frozen=True) class HouseholdPage: + """One page of households in creation order; ``has_more`` says if more follow.""" + items: tuple[HouseholdRead, ...] offset: int limit: int @@ -109,5 +129,7 @@ class HouseholdPage: @dataclass(frozen=True) class NativeHouseholdCreation: + """A created-or-found household, with ``created`` choosing 201 or 200.""" + item: HouseholdRead created: bool diff --git a/policyengine_api/services/v2/households/validators.py b/policyengine_api/services/v2/households/validators.py index 4f3e80fec..072faa299 100644 --- a/policyengine_api/services/v2/households/validators.py +++ b/policyengine_api/services/v2/households/validators.py @@ -52,7 +52,6 @@ def require_json_value( containers: frozenset[int] = frozenset(), ) -> object: """Reject values that are not standards-compliant JSON.""" - if depth > MAXIMUM_JSON_NESTING: raise ValueError("JSON values must not exceed 100 nested containers") if value is None or type(value) in {str, bool, int}: @@ -109,7 +108,6 @@ def normalize_household_document( document: dict[str, Any], ) -> dict[str, Any]: """Validate and normalize one country-specific household document.""" - collections = SUPPORTED_ENTITY_COLLECTIONS.get(country_id) if collections is None: raise HouseholdValidationError(f"unsupported country_id {country_id!r}") @@ -208,7 +206,6 @@ def verify_legacy_household_mapping( expected_household_id: UUID | None = None, ) -> None: """Validate one stored immutable legacy mapping without executing SQL.""" - if ( mapping.fingerprint_version != fingerprint_version or mapping.fingerprint_sha256 != fingerprint_sha256 diff --git a/policyengine_api/services/v2/identity/database_connectors/creates.py b/policyengine_api/services/v2/identity/database_connectors/creates.py index 83d33e116..d7e8a26b7 100644 --- a/policyengine_api/services/v2/identity/database_connectors/creates.py +++ b/policyengine_api/services/v2/identity/database_connectors/creates.py @@ -15,6 +15,7 @@ def create_provisioned_user(session: Session, *, primary_country: str) -> User: + """Insert a standard user for a new sign-in identity and return it with its ID.""" user = User(primary_country=primary_country, user_type=UserType.STANDARD.value) session.add(user) session.flush() @@ -26,7 +27,6 @@ def create_unauthenticated_user(session: Session, *, user_id: UUID) -> bool: Returns whether this call inserted the row. """ - return ( session.execute( insert(User) @@ -48,7 +48,6 @@ def create_user_identity( seen_at: datetime, ) -> UUID | None: """Insert one identity, returning its user unless another request won.""" - return session.execute( insert(UserIdentity) .values( diff --git a/policyengine_api/services/v2/identity/database_connectors/deletes.py b/policyengine_api/services/v2/identity/database_connectors/deletes.py index 9a9c26068..6efe0d4e9 100644 --- a/policyengine_api/services/v2/identity/database_connectors/deletes.py +++ b/policyengine_api/services/v2/identity/database_connectors/deletes.py @@ -19,7 +19,6 @@ def delete_unauthenticated_user(session: Session, *, user_id: UUID) -> None: The statement matches only an unauthenticated row, so it never deletes a standard user, including one converted by a concurrent transaction. """ - session.execute( delete(User).where( col(User.id) == user_id, diff --git a/policyengine_api/services/v2/identity/database_connectors/reads.py b/policyengine_api/services/v2/identity/database_connectors/reads.py index fe5c332e5..e7ffab6f2 100644 --- a/policyengine_api/services/v2/identity/database_connectors/reads.py +++ b/policyengine_api/services/v2/identity/database_connectors/reads.py @@ -18,6 +18,7 @@ def read_user_identity( issuer: str, subject: str, ) -> UserIdentity | None: + """Return the identity for this issuer and subject, or None if never seen.""" return session.exec( select(UserIdentity).where( UserIdentity.issuer == issuer, @@ -38,7 +39,6 @@ def read_user( together but waits for a conversion holding ``FOR UPDATE``, then re-reads the row as that conversion committed it. """ - statement = select(User).where(User.id == user_id) if for_share: statement = statement.with_for_update(read=True) @@ -57,7 +57,6 @@ def read_unauthenticated_user_for_update( concurrent transaction is re-checked once that transaction commits, so it is no longer returned. """ - return session.exec( select(User) .where( diff --git a/policyengine_api/services/v2/identity/database_connectors/updates.py b/policyengine_api/services/v2/identity/database_connectors/updates.py index b23ff2271..64368c389 100644 --- a/policyengine_api/services/v2/identity/database_connectors/updates.py +++ b/policyengine_api/services/v2/identity/database_connectors/updates.py @@ -36,6 +36,7 @@ def update_user_identity_activity( email: str | None, seen_at: datetime, ) -> UserIdentity: + """Record the identity's current email and the time it was last seen.""" identity.email = email identity.last_seen_at = seen_at session.add(identity) @@ -49,6 +50,7 @@ def update_user_to_standard( *, primary_country: str, ) -> User: + """Make an unauthenticated user standard, keeping its ID and associations.""" user.user_type = UserType.STANDARD.value user.primary_country = primary_country session.add(user) @@ -68,7 +70,6 @@ def reassign_user_associations( the moved rows (a report's ``user_policy_id`` and ``baseline_user_policy_id``) stay valid. """ - for table in USER_ASSOCIATION_TABLES: session.execute( update(table) diff --git a/policyengine_api/services/v2/identity/database_session.py b/policyengine_api/services/v2/identity/database_session.py index 1f00c0ddd..0163c40ff 100644 --- a/policyengine_api/services/v2/identity/database_session.py +++ b/policyengine_api/services/v2/identity/database_session.py @@ -19,7 +19,6 @@ def provisioning_savepoint(session: Session) -> Iterator[SessionTransaction]: The caller rolls the yielded savepoint back when another request created the identity first; otherwise it is released when the block exits. """ - savepoint = session.begin_nested() try: yield savepoint diff --git a/policyengine_api/services/v2/identity/services.py b/policyengine_api/services/v2/identity/services.py index a96ad92f5..46d1cf1cd 100644 --- a/policyengine_api/services/v2/identity/services.py +++ b/policyengine_api/services/v2/identity/services.py @@ -91,9 +91,11 @@ class StandardAccountProvisioning: """Provision the standard user of a verified sign-in identity on any request.""" def __init__(self, principal: UserPrincipal) -> None: + """Act for the verified sign-in identity ``principal``.""" self._principal = principal def find_in_session(self, session: Session) -> ResolvedUser | None: + """Return the identity's standard user, or None before its first request.""" identity = read_user_identity( session, issuer=self._principal.issuer, @@ -110,6 +112,13 @@ def resolve_in_session( country_id: str, now: datetime | None = None, ) -> ResolvedUser: + """Return the identity's user, provisioning it with ``country_id`` if absent. + + A known identity's email and last-seen time are written at most hourly, + or when the email changes. When a concurrent first request wins the + ``UNIQUE(issuer, subject)`` race, its user is returned instead; raise + IdentityProvisioningError if that user's identity cannot be re-read. + """ principal = self._principal seen_at = now or utc_now() identity = read_user_identity( @@ -167,6 +176,7 @@ def resolve_for_read_in_session( *, country_id: str, ) -> ResolvedUser | None: + """Return the identity's user, provisioning it as on any request.""" # A signed-in account exists from its first request, read or write. return self.resolve_in_session(session, country_id=country_id) @@ -175,9 +185,14 @@ class UnauthenticatedAccountProvisioning: """Provision a signed-out browser's unauthenticated user on its first write.""" def __init__(self, principal: UnauthenticatedPrincipal) -> None: + """Act for the signed-out browser whose UUID ``principal`` carries.""" self._principal = principal def find_in_session(self, session: Session) -> ResolvedUser | None: + """Return the browser's user, or None if it has not saved anything yet. + + Raise StandardUserRequiresSignInError if the UUID names a standard user. + """ user = require_unauthenticated_user(read_user(session, self._principal.user_id)) return None if user is None else unauthenticated_resolved_user(user) @@ -188,6 +203,12 @@ def resolve_in_session( country_id: str, now: datetime | None = None, ) -> ResolvedUser: + """Return the browser's user, inserting it on the browser's first write. + + The user is held ``FOR SHARE`` until the transaction ends, and the insert + is retried once if the user is deleted before it is re-read. Raise + StandardUserRequiresSignInError if the UUID names a standard user. + """ # An unauthenticated user has no primary country or sign-in activity. del country_id, now user_id = self._principal.user_id @@ -221,6 +242,7 @@ def resolve_for_read_in_session( *, country_id: str, ) -> ResolvedUser | None: + """Return the browser's existing user, or None; reads never provision one.""" # Reads never provision: an absent user simply owns no records. del country_id return self.find_in_session(session) @@ -230,7 +252,6 @@ def account_provisioning_for( principal: AccountPrincipal, ) -> AccountProvisioningStrategy: """Select the provisioning strategy for a principal's type.""" - if isinstance(principal, UnauthenticatedPrincipal): return UnauthenticatedAccountProvisioning(principal) return StandardAccountProvisioning(principal) @@ -241,7 +262,6 @@ def find_principal_user_in_session( principal: AccountPrincipal, ) -> ResolvedUser | None: """Return the existing user for a principal without provisioning one.""" - return account_provisioning_for(principal).find_in_session(session) @@ -253,7 +273,6 @@ def resolve_principal_user_in_session( now: datetime | None = None, ) -> ResolvedUser: """Return the principal's v2 user, provisioning it on first use.""" - return account_provisioning_for(principal).resolve_in_session( session, country_id=country_id, now=now ) @@ -270,7 +289,6 @@ def resolve_principal_user_for_read_in_session( A standard principal is provisioned as on any request. An unauthenticated principal is not: ``None`` means it owns no records yet. """ - return account_provisioning_for(principal).resolve_for_read_in_session( session, country_id=country_id ) @@ -303,7 +321,6 @@ def convert_unauthenticated_user_in_session( ``UNIQUE(issuer, subject)``: when it wins, this conversion merges into its user instead. """ - seen_at = now or utc_now() standard_provisioning = StandardAccountProvisioning(principal) unauthenticated_user = read_unauthenticated_user_for_update( diff --git a/policyengine_api/services/v2/identity/transformations.py b/policyengine_api/services/v2/identity/transformations.py index 3ef5f6265..e5e602d6c 100644 --- a/policyengine_api/services/v2/identity/transformations.py +++ b/policyengine_api/services/v2/identity/transformations.py @@ -20,7 +20,6 @@ def _aware(value: datetime) -> datetime: def refreshed_identity_email(identity: UserIdentity, email: str | None) -> str | None: """Keep a stored verified email when the current token carries none.""" - return identity.email if email is None else email @@ -31,21 +30,21 @@ def identity_refresh_required( now: datetime, ) -> bool: """Write identity activity at most hourly, or when the email changes.""" - if refreshed_identity_email(identity, email) != identity.email: return True return _aware(now) - _aware(identity.last_seen_at) >= LAST_SEEN_REFRESH_INTERVAL def standard_resolved_user(user_id: UUID, *, email: str | None) -> ResolvedUser: + """Describe a signed-in identity's standard user and its verified email.""" return ResolvedUser(user_id=user_id, email=email, user_type=UserType.STANDARD) def unauthenticated_resolved_user(user: User) -> ResolvedUser: + """Describe a signed-out browser's user, which never has an email.""" return ResolvedUser(user_id=user.id, email=None, user_type=UserType.UNAUTHENTICATED) def converted_primary_country(user: User, *, country_id: str) -> str: - """A converted user keeps any primary country, or takes the request's.""" - + """Return the user's primary country, or the request's if it has none.""" return user.primary_country or country_id diff --git a/policyengine_api/services/v2/identity/types.py b/policyengine_api/services/v2/identity/types.py index 10a77341f..75132d614 100644 --- a/policyengine_api/services/v2/identity/types.py +++ b/policyengine_api/services/v2/identity/types.py @@ -81,6 +81,8 @@ def resolve_for_read_in_session( class AccountConversionOutcome(StrEnum): + """What signing in did with a browser's unauthenticated user.""" + # The unauthenticated user became the signed-in identity's user in place. CONVERTED = "converted" # Its records moved to the identity's existing user, and it was deleted. @@ -91,5 +93,7 @@ class AccountConversionOutcome(StrEnum): @dataclass(frozen=True) class AccountConversion: + """The signed-in identity's user after sign-in, and how it was reached.""" + owner: ResolvedUser outcome: AccountConversionOutcome diff --git a/policyengine_api/services/v2/identity/validators.py b/policyengine_api/services/v2/identity/validators.py index 9a875a7ee..e2e3921a4 100644 --- a/policyengine_api/services/v2/identity/validators.py +++ b/policyengine_api/services/v2/identity/validators.py @@ -18,6 +18,10 @@ class StandardUserRequiresSignInError(PermissionError): def require_concurrent_identity(identity: UserIdentity | None) -> UserIdentity: + """Return the identity a concurrent request stored, or raise if it is absent. + + Raise IdentityProvisioningError when a lost provisioning race finds no row. + """ if identity is None: raise IdentityProvisioningError( "user identity provisioning conflict did not resolve to a stored row" @@ -26,12 +30,12 @@ def require_concurrent_identity(identity: UserIdentity | None) -> UserIdentity: def is_unauthenticated_user(user: User) -> bool: + """Whether the user belongs to a signed-out browser rather than an identity.""" return user.user_type == UserType.UNAUTHENTICATED def require_unauthenticated_user(user: User | None) -> User | None: """Pass through an absent or unauthenticated user; refuse a standard one.""" - if user is not None and not is_unauthenticated_user(user): raise StandardUserRequiresSignInError( "This browser's saved work now belongs to a signed-in account; " @@ -41,6 +45,7 @@ def require_unauthenticated_user(user: User | None) -> User | None: def require_provisioned_user(user: User | None) -> User: + """Return the provisioned user, or raise IdentityProvisioningError if None.""" if user is None: raise IdentityProvisioningError( "unauthenticated user provisioning did not resolve to a stored row" diff --git a/policyengine_api/services/v2/legacy_refs/database_connectors/reads.py b/policyengine_api/services/v2/legacy_refs/database_connectors/reads.py index f2dbb8689..d9ca581f4 100644 --- a/policyengine_api/services/v2/legacy_refs/database_connectors/reads.py +++ b/policyengine_api/services/v2/legacy_refs/database_connectors/reads.py @@ -42,7 +42,6 @@ def read_existing_legacy_ids( legacy_ids: Collection[int], ) -> frozenset[int]: """Return the requested IDs that exist for the country, in bounded batches.""" - id_column, country_column = _identity_columns(kind) requested = sorted(set(legacy_ids)) existing: set[int] = set() diff --git a/policyengine_api/services/v2/legacy_refs/database_session.py b/policyengine_api/services/v2/legacy_refs/database_session.py index 404774fca..761abbd1d 100644 --- a/policyengine_api/services/v2/legacy_refs/database_session.py +++ b/policyengine_api/services/v2/legacy_refs/database_session.py @@ -15,10 +15,12 @@ class LegacyReferenceDatabaseSession: """Open one v1 transaction that is always rolled back and closed.""" def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory``, the v1 database's, to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a v1 session whose transaction is rolled back on exit.""" with self._session_factory() as session: transaction = session.begin() try: diff --git a/policyengine_api/services/v2/legacy_refs/services.py b/policyengine_api/services/v2/legacy_refs/services.py index e67670c46..4da194424 100644 --- a/policyengine_api/services/v2/legacy_refs/services.py +++ b/policyengine_api/services/v2/legacy_refs/services.py @@ -31,6 +31,7 @@ class V1LegacyReferenceService: """Batched, read-only v1 existence checks.""" def __init__(self, database_session: LegacyReferenceDatabaseSession) -> None: + """Open every v1 read through ``database_session``.""" self._database_session = database_session def existing_ids( @@ -40,7 +41,6 @@ def existing_ids( requested: Mapping[LegacyRecordKind, Collection[int]], ) -> dict[LegacyRecordKind, frozenset[int]]: """Return existing IDs per kind using one short v1 transaction.""" - selected = {kind: ids for kind, ids in requested.items() if ids} if not selected: return {kind: frozenset() for kind in requested} @@ -63,7 +63,6 @@ def require_existing( requested: Mapping[LegacyRecordKind, Collection[int]], ) -> None: """Raise ``LegacyReferenceNotFoundError`` for any missing reference.""" - existing = self.existing_ids(country_id=country_id, requested=requested) for kind, ids in requested.items(): require_legacy_references( @@ -82,6 +81,7 @@ class DeferredLegacyReferenceService: """ def __init__(self, build: Callable[[], LegacyReferenceChecker]) -> None: + """Call ``build`` at the first check and reuse its service afterwards.""" self._build = build self._service: LegacyReferenceChecker | None = None @@ -96,6 +96,7 @@ def existing_ids( country_id: str, requested: Mapping[LegacyRecordKind, Collection[int]], ) -> dict[LegacyRecordKind, frozenset[int]]: + """Return the existing IDs per kind from the built v1 reference service.""" return self._resolved().existing_ids(country_id=country_id, requested=requested) def require_existing( @@ -104,4 +105,5 @@ def require_existing( country_id: str, requested: Mapping[LegacyRecordKind, Collection[int]], ) -> None: + """Raise ``LegacyReferenceNotFoundError`` for any missing reference.""" self._resolved().require_existing(country_id=country_id, requested=requested) diff --git a/policyengine_api/services/v2/legacy_refs/types.py b/policyengine_api/services/v2/legacy_refs/types.py index 5793d0cdf..2c92b6fdb 100644 --- a/policyengine_api/services/v2/legacy_refs/types.py +++ b/policyengine_api/services/v2/legacy_refs/types.py @@ -21,6 +21,8 @@ class LegacyRecordKind(StrEnum): + """Kind of v1 artifact that a v2 user association can reference by ID.""" + POLICY = "policy" HOUSEHOLD = "household" SIMULATION = "simulation" diff --git a/policyengine_api/services/v2/legacy_refs/validators.py b/policyengine_api/services/v2/legacy_refs/validators.py index 8c0c8637d..7078dbbb4 100644 --- a/policyengine_api/services/v2/legacy_refs/validators.py +++ b/policyengine_api/services/v2/legacy_refs/validators.py @@ -21,6 +21,7 @@ def require_legacy_references( requested: Collection[int], existing: Collection[int], ) -> None: + """Raise LegacyReferenceNotFoundError naming every requested ID not in existing.""" missing = sorted(set(requested) - set(existing)) if missing: rendered = ", ".join(str(legacy_id) for legacy_id in missing) diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads.py b/policyengine_api/services/v2/metadata/database_connectors/reads.py index 48690c061..bd91984b2 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads.py @@ -23,6 +23,11 @@ def read_metadata_catalog( running_policyengine_version: str, policyengine_version: str | None, ) -> SelectedCatalog: + """Return the country's catalog for policyengine_version, else the running version. + + Raise the catalog_selection errors for an unsupported country, an + unpublished version or an unreachable database. + """ return select_catalog( session, country_id=country_id, @@ -35,6 +40,10 @@ def read_rows( session: Session, statement: Any, # noqa: ANN401 # Session.exec overloads omit core SQLAlchemy selects ) -> list[Any]: + """Execute statement and return every row. + + Raise MetadataCatalogUnavailableError when the database query fails. + """ try: return list(session.exec(statement).all()) except SQLAlchemyError as error: diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py b/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py index 27a197bae..cd5dbc276 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_datasets.py @@ -16,6 +16,10 @@ def read_datasets( session: Session, *, model_version_id: UUID, offset: int, limit: int ) -> list[Dataset]: + """Return a page of a model version's logical input datasets in name order. + + Fetch limit + 1 rows so the caller can tell whether another page exists. + """ return read_rows( session, select(Dataset) @@ -33,6 +37,7 @@ def read_datasets( def read_dataset( session: Session, *, model_version_id: UUID, dataset_id: UUID ) -> Dataset | None: + """Return the model version's logical input dataset with dataset_id, or None.""" rows = read_rows( session, select(Dataset).where( diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py b/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py index 230793152..d29390ef5 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_parameter_tree.py @@ -150,6 +150,13 @@ def read_parameter_children( offset: int, limit: int, ) -> list[Any]: + """Return a page of the direct children of parent_path in a model version's tree. + + Each row has the child's path, label, type (node or parameter), child count + and, for parameters, the parameter's summary columns. dialect selects SQLite + or PostgreSQL string functions; limit + 1 rows are fetched so the caller can + tell whether another page exists. + """ return read_rows( session, _parameter_children_statement( diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py b/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py index 848804c66..85fa4dc0d 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_parameters.py @@ -27,6 +27,12 @@ def read_parameters( limit: int, search: str | None, ) -> list[Parameter]: + """Return a page of a model version's parameters in name order. + + search matches name, label or description as a literal, case-insensitive + substring. Fetch limit + 1 rows so the caller can tell whether another page + exists. + """ statement = select(Parameter).where( col(Parameter.tax_benefit_model_version_id) == model_version_id ) @@ -48,6 +54,7 @@ def read_parameters( def read_parameter( session: Session, *, model_version_id: UUID, parameter_id: UUID ) -> Parameter | None: + """Return the model version's parameter with parameter_id, or None.""" rows = read_rows( session, select(Parameter).where( @@ -67,6 +74,13 @@ def read_parameter_values( offset: int, limit: int, ) -> list[ParameterValue]: + """Return a page of a model version's current-law parameter values. + + Values tied to a policy or dynamic are excluded. parameter_id keeps one + parameter's values and selected_day keeps those in effect on that day. Rows + are ordered by parameter name, then newest first, and limit + 1 are fetched + so the caller can tell whether another page exists. + """ statement = ( select(ParameterValue) .join(Parameter, col(Parameter.id) == col(ParameterValue.parameter_id)) @@ -101,6 +115,7 @@ def read_parameter_values( def read_parameter_value( session: Session, *, model_version_id: UUID, value_id: UUID ) -> ParameterValue | None: + """Return the model version's current-law value with value_id, or None.""" rows = read_rows( session, select(ParameterValue) diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py b/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py index 719f77a20..e7dacf1a9 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_regions.py @@ -21,6 +21,11 @@ def read_regions( offset: int | None = None, limit: int | None = None, ) -> list[Region]: + """Return a model version's regions in code order, optionally of one type. + + Without a limit every region is returned; with one, limit + 1 rows are + fetched so the caller can tell whether another page exists. + """ statement = select(Region).where( col(Region.tax_benefit_model_version_id) == model_version_id ) @@ -41,6 +46,7 @@ def read_region( region_id: UUID | None = None, region_code: str | None = None, ) -> Region | None: + """Return the model version's region matching the given id and code, or None.""" statement = select(Region).where( col(Region.tax_benefit_model_version_id) == model_version_id ) @@ -55,6 +61,10 @@ def read_region( def read_input_dataset( session: Session, *, model_version_id: UUID, dataset_id: UUID ) -> Dataset | None: + """Return the model version's logical input dataset with dataset_id, or None. + + Used to load a region's default dataset for the economy options. + """ rows = read_rows( session, select(Dataset).where( diff --git a/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py b/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py index f833c2257..4d84be396 100644 --- a/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py +++ b/policyengine_api/services/v2/metadata/database_connectors/reads_variables.py @@ -26,6 +26,12 @@ def read_variables( limit: int, search: str | None, ) -> list[Variable]: + """Return a page of a model version's variables in name order. + + search matches name, label or description as a literal, case-insensitive + substring. Fetch limit + 1 rows so the caller can tell whether another page + exists. + """ statement = select(Variable).where( col(Variable.tax_benefit_model_version_id) == model_version_id ) @@ -47,6 +53,7 @@ def read_variables( def read_variable( session: Session, *, model_version_id: UUID, variable_id: UUID ) -> Variable | None: + """Return the model version's variable with variable_id, or None.""" rows = read_rows( session, select(Variable).where( diff --git a/policyengine_api/services/v2/metadata/database_session.py b/policyengine_api/services/v2/metadata/database_session.py index e0e1fb278..207b7d07e 100644 --- a/policyengine_api/services/v2/metadata/database_session.py +++ b/policyengine_api/services/v2/metadata/database_session.py @@ -12,15 +12,19 @@ class MetadataDatabaseSession: """Own the request-scoped read session used by metadata services.""" def __init__(self, session: Session) -> None: + """Wrap one request's session; the service reads through it until close.""" self._session = session @property def session(self) -> Session: + """The open request session that every metadata query runs on.""" return self._session @property def dialect_name(self) -> str: + """The bound dialect's name; it picks SQLite or PostgreSQL string functions.""" return self._session.get_bind().dialect.name def close(self) -> None: + """Close the request's session.""" self._session.close() diff --git a/policyengine_api/services/v2/metadata/services.py b/policyengine_api/services/v2/metadata/services.py index 948743543..8082c6fad 100644 --- a/policyengine_api/services/v2/metadata/services.py +++ b/policyengine_api/services/v2/metadata/services.py @@ -77,7 +77,13 @@ class V2MetadataService: - """Sequence metadata reads through a request-scoped database session.""" + """Sequence metadata reads through a request-scoped database session. + + Each read selects the country's catalog for the requested PolicyEngine.py + version, or for the running version when none is given. List methods page + with offset and a limit of 1-500 and report has_more; get methods raise + MetadataResourceNotFoundError when the selected catalog lacks the item. + """ def __init__( self, @@ -85,12 +91,14 @@ def __init__( *, running_policyengine_version: str, ) -> None: + """Bind to a request's session and the running PolicyEngine.py version.""" self._database_session = database_session self._running_policyengine_version = validate_policyengine_version( running_policyengine_version ) def close(self) -> None: + """Close the request's database session.""" self._database_session.close() def _select_catalog( @@ -122,6 +130,7 @@ def list_models( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataModel]: + """Return the catalog's only model on the first page; later pages are empty.""" selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -134,6 +143,7 @@ def get_model( model_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataModel]: + """Return the selected catalog's model if its id is model_id.""" selected = self._select_catalog(country_id, policyengine_version) item = require_metadata_resource( metadata_model(selected) if selected.model.id == model_id else None, @@ -146,6 +156,7 @@ def get_model( def get_model_by_country( self, country_id: str, policyengine_version: str | None = None ) -> MetadataModelSelectionResult: + """Return the country's model together with the selected model version.""" return metadata_model_selection( self._select_catalog(country_id, policyengine_version) ) @@ -158,6 +169,7 @@ def list_model_versions( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataModelVersionDetail]: + """Return the selected version on the first page; later pages are empty.""" selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -170,6 +182,7 @@ def get_model_version( version_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataModelVersionDetail]: + """Return the selected model version if its id is version_id.""" selected = self._select_catalog(country_id, policyengine_version) item = require_metadata_resource( ( @@ -192,6 +205,11 @@ def list_variables( limit: int = 100, search: str | None = None, ) -> MetadataPageResult[MetadataVariable]: + """Return a page of the catalog's variables in name order. + + search keeps variables whose name, label or description contains it, + ignoring case. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -215,6 +233,7 @@ def get_variable( variable_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataVariable]: + """Return the variable with variable_id if it is in the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_variable( @@ -238,6 +257,11 @@ def list_parameters( limit: int = 100, search: str | None = None, ) -> MetadataPageResult[MetadataParameterSummary]: + """Return a page of the catalog's parameters in name order. + + search keeps parameters whose name, label or description contains it, + ignoring case. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -261,6 +285,7 @@ def get_parameter( parameter_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataParameterSummary]: + """Return the parameter with parameter_id if it is in the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_parameter( @@ -284,6 +309,11 @@ def list_parameter_children( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataParameterChild]: + """Return a page of the direct children of parent_path in the parameter tree. + + An empty parent_path lists the top-level children. A child is a parameter + when it is a leaf with a parameter row, and a node otherwise. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -313,6 +343,12 @@ def list_parameter_values( limit: int = 100, now: datetime | None = None, ) -> MetadataPageResult[MetadataCanonicalParameterValue]: + """Return a page of current-law parameter values, newest first per parameter. + + Values belonging to a reform or dynamic are never listed. parameter_id keeps + one parameter's values, and current=True keeps those in effect on the UTC + day of now (today by default). + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -338,6 +374,10 @@ def get_parameter_value( value_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataCanonicalParameterValue]: + """Return the current-law value with value_id if it is in the selected catalog. + + Values belonging to a reform or dynamic are never returned. + """ selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_parameter_value( @@ -360,6 +400,10 @@ def list_datasets( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataDataset]: + """Return a page of the catalog's logical input datasets in name order. + + Output datasets and datasets with a storage path are never listed. + """ selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -382,6 +426,7 @@ def get_dataset( dataset_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataDataset]: + """Return the logical input dataset with dataset_id in the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_dataset( @@ -405,6 +450,7 @@ def list_regions( offset: int = 0, limit: int = 100, ) -> MetadataPageResult[MetadataRegion]: + """Return a page of the catalog's regions by code, optionally of one type.""" selected = self._select_paginated_catalog( country_id, policyengine_version, offset=offset, limit=limit ) @@ -428,6 +474,7 @@ def get_region( region_id: UUID, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataRegion]: + """Return the region with region_id from the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_region( @@ -448,6 +495,7 @@ def get_region_by_code( region_code: str, policyengine_version: str | None = None, ) -> MetadataDetailResult[MetadataRegion]: + """Return the region with region_code from the selected catalog.""" selected = self._select_catalog(country_id, policyengine_version) row = require_metadata_resource( read_region( @@ -465,6 +513,11 @@ def get_region_by_code( def get_economy_options( self, country_id: str, policyengine_version: str | None = None ) -> MetadataEconomyOptionsResult: + """Return the region, year and dataset choices for an economy calculation. + + Raise MetadataCatalogUnavailableError when the catalog lacks the national + region, its default input dataset, the current-law id or the years. + """ selected = self._select_catalog(country_id, policyengine_version) regions = read_regions( self._database_session.session, diff --git a/policyengine_api/services/v2/metadata/transformations.py b/policyengine_api/services/v2/metadata/transformations.py index e17a9fd57..588869af4 100644 --- a/policyengine_api/services/v2/metadata/transformations.py +++ b/policyengine_api/services/v2/metadata/transformations.py @@ -44,6 +44,11 @@ def page_result( offset: int, limit: int, ) -> MetadataPageResult[ResourceT]: + """Wrap fetched rows as one page of the selected catalog. + + rows holds up to limit + 1 items; an extra item only sets has_more and is + dropped from the page. + """ return MetadataPageResult( policyengine_version=selected.policyengine_version, items=rows[:limit], @@ -54,6 +59,7 @@ def page_result( def metadata_model(selected: SelectedCatalog) -> MetadataModel: + """Describe the catalog's model, using the selected version's description.""" return MetadataModel( id=selected.model.id, name=selected.model.name, @@ -62,6 +68,7 @@ def metadata_model(selected: SelectedCatalog) -> MetadataModel: def metadata_model_version(selected: SelectedCatalog) -> MetadataModelVersionDetail: + """Describe the selected model version, with its current-law id and years.""" return MetadataModelVersionDetail( id=selected.model_version.id, model_id=selected.model.id, @@ -73,6 +80,7 @@ def metadata_model_version(selected: SelectedCatalog) -> MetadataModelVersionDet def metadata_model_selection(selected: SelectedCatalog) -> MetadataModelSelectionResult: + """Pair the catalog's model with its selected version.""" return MetadataModelSelectionResult( policyengine_version=selected.policyengine_version, model=metadata_model(selected), @@ -81,6 +89,7 @@ def metadata_model_selection(selected: SelectedCatalog) -> MetadataModelSelectio def metadata_variable(variable: Variable) -> MetadataVariable: + """Return the API representation of a stored variable.""" return MetadataVariable( id=variable.id, name=variable.name, @@ -96,6 +105,7 @@ def metadata_variable(variable: Variable) -> MetadataVariable: def metadata_parameter(parameter: Parameter) -> MetadataParameterSummary: + """Return the API representation of a stored parameter.""" return MetadataParameterSummary( id=parameter.id, name=parameter.name, @@ -107,6 +117,7 @@ def metadata_parameter(parameter: Parameter) -> MetadataParameterSummary: def metadata_parameter_value(value: ParameterValue) -> MetadataCanonicalParameterValue: + """Return the API representation of a stored value, exposing value_json as value.""" return MetadataCanonicalParameterValue( id=value.id, parameter_id=value.parameter_id, @@ -117,6 +128,7 @@ def metadata_parameter_value(value: ParameterValue) -> MetadataCanonicalParamete def utc_day_start(selected_time: datetime) -> datetime: + """Return midnight UTC on the day of selected_time, treating a naive time as UTC.""" if selected_time.tzinfo is None: selected_time = selected_time.replace(tzinfo=UTC) return selected_time.astimezone(UTC).replace( @@ -125,6 +137,7 @@ def utc_day_start(selected_time: datetime) -> datetime: def metadata_dataset(dataset: Dataset) -> MetadataDataset: + """Return the API representation of a stored logical input dataset.""" return MetadataDataset( id=dataset.id, name=dataset.name, @@ -134,6 +147,7 @@ def metadata_dataset(dataset: Dataset) -> MetadataDataset: def metadata_region(region: Region) -> MetadataRegion: + """Return the API representation of a stored region.""" return MetadataRegion( id=region.id, code=region.code, @@ -151,6 +165,11 @@ def metadata_region(region: Region) -> MetadataRegion: def metadata_parameter_children(rows: list[Any]) -> list[MetadataParameterChild]: + """Build parameter-tree children from read_parameter_children rows. + + Parameter leaves carry the parameter's summary, and a child with no label is + labelled with the last segment of its path. + """ items = [] for row in rows: parameter = None @@ -178,6 +197,10 @@ def metadata_parameter_children(rows: list[Any]) -> list[MetadataParameterChild] def metadata_economy_options( selected: SelectedCatalog, *, regions: list[Region], national_dataset: Dataset ) -> MetadataEconomyOptionsResult: + """Build economy options from every region, each year and the national dataset. + + The national dataset is the only dataset choice, and it is the default. + """ return MetadataEconomyOptionsResult( policyengine_version=selected.policyengine_version, current_law_id=selected.model_version.current_law_id, diff --git a/policyengine_api/services/v2/metadata/types.py b/policyengine_api/services/v2/metadata/types.py index fef5b7894..b74333291 100644 --- a/policyengine_api/services/v2/metadata/types.py +++ b/policyengine_api/services/v2/metadata/types.py @@ -11,10 +11,14 @@ class StrictResponseModel(BaseModel): + """Base for metadata results that rejects unknown fields and non-finite numbers.""" + model_config = ConfigDict(extra="forbid", allow_inf_nan=False) class MetadataRegionType(StrEnum): + """The geographic level of a region.""" + NATIONAL = "national" COUNTRY = "country" STATE = "state" @@ -26,12 +30,16 @@ class MetadataRegionType(StrEnum): class MetadataModel(StrictResponseModel): + """A country's tax-benefit model in the selected PolicyEngine.py catalog.""" + id: UUID name: str description: str | None class MetadataModelVersion(StrictResponseModel): + """One published version of a country's tax-benefit model.""" + id: UUID model_id: UUID version: str @@ -39,6 +47,12 @@ class MetadataModelVersion(StrictResponseModel): class MetadataVariable(StrictResponseModel): + """A variable the tax-benefit model computes or accepts as input. + + adds and subtracts, when present, list the variables summed into or + subtracted from this one. + """ + id: UUID name: str label: str | None @@ -52,6 +66,11 @@ class MetadataVariable(StrictResponseModel): class MetadataParameterSummary(StrictResponseModel): + """A policy parameter, described without its values. + + Its current-law values are listed by GET /v2/parameter-values. + """ + id: UUID name: str label: str | None @@ -61,6 +80,12 @@ class MetadataParameterSummary(StrictResponseModel): class MetadataCanonicalParameterValue(StrictResponseModel): + """One current-law value of a parameter. + + The value applies from start_date until end_date, or indefinitely when + end_date is null. + """ + id: UUID parameter_id: UUID value: JsonValue @@ -69,6 +94,12 @@ class MetadataCanonicalParameterValue(StrictResponseModel): class MetadataDataset(StrictResponseModel): + """A named input dataset available to the selected model version. + + These datasets are logical inputs, so storage_path is always null and + is_output_dataset is always false. + """ + id: UUID name: str description: str | None @@ -78,6 +109,12 @@ class MetadataDataset(StrictResponseModel): class MetadataRegion(StrictResponseModel): + """A region the model can simulate, such as a country, state or constituency. + + When requires_filter is true, the region is simulated by filtering its + default dataset on filter_field and filter_value. + """ + id: UUID code: str label: str @@ -93,28 +130,41 @@ class MetadataRegion(StrictResponseModel): class MetadataRegionOption(StrictResponseModel): + """A region that can be chosen for an economy calculation; name is its code.""" + name: str label: str type: MetadataRegionType class MetadataTimePeriodOption(StrictResponseModel): + """A year that can be chosen for an economy calculation.""" + name: int label: str class MetadataDatasetOption(StrictResponseModel): + """The national dataset, offered as the default for an economy calculation.""" + name: str label: str default: Literal[True] = True class MetadataModelVersionDetail(MetadataModelVersion): + """A model version, with its current-law policy id and the years offered.""" + current_law_id: int metadata_time_periods: list[int] class MetadataParameterChild(StrictResponseModel): + """A direct child in the parameter tree. + + A node has a child_count; a parameter leaf carries the parameter's summary. + """ + path: str label: str type: Literal["node", "parameter"] @@ -126,6 +176,8 @@ class MetadataParameterChild(StrictResponseModel): class MetadataPageResult(StrictResponseModel, Generic[ResourceT]): + """One page of a metadata collection; has_more says whether a later page exists.""" + policyengine_version: str items: list[ResourceT] offset: int @@ -134,17 +186,27 @@ class MetadataPageResult(StrictResponseModel, Generic[ResourceT]): class MetadataDetailResult(StrictResponseModel, Generic[ResourceT]): + """One metadata item, with the PolicyEngine.py version it was read from.""" + policyengine_version: str item: ResourceT class MetadataModelSelectionResult(StrictResponseModel): + """A country's model together with the model version selected for the request.""" + policyengine_version: str model: MetadataModel model_version: MetadataModelVersionDetail class MetadataEconomyOptionsResult(StrictResponseModel): + """The choices for configuring an economy calculation. + + It lists the selectable regions, years and datasets, and the id of the + current-law policy to use as a baseline. + """ + policyengine_version: str current_law_id: int region: list[MetadataRegionOption] diff --git a/policyengine_api/services/v2/metadata/validators.py b/policyengine_api/services/v2/metadata/validators.py index 6172d2018..036722af4 100644 --- a/policyengine_api/services/v2/metadata/validators.py +++ b/policyengine_api/services/v2/metadata/validators.py @@ -22,6 +22,10 @@ class InvalidMetadataPageError(ValueError): def validate_metadata_page(offset: int, limit: int) -> tuple[int, int]: + """Return offset and limit unchanged, or raise InvalidMetadataPageError. + + offset must be at least 0 and limit between 1 and 500. + """ if offset < 0: raise InvalidMetadataPageError("offset must be at least 0") if not 1 <= limit <= 500: @@ -35,6 +39,7 @@ def validate_metadata_page(offset: int, limit: int) -> tuple[int, int]: def require_metadata_resource( resource: ResourceT | None, *, description: str ) -> ResourceT: + """Return resource, or raise MetadataResourceNotFoundError naming description.""" if resource is None: raise MetadataResourceNotFoundError(f"{description} was not found") return resource @@ -48,7 +53,6 @@ def validate_economy_options( national_dataset: Dataset | None, ) -> Dataset: """Validate preloaded records required by the economy-options response.""" - if not any(region.code == country_id for region in regions): raise MetadataCatalogUnavailableError( f"the {country_id} national v2 region is absent" diff --git a/policyengine_api/services/v2/policies/database_connectors/creates.py b/policyengine_api/services/v2/policies/database_connectors/creates.py index ccebd0bb3..18e81b853 100644 --- a/policyengine_api/services/v2/policies/database_connectors/creates.py +++ b/policyengine_api/services/v2/policies/database_connectors/creates.py @@ -21,6 +21,12 @@ def create_policy( canonicalization_version: int, content_hash: str, ) -> UUID | None: + """Insert a policy unless one with the same canonical version and hash exists. + + Return the new id, or None when the unique content-hash constraint made the + insert a no-op; the caller then reads the existing row. Parameter values are + inserted separately by create_parameter_values. + """ policy_id = uuid4() statement = ( insert(Policy) @@ -44,6 +50,7 @@ def create_parameter_values( policy_id: UUID, policy_input: ResolvedPolicyCreationInput, ) -> None: + """Insert a new policy's parameter values and flush them so later reads see them.""" session.add_all( [ ParameterValue( @@ -68,6 +75,11 @@ def create_legacy_policy_mapping( source_policy_hash: str, policy_id: UUID, ) -> UUID | None: + """Insert the mapping from a v1 policy id to its v2 policy. + + Return the new mapping id, or None when a mapping for this country and v1 id + already exists (for example, a concurrent copy won the race). + """ return session.execute( insert(LegacyPolicyMapping) .values( diff --git a/policyengine_api/services/v2/policies/database_connectors/reads.py b/policyengine_api/services/v2/policies/database_connectors/reads.py index 244bb9eeb..6c649b384 100644 --- a/policyengine_api/services/v2/policies/database_connectors/reads.py +++ b/policyengine_api/services/v2/policies/database_connectors/reads.py @@ -30,6 +30,10 @@ def read_policy_catalog( policyengine_version: str | None = None, running_policyengine_version: str = POLICYENGINE_VERSION, ) -> SelectedCatalog: + """Select the country's metadata catalog for the requested PolicyEngine version. + + Without ``policyengine_version``, select the catalog of the running version. + """ return select_catalog( session, country_id=country_id, @@ -41,6 +45,7 @@ def read_policy_catalog( def read_version_parameter_ids( session: Session, *, model_version_id: UUID, requested_ids: set[UUID] ) -> set[UUID]: + """Return the requested parameter ids that exist in this model version.""" if not requested_ids: return set() return set( @@ -56,6 +61,7 @@ def read_version_parameter_ids( def read_parameters_by_name( session: Session, *, model_version_id: UUID, names: set[str] ) -> dict[str, Parameter]: + """Return this model version's parameters with the given names, keyed by name.""" if not names: return {} parameters = session.exec( @@ -70,6 +76,7 @@ def read_parameters_by_name( def read_policy_by_content_identity( session: Session, *, canonicalization_version: int, content_hash: str ) -> Policy | None: + """Return the policy with this canonical version and content hash, or None.""" return session.exec( select(Policy).where( Policy.canonicalization_version == canonicalization_version, @@ -81,10 +88,12 @@ def read_policy_by_content_identity( def read_model_version( session: Session, model_version_id: UUID ) -> TaxBenefitModelVersion | None: + """Return the model version with this id, or None.""" return session.get(TaxBenefitModelVersion, model_version_id) def read_parameter_values(session: Session, policy_id: UUID) -> list[ParameterValue]: + """Return every parameter value stored for one policy, in no particular order.""" return list( session.exec( select(ParameterValue).where(ParameterValue.policy_id == policy_id) @@ -95,6 +104,11 @@ def read_parameter_values(session: Session, policy_id: UUID) -> list[ParameterVa def read_legacy_policy_mapping( session: Session, *, country_id: str, legacy_policy_id: int, lock: bool ) -> LegacyPolicyMapping | None: + """Return the v2 mapping for a v1 policy id, or None. + + With ``lock``, the row is selected FOR UPDATE so concurrent copies of the same + v1 policy serialize on it. + """ statement = select(LegacyPolicyMapping).where( LegacyPolicyMapping.country_id == country_id, LegacyPolicyMapping.legacy_policy_id == legacy_policy_id, @@ -107,6 +121,7 @@ def read_legacy_policy_mapping( def read_policy_row( session: Session, *, country_id: str, policy_id: UUID ) -> Policy | None: + """Return the policy with this id in this country, or None.""" return session.exec( select(Policy).where(Policy.id == policy_id, Policy.country_id == country_id) ).one_or_none() @@ -120,6 +135,10 @@ def read_policy_rows( offset: int, limit: int, ) -> list[Policy]: + """Return up to ``limit + 1`` policies in creation order, starting at ``offset``. + + The extra row lets the caller report whether another page exists. + """ statement = select(Policy).where(Policy.country_id == country_id) if tax_benefit_model_id is not None: statement = statement.where(Policy.tax_benefit_model_id == tax_benefit_model_id) @@ -135,6 +154,7 @@ def read_policy_rows( def read_parameter_values_with_names( session: Session, policy_ids: list[UUID] ) -> list[tuple[ParameterValue, str]]: + """Return the policies' values with parameter names, by name then start date.""" if not policy_ids: return [] return list( diff --git a/policyengine_api/services/v2/policies/database_session.py b/policyengine_api/services/v2/policies/database_session.py index af57e150e..dbc02bdab 100644 --- a/policyengine_api/services/v2/policies/database_session.py +++ b/policyengine_api/services/v2/policies/database_session.py @@ -16,14 +16,17 @@ class PolicyDatabaseSession: """Open read sessions and transaction-scoped sessions for policy services.""" def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/policies/services.py b/policyengine_api/services/v2/policies/services.py index 0fcb68c23..9a8e0d3b2 100644 --- a/policyengine_api/services/v2/policies/services.py +++ b/policyengine_api/services/v2/policies/services.py @@ -70,6 +70,11 @@ def resolve_policy_creation_input( policyengine_version: str | None = None, running_policyengine_version: str = POLICYENGINE_VERSION, ) -> ResolvedPolicyCreationInput: + """Bind policy content to the selected catalog's model version. + + Raise PolicyCatalogValidationError if the model or any parameter id is not + part of the selected catalog. + """ selected = read_policy_catalog( session, policy_input.country_id, @@ -99,6 +104,11 @@ def create_resolved_policy( [ResolvedPolicyCreationInput], CanonicalPolicyContent ] = canonicalize_policy, ) -> PolicyCreationResult: + """Store a resolved policy, or return the existing one with identical content. + + Raise PolicyCreationIntegrityError if the hash conflict finds no usable row, + and PolicyContentHashCollisionError if the stored row's canonical bytes differ. + """ content = canonicalizer(policy_input) created_policy_id = create_policy( session, @@ -143,6 +153,7 @@ def create_resolved_policy( def read_complete_policy( session: Session, *, country_id: str, policy_id: UUID ) -> PolicyRead: + """Return one policy of this country, or raise PolicyNotFoundError.""" policy = read_policy_row(session, country_id=country_id, policy_id=policy_id) if policy is None: raise PolicyNotFoundError(f"policy {policy_id} was not found") @@ -160,6 +171,7 @@ def read_policy_page( offset: int, limit: int, ) -> PolicyPage: + """Return a page of policies with their values, optionally for one model.""" rows = read_policy_rows( session, country_id=country_id, @@ -182,6 +194,12 @@ def mirror_legacy_policy_in_session( running_policyengine_version: str = POLICYENGINE_VERSION, country_package_versions: Mapping[str, str] = COUNTRY_PACKAGE_VERSIONS, ) -> LegacyPolicyPersistenceResult: + """Copy one committed v1 policy into v2 and map its v1 id to the result. + + The policy is translated against the running catalog. Replaying the same v1 + policy is idempotent. Raise LegacyPolicyMappingIntegrityError if the v1 id is + already mapped with a different source hash or to a different v2 policy. + """ existing = read_legacy_policy_mapping( session, country_id=snapshot.country_id, @@ -268,12 +286,20 @@ def __init__( *, running_policyengine_version: str = POLICYENGINE_VERSION, ) -> None: + """Open sessions through ``database_session``. + + ``running_policyengine_version`` selects the catalog when a request names none. + """ self._database_session = database_session self._running_policyengine_version = running_policyengine_version def create_policy( self, policy_input: NativePolicyCreationInput ) -> NativePolicyCreation: + """Store native policy content, deduplicated by content hash. + + Return the stored policy and whether this call created it. + """ content = PolicyCreationInput.model_validate( policy_input.model_dump(exclude={"policyengine_version"}) ) @@ -293,6 +319,7 @@ def create_policy( return NativePolicyCreation(item=item, created=persisted.created) def get_policy(self, *, country_id: str, policy_id: UUID) -> PolicyRead: + """Return one policy of this country, or raise PolicyNotFoundError.""" with self._database_session.read() as session: return read_complete_policy( session, country_id=country_id, policy_id=policy_id @@ -306,6 +333,7 @@ def list_policies( offset: int = 0, limit: int = 100, ) -> PolicyPage: + """Return a page of policies, optionally for one tax-benefit model.""" with self._database_session.read() as session: return read_policy_page( session, @@ -318,6 +346,7 @@ def list_policies( def mirror_legacy_policy( self, snapshot: LegacyPolicySnapshot ) -> LegacyPolicyPersistenceResult: + """Copy one committed v1 policy into v2 in its own transaction.""" with self._database_session.transaction() as session: return mirror_legacy_policy_in_session( session, diff --git a/policyengine_api/services/v2/policies/transformations.py b/policyengine_api/services/v2/policies/transformations.py index f80e04390..ef124f337 100644 --- a/policyengine_api/services/v2/policies/transformations.py +++ b/policyengine_api/services/v2/policies/transformations.py @@ -75,11 +75,17 @@ def _canonical_json( def canonical_utc_datetime(value: datetime) -> str: + """Render a datetime in UTC with microseconds and a ``Z`` suffix, as hashed.""" utc_value = value.astimezone(UTC) return utc_value.isoformat(timespec="microseconds").replace("+00:00", "Z") def canonical_policy_document(policy_input: ResolvedPolicyCreationInput) -> bytes: + """Serialize a resolved policy to the canonical bytes that its content hash covers. + + Values are sorted by parameter id and dates, and keys are sorted, so the same + policy always yields the same bytes regardless of input order. + """ parameter_values = sorted( policy_input.parameter_values, key=lambda value: ( @@ -113,6 +119,7 @@ def canonical_policy_document(policy_input: ResolvedPolicyCreationInput) -> byte def canonicalize_policy( policy_input: ResolvedPolicyCreationInput, ) -> CanonicalPolicyContent: + """Return the canonical bytes, SHA-256 hash and canonicalization version.""" document = canonical_policy_document(policy_input) return CanonicalPolicyContent( version=POLICY_CANONICALIZATION_VERSION, @@ -126,6 +133,7 @@ def stored_policy_creation_input( model_version: TaxBenefitModelVersion, values: list[ParameterValue], ) -> ResolvedPolicyCreationInput: + """Rebuild a stored policy's input so its canonical bytes can be compared.""" return ResolvedPolicyCreationInput.model_validate( { "country_id": policy.country_id, @@ -149,6 +157,7 @@ def policy_parameter_values_by_policy( policy_ids: list[UUID], rows: list[tuple[ParameterValue, str]], ) -> dict[UUID, tuple[PolicyParameterValueRead, ...]]: + """Group named value rows by policy id; every requested policy gets an entry.""" grouped: dict[UUID, list[PolicyParameterValueRead]] = { policy_id: [] for policy_id in policy_ids } @@ -172,6 +181,7 @@ def policy_read( policy: Policy, values: dict[UUID, tuple[PolicyParameterValueRead, ...]], ) -> PolicyRead: + """Convert a stored policy and its grouped values to the service's read model.""" return PolicyRead( id=policy.id, country_id=policy.country_id, @@ -190,6 +200,10 @@ def policy_page( offset: int, limit: int, ) -> PolicyPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ policies = rows[:limit] return PolicyPage( items=tuple(policy_read(policy, values) for policy in policies), @@ -210,6 +224,11 @@ def _utc_midnight(value: str) -> datetime: def parse_legacy_period(value: str) -> tuple[datetime, datetime]: + """Return the UTC-midnight start and end dates of a v1 period key. + + Accept a ``start.end`` ISO date range or a PolicyEngine period such as + ``2024``. Raise LegacyPolicyTranslationError for anything else. + """ if not value or value != value.strip(): raise LegacyPolicyTranslationError("legacy period must be non-empty") if "." in value: @@ -242,6 +261,12 @@ def translate_legacy_policy( parameters: Mapping[str, Parameter], country_package_versions: Mapping[str, str], ) -> ResolvedPolicyCreationInput: + """Translate a v1 policy into policy content bound to the running catalog. + + Raise LegacyPolicyTranslationError if the v1 policy was saved under another + country package version, names a parameter the catalog lacks, or has + invalid or conflicting periods. + """ expected_api_version = country_package_versions.get(snapshot.country_id) if expected_api_version is None or snapshot.api_version != expected_api_version: raise LegacyPolicyTranslationError( diff --git a/policyengine_api/services/v2/policies/types.py b/policyengine_api/services/v2/policies/types.py index ef1f2823a..fcf1647c3 100644 --- a/policyengine_api/services/v2/policies/types.py +++ b/policyengine_api/services/v2/policies/types.py @@ -45,6 +45,7 @@ class PolicyParameterValueInput(StrictPolicyInput): @model_validator(mode="after") def validate_period(self) -> PolicyParameterValueInput: + """Reject a value whose end date precedes its start date.""" if self.end_date is not None and self.end_date < self.start_date: raise ValueError("end_date must not precede start_date") return self @@ -66,6 +67,7 @@ def reject_duplicate_effective_values( cls, values: list[PolicyParameterValueInput], ) -> list[PolicyParameterValueInput]: + """Reject two values for the same parameter with the same start date.""" identities: set[tuple[UUID, datetime]] = set() for value in values: identity = (value.parameter_id, value.start_date) @@ -103,6 +105,7 @@ class LegacyPolicySnapshot(StrictPolicyInput): @field_validator("policy_json") @classmethod def require_parameter_mapping(cls, value: object) -> object: + """Reject a v1 policy_json that is not an object keyed by parameter path.""" if type(value) is not dict: raise ValueError("legacy policy_json must be a parameter-path object") return value @@ -136,6 +139,8 @@ class LegacyPolicyPersistenceResult: @dataclass(frozen=True) class PolicyParameterValueRead: + """One stored parameter value of a policy, with the parameter's name.""" + id: UUID parameter_id: UUID parameter_name: str @@ -146,6 +151,8 @@ class PolicyParameterValueRead: @dataclass(frozen=True) class PolicyRead: + """A stored v2 policy with its parameter values, as returned to routes.""" + id: UUID country_id: str tax_benefit_model_id: UUID @@ -157,6 +164,8 @@ class PolicyRead: @dataclass(frozen=True) class PolicyPage: + """One page of policies in creation order; ``has_more`` says if more follow.""" + items: tuple[PolicyRead, ...] offset: int limit: int diff --git a/policyengine_api/services/v2/policies/validators.py b/policyengine_api/services/v2/policies/validators.py index 0ceabb046..d0ecf8120 100644 --- a/policyengine_api/services/v2/policies/validators.py +++ b/policyengine_api/services/v2/policies/validators.py @@ -52,7 +52,6 @@ def require_json_value( containers: frozenset[int] = frozenset(), ) -> object: """Reject values that cannot be represented as standards-compliant JSON.""" - if depth > MAXIMUM_JSON_NESTING: raise ValueError("JSON values must not exceed 100 nested containers") if value is None or type(value) in {str, bool, int}: @@ -80,7 +79,6 @@ def require_json_value( def normalize_utc(value: datetime) -> datetime: """Require a time-zone-aware value and normalize it to UTC.""" - if value.tzinfo is None or value.utcoffset() is None: raise ValueError("effective dates must include a UTC offset") return value.astimezone(UTC) @@ -93,7 +91,6 @@ def validate_policy_catalog( resolved_parameter_ids: set[UUID], ) -> ResolvedPolicyCreationInput: """Validate preloaded catalog records and bind them to policy content.""" - from policyengine_api.services.v2.policies.types import ResolvedPolicyCreationInput if policy_input.tax_benefit_model_id != selected.model.id: @@ -123,7 +120,6 @@ def verify_legacy_policy_mapping( expected_policy_id: UUID | None = None, ) -> None: """Validate an existing mapping without performing SQL.""" - if mapping.source_policy_hash != source_policy_hash: raise LegacyPolicyMappingIntegrityError( "legacy policy identity was presented with a different source hash" diff --git a/policyengine_api/services/v2/reports/types.py b/policyengine_api/services/v2/reports/types.py index ef4d997f4..cd6d8d2b1 100644 --- a/policyengine_api/services/v2/reports/types.py +++ b/policyengine_api/services/v2/reports/types.py @@ -22,6 +22,8 @@ class ReportAggregate(StrEnum): + """An aggregate result section a report can request.""" + BUDGET = "budget" POVERTY = "poverty" INEQUALITY = "inequality" @@ -32,6 +34,12 @@ class ReportAggregate(StrEnum): class ReportExecutionInput(StrictContractModel): + """The complete input for one report run. + + It pairs a baseline and a reform simulation that differ only in policy, and + lists the aggregates to compute from them. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID baseline: SimulationExecutionInput @@ -43,6 +51,13 @@ class ReportExecutionInput(StrictContractModel): @model_validator(mode="after") def require_aligned_simulations(self) -> ReportExecutionInput: + """Reject a report whose two simulations cannot be compared. + + Baseline and reform must carry their own roles and this report's + evaluation_id, have distinct execution ids, and share population, year, + geography, options, requested output and bundle. Requested aggregates + must be unique. + """ if self.baseline.role is not SimulationRole.BASELINE: raise ValueError("baseline input must use the baseline role") if self.reform.role is not SimulationRole.REFORM: @@ -72,6 +87,12 @@ def require_aligned_simulations(self) -> ReportExecutionInput: class AggregateReportArtifactDescriptor(StrictContractModel): + """The receipt for one stored aggregate report. + + It names the artifact and the bundle, and records the digests of the + baseline and reform simulation artifacts it was computed from. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID artifact: ArtifactReference @@ -100,6 +121,7 @@ def require_unique_aggregates( cls, value: tuple[ReportAggregate, ...], ) -> tuple[ReportAggregate, ...]: + """Reject a payload that lists any aggregate twice.""" if len(value) != len(set(value)): raise ValueError("requested aggregates must be unique") return value @@ -107,7 +129,6 @@ def require_unique_aggregates( @model_validator(mode="after") def validate_spm_provenance(self) -> AggregateReportArtifactPayload: """Reject duplicate or noncanonical SPM data in aggregate artifacts.""" - if "spm_config" in self.result: raise ValueError("aggregate result must not include legacy spm_config") if "spm_provenance" in self.result: diff --git a/policyengine_api/services/v2/simulations/types.py b/policyengine_api/services/v2/simulations/types.py index d1e466b13..7a19b9793 100644 --- a/policyengine_api/services/v2/simulations/types.py +++ b/policyengine_api/services/v2/simulations/types.py @@ -46,12 +46,16 @@ class StrictContractModel(BaseModel): class SimulationRole(StrEnum): + """The part a simulation plays: one side of a comparison, or a standalone run.""" + BASELINE = "baseline" REFORM = "reform" STANDALONE = "standalone" class ArtifactMediaType(StrEnum): + """Media type of a worker output: JSON for reports, Parquet for simulations.""" + JSON = "application/json" PARQUET = "application/vnd.apache.parquet" @@ -90,12 +94,16 @@ class SimulationParquetPayloadContract(StrictContractModel): class DatasetArtifactMediaType(StrEnum): + """Media type of an input dataset artifact.""" + HDF5 = "application/x-hdf5" JSON = "application/json" PARQUET = "application/vnd.apache.parquet" class DatasetProvenance(StrictContractModel): + """The input dataset a simulation runs on, with its revision and data package.""" + identity: ContractText uri: Annotated[str, Field(min_length=1, max_length=2048)] artifact_revision: ContractText @@ -104,6 +112,12 @@ class DatasetProvenance(StrictContractModel): class BundleProvenance(StrictContractModel): + """The PolicyEngine.py bundle a simulation runs against. + + It pins the package versions, the input dataset and the bundle manifest + digest, so a result can be traced to an exact release. + """ + policyengine_version: ContractText country_package_name: ContractText country_package_version: ContractText @@ -112,6 +126,8 @@ class BundleProvenance(StrictContractModel): class ArtifactReference(StrictContractModel): + """A private gs:// object, addressed by its SHA-256 digest and size.""" + uri: StorageUri media_type: ArtifactMediaType content_sha256: Sha256Digest @@ -119,6 +135,11 @@ class ArtifactReference(StrictContractModel): class DatasetArtifactReference(StrictContractModel): + """An input dataset object in GCS or Hugging Face, addressed by its SHA-256 digest. + + size_bytes is optional because a dataset's size may not be known. + """ + uri: DatasetArtifactUri media_type: DatasetArtifactMediaType content_sha256: Sha256Digest @@ -126,11 +147,15 @@ class DatasetArtifactReference(StrictContractModel): class DatasetPopulationInput(StrictContractModel): + """A population supplied as an input dataset artifact.""" + kind: Literal["dataset"] = "dataset" artifact: DatasetArtifactReference class HouseholdPopulationInput(StrictContractModel): + """A population supplied as one household document.""" + kind: Literal["household"] = "household" document: dict[str, JsonValue] @@ -142,6 +167,12 @@ class HouseholdPopulationInput(StrictContractModel): class GeographySelection(StrictContractModel): + """The country and region to simulate. + + filter_field and filter_value, set together, restrict the population to the + region; filter_strategy, if set, names how the filter is applied. + """ + country: CountryId region: ContractText filter_field: ContractText | None = None @@ -150,6 +181,7 @@ class GeographySelection(StrictContractModel): @model_validator(mode="after") def require_complete_filter(self) -> GeographySelection: + """Require filter_field with filter_value, and both for filter_strategy.""" if (self.filter_field is None) != (self.filter_value is None): raise ValueError("filter_field and filter_value must be supplied together") if self.filter_strategy is not None and self.filter_field is None: @@ -158,18 +190,27 @@ def require_complete_filter(self) -> GeographySelection: class RequestedSimulationOutput(StrictContractModel): + """The variables the worker writes to a simulation's output, each named once.""" + schema_version: ContractVersion = 1 variables: Annotated[tuple[ContractText, ...], Field(min_length=1)] @field_validator("variables") @classmethod def require_unique_variables(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Reject a variable list that names any variable twice.""" if len(value) != len(set(value)): raise ValueError("requested variables must be unique") return value class SimulationExecutionInput(StrictContractModel): + """The complete input for one simulation in a comparison run. + + It fixes the policy, population, year, geography, options, requested + outputs and the bundle to run against. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID simulation_execution_id: UUID @@ -184,6 +225,8 @@ class SimulationExecutionInput(StrictContractModel): class RowIdentity(StrictContractModel): + """The identifier columns, row count and digest that fix the rows of an output.""" + schema_version: ContractVersion = 1 identifier_columns: Annotated[tuple[ContractText, ...], Field(min_length=1)] row_count: Annotated[int, Field(ge=0)] @@ -192,12 +235,19 @@ class RowIdentity(StrictContractModel): @field_validator("identifier_columns") @classmethod def require_unique_columns(cls, value: tuple[str, ...]) -> tuple[str, ...]: + """Reject identifier columns that name any column twice.""" if len(value) != len(set(value)): raise ValueError("row identifier columns must be unique") return value class SimulationArtifactDescriptor(StrictContractModel): + """The receipt for one stored simulation output. + + It names the artifact, its row identity and bundle, and any detached SPM + calculation receipt needed to reproduce an aggregate from it. + """ + contract_version: ContractVersion = 1 evaluation_id: UUID simulation_execution_id: UUID diff --git a/policyengine_api/services/v2/user_accounts/database_connectors/deletes.py b/policyengine_api/services/v2/user_accounts/database_connectors/deletes.py index 102ba2c96..31a1a7ef5 100644 --- a/policyengine_api/services/v2/user_accounts/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_accounts/database_connectors/deletes.py @@ -18,6 +18,7 @@ def delete_account_legacy_user_mappings(session: Session, *, user_id: UUID) -> None: + """Delete the v1 user mappings that would otherwise block deleting the user.""" session.execute( delete(LegacyUserMapping).where(col(LegacyUserMapping.user_id) == user_id) ) @@ -25,5 +26,6 @@ def delete_account_legacy_user_mappings(session: Session, *, user_id: UUID) -> N def delete_account_user(session: Session, *, user_id: UUID) -> None: + """Delete the user; its identities and every association cascade with it.""" session.execute(delete(User).where(col(User.id) == user_id)) session.flush() diff --git a/policyengine_api/services/v2/user_accounts/database_connectors/reads.py b/policyengine_api/services/v2/user_accounts/database_connectors/reads.py index 8d89d0d8d..dbff75726 100644 --- a/policyengine_api/services/v2/user_accounts/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_accounts/database_connectors/reads.py @@ -13,6 +13,7 @@ def read_account_user(session: Session, user_id: UUID) -> User | None: + """Return the v2 user with this id, or None.""" return session.get(User, user_id) @@ -22,6 +23,7 @@ def read_identity_subjects( user_id: UUID, issuer: str, ) -> list[str]: + """Return the user's sign-in subjects from ``issuer``, oldest identity first.""" return list( session.exec( select(UserIdentity.subject) diff --git a/policyengine_api/services/v2/user_accounts/database_connectors/updates.py b/policyengine_api/services/v2/user_accounts/database_connectors/updates.py index 545899bad..cd6a0c99f 100644 --- a/policyengine_api/services/v2/user_accounts/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_accounts/database_connectors/updates.py @@ -17,6 +17,7 @@ def update_account_user( user: User, values: Mapping[str, Any], ) -> User: + """Set only the given fields on the user and return it reloaded.""" for field_name, value in values.items(): setattr(user, field_name, value) session.add(user) diff --git a/policyengine_api/services/v2/user_accounts/database_session.py b/policyengine_api/services/v2/user_accounts/database_session.py index b4ecb8920..e312415ef 100644 --- a/policyengine_api/services/v2/user_accounts/database_session.py +++ b/policyengine_api/services/v2/user_accounts/database_session.py @@ -13,15 +13,20 @@ class AccountDatabaseSession: + """Session lifetimes for account services: reads and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_accounts/services.py b/policyengine_api/services/v2/user_accounts/services.py index 77c4cd11c..50f2d86cb 100644 --- a/policyengine_api/services/v2/user_accounts/services.py +++ b/policyengine_api/services/v2/user_accounts/services.py @@ -97,6 +97,11 @@ def __init__( *, auth0_user_deleter_factory: Callable[[], Auth0UserDeleter | None] | None = None, ) -> None: + """Open every session through ``database_session``. + + ``auth0_user_deleter_factory`` is called only when a standard account is + deleted; without it, or when it returns None, the Auth0 user is kept. + """ self._database_session = database_session self._auth0_user_deleter_factory = auth0_user_deleter_factory @@ -106,7 +111,6 @@ def _prepared_auth0_user_deleter(self) -> Auth0UserDeleter | None: Only account deletion needs the Management API, so its settings are read here rather than for every /v2/me request. """ - if self._auth0_user_deleter_factory is None: return None try: @@ -126,6 +130,11 @@ def _prepared_auth0_user_deleter(self) -> Auth0UserDeleter | None: def get_account( self, *, principal: AccountPrincipal, country_id: str ) -> AccountRead: + """Return the caller's account, provisioning a signed-in caller's user. + + Raise AccountNotFoundError for a signed-out browser that has not saved + anything yet. + """ with self._database_session.transaction() as session: owner = require_resolved_account( resolve_principal_user_for_read_in_session( @@ -142,6 +151,10 @@ def update_account( country_id: str, account_input: AccountUpdateInput, ) -> AccountRead: + """Change only the supplied account fields and return the account. + + Raise AccountEditRequiresSignInError for a signed-out browser. + """ signed_in = require_editable_account(principal) with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( @@ -158,6 +171,10 @@ def update_account( def export_account( self, *, principal: AccountPrincipal, country_id: str ) -> AccountExport: + """Return the caller's account and every association it owns, in all countries. + + A signed-out browser with no saved work gets an empty export with no user. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -186,6 +203,11 @@ def upgrade_account( unauthenticated_user_id: UUID, country_id: str, ) -> AccountUpgrade: + """Move a browser's unauthenticated user into the signed-in account. + + The browser's user becomes the account, merges into an existing one, or + is absent (outcome ``none``); see convert_unauthenticated_user_in_session. + """ with self._database_session.transaction() as session: conversion = convert_unauthenticated_user_in_session( session, @@ -206,6 +228,13 @@ def upgrade_account( return upgrade def delete_account(self, *, principal: AccountPrincipal) -> AccountDeletion: + """Delete the caller's user, its associations and, if configured, Auth0 users. + + Deleting an absent account succeeds with no user ID. Raise + AccountDeletionUnavailableError, before anything is deleted, when the + Auth0 deleter cannot authenticate, and AccountIdentityDeletionError when + the data was deleted but an Auth0 user was not. + """ if isinstance(principal, UnauthenticatedPrincipal): return self._delete_unauthenticated_account(principal) return self._delete_standard_account(principal) @@ -214,7 +243,6 @@ def _delete_unauthenticated_account( self, principal: UnauthenticatedPrincipal ) -> AccountDeletion: """Delete a browser's unauthenticated user; there is no Auth0 user.""" - with self._database_session.transaction() as session: owner = find_principal_user_in_session(session, principal) if owner is not None: diff --git a/policyengine_api/services/v2/user_accounts/transformations.py b/policyengine_api/services/v2/user_accounts/transformations.py index cf3353304..11134208f 100644 --- a/policyengine_api/services/v2/user_accounts/transformations.py +++ b/policyengine_api/services/v2/user_accounts/transformations.py @@ -14,6 +14,7 @@ def account_read(user: User, *, email: str | None) -> AccountRead: + """Convert a user row to the account read model, preferring the identity's email.""" return AccountRead( id=user.id, user_type=user.user_type, @@ -26,4 +27,5 @@ def account_read(user: User, *, email: str | None) -> AccountRead: def account_update_values(account_input: AccountUpdateInput) -> dict[str, Any]: + """Return only the fields the caller supplied, including explicit nulls.""" return account_input.model_dump(exclude_unset=True) diff --git a/policyengine_api/services/v2/user_accounts/types.py b/policyengine_api/services/v2/user_accounts/types.py index f5f2a3a45..38a6ee1a1 100644 --- a/policyengine_api/services/v2/user_accounts/types.py +++ b/policyengine_api/services/v2/user_accounts/types.py @@ -23,16 +23,21 @@ class StrictAccountInput(BaseModel): + """Reject fields outside the reviewed account contract; inputs are immutable.""" + model_config = ConfigDict(extra="forbid", frozen=True) class AccountUpdateInput(StrictAccountInput): + """Account fields to change; only the supplied ones are updated.""" + first_name: PersonName | None = None last_name: PersonName | None = None primary_country: CountryId | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> AccountUpdateInput: + """Require a supplied field, and reject an explicit null primary_country.""" supplied = self.model_fields_set.intersection( {"first_name", "last_name", "primary_country"} ) @@ -48,6 +53,8 @@ def require_valid_supplied_field(self) -> AccountUpdateInput: @dataclass(frozen=True) class AccountRead: + """A user's account as the service returns it to routes.""" + id: UUID user_type: str # Null only for an unauthenticated user. @@ -60,6 +67,8 @@ class AccountRead: @dataclass(frozen=True) class AccountExport: + """A user's account and every association it owns, across countries.""" + # Null when an unauthenticated principal has not saved anything yet. user: AccountRead | None policies: tuple[UserPolicyRead, ...] @@ -78,6 +87,8 @@ class AccountUpgrade: @dataclass(frozen=True) class AccountDeletion: + """The deleted user, if any, and the Auth0 subjects removed with it.""" + user_id: UUID | None auth0_subjects: tuple[str, ...] @@ -88,4 +99,6 @@ class Auth0UserDeleter(Protocol): def prepare_deletion(self) -> None: """Obtain working credentials; raise before any data is deleted.""" - def delete_user(self, subject: str) -> None: ... + def delete_user(self, subject: str) -> None: + """Delete one Auth0 user; an already-absent user counts as deleted.""" + ... diff --git a/policyengine_api/services/v2/user_accounts/validators.py b/policyengine_api/services/v2/user_accounts/validators.py index 9c246e262..0fe6cd7e7 100644 --- a/policyengine_api/services/v2/user_accounts/validators.py +++ b/policyengine_api/services/v2/user_accounts/validators.py @@ -31,20 +31,24 @@ class AccountDeletionUnavailableError(RuntimeError): def require_resolved_account(owner: ResolvedUser | None) -> ResolvedUser: - """An unauthenticated principal has no account until its first write.""" + """Return the resolved user, or raise AccountNotFoundError if there is none. + An unauthenticated principal has no account until its first write. + """ if owner is None: raise AccountNotFoundError("The account was not found") return owner def require_account(user: User | None) -> User: + """Return the stored user, or raise AccountNotFoundError if it is None.""" if user is None: raise AccountNotFoundError("The account was not found") return user def require_editable_account(principal: AccountPrincipal) -> UserPrincipal: + """Return a signed-in principal, or raise AccountEditRequiresSignInError.""" if not isinstance(principal, UserPrincipal): raise AccountEditRequiresSignInError("Sign in to edit your account details.") return principal diff --git a/policyengine_api/services/v2/user_households/database_connectors/creates.py b/policyengine_api/services/v2/user_households/database_connectors/creates.py index c3a9b494e..2ab2e6c10 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/creates.py +++ b/policyengine_api/services/v2/user_households/database_connectors/creates.py @@ -18,6 +18,7 @@ def create_user_household( session: Session, association_input: UserHouseholdCreationInput, ) -> UserHouseholdAssociation: + """Insert an association and return it with its database-generated fields loaded.""" association = UserHouseholdAssociation(**association_input.model_dump()) session.add(association) session.flush() diff --git a/policyengine_api/services/v2/user_households/database_connectors/deletes.py b/policyengine_api/services/v2/user_households/database_connectors/deletes.py index 800a3e645..b18a2ccfc 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_households/database_connectors/deletes.py @@ -9,5 +9,6 @@ def delete_user_household( session: Session, association: UserHouseholdAssociation, ) -> None: + """Delete an association row; the household it points to is kept.""" session.delete(association) session.flush() diff --git a/policyengine_api/services/v2/user_households/database_connectors/reads.py b/policyengine_api/services/v2/user_households/database_connectors/reads.py index f88ec8232..f2210a306 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_households/database_connectors/reads.py @@ -21,6 +21,7 @@ def read_household_for_association( session: Session, household_id: UUID, ) -> Household | None: + """Return the household with this id in any country, or None.""" return session.exec( select(Household).where(Household.id == household_id) ).one_or_none() @@ -33,6 +34,7 @@ def read_owned_user_household_row( country_id: str, association_id: UUID, ) -> UserHouseholdAssociation | None: + """Return the user's association with this id in this country, or None.""" return session.exec( select(UserHouseholdAssociation).where( UserHouseholdAssociation.id == association_id, @@ -52,6 +54,12 @@ def read_user_household_rows( limit: int, legacy_household_id: int | None = None, ) -> list[UserHouseholdAssociation]: + """Return up to ``limit + 1`` of a user's associations in creation order. + + With ``household_id`` or ``legacy_household_id``, only associations to that + household are returned. The extra row lets the caller report whether another + page exists. + """ statement = select(UserHouseholdAssociation).where( UserHouseholdAssociation.country_id == country_id, UserHouseholdAssociation.user_id == user_id, @@ -79,6 +87,7 @@ def read_user_household_rows( def read_all_user_household_rows( session: Session, *, user_id: UUID ) -> list[UserHouseholdAssociation]: + """Return every association the user owns, in any country, in creation order.""" return list( session.exec( select(UserHouseholdAssociation) @@ -98,7 +107,6 @@ def read_copied_household_ids( legacy_household_ids: Collection[int], ) -> dict[int, UUID]: """Return v2 household UUIDs for v1 households that Stage 11 copied.""" - if not legacy_household_ids: return {} rows = session.exec( diff --git a/policyengine_api/services/v2/user_households/database_connectors/updates.py b/policyengine_api/services/v2/user_households/database_connectors/updates.py index 146c45b42..ed6cfb989 100644 --- a/policyengine_api/services/v2/user_households/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_households/database_connectors/updates.py @@ -25,6 +25,10 @@ def update_user_household( repointed_household_id: UUID | None = None, repoint: bool = False, ) -> UserHouseholdAssociation: + """Apply the supplied fields, bump ``updated_at`` and return the row. + + With ``repoint``, also set ``household_id`` to ``repointed_household_id``. + """ for field_name, value in association_input.model_dump(exclude_unset=True).items(): setattr(association, field_name, value) if repoint: diff --git a/policyengine_api/services/v2/user_households/database_session.py b/policyengine_api/services/v2/user_households/database_session.py index ffa531976..2f7fae41d 100644 --- a/policyengine_api/services/v2/user_households/database_session.py +++ b/policyengine_api/services/v2/user_households/database_session.py @@ -13,15 +13,20 @@ class UserHouseholdDatabaseSession: + """Session lifetimes for user-household services: reads and transactions.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_households/services.py b/policyengine_api/services/v2/user_households/services.py index a35066d9e..65e09c0d0 100644 --- a/policyengine_api/services/v2/user_households/services.py +++ b/policyengine_api/services/v2/user_households/services.py @@ -72,7 +72,6 @@ def read_owned_user_household( association_id: UUID, ) -> UserHouseholdAssociation: """``user_id`` is None for an unauthenticated user that owns nothing yet.""" - return require_user_household( None if user_id is None @@ -94,7 +93,6 @@ def resolve_user_household_reference( legacy_household_id: int | None, ) -> UUID | None: """Validate an explicit v2 household and combine it with the v1 mapping.""" - if household_id is not None: validate_association_household( country_id=country_id, @@ -120,7 +118,6 @@ def read_all_user_households( session: Session, *, user_id: UUID ) -> tuple[UserHouseholdRead, ...]: """Return every household association owned by a user, across countries.""" - return tuple( user_household_read(row) for row in read_all_user_household_rows(session, user_id=user_id) @@ -135,6 +132,11 @@ def __init__( database_session: UserHouseholdDatabaseSession, legacy_references: LegacyReferenceChecker | None = None, ) -> None: + """Open sessions through ``database_session`` and check v1 ids with the checker. + + Without ``legacy_references``, an operation that names a v1 household + raises LegacyReferenceCheckUnavailableError. + """ self._database_session = database_session self._legacy_references = legacy_references @@ -154,6 +156,14 @@ def create_user_household( principal: AccountPrincipal, association_input: UserHouseholdCreateFields, ) -> UserHouseholdRead: + """Save a household for the principal, provisioning their v2 user if needed. + + A ``legacy_household_id`` must exist in v1 (LegacyReferenceNotFoundError) + and links the v2 copy of that household when Stage 11 made one. An explicit + ``household_id`` must exist in this country and match that copy + (AssociationHouseholdNotFoundError, AssociationCountryConflictError or + AssociationReferenceConflictError). + """ if association_input.legacy_household_id is not None: self._require_legacy_household( country_id=association_input.country_id, @@ -189,6 +199,11 @@ def get_user_household( country_id: str, association_id: UUID, ) -> UserHouseholdRead: + """Return the principal's association in this country. + + Raise UserHouseholdNotFoundError when it is absent or another user's. A + signed-out browser that has not saved anything yet owns nothing. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -212,6 +227,11 @@ def list_user_households( offset: int = 0, limit: int = 100, ) -> UserHouseholdPage: + """Return a page of the principal's associations in this country. + + ``household_id`` and ``legacy_household_id`` narrow it to one household. A + signed-out browser that has not saved anything yet gets an empty page. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -237,6 +257,12 @@ def patch_user_household( association_id: UUID, association_input: UserHouseholdUpdateInput, ) -> UserHouseholdRead: + """Update the supplied fields of the principal's association in this country. + + A new ``legacy_household_id`` must exist in v1 and repoints ``household_id`` + to the v2 copy of that household, or to None when there is no copy. Raise + UserHouseholdNotFoundError when the association is not the principal's. + """ repoint = "legacy_household_id" in association_input.model_fields_set if repoint: assert association_input.legacy_household_id is not None @@ -280,6 +306,7 @@ def delete_user_household( country_id: str, association_id: UUID, ) -> None: + """Delete the principal's association, or raise UserHouseholdNotFoundError.""" with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( session, principal, country_id=country_id diff --git a/policyengine_api/services/v2/user_households/transformations.py b/policyengine_api/services/v2/user_households/transformations.py index 7b4351159..2a7122fa6 100644 --- a/policyengine_api/services/v2/user_households/transformations.py +++ b/policyengine_api/services/v2/user_households/transformations.py @@ -20,6 +20,7 @@ def user_household_read( association: UserHouseholdAssociation, ) -> UserHouseholdRead: + """Convert a stored association row to the service's read model.""" return UserHouseholdRead( id=association.id, country_id=association.country_id, @@ -39,6 +40,10 @@ def user_household_page( offset: int, limit: int, ) -> UserHouseholdPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ return UserHouseholdPage( items=tuple(user_household_read(row) for row in rows[:limit]), offset=offset, @@ -53,6 +58,7 @@ def owned_user_household_input( user_id: UUID, household_id: UUID | None, ) -> UserHouseholdCreationInput: + """Add the session's owner and the resolved v2 household to the caller's fields.""" return UserHouseholdCreationInput( **fields.model_dump(exclude={"household_id"}), user_id=user_id, diff --git a/policyengine_api/services/v2/user_households/types.py b/policyengine_api/services/v2/user_households/types.py index 1f6a0f361..6298b3999 100644 --- a/policyengine_api/services/v2/user_households/types.py +++ b/policyengine_api/services/v2/user_households/types.py @@ -35,6 +35,7 @@ class UserHouseholdCreateFields(StrictUserHouseholdInput): @model_validator(mode="after") def require_household_reference(self) -> UserHouseholdCreateFields: + """Require a v2 household_id, a v1 legacy_household_id, or both.""" if self.household_id is None and self.legacy_household_id is None: raise ValueError( "At least one of legacy_household_id or household_id is required" @@ -54,6 +55,7 @@ class UserHouseholdCreationInput(StrictUserHouseholdInput): @model_validator(mode="after") def require_household_reference(self) -> UserHouseholdCreationInput: + """Require a v2 household_id, a v1 legacy_household_id, or both.""" if self.household_id is None and self.legacy_household_id is None: raise ValueError( "At least one of legacy_household_id or household_id is required" @@ -62,12 +64,19 @@ def require_household_reference(self) -> UserHouseholdCreationInput: class UserHouseholdUpdateInput(StrictUserHouseholdInput): + """A partial update of a saved household; only supplied fields are changed. + + At least one field must be supplied. ``legacy_household_id`` may point the + association at another v1 household but cannot be cleared. + """ + legacy_household_id: LegacyRecordId | None = None name: AssociationName | None = None description: AssociationDescription | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> UserHouseholdUpdateInput: + """Require a supplied field, and reject an explicit null legacy_household_id.""" mutable_fields = {"legacy_household_id", "name", "description"} supplied = self.model_fields_set.intersection(mutable_fields) if not supplied: @@ -82,6 +91,8 @@ def require_valid_supplied_field(self) -> UserHouseholdUpdateInput: @dataclass(frozen=True) class UserHouseholdRead: + """A stored user-household association as the service returns it to routes.""" + id: UUID country_id: str user_id: UUID @@ -95,6 +106,8 @@ class UserHouseholdRead: @dataclass(frozen=True) class UserHouseholdPage: + """One page of associations in creation order; ``has_more`` says if more follow.""" + items: tuple[UserHouseholdRead, ...] offset: int limit: int diff --git a/policyengine_api/services/v2/user_households/validators.py b/policyengine_api/services/v2/user_households/validators.py index 7d40d4103..de1f29aa4 100644 --- a/policyengine_api/services/v2/user_households/validators.py +++ b/policyengine_api/services/v2/user_households/validators.py @@ -36,6 +36,7 @@ def require_user_household( *, association_id: UUID, ) -> UserHouseholdAssociation: + """Return the association, or raise UserHouseholdNotFoundError if it is None.""" if association is None: raise UserHouseholdNotFoundError( f"user-household association {association_id} was not found" @@ -49,6 +50,10 @@ def validate_association_household( household_id: UUID, household: Household | None, ) -> None: + """Check that a referenced household exists and is in the association's country. + + Raise AssociationHouseholdNotFoundError or AssociationCountryConflictError. + """ if household is None: raise AssociationHouseholdNotFoundError( f"household {household_id} was not found" @@ -65,7 +70,6 @@ def resolve_association_household_id( mapped_household_id: UUID | None, ) -> UUID | None: """Combine an explicit v2 reference with the v1 mapping, if both exist.""" - if ( requested_household_id is not None and mapped_household_id is not None diff --git a/policyengine_api/services/v2/user_policies/database_connectors/creates.py b/policyengine_api/services/v2/user_policies/database_connectors/creates.py index 822abafd3..4ed880968 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/creates.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/creates.py @@ -23,6 +23,7 @@ def create_user_policy( session: Session, association_input: UserPolicyCreationInput ) -> UserPolicy: + """Insert an association and return it with its database-generated fields loaded.""" association = UserPolicy(**association_input.model_dump()) session.add(association) session.flush() @@ -31,6 +32,7 @@ def create_user_policy( def create_transition_user(session: Session, *, primary_country: str) -> User: + """Insert a v2 user to stand in for a v1 user that has no v2 user yet.""" user = User(primary_country=primary_country) session.add(user) session.flush() @@ -40,6 +42,11 @@ def create_transition_user(session: Session, *, primary_country: str) -> User: def create_legacy_user_mapping( session: Session, *, legacy_user_id: str, user_id: UUID ) -> UUID | None: + """Map a v1 user id to a v2 user. + + Return the v2 user id, or None when the v1 id is already mapped (for example, + a concurrent copy won the race). + """ return session.execute( insert(LegacyUserMapping) .values(legacy_user_id=legacy_user_id, user_id=user_id) @@ -58,6 +65,11 @@ def create_legacy_user_policy_mapping( fingerprint_version: int, fingerprint: str, ) -> UUID | None: + """Map a v1 saved policy to its v2 association at ``source_revision``. + + Return the new mapping id, or None when this v1 saved policy is already + mapped in this country. + """ return session.execute( insert(LegacyUserPolicyMapping) .values( diff --git a/policyengine_api/services/v2/user_policies/database_connectors/deletes.py b/policyengine_api/services/v2/user_policies/database_connectors/deletes.py index 7b62bb24b..d9f2bc6be 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/deletes.py @@ -11,10 +11,12 @@ def delete_user_policy(session: Session, association: UserPolicy) -> None: + """Delete an association row; the policy it points to is kept.""" session.delete(association) session.flush() def delete_transition_user(session: Session, user: User) -> None: + """Delete a transition user that lost a race to map its v1 user id.""" session.delete(user) session.flush() diff --git a/policyengine_api/services/v2/user_policies/database_connectors/reads.py b/policyengine_api/services/v2/user_policies/database_connectors/reads.py index 0389864b8..908b370a1 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/reads.py @@ -21,16 +21,19 @@ def read_user(session: Session, user_id: UUID) -> User | None: + """Return the v2 user with this id, or None.""" return session.get(User, user_id) def read_policy_for_association(session: Session, policy_id: UUID) -> Policy | None: + """Return the policy with this id in any country, or None.""" return session.exec(select(Policy).where(Policy.id == policy_id)).one_or_none() def read_legacy_user_mapping( session: Session, legacy_user_id: str, *, lock: bool ) -> LegacyUserMapping | None: + """Return the mapping for a v1 user id, or None; ``lock`` selects it FOR UPDATE.""" statement = select(LegacyUserMapping).where( LegacyUserMapping.legacy_user_id == legacy_user_id ) @@ -46,6 +49,11 @@ def read_legacy_user_policy_mapping( legacy_user_policy_id: int, lock: bool, ) -> LegacyUserPolicyMapping | None: + """Return the mapping for a v1 saved policy in this country, or None. + + With ``lock``, the row is selected FOR UPDATE so concurrent revisions of the + same v1 saved policy serialize on it. + """ statement = select(LegacyUserPolicyMapping).where( LegacyUserPolicyMapping.country_id == country_id, LegacyUserPolicyMapping.legacy_user_policy_id == legacy_user_policy_id, @@ -58,6 +66,7 @@ def read_legacy_user_policy_mapping( def read_mapped_user_policy( session: Session, mapping: LegacyUserPolicyMapping ) -> UserPolicy | None: + """Return the association a mapping points to, in the mapping's country, or None.""" return session.exec( select(UserPolicy).where( UserPolicy.id == mapping.user_policy_id, @@ -73,6 +82,7 @@ def read_owned_user_policy_row( country_id: str, association_id: UUID, ) -> UserPolicy | None: + """Return the user's association with this id in this country, or None.""" return session.exec( select(UserPolicy).where( UserPolicy.id == association_id, @@ -92,6 +102,12 @@ def read_user_policy_rows( limit: int, legacy_policy_id: int | None = None, ) -> list[UserPolicy]: + """Return up to ``limit + 1`` of a user's associations in creation order. + + With ``policy_id`` or ``legacy_policy_id``, only associations to that policy + are returned. The extra row lets the caller report whether another page + exists. + """ statement = select(UserPolicy).where( UserPolicy.country_id == country_id, UserPolicy.user_id == user_id, @@ -110,6 +126,7 @@ def read_user_policy_rows( def read_all_user_policy_rows(session: Session, *, user_id: UUID) -> list[UserPolicy]: + """Return every association the user owns, in any country, in creation order.""" return list( session.exec( select(UserPolicy) @@ -120,6 +137,7 @@ def read_all_user_policy_rows(session: Session, *, user_id: UUID) -> list[UserPo def read_user_policy_mirror_exists(session: Session, *, association_id: UUID) -> bool: + """Return whether the Stage 10 v1 saved-policy mirror maintains the association.""" return ( session.exec( select(LegacyUserPolicyMapping.id).where( @@ -137,7 +155,6 @@ def read_mirrored_policy_ids( legacy_policy_ids: Collection[int], ) -> dict[int, UUID]: """Return v2 policy UUIDs for v1 policies that the Stage 10 mirror copied.""" - if not legacy_policy_ids: return {} rows = session.exec( diff --git a/policyengine_api/services/v2/user_policies/database_connectors/updates.py b/policyengine_api/services/v2/user_policies/database_connectors/updates.py index c6fd5dc16..151ac8499 100644 --- a/policyengine_api/services/v2/user_policies/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_policies/database_connectors/updates.py @@ -23,6 +23,10 @@ def update_user_policy( repointed_policy_id: UUID | None = None, repoint: bool = False, ) -> UserPolicy: + """Apply the supplied fields, bump ``updated_at`` and return the row. + + With ``repoint``, also set ``policy_id`` to ``repointed_policy_id``. + """ for field_name, value in association_input.model_dump(exclude_unset=True).items(): setattr(association, field_name, value) if repoint: @@ -44,6 +48,10 @@ def update_legacy_user_policy_state( fingerprint: str, source_revision: int, ) -> None: + """Record a newly applied v1 revision on its mapping. + + With ``update_name``, also set the association's name to the v1 reform label. + """ if update_name: association.name = reform_label association.updated_at = utc_now() diff --git a/policyengine_api/services/v2/user_policies/database_session.py b/policyengine_api/services/v2/user_policies/database_session.py index 78d084ac9..8ec71003f 100644 --- a/policyengine_api/services/v2/user_policies/database_session.py +++ b/policyengine_api/services/v2/user_policies/database_session.py @@ -16,14 +16,17 @@ class UserPolicyDatabaseSession: """Open read sessions and transactions for user-policy services.""" def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_policies/services.py b/policyengine_api/services/v2/user_policies/services.py index 97b05a737..a6714b330 100644 --- a/policyengine_api/services/v2/user_policies/services.py +++ b/policyengine_api/services/v2/user_policies/services.py @@ -99,6 +99,7 @@ def read_user_policy_page( limit: int, legacy_policy_id: int | None = None, ) -> UserPolicyPage: + """Return a page of a user's associations, optionally for one policy.""" rows = read_user_policy_rows( session, country_id=country_id, @@ -119,7 +120,6 @@ def read_owned_user_policy( association_id: UUID, ) -> UserPolicy: """``user_id`` is None for an unauthenticated user that owns nothing yet.""" - return require_user_policy( None if user_id is None @@ -141,7 +141,6 @@ def resolve_user_policy_reference( legacy_policy_id: int | None, ) -> UUID | None: """Validate an explicit v2 policy and combine it with the v1 mapping.""" - if policy_id is not None: validate_association_policy( country_id=country_id, @@ -167,7 +166,6 @@ def read_all_user_policies( session: Session, *, user_id: UUID ) -> tuple[UserPolicyRead, ...]: """Return every association owned by a user, across countries.""" - return tuple( user_policy_read(row) for row in read_all_user_policy_rows(session, user_id=user_id) @@ -180,6 +178,11 @@ def resolve_legacy_user_id( legacy_user_id: str, primary_country: str, ) -> UUID: + """Return the v2 user id for a v1 user id, creating a transition user if needed. + + If a concurrent transaction maps the v1 id first, the new user is deleted and + the winner's user is returned. + """ existing = read_legacy_user_mapping(session, legacy_user_id, lock=True) if existing is not None: return validate_existing_legacy_user_mapping( @@ -218,6 +221,12 @@ def apply_existing_legacy_user_policy_mapping( changed_fields: frozenset[str], source_revision: int, ) -> LegacyUserPolicyPersistenceResult: + """Apply one v1 revision to an already-mapped association. + + Stale and replayed revisions change nothing. The next revision renames the + association only if the v1 reform label changed. Raise + LegacyUserPolicyIntegrityError on conflicting or skipped revisions. + """ action, association = validate_existing_legacy_user_policy_mapping( mapping, association=read_mapped_user_policy(session, mapping), @@ -274,6 +283,13 @@ def persist_legacy_user_policy_mapping( source_revision: int, changed_fields: frozenset[str], ) -> LegacyUserPolicyPersistenceResult: + """Create the association and mapping for a v1 saved policy, or apply a revision. + + The mapping row is locked, so revisions of one v1 saved policy serialize. If + a concurrent transaction maps it first, the new association is deleted and + the revision is applied to the winner's. Raise LegacyUserPolicyIntegrityError + when the projection has no v2 policy or the conflict leaves no mapping. + """ policy_id = require_mirrored_policy_id(projection) existing = read_legacy_user_policy_mapping( session, @@ -345,6 +361,11 @@ def mirror_legacy_user_policy_in_session( source_revision: int, changed_fields: frozenset[str] = frozenset(), ) -> LegacyUserPolicyPersistenceResult: + """Copy one revision of a v1 saved policy into v2. + + Copy its reform policy, resolve the v1 user to a v2 user, then create or + update the association. Only the v1 reform label is copied, as the name. + """ validate_legacy_user_policy_input( snapshot, reform_snapshot, @@ -388,6 +409,11 @@ def __init__( database_session: UserPolicyDatabaseSession, legacy_references: LegacyReferenceChecker | None = None, ) -> None: + """Open sessions through ``database_session`` and check v1 ids with the checker. + + Without ``legacy_references``, an operation that names a v1 policy raises + LegacyReferenceCheckUnavailableError. + """ self._database_session = database_session self._legacy_references = legacy_references @@ -405,6 +431,14 @@ def create_user_policy( principal: AccountPrincipal, association_input: UserPolicyCreateFields, ) -> UserPolicyRead: + """Save a policy for the principal, provisioning their v2 user if needed. + + A ``legacy_policy_id`` must exist in v1 (LegacyReferenceNotFoundError) and + links the v2 copy of that policy when the Stage 10 mirror made one. An + explicit ``policy_id`` must exist in this country and match that copy + (AssociationPolicyNotFoundError, AssociationCountryConflictError or + AssociationReferenceConflictError). + """ if association_input.legacy_policy_id is not None: self._require_legacy_policy( country_id=association_input.country_id, @@ -440,6 +474,11 @@ def get_user_policy( country_id: str, association_id: UUID, ) -> UserPolicyRead: + """Return the principal's association in this country. + + Raise UserPolicyNotFoundError when it is absent or another user's. A + signed-out browser that has not saved anything yet owns nothing. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -463,6 +502,11 @@ def list_user_policies( offset: int = 0, limit: int = 100, ) -> UserPolicyPage: + """Return a page of the principal's associations in this country. + + ``policy_id`` and ``legacy_policy_id`` narrow it to one policy. A signed-out + browser that has not saved anything yet gets an empty page. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -487,6 +531,13 @@ def patch_user_policy( association_id: UUID, association_input: UserPolicyUpdateInput, ) -> UserPolicyRead: + """Update the supplied fields of the principal's association in this country. + + A new ``legacy_policy_id`` must exist in v1 and repoints ``policy_id`` to the + v2 copy of that policy, or to None when there is no copy. Raise + UserPolicyNotFoundError when the association is not the principal's, and + MirrorOwnedAssociationError when the v1 saved-policy mirror maintains it. + """ repoint = "legacy_policy_id" in association_input.model_fields_set if repoint: assert association_input.legacy_policy_id is not None @@ -534,6 +585,11 @@ def delete_user_policy( country_id: str, association_id: UUID, ) -> None: + """Delete the principal's association in this country. + + Raise UserPolicyNotFoundError when it is not the principal's, and + MirrorOwnedAssociationError when the v1 saved-policy mirror maintains it. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( session, principal, country_id=country_id @@ -558,6 +614,7 @@ def mirror_legacy_user_policy( source_revision: int, changed_fields: frozenset[str], ) -> LegacyUserPolicyPersistenceResult: + """Copy one v1 saved-policy revision into v2 in its own transaction.""" with self._database_session.transaction() as session: return mirror_legacy_user_policy_in_session( session, diff --git a/policyengine_api/services/v2/user_policies/transformations.py b/policyengine_api/services/v2/user_policies/transformations.py index 1870f1e7b..5715fde3a 100644 --- a/policyengine_api/services/v2/user_policies/transformations.py +++ b/policyengine_api/services/v2/user_policies/transformations.py @@ -23,6 +23,7 @@ def user_policy_read(association: UserPolicy) -> UserPolicyRead: + """Convert a stored association row to the service's read model.""" return UserPolicyRead( id=association.id, country_id=association.country_id, @@ -41,6 +42,10 @@ def user_policy_read(association: UserPolicy) -> UserPolicyRead: def user_policy_page( rows: list[UserPolicy], *, offset: int, limit: int ) -> UserPolicyPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ return UserPolicyPage( items=tuple(user_policy_read(row) for row in rows[:limit]), offset=offset, @@ -51,7 +56,6 @@ def user_policy_page( def fingerprint_legacy_user_policy(snapshot: LegacyUserPolicySnapshot) -> str: """Hash every committed source field through deterministic JSON.""" - document = { "fingerprint_version": USER_POLICY_FINGERPRINT_VERSION, **snapshot.model_dump(mode="json"), @@ -73,7 +77,6 @@ def project_legacy_user_policy( policy_id: UUID, ) -> UserPolicyCreationInput: """Map v1 presentation data onto an association, never core policy content.""" - return UserPolicyCreationInput( country_id=snapshot.country_id, user_id=user_id, @@ -91,7 +94,6 @@ def owned_user_policy_input( policy_id: UUID | None, ) -> UserPolicyCreationInput: """Attach the session owner and the resolved v2 policy reference.""" - return UserPolicyCreationInput( **fields.model_dump(exclude={"policy_id"}), user_id=user_id, @@ -105,4 +107,9 @@ def legacy_name_requires_update( reform_label: str | None, changed_fields: frozenset[str], ) -> bool: + """Return whether a v1 revision should rename the association. + + Only a changed v1 reform label that differs from the current name renames it, + so a v2-side rename survives v1 revisions that leave the label alone. + """ return "reform_label" in changed_fields and association.name != reform_label diff --git a/policyengine_api/services/v2/user_policies/types.py b/policyengine_api/services/v2/user_policies/types.py index 5bfea8a5f..213ad45bd 100644 --- a/policyengine_api/services/v2/user_policies/types.py +++ b/policyengine_api/services/v2/user_policies/types.py @@ -53,6 +53,7 @@ class UserPolicyCreateFields(StrictAssociationInput): @model_validator(mode="after") def require_policy_reference(self) -> UserPolicyCreateFields: + """Require a v2 policy_id, a v1 legacy_policy_id, or both.""" if self.policy_id is None and self.legacy_policy_id is None: raise ValueError( "At least one of legacy_policy_id or policy_id is required" @@ -75,6 +76,7 @@ class UserPolicyCreationInput(StrictAssociationInput): @model_validator(mode="after") def require_policy_reference(self) -> UserPolicyCreationInput: + """Require a v2 policy_id, a v1 legacy_policy_id, or both.""" if self.policy_id is None and self.legacy_policy_id is None: raise ValueError( "At least one of legacy_policy_id or policy_id is required" @@ -83,6 +85,12 @@ def require_policy_reference(self) -> UserPolicyCreationInput: class UserPolicyUpdateInput(StrictAssociationInput): + """A partial update of a saved policy; only supplied fields are changed. + + At least one field must be supplied. ``legacy_policy_id`` may point the + association at another v1 policy but cannot be cleared. + """ + legacy_policy_id: LegacyRecordId | None = None name: AssociationName | None = None description: AssociationDescription | None = None @@ -91,6 +99,7 @@ class UserPolicyUpdateInput(StrictAssociationInput): @model_validator(mode="after") def require_supplied_field(self) -> UserPolicyUpdateInput: + """Require a supplied field, and reject an explicit null legacy_policy_id.""" supplied = self.model_fields_set.intersection( {"legacy_policy_id", "name", "description", "baseline", "provenance"} ) @@ -126,6 +135,12 @@ class LegacyUserPolicySnapshot(StrictAssociationInput): class LegacyUserPolicyMappingAction(StrEnum): + """What to do with a v1 saved-policy revision whose row is already mapped. + + STALE is older than the applied revision and is ignored, REPLAY repeats the + applied revision, and UPDATE is the next revision, which is applied. + """ + STALE = "stale" REPLAY = "replay" UPDATE = "update" @@ -133,6 +148,8 @@ class LegacyUserPolicyMappingAction(StrEnum): @dataclass(frozen=True) class UserPolicyRead: + """A stored user-policy association as the service returns it to routes.""" + id: UUID country_id: str user_id: UUID @@ -148,6 +165,8 @@ class UserPolicyRead: @dataclass(frozen=True) class UserPolicyPage: + """One page of associations in creation order; ``has_more`` says if more follow.""" + items: tuple[UserPolicyRead, ...] offset: int limit: int @@ -156,6 +175,12 @@ class UserPolicyPage: @dataclass(frozen=True) class LegacyUserPolicyPersistenceResult: + """Outcome of copying one v1 saved-policy revision into v2. + + Names the v2 association and policy, and whether this call created the + association, renamed it, or created the v1-to-v2 mapping. + """ + association_id: UUID policy_id: UUID association_created: bool diff --git a/policyengine_api/services/v2/user_policies/validators.py b/policyengine_api/services/v2/user_policies/validators.py index bf07d1b50..5a60854a1 100644 --- a/policyengine_api/services/v2/user_policies/validators.py +++ b/policyengine_api/services/v2/user_policies/validators.py @@ -52,6 +52,7 @@ class LegacyUserPolicyIntegrityError(RuntimeError): def require_user_policy( association: UserPolicy | None, *, association_id: UUID ) -> UserPolicy: + """Return the association, or raise UserPolicyNotFoundError if it is None.""" if association is None: raise UserPolicyNotFoundError( f"user-policy association {association_id} was not found" @@ -65,6 +66,10 @@ def validate_association_policy( policy_id: UUID, policy: Policy | None, ) -> None: + """Check that a referenced policy exists and is in the association's country. + + Raise AssociationPolicyNotFoundError or AssociationCountryConflictError. + """ if policy is None: raise AssociationPolicyNotFoundError(f"policy {policy_id} was not found") if policy.country_id != country_id: @@ -79,7 +84,6 @@ def resolve_association_policy_id( mapped_policy_id: UUID | None, ) -> UUID | None: """Combine an explicit v2 reference with the v1 mapping, if both exist.""" - if ( requested_policy_id is not None and mapped_policy_id is not None @@ -94,6 +98,7 @@ def resolve_association_policy_id( def validate_existing_legacy_user_mapping( mapping: LegacyUserMapping, *, user: User | None ) -> UUID: + """Return the mapped v2 user id, or raise if that user no longer exists.""" if user is None: raise LegacyUserPolicyIntegrityError( "legacy user mapping has no referenced v2 user" @@ -104,6 +109,7 @@ def validate_existing_legacy_user_mapping( def validate_legacy_user_mapping_conflict( mapping: LegacyUserMapping | None, *, user: User | None ) -> UUID: + """Return the user id of the mapping that won the race, or raise if unusable.""" if mapping is None or user is None: raise LegacyUserPolicyIntegrityError( "legacy user mapping conflict did not resolve to a v2 user" @@ -117,6 +123,7 @@ def validate_legacy_user_policy_input( *, source_revision: int, ) -> None: + """Require a positive revision and a saved policy that uses the reform snapshot.""" if source_revision <= 0: raise LegacyUserPolicyIntegrityError( "legacy user-policy source revision must be positive" @@ -141,6 +148,13 @@ def validate_existing_legacy_user_policy_mapping( source_revision: int, fingerprint_version: int, ) -> tuple[LegacyUserPolicyMappingAction, UserPolicy]: + """Decide how to apply a v1 revision to an existing mapping. + + Return the action with the mapped association. Raise + LegacyUserPolicyIntegrityError if the association is missing or its user, + policy or country changed, if the fingerprint version differs, if a replay + has a different fingerprint, or if a revision was skipped. + """ if association is None: raise LegacyUserPolicyIntegrityError( "legacy user-policy mapping has no association" @@ -173,8 +187,10 @@ def validate_existing_legacy_user_policy_mapping( def require_mirrored_policy_id(projection: UserPolicyCreationInput) -> UUID: - """The Stage 10 mirror always links a saved policy to its v2 policy.""" + """Return the projection's v2 policy, which the Stage 10 mirror always sets. + Raise LegacyUserPolicyIntegrityError when it is missing. + """ if projection.policy_id is None: raise LegacyUserPolicyIntegrityError( "mirrored saved policy has no v2 policy reference" diff --git a/policyengine_api/services/v2/user_reports/database_connectors/creates.py b/policyengine_api/services/v2/user_reports/database_connectors/creates.py index 887270592..2dc5d838e 100644 --- a/policyengine_api/services/v2/user_reports/database_connectors/creates.py +++ b/policyengine_api/services/v2/user_reports/database_connectors/creates.py @@ -16,6 +16,7 @@ def create_user_report( session: Session, association_input: UserReportCreationInput, ) -> UserReportAssociation: + """Insert an association and return it with its database-generated fields loaded.""" association = UserReportAssociation(**association_input.model_dump()) session.add(association) session.flush() diff --git a/policyengine_api/services/v2/user_reports/database_connectors/deletes.py b/policyengine_api/services/v2/user_reports/database_connectors/deletes.py index 23741dbbf..878675cd6 100644 --- a/policyengine_api/services/v2/user_reports/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_reports/database_connectors/deletes.py @@ -14,5 +14,6 @@ def delete_user_report( session: Session, association: UserReportAssociation, ) -> None: + """Delete an association row; the report and the saved policies it links stay.""" session.delete(association) session.flush() diff --git a/policyengine_api/services/v2/user_reports/database_connectors/reads.py b/policyengine_api/services/v2/user_reports/database_connectors/reads.py index 75c4539e8..b583bf5ab 100644 --- a/policyengine_api/services/v2/user_reports/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_reports/database_connectors/reads.py @@ -19,6 +19,7 @@ def read_owned_user_report_row( country_id: str, association_id: UUID, ) -> UserReportAssociation | None: + """Return the user's association with this id in this country, or None.""" return session.exec( select(UserReportAssociation).where( UserReportAssociation.id == association_id, @@ -35,7 +36,6 @@ def read_user_report_row( association_id: UUID, ) -> UserReportAssociation | None: """Read a row whoever owns it: a saved report's ID is its shareable link.""" - return session.exec( select(UserReportAssociation).where( UserReportAssociation.id == association_id, @@ -53,6 +53,11 @@ def read_user_report_rows( offset: int, limit: int, ) -> list[UserReportAssociation]: + """Return up to ``limit + 1`` of a user's associations in creation order. + + With ``legacy_report_id``, only associations to that v1 report output are + returned. The extra row lets the caller report whether another page exists. + """ statement = select(UserReportAssociation).where( UserReportAssociation.country_id == country_id, UserReportAssociation.user_id == user_id, @@ -76,6 +81,7 @@ def read_user_report_rows( def read_all_user_report_rows( session: Session, *, user_id: UUID ) -> list[UserReportAssociation]: + """Return every association the user owns, in every country, in creation order.""" return list( session.exec( select(UserReportAssociation) @@ -95,6 +101,7 @@ def read_owned_user_policy_row( country_id: str, user_policy_id: UUID, ) -> UserPolicy | None: + """Return the user's saved policy with this id in this country, or None.""" return session.exec( select(UserPolicy).where( UserPolicy.id == user_policy_id, @@ -111,6 +118,7 @@ def read_owned_user_policy_exists( country_id: str, user_policy_id: UUID, ) -> bool: + """Return whether the user has a saved policy with this id in this country.""" return ( session.exec( select(UserPolicy.id).where( diff --git a/policyengine_api/services/v2/user_reports/database_connectors/updates.py b/policyengine_api/services/v2/user_reports/database_connectors/updates.py index 4b6b13de5..4089c3373 100644 --- a/policyengine_api/services/v2/user_reports/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_reports/database_connectors/updates.py @@ -19,6 +19,7 @@ def update_user_report( association: UserReportAssociation, values: Mapping[str, Any], ) -> UserReportAssociation: + """Set the given column values, bump ``updated_at`` and return the row.""" for field_name, value in values.items(): setattr(association, field_name, value) association.updated_at = utc_now() diff --git a/policyengine_api/services/v2/user_reports/database_session.py b/policyengine_api/services/v2/user_reports/database_session.py index b2102842b..98694f0a9 100644 --- a/policyengine_api/services/v2/user_reports/database_session.py +++ b/policyengine_api/services/v2/user_reports/database_session.py @@ -13,15 +13,20 @@ class UserReportDatabaseSession: + """Open read sessions and transactions for user-report services.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_reports/services.py b/policyengine_api/services/v2/user_reports/services.py index 000967287..f3e218bb8 100644 --- a/policyengine_api/services/v2/user_reports/services.py +++ b/policyengine_api/services/v2/user_reports/services.py @@ -82,7 +82,6 @@ def read_owned_user_report( association_id: UUID, ) -> UserReportAssociation: """``user_id`` is None for an unauthenticated user that owns nothing yet.""" - return require_user_report( None if user_id is None @@ -103,6 +102,10 @@ def validate_linked_user_policy( country_id: str, user_policy_id: UUID | None, ) -> None: + """Check that a saved-reform link names the owner's saved policy in this country. + + A None link passes; otherwise raise AssociationLinkNotFoundError. + """ if user_policy_id is None: return require_owned_user_policy( @@ -120,7 +123,6 @@ def read_all_user_reports( session: Session, *, user_id: UUID ) -> tuple[UserReportRead, ...]: """Return every report association owned by a user, across countries.""" - return tuple( user_report_read(row) for row in read_all_user_report_rows(session, user_id=user_id) @@ -135,6 +137,12 @@ def __init__( database_session: UserReportDatabaseSession, legacy_references: LegacyReferenceChecker | None = None, ) -> None: + """Open sessions through ``database_session``. + + ``legacy_references`` checks v1 report outputs before a write. Without + it, creating a saved report or repointing one raises + LegacyReferenceCheckUnavailableError. + """ self._database_session = database_session self._legacy_references = legacy_references @@ -152,6 +160,13 @@ def create_user_report( principal: AccountPrincipal, association_input: UserReportCreateFields, ) -> UserReportRead: + """Save a v1 report for the principal, provisioning its v2 user if needed. + + The v1 report is checked first, in a read-only v1 transaction, and + LegacyReferenceNotFoundError is raised if it does not exist in the + country. Each saved-reform link must name the owner's saved policy in + the same country, or AssociationLinkNotFoundError is raised. + """ self._require_legacy_report( country_id=association_input.country_id, legacy_report_id=association_input.legacy_report_id, @@ -186,6 +201,11 @@ def get_user_report( country_id: str, association_id: UUID, ) -> UserReportRead: + """Return the principal's saved report, or raise UserReportNotFoundError. + + The report must be in this country. A signed-out browser that has never + saved anything owns nothing, and is not provisioned by the read. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -206,7 +226,6 @@ def get_shared_user_report( association_id: UUID, ) -> SharedUserReportRead: """Read a saved report for anyone holding its link; no credential.""" - with self._database_session.transaction() as session: association = require_user_report( read_user_report_row( @@ -240,6 +259,12 @@ def list_user_reports( offset: int = 0, limit: int = 100, ) -> UserReportPage: + """Return a page of the principal's saved reports in this country. + + With ``legacy_report_id``, only saves of that v1 report output are listed. + A signed-out browser that has never saved anything gets an empty page and + is not provisioned. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -264,6 +289,13 @@ def patch_user_report( association_id: UUID, association_input: UserReportUpdateInput, ) -> UserReportRead: + """Update the supplied fields of one of the principal's saved reports. + + A new ``legacy_report_id`` is checked in v1 first, as on create, and a + supplied saved-reform link must name the owner's saved policy in this + country. Raise UserReportNotFoundError if the principal has no such + report in this country. + """ if association_input.legacy_report_id is not None: self._require_legacy_report( country_id=country_id, @@ -302,6 +334,10 @@ def delete_user_report( country_id: str, association_id: UUID, ) -> None: + """Delete one of the principal's saved reports in this country. + + Raise UserReportNotFoundError if the principal has no such report. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( session, principal, country_id=country_id diff --git a/policyengine_api/services/v2/user_reports/transformations.py b/policyengine_api/services/v2/user_reports/transformations.py index 7643a2353..b029e0727 100644 --- a/policyengine_api/services/v2/user_reports/transformations.py +++ b/policyengine_api/services/v2/user_reports/transformations.py @@ -23,6 +23,7 @@ def user_report_read(association: UserReportAssociation) -> UserReportRead: + """Convert a stored association row to the service's read model.""" return UserReportRead( id=association.id, country_id=association.country_id, @@ -48,6 +49,12 @@ def shared_user_report_read( *, baseline_reform: UserPolicy | None, ) -> SharedUserReportRead: + """Build the view of a saved report that anyone holding its link may read. + + The owner, the description and the saved-reform IDs are left out; + ``baseline_reform`` is the saved reform the report was compared with, which + the view names, or None. + """ return SharedUserReportRead( id=association.id, country_id=association.country_id, @@ -74,6 +81,10 @@ def user_report_page( offset: int, limit: int, ) -> UserReportPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ return UserReportPage( items=tuple(user_report_read(row) for row in rows[:limit]), offset=offset, @@ -87,6 +98,7 @@ def owned_user_report_input( *, user_id: UUID, ) -> UserReportCreationInput: + """Add the resolved owner and serialize the provision and validation snapshots.""" return UserReportCreationInput( **fields.model_dump(exclude={"provisions", "validation"}), provisions=stored_provisions(fields.provisions), @@ -99,7 +111,6 @@ def user_report_update_values( association_input: UserReportUpdateInput, ) -> dict[str, Any]: """Return explicitly supplied fields with JSON documents normalized.""" - values = association_input.model_dump( exclude_unset=True, exclude={"provisions", "validation"}, diff --git a/policyengine_api/services/v2/user_reports/types.py b/policyengine_api/services/v2/user_reports/types.py index 615241224..149c49ddf 100644 --- a/policyengine_api/services/v2/user_reports/types.py +++ b/policyengine_api/services/v2/user_reports/types.py @@ -48,6 +48,12 @@ class StrictUserReportModel(StrictAssociationModel): class ReportProvisionInterval(StrictUserReportModel): + """One dated value in a provision's schedule, from ``start_date`` to ``end_date``. + + Dates are ``YYYY-MM-DD`` strings and ``value`` is any finite JSON value. + ``parameter`` is optional, and is left out of the stored report when unset. + """ + # The contract makes this key optional rather than nullable. parameter: str | None = Field(default=None, exclude_if=lambda value: value is None) start_date: IsoDate @@ -56,6 +62,15 @@ class ReportProvisionInterval(StrictUserReportModel): class ReportProvision(StrictUserReportModel): + """One provision of the report's reform: a parameter's baseline and reform values. + + ``path`` is the parameter path and ``breadcrumb`` its readable label, and + ``intervals`` optionally gives a dated schedule. ``growth`` says how a value + set from one year on moves in later years (``fixed``, ``current_law`` or an + index's parameter path); the other ``growth_*`` fields say when it moves and + how grown values round. They are absent for a hand-set schedule. + """ + path: str breadcrumb: str unit: str | None = None @@ -73,6 +88,8 @@ class ReportProvision(StrictUserReportModel): class ReportCalibrationWorst(StrictUserReportModel): + """The calibration target a match fits worst, with its source and geography.""" + name: LongLabelText source: LongLabelText geography: ShortText @@ -80,6 +97,13 @@ class ReportCalibrationWorst(StrictUserReportModel): class ReportCalibrationMatch(StrictUserReportModel): + """Calibration of one variable: how closely its targets are met. + + ``ring`` is ``primary``, ``mechanism`` or ``downstream``. + ``mean_abs_relative_error`` is averaged over the variable's ``target_count`` + targets, and ``worst`` is the target with the largest error. + """ + variable: LabelText depth: int ring: Literal["primary", "mechanism", "downstream"] @@ -89,6 +113,8 @@ class ReportCalibrationMatch(StrictUserReportModel): class ReportCalibration(StrictUserReportModel): + """Calibration results for one geography, with at most 500 variable matches.""" + release_id: LabelText | None = None geography: ShortText reached_count: int @@ -99,10 +125,18 @@ class ReportCalibration(StrictUserReportModel): class ReportScorecard(StrictUserReportModel): + """The programs on the report's validation scorecard, at most 200 names.""" + programs: Annotated[list[LabelText], Field(max_length=MAXIMUM_SCORECARD_PROGRAMS)] class ReportValidation(StrictUserReportModel): + """The validation snapshot saved with a report. + + ``calibration`` and ``scorecard`` are optional, and the stored snapshot keeps + only the keys supplied. + """ + matched_at: ShortText map_model_version: LabelText calibration: ReportCalibration | None = None @@ -116,7 +150,6 @@ def stored_provisions( provisions: list[ReportProvision] | None, ) -> list[StoredJSONObject] | None: """Serialize provisions exactly as supplied, enforcing the 64 KiB cap.""" - if provisions is None: return None stored = [ @@ -132,6 +165,7 @@ def stored_provisions( def stored_validation( validation: ReportValidation | None, ) -> StoredJSONObject | None: + """Serialize the validation snapshot with only the keys supplied, or return None.""" if validation is None: return None return validation.model_dump(mode="json", exclude_unset=True) @@ -179,12 +213,20 @@ class UserReportCreationInput(StrictUserReportModel): class UserReportUpdateInput(_ReportPresentationFields): + """A partial update of a saved report; only supplied fields change. + + At least one field must be supplied. ``legacy_report_id`` repoints the saved + report to another v1 report output and cannot be cleared; the saved-reform + links are cleared with null. + """ + legacy_report_id: LegacyRecordId | None = None user_policy_id: ResourceId | None = None baseline_user_policy_id: ResourceId | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> UserReportUpdateInput: + """Require a supplied field and reject a null ``legacy_report_id``.""" mutable_fields = { "legacy_report_id", "name", @@ -207,6 +249,12 @@ def require_valid_supplied_field(self) -> UserReportUpdateInput: @dataclass(frozen=True) class UserReportRead: + """A stored user-report association as the service returns it to its owner. + + Unlike SharedUserReportRead, it includes the owner, the description and the + owner's saved-reform links. + """ + id: UUID country_id: str user_id: UUID @@ -256,6 +304,8 @@ class SharedUserReportRead: @dataclass(frozen=True) class UserReportPage: + """One page of associations in creation order; ``has_more`` says if more follow.""" + items: tuple[UserReportRead, ...] offset: int limit: int diff --git a/policyengine_api/services/v2/user_reports/validators.py b/policyengine_api/services/v2/user_reports/validators.py index 7b9cdc542..a66e29e7e 100644 --- a/policyengine_api/services/v2/user_reports/validators.py +++ b/policyengine_api/services/v2/user_reports/validators.py @@ -24,6 +24,7 @@ def require_user_report( *, association_id: UUID, ) -> UserReportAssociation: + """Return the association, or raise UserReportNotFoundError if it is None.""" if association is None: raise UserReportNotFoundError( f"user-report association {association_id} was not found" @@ -32,6 +33,11 @@ def require_user_report( def require_owned_user_policy(owned: bool, *, user_policy_id: UUID) -> None: + """Raise AssociationLinkNotFoundError unless a linked saved policy is ``owned``. + + ``owned`` says whether the report's owner has that saved policy in the + report's country. + """ if not owned: raise AssociationLinkNotFoundError( f"user-policy association {user_policy_id} was not found for this " diff --git a/policyengine_api/services/v2/user_simulations/database_connectors/creates.py b/policyengine_api/services/v2/user_simulations/database_connectors/creates.py index 2313ebcd8..9f91e6604 100644 --- a/policyengine_api/services/v2/user_simulations/database_connectors/creates.py +++ b/policyengine_api/services/v2/user_simulations/database_connectors/creates.py @@ -18,6 +18,7 @@ def create_user_simulation( session: Session, association_input: UserSimulationCreationInput, ) -> UserSimulationAssociation: + """Insert an association and return it with its database-generated fields loaded.""" association = UserSimulationAssociation(**association_input.model_dump()) session.add(association) session.flush() diff --git a/policyengine_api/services/v2/user_simulations/database_connectors/deletes.py b/policyengine_api/services/v2/user_simulations/database_connectors/deletes.py index e7d6f4a4e..af7b8a11c 100644 --- a/policyengine_api/services/v2/user_simulations/database_connectors/deletes.py +++ b/policyengine_api/services/v2/user_simulations/database_connectors/deletes.py @@ -14,5 +14,6 @@ def delete_user_simulation( session: Session, association: UserSimulationAssociation, ) -> None: + """Delete an association row; the simulation it names is not touched.""" session.delete(association) session.flush() diff --git a/policyengine_api/services/v2/user_simulations/database_connectors/reads.py b/policyengine_api/services/v2/user_simulations/database_connectors/reads.py index eb2bf2e9d..c33b7c57e 100644 --- a/policyengine_api/services/v2/user_simulations/database_connectors/reads.py +++ b/policyengine_api/services/v2/user_simulations/database_connectors/reads.py @@ -19,6 +19,7 @@ def read_owned_user_simulation_row( country_id: str, association_id: UUID, ) -> UserSimulationAssociation | None: + """Return the user's association with this id in this country, or None.""" return session.exec( select(UserSimulationAssociation).where( UserSimulationAssociation.id == association_id, @@ -37,6 +38,11 @@ def read_user_simulation_rows( offset: int, limit: int, ) -> list[UserSimulationAssociation]: + """Return up to ``limit + 1`` of a user's associations in creation order. + + With ``legacy_simulation_id``, only associations to that v1 simulation are + returned. The extra row lets the caller report whether another page exists. + """ statement = select(UserSimulationAssociation).where( UserSimulationAssociation.country_id == country_id, UserSimulationAssociation.user_id == user_id, @@ -60,6 +66,7 @@ def read_user_simulation_rows( def read_all_user_simulation_rows( session: Session, *, user_id: UUID ) -> list[UserSimulationAssociation]: + """Return every association the user owns, in every country, in creation order.""" return list( session.exec( select(UserSimulationAssociation) diff --git a/policyengine_api/services/v2/user_simulations/database_connectors/updates.py b/policyengine_api/services/v2/user_simulations/database_connectors/updates.py index 4bd87cb4b..2e11cb6fe 100644 --- a/policyengine_api/services/v2/user_simulations/database_connectors/updates.py +++ b/policyengine_api/services/v2/user_simulations/database_connectors/updates.py @@ -20,6 +20,7 @@ def update_user_simulation( association: UserSimulationAssociation, association_input: UserSimulationUpdateInput, ) -> UserSimulationAssociation: + """Apply only the supplied fields, bump ``updated_at`` and return the row.""" for field_name, value in association_input.model_dump(exclude_unset=True).items(): setattr(association, field_name, value) association.updated_at = utc_now() diff --git a/policyengine_api/services/v2/user_simulations/database_session.py b/policyengine_api/services/v2/user_simulations/database_session.py index 8048575c0..d229a3a17 100644 --- a/policyengine_api/services/v2/user_simulations/database_session.py +++ b/policyengine_api/services/v2/user_simulations/database_session.py @@ -13,15 +13,20 @@ class UserSimulationDatabaseSession: + """Open read sessions and transactions for user-simulation services.""" + def __init__(self, session_factory: sessionmaker[Session]) -> None: + """Use ``session_factory`` to open every session.""" self._session_factory = session_factory @contextmanager def read(self) -> Iterator[Session]: + """Yield a session for reads; nothing is committed.""" with self._session_factory() as session: yield session @contextmanager def transaction(self) -> Iterator[Session]: + """Yield a session whose transaction commits on exit and rolls back on error.""" with self._session_factory.begin() as session: yield session diff --git a/policyengine_api/services/v2/user_simulations/services.py b/policyengine_api/services/v2/user_simulations/services.py index f3c0a28d1..cadac3f5c 100644 --- a/policyengine_api/services/v2/user_simulations/services.py +++ b/policyengine_api/services/v2/user_simulations/services.py @@ -69,7 +69,6 @@ def read_owned_user_simulation( association_id: UUID, ) -> UserSimulationAssociation: """``user_id`` is None for an unauthenticated user that owns nothing yet.""" - return require_user_simulation( None if user_id is None @@ -87,7 +86,6 @@ def read_all_user_simulations( session: Session, *, user_id: UUID ) -> tuple[UserSimulationRead, ...]: """Return every simulation association owned by a user, across countries.""" - return tuple( user_simulation_read(row) for row in read_all_user_simulation_rows(session, user_id=user_id) @@ -102,6 +100,12 @@ def __init__( database_session: UserSimulationDatabaseSession, legacy_references: LegacyReferenceChecker | None = None, ) -> None: + """Open sessions through ``database_session``. + + ``legacy_references`` checks v1 simulations before a write. Without it, + creating a saved simulation or repointing one raises + LegacyReferenceCheckUnavailableError. + """ self._database_session = database_session self._legacy_references = legacy_references @@ -121,6 +125,12 @@ def create_user_simulation( principal: AccountPrincipal, association_input: UserSimulationCreateFields, ) -> UserSimulationRead: + """Save a v1 simulation for the principal, provisioning its v2 user if needed. + + The v1 simulation is checked first, in a read-only v1 transaction, and + LegacyReferenceNotFoundError is raised if it does not exist in the + country. + """ self._require_legacy_simulation( country_id=association_input.country_id, legacy_simulation_id=association_input.legacy_simulation_id, @@ -148,6 +158,11 @@ def get_user_simulation( country_id: str, association_id: UUID, ) -> UserSimulationRead: + """Return the principal's association, or raise UserSimulationNotFoundError. + + The association must be in this country. A signed-out browser that has + never saved anything owns nothing, and is not provisioned by the read. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -170,6 +185,12 @@ def list_user_simulations( offset: int = 0, limit: int = 100, ) -> UserSimulationPage: + """Return a page of the principal's associations in this country. + + With ``legacy_simulation_id``, only saves of that v1 simulation are + listed. A signed-out browser that has never saved anything gets an empty + page and is not provisioned. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_for_read_in_session( session, principal, country_id=country_id @@ -194,6 +215,12 @@ def patch_user_simulation( association_id: UUID, association_input: UserSimulationUpdateInput, ) -> UserSimulationRead: + """Update the supplied fields of one of the principal's associations. + + A new ``legacy_simulation_id`` is checked in v1 first, as on create. Raise + UserSimulationNotFoundError if the principal has no such association in + this country. + """ if association_input.legacy_simulation_id is not None: self._require_legacy_simulation( country_id=country_id, @@ -223,6 +250,10 @@ def delete_user_simulation( country_id: str, association_id: UUID, ) -> None: + """Delete one of the principal's associations in this country. + + Raise UserSimulationNotFoundError if the principal has no such association. + """ with self._database_session.transaction() as session: owner = resolve_principal_user_in_session( session, principal, country_id=country_id diff --git a/policyengine_api/services/v2/user_simulations/transformations.py b/policyengine_api/services/v2/user_simulations/transformations.py index 0b65a5f51..75216a5b6 100644 --- a/policyengine_api/services/v2/user_simulations/transformations.py +++ b/policyengine_api/services/v2/user_simulations/transformations.py @@ -20,6 +20,7 @@ def user_simulation_read( association: UserSimulationAssociation, ) -> UserSimulationRead: + """Convert a stored association row to the service's read model.""" return UserSimulationRead( id=association.id, country_id=association.country_id, @@ -39,6 +40,10 @@ def user_simulation_page( offset: int, limit: int, ) -> UserSimulationPage: + """Build a page from up to ``limit + 1`` rows. + + The extra row is not returned; it only sets ``has_more``. + """ return UserSimulationPage( items=tuple(user_simulation_read(row) for row in rows[:limit]), offset=offset, @@ -52,4 +57,5 @@ def owned_user_simulation_input( *, user_id: UUID, ) -> UserSimulationCreationInput: + """Add the owner resolved from the session principal to the caller's fields.""" return UserSimulationCreationInput(**fields.model_dump(), user_id=user_id) diff --git a/policyengine_api/services/v2/user_simulations/types.py b/policyengine_api/services/v2/user_simulations/types.py index 4fde05452..af75d4b3c 100644 --- a/policyengine_api/services/v2/user_simulations/types.py +++ b/policyengine_api/services/v2/user_simulations/types.py @@ -34,6 +34,8 @@ class UserSimulationCreateFields(StrictUserSimulationInput): class UserSimulationCreationInput(StrictUserSimulationInput): + """Complete persisted association values, including the resolved owner.""" + country_id: CountryId user_id: UserId legacy_simulation_id: LegacyRecordId @@ -42,12 +44,19 @@ class UserSimulationCreationInput(StrictUserSimulationInput): class UserSimulationUpdateInput(StrictUserSimulationInput): + """A partial update of a saved simulation; only supplied fields change. + + At least one field must be supplied. ``legacy_simulation_id`` repoints the + saved simulation to another v1 simulation and cannot be cleared. + """ + legacy_simulation_id: LegacyRecordId | None = None name: AssociationName | None = None description: AssociationDescription | None = None @model_validator(mode="after") def require_valid_supplied_field(self) -> UserSimulationUpdateInput: + """Require a supplied field and reject a null ``legacy_simulation_id``.""" supplied = self.model_fields_set.intersection( {"legacy_simulation_id", "name", "description"} ) @@ -63,6 +72,12 @@ def require_valid_supplied_field(self) -> UserSimulationUpdateInput: @dataclass(frozen=True) class UserSimulationRead: + """A stored user-simulation association as the service returns it to routes. + + Stage 12A saves only v1 simulations, so ``legacy_simulation_id`` names the + simulation and ``simulation_id`` is a logical reference with no foreign key. + """ + id: UUID country_id: str user_id: UUID @@ -76,6 +91,8 @@ class UserSimulationRead: @dataclass(frozen=True) class UserSimulationPage: + """One page of associations in creation order; ``has_more`` says if more follow.""" + items: tuple[UserSimulationRead, ...] offset: int limit: int diff --git a/policyengine_api/services/v2/user_simulations/validators.py b/policyengine_api/services/v2/user_simulations/validators.py index 4876b3b14..c98d0b54c 100644 --- a/policyengine_api/services/v2/user_simulations/validators.py +++ b/policyengine_api/services/v2/user_simulations/validators.py @@ -23,6 +23,7 @@ def require_user_simulation( *, association_id: UUID, ) -> UserSimulationAssociation: + """Return the association, or raise UserSimulationNotFoundError if it is None.""" if association is None: raise UserSimulationNotFoundError( f"user-simulation association {association_id} was not found" diff --git a/policyengine_api/spm.py b/policyengine_api/spm.py index 9919aa729..0222b472a 100644 --- a/policyengine_api/spm.py +++ b/policyengine_api/spm.py @@ -93,12 +93,14 @@ def serialize_selection( # noqa: ANN201 # pydantic would build the output sche @field_validator("as_of") @classmethod def validate_as_of(cls, value: str | None) -> str | None: + """Reject an ``as_of`` that is not an exact ISO date such as 2025-01-31.""" if value is not None and date.fromisoformat(value).isoformat() != value: raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") return value @model_validator(mode="after") def validate_location(self) -> Self: + """Require a non-blank ``geography_id`` for a metro; forbid it otherwise.""" if self.geography_kind == "metro": if not self.geography_id or not self.geography_id.strip(): raise ValueError("An SPM area selection requires geography_id") @@ -121,12 +123,14 @@ class SPMResolvedConfiguration(BaseModel): @field_validator("as_of") @classmethod def validate_as_of(cls, value: str | None) -> str | None: + """Reject an ``as_of`` that is not an exact ISO date such as 2025-01-31.""" if value is not None and date.fromisoformat(value).isoformat() != value: raise ValueError("as_of must be an ISO calendar date (YYYY-MM-DD)") return value @model_validator(mode="after") def validate_location(self) -> SPMResolvedConfiguration: + """Require a non-blank ``geography_id`` for a metro; forbid it otherwise.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("An SPM area selection requires geography_id") @@ -153,6 +157,7 @@ class SPMProvenance(BaseModel): @field_validator("years") @classmethod def validate_years(cls, years: tuple[str, ...]) -> tuple[str, ...]: + """Require sorted, unique four-digit year strings.""" if any( len(year) != 4 or not year.isascii() or not year.isdigit() for year in years ): @@ -163,6 +168,7 @@ def validate_years(cls, years: tuple[str, ...]) -> tuple[str, ...]: @model_validator(mode="after") def validate_location(self) -> SPMProvenance: + """Require a non-blank ``geography_id`` on a metro receipt, else forbid it.""" if self.geography_kind == "metro": if self.geography_id is None or not self.geography_id.strip(): raise ValueError("A metro SPM receipt requires geography_id") @@ -202,6 +208,7 @@ class SPMComparisonProvenance(BaseModel): @model_validator(mode="after") def validate_compatible_receipts(self) -> SPMComparisonProvenance: + """Reject a comparison whose baseline and reform used different SPM receipts.""" if self.baseline.receipt != self.reform.receipt: raise ValueError("baseline and reform SPM receipts differ") return self @@ -216,7 +223,6 @@ def build_spm_provenance( runtime_versions: SPMRuntimeVersions, ) -> SPMProvenance: """Build one compact receipt from typed calculation inputs.""" - if selection.forecast_content_sha256 != forecast_sha256: raise ValueError("SPM receipt differs from the SPM configuration") if selection.scenario is None: @@ -247,7 +253,6 @@ def build_spm_calculation_provenance( receipt: Mapping[str, object], ) -> SPMProvenance: """Reduce the country-owned rich receipt to the public compact contract.""" - forecast_sha256 = _required_string(receipt, "forecast_sha256") scenario = _required_string(receipt, "scenario") geography_kind = _required_string(receipt, "geography_kind") @@ -284,7 +289,6 @@ def resolved_spm_configuration( provenance: SPMProvenance, ) -> SPMResolvedConfiguration: """Derive the complete resolved selection from its compact receipt.""" - return SPMResolvedConfiguration( forecast_content_sha256=provenance.forecast_sha256, scenario=provenance.scenario, @@ -313,7 +317,6 @@ def build_spm_comparison_provenance( reform_receipts: Sequence[SPMProvenance], ) -> SPMComparisonProvenance: """Collapse matching child receipts into one comparison receipt.""" - return SPMComparisonProvenance( schema_version="canonical-spm-comparison-v2", baseline=_collapse_execution_receipts("baseline", baseline_receipts), @@ -325,11 +328,13 @@ class SPMValidationError(ValueError): """An SPM request/configuration error safe for the API's 4xx envelope.""" def __init__(self, code: str, message: str) -> None: + """Record a stable machine-readable error code and the client-facing message.""" self.code = code self.message = message super().__init__(message) def to_dict(self) -> dict[str, str]: + """Return the ``{"code", "message"}`` entry for a response's ``errors`` list.""" return {"code": self.code, "message": self.message} @@ -425,7 +430,6 @@ def validate_spm_calculation_provenance( provenance: object, ) -> SPMProvenance: """Validate the sole transported compact receipt for a completed result.""" - return SPMProvenance.model_validate(provenance) @@ -585,11 +589,12 @@ def normalize_spm_selection( def _validate_metro_selection(forecast: SPMForecast, selection: SPMSelection) -> None: - """A selected area must exist in the artifact, without guessing a year. + """Raise SPM_GEOGRAPHY_UNAVAILABLE unless the selected metro is in the artifact. - Area definitions may change between years. Accept an area available in any - covered year; the country owns validation for the actual calculation year. - This reads geography metadata and creates no household calculation receipt. + Area definitions may change between years, so this does not guess a year: it + accepts an area available in any covered year, and the country owns + validation for the actual calculation year. This reads geography metadata + and creates no household calculation receipt. """ for year in forecast.years: try: diff --git a/policyengine_api/utils/__init__.py b/policyengine_api/utils/__init__.py index 1aeb79c92..85ca5a173 100644 --- a/policyengine_api/utils/__init__.py +++ b/policyengine_api/utils/__init__.py @@ -1,3 +1,5 @@ +"""Shared helpers for caching, JSON values, current-law lookup and singletons.""" + from .cache_utils import * # noqa: F403 from .get_current_law import get_current_law_policy_id as get_current_law_policy_id from .json import * # noqa: F403 diff --git a/policyengine_api/utils/budget_window.py b/policyengine_api/utils/budget_window.py index 6acbd5b3a..d827c66eb 100644 --- a/policyengine_api/utils/budget_window.py +++ b/policyengine_api/utils/budget_window.py @@ -1,3 +1,5 @@ +"""Limits and request normalisation for multi-year budget-window calculations.""" + from dataclasses import dataclass from typing import Literal @@ -8,6 +10,8 @@ @dataclass(frozen=True) class BudgetWindowRequestSetup: + """A validated budget-window request, expanded into its years and cache period.""" + start_year: str window_size: int years: list[str] @@ -15,11 +19,13 @@ class BudgetWindowRequestSetup: def build_budget_window_years(*, start_year: str, window_size: int) -> list[str]: + """Return ``window_size`` consecutive years from ``start_year``, as strings.""" start_year_int = int(start_year) return [str(start_year_int + index) for index in range(window_size)] def build_budget_window_time_period(*, start_year: str, window_size: int) -> str: + """Return the ``budget_window::`` period naming the whole window.""" return f"budget_window:{start_year}:{window_size}" @@ -29,6 +35,11 @@ def build_budget_window_request_setup( window_size: int, target: Literal["general", "cliff"], ) -> BudgetWindowRequestSetup: + """Validate a budget-window request and expand it into its years. + + Raises ValueError for a target other than ``general``, a window outside + 1..BUDGET_WINDOW_MAX_YEARS, or a last year after BUDGET_WINDOW_MAX_END_YEAR. + """ if target != "general": raise ValueError("Budget-window calculations only support target='general'") diff --git a/policyengine_api/utils/cache_utils.py b/policyengine_api/utils/cache_utils.py index cbd02da9d..c5f8bf516 100644 --- a/policyengine_api/utils/cache_utils.py +++ b/policyengine_api/utils/cache_utils.py @@ -10,9 +10,10 @@ def make_cache_key(*args: object, **kwargs: object) -> str: # pylint: disable=unused-argument - """make a hash to uniquely identify a cache entry. - keep it fast, adding overhead to try to add some minor chance of a - cache hit is not worth it. + """Hash the current request's full path and JSON or form body into a cache key. + + The arguments are ignored. Keep it fast: adding overhead for a minor extra + chance of a cache hit is not worth it. Use a cryptographic digest (SHA-256) rather than the builtin `hash()`, whose output depends on PYTHONHASHSEED and is therefore diff --git a/policyengine_api/utils/deprecated_inputs.py b/policyengine_api/utils/deprecated_inputs.py index 1e1af7314..85cf7ee0b 100644 --- a/policyengine_api/utils/deprecated_inputs.py +++ b/policyengine_api/utils/deprecated_inputs.py @@ -38,6 +38,7 @@ class DeprecatedVariableWarning: @property def message(self) -> str: + """The user-facing warning naming the ignored input, its location and hint.""" location = f"`{self.entity_plural}/{self.entity_id}`" if self.entity_plural == "axes": location = f"`axes[{self.entity_id}].name`" diff --git a/policyengine_api/utils/get_current_law.py b/policyengine_api/utils/get_current_law.py index 929990461..1853570fc 100644 --- a/policyengine_api/utils/get_current_law.py +++ b/policyengine_api/utils/get_current_law.py @@ -1,4 +1,11 @@ +"""Lookup of each country's built-in current-law policy ID.""" + + def get_current_law_policy_id(country_id: str) -> int: + """Return the ID of the stored policy that represents current law in a country. + + Raises KeyError for a country without one (Israel has none here). + """ return { "uk": 1, "us": 2, diff --git a/policyengine_api/utils/input_validation.py b/policyengine_api/utils/input_validation.py index a74934197..af691f508 100644 --- a/policyengine_api/utils/input_validation.py +++ b/policyengine_api/utils/input_validation.py @@ -5,6 +5,11 @@ @dataclass(frozen=True) class UnrecognizedInput: + """One payload name the country model does not recognise, and where it appears. + + The entity plurals are set only for a variable placed on the wrong entity. + """ + input_type: str name: str path: str @@ -13,6 +18,7 @@ class UnrecognizedInput: @property def message(self) -> str: + """A one-sentence, user-facing description of the problem and where it is.""" if self.input_type == "household_entity": return ( f"Unrecognized household entity group `{self.name}` at `{self.path}`." @@ -32,6 +38,7 @@ def message(self) -> str: return f"Unrecognized household variable `{self.name}` at `{self.path}`." def to_dict(self) -> dict: + """Return the JSON error entry, omitting entity plurals that are not set.""" data = { "type": self.input_type, "name": self.name, @@ -51,7 +58,6 @@ def find_unrecognized_inputs( metadata: dict, ) -> list[UnrecognizedInput]: """Return calculate payload names that the country package cannot resolve.""" - return [ *find_unrecognized_household_inputs(household, metadata), *find_unrecognized_policy_inputs(policy, metadata), @@ -62,6 +68,12 @@ def find_unrecognized_household_inputs( household: dict, metadata: dict, ) -> list[UnrecognizedInput]: + """Find unknown entity groups, variables and axis variables in a household. + + Also reports a known variable placed on the wrong entity. Role membership + keys (such as ``members``) are not variables and are skipped, as are + sections that are not objects. + """ if not isinstance(household, dict): return [] @@ -135,6 +147,7 @@ def find_unrecognized_policy_inputs( policy: dict, metadata: dict, ) -> list[UnrecognizedInput]: + """Find policy keys that are not parameters in the country metadata.""" if not isinstance(policy, dict): return [] @@ -151,6 +164,7 @@ def find_unrecognized_policy_inputs( def format_unrecognized_inputs_message(errors: list[UnrecognizedInput]) -> str: + """Join the problems' messages into one error message for the response.""" return "Unrecognized calculate input(s): " + "; ".join( error.message for error in errors ) diff --git a/policyengine_api/utils/json.py b/policyengine_api/utils/json.py index 54f54e8d7..668803ef8 100644 --- a/policyengine_api/utils/json.py +++ b/policyengine_api/utils/json.py @@ -1,3 +1,5 @@ +"""Content hashing and JSON-safe conversion of policy, household and metadata values.""" + from __future__ import annotations import base64 @@ -13,6 +15,10 @@ def make_hashable( o: object, ) -> Any: # noqa: ANN401 # returns scalars unchanged, and callers sort them + """Convert nested lists, dicts and sets into tuples, sorting dict items and sets. + + Equal values therefore convert identically whatever their key or member order. + """ if isinstance(o, (tuple, list)): return tuple(make_hashable(e) for e in o) @@ -26,12 +32,20 @@ def make_hashable( def hash_object(o: object) -> str: + """Return a base64 SHA-256 digest of a value that ignores dict key order. + + It is the stored policy and household hash, and keys cached calculations. + """ hasher = hashlib.sha256() hasher.update(repr(make_hashable(o)).encode()) return base64.b64encode(hasher.digest()).decode() def get_safe_json(value: object) -> JSONValue: + """Return a JSON-safe copy of a metadata value, writing infinities as strings. + + Values of any other type (dates, None, objects) become None. + """ # This function is used when constructing metadata, # which will be consumed by JS, hence Infinity, not .inf if isinstance(value, (int, float)): diff --git a/policyengine_api/utils/payload_validators/__init__.py b/policyengine_api/utils/payload_validators/__init__.py index 6da6b53cb..ba014323d 100644 --- a/policyengine_api/utils/payload_validators/__init__.py +++ b/policyengine_api/utils/payload_validators/__init__.py @@ -1,3 +1,5 @@ +"""Validators for Flask route inputs: country IDs and policy and household payloads.""" + from .validate_country import validate_country as validate_country from .validate_household_payload import ( validate_household_payload as validate_household_payload, diff --git a/policyengine_api/utils/payload_validators/validate_country.py b/policyengine_api/utils/payload_validators/validate_country.py index 9fc4130c1..33f2e7a15 100644 --- a/policyengine_api/utils/payload_validators/validate_country.py +++ b/policyengine_api/utils/payload_validators/validate_country.py @@ -1,3 +1,5 @@ +"""Flask view decorator that rejects unsupported country IDs.""" + from __future__ import annotations from functools import wraps @@ -20,13 +22,19 @@ def validate_country( func: Callable[..., _ViewResult], ) -> Callable[..., _ViewResult | Response]: - """Validate that a country ID is valid. If not, return a 400 response. + """Wrap a Flask view so an unsupported ``country_id`` gets 400 before the view runs. + + The 400 body is the shared JSON error payload, sent with the legacy v1 + content type that native FastAPI routes must match. Args: - country_id (str): The country ID to validate. + func (Callable): The view to wrap; it takes ``country_id`` as its first + argument. Returns: - Response(400) if country is not valid, else continues + Callable: The wrapped view, which returns a 400 response for an + unsupported country and the view's own result otherwise. + """ @wraps(func) diff --git a/policyengine_api/utils/payload_validators/validate_household_payload.py b/policyengine_api/utils/payload_validators/validate_household_payload.py index 19a80a672..750c9874a 100644 --- a/policyengine_api/utils/payload_validators/validate_household_payload.py +++ b/policyengine_api/utils/payload_validators/validate_household_payload.py @@ -1,3 +1,5 @@ +"""Validation of the body of a v1 household create request.""" + from __future__ import annotations from typing import TYPE_CHECKING @@ -9,14 +11,14 @@ def validate_household_payload( payload: Mapping[str, object], ) -> tuple[bool, str | None]: - """ - Validate the payload for a POST request to set a household's input data. + """Validate the payload for a POST request to set a household's input data. Args: payload (dict): The payload to validate. Returns: tuple[bool, str]: A tuple containing a boolean indicating whether the payload is valid and a message. + """ # Check that all required keys are present required_keys = ["data"] diff --git a/policyengine_api/utils/payload_validators/validate_set_policy_payload.py b/policyengine_api/utils/payload_validators/validate_set_policy_payload.py index 6bb9002fd..9c9d2f098 100644 --- a/policyengine_api/utils/payload_validators/validate_set_policy_payload.py +++ b/policyengine_api/utils/payload_validators/validate_set_policy_payload.py @@ -1,4 +1,11 @@ +"""Validation of the body of a v1 policy create request.""" + + def validate_set_policy_payload(payload: dict) -> tuple[bool, str | None]: + """Check that a policy body has a ``data`` object and an optional string ``label``. + + Return ``(True, None)`` when valid, or ``(False, message)`` naming the problem. + """ # Check that all required keys are present required_keys = ["data"] missing_keys = [key for key in required_keys if key not in payload] diff --git a/policyengine_api/utils/singleton.py b/policyengine_api/utils/singleton.py index ecba9af47..bc801566f 100644 --- a/policyengine_api/utils/singleton.py +++ b/policyengine_api/utils/singleton.py @@ -1,7 +1,16 @@ +"""Metaclass for classes that have a single shared instance per process.""" + + class Singleton(type): + """Metaclass whose classes build one instance and return it on every later call.""" + _instances = {} def __call__(cls, *args: object, **kwargs: object) -> object: + """Return the shared instance, creating it on the first call. + + Arguments to later calls are ignored. + """ if cls not in cls._instances: cls._instances[cls] = super().__call__(*args, **kwargs) return cls._instances[cls] diff --git a/policyengine_api/utils/streaming_json.py b/policyengine_api/utils/streaming_json.py index f0589a525..488771a09 100644 --- a/policyengine_api/utils/streaming_json.py +++ b/policyengine_api/utils/streaming_json.py @@ -20,11 +20,13 @@ def should_stream_body(body: bytes) -> bool: + """Return True when a body is large enough that it must be sent chunked.""" return len(body) >= STREAMING_THRESHOLD_BYTES def iter_body_chunks( body: bytes, chunk_size: int = STREAMING_CHUNK_BYTES ) -> Iterator[bytes]: + """Yield a body in fixed-size slices for a chunked response.""" for start in range(0, len(body), chunk_size): yield body[start : start + chunk_size] diff --git a/policyengine_api/worker_spm.py b/policyengine_api/worker_spm.py index 498b5c5af..f0e446830 100644 --- a/policyengine_api/worker_spm.py +++ b/policyengine_api/worker_spm.py @@ -28,7 +28,9 @@ def get_spm_capability( /, *, policyengine_version: str | None, - ) -> object: ... + ) -> object: + """Return the worker's SPM capability document, or None if it has none.""" + ... def validate_worker_spm( @@ -39,6 +41,13 @@ def validate_worker_spm( policyengine_version: str | None = None, model_version: str | None = None, ) -> dict | None: + """Resolve an SPM selection and confirm the selected worker can execute it. + + Return the resolved selection, or None for a country without SPM support. + Raise SPMValidationError (SPM_CONFIGURATION_UNAVAILABLE) + when the worker advertises no canonical SPM capability, a different + contract, or a different forecast artifact from the API's. + """ country_id = country_id.lower() if isinstance(country_id, str) else country_id resolved = normalize_spm_selection(country_id, selection) if resolved is None: @@ -79,6 +88,14 @@ def validate_worker_result( expected_year: str | int | None = None, expected_years: list[str] | None = None, ) -> None: + """Check that a worker's economy result carries SPM receipts matching the request. + + An annual result must hold a comparison receipt whose baseline and reform + resolve to ``selection`` and cover ``expected_year``; a budget-window result + must have one valid annual result per requested year. Raise + SPMValidationError (SPM_CONFIGURATION_UNAVAILABLE) on any mismatch. Do nothing + when no selection was requested. + """ if selection is None: return try: diff --git a/policyengine_api/worker_spm_release.py b/policyengine_api/worker_spm_release.py index 17d430325..d5bee00c8 100644 --- a/policyengine_api/worker_spm_release.py +++ b/policyengine_api/worker_spm_release.py @@ -22,6 +22,7 @@ class RegistryCapability: """Expose one registry document through the request validator interface.""" def __init__(self, registry: dict) -> None: + """Wrap a parsed worker registry document.""" self.registry = registry def get_spm_capability( @@ -31,6 +32,11 @@ def get_spm_capability( *, policyengine_version: str | None, ) -> object: + """Return the registry's SPM capability for ``policyengine_version``, or None. + + The registry keys capabilities by PolicyEngine.py version alone, so + ``country`` and ``model_version`` are not consulted. + """ capabilities = self.registry.get("spm_capabilities") if not isinstance(capabilities, dict): return None diff --git a/pyproject.toml b/pyproject.toml index 62f3003a3..09d168f06 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -92,7 +92,12 @@ select = [ "TC", # Group 3: annotations. "ANN", + # Group 4: docstrings. + "D", ] +# Mutually exclusive pydocstyle pairs: keep D211 (no blank line before a class +# docstring) and D212 (summary on the first line). +ignore = ["D203", "D213"] [tool.ruff.lint.per-file-ignores] # Alembic revisions must stay exactly as Alembic generates them diff --git a/rules/fixtures/api.v2.user-services-no-sql/bad_services.py b/rules/fixtures/api.v2.user-services-no-sql/bad_services.py index 3ca5e0b0e..1fb46557b 100644 --- a/rules/fixtures/api.v2.user-services-no-sql/bad_services.py +++ b/rules/fixtures/api.v2.user-services-no-sql/bad_services.py @@ -6,4 +6,5 @@ def list_user_policies(session: Session, user_id: str) -> list[UserPolicy]: + """List the user's policies by querying the rows itself.""" return list(session.exec(select(UserPolicy).where(UserPolicy.user_id == user_id))) diff --git a/rules/fixtures/api.v2.user-services-no-sql/good_services.py b/rules/fixtures/api.v2.user-services-no-sql/good_services.py index 13257bf05..c09a0099e 100644 --- a/rules/fixtures/api.v2.user-services-no-sql/good_services.py +++ b/rules/fixtures/api.v2.user-services-no-sql/good_services.py @@ -7,4 +7,5 @@ def list_user_policies(session: Session, user_id: str) -> list[UserPolicy]: + """List the user's policies through the read connector.""" return reads.list_user_policies(session, user_id=user_id) diff --git a/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py b/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py index a093c9b3c..9680bbdcd 100644 --- a/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py +++ b/rules/fixtures/api.v2.validators-transformations-no-db/bad_validators.py @@ -6,6 +6,7 @@ def validate_policy_exists(session: Session, policy_id: int) -> Policy: + """Load the policy and fail if it doesn't exist.""" policy = session.exec(select(Policy).where(Policy.id == policy_id)).first() if policy is None: raise ValueError(policy_id) diff --git a/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py b/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py index 5b04ea204..aae3fa903 100644 --- a/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py +++ b/rules/fixtures/api.v2.validators-transformations-no-db/good_validators.py @@ -4,6 +4,7 @@ def validate_policy_exists(policy: Policy | None, policy_id: int) -> Policy: + """Fail if the already-loaded policy is missing.""" if policy is None: raise ValueError(policy_id) return policy diff --git a/rules/registry.toml b/rules/registry.toml index 2ac1579fc..11f342248 100644 --- a/rules/registry.toml +++ b/rules/registry.toml @@ -700,7 +700,7 @@ autofix = true [rule.enforcer] kind = "ruff" mode = "lint" -codes = ["ANN", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] +codes = ["ANN", "D", "E4", "E7", "E9", "F", "FLY", "FURB", "I", "PIE", "PYI", "RET", "RSE", "TC", "UP", "W"] [[rule.ci]] workflow = ".github/workflows/pr.yml" job = "lint" diff --git a/scripts/capture_migration_baseline.py b/scripts/capture_migration_baseline.py index 75e37dcd0..674871fcf 100644 --- a/scripts/capture_migration_baseline.py +++ b/scripts/capture_migration_baseline.py @@ -24,6 +24,8 @@ @dataclass(frozen=True) class RouteProbe: + """One API request to time, and the status codes that count as success.""" + name: str method: str path: str @@ -32,6 +34,8 @@ class RouteProbe: @dataclass(frozen=True) class ProbeResult: + """One timed request, or, when ``completed`` is set, one whole simulation run.""" + name: str method: str path: str @@ -105,6 +109,7 @@ def _request_probe( def run_probes(base_url: str, repetitions: int) -> list[ProbeResult]: + """Time each default probe ``repetitions`` times; transport errors are failures.""" results: list[ProbeResult] = [] with httpx.Client(base_url=base_url.rstrip("/"), timeout=90.0) as client: for _ in range(repetitions): @@ -127,6 +132,12 @@ def run_simulation_gateway_probe( poll_timeout_seconds: float, poll_interval_seconds: float, ) -> list[ProbeResult]: + """Submit one economy comparison to the gateway and poll it until it finishes. + + Return each submit and poll request followed by one completion row, whose + ``completed`` flag is False on failure or when the timeout (counted from + submission) expires. + """ results: list[ProbeResult] = [] started_at = time.perf_counter() @@ -220,6 +231,11 @@ def run_simulation_gateway_probe( def summarize(results: Iterable[ProbeResult]) -> dict: + """Summarise request latency, error rate and simulation completion times. + + Completion rows are counted apart from request rows, and latency + percentiles use successful requests only. + """ rows = list(results) request_rows = [row for row in rows if row.completed is None] latencies = [row.latency_ms for row in request_rows if row.ok] @@ -255,6 +271,7 @@ def summarize(results: Iterable[ProbeResult]) -> dict: def main(argv: list[str] | None = None) -> int: + """Print a JSON baseline summary; return 0 even when probes fail or are skipped.""" parser = argparse.ArgumentParser() parser.add_argument("--base-url", default=os.environ.get("API_BASE_URL")) parser.add_argument( diff --git a/scripts/compare_migration_baseline.py b/scripts/compare_migration_baseline.py index affda18a5..d3f320526 100644 --- a/scripts/compare_migration_baseline.py +++ b/scripts/compare_migration_baseline.py @@ -31,6 +31,11 @@ @dataclass(frozen=True) class ProbeStats: + """Error rate and p95 of one probe, or of the simulation completion pseudo-probe. + + For the completion pseudo-probe ``p95_latency_ms`` holds completion time. + """ + name: str request_count: int failure_count: int @@ -138,6 +143,7 @@ def compare( def main(argv: list[str] | None = None) -> int: + """Print the per-probe gate report; return 1 if any baseline probe fails.""" parser = argparse.ArgumentParser( description="Gate a candidate baseline JSON against a reference " "baseline JSON, per probe." diff --git a/scripts/export_migration_contracts.py b/scripts/export_migration_contracts.py index 9a0bd41d0..01e2d0c09 100644 --- a/scripts/export_migration_contracts.py +++ b/scripts/export_migration_contracts.py @@ -22,7 +22,6 @@ def build_payload() -> dict[str, Any]: """Build the migration contract payload from typed registries.""" - route_groups = [] for group in ROUTE_GROUPS: row = asdict(group) @@ -68,13 +67,11 @@ def build_payload() -> dict[str, Any]: def json_text(payload: dict[str, Any]) -> str: """Return the canonical JSON text for a migration contract payload.""" - return json.dumps(payload, indent=2, sort_keys=True) + "\n" def render_markdown(payload: dict[str, Any]) -> str: """Render migration contracts as reviewer- and agent-readable Markdown.""" - lines = [ "# Migration Contracts", "", @@ -147,7 +144,6 @@ def write_outputs( markdown_path: Path = DEFAULT_MARKDOWN, ) -> None: """Write generated migration contract artifacts.""" - json_path.parent.mkdir(parents=True, exist_ok=True) markdown_path.parent.mkdir(parents=True, exist_ok=True) json_path.write_text(json_text(payload)) @@ -155,6 +151,7 @@ def write_outputs( def main(argv: list[str] | None = None) -> int: + """Regenerate the migration contract JSON and Markdown, overwriting both files.""" parser = argparse.ArgumentParser() parser.add_argument("--json", default=str(DEFAULT_JSON)) parser.add_argument("--markdown", default=str(DEFAULT_MARKDOWN)) diff --git a/scripts/guards/migration_contracts.py b/scripts/guards/migration_contracts.py index 8156f828e..3018357a5 100644 --- a/scripts/guards/migration_contracts.py +++ b/scripts/guards/migration_contracts.py @@ -139,6 +139,7 @@ def _check_generated_artifacts(payload: dict[str, Any]) -> list[str]: def check() -> list[str]: + """Return every route-group, workflow and stale-artifact violation.""" payload = export_migration_contracts.build_payload() return [ *_check_route_groups(payload), @@ -148,6 +149,7 @@ def check() -> list[str]: def main() -> int: + """Print the guard's violations and return 1 if there are any.""" violations = check() if not violations: print("migration-contracts guard passed") diff --git a/scripts/measure_cloud_run_runtime.py b/scripts/measure_cloud_run_runtime.py index e269315b0..3d1c049d3 100644 --- a/scripts/measure_cloud_run_runtime.py +++ b/scripts/measure_cloud_run_runtime.py @@ -233,6 +233,7 @@ async def _run(args: argparse.Namespace) -> dict: def main() -> int: + """Wait for readiness, run the load cell, and print its JSON summary to stdout.""" parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--base-url", required=True) parser.add_argument("--cell-label", required=True) diff --git a/scripts/process_v1_household_mirror_event.py b/scripts/process_v1_household_mirror_event.py index 8ffd8392e..61908eba0 100644 --- a/scripts/process_v1_household_mirror_event.py +++ b/scripts/process_v1_household_mirror_event.py @@ -45,7 +45,6 @@ def validate_command_environment( environ: Mapping[str, str], ) -> dict[str, str]: """Require matching source, destination, and operator declarations.""" - deployment_environment = _required(environ, "DEPLOYMENT_ENVIRONMENT") supabase_environment = _required(environ, "V2_SUPABASE_ENVIRONMENT") if deployment_environment != selected_environment: @@ -110,7 +109,6 @@ def _proxy_event_service( environ: Mapping[str, str], ) -> tuple[HouseholdService | None, Engine | None]: """Use the workflow-authenticated Cloud SQL proxy when it is selected.""" - raw_port = environ.get("POLICYENGINE_DB_PROXY_PORT", "").strip() if not raw_port: return None, None @@ -162,7 +160,6 @@ def _record_command_outcome( error_type: str | None = None, ) -> None: """Write an audit record without source data or credentials.""" - logger.log_struct( { "message": "Explicit v1 household create-event processing completed", @@ -197,6 +194,11 @@ def main( *, environ: Mapping[str, str] | None = None, ) -> int: + """Mirror the selected pending event and print a JSON outcome; return 1 on failure. + + Error output names only configuration problems; any other failure is reported + by exception type, so source data and credentials never reach the log. + """ args = _parser().parse_args(argv) if args.legacy_household_id < 0: _parser().error("--legacy-household-id must be non-negative") diff --git a/scripts/rules/adapters.py b/scripts/rules/adapters.py index b22bd112a..19a7ee979 100644 --- a/scripts/rules/adapters.py +++ b/scripts/rules/adapters.py @@ -18,6 +18,10 @@ def check_adapters(agents: Agents, root: Path) -> list[Violation]: + """Return a violation for each way an adapter stops being a short pointer. + + Missing adapter files are skipped rather than reported. + """ violations = [] instructions = Path(agents.instructions).name for adapter in agents.adapters: diff --git a/scripts/rules/baseline.py b/scripts/rules/baseline.py index f1a8aeb4b..5bc3feb90 100644 --- a/scripts/rules/baseline.py +++ b/scripts/rules/baseline.py @@ -18,12 +18,16 @@ class BaseUnavailable(Exception): def default_base_ref() -> str | None: + """Return ``origin/`` on a pull-request run, otherwise None.""" base = os.environ.get("GITHUB_BASE_REF") return f"origin/{base}" if base else None def load_base_registry(root: Path, ref: str) -> Registry | None: - """The registry as it is on ``ref``, or None if the base has no registry yet.""" + """Return the registry as it is on ``ref``, or None if the base has none yet. + + Raise BaseUnavailable when ``ref`` does not name a commit. + """ verify = subprocess.run( ["git", "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"], cwd=root, diff --git a/scripts/rules/ci_scripts.py b/scripts/rules/ci_scripts.py index 5a0a6a066..0a0edecd2 100644 --- a/scripts/rules/ci_scripts.py +++ b/scripts/rules/ci_scripts.py @@ -36,6 +36,8 @@ @dataclass class Findings: + """What one ci-scripts run found, as baseline-style entries for each check.""" + multiline: set[str] = field(default_factory=set) untested: set[str] = field(default_factory=set) shellcheck: set[str] | None = None # None when shellcheck could not run @@ -43,8 +45,11 @@ class Findings: def multiline_runs(config: CiScripts, root: Path) -> set[str]: - """``workflow :: job :: step`` for every run: block with more than one line of - code, or with commands chained by ``&&``, ``||`` or ``;``.""" + """Return ``workflow :: job :: step`` for each run: block that is not one command. + + A block fails when it has more than one line of code, or chains commands with + ``&&``, ``||`` or ``;``. + """ found = set() for workflow in sorted(root.glob(config.workflows)): relative = workflow.relative_to(root).as_posix() @@ -66,7 +71,7 @@ def called_scripts(config: CiScripts, root: Path) -> set[str]: def untested_scripts(config: CiScripts, root: Path) -> set[str]: - """Called scripts whose file name appears in no test file.""" + """Return the called scripts whose file name appears in no test file.""" tests = [ path.read_text(errors="replace") for pattern in config.test_globs @@ -101,6 +106,7 @@ def shellcheck_failures(config: CiScripts, root: Path) -> set[str] | None: def find(config: CiScripts, root: Path) -> Findings: + """Run the three checks on freshly parsed workflows, noting a skipped shellcheck.""" clear_caches() findings = Findings( multiline=multiline_runs(config, root), @@ -138,6 +144,10 @@ def _rule_for(registry: Registry, baseline: str, check: str) -> str | None: def compare(registry: Registry, findings: Findings, root: Path) -> list[Violation]: + """Report findings missing from their baseline and baseline entries now fixed. + + A check that could not run is skipped rather than reported. + """ config = registry.ci_scripts assert config is not None checks = ( @@ -204,6 +214,10 @@ def write_baselines(registry: Registry, findings: Findings, root: Path) -> None: def run_check(root: Path = REPO_ROOT) -> tuple[list[Violation], list[str]]: + """Load the registry and return the ci-scripts violations and notes. + + A missing shellcheck is a violation on GitHub Actions and only a note elsewhere. + """ registry, violations = load_registry(root) if registry is None or registry.ci_scripts is None: return violations, [ diff --git a/scripts/rules/cli.py b/scripts/rules/cli.py index 9322bb086..a4134cacb 100644 --- a/scripts/rules/cli.py +++ b/scripts/rules/cli.py @@ -24,6 +24,10 @@ def _print(violations: list, notes: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int: + """Run one rules command and return 1 if it found violations. + + ``render`` refuses to write anything while the registry fails validation. + """ parser = argparse.ArgumentParser( prog="python -m scripts.rules", description=__doc__ ) diff --git a/scripts/rules/discover.py b/scripts/rules/discover.py index 323425618..a352c05ef 100644 --- a/scripts/rules/discover.py +++ b/scripts/rules/discover.py @@ -60,6 +60,7 @@ def _discovered_tests(registry: Registry, root: Path) -> set[str]: def discover_hidden(registry: Registry, root: Path) -> list[Violation]: + """Return a hidden-rule violation for each check no rule references or ignores.""" violations = [] rules = registry.rules diff --git a/scripts/rules/guard.py b/scripts/rules/guard.py index c68fce02b..3efa3d4f5 100644 --- a/scripts/rules/guard.py +++ b/scripts/rules/guard.py @@ -38,6 +38,8 @@ @dataclass class GuardResult: + """Guard violations, which fail the run, and notes, which only inform.""" + violations: list[Violation] = field(default_factory=list) notes: list[str] = field(default_factory=list) @@ -75,6 +77,11 @@ def run_guard( today: date | None = None, compare_base: bool = True, ) -> GuardResult: + """Run every registry check, then compare baselines with the base ref. + + Stop after the schema errors if the registry does not load. An unreadable + base ref is a violation on GitHub Actions and only a note elsewhere. + """ clear_caches() result = GuardResult() registry, schema_violations = load_registry(root) diff --git a/scripts/rules/load.py b/scripts/rules/load.py index 83b455e2e..936a3a9f4 100644 --- a/scripts/rules/load.py +++ b/scripts/rules/load.py @@ -389,6 +389,7 @@ def parse(raw: dict[str, Any]) -> tuple[Registry | None, list[Violation]]: def load_text(text: str) -> tuple[Registry | None, list[Violation]]: + """Parse registry TOML text, or return None and one violation if it is not TOML.""" try: raw = tomllib.loads(text) except tomllib.TOMLDecodeError as error: @@ -404,6 +405,7 @@ def load_text(text: str) -> tuple[Registry | None, list[Violation]]: def load_registry(root: Path) -> tuple[Registry | None, list[Violation]]: + """Load the registry under ``root``, or return None and a violation if missing.""" path = root / REGISTRY_PATH if not path.exists(): return None, [ diff --git a/scripts/rules/model.py b/scripts/rules/model.py index 6d8f9c365..401becaaa 100644 --- a/scripts/rules/model.py +++ b/scripts/rules/model.py @@ -12,12 +12,16 @@ class Tier(StrEnum): + """How mechanically a rule can be checked; judgment-tier rules never block.""" + MECHANICAL = "mechanical" METRIC = "metric" JUDGMENT = "judgment" class Status(StrEnum): + """How a rule is enforced today; only enforced and ratcheted rules can block.""" + ENFORCED = "enforced" RATCHETED = "ratcheted" ADVISORY = "advisory" @@ -25,6 +29,8 @@ class Status(StrEnum): class Kind(StrEnum): + """The checker behind a rule; it decides which ``[rule.enforcer]`` fields apply.""" + PYTEST = "pytest" QUALITY_GUARD = "quality-guard" RUFF = "ruff" @@ -47,12 +53,15 @@ class Violation: fix: str def __str__(self) -> str: + """Format as ``[rule] check: message. Fix: fix.``; rule defaults to registry.""" subject = self.rule_id or "registry" return f"[{subject}] {self.check}: {self.message}. Fix: {self.fix}." @dataclass(frozen=True) class Category: + """A ``[[category]]`` grouping rules, in declared order, in the rendered views.""" + id: str title: str @@ -67,6 +76,8 @@ class Ignored: @dataclass(frozen=True) class Discovery: + """Where the guard looks for unregistered checks, and which ones it may skip.""" + workflows: tuple[str, ...] = (".github/workflows/pr.yml",) guard_test_globs: tuple[str, ...] = () extra_guard_tests: tuple[str, ...] = () @@ -97,6 +108,8 @@ class CiScripts: @dataclass(frozen=True) class Enforcer: + """What checks a rule: its ``kind`` plus the raw ``[rule.enforcer]`` fields.""" + kind: Kind fields: Mapping[str, Any] = field(default_factory=dict) @@ -105,6 +118,7 @@ def get( key: str, default: object = None, ) -> Any: # noqa: ANN401 # a raw TOML value that callers read by its shape + """Return a raw enforcer field, or ``default`` when the registry omits it.""" return self.fields.get(key, default) @property @@ -153,6 +167,8 @@ def expects_proof(self) -> bool: @dataclass(frozen=True) class CiRef: + """A ``[[rule.ci]]`` reference: a workflow job whose run text has ``invocation``.""" + workflow: str job: str invocation: str @@ -160,17 +176,26 @@ class CiRef: @dataclass(frozen=True) class Proof: + """The raw ``[rule.proof]`` fields; which keys are required depends on the tier.""" + fields: Mapping[str, Any] @dataclass(frozen=True) class Baseline: + """A ratcheted rule's file of tolerated violations and the number it lists. + + The count must match the file and may never grow relative to the base branch. + """ + file: str count: int @dataclass(frozen=True) class RuleException: + """A temporary carve-out of ``path`` from a rule; the guard fails once expired.""" + path: str reason: str expires: date @@ -178,6 +203,8 @@ class RuleException: @dataclass(frozen=True) class Rule: + """One ``[[rule]]`` registry entry; rules-registry.md explains each field.""" + id: str category: str statement: str @@ -205,6 +232,8 @@ def blocking(self) -> bool: @dataclass(frozen=True) class Registry: + """The parsed contents of rules/registry.toml.""" + schema_version: int repo: str categories: tuple[Category, ...] @@ -214,6 +243,7 @@ class Registry: ci_scripts: CiScripts | None = None def rule(self, rule_id: str) -> Rule | None: + """Return the rule with ``rule_id``, or None if the registry has none.""" for rule in self.rules: if rule.id == rule_id: return rule diff --git a/scripts/rules/render.py b/scripts/rules/render.py index 97815bf30..2b70361e4 100644 --- a/scripts/rules/render.py +++ b/scripts/rules/render.py @@ -52,7 +52,7 @@ def _proof_label(rule: Rule) -> str: def _enforced_by(rule: Rule) -> str: - """What enforces a rule, naming the CI job when the job itself is the check.""" + """Describe what enforces a rule, naming the CI job when the job is the check.""" if rule.enforcer.kind is Kind.WORKFLOW_JOB: jobs = ", ".join(f"`{ref.job}`" for ref in rule.ci) return f"CI job {jobs}" @@ -121,6 +121,7 @@ def _bullets(items: list[str]) -> list[str]: def render_rules_md(registry: Registry) -> str: + """Render rules.md: the tier summary, the gaps, then each category's rules.""" rules = sorted(registry.rules, key=lambda r: r.id) out = [ "# Rules", @@ -211,6 +212,11 @@ def _agents_line(rule: Rule) -> str: def render_agents_block(registry: Registry) -> str: + """Render the AGENTS.md rules block, markers included. + + Past ``AGENTS_COLLAPSE_AT`` rules, each category lists its checked rules on + one line and spells out only the doc-only ones. + """ out = [ BEGIN, "", @@ -245,6 +251,7 @@ def splice_block(text: str, block: str) -> str: def extract_block(text: str) -> str | None: + """Return the generated rules block in ``text``, markers included, or None.""" start, end = text.find(BEGIN), text.find(END) if start == -1 or end == -1 or end < start: return None @@ -252,6 +259,7 @@ def extract_block(text: str) -> str | None: def write_outputs(registry: Registry, root: Path) -> list[Path]: + """Write rules.md and splice the rules block into AGENTS.md; return both paths.""" rules_md = root / RULES_MD rules_md.parent.mkdir(parents=True, exist_ok=True) rules_md.write_text(render_rules_md(registry)) diff --git a/scripts/rules/resolve.py b/scripts/rules/resolve.py index cddb6536d..ed2bba559 100644 --- a/scripts/rules/resolve.py +++ b/scripts/rules/resolve.py @@ -34,6 +34,7 @@ def _read(path: Path) -> str | None: def clear_caches() -> None: + """Forget cached file reads and parses, so the next check sees the files afresh.""" _read.cache_clear() workflow_jobs.cache_clear() test_names.cache_clear() @@ -41,6 +42,7 @@ def clear_caches() -> None: @cache def workflow_jobs(path: Path) -> tuple[dict[str, JobBlock], tuple[str, ...]]: + """Return a workflow's jobs by id and its parse problems; cached per path.""" text = _read(path) if text is None: return {}, (f"{path.name} is missing",) @@ -97,7 +99,10 @@ def node_exists(root: Path, node: str) -> str | None: def guard_names(root: Path, runner: str = DEFAULT_RUNNER) -> list[str] | None: - """String names in the ``GUARDS`` tuple of the quality-guard runner.""" + """Return the string names in the runner's ``GUARDS`` tuple. + + Return None if the runner is missing, and an empty list if it has no ``GUARDS``. + """ text = _read(root / runner) if text is None: return None @@ -131,7 +136,7 @@ def _toml(path: Path) -> dict[str, Any]: def ruff_lint_config(root: Path) -> dict[str, Any]: - """The ``lint`` table from ruff.toml, or ``[tool.ruff.lint]`` in pyproject.""" + """Return ruff.toml's ``lint`` table, else ``[tool.ruff.lint]`` in pyproject.""" ruff_toml = _toml(root / "ruff.toml") if ruff_toml: return ruff_toml.get("lint", {}) @@ -141,12 +146,14 @@ def ruff_lint_config(root: Path) -> dict[str, Any]: def ruff_selected(root: Path) -> tuple[str, ...]: + """Return ruff's ``select`` (its default when unset) plus ``extend-select``.""" lint = ruff_lint_config(root) selected = tuple(lint.get("select", RUFF_DEFAULT_SELECT)) return selected + tuple(lint.get("extend-select", ())) def mypy_files(root: Path) -> list[str]: + """Return ``[tool.mypy].files`` as a list, even when it is a single string.""" files = ( _toml(root / "pyproject.toml").get("tool", {}).get("mypy", {}).get("files", []) ) @@ -323,6 +330,7 @@ def _resolve_workflow_job(rule: Rule, root: Path) -> list[Violation]: def resolve_enforcer(rule: Rule, root: Path) -> list[Violation]: + """Check that the rule's enforcer exists where its kind says it should.""" return RESOLVERS[rule.enforcer.kind](rule, root) @@ -332,7 +340,7 @@ def resolve_enforcer(rule: Rule, root: Path) -> list[Violation]: def _job_for( rule_id: str, ref: CiRef, root: Path, check: str, field: str ) -> tuple[JobBlock | None, list[Violation]]: - """The job a CI reference names, if it exists and runs the invocation.""" + """Return the job a CI reference names if it runs the invocation, plus problems.""" def v(message: str, fix: str) -> Violation: return Violation(rule_id, check, message, fix) @@ -386,12 +394,13 @@ def resolve_ci(rule: Rule, root: Path) -> list[Violation]: def resolve_proof_ci(rule: Rule, root: Path) -> list[Violation]: - """A proof test must run in CI, or nothing shows the checker still works. + """Check that a CI job collects the rule's proof test. - The proof test itself asserts that the bad fixture is rejected and the good - one accepted; the guard checks that some CI job collects it. A proof runs - in the rule's own CI job when the rule is pytest-enforced; otherwise the - proof names its job with ``workflow``, ``job`` and ``invocation``. + Without that, nothing shows the checker still works. The proof test itself + asserts that the bad fixture is rejected and the good one accepted; the + guard checks that some CI job collects it. A proof runs in the rule's own + CI job when the rule is pytest-enforced; otherwise the proof names its job + with ``workflow``, ``job`` and ``invocation``. """ if rule.proof is None or not isinstance(rule.proof.fields.get("test"), str): return [] @@ -434,7 +443,11 @@ def _test_coverage( root: Path, check: str, ) -> list[Violation]: - """The job's test command actually collects these root-relative test nodes.""" + """Check that the job's test command collects these root-relative test nodes. + + Only ``make`` targets and ``unittest discover`` commands are inspected; any + other command is trusted. + """ violations = [] collected: list[str] | None = None described = "" diff --git a/scripts/rules/validate.py b/scripts/rules/validate.py index f7d117d65..406b40ef6 100644 --- a/scripts/rules/validate.py +++ b/scripts/rules/validate.py @@ -55,6 +55,7 @@ def baseline_entries(path: Path) -> int: + """Count the entries in a baseline file, skipping blank and ``#`` comment lines.""" return sum( 1 for line in path.read_text().splitlines() @@ -242,6 +243,10 @@ def _validate_proof(rule: Rule, root: Path) -> list[Violation]: def validate(registry: Registry, root: Path, today: date) -> list[Violation]: + """Check the registry's fields and file references against the tree at ``root``. + + ``today`` decides whether review dates and exceptions have expired. + """ violations: list[Violation] = [] def v(rule_id: str | None, message: str, fix: str, check: str = "schema") -> None: diff --git a/scripts/rules/workflows.py b/scripts/rules/workflows.py index a432455fd..5c7c039b4 100644 --- a/scripts/rules/workflows.py +++ b/scripts/rules/workflows.py @@ -32,11 +32,14 @@ class RunBlock: @property def single_line(self) -> bool: + """Whether the block is at most one line of code with no chained commands.""" return self.code_lines <= 1 and not self.chained @dataclass(frozen=True) class JobBlock: + """One workflow job: its run text, step names, ``uses:`` references and matrix.""" + id: str name: str run_text: str diff --git a/scripts/run_quality_guards.py b/scripts/run_quality_guards.py index 5da7c5fb5..a89b7bd3b 100644 --- a/scripts/run_quality_guards.py +++ b/scripts/run_quality_guards.py @@ -24,6 +24,7 @@ def main() -> int: + """Run every guard, print its violations, and return 1 if any guard failed.""" failed = False for name, check in GUARDS: violations = check() diff --git a/scripts/smoke_test_local_boot.py b/scripts/smoke_test_local_boot.py index 5dc8fa2aa..eca2fd272 100644 --- a/scripts/smoke_test_local_boot.py +++ b/scripts/smoke_test_local_boot.py @@ -26,18 +26,21 @@ def find_free_port() -> int: + """Return a localhost TCP port that was free when checked (nothing reserves it).""" with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: sock.bind(("127.0.0.1", 0)) return int(sock.getsockname()[1]) def read_text(path: Path) -> str: + """Return the file's text with undecodable bytes replaced, or "" if missing.""" if not path.exists(): return "" return path.read_text(encoding="utf-8", errors="replace") def fetch(url: str) -> tuple[int, str]: + """Return the status and body of a GET; error statuses raise urllib's HTTPError.""" request = urllib.request.Request(url, method="GET") with urllib.request.urlopen(request, timeout=HTTP_TIMEOUT_SECONDS) as response: body = response.read().decode("utf-8", errors="replace") @@ -45,11 +48,16 @@ def fetch(url: str) -> tuple[int, str]: def print_section(title: str, body: str) -> None: + """Print a titled block of captured output, or "(empty)" when there is none.""" print(f"\n=== {title} ===") print(body.rstrip() if body else "(empty)") def main() -> int: + """Boot the Flask app on a free port and return 1 unless both health checks pass. + + On failure print the app's captured stdout and stderr; always stop the app. + """ port = find_free_port() stdout_file = tempfile.NamedTemporaryFile( mode="w+", diff --git a/scripts/v1_database_migration.py b/scripts/v1_database_migration.py index e57194a3e..aa63effc3 100644 --- a/scripts/v1_database_migration.py +++ b/scripts/v1_database_migration.py @@ -28,6 +28,12 @@ class DatabaseState(StrEnum): + """Where a v1 database stands against the Alembic revision scripts. + + UNVERSIONED has no version table and INVALID has an empty one; neither is + ever stamped or upgraded automatically. + """ + UNVERSIONED = "unversioned" INVALID = "invalid" PENDING = "pending" @@ -43,7 +49,6 @@ def build_database_url( database: str, ) -> URL: """Build a SQLAlchemy URL without stringifying its credentials.""" - return URL.create( drivername="mysql+pymysql", username=username, @@ -60,6 +65,10 @@ def classify_database_state( current_heads: set[str], script_heads: set[str] | None = None, ) -> DatabaseState: + """Classify a database by its version table and heads. + + Without ``script_heads`` a versioned database is reported as pending, never head. + """ if not version_table_exists: return DatabaseState.UNVERSIONED if not current_heads: @@ -71,7 +80,6 @@ def classify_database_state( def describe_metadata_difference(difference: object) -> str: """Return a stable, data-free description of an Alembic metadata diff.""" - item = difference if isinstance(item, list) and len(item) == 1: item = item[0] @@ -104,6 +112,7 @@ def _script_heads(config: Config) -> set[str]: def database_state(connection: Connection) -> DatabaseState: + """Return the database's Alembic state against the repository's revision heads.""" inspector = inspect(connection) version_table_exists = "alembic_version" in inspector.get_table_names() context = MigrationContext.configure(connection) @@ -116,6 +125,10 @@ def database_state(connection: Connection) -> DatabaseState: def metadata_differences(connection: Connection) -> set[str]: + """Describe every difference between the live schema and the v1 models. + + Raise ValueError for a kind of difference that has no data-free description. + """ context = MigrationContext.configure( connection, opts={"compare_type": True, "compare_server_default": True}, @@ -127,6 +140,7 @@ def metadata_differences(connection: Connection) -> set[str]: def verify_head_schema(connection: Connection) -> None: + """Raise RuntimeError unless the database is at the head with no schema drift.""" state = database_state(connection) if state is not DatabaseState.HEAD: raise RuntimeError(f"Database is not at the Alembic head; found {state}") @@ -152,6 +166,11 @@ def _release_lock(connection: Connection) -> None: def upgrade_database(connection: Connection, *, backup_id: str) -> None: + """Upgrade a versioned database to head under a MySQL named lock, then verify it. + + A database already at head is only verified, so ``backup_id`` is required + only when an upgrade will actually run. + """ _acquire_lock(connection) try: state = database_state(connection) @@ -211,6 +230,11 @@ def _parse_args(argv: list[str] | None = None) -> argparse.Namespace: def main(argv: list[str] | None = None) -> int: + """Run one --mode against the v1 database: state, verify-head or upgrade. + + Unless a database URL is set, the read-only modes connect through the proxy + as the schema reader and upgrade as the migrator, in one transaction. + """ args = _parse_args(argv) engine = create_engine(_database_target(args.mode), poolclass=NullPool) try: diff --git a/tests/conftest.py b/tests/conftest.py index 5bf715308..7b77c0c8f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -21,7 +21,6 @@ @pytest.fixture def mock_v2_runtime_settings(): """Avoid external v2 configuration access in ordinary API route tests.""" - with patch( "policyengine_api.data.v2.settings.load_v2_runtime_database_settings", return_value=object(), diff --git a/tests/env_variables/test_environment_variables.py b/tests/env_variables/test_environment_variables.py index 631f6baca..de4175d7b 100644 --- a/tests/env_variables/test_environment_variables.py +++ b/tests/env_variables/test_environment_variables.py @@ -27,7 +27,6 @@ class TestEnvironmentVariables: ) def test_hugging_face_token(self): """Test if HUGGING_FACE_TOKEN is valid by querying Hugging Face API.""" - token = os.getenv("HUGGING_FACE_TOKEN") assert token is not None, "HUGGING_FACE_TOKEN is not set" @@ -89,7 +88,6 @@ def test_hugging_face_token(self): ) def test_github_microdata_auth_token(self): """Test if POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN is valid by querying GitHub user API.""" - token = os.getenv("POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN") assert token is not None, "POLICYENGINE_GITHUB_MICRODATA_AUTH_TOKEN is not set" diff --git a/tests/fixtures/integration/simulations.py b/tests/fixtures/integration/simulations.py index 22845e338..383976941 100644 --- a/tests/fixtures/integration/simulations.py +++ b/tests/fixtures/integration/simulations.py @@ -1,6 +1,4 @@ -""" -Test fixtures and constants for axes calculation tests -""" +"""Households, axes configurations and builders for the axes calculation tests.""" from policyengine_api.services.household_calculation_service import ( add_yearly_variables, @@ -45,21 +43,21 @@ # Pytest fixtures def create_base_household(): - """Returns a copy of the base household configuration""" + """Return a deep copy of BASE_HOUSEHOLD that a test may mutate freely.""" import copy return copy.deepcopy(BASE_HOUSEHOLD) def create_small_axes(): - """Returns the small axes configuration for faster tests""" + """Return a deep copy of the small axes configuration, which keeps tests fast.""" import copy return copy.deepcopy(SMALL_AXES_CONFIG) def create_household_with_axes(base_household, axes_config): - """Helper function to add axes to a household""" + """Return a deep copy of ``base_household`` with ``axes_config`` as its axes.""" import copy household = copy.deepcopy(base_household) @@ -68,6 +66,6 @@ def create_household_with_axes(base_household, axes_config): def setup_small_axes_household(base_household, small_axes_config): - """Fixture to setup a household with small axes for testing""" + """Add the axes and the test country's yearly variables to a base household.""" household_with_axes = create_household_with_axes(base_household, small_axes_config) return add_yearly_variables(household_with_axes, TEST_COUNTRY_ID) diff --git a/tests/fixtures/libs/simulation_entrypoint.py b/tests/fixtures/libs/simulation_entrypoint.py index de5e508f2..6cbdc5961 100644 --- a/tests/fixtures/libs/simulation_entrypoint.py +++ b/tests/fixtures/libs/simulation_entrypoint.py @@ -1,5 +1,4 @@ -""" -Test fixtures for SimulationAPIModal. +"""Test fixtures for SimulationAPIModal. This module provides mock data, fixtures, and helper functions for testing the Modal simulation API client. @@ -163,8 +162,7 @@ def create_mock_httpx_response( json_data: dict = None, headers: dict | None = None, ): - """ - Helper function to create a mock httpx response. + """Build a mock httpx response for the Simulation Entrypoint client tests. Parameters ---------- @@ -172,11 +170,14 @@ def create_mock_httpx_response( HTTP status code for the response. json_data : dict JSON data to return from response.json(). + headers : dict or None + Response headers; defaults to a mock observability-id header. Returns ------- MagicMock A mock httpx response object. + """ mock_response = MagicMock() mock_response.status_code = status_code @@ -211,8 +212,7 @@ def mock_modal_env_url(): @pytest.fixture def mock_httpx_client(): - """ - Mock httpx.Client for testing SimulationAPIModal. + """Mock httpx.Client for testing SimulationAPIModal. Returns a mock client that can be configured for different responses. """ diff --git a/tests/fixtures/services/economy_service.py b/tests/fixtures/services/economy_service.py index 0ec3481ee..418def7ce 100644 --- a/tests/fixtures/services/economy_service.py +++ b/tests/fixtures/services/economy_service.py @@ -223,7 +223,7 @@ def create_mock_reform_impact( time_period=MOCK_TIME_PERIOD, message=None, ): - """Helper function to create mock reform impact records.""" + """Build a ReformImpact record populated with the shared mock identifiers.""" default_reform_impact_json = json.dumps( { **MOCK_REFORM_IMPACT_DATA, @@ -266,8 +266,7 @@ def create_mock_modal_execution( policyengine_bundle=None, run_progress=None, ): - """ - Helper function to create mock Modal execution objects. + """Build a mock Modal simulation execution. Parameters ---------- @@ -279,11 +278,16 @@ def create_mock_modal_execution( The simulation result if complete. error : str or None The error message if failed. + policyengine_bundle : dict or None + The runtime bundle the execution reports; defaults to MOCK_RUNTIME_BUNDLE. + run_progress : SimulationRunProgress or None + The stage fields the execution reports while it runs. Returns ------- MagicMock A mock ModalSimulationExecution object. + """ mock_execution = MagicMock() mock_execution.job_id = job_id @@ -308,7 +312,7 @@ def create_mock_budget_window_batch_execution( result=None, error=None, ): - """Helper function to create mock batch execution objects.""" + """Build a mock budget-window batch execution; year lists default to empty.""" mock_execution = MagicMock() mock_execution.batch_job_id = batch_job_id mock_execution.name = batch_job_id diff --git a/tests/fixtures/v2_catalog.py b/tests/fixtures/v2_catalog.py index 948b9d0e7..6e69f789b 100644 --- a/tests/fixtures/v2_catalog.py +++ b/tests/fixtures/v2_catalog.py @@ -62,7 +62,6 @@ def _region( def bundle(*, policyengine_version: str = POLICYENGINE_VERSION) -> dict: """Return a minimal packaged-bundle manifest.""" - return { "bundle_version": policyengine_version, "policyengine_version": policyengine_version, @@ -158,7 +157,6 @@ def source_models( policyengine_version: str = POLICYENGINE_VERSION, ) -> dict[str, SimpleNamespace]: """Return US and UK public-model fixtures.""" - us_default_uri = "hf://policyengine/us/populace_us_2024.h5@fixture" uk_default_uri = "hf://policyengine/uk/enhanced_frs_2024_25.h5@fixture" return { @@ -218,7 +216,6 @@ def source_models( def installed_version(name: str) -> str: """Return fixture observed versions.""" - return DEPENDENCY_VERSIONS[name] @@ -227,7 +224,6 @@ def normalized_catalog( policyengine_version: str = POLICYENGINE_VERSION, ) -> NormalizedCatalog: """Return the complete deterministic normalized fixture.""" - return extract_catalog( bundle=bundle(policyengine_version=policyengine_version), policyengine_version=policyengine_version, diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index dc838e439..9f4b532f7 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -17,7 +17,6 @@ @pytest.fixture(scope="session") def disposable_v1_database_url() -> str: """Return a validated local MySQL URL for destructive integration tests.""" - from sqlalchemy.engine import make_url database_url = os.environ.get("ALEMBIC_DATABASE_URL", "") @@ -39,7 +38,6 @@ def disposable_v1_database_url() -> str: @pytest.fixture(scope="session") def disposable_v2_database_url() -> str: """Return a validated disposable PostgreSQL URL for integration tests.""" - from policyengine_api.data.v2.migration_target import ( V2_ALEMBIC_DISPOSABLE_TEST, load_v2_alembic_settings, diff --git a/tests/integration/test_alembic_mysql_lifecycle.py b/tests/integration/test_alembic_mysql_lifecycle.py index a4935ec1c..6817297b0 100644 --- a/tests/integration/test_alembic_mysql_lifecycle.py +++ b/tests/integration/test_alembic_mysql_lifecycle.py @@ -43,7 +43,6 @@ def _deployed_question_table() -> Table: """Describe the orphaned production table without adding it to ORM metadata.""" - metadata = MetaData() return Table( "question", diff --git a/tests/integration/test_alembic_v2_lifecycle.py b/tests/integration/test_alembic_v2_lifecycle.py index e7840bba7..9baff0f02 100644 --- a/tests/integration/test_alembic_v2_lifecycle.py +++ b/tests/integration/test_alembic_v2_lifecycle.py @@ -456,7 +456,6 @@ def _row_values( ``head_table`` adds the columns head newly requires; only rows of the parent's read-only tables, which new code writes, may need them. """ - values = { column.name: _test_value(table, column, sequence) for column in table.columns @@ -495,7 +494,6 @@ def _insert_and_read_back( Rows of the parent's read-only tables are inserted through head's table, as only new code writes them. """ - target = insert_table if insert_table is not None else table connection.execute(target.insert().values(values)) # Selecting the reflected table reads every column the parent maps. @@ -508,7 +506,6 @@ def _insert_and_read_back( def _rewrite_mapped_columns(connection, table: sa.Table, row: dict) -> None: """Update every non-key column the parent maps, as an ORM flush can.""" - key_columns = {column.name for column in table.primary_key.columns} | { element.parent.name for foreign_key in table.foreign_key_constraints @@ -537,7 +534,6 @@ def _assert_user_deletion_cascades( sequence: Iterator[int], ) -> None: """Delete, through the parent's users table, a user who owns Stage 12A rows.""" - users = parent.tables["users"] owner = _insert_and_read_back( connection, users, _row_values(users, stored_rows, sequence) @@ -589,7 +585,6 @@ def test_stage_12a_schema_serves_the_previous_revision() -> None: cascades run. Head may not give any table the parent revision writes a new required column without a server default. """ - database_url = _disposable_url() config = _config() engine = create_engine(database_url) @@ -889,7 +884,6 @@ def test_baseline_region_default_enforces_same_model_dataset() -> None: def test_stage_12a_reshape_refuses_retained_dormant_association_rows() -> None: """A retained pre-12A report association stops the reshape unchanged.""" - database_url = _disposable_url() config = _config() engine = create_engine(database_url) diff --git a/tests/integration/test_cloud_run_candidate.py b/tests/integration/test_cloud_run_candidate.py index beb497efc..25eb77aab 100644 --- a/tests/integration/test_cloud_run_candidate.py +++ b/tests/integration/test_cloud_run_candidate.py @@ -70,7 +70,6 @@ def test_cloud_run_candidate_stage6_read_route_contracts(api_client): def test_cloud_run_candidate_user_routes_fail_closed_without_a_token(api_client): """User-owned routes must never become readable without a credential.""" - for path in ("/v2/me", "/v2/user-policies", "/v2/user-reports"): for headers in ({}, {"PolicyEngine-Unauthenticated-User": "not-a-uuid"}): response = api_client.get( diff --git a/tests/integration/test_live_v2_policies.py b/tests/integration/test_live_v2_policies.py index 9a94f142f..b80792a2f 100644 --- a/tests/integration/test_live_v2_policies.py +++ b/tests/integration/test_live_v2_policies.py @@ -109,7 +109,6 @@ def _v1_id(api_client, path: str, body: dict, field: str) -> int: def _create_v1_artifacts(api_client) -> dict[str, int]: """Create (or reuse deduplicated) v1 artifacts without running compute.""" - policy_id = _v1_id(api_client, "policy", {"data": {}}, "policy_id") household_id = _v1_id( api_client, @@ -149,8 +148,7 @@ def _delete_account(api_client, headers: dict[str, str]) -> None: def _browser_headers() -> dict[str, str]: - """A fresh signed-out browser; its UUID never leaves the request headers.""" - + """Return headers for a fresh signed-out browser, whose UUID stays in them.""" return {UNAUTHENTICATED_USER_HEADER: str(uuid4())} @@ -409,7 +407,6 @@ def test_live_unauthenticated_users_round_trip_and_upgrade( Every run uses fresh browser UUIDs and starts and ends by deleting the test user's account, so it leaves nothing behind and can run repeatedly. """ - params = {"country_id": COUNTRY_ID} legacy_policy_id = _v1_id(api_client, "policy", {"data": {}}, "policy_id") converted_browser = _browser_headers() diff --git a/tests/integration/test_runtime_cache_redis.py b/tests/integration/test_runtime_cache_redis.py index ba332c7b7..941b4953d 100644 --- a/tests/integration/test_runtime_cache_redis.py +++ b/tests/integration/test_runtime_cache_redis.py @@ -94,9 +94,12 @@ def test_cached_household_calculation_is_shared_between_connections( def test_cached_calculation_receipts_survive_a_real_round_trip(redis_pair) -> None: - """The selection is part of cache identity and its receipts travel with the + """Key cached calculations by SPM selection and keep their receipts. + + The selection is part of cache identity and its receipts travel with the result, so a replay under the same measurement reads back a complete receipt - and a replay under a different one misses.""" + and a replay under a different one misses. + """ first, second, namespace = redis_pair selection = { "forecast_content_sha256": "c" * 64, diff --git a/tests/integration/test_simulations.py b/tests/integration/test_simulations.py index 7bc19e2e3..64c53aa2e 100644 --- a/tests/integration/test_simulations.py +++ b/tests/integration/test_simulations.py @@ -14,7 +14,7 @@ class TestSimsWithAxes: def test__given_any_number_of_axes__sim_returns_valid_arrays( self, ): # , patched_add_yearly_variables, patched_countries_get): - """Integration test with small axes for speed""" + """Calculate a household with small axes, which keeps the test fast.""" base_household = create_base_household() small_axes_config = create_small_axes() household_with_axes = setup_small_axes_household( diff --git a/tests/integration/test_v1_schema_metadata_compatibility.py b/tests/integration/test_v1_schema_metadata_compatibility.py index ed16a0247..f08fdc6ae 100644 --- a/tests/integration/test_v1_schema_metadata_compatibility.py +++ b/tests/integration/test_v1_schema_metadata_compatibility.py @@ -10,7 +10,6 @@ def compare_existing_schema(database_url: str) -> list[str]: """Return metadata drift without mutating the target database.""" - engine = create_engine(database_url) try: with engine.connect() as connection: diff --git a/tests/integration/test_v2_user_identity_persistence.py b/tests/integration/test_v2_user_identity_persistence.py index eaf74a249..e0f0f694d 100644 --- a/tests/integration/test_v2_user_identity_persistence.py +++ b/tests/integration/test_v2_user_identity_persistence.py @@ -148,8 +148,7 @@ def _principal(sessions, subject: str) -> UserPrincipal: def _browser(sessions) -> UnauthenticatedPrincipal: - """A fresh signed-out browser whose user the fixture removes afterwards.""" - + """Return a fresh signed-out browser whose user the fixture removes afterwards.""" user_id = uuid4() sessions.kw["info"]["browser_user_ids"].append(user_id) return UnauthenticatedPrincipal(user_id=user_id) @@ -425,7 +424,6 @@ def test_merge_moves_every_association_and_keeps_links_and_constraints( def _hold_a_conversion(sessions, *, browser, principal) -> Session: """Convert or merge ``browser`` in an open transaction holding its lock.""" - session = sessions() session.begin() convert_unauthenticated_user_in_session( diff --git a/tests/to_refactor/conftest.py b/tests/to_refactor/conftest.py index fcd3b7580..b33f66ebf 100644 --- a/tests/to_refactor/conftest.py +++ b/tests/to_refactor/conftest.py @@ -32,7 +32,6 @@ def _seed_current_law_policies(engine: Engine) -> None: @pytest.fixture(scope="session", autouse=True) def isolated_legacy_v1_database() -> Iterator[None]: """Bind legacy tests to explicit in-memory persistence, never Cloud SQL.""" - engine = create_engine( "sqlite+pysqlite:///:memory:", connect_args={"check_same_thread": False}, diff --git a/tests/to_refactor/python/test_error_routes.py b/tests/to_refactor/python/test_error_routes.py index 09c2058da..34b24fbc4 100644 --- a/tests/to_refactor/python/test_error_routes.py +++ b/tests/to_refactor/python/test_error_routes.py @@ -26,7 +26,7 @@ def client(app): def test_404_handler(app, client): - """Test 404 Not Found error handling""" + """Map NotFound to a 404 error envelope that keeps its message.""" @app.route("/nonexistent") def nonexistent(): @@ -42,7 +42,7 @@ def nonexistent(): def test_400_handler(app, client): - """Test 400 Bad Request error handling""" + """Map BadRequest to a 400 error envelope that keeps its message.""" @app.route("/bad-request") def bad_request(): @@ -58,7 +58,7 @@ def bad_request(): def test_401_handler(app, client): - """Test 401 Unauthorized error handling""" + """Map Unauthorized to a 401 error envelope that keeps its message.""" @app.route("/unauthorized") def unauthorized(): @@ -74,7 +74,7 @@ def unauthorized(): def test_403_handler(app, client): - """Test 403 Forbidden error handling""" + """Map Forbidden to a 403 error envelope that keeps its message.""" @app.route("/forbidden") def forbidden(): @@ -90,7 +90,7 @@ def forbidden(): def test_500_handler(app, client): - """Test 500 Internal Server Error handling""" + """Map InternalServerError to a 500 error envelope that keeps its message.""" @app.route("/server-error") def server_error(): @@ -106,7 +106,7 @@ def server_error(): def test_generic_exception_handler(app, client): - """Test handling of generic exceptions""" + """Map any other exception to a 500 error envelope that keeps its message.""" @app.route("/generic-error") def generic_error(): diff --git a/tests/to_refactor/python/test_validate_country.py b/tests/to_refactor/python/test_validate_country.py index f3495ffe8..389603d37 100644 --- a/tests/to_refactor/python/test_validate_country.py +++ b/tests/to_refactor/python/test_validate_country.py @@ -5,17 +5,20 @@ @validate_country def foo(country_id, other): - """ - A simple dummy test method for validation testing. Must be defined outside of the class (or within - the test functions themselves) due to complications with the `self` parameter for class methods. + """Return "bar"; a stand-in view for exercising the decorator. + + It must be defined outside the class (or within the test functions themselves) + because the decorator treats the first argument as the country ID, and on a + method that would be `self`. """ return "bar" class TestValidateCountry: - """ - Test that the @validate_country decorator returns 404 if the country does not exist, otherwise - continues execution of the function. + """Tests for the @validate_country decorator. + + The decorator returns a 400 response if the country does not exist, and + otherwise calls the wrapped function. """ def test_valid_country(self): diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 7c8bd744e..7b2688dbf 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -26,7 +26,6 @@ def test_engine(): @pytest.fixture(autouse=True) def isolated_orm_database(test_engine, monkeypatch): """Bind runtime factories to a clean in-memory ORM database per test.""" - monkeypatch.setattr(orm, "get_v1_engine", lambda: test_engine) orm.clear_v1_session_factories() factory = orm.get_v1_session_factory() diff --git a/tests/unit/libs/test_auth0_management.py b/tests/unit/libs/test_auth0_management.py index 5ec317676..3e1daafbd 100644 --- a/tests/unit/libs/test_auth0_management.py +++ b/tests/unit/libs/test_auth0_management.py @@ -163,7 +163,6 @@ def record(request: httpx.Request) -> httpx.Response: def _tokens_then(delete_statuses: list[int]): """Issue numbered tokens and answer deletions with the given statuses.""" - issued_tokens: list[str] = [] remaining_statuses = list(delete_statuses) diff --git a/tests/unit/libs/test_simulation_entrypoint.py b/tests/unit/libs/test_simulation_entrypoint.py index bb0c452b2..bfe36c7e2 100644 --- a/tests/unit/libs/test_simulation_entrypoint.py +++ b/tests/unit/libs/test_simulation_entrypoint.py @@ -1,5 +1,4 @@ -""" -Unit tests for SimulationAPIModal class. +"""Unit tests for SimulationAPIModal class. Tests the selectable simulation API HTTP client functionality including job submission, status polling, and error handling. diff --git a/tests/unit/routes/test_metadata_routes_streaming.py b/tests/unit/routes/test_metadata_routes_streaming.py index c9d8b5c9e..ed2423c14 100644 --- a/tests/unit/routes/test_metadata_routes_streaming.py +++ b/tests/unit/routes/test_metadata_routes_streaming.py @@ -28,7 +28,6 @@ def _load_metadata_blueprint_with_fake_service(): import, so the route module is loaded with the fake installed first and the module cache is restored afterwards. """ - sentinel = object() original_route_module = sys.modules.get(ROUTE_MODULE, sentinel) original_service_module = sys.modules.get(SERVICE_MODULE, sentinel) diff --git a/tests/unit/services/test_economy_comparison.py b/tests/unit/services/test_economy_comparison.py index 3982731b3..07eba32dc 100644 --- a/tests/unit/services/test_economy_comparison.py +++ b/tests/unit/services/test_economy_comparison.py @@ -738,7 +738,7 @@ def _make_economy( people=None, decile_key="household_income_decile", ): - """Helper to build baseline/reform dicts for intra_decile tests.""" + """Build a baseline or reform economy dict for the intra-decile tests.""" n = len(incomes) return { "household_net_income": np.array(incomes, dtype=float), @@ -757,8 +757,10 @@ def test__income_change_formula_exact(self): class TestIntraDecileImpact: - """Tests for the intra_decile_impact function — verifying correct - percentage change calculation and bucket assignment.""" + """Tests for intra_decile_impact. + + They verify the percentage change calculation and the bucket assignment. + """ def test__5pct_gain_classified_below_5pct(self): """A uniform 5% income gain must NOT land in 'Gain more than 5%'. @@ -833,8 +835,10 @@ def test__no_change_classified_correctly(self): assert pct == 1.0 def test__near_zero_baseline_no_division_error(self): - """Households with near-zero baseline income should not cause - division errors — the floor of 1 handles this.""" + """Avoid division errors for households with zero baseline income. + + The floor of 1 on the denominator handles this. + """ baseline = _make_economy( incomes=[0.0] * 10, deciles=list(range(1, 11)), @@ -850,9 +854,11 @@ def test__near_zero_baseline_no_division_error(self): assert abs(total - 1.0) < 1e-9, f"Proportions should sum to 1, got {total}" def test__zero_baseline_uses_floor_of_one(self): - """When baseline income is 0, the max(B, 1) floor means the - effective denominator is 1. A $0 -> $100 change should give - income_change = 100/1 = 100 (10000%), landing in >5%.""" + """Divide by the floor of 1 when baseline income is 0. + + The max(B, 1) floor makes the effective denominator 1, so a $0 -> $100 + change gives income_change = 100/1 = 100 (10000%), landing in >5%. + """ baseline = _make_economy( incomes=[0.0] * 10, deciles=list(range(1, 11)), @@ -872,8 +878,10 @@ def test__zero_baseline_uses_floor_of_one(self): assert not np.isinf(result["all"][label]) def test__negative_baseline_handled(self): - """Households with negative baseline income should be handled - by the max(B, 1) floor without producing NaN or Inf.""" + """Handle negative baseline income without producing NaN or Inf. + + The max(B, 1) floor applies to negative incomes too. + """ baseline = _make_economy( incomes=[-500.0] * 10, deciles=list(range(1, 11)), @@ -889,8 +897,11 @@ def test__negative_baseline_handled(self): assert not np.isinf(result["all"][label]) def test__2pct_gain_is_not_doubled(self): - """A 2% gain stays in <5% even if doubled (2*2%=4% < 5%), so this - test alone would not catch the doubling bug.""" + """Keep a 2% gain in the <5% bucket. + + A 2% gain stays in <5% even if doubled (2*2%=4% < 5%), so this test + alone would not catch the doubling bug. + """ baseline = _make_economy( incomes=[50000.0] * 10, deciles=list(range(1, 11)), @@ -907,9 +918,11 @@ def test__2pct_gain_is_not_doubled(self): assert pct == 1.0, "2% gain not classified as <5%" def test__4pct_gain_not_doubled_into_above_5pct(self): - """A 4% gain must stay in <5%. With the doubling bug, 4% * 2 = 8% - would incorrectly land in >5%. This is the tightest regression - test for the doubling bug on the gain side.""" + """Keep a 4% gain in the <5% bucket. + + With the doubling bug, 4% * 2 = 8% would incorrectly land in >5%. This + is the tightest regression test for the doubling bug on the gain side. + """ baseline = _make_economy( incomes=[10000.0] * 10, deciles=list(range(1, 11)), @@ -943,8 +956,11 @@ def test__all_field_averages_deciles(self): class TestIntraWealthDecileImpact: - """Tests for intra_wealth_decile_impact — same formula, keyed by - wealth decile instead of income decile.""" + """Tests for intra_wealth_decile_impact. + + It uses the same formula as intra_decile_impact, keyed by wealth decile + instead of income decile. + """ def test__5pct_gain_classified_below_5pct(self): """Regression test: 5% gain must not be doubled into >5% bucket.""" diff --git a/tests/unit/services/test_metadata_service.py b/tests/unit/services/test_metadata_service.py index a8e3bda7d..25927e75f 100644 --- a/tests/unit/services/test_metadata_service.py +++ b/tests/unit/services/test_metadata_service.py @@ -65,14 +65,14 @@ def test_get_metadata_empty_country_id(self): def test_verify_metadata_for_given_country( self, country_id, current_law_id, test_regions ): - """ - Verifies metadata for a specific country contains expected values. + """Verifies metadata for a specific country contains expected values. Args: service: The MetadataService fixture country_id: Country identifier string current_law_id: Expected current law ID for this country test_regions: List of region codes that should be present for this country + """ # Get metadata for the country service = MetadataService() @@ -158,10 +158,7 @@ def test_default_dataset_uses_user_facing_label(self, country_id, expected_label ], ) def test_verify_region_types_for_given_country(self, country_id, expected_types): - """ - Verifies that all regions for UK and US have a 'type' field - with valid values. - """ + """Give every UK and US region a 'type' field with a valid value.""" service = MetadataService() metadata = service.get_metadata(country_id) diff --git a/tests/unit/test_country_spm.py b/tests/unit/test_country_spm.py index 3f618b7b0..2b7666539 100644 --- a/tests/unit/test_country_spm.py +++ b/tests/unit/test_country_spm.py @@ -46,7 +46,7 @@ def requested_household(variable="spm_unit_spm_threshold", *, axes=False): def _country_raising(monkeypatch, error): - """A country whose every variable calculation raises `error`.""" + """Return a country whose every variable calculation raises `error`.""" def calculate(variable, period): # This also represents a requested resource with a nested SPM dependency. @@ -94,8 +94,11 @@ def test_country_never_swallows_spm_input_errors_for_a_chosen_measurement( def test_inherited_measurement_leaves_dependent_variables_unavailable( monkeypatch, code, axes ): - """Nobody chose this measurement, so its dependants behave like any other - variable the model cannot compute: null, not a rejected request.""" + """Leave variables that depend on an inherited SPM measurement null. + + Nobody chose this measurement, so its dependants behave like any other + variable the model cannot compute: null, not a rejected request. + """ country, _ = _failing_country(monkeypatch, code) # The requested cell is seeded null, so returning at all is the observable # change; the sibling test below pins that the rest of the result survives. diff --git a/tests/unit/test_workflow_guard_scripts.py b/tests/unit/test_workflow_guard_scripts.py index 9283cc15b..0763e7f15 100644 --- a/tests/unit/test_workflow_guard_scripts.py +++ b/tests/unit/test_workflow_guard_scripts.py @@ -49,7 +49,7 @@ def test_require_git_ref_fails_with_an_annotation_on_another_ref() -> None: def _fake_gcloud(tmp_path: Path, rows: list[str], *, delete_exit: int = 0) -> Path: - """A gcloud that lists ``rows`` and records every other call.""" + """Return a bin dir whose fake gcloud lists ``rows`` and logs every other call.""" bin_dir = tmp_path / "bin" bin_dir.mkdir() listing = tmp_path / "listing.txt" diff --git a/tests/unit/utils/test_cache_utils.py b/tests/unit/utils/test_cache_utils.py index ddf0fac1c..0b75db3ad 100644 --- a/tests/unit/utils/test_cache_utils.py +++ b/tests/unit/utils/test_cache_utils.py @@ -63,9 +63,11 @@ def test_make_cache_key_is_sha256_hex(): def test_make_cache_key_stable_across_processes(): - """Two independent Python processes must produce the same cache - key for the same inputs, even though they use different - PYTHONHASHSEED values.""" + """Produce the same cache key in two independent Python processes. + + The processes use different PYTHONHASHSEED values, so the key must not + depend on Python's randomised string hashing. + """ script = textwrap.dedent( """ from flask import Flask diff --git a/tests/unit/v2/test_unauthenticated_user_service.py b/tests/unit/v2/test_unauthenticated_user_service.py index 0c2cf7f76..449ac5e6e 100644 --- a/tests/unit/v2/test_unauthenticated_user_service.py +++ b/tests/unit/v2/test_unauthenticated_user_service.py @@ -201,7 +201,6 @@ def _users(services: Services) -> list[User]: @contextmanager def _recorded_logs(*loggers: logging.Logger) -> Iterator[list[logging.LogRecord]]: """Record INFO logs; Alembic tests may have disabled these loggers.""" - records: list[logging.LogRecord] = [] class RecordingHandler(logging.Handler): diff --git a/tests/unit/v2/test_user_route_auth.py b/tests/unit/v2/test_user_route_auth.py index 882760986..6c57c6c8e 100644 --- a/tests/unit/v2/test_user_route_auth.py +++ b/tests/unit/v2/test_user_route_auth.py @@ -69,7 +69,6 @@ def _services() -> tuple[dict[str, RecordingService], dict[str, Any]]: def _user_operations() -> list[tuple[str, str]]: """Every signed-in operation, read from the built application's OpenAPI.""" - _services_by_name, factories = _services() client, _flask_calls = user_route_client(require_user=None, **factories) paths = client.get("/v2/openapi.json").json()["paths"] diff --git a/tests/unit/v2/user_route_support.py b/tests/unit/v2/user_route_support.py index 2f328d74a..c60298418 100644 --- a/tests/unit/v2/user_route_support.py +++ b/tests/unit/v2/user_route_support.py @@ -49,7 +49,6 @@ def user_route_client( Passing ``require_user=None`` keeps the fail-closed default authenticator. """ - flask_calls = {"count": 0} flask_app = Flask(__name__) diff --git a/tests/unit/v2/user_service_support.py b/tests/unit/v2/user_service_support.py index 15ed550c0..ce54237e3 100644 --- a/tests/unit/v2/user_service_support.py +++ b/tests/unit/v2/user_service_support.py @@ -111,7 +111,6 @@ class V2Store: @contextmanager def v2_store() -> Iterator[V2Store]: """Create an isolated SQLite v2 schema with mirrored v1 records.""" - engine = create_engine("sqlite://") @sa.event.listens_for(engine, "connect")