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
 }

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

@@ -0,0 +1,34 @@
+"""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),
+    )
+
+
+def downgrade():
+    pass

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:
@@ -68,18 +71,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="dynamic"),
+        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 +104,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 +166,21 @@ def to_dict(self) -> dict:
             type=self.type if self.type in ("forecaster", "scheduler") else "other",
             description=self.description,
         )
+
+    def get_attribute(self, attribute: str, default: Any = None) -> Any:
+        """Looks for the attribute on the DataSource.
+        If not found, returns the default.
+        """
+        if hasattr(self, attribute):
+            return getattr(self, attribute)
+        if attribute in self.attributes:
+            return self.attributes[attribute]
+
+        return default
+
+    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

@@ -2,6 +2,9 @@
 
 from flask import current_app
 
+import json
+import hashlib
+
 from flexmeasures import User
 from flexmeasures.data import db
 from flexmeasures.data.models.data_sources import DataSource
@@ -13,15 +16,22 @@ 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:
+    attributes_hash = None
     if is_user(source):
         source_type = "user"
     query = DataSource.query.filter(DataSource.type == source_type)
     if model is not None:
         query = query.filter(DataSource.model == model)
     if version is not None:
         query = query.filter(DataSource.version == version)
+    if attributes is not None:
+        attributes_hash = hashlib.sha256(
+            json.dumps(attributes).encode("utf-8")
+        ).digest()
+        query = query.filter(DataSource.attributes_hash == attributes_hash)
     if is_user(source):
         query = query.filter(DataSource.user == source)
     elif isinstance(source, str):
@@ -36,7 +46,12 @@ 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,
+                attributes_hash=attributes_hash,
             )
         current_app.logger.info(f"Setting up {_source} as new data source...")
         db.session.add(_source)