Add optional model and version to DataSource

flexmeasures/data/migrations/versions/30fe2267e7d5_add_optional_DataSource_columns_for_model_and_version.py

@@ -0,0 +1,38 @@
+"""add optional DataSource columns for model and version
+
+Revision ID: 30fe2267e7d5
+Revises: 96f2db5bed30
+Create Date: 2021-10-11 10:54:24.348371
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "30fe2267e7d5"
+down_revision = "96f2db5bed30"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.add_column(
+        "data_source", sa.Column("model", sa.String(length=80), nullable=True)
+    )
+    op.add_column(
+        "data_source", sa.Column("version", sa.String(length=17), nullable=True)
+    )
+    op.create_unique_constraint(
+        "_data_source_name_user_id_model_version_key",
+        "data_source",
+        ["name", "user_id", "model", "version"],
+    )
+
+
+def downgrade():
+    op.drop_constraint(
+        "_data_source_name_user_id_model_version_key", "data_source", type_="unique"
+    )
+    op.drop_column("data_source", "version")
+    op.drop_column("data_source", "model")

flexmeasures/data/models/data_sources.py

@@ -11,16 +11,24 @@ class DataSource(db.Model, tb.BeliefSourceDBMixin):
     """Each data source is a data-providing entity."""
 
     __tablename__ = "data_source"
+    __table_args__ = (db.UniqueConstraint("name", "user_id", "model", "version"),)
 
     # The type of data source (e.g. user, forecasting script or scheduling script)
     type = db.Column(db.String(80), default="")
-    # The id of the source (can link e.g. to fm_user table)
+
+    # The id of the user source (can link e.g. to fm_user table)
     user_id = db.Column(
         db.Integer, db.ForeignKey("fm_user.id"), nullable=True, unique=True
     )
-
     user = db.relationship("User", backref=db.backref("data_source", lazy=True))
 
+    # The model and version of a script source
+    model = db.Column(db.String(80), nullable=True)
+    version = db.Column(
+        db.String(17),  # length supports up to version 999.999.999dev999
+        nullable=True,
+    )
+
     def __init__(
         self,
         name: Optional[str] = None,
@@ -54,6 +62,23 @@ def label(self):
         else:
             return f"data from {self.name}"
 
+    @property
+    def description(self):
+        """Extended description
+
+        For example:
+
+            >>> DataSource("Seita", type="forecasting script", model="naive", version="1.2").description
+            <<< "Seita's naive model v1.2.0"
+
+        """
+        descr = self.name
+        if self.model:
+            descr += f"'s {self.model} model"
+            if self.version:
+                descr += f" v{self.version}"
+        return descr
+
     def __repr__(self):
         return "<Data source %r (%s)>" % (self.id, self.label)
 

flexmeasures/data/utils.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Optional
 
 import click
 
@@ -16,20 +16,36 @@ def save_to_session(objects: List[db.Model], overwrite: bool = False):
 
 
 def get_data_source(
-    data_source_name: str, data_source_type: str = "script"
+    data_source_name: str,
+    data_source_model: Optional[str] = None,
+    data_source_version: Optional[str] = None,
+    data_source_type: str = "script",
 ) -> DataSource:
     """Make sure we have a data source. Create one if it doesn't exist, and add to session.
     Meant for scripts that may run for the first time.
-    It should probably not be used in the middle of a transaction, because we commit to the session."""
+    """
 
     data_source = DataSource.query.filter_by(
-        name=data_source_name, type=data_source_type
+        name=data_source_name,
+        model=data_source_model,
+        version=data_source_version,
+        type=data_source_type,
     ).one_or_none()
     if data_source is None:
-        data_source = DataSource(name=data_source_name, type=data_source_type)
+        data_source = DataSource(
+            name=data_source_name,
+            model=data_source_model,
+            version=data_source_version,
+            type=data_source_type,
+        )
         db.session.add(data_source)
         db.session.flush()  # populate the primary key attributes (like id) without committing the transaction
-        click.echo(
-            f'Session updated with new {data_source_type} data source named "{data_source_name}".'
-        )
+        if data_source_model is None:
+            click.echo(
+                f'Session updated with new {data_source_type} data source named "{data_source_name}".'
+            )
+        else:
+            click.echo(
+                f'Session updated with new {data_source_type} data source named "{data_source_name}" ({data_source_model} model v{data_source_version}).'
+            )
     return data_source