feat: save/fetch Reporter as DataSource

.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/app.py

@@ -5,6 +5,7 @@
 from __future__ import annotations
 
 import time
+from copy import copy
 import os
 from pathlib import Path
 from datetime import date
@@ -121,8 +122,14 @@ def create(  # noqa C901
     from flexmeasures.utils.coding_utils import get_classes_module
     from flexmeasures.data.models import reporting, planning
 
-    app.reporters = get_classes_module("flexmeasures.data.models", reporting.Reporter)
-    app.schedulers = get_classes_module("flexmeasures.data.models", planning.Scheduler)
+    reporters = get_classes_module("flexmeasures.data.models", reporting.Reporter)
+    schedulers = get_classes_module("flexmeasures.data.models", planning.Scheduler)
+
+    app.reporters = reporters
+    app.schedulers = schedulers
+
+    app.data_generators = copy(reporters)  # use copy to avoid mutating app.reporters
+    app.data_generators.update(schedulers)
 
     # add auth policy
 

flexmeasures/cli/tests/conftest.py

@@ -72,7 +72,7 @@ def setup_dummy_data(db, app):
 
 @pytest.fixture(scope="module")
 @pytest.mark.skip_github
-def reporter_config_raw(app, db, setup_dummy_data):
+def reporter_config(app, db, setup_dummy_data):
     """
     This reporter_config defines the operations to add up the
     values of the sensors 1 and 2 and resamples the result to a
@@ -81,8 +81,8 @@ def reporter_config_raw(app, db, setup_dummy_data):
 
     sensor1, sensor2, report_sensor = setup_dummy_data
 
-    reporter_config_raw = dict(
-        beliefs_search_configs=[dict(sensor=sensor1.id), dict(sensor=sensor2.id)],
+    reporter_config = dict(
+        input_variables=["sensor_1", "sensor_2"],
         transformations=[
             dict(
                 df_input="sensor_1",
@@ -95,4 +95,4 @@ def reporter_config_raw(app, db, setup_dummy_data):
         final_df_output="df_agg",
     )
 
-    return reporter_config_raw
+    return reporter_config

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,20 +1,60 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+import json
+from typing import TYPE_CHECKING, Any, Type
+from sqlalchemy.ext.mutable import MutableDict
 
 import timely_beliefs as tb
 
 from flexmeasures.data import db
 from flask import current_app
+import hashlib
+
+from marshmallow import Schema
 
 
 if TYPE_CHECKING:
     from flexmeasures.data.models.user import User
 
 
-class DataGeneratorMixin:
+class DataGenerator:
     _data_source: DataSource | None = None
 
+    _config: dict = None
+
+    _inputs_schema: Type[Schema] | None = None
+    _config_schema: Type[Schema] | None = None
+
+    def __init__(self, config: dict | None = None, **kwargs) -> None:
+        if config is None:
+            _config = kwargs
+        else:
+            if self._config_schema:
+                _config = self._config_schema.load(config)
+            else:
+                _config = config
+
+        self._config = _config
+
+    def _compute(self, **kwargs):
+        raise NotImplementedError()
+
+    def compute(self, inputs: dict = None, **kwargs):
+        if inputs is None:
+            _inputs = kwargs
+            self.validate_deserialized_inputs(_inputs)
+        else:
+            if self._inputs_schema:
+                _inputs = self._inputs_schema.load(inputs)
+
+            else:  # skip validation
+                _inputs = inputs
+
+        return self._compute(**_inputs)
+
+    def validate_deserialized_inputs(self, inputs: dict):
+        self._inputs_schema.load(self._inputs_schema.dump(inputs))
+
     @classmethod
     def get_data_source_info(cls: type) -> dict:
         """
@@ -23,33 +63,35 @@ def get_data_source_info(cls: type) -> dict:
         See for instance get_data_source_for_job().
         """
         source_info = dict(
-            name=current_app.config.get("FLEXMEASURES_DEFAULT_DATASOURCE")
+            source=current_app.config.get("FLEXMEASURES_DEFAULT_DATASOURCE")
         )  # default
 
         from flexmeasures.data.models.planning import Scheduler
         from flexmeasures.data.models.reporting import Reporter
 
         if issubclass(cls, Reporter):
-            source_info["type"] = "reporter"
+            source_info["source_type"] = "reporter"
         elif issubclass(cls, Scheduler):
-            source_info["type"] = "scheduler"
+            source_info["source_type"] = "scheduler"
         else:
-            source_info["type"] = "undefined"
+            source_info["source_type"] = "undefined"
+
+        source_info["model"] = cls.__name__
 
         return source_info
 
     @property
-    def data_source(self):
+    def data_source(self) -> "DataSource | None":
         from flexmeasures.data.services.data_sources import get_or_create_source
 
         if self._data_source is None:
             data_source_info = self.get_data_source_info()
-
-            self._data_source = get_or_create_source(
-                source=data_source_info.get("name"),
-                source_type=data_source_info.get("type"),
+            data_source_info["attributes"] = dict(
+                config=self._config_schema.dump(self._config)
             )
 
+            self._data_source = get_or_create_source(**data_source_info)
+
         return self._data_source
 
 
@@ -68,18 +110,32 @@ 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,
+    )
+
+    _data_generator: DataGenerator | None = None
+
     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,9 +145,53 @@ 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)
 
+    @property
+    def data_generator(self):
+        if self._data_generator:
+            return self._data_generator
+
+        data_generator = None
+
+        if self.type not in ["scheduler", "forecaster", "reporter"]:
+            current_app.logger.warning(
+                "Only the classes Scheduler, Forecaster and Reporters are DataGenerator's."
+            )
+            return None
+
+        if not self.model:
+            current_app.logger.warning(
+                "There's no DataGenerator class defined in this DataSource."
+            )
+            return None
+
+        if self.model not in current_app.data_generators:
+            current_app.logger.warning(
+                "DataGenerator `{self.model}` not registered in this FlexMeasures instance."
+            )
+            return None
+
+        # fetch DataGenerator details
+        data_generator_details = self.attributes.get("data_generator", {})
+        config = data_generator_details.get("config", {})
+
+        # create DataGenerator class and assign the current DataSource (self) as its source
+        data_generator = current_app.data_generators[self.model](config=config)
+        data_generator._data_source = self
+
+        self._data_generator = data_generator
+
+        return self._data_generator
+
     @property
     def label(self):
         """Human-readable label (preferably not starting with a capital letter, so it can be used in a sentence)."""
@@ -144,3 +244,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