Enforce stage 2 lineage invariants

Author: anth-volk

policyengine_api/data/data.py

@@ -1,5 +1,5 @@
 import sqlite3
-from policyengine_api.constants import REPO, VERSION, COUNTRY_PACKAGE_VERSIONS
+from policyengine_api.constants import REPO, COUNTRY_PACKAGE_VERSIONS
 from policyengine_api.utils import hash_object
 from pathlib import Path
 from dotenv import load_dotenv
@@ -41,6 +41,29 @@ def fetchall(self):
         return remaining
 
 
+class _TransactionProxy:
+    """Execute queries against an existing connection inside a transaction."""
+
+    def __init__(self, connection, local: bool):
+        self._connection = connection
+        self._local = local
+
+    def query(self, *query):
+        if self._local:
+            cursor = self._connection.cursor()
+            return cursor.execute(*query)
+
+        query = list(query)
+        main_query = query[0].replace("?", "%s")
+        query[0] = main_query
+        params = query[1] if len(query) > 1 else None
+        if params is not None:
+            result = self._connection.exec_driver_sql(main_query, params)
+        else:
+            result = self._connection.exec_driver_sql(main_query)
+        return _ResultProxy(result)
+
+
 class PolicyEngineDatabase:
     """
     A wrapper around the database connection.
@@ -50,6 +73,13 @@ class PolicyEngineDatabase:
 
     household_cache: dict = {}
 
+    @staticmethod
+    def _dict_factory(cursor, row):
+        d = {}
+        for idx, col in enumerate(cursor.description):
+            d[col[0]] = row[idx]
+        return d
+
     def __init__(
         self,
         local: bool = False,
@@ -91,7 +121,7 @@ def _close_pool(self):
         try:
             self.pool.dispose()
             self.connector.close()
-        except:
+        except Exception:
             pass
 
     def _execute_remote(self, query_args):
@@ -110,17 +140,22 @@ def _execute_remote(self, query_args):
             # connection context closing
             return _ResultProxy(result)
 
+    def _execute_remote_transaction(self, callback):
+        with self.pool.connect() as conn:
+            transaction = conn.begin()
+            proxy = _TransactionProxy(conn, local=False)
+            try:
+                result = callback(proxy)
+                transaction.commit()
+                return result
+            except Exception:
+                transaction.rollback()
+                raise
+
     def query(self, *query):
         if self.local:
             with sqlite3.connect(self.db_url) as conn:
-
-                def dict_factory(cursor, row):
-                    d = {}
-                    for idx, col in enumerate(cursor.description):
-                        d[col[0]] = row[idx]
-                    return d
-
-                conn.row_factory = dict_factory
+                conn.row_factory = self._dict_factory
                 cursor = conn.cursor()
                 return cursor.execute(*query)
         else:
@@ -134,14 +169,44 @@ def dict_factory(cursor, row):
             except (
                 sqlalchemy.exc.InterfaceError,
                 sqlalchemy.exc.OperationalError,
-            ) as e:
+            ):
                 try:
                     self._close_pool()
                     self._create_pool()
                     return self._execute_remote(query)
                 except Exception as e:
                     raise e
 
+    def transaction(self, callback):
+        if self.local:
+            connection = getattr(self, "_connection", None)
+            owns_connection = connection is None
+            if owns_connection:
+                connection = sqlite3.connect(self.db_url)
+            connection.row_factory = self._dict_factory
+            try:
+                connection.execute("BEGIN IMMEDIATE")
+                proxy = _TransactionProxy(connection, local=True)
+                result = callback(proxy)
+                connection.commit()
+                return result
+            except Exception:
+                connection.rollback()
+                raise
+            finally:
+                if owns_connection:
+                    connection.close()
+
+        try:
+            return self._execute_remote_transaction(callback)
+        except (
+            sqlalchemy.exc.InterfaceError,
+            sqlalchemy.exc.OperationalError,
+        ):
+            self._close_pool()
+            self._create_pool()
+            return self._execute_remote_transaction(callback)
+
     def initialize(self):
         """
         Create the database tables.
