Add endpoint for getting asset annotations

flexmeasures/api/dev/sensors.py

@@ -1,5 +1,8 @@
+from __future__ import annotations
+
 import json
 import warnings
+from typing import Any
 
 from flask_classful import FlaskView, route
 from flask_security import current_user
@@ -20,6 +23,10 @@
 from flexmeasures.data.models.time_series import Sensor
 from flexmeasures.data.services.annotations import prepare_annotations_for_chart
 from flexmeasures.ui.utils.view_utils import set_session_variables
+from flexmeasures.data.models.annotations import (
+    SensorAnnotationRelationship,
+    GenericAssetAnnotationRelationship,
+)
 
 
 class SensorAPI(FlaskView):
@@ -132,20 +139,13 @@ def get_chart_annotations(self, id: int, sensor: Sensor, **kwargs):
 
         Annotations for use in charts (in case you have the chart specs already).
         """
-        event_starts_after = kwargs.get("event_starts_after", None)
-        event_ends_before = kwargs.get("event_ends_before", None)
-        df = sensor.generic_asset.search_annotations(
-            annotations_after=event_starts_after,
-            annotations_before=event_ends_before,
-            as_frame=True,
+        df = get_annotations_data(
+            sensor_id=id,
+            asset_or_sensor=sensor,
+            relationship_module=SensorAnnotationRelationship,
+            kwargs=kwargs,
         )
 
-        # Wrap and stack annotations
-        df = prepare_annotations_for_chart(df)
-
-        # Return JSON records
-        df = df.reset_index()
-        df["source"] = df["source"].astype(str)
         return df.to_json(orient="records")
 
     @route("/<id>", strict_slashes=False)
@@ -186,6 +186,70 @@ def get(self, id: int, asset: GenericAsset):
         attributes = ["name", "timezone", "timerange_of_sensors_to_show"]
         return {attr: getattr(asset, attr) for attr in attributes}
 
+    @route("/<id>/chart_annotations", strict_slashes=False)
+    @use_kwargs(
+        {"asset": AssetIdField(data_key="id")},
+        location="path",
+    )
+    @use_kwargs(
+        {
+            "event_starts_after": AwareDateTimeField(format="iso", required=False),
+            "event_ends_before": AwareDateTimeField(format="iso", required=False),
+            "beliefs_after": AwareDateTimeField(format="iso", required=False),
+            "beliefs_before": AwareDateTimeField(format="iso", required=False),
+        },
+        location="query",
+    )
+    @permission_required_for_context("read", ctx_arg_name="asset")
+    def get_chart_annotations(self, id: int, asset: GenericAsset, **kwargs):
+        """GET from /asset/<id>/chart_annotations
+
+        .. :quickref: Chart; Download annotations for use in charts
+
+        Annotations for use in charts (in case you have the chart specs already).
+        """
+        df = get_annotations_data(
+            sensor_id=None,
+            asset_or_sensor=asset,
+            relationship_module=GenericAssetAnnotationRelationship,
+            kwargs=kwargs,
+        )
+
+        return df.to_json(orient="records")
+
+
+def get_annotations_data(
+    sensor_id: int | None,
+    asset_or_sensor: GenericAsset | Sensor,
+    relationship_module: Any,
+    kwargs,
+):
+    """
+    This function fetches a sensor or an asset annotations data
+    """
+    event_starts_after = kwargs.get("event_starts_after", None)
+    event_ends_before = kwargs.get("event_ends_before", None)
+    if asset_or_sensor is GenericAsset:
+        asset_or_sensor_class = asset_or_sensor.generic_asset
+    else:
+        asset_or_sensor_class = asset_or_sensor
+
+    df = asset_or_sensor_class.search_annotations(
+        annotations_after=event_starts_after,
+        annotations_before=event_ends_before,
+        as_frame=True,
+        sensor_id=sensor_id,
+        relationship_module=relationship_module,
+    )
+
+    # Wrap and stack annotations
+    df = prepare_annotations_for_chart(df)
+
+    # Return JSON records
+    df = df.reset_index()
+    df["source"] = df["source"].astype(str)
+    return df
+
 
 def get_sensor_or_abort(id: int) -> Sensor:
     """

flexmeasures/data/models/generic_assets.py

@@ -15,7 +15,11 @@
 from timely_beliefs import BeliefsDataFrame, utils as tb_utils
 
 from flexmeasures.data import db
-from flexmeasures.data.models.annotations import Annotation, to_annotation_frame
+from flexmeasures.data.models.annotations import (
+    Annotation,
+    to_annotation_frame,
+    GenericAssetAnnotationRelationship,
+)
 from flexmeasures.data.models.charts import chart_type_to_chart_specs
 from flexmeasures.data.models.data_sources import DataSource
 from flexmeasures.data.models.parsing_utils import parse_source_arg
