From 04537466584da7cea1233f2ff58f1655102f6275 Mon Sep 17 00:00:00 2001 From: "F.N. Claessen" Date: Tue, 27 Aug 2024 20:18:51 +0200 Subject: [PATCH] fix: save new sources on bulk save Signed-off-by: F.N. Claessen --- timely_beliefs/beliefs/classes.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/timely_beliefs/beliefs/classes.py b/timely_beliefs/beliefs/classes.py index 3d4bdbf..c886a43 100644 --- 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"])