fix: recommendation database persistence (id set/get)

Author: glichtner

execution_engine/converter/time_from_event/abstract.py

@@ -26,7 +26,7 @@ def _wrap_criteria_with_factory(
     Recursively wraps all Criterion instances within a combination using the factory.
     """
     # Create a new combination of the same type with the same operator
-    new_combo = combo.__class__(operator=combo.operator, category=combo.category)
+    new_combo = combo.__class__(operator=combo.operator)
 
     # Loop through all elements
     for element in combo:

execution_engine/execution_engine.py

@@ -141,7 +141,7 @@ def execute(
             # database, its _id slot is not None. Otherwise, register
             # the recommendation to store it into the database and
             # assign an id.
-            if recommendation._id is None:
+            if not recommendation.is_persisted():
                 self.register_recommendation(recommendation)
             run_id = self.register_run(
                 recommendation, start_datetime=start_datetime, end_datetime=end_datetime
@@ -193,14 +193,34 @@ def load_recommendation_from_database(
             recommendation = cohort.Recommendation.from_json(
                 rec_db.recommendation_json.decode()
             )
-            # All objects in the deserialized object graph must have
-            # an id.
-            assert recommendation._id is not None
-            assert recommendation._base_criterion._id is not None
-            for pi_pair in recommendation._pi_pairs:
-                assert pi_pair._id is not None
+
+            # we'll run registration again to set all ids
+            recommendation.set_id(rec_db.recommendation_id)
+
+            assert (
+                recommendation.base_criterion is not None
+            ), "Base Criterion must be set"
+            self.register_criterion(recommendation.base_criterion)
+
+            for pi_pair in recommendation.population_intervention_pairs():
+                self.register_population_intervention_pair(
+                    pi_pair, rec_db.recommendation_id
+                )
+
+                for criterion in pi_pair.flatten():
+                    self.register_criterion(criterion)
+
+            # All objects in the deserialized object graph must have an id.
+            assert recommendation.id is not None
+            assert recommendation.base_criterion is not None
+            assert recommendation.base_criterion.id is not None
+
+            for pi_pair in recommendation.population_intervention_pairs():
+                assert pi_pair.id is not None
+
                 for criterion in pi_pair.flatten():
-                    assert criterion._id is not None
+                    assert criterion.id is not None
+
             return recommendation
 
         return None
@@ -229,7 +249,12 @@ def register_recommendation(self, recommendation: cohort.Recommendation) -> None
             rec_db = con.execute(query).fetchone()
 
             if rec_db is not None:
-                recommendation.id = rec_db.recommendation_id
+                if recommendation.is_persisted():
+                    assert (
+                        recommendation.id == rec_db.recommendation_id
+                    ), "Recommendation IDs do not match"
+                else:
+                    recommendation.set_id(rec_db.recommendation_id)
             else:
                 query = (
                     insert(recommendation_table)
@@ -247,11 +272,17 @@ def register_recommendation(self, recommendation: cohort.Recommendation) -> None
                 )
 
                 result = con.execute(query)
-                recommendation.id = result.fetchone().recommendation_id
+                recommendation.set_id(result.fetchone().recommendation_id)
+
         # Register all child objects. After that, the recommendation
         # and all child objects have valid ids (either restored or
         # fresh).
-        self.register_criterion(recommendation._base_criterion)
+
+        if recommendation.base_criterion is None:
+            raise ValueError("Base Criterion must be set when storing recommendation")
+
+        self.register_criterion(recommendation.base_criterion)
+
         for pi_pair in recommendation.population_intervention_pairs():
             self.register_population_intervention_pair(
                 pi_pair, recommendation_id=recommendation.id
@@ -260,12 +291,14 @@ def register_recommendation(self, recommendation: cohort.Recommendation) -> None
                 self.register_criterion(criterion)
 
         assert recommendation.id is not None
-        # TODO(jmoringe): mypy doesn't like this one. Not sure why.
-        # assert recommendation._base_criterion._id is not None
-        for pi_pair in recommendation._pi_pairs:
-            assert pi_pair._id is not None
+        assert recommendation.base_criterion is not None
+        assert recommendation.base_criterion.id is not None
+
+        for pi_pair in recommendation.population_intervention_pairs():
+            assert pi_pair.id is not None
+
             for criterion in pi_pair.flatten():
-                assert criterion._id is not None
+                assert criterion.id is not None
 
         # Update the recommendation in the database with the final
         # JSON representation and execution graph (now that
@@ -306,7 +339,12 @@ def register_population_intervention_pair(
             pi_pair_db = con.execute(query).fetchone()
 
             if pi_pair_db is not None:
-                pi_pair.id = pi_pair_db.pi_pair_id
+                if pi_pair.is_persisted():
+                    assert (
+                        pi_pair.id == pi_pair_db.pi_pair_id
+                    ), "Population/Intervention Pair IDs do not match"
+                else:
+                    pi_pair.set_id(pi_pair_db.pi_pair_id)
             else:
                 query = (
                     insert(result_db.PopulationInterventionPair)
@@ -320,7 +358,7 @@ def register_population_intervention_pair(
                 )
 
                 result = con.execute(query)
-                pi_pair.id = result.fetchone().pi_pair_id
+                pi_pair.set_id(result.fetchone().pi_pair_id)
 
     def register_criterion(self, criterion: Criterion) -> None:
         """
@@ -337,7 +375,12 @@ def register_criterion(self, criterion: Criterion) -> None:
             criterion_db = con.execute(query).fetchone()
 
             if criterion_db is not None:
-                criterion.id = criterion_db.criterion_id
+                if criterion.is_persisted():
+                    assert (
+                        criterion.id == criterion_db.criterion_id
+                    ), "Criterion IDs do not match"
+                else:
+                    criterion.set_id(criterion_db.criterion_id)
             else:
                 query = (
                     insert(result_db.Criterion)
@@ -349,7 +392,7 @@ def register_criterion(self, criterion: Criterion) -> None:
                 )
 
                 result = con.execute(query)
-                criterion.id = result.fetchone().criterion_id
+                criterion.set_id(result.fetchone().criterion_id)
 
         assert criterion.id is not None
 

execution_engine/omop/cohort/population_intervention_pair.py

@@ -38,7 +38,6 @@ class PopulationInterventionPair(Serializable):
     (e.g. "has condition X and lab value Y >= Z").
     """
 
-    _id: int | None
     _name: str
     _population: CriterionCombination
     _intervention: CriterionCombination
@@ -50,9 +49,7 @@ def __init__(
         base_criterion: Criterion,
         population: CriterionCombination | None = None,
         intervention: CriterionCombination | None = None,
-        id: int | None = None,
     ) -> None:
-        self._id = id
         self._name = name
         self._url = url
         self._base_criterion = base_criterion

execution_engine/omop/cohort/recommendation.py

@@ -51,7 +51,6 @@ def __init__(
         url: str,
         version: str,
         description: str,
-        recommendation_id: int | None = None,
         package_version: str | None = None,
     ) -> None:
         self._pi_pairs: list[cohort.PopulationInterventionPair] = pi_pairs
@@ -62,8 +61,6 @@ def __init__(
         self._version: str = version
         self._description: str = description
         self._package_version: str | None = package_version
-        # The id is used in the recommendation_id field in the result tables.
-        self._id = recommendation_id
 
     def __repr__(self) -> str:
         """
@@ -131,6 +128,13 @@ def description(self) -> str:
         """
         return self._description
 
+    @property
+    def base_criterion(self) -> Criterion | None:
+        """
+        Get the base criterion of the recommendation.
+        """
+        return self._base_criterion
+
     def execution_graph(self) -> ExecutionGraph:
         """
         Get the execution maps of the full recommendation.

execution_engine/omop/serializable.py

@@ -9,22 +9,42 @@ class Serializable(ABC):
     """
 
     _id: int | None = None
+    """
+    The id is used in the database tables.
+    """
+
+    def set_id(self, value: int) -> None:
+        """
+        Assigns the database ID to the object. This can only be done once.
+
+        This ID corresponds to the primary key in the database and is set
+        when the object is persisted.
+
+        :param value: The database ID assigned to the object.
+        :raises ValueError: If the ID has already been set.
+        """
+        if self._id is not None:
+            raise ValueError("Database ID has already been set!")
+        self._id = value
 
     @property
     def id(self) -> int:
         """
-        Get the id of the object (used in the database).
+        Retrieves the database ID of the object.
+
+        This ID is only available after the object has been stored in the database.
+
+        :return: The database ID, or None if the object has not been stored yet.
         """
         if self._id is None:
-            raise ValueError("Id not set")
+            raise ValueError("Database ID has not been set yet!")
         return self._id
 
-    @id.setter
-    def id(self, value: int) -> None:
+    def is_persisted(self) -> bool:
         """
-        Set the id of the object (used in the database).
+        Returns True if the object has been stored in the database.
         """
-        self._id = value
+        return self._id is not None
 
     @abstractmethod
     def dict(self) -> dict:

tests/execution_engine/omop/cohort/test_cohort_recommendation.py

@@ -1,5 +1,6 @@
 import pytest
 
+from execution_engine.omop.cohort import PopulationInterventionPair
 from execution_engine.omop.cohort.recommendation import Recommendation
 from execution_engine.omop.concepts import Concept
 from execution_engine.omop.criterion.visit_occurrence import ActivePatients
@@ -58,3 +59,79 @@ def test_serialization_with_active_patients(self, concept):
         json = original.json()
         deserialized = Recommendation.from_json(json)
         assert original == deserialized
+
+    def test_database_serialization(self, concept):
+        from execution_engine.execution_engine import ExecutionEngine
+        from execution_engine.omop.criterion.factory import register_criterion_class
+
+        register_criterion_class("MockCriterion", MockCriterion)
+
+        e = ExecutionEngine(verbose=False)
+
+        url = "http://example.com/fhir/recommendation"
+
+        base_criterion = MockCriterion("c")
+        population_criterion = MockCriterion("p")
+        intervention_criterion = MockCriterion("i)")
+
+        pi_pair = PopulationInterventionPair(
+            name="foo",
+            url="foo",
+            base_criterion=base_criterion,
+            population=population_criterion,
+            intervention=intervention_criterion,
+        )
+
+        pi_pairs = [pi_pair]
+
+        recommendation = Recommendation(
+            pi_pairs=pi_pairs,
+            base_criterion=ActivePatients(),
+            name="foo",
+            title="bar",
+            url=url,
+            version="1.0",
+            description="test",
+            package_version="1.0",
+        )
+
+        with pytest.raises(ValueError, match=r"Database ID has not been set yet!"):
+            assert recommendation.id is None
+
+        with pytest.raises(ValueError, match=r"Database ID has not been set yet!"):
+            assert recommendation.base_criterion.id is None
+
+        for pi_pair in recommendation.population_intervention_pairs():
+            with pytest.raises(ValueError, match=r"Database ID has not been set yet!"):
+                assert pi_pair.id is None
+
+            for criterion in pi_pair.flatten():
+                with pytest.raises(
+                    ValueError, match=r"Database ID has not been set yet!"
+                ):
+                    assert criterion.id is None
+
+        e.register_recommendation(recommendation)
+
+        assert recommendation.id is not None
+        assert recommendation.base_criterion.id is not None
+        for pi_pair in recommendation.population_intervention_pairs():
+            assert pi_pair.id is not None
+            for criterion in pi_pair.flatten():
+                assert criterion.id is not None
+
+        rec_loaded = e.load_recommendation_from_database(url)
+
+        assert recommendation == rec_loaded
+        assert rec_loaded.id == recommendation.id
+        assert len(rec_loaded._pi_pairs) == len(pi_pairs)
+
+        for pi_pair_loaded, pi_pair in zip(rec_loaded._pi_pairs, pi_pairs):
+            assert pi_pair_loaded.id == pi_pair.id
+
+            assert len(pi_pair_loaded.flatten()) == len(pi_pair.flatten())
+
+            for criterion_loaded, criterion in zip(
+                pi_pair_loaded.flatten(), pi_pair.flatten()
+            ):
+                assert criterion.id == criterion_loaded.id

tests/execution_engine/omop/criterion/combination/__init__.py

@@ -41,7 +41,7 @@ def from_dict(cls, data: Dict[str, Any]) -> Self:
         """
         Create an object from a dictionary.
         """
-        return cls(**data)
+        return cls()
 
     def description(self) -> str:
         """
@@ -53,6 +53,4 @@ def dict(self) -> dict:
         """
         Get a dictionary representation of the object.
         """
-        return {
-            "class": self.__class__.__name__,
-        }
+        return {}