fix: Update simulations table and PATCH endpoint to contain output structures present in report_outputs

Author: anth-volk

changelog_entry.yaml

@@ -0,0 +1,4 @@
+- bump: patch
+  changes:
+    changed:
+    - Updated simulation table and PATCH endpoint to contain output structures present in report_outputs.

policyengine_api/data/initialise.sql

@@ -112,9 +112,9 @@ CREATE TABLE IF NOT EXISTS simulations (
     population_id VARCHAR(255) NOT NULL,
     population_type VARCHAR(50) NOT NULL,
     policy_id INT NOT NULL,
-    -- output_json stores calculation results for household simulations only
-    -- For geography simulations, outputs are stored in report_outputs table
-    output_json JSON DEFAULT NULL
+    status VARCHAR(32) NOT NULL DEFAULT 'pending',
+    output JSON DEFAULT NULL,
+    error_message TEXT DEFAULT NULL
 );
 
 CREATE TABLE IF NOT EXISTS report_outputs (

policyengine_api/data/initialise_local.sql

@@ -121,9 +121,9 @@ CREATE TABLE IF NOT EXISTS simulations (
     population_id VARCHAR(255) NOT NULL,
     population_type VARCHAR(50) NOT NULL,
     policy_id INT NOT NULL,
-    -- output_json stores calculation results for household simulations only
-    -- For geography simulations, outputs are stored in report_outputs table
-    output_json JSON DEFAULT NULL
+    status VARCHAR(32) NOT NULL DEFAULT 'pending',
+    output JSON DEFAULT NULL,
+    error_message TEXT DEFAULT NULL
 );
 
 CREATE TABLE IF NOT EXISTS report_outputs (

policyengine_api/routes/simulation_routes.py

@@ -139,34 +139,37 @@ def get_simulation(country_id: str, simulation_id: int) -> Response:
 @validate_country
 def update_simulation(country_id: str) -> Response:
     """
-    Update a simulation record with calculation output.
+    Update a simulation record with results or error.
 
     Args:
         country_id (str): The country ID.
 
     Request body can contain:
         - id (int): The simulation ID.
-        - output_json (str): The calculation output as JSON string (for household simulations)
+        - status (str): The new status ('complete' or 'error')
+        - 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")
 
-    # Extract fields
+    # Extract optional fields
+    status = payload.get("status")
     simulation_id = payload.get("id")
-    output_json = payload.get("output_json")
+    output = payload.get("output")
+    error_message = payload.get("error_message")
     print(f"Updating simulation #{simulation_id} for country {country_id}")
 
-    # Validate that id is provided
-    if simulation_id is None:
-        raise BadRequest("id is required")
-    if not isinstance(simulation_id, int):
-        raise BadRequest("id must be an integer")
+    # Validate status if provided
+    if status is not None and status not in ["pending", "complete", "error"]:
+        raise BadRequest("status must be 'pending', 'complete', or 'error'")
 
-    # Validate that at least one field is being updated
-    if output_json is None:
-        raise BadRequest("output_json must be provided for update")
+    # Validate that complete status has output
+    if status == "complete" and output is None:
+        raise BadRequest("output is required when status is 'complete'")
 
     try:
         # First check if the simulation exists
@@ -177,10 +180,12 @@ def update_simulation(country_id: str) -> Response:
             raise NotFound(f"Simulation #{simulation_id} not found.")
 
         # Update the simulation
-        success = simulation_service.update_simulation_output(
+        success = simulation_service.update_simulation(
             country_id=country_id,
             simulation_id=simulation_id,
-            output_json=output_json,
+            status=status,
+            output=output,
+            error_message=error_message,
         )
 
         if not success:

policyengine_api/services/simulation_service.py

@@ -56,7 +56,7 @@ def create_simulation(
         policy_id: int,
     ) -> dict:
         """
-        Create a new simulation record.
+        Create a new simulation record with pending status.
 
         Args:
             country_id (str): The country ID.
@@ -72,13 +72,14 @@ def create_simulation(
 
         try:
             database.query(
-                "INSERT INTO simulations (country_id, api_version, population_id, population_type, policy_id) VALUES (?, ?, ?, ?, ?)",
+                "INSERT INTO simulations (country_id, api_version, population_id, population_type, policy_id, status) VALUES (?, ?, ?, ?, ?, ?)",
                 (
                     country_id,
                     api_version,
                     population_id,
                     population_type,
                     policy_id,
+                    "pending",
                 ),
             )
 
@@ -135,24 +136,28 @@ def get_simulation(
             )
             raise e
 
-    def update_simulation_output(
+    def update_simulation(
         self,
         country_id: str,
         simulation_id: int,
-        output_json: str | None = None,
+        status: str | None = None,
+        output: str | None = None,
+        error_message: str | None = None,
     ) -> bool:
         """
-        Update a simulation record with calculation output.
+        Update a simulation record with results or error.
 
         Args:
             country_id (str): The country ID.
             simulation_id (int): The simulation ID.
-            output_json (str | None): The output as JSON string (for household simulations).
+            status (str | None): The new status ('complete' or 'error').
+            output (str | None): The result output as JSON string (for complete status).
+            error_message (str | None): The error message (for error status).
 
         Returns:
             bool: True if update was successful.
         """
-        print(f"Updating simulation {simulation_id} with output")
+        print(f"Updating simulation {simulation_id}")
         # Automatically update api_version on every update to latest
         api_version: str = COUNTRY_PACKAGE_VERSIONS.get(country_id)
 
@@ -161,10 +166,18 @@ def update_simulation_output(
             update_fields = []
             update_values = []
 
-            if output_json is not None:
-                update_fields.append("output_json = ?")
+            if status is not None:
+                update_fields.append("status = ?")
+                update_values.append(status)
+
+            if output is not None:
+                update_fields.append("output = ?")
                 # Output is already a JSON string from frontend
-                update_values.append(output_json)
+                update_values.append(output)
+
+            if error_message is not None:
+                update_fields.append("error_message = ?")
+                update_values.append(error_message)
 
             # Always update API version
             update_fields.append("api_version = ?")

tests/fixtures/services/simulation_fixtures.py

@@ -22,14 +22,15 @@ def existing_simulation_record(test_db):
     """Insert an existing simulation record into the database."""
     test_db.query(
         """INSERT INTO simulations
-        (country_id, api_version, population_id, population_type, policy_id)
-        VALUES (?, ?, ?, ?, ?)""",
+        (country_id, api_version, population_id, population_type, policy_id, status)
+        VALUES (?, ?, ?, ?, ?, ?)""",
         (
             valid_simulation_data["country_id"],
             valid_simulation_data["api_version"],
             valid_simulation_data["population_id"],
             valid_simulation_data["population_type"],
             valid_simulation_data["policy_id"],
+            "pending",
         ),
     )