Merge 3d373df into af7f411

.vscode/settings.json

@@ -13,5 +13,10 @@
     "python.linting.pylintEnabled": false,
     "python.linting.flake8Enabled": true,
     "workbench.editor.wrapTabs": true,
-    "python.formatting.provider": "black"
+    "python.formatting.provider": "black",
+    "python.testing.pytestArgs": [
+        "flexmeasures"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
 }

flexmeasures/data/migrations/versions/2ac7fb39ce0c_add_attribute_column_to_data_source.py

@@ -0,0 +1,51 @@
+"""add attribute column to data source
+
+Revision ID: 2ac7fb39ce0c
+Revises: d814c0688ae0
+Create Date: 2023-06-05 23:41:31.788961
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "2ac7fb39ce0c"
+down_revision = "d814c0688ae0"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # add the column `attributes` to the table `data_source`
+    op.add_column(
+        "data_source",
+        sa.Column("attributes", sa.JSON(), nullable=True, default={}),
+    )
+
+    # add the column `attributes_hash` to the table `data_source`
+    op.add_column(
+        "data_source",
+        sa.Column("attributes_hash", sa.LargeBinary(length=256), nullable=True),
+    )
+
+    # remove previous uniqueness constraint and add a new that takes attributes_hash into account
+    op.drop_constraint(op.f("data_source_name_key"), "data_source", type_="unique")
+    op.create_unique_constraint(
+        "data_source_name_key",
+        "data_source",
+        ["name", "user_id", "model", "version", "attributes_hash"],
+    )
+
+
+def downgrade():
+
+    op.drop_constraint("data_source_name_key", "data_source", type_="unique")
+    op.create_unique_constraint(
+        "data_source_name_key",
+        "data_source",
+        ["name", "user_id", "model", "version"],
+    )
+
+    op.drop_column("data_source", "attributes")
+    op.drop_column("data_source", "attributes_hash")

flexmeasures/data/models/data_sources.py

@@ -1,11 +1,14 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import json
+from typing import TYPE_CHECKING, Any
+from sqlalchemy.ext.mutable import MutableDict
 
 import timely_beliefs as tb
 
 from flexmeasures.data import db
 from flask import current_app
+import hashlib
 
 
 if TYPE_CHECKING:
@@ -57,7 +60,9 @@ 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"),)
+    __table_args__ = (
+        db.UniqueConstraint("name", "user_id", "model", "version", "attributes_hash"),
+    )
 
     # The type of data source (e.g. user, forecaster or scheduler)
     type = db.Column(db.String(80), default="")
@@ -68,18 +73,30 @@ class DataSource(db.Model, tb.BeliefSourceDBMixin):
     )
     user = db.relationship("User", backref=db.backref("data_source", lazy=True))
 
+    attributes = db.Column(MutableDict.as_mutable(db.JSON), nullable=False, default={})
+
+    attributes_hash = db.Column(db.LargeBinary(length=256))
+
     # 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,
     )
 
+    sensors = db.relationship(
+        "Sensor",
+        secondary="timed_belief",
+        backref=db.backref("data_sources", lazy="select"),
+        viewonly=True,
+    )
+
     def __init__(
         self,
         name: str | None = None,
         type: str | None = None,
         user: User | None = None,
+        attributes: dict | None = None,
         **kwargs,
     ):
         if user is not None:
@@ -89,6 +106,13 @@ def __init__(
         elif user is None and type == "user":
             raise TypeError("A data source cannot have type 'user' but no user set.")
         self.type = type
+
+        if attributes is not None:
+            self.attributes = attributes
+            self.attributes_hash = hashlib.sha256(
+                json.dumps(attributes).encode("utf-8")
+            ).digest()
+
         tb.BeliefSourceDBMixin.__init__(self, name=name)
         db.Model.__init__(self, **kwargs)
 
@@ -144,3 +168,18 @@ def to_dict(self) -> dict:
             type=self.type if self.type in ("forecaster", "scheduler") else "other",
             description=self.description,
         )
+
+    @staticmethod
+    def hash_attributes(attributes: dict) -> str:
+        return hashlib.sha256(json.dumps(attributes).encode("utf-8")).digest()
+
+    def get_attribute(self, attribute: str, default: Any = None) -> Any:
+        """Looks for the attribute on the DataSource."""
+        return self.attributes.get(attribute)
+
+    def has_attribute(self, attribute: str) -> bool:
+        return attribute in self.attributes
+
+    def set_attribute(self, attribute: str, value):
+        if self.has_attribute(attribute):
+            self.attributes[attribute] = value

flexmeasures/data/services/data_sources.py

@@ -13,6 +13,7 @@ def get_or_create_source(
     source_type: str | None = None,
     model: str | None = None,
     version: str | None = None,
+    attributes: dict | None = None,
     flush: bool = True,
 ) -> DataSource:
     if is_user(source):
@@ -22,6 +23,10 @@ def get_or_create_source(
         query = query.filter(DataSource.model == model)
     if version is not None:
         query = query.filter(DataSource.version == version)
+    if attributes is not None:
+        query = query.filter(
+            DataSource.attributes_hash == DataSource.hash_attributes(attributes)
+        )
     if is_user(source):
         query = query.filter(DataSource.user == source)
     elif isinstance(source, str):
@@ -36,7 +41,11 @@ def get_or_create_source(
             if source_type is None:
                 raise TypeError("Please specify a source type")
             _source = DataSource(
-                name=source, model=model, version=version, type=source_type
+                name=source,
+                model=model,
+                version=version,
+                type=source_type,
+                attributes=attributes,
             )
         current_app.logger.info(f"Setting up {_source} as new data source...")
         db.session.add(_source)