@@ -355,6 +359,8 @@ def search_annotations(
         annotation_type: str = None,
         include_account_annotations: bool = False,
         as_frame: bool = False,
+        sensor_id: int = None,
+        relationship_module: Any = GenericAssetAnnotationRelationship,
     ) -> list[Annotation] | pd.DataFrame:
         """Return annotations assigned to this asset, and optionally, also those assigned to the asset's account.
 
@@ -366,11 +372,12 @@ def search_annotations(
         parsed_sources = parse_source_arg(source)
         annotations = db.session.scalars(
             query_asset_annotations(
-                asset_id=self.id,
+                asset_or_sensor_id=self.id if sensor_id is None else sensor_id,
                 annotations_after=annotations_after,
                 annotations_before=annotations_before,
                 sources=parsed_sources,
                 annotation_type=annotation_type,
+                relationship_module=relationship_module,
             )
         ).all()
         if include_account_annotations and self.owner is not None:

flexmeasures/data/queries/annotations.py

@@ -8,26 +8,33 @@
 from flexmeasures.data.models.annotations import (
     Annotation,
     GenericAssetAnnotationRelationship,
+    SensorAnnotationRelationship,
 )
 from flexmeasures.data.models.data_sources import DataSource
 
 
 def query_asset_annotations(
-    asset_id: int,
+    asset_or_sensor_id: int,
+    relationship_module,
     annotations_after: datetime | None = None,
     annotations_before: datetime | None = None,
     sources: list[DataSource] | None = None,
     annotation_type: str | None = None,
 ) -> Query:
     """Match annotations assigned to the given asset."""
-    query = (
-        select(Annotation)
-        .join(GenericAssetAnnotationRelationship)
-        .filter(
-            GenericAssetAnnotationRelationship.generic_asset_id == asset_id,
+    query = select(Annotation)
+    if relationship_module is GenericAssetAnnotationRelationship:
+        query = query.join(
+            GenericAssetAnnotationRelationship,
             GenericAssetAnnotationRelationship.annotation_id == Annotation.id,
+        ).filter(
+            GenericAssetAnnotationRelationship.generic_asset_id == asset_or_sensor_id
         )
-    )
+    else:
+        query = query.join(
+            SensorAnnotationRelationship,
+            SensorAnnotationRelationship.annotation_id == Annotation.id,
+        ).filter(SensorAnnotationRelationship.sensor_id == asset_or_sensor_id)
 
     if annotations_after is not None:
         query = query.filter(

flexmeasures/ui/templates/base.html

@@ -321,7 +321,7 @@
     async function embedAndLoad(chartSpecsPath, elementId, datasetName, previousResult, startDate, endDate) {
         var combineLegend = 'true';
 
-        await vegaEmbed('#'+elementId, chartSpecsPath + 'dataset_name=' + datasetName + '&combine_legend='+ combineLegend + '&width=container&include_sensor_annotations=false&include_asset_annotations=false&chart_type=' + chartType, {{ chart_options | safe }})
+        await vegaEmbed('#'+elementId, chartSpecsPath + 'dataset_name=' + datasetName + '&combine_legend='+ combineLegend + '&width=container&include_sensor_annotations=true&include_asset_annotations=true&chart_type=' + chartType, {{ chart_options | safe }})
         .then(function (result) {
 
             // Create a custom menu item for exporting to CSV
@@ -445,15 +445,15 @@
             })
             .then(function(response) { return response.json(); }),
 
-            /**
+
             // Fetch annotations
-            fetch(dataPath + '/chart_annotations?event_starts_after=' + queryStartDate + '&event_ends_before=' + queryEndDate, {
+            fetch(dataDevPath + '/chart_annotations?event_starts_after=' + queryStartDate + '&event_ends_before=' + queryEndDate, {
                 method: "GET",
                 headers: {"Content-Type": "application/json"},
                 signal: signal,
             })
             .then(function(response) { return response.json(); }),
-            */
+
 
             // Embed chart
             embedAndLoad(chartSpecsPath + 'event_starts_after=' + queryStartDate + '&event_ends_before=' + queryEndDate + '&', elementId, datasetName, previousResult, startDate, endDate),
@@ -462,9 +462,9 @@
             vegaView.change(datasetName, vega.changeset().remove(vega.truthy).insert(result[0])).resize().run();
             previousResult = result[0];
             checkSourceMasking(previousResult);
-            /**
+
             vegaView.change(datasetName + '_annotations', vega.changeset().remove(vega.truthy).insert(result[1])).resize().run();
-            */
+
         }).catch(console.error);
     });