Overwrite beliefs

flexmeasures/data/models/time_series.py

@@ -110,17 +110,29 @@ def search(
         )
 
     @classmethod
-    def add(cls, bdf: tb.BeliefsDataFrame, commit_transaction: bool = True):
+    def add(
+        cls,
+        bdf: tb.BeliefsDataFrame,
+        expunge_session: bool = True,
+        allow_overwrite: bool = False,
+        commit_transaction: bool = True,
+    ):
         """Add a BeliefsDataFrame as timed beliefs in the database.
 
         :param bdf: the BeliefsDataFrame to be persisted
+        :param expunge_session:     if True, all instances are removed from the session before adding beliefs
+                                    (make sure you flush any other new objects)
+        :param allow_overwrite:     if True, new objects are merged
+                                    if False, objects are added to the session
         :param commit_transaction: if True, the session is committed
                                    if False, you can still add other data to the session
                                    and commit it all within an atomic transaction
         """
         return cls.add_to_session(
             session=db.session,
             beliefs_data_frame=bdf,
+            expunge_session=expunge_session,
+            allow_overwrite=allow_overwrite,
             commit_transaction=commit_transaction,
         )
 

flexmeasures/data/scripts/cli_tasks/data_add.py

@@ -10,6 +10,7 @@
 from flask_security.utils import hash_password
 import click
 import getpass
+from sqlalchemy.exc import IntegrityError
 import timely_beliefs as tb
 
 from flexmeasures.data import db
@@ -226,8 +227,17 @@ def add_initial_structure():
     type=click.FloatRange(0, 1),
     help="Cumulative probability in the range [0, 1].",
 )
+@click.option(
+    "--allow-overwrite/--do-not-allow-overwrite",
+    default=False,
+    help="Allow overwriting possibly already existing data.",
+)
 def add_beliefs(
-    file: str, sensor_id: int, horizon: Optional[int] = None, cp: Optional[float] = None
+    file: str,
+    sensor_id: int,
+    horizon: Optional[int] = None,
+    cp: Optional[float] = None,
+    allow_overwrite: bool = False,
 ):
     """Add sensor data from a csv file.
 
@@ -257,6 +267,7 @@ def add_beliefs(
         print("SETTING UP CLI SCRIPT AS NEW DATA SOURCE...")
         source = DataSource(name="Seita", type="CLI script")
         db.session.add(source)
+        db.session.flush()  # assigns id
     bdf = tb.read_csv(
         file,
         sensor,
@@ -272,8 +283,26 @@ def add_beliefs(
             )
         ),
     )
-    TimedBelief.add(bdf, commit_transaction=False)
-    db.session.commit()
+    try:
+        # Disallowing overwriting can be much more efficient, through a bulk insert
+        TimedBelief.add(
+            bdf, expunge_session=True, allow_overwrite=False, commit_transaction=True
+        )
+    except IntegrityError as e:
+        if allow_overwrite:
+            print(
+                f"Attempting to overwrite data to bypass the following error: {e.orig}"
+            )
+            db.session.rollback()
+            TimedBelief.add(
+                bdf, expunge_session=True, allow_overwrite=True, commit_transaction=True
+            )
+        else:
+            print(
+                f"Failed to create beliefs due to the following error: {e.orig}\n"
+                f"As a possible workaround, use the --allow-overwrite flag."
+            )
+            return
     print(f"Successfully created beliefs\n{bdf}")