Allow reading beliefs from xlsx files, and from data files with multiple headers

flexmeasures/data/migrations/versions/04f0e2d2924a_add_source_id_as_primary_key_for_timed_beliefs.py

@@ -0,0 +1,39 @@
+"""add source id as primary key for timed beliefs
+
+Revision ID: 04f0e2d2924a
+Revises: e62ac5f519d7
+Create Date: 2021-04-10 13:53:22.561718
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "04f0e2d2924a"
+down_revision = "e62ac5f519d7"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.drop_constraint("timed_belief_pkey", "timed_belief")
+    op.create_primary_key(
+        "timed_belief_pkey",
+        "timed_belief",
+        [
+            "event_start",
+            "belief_horizon",
+            "cumulative_probability",
+            "sensor_id",
+            "source_id",
+        ],
+    )
+
+
+def downgrade():
+    op.drop_constraint("timed_belief_pkey", "timed_belief")
+    op.create_primary_key(
+        "timed_belief_pkey",
+        "timed_belief",
+        ["event_start", "belief_horizon", "cumulative_probability", "sensor_id"],
+    )

flexmeasures/data/scripts/cli_tasks/data_add.py

@@ -220,6 +220,12 @@ def add_initial_structure():
     type=click.IntRange(min=1),
     help="Sensor to which the beliefs pertain.",
 )
+@click.option(
+    "--source",
+    required=True,
+    type=str,
+    help="Source of the beliefs (an existing source id or name, or a new name).",
+)
 @click.option(
     "--horizon",
     required=False,
@@ -274,6 +280,7 @@ def add_initial_structure():
 def add_beliefs(
     file: str,
     sensor_id: int,
+    source: str,
     horizon: Optional[int] = None,
     cp: Optional[float] = None,
     allow_overwrite: bool = False,
@@ -305,16 +312,21 @@ def add_beliefs(
     if sensor is None:
         print(f"Failed to create beliefs: no sensor found with id {sensor_id}.")
         return
-    source = (
-        DataSource.query.filter(DataSource.name == "Seita")
-        .filter(DataSource.type == "CLI script")
-        .one_or_none()
-    )
-    if not source:
-        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
+    query = DataSource.query.filter(DataSource.type == "CLI script")
+    if source.isdigit():
+        query = query.filter(DataSource.id == int(source))
+        _source = query.one_or_none()
+        if not _source:
+            print(f"Failed to find source {source}.")
+            return
+    else:
+        query = query.filter(DataSource.name == source)
+        _source = query.one_or_none()
+        if not _source:
+            print(f"Setting up '{source}' as new data source...")
+            _source = DataSource(name=source, type="CLI script")
+            db.session.add(_source)
+            db.session.flush()  # assigns id
 
     # Set up optional parameters for read_csv
     kwargs = dict()
@@ -330,7 +342,7 @@ def add_beliefs(
     bdf = tb.read_csv(
         file,
         sensor,
-        source=source,
+        source=_source,
         cumulative_probability=cp,
         header=None,
         skiprows=skiprows,