feat: add attributes column to data_source table

.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
 }

documentation/changelog.rst

@@ -12,6 +12,7 @@ New features
 * Allow deleting multiple sensors with a single call to ``flexmeasures delete sensor`` by passing the ``--id`` option multiple times [see `PR #734 <https://www.github.com/FlexMeasures/flexmeasures/pull/734>`_]
 * Make it a lot easier to read off the color legend on the asset page, especially when showing many sensors, as they will now be ordered from top to bottom in the same order as they appear in the chart (as defined in the ``sensors_to_show`` attribute), rather than alphabetically [see `PR #742 <https://www.github.com/FlexMeasures/flexmeasures/pull/742>`_]
 * Having percentages within the [0, 100] domain is such a common use case that we now always include it in sensor charts with % units, making it easier to read off individual charts and also to compare across charts [see `PR #739 <https://www.github.com/FlexMeasures/flexmeasures/pull/739>`_]
+* DataSource table now allows storing arbitrary attributes as a JSON (without content validation), similar to the Sensor and GenericAsset tables [see `PR #750 <https://www.github.com/FlexMeasures/flexmeasures/pull/750>`_]
 
 Bugfixes
 -----------

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,17 @@ 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 in the DataSource's attributes column."""
+        return self.attributes.get(attribute, default)
+
+    def has_attribute(self, attribute: str) -> bool:
+        return attribute in self.attributes
+
+    def set_attribute(self, attribute: str, value):
+        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)