Migrate policy endpoints to new API structure

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: minor
+  changes:
+    changed:
+    - Refactored policy endpoints to use new API structure

policyengine_api/api.py

@@ -19,14 +19,13 @@
 from policyengine_api.routes.simulation_analysis_routes import (
     simulation_analysis_bp,
 )
+from policyengine_api.routes.policy_routes import policy_bp
 from policyengine_api.routes.tracer_analysis_routes import tracer_analysis_bp
 from policyengine_api.routes.metadata_routes import metadata_bp
 from policyengine_api.routes.user_profile_routes import user_profile_bp
 
 from .endpoints import (
     get_home,
-    get_policy,
-    set_policy,
     get_policy_search,
     get_household_under_policy,
     get_calculate,
@@ -61,9 +60,8 @@
 
 app.register_blueprint(household_bp)
 
-app.route("/<country_id>/policy/<policy_id>", methods=["GET"])(get_policy)
-
-app.route("/<country_id>/policy", methods=["POST"])(set_policy)
+# Routes for getting and setting a "policy" record
+app.register_blueprint(policy_bp)
 
 app.route("/<country_id>/policies", methods=["GET"])(get_policy_search)
 

policyengine_api/endpoints/__init__.py

@@ -4,8 +4,6 @@
     get_calculate,
 )
 from .policy import (
-    get_policy,
-    set_policy,
     get_policy_search,
     set_user_policy,
     get_user_policy,

policyengine_api/endpoints/policy.py

@@ -6,169 +6,6 @@
 from flask import Response, request
 
 
-@validate_country
-def get_policy(country_id: str, policy_id: int) -> dict:
-    """
-    Get policy data for a given country and policy ID.
-
-    Args:
-        country_id (str): The country ID.
-        policy_id (int): The policy ID.
-
-    Returns:
-        dict: The policy record.
-    """
-
-    # Get the policy record for a given policy ID.
-    row = database.query(
-        f"SELECT * FROM policy WHERE country_id = ? AND id = ?",
-        (country_id, policy_id),
-    ).fetchone()
-    if row is None:
-        response = dict(
-            status="error",
-            message=f"Policy #{policy_id} not found.",
-        )
-        return Response(
-            json.dumps(response),
-            status=404,
-            mimetype="application/json",
-        )
-    policy = dict(row)
-    policy["policy_json"] = json.loads(policy["policy_json"])
-    return dict(
-        status="ok",
-        message=None,
-        result=policy,
-    )
-
-
-@validate_country
-def set_policy(
-    country_id: str,
-) -> dict:
-    """
-    Set policy data for a given country and policy. If the policy already exists,
-    fail quietly by returning a 200, but passing a warning message and the previously
-    created policy
-
-    Args:
-        country_id (str): The country ID.
-    """
-
-    payload = request.json
-    label = payload.pop("label", None)
-    policy_json = payload.pop("data", None)
-    policy_hash = hash_object(policy_json)
-    api_version = COUNTRY_PACKAGE_VERSIONS.get(country_id)
-
-    # Check if policy already exists.
-    try:
-        # The following code is a workaround to the fact that
-        # SQLite's cursor method does not properly convert
-        # 'WHERE x = None' to 'WHERE x IS NULL'; though SQLite
-        # supports searching and setting with 'WHERE x IS y',
-        # the production MySQL does not, requiring this
-
-        # This workaround should be removed if and when a proper
-        # ORM package is added to the API, and this package's
-        # sanitization methods should be utilized instead
-        label_value = "IS NULL" if not label else "= ?"
-        args = [country_id, policy_hash]
-        if label:
-            args.append(label)
-
-        row = database.query(
-            f"SELECT * FROM policy WHERE country_id = ? AND policy_hash = ? AND label {label_value}",
-            tuple(args),
-        ).fetchone()
-    except Exception as e:
-        return Response(
-            json.dumps(
-                {
-                    "message": f"Internal database error: {e}; please try again later."
-                }
-            ),
-            status=500,
-            mimetype="application/json",
-        )
-
-    code = None
-    message = None
-    status = None
-    policy_id = None
-
-    if row is not None:
-        policy_id = str(row["id"])
-        message = (
-            "Warning: Record created previously with this label. To create "
-            + "a new record, change the submitted data's country ID, policy "
-            + "parameters, or label, and emit the request again"
-        )
-        status = "ok"
-        code = 200
-
-    else:
-        message = None
-        status = "ok"
-        code = 201
-
-        try:
-            database.query(
-                f"INSERT INTO policy (country_id, policy_json, policy_hash, label, api_version) VALUES (?, ?, ?, ?, ?)",
-                (
-                    country_id,
-                    json.dumps(policy_json),
-                    policy_hash,
-                    label,
-                    api_version,
-                ),
-            )
-
-            # The following code is a workaround to the fact that
-            # SQLite's cursor method does not properly convert
-            # 'WHERE x = None' to 'WHERE x IS NULL'; though SQLite
-            # supports searching and setting with 'WHERE x IS y',
-            # the production MySQL does not, requiring this
-
-            # This workaround should be removed if and when a proper
-            # ORM package is added to the API, and this package's
-            # sanitization methods should be utilized instead
-            label_value = "IS NULL" if not label else "= ?"
-            args = [country_id, policy_hash]
-            if label:
-                args.append(label)
-
-            policy_id = database.query(
-                f"SELECT id FROM policy WHERE country_id = ? AND policy_hash = ? AND label {label_value}",
-                (tuple(args)),
-            ).fetchone()["id"]
-
-        except Exception as e:
-            return Response(
-                json.dumps(
-                    {
-                        "message": f"Internal database error: {e}; please try again later."
-                    }
-                ),
-                status=500,
-                mimetype="application/json",
-            )
-
-    response_body = dict(
-        status=status,
-        message=message,
-        result=dict(
-            policy_id=policy_id,
-        ),
-    )
-    return Response(
-        json.dumps(response_body),
-        status=code,
-        mimetype="application/json",
-    )
-
-
 @validate_country
 def get_policy_search(country_id: str) -> dict:
     """

policyengine_api/routes/economy_routes.py

@@ -15,7 +15,9 @@
     "/<country_id>/economy/<int:policy_id>/over/<int:baseline_policy_id>",
     methods=["GET"],
 )
-def get_economic_impact(country_id, policy_id, baseline_policy_id):
+def get_economic_impact(
+    country_id: str, policy_id: int, baseline_policy_id: int
+):
 
     policy_id = int(policy_id or get_current_law_policy_id(country_id))
     baseline_policy_id = int(

policyengine_api/routes/policy_routes.py

@@ -0,0 +1,81 @@
+from flask import Blueprint, Response, request
+import json
+
+from policyengine_api.services.policy_service import PolicyService
+from werkzeug.exceptions import NotFound, BadRequest
+from policyengine_api.utils.payload_validators import (
+    validate_country,
+    validate_set_policy_payload,
+)
+
+policy_bp = Blueprint("policy", __name__)
+policy_service = PolicyService()
+
+
+@policy_bp.route("/<country_id>/policy/<int:policy_id>", 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.
+
+    Args:
+        country_id (str)
+        policy_id (int | str)
+
+    Returns:
+        Response: A Flask response object containing the
+        policy data in JSON format
+    """
+
+    # Specifically cast policy_id to an integer
+    policy_id = int(policy_id)
+
+    policy: dict | None = policy_service.get_policy(country_id, policy_id)
+
+    if policy is None:
+        raise NotFound(f"Policy #{policy_id} not found.")
+
+    return Response(
+        json.dumps({"status": "ok", "message": None, "result": policy}),
+        status=200,
+    )
+
+
+@policy_bp.route("/<country_id>/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.
+
+    Args:
+        country_id (str)
+    """
+
+    payload = request.json
+
+    is_payload_valid, message = validate_set_policy_payload(payload)
+    if not is_payload_valid:
+        raise BadRequest(f"Invalid JSON data; details: {message}")
+
+    label = payload.pop("label", None)
+    policy_json = payload.pop("data", None)
+
+    policy_id, message, is_existing_policy = policy_service.set_policy(
+        country_id,
+        label,
+        policy_json,
+    )
+
+    response_body = dict(
+        status="ok",
+        message=message,
+        result=dict(
+            policy_id=policy_id,
+        ),
+    )
+
+    code = 200 if is_existing_policy else 201
+    return Response(
+        json.dumps(response_body), status=code, mimetype="application/json"
+    )

policyengine_api/services/economy_service.py

@@ -21,15 +21,18 @@ class EconomyService:
 
     def get_economic_impact(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options,
-        api_version,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options: dict,
+        api_version: str,
     ):
+        """
+        Calculate the society-wide economic impact of a policy reform.
+        """
         try:
             # Note for anyone modifying options_hash: redis-queue treats ":" as a namespace
             # delimiter; don't use colons in options_hash
@@ -172,15 +175,19 @@ def get_economic_impact(
 
     def _get_previous_impacts(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options_hash,
-        api_version,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options_hash: str,
+        api_version: str,
     ):
+        """
+        Fetch any previous simulation runs for the given policy reform.
+        """
+
         previous_impacts = reform_impacts_service.get_all_reform_impacts(
             country_id,
             policy_id,
@@ -203,16 +210,19 @@ def _get_previous_impacts(
 
     def _set_impact_computing(
         self,
-        country_id,
-        policy_id,
-        baseline_policy_id,
-        region,
-        dataset,
-        time_period,
-        options,
-        options_hash,
-        api_version,
+        country_id: str,
+        policy_id: int,
+        baseline_policy_id: int,
+        region: str,
+        dataset: str,
+        time_period: str,
+        options: dict,
+        options_hash: str,
+        api_version: str,
     ):
+        """
+        In the reform_impact table, set the status of the impact to "computing".
+        """
         try:
             reform_impacts_service.set_reform_impact(
                 country_id,