feat: bulk save objects ignoring conflicts

timely_beliefs/beliefs/classes.py

@@ -28,6 +28,7 @@
     func,
     select,
 )
+from sqlalchemy.dialects.postgresql import insert
 from sqlalchemy.ext.declarative import declared_attr
 from sqlalchemy.ext.hybrid import hybrid_method, hybrid_property
 from sqlalchemy.orm import Session, backref, declarative_mixin, relationship
@@ -275,22 +276,48 @@ def add_to_session(
                                     and commit it all within an atomic transaction
         """
         # Belief timing is stored as the belief horizon rather than as the belief time
-        belief_records = (
-            beliefs_data_frame.convert_index_from_belief_time_to_horizon()
-            .reset_index()
-            .to_dict("records")
+        beliefs_data_frame = (
+            beliefs_data_frame.convert_index_from_belief_time_to_horizon().reset_index()
         )
-        beliefs = [cls(sensor=beliefs_data_frame.sensor, **d) for d in belief_records]
+        beliefs = [
+            cls(sensor=beliefs_data_frame.sensor, **d)
+            for d in beliefs_data_frame.to_dict("records")
+        ]
+
         if expunge_session:
             session.expunge_all()
-        if not allow_overwrite:
-            if bulk_save_objects:
-                session.bulk_save_objects(beliefs)
+
+        if bulk_save_objects:
+            # serialize source and sensor
+            beliefs_data_frame["source_id"] = beliefs_data_frame["source"].apply(
+                lambda x: x.id
+            )
+            beliefs_data_frame["sensor_id"] = beliefs_data_frame.sensor.id
+            beliefs_data_frame = beliefs_data_frame.drop(columns=["source"])
+
+            smt = insert(cls).values(beliefs_data_frame.to_dict("records"))
+
+            if allow_overwrite:
+                smt = smt.on_conflict_do_update(
+                    index_elements=[
+                        "event_start",
+                        "belief_horizon",
+                        "source_id",
+                        "sensor_id",
+                        "cumulative_probability",
+                    ],
+                    set_=dict(event_value=smt.excluded.event_value),
+                )
+
+            session.execute(smt)
+
+        else:
+            if allow_overwrite:
+                for belief in beliefs:
+                    session.merge(belief)
             else:
                 session.add_all(beliefs)
-        else:
-            for belief in beliefs:
-                session.merge(belief)
+
         if commit_transaction:
             session.commit()