fix: save new sources on bulk save (#186)

Signed-off-by: F.N. Claessen <[email protected]>
SeitaBV · Aug 30, 2024 · 93ef8b1 · 93ef8b1
1 parent 60a2793
commit 93ef8b1
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/timely_beliefs/beliefs/classes.py b/timely_beliefs/beliefs/classes.py
@@ -295,10 +295,23 @@ def add_to_session(
             session.expunge_all()
 
         if bulk_save_objects:
-            # serialize source and sensor
+            # serialize sources and sensor, while adding new sources
+
+            # serialize sources
             beliefs_data_frame["source_id"] = beliefs_data_frame["source"].apply(
                 lambda x: x.id
             )
+
+            # Add new sources
+            newbies = pd.isnull(beliefs_data_frame["source_id"])
+            if any(newbies):
+                session.add_all(beliefs_data_frame.loc[newbies, "source"])
+                session.flush()  # assign IDs
+                beliefs_data_frame.loc[newbies, "source_id"] = beliefs_data_frame.loc[
+                    newbies, "source"
+                ].apply(lambda x: x.id)
+
+            # serialize sensor
             beliefs_data_frame["sensor_id"] = beliefs_data_frame.sensor.id
             beliefs_data_frame = beliefs_data_frame.drop(columns=["source"])