@@ -175,7 +240,7 @@ def initialize(self):
             range(1, 1 + len(COUNTRY_PACKAGE_VERSIONS)),
         ):
             self.query(
-                f"INSERT INTO policy (id, country_id, label, api_version, policy_json, policy_hash) VALUES (?, ?, ?, ?, ?, ?)",
+                "INSERT INTO policy (id, country_id, label, api_version, policy_json, policy_hash) VALUES (?, ?, ?, ?, ?, ?)",
                 (
                     policy_id,
                     country_id,

policyengine_api/services/report_output_alias_service.py

@@ -4,12 +4,16 @@
 
 
 class ReportOutputAliasService:
-    def _report_output_exists(self, report_output_id: int) -> bool:
+    def _get_report_output_row(self, report_output_id: int) -> dict | None:
         row: Row | None = database.query(
-            "SELECT id FROM report_outputs WHERE id = ?",
+            """
+            SELECT id, country_id, simulation_1_id, simulation_2_id, year
+            FROM report_outputs
+            WHERE id = ?
+            """,
             (report_output_id,),
         ).fetchone()
-        return row is not None
+        return dict(row) if row is not None else None
 
     def get_alias(self, legacy_report_output_id: int) -> dict | None:
         row: Row | None = database.query(
@@ -27,7 +31,7 @@ def resolve_canonical_report_output_id(
         alias = self.get_alias(requested_report_output_id)
         if alias is not None:
             canonical_report_output_id = alias["canonical_report_output_id"]
-            if not self._report_output_exists(canonical_report_output_id):
+            if self._get_report_output_row(canonical_report_output_id) is None:
                 raise ValueError(
                     "Alias points to missing canonical report output "
                     f"#{canonical_report_output_id}"
@@ -45,29 +49,49 @@ def set_alias(
         legacy_report_output_id: int,
         canonical_report_output_id: int,
     ) -> bool:
-        if not self._report_output_exists(canonical_report_output_id):
+        legacy_report_output = self._get_report_output_row(legacy_report_output_id)
+        if legacy_report_output is None:
+            raise ValueError(
+                f"Legacy report output #{legacy_report_output_id} not found"
+            )
+
+        canonical_report_output = self._get_report_output_row(
+            canonical_report_output_id
+        )
+        if canonical_report_output is None:
             raise ValueError(
                 f"Canonical report output #{canonical_report_output_id} not found"
             )
+        if legacy_report_output_id == canonical_report_output_id:
+            raise ValueError("Legacy and canonical report outputs must be different")
 
         existing_alias = self.get_alias(legacy_report_output_id)
-        if existing_alias is None:
-            database.query(
-                """
-                INSERT INTO legacy_report_output_aliases
-                (legacy_report_output_id, canonical_report_output_id)
-                VALUES (?, ?)
-                """,
-                (legacy_report_output_id, canonical_report_output_id),
-            )
-            return True
+        if existing_alias is not None:
+            if (
+                existing_alias["canonical_report_output_id"]
+                == canonical_report_output_id
+            ):
+                return True
 
-        if existing_alias["canonical_report_output_id"] == canonical_report_output_id:
-            return True
+            raise ValueError(
+                "Legacy report output alias already points to canonical report output "
+                f"#{existing_alias['canonical_report_output_id']}"
+            )
 
-        raise ValueError(
-            "Legacy report output alias already points to canonical report output "
-            f"#{existing_alias['canonical_report_output_id']}"
+        logical_key = ("country_id", "simulation_1_id", "simulation_2_id", "year")
+        if any(
+            legacy_report_output[field] != canonical_report_output[field]
+            for field in logical_key
+        ):
+            raise ValueError(
+                "Legacy and canonical report outputs must describe the same report"
+            )
+        database.query(
+            """
+            INSERT INTO legacy_report_output_aliases
+            (legacy_report_output_id, canonical_report_output_id)
+            VALUES (?, ?)
+            """,
+            (legacy_report_output_id, canonical_report_output_id),
         )
-
         return True