feat: add attributes column to data_source table (#750)

* feat: revision to add `attributes` column to the `data_source` table Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add `attributes` column to the DataSource model Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add sensors relationship in DataSource Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * fix: make sensors relationship viewonly Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add helper methods to DataSource Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add attributes hash Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add attributes to the function get_or_create_source Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add attribute hash to get_or_create_source Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * changing backref from "dynamic" to "select" Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: add hash_attributes static method Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * fix: use hash_attributes static method Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * feat: adding attributes_hash to the DataSource unique constraint list Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * fix: add constraint to migration and downgrade Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * fix: only returning keys from the attributes field Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * docs: fix docstring Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * fix: use default value Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * fix: allow creating new attributes with the method `set_attributes` Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * docs: add changelog entry Signed-off-by: Victor Garcia Reolid <victor@seita.nl> * docs: add db upgrade warning Signed-off-by: Victor Garcia Reolid <victor@seita.nl> --------- Signed-off-by: Victor Garcia Reolid <victor@seita.nl>
FlexMeasures · Jul 6, 2023 · 3772b08 · 3772b08
1 parent 9291830
commit 3772b08
Show file tree

Hide file tree

Showing 5 changed files with 110 additions and 4 deletions.
diff --git a/.vscode/settings.json b/.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
 }
diff --git a/documentation/changelog.rst b/documentation/changelog.rst
@@ -6,12 +6,15 @@ FlexMeasures Changelog
 v0.15.0 | July XX, 2023
 ============================
 
+.. warning:: Upgrading to this version requires running ``flexmeasures db upgrade`` (you can create a backup first with ``flexmeasures db-ops dump``).
+
 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
 -----------

diff --git a/flexmeasures/data/migrations/versions/2ac7fb39ce0c_add_attribute_column_to_data_source.py b/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")
diff --git a/flexmeasures/data/models/data_sources.py b/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
diff --git a/flexmeasures/data/services/data_sources.py b/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)