Skip to content

Commit d2a9c2c

Browse files
anth-volkanth-volk
authored
Merge pull request #109 from PolicyEngine/feat/standardize-country-id
Standardize all endpoints on country_id with latest-version defaulting
2 parents 6092796 + 9fd3030 commit d2a9c2c

39 files changed

Lines changed: 1067 additions & 370 deletions

‎CHANGELOG.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,14 @@
1+
0.4.0 (2026-03-13)
2+
3+
# Added
4+
5+
- Standardize all endpoints on `country_id` instead of `tax_benefit_model_name` (#109)
6+
- Default metadata endpoints (variables, parameters, datasets) to latest model version with optional version pinning (#109)
7+
- Dual policy IDs (`baseline_policy_id` / `reform_policy_id`) on reports and `EXECUTION_DEFERRED` report status (#109)
8+
- Auto-start `EXECUTION_DEFERRED` reports on GET endpoint access (#109)
9+
- Convert VARCHAR enum columns to native PostgreSQL enums with `values_callable` (#109)
10+
11+
112
0.3.1 (2026-03-11)
213

314
# Fixed

‎alembic/versions/20260309_62385cd8049d_rename_tax_benefit_model_name_to_.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
"""rename_tax_benefit_model_name_to_country_id
2+
3+
Revision ID: 62385cd8049d
4+
Revises: 886921687770
5+
Create Date: 2026-03-09 16:48:30.899791
6+
7+
"""
8+
9+
from typing import Sequence, Union
10+
11+
import sqlalchemy as sa
12+
import sqlmodel.sql.sqltypes
13+
14+
from alembic import op
15+
16+
# revision identifiers, used by Alembic.
17+
revision: str = "62385cd8049d"
18+
down_revision: Union[str, Sequence[str], None] = "886921687770"
19+
branch_labels: Union[str, Sequence[str], None] = None
20+
depends_on: Union[str, Sequence[str], None] = None
21+
22+
23+
def upgrade() -> None:
24+
"""Upgrade schema: rename tax_benefit_model_name → country_id with data migration."""
25+
# 1. Add country_id columns (nullable initially)
26+
op.add_column(
27+
"households",
28+
sa.Column("country_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
29+
)
30+
op.add_column(
31+
"household_jobs",
32+
sa.Column("country_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
33+
)
34+
35+
# 2. Populate country_id from tax_benefit_model_name
36+
op.execute("""
37+
UPDATE households SET country_id = CASE
38+
WHEN tax_benefit_model_name LIKE '%_us' OR tax_benefit_model_name LIKE '%-us' THEN 'us'
39+
WHEN tax_benefit_model_name LIKE '%_uk' OR tax_benefit_model_name LIKE '%-uk' THEN 'uk'
40+
ELSE 'us'
41+
END
42+
""")
43+
op.execute("""
44+
UPDATE household_jobs SET country_id = CASE
45+
WHEN tax_benefit_model_name LIKE '%_us' OR tax_benefit_model_name LIKE '%-us' THEN 'us'
46+
WHEN tax_benefit_model_name LIKE '%_uk' OR tax_benefit_model_name LIKE '%-uk' THEN 'uk'
47+
ELSE 'us'
48+
END
49+
""")
50+
51+
# 3. Make country_id non-nullable
52+
op.alter_column("households", "country_id", nullable=False)
53+
op.alter_column("household_jobs", "country_id", nullable=False)
54+
55+
# 4. Drop old columns
56+
op.drop_column("households", "tax_benefit_model_name")
57+
op.drop_column("household_jobs", "tax_benefit_model_name")
58+
59+
60+
def downgrade() -> None:
61+
"""Downgrade schema: restore tax_benefit_model_name from country_id."""
62+
# 1. Re-add tax_benefit_model_name columns (nullable initially)
63+
op.add_column(
64+
"households", sa.Column("tax_benefit_model_name", sa.VARCHAR(), nullable=True)
65+
)
66+
op.add_column(
67+
"household_jobs",
68+
sa.Column("tax_benefit_model_name", sa.VARCHAR(), nullable=True),
69+
)
70+
71+
# 2. Populate from country_id
72+
op.execute(
73+
"UPDATE households SET tax_benefit_model_name = 'policyengine_' || country_id"
74+
)
75+
op.execute(
76+
"UPDATE household_jobs SET tax_benefit_model_name = 'policyengine_' || country_id"
77+
)
78+
79+
# 3. Make non-nullable
80+
op.alter_column("households", "tax_benefit_model_name", nullable=False)
81+
op.alter_column("household_jobs", "tax_benefit_model_name", nullable=False)
82+
83+
# 4. Drop country_id columns
84+
op.drop_column("households", "country_id")
85+
op.drop_column("household_jobs", "country_id")

‎alembic/versions/20260310_f887cb5490bc_add_execution_deferred_to_reportstatus.py‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
"""add_execution_deferred_to_reportstatus
2+
3+
Revision ID: f887cb5490bc
4+
Revises: 62385cd8049d
5+
Create Date: 2026-03-10 21:27:32.072364
6+
7+
"""
8+
9+
from typing import Sequence, Union
10+
11+
from alembic import op
12+
13+
# revision identifiers, used by Alembic.
14+
revision: str = "f887cb5490bc"
15+
down_revision: Union[str, Sequence[str], None] = "62385cd8049d"
16+
branch_labels: Union[str, Sequence[str], None] = None
17+
depends_on: Union[str, Sequence[str], None] = None
18+
19+
20+
def upgrade() -> None:
21+
"""Add EXECUTION_DEFERRED value to the reportstatus enum."""
22+
op.execute("ALTER TYPE reportstatus ADD VALUE IF NOT EXISTS 'EXECUTION_DEFERRED'")
23+
24+
25+
def downgrade() -> None:
26+
"""Downgrade: PostgreSQL does not support removing enum values.
27+
28+
The 'EXECUTION_DEFERRED' value will remain in the enum type.
29+
To fully remove it, drop and recreate the type (requires migrating data).
30+
"""
31+
pass

‎alembic/versions/20260311_dac22a838dda_convert_varchar_enums_to_native_pg_enums.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
"""convert_varchar_enums_to_native_pg_enums
2+
3+
Revision ID: dac22a838dda
4+
Revises: f887cb5490bc
5+
Create Date: 2026-03-11 01:37:08.928795
6+
7+
"""
8+
9+
from typing import Sequence, Union
10+
11+
from alembic import op
12+
13+
# revision identifiers, used by Alembic.
14+
revision: str = "dac22a838dda"
15+
down_revision: Union[str, Sequence[str], None] = "f887cb5490bc"
16+
branch_labels: Union[str, Sequence[str], None] = None
17+
depends_on: Union[str, Sequence[str], None] = None
18+
19+
20+
def upgrade() -> None:
21+
"""Convert VARCHAR enum columns to native PostgreSQL enum types.
22+
23+
The enum types may already exist with UPPERCASE values (created by
24+
SQLAlchemy's default create_all behavior). Since the columns are still
25+
VARCHAR, the types are unused — drop and recreate with lowercase values
26+
matching the data and the values_callable convention.
27+
"""
28+
# Drop any pre-existing enum types (unused — columns are still VARCHAR)
29+
op.execute("DROP TYPE IF EXISTS regiontype CASCADE")
30+
op.execute("DROP TYPE IF EXISTS reporttype CASCADE")
31+
op.execute("DROP TYPE IF EXISTS deciletype CASCADE")
32+
33+
# Create PG enum types with lowercase values
34+
op.execute("""
35+
CREATE TYPE regiontype AS ENUM (
36+
'national', 'country', 'state', 'congressional_district',
37+
'constituency', 'local_authority', 'city', 'place'
38+
)
39+
""")
40+
op.execute("""
41+
CREATE TYPE reporttype AS ENUM (
42+
'economy_comparison', 'household_comparison', 'household_single'
43+
)
44+
""")
45+
op.execute("CREATE TYPE deciletype AS ENUM ('income', 'wealth')")
46+
47+
# Alter columns from VARCHAR to enum.
48+
# LOWER() handles any databases where values were previously uppercased.
49+
op.execute("""
50+
ALTER TABLE regions
51+
ALTER COLUMN region_type TYPE regiontype
52+
USING LOWER(region_type)::regiontype
53+
""")
54+
op.execute("""
55+
ALTER TABLE reports
56+
ALTER COLUMN report_type TYPE reporttype
57+
USING LOWER(report_type)::reporttype
58+
""")
59+
# decile_type has a VARCHAR default that must be dropped before type change
60+
op.execute("ALTER TABLE intra_decile_impacts ALTER COLUMN decile_type DROP DEFAULT")
61+
op.execute("""
62+
ALTER TABLE intra_decile_impacts
63+
ALTER COLUMN decile_type TYPE deciletype
64+
USING LOWER(decile_type)::deciletype
65+
""")
66+
op.execute(
67+
"ALTER TABLE intra_decile_impacts ALTER COLUMN decile_type SET DEFAULT 'income'::deciletype"
68+
)
69+
70+
71+
def downgrade() -> None:
72+
"""Revert native PG enum columns back to VARCHAR."""
73+
op.execute("""
74+
ALTER TABLE regions
75+
ALTER COLUMN region_type TYPE VARCHAR
76+
USING region_type::text
77+
""")
78+
op.execute("""
79+
ALTER TABLE reports
80+
ALTER COLUMN report_type TYPE VARCHAR
81+
USING report_type::text
82+
""")
83+
op.execute("""
84+
ALTER TABLE intra_decile_impacts
85+
ALTER COLUMN decile_type TYPE VARCHAR
86+
USING decile_type::text
87+
""")
88+
89+
# Drop the PG enum types
90+
op.execute("DROP TYPE IF EXISTS regiontype")
91+
op.execute("DROP TYPE IF EXISTS reporttype")
92+
op.execute("DROP TYPE IF EXISTS deciletype")

‎scripts/seed_regions.py‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
RegionDatasetLink,
3434
TaxBenefitModel,
3535
)
36+
from policyengine_api.models.region import RegionType # noqa: E402
3637

3738

3839
def _group_us_datasets(
@@ -195,7 +196,7 @@ def seed_us_regions(
195196
db_region = Region(
196197
code=pe_region.code,
197198
label=pe_region.label,
198-
region_type=pe_region.region_type,
199+
region_type=RegionType(pe_region.region_type),
199200
requires_filter=pe_region.requires_filter,
200201
filter_field=pe_region.filter_field,
201202
filter_value=pe_region.filter_value,
@@ -293,7 +294,7 @@ def seed_uk_regions(session: Session) -> tuple[int, int, int]:
293294
db_region = Region(
294295
code=pe_region.code,
295296
label=pe_region.label,
296-
region_type=pe_region.region_type,
297+
region_type=RegionType(pe_region.region_type),
297298
requires_filter=pe_region.requires_filter,
298299
filter_field=pe_region.filter_field,
299300
filter_value=pe_region.filter_value,

‎src/policyengine_api/agent_sandbox.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -53,9 +53,9 @@ def configure_logfire(traceparent: str | None = None):
5353
5454
## CRITICAL: Always filter by country
5555
56-
When searching for parameters or datasets, ALWAYS include tax_benefit_model_name:
57-
- "policyengine-uk" for UK questions
58-
- "policyengine-us" for US questions
56+
When searching for parameters or datasets, ALWAYS include country_id:
57+
- "uk" for UK questions
58+
- "us" for US questions
5959
6060
Parameters and datasets from both countries are in the same database. Without the filter, you'll get mixed results and waste turns finding the right ones.
6161
@@ -66,14 +66,14 @@ def configure_logfire(traceparent: str | None = None):
6666
- Poll GET /household/calculate/{job_id} until completed
6767
6868
2. **Parameter lookup**:
69-
- GET /parameters/?search=...&tax_benefit_model_name=policyengine-uk (ALWAYS include country filter)
69+
- GET /parameters/?search=...&country_id=uk (ALWAYS include country filter)
7070
- GET /parameter-values/?parameter_id=...&current=true for the current value
7171
7272
3. **Economic impact analysis** (budget impact, decile impacts):
73-
- GET /parameters/?search=...&tax_benefit_model_name=policyengine-uk to find parameter_id
73+
- GET /parameters/?search=...&country_id=uk to find parameter_id
7474
- POST /policies/ to create reform with parameter_values
75-
- GET /datasets/?tax_benefit_model_name=policyengine-uk to find dataset_id
76-
- POST /analysis/economic-impact with tax_benefit_model_name, policy_id and dataset_id
75+
- GET /datasets/?country_id=uk to find dataset_id
76+
- POST /analysis/economic-impact with country_id, policy_id and dataset_id
7777
- GET /analysis/economic-impact/{report_id} for results (includes decile_impacts and program_statistics)
7878
7979
## Response formatting

0 commit comments

Comments
 (0)