Annotations (#343)

Introduce a database model for annotations and their relationship to assets and sensors. Also add a CLI option to load public holidays and store them as annotations. * Add Annotation db model, including relationships Signed-off-by: F.N. Claessen <felix@seita.nl> * Add representation of Annotation objects Signed-off-by: F.N. Claessen <felix@seita.nl> * Add CLI command to add holidays Signed-off-by: F.N. Claessen <felix@seita.nl> * Add migration Signed-off-by: F.N. Claessen <felix@seita.nl> * One source per country Signed-off-by: F.N. Claessen <felix@seita.nl> * Refactor migration Signed-off-by: F.N. Claessen <felix@seita.nl> * Add annotations_assets table Signed-off-by: F.N. Claessen <felix@seita.nl> * Improve inflection of plural Signed-off-by: F.N. Claessen <felix@seita.nl> * Add a CLI option to relate new holidays annotations to assets Signed-off-by: F.N. Claessen <felix@seita.nl> * Add a CLI option to relate new holidays annotations to all assets of an account Signed-off-by: F.N. Claessen <felix@seita.nl> * Add annotations_sensors table Signed-off-by: F.N. Claessen <felix@seita.nl> * black Signed-off-by: F.N. Claessen <felix@seita.nl> * Fix downgrade Signed-off-by: F.N. Claessen <felix@seita.nl> * Fix table name Signed-off-by: F.N. Claessen <felix@seita.nl> * Add downgrade prompt Signed-off-by: F.N. Claessen <felix@seita.nl> * black Signed-off-by: F.N. Claessen <felix@seita.nl> * Fix docstring Signed-off-by: F.N. Claessen <felix@seita.nl> * Add dependency Signed-off-by: F.N. Claessen <felix@seita.nl> * Update docstring Signed-off-by: F.N. Claessen <felix@seita.nl> * Add missing enum name Signed-off-by: F.N. Claessen <felix@seita.nl> * Clarify use of workalendar registry Signed-off-by: F.N. Claessen <felix@seita.nl> * Use recommended way of setting up UniqueConstraint Signed-off-by: F.N. Claessen <felix@seita.nl> * Fix setup of many-to-many relationships between annotations, generic assets and sensors Signed-off-by: F.N. Claessen <felix@seita.nl> * Refactor loop Signed-off-by: F.N. Claessen <felix@seita.nl> * Log the string representation of the newly created Source object, which print its description instead of just its name Signed-off-by: F.N. Claessen <felix@seita.nl> * flake8 Signed-off-by: F.N. Claessen <felix@seita.nl> * Add changelog warning to upgrade the database Signed-off-by: F.N. Claessen <felix@seita.nl> * Typos Signed-off-by: F.N. Claessen <felix@seita.nl> * Changelog entry Signed-off-by: F.N. Claessen <felix@seita.nl> * Refactor: move new annotation classes to dedicated module Signed-off-by: F.N. Claessen <felix@seita.nl> * Refactor: factor out import statement Signed-off-by: F.N. Claessen <felix@seita.nl>
FlexMeasures · Feb 2, 2022 · 3ef89eb · 3ef89eb
1 parent b1aeca3
commit 3ef89eb
Show file tree

Hide file tree

Showing 10 changed files with 265 additions and 21 deletions.
diff --git a/documentation/changelog.rst b/documentation/changelog.rst
@@ -5,14 +5,15 @@ FlexMeasures Changelog
 v0.9.0 | February XX, 2022
 ===========================
 
+.. warning:: Upgrading to this version requires running ``flexmeasures db upgrade`` (you can create a backup first with ``flexmeasures db-ops dump``).
+
 New features
 -----------
 * Three new CLI commands for cleaning up your database: delete 1) unchanged beliefs, 2) NaN values or 3) a sensor and all of its time series data [see `PR #328 <http://www.github.com/FlexMeasures/flexmeasures/pull/328>`_]
 * Add CLI option to pass a data unit when reading in time series data from CSV, so data can automatically be converted to the sensor unit [see `PR #341 <http://www.github.com/FlexMeasures/flexmeasures/pull/341>`_]
-
-* Add CLI-commands ``flexmeasures add sensor``, ``flexmeasures add asset-type``, ``flexmeasures add beliefs`` (which were experimental features before). [see `PR #337 <http://www.github.com/FlexMeasures/flexmeasures/pull/337>`_]
-
-* add CLI comands for showing data [see `PR #339 <http://www.github.com/FlexMeasures/flexmeasures/pull/339>`_]
+* Add CLI commands ``flexmeasures add sensor``, ``flexmeasures add asset-type``, ``flexmeasures add beliefs`` (which were experimental features before). [see `PR #337 <http://www.github.com/FlexMeasures/flexmeasures/pull/337>`_]
+* Add CLI commands for showing data [see `PR #339 <http://www.github.com/FlexMeasures/flexmeasures/pull/339>`_]
+* Add CLI command for attaching annotations to assets: ``flexmeasures add holidays`` adds public holidays [see `PR #343 <http://www.github.com/FlexMeasures/flexmeasures/pull/343>`_]
 
 Bugfixes
 -----------

diff --git a/flexmeasures/cli/data_add.py b/flexmeasures/cli/data_add.py
@@ -12,12 +12,17 @@
 import getpass
 from sqlalchemy.exc import IntegrityError
 import timely_beliefs as tb
+from workalendar.registry import registry as workalendar_registry
 
 from flexmeasures.data import db
 from flexmeasures.data.services.forecasting import create_forecasting_jobs
 from flexmeasures.data.services.users import create_user
 from flexmeasures.data.models.user import Account, AccountRole, RolesAccounts
-from flexmeasures.data.models.time_series import Sensor, TimedBelief
+from flexmeasures.data.models.time_series import (
+    Sensor,
+    TimedBelief,
+)
+from flexmeasures.data.models.annotations import Annotation
 from flexmeasures.data.schemas.sensors import SensorSchema
 from flexmeasures.data.schemas.generic_assets import (
     GenericAssetSchema,
@@ -30,6 +35,7 @@
     get_or_create_source,
     get_source_or_none,
 )
+from flexmeasures.utils import flexmeasures_inflection
 from flexmeasures.utils.time_utils import server_now
 from flexmeasures.utils.unit_utils import convert_units
 
@@ -497,6 +503,77 @@ def add_beliefs(
             print("As a possible workaround, use the --allow-overwrite flag.")
 
 
+@fm_add_data.command("holidays")
+@with_appcontext
+@click.option(
+    "--year",
+    type=click.INT,
+    help="The year for which to look up holidays",
+)
+@click.option(
+    "--country",
+    "countries",
+    type=click.STRING,
+    multiple=True,
+    help="The ISO 3166-1 country/region or ISO 3166-2 sub-region for which to look up holidays (such as US, BR and DE). This argument can be given multiple times.",
+)
+@click.option(
+    "--asset-id",
+    "generic_asset_ids",
+    type=click.INT,
+    multiple=True,
+    help="Add annotations to this asset. Follow up with the asset's ID. This argument can be given multiple times.",
+)
+@click.option(
+    "--account-id",
+    "account_ids",
+    type=click.INT,
+    multiple=True,
+    help="Add annotations to all assets of this account. Follow up with the account's ID. This argument can be given multiple times.",
+)
+def add_holidays(
+    year: int,
+    countries: List[str],
+    generic_asset_ids: List[int],
+    account_ids: List[int],
+):
+    """Add holiday annotations to assets."""
+    calendars = workalendar_registry.get_calendars(countries)
+    num_holidays = {}
+    asset_query = db.session.query(GenericAsset)
+    if generic_asset_ids:
+        asset_query = asset_query.filter(GenericAsset.id.in_(generic_asset_ids))
+    if account_ids:
+        asset_query = asset_query.filter(GenericAsset.account_id.in_(account_ids))
+    assets = asset_query.all()
+    annotations = []
+    for country, calendar in calendars.items():
+        _source = get_or_create_source(
+            "workalendar", model=country, source_type="CLI script"
+        )
+        holidays = calendar().holidays(year)
+        for holiday in holidays:
+            start = pd.Timestamp(holiday[0])
+            end = start + pd.offsets.DateOffset(days=1)
+            annotations.append(
+                Annotation(
+                    name=holiday[1],
+                    start=start,
+                    end=end,
+                    source=_source,
+                    type="holiday",
+                )
+            )
+        num_holidays[country] = len(holidays)
+    db.session.add_all(annotations)
+    for asset in assets:
+        asset.annotations += annotations
+    db.session.commit()
+    print(
+        f"Successfully added holidays to {len(assets)} {flexmeasures_inflection.pluralize('asset', len(assets))}:\n{num_holidays}"
+    )
+
+
 @fm_add_data.command("forecasts")
 @with_appcontext
 @click.option(

diff --git a/flexmeasures/data/migrations/versions/7f8b8920355f_create_annotation_table.py b/flexmeasures/data/migrations/versions/7f8b8920355f_create_annotation_table.py
@@ -0,0 +1,81 @@
+"""create annotation table
+
+Revision ID: 7f8b8920355f
+Revises: c1d316c60985
+Create Date: 2022-01-29 20:23:29.996133
+
+"""
+from alembic import op
+import click
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = "7f8b8920355f"
+down_revision = "c1d316c60985"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    create_annotation_table()
+    create_annotation_asset_relationship_table()
+    create_annotation_sensor_relationship_table()
+
+
+def downgrade():
+    click.confirm(
+        "This downgrade drops the tables 'annotations_assets', 'annotations_sensors' and 'annotation'. Continue?",
+        abort=True,
+    )
+    op.drop_table("annotations_assets")
+    op.drop_table("annotations_sensors")
+    op.drop_constraint(op.f("annotation_name_key"), "annotation", type_="unique")
+    op.drop_table("annotation")
+    op.execute("DROP TYPE annotation_type;")
+
+
+def create_annotation_sensor_relationship_table():
+    op.create_table(
+        "annotations_sensors",
+        sa.Column("id", sa.Integer(), primary_key=True),
+        sa.Column("sensor_id", sa.Integer()),
+        sa.Column("annotation_id", sa.Integer()),
+        sa.ForeignKeyConstraint(("sensor_id",), ["sensor.id"]),
+        sa.ForeignKeyConstraint(("annotation_id",), ["annotation.id"]),
+    )
+
+
+def create_annotation_asset_relationship_table():
+    op.create_table(
+        "annotations_assets",
+        sa.Column("id", sa.Integer(), primary_key=True),
+        sa.Column("generic_asset_id", sa.Integer()),
+        sa.Column("annotation_id", sa.Integer()),
+        sa.ForeignKeyConstraint(("generic_asset_id",), ["generic_asset.id"]),
+        sa.ForeignKeyConstraint(("annotation_id",), ["annotation.id"]),
+    )
+
+
+def create_annotation_table():
+    op.create_table(
+        "annotation",
+        sa.Column(
+            "id", sa.Integer(), nullable=False, autoincrement=True, primary_key=True
+        ),
+        sa.Column("name", sa.String(255), nullable=False),
+        sa.Column("start", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("end", sa.DateTime(timezone=True), nullable=False),
+        sa.Column("source_id", sa.Integer(), nullable=False),
+        sa.Column(
+            "type",
+            sa.Enum("alert", "holiday", "label", name="annotation_type"),
+            nullable=False,
+        ),
+        sa.ForeignKeyConstraint(("source_id",), ["data_source.id"]),
+    )
+    op.create_unique_constraint(
+        op.f("annotation_name_key"),
+        "annotation",
+        ["name", "start", "source_id", "type"],
+    )
diff --git a/flexmeasures/data/models/annotations.py b/flexmeasures/data/models/annotations.py
@@ -0,0 +1,64 @@
+from datetime import timedelta
+
+from flexmeasures.data import db
+from flexmeasures.data.models.data_sources import DataSource
+
+
+class Annotation(db.Model):
+    """An annotation is a nominal value that applies to a specific time or time span.
+
+    Examples of annotation types:
+        - user annotation: annotation.type == "label" and annotation.source.type == "user"
+        - unresolved alert: annotation.type == "alert"
+        - resolved alert: annotation.type == "label" and annotation.source.type == "alerting script"
+        - organisation holiday: annotation.type == "holiday" and annotation.source.type == "user"
+        - public holiday: annotation.type == "holiday" and annotation.source.name == "workalendar"
+    """
+
+    id = db.Column(db.Integer, nullable=False, autoincrement=True, primary_key=True)
+    name = db.Column(db.String(255), nullable=False)
+    start = db.Column(db.DateTime(timezone=True), nullable=False)
+    end = db.Column(db.DateTime(timezone=True), nullable=False)
+    source_id = db.Column(db.Integer, db.ForeignKey(DataSource.__tablename__ + ".id"))
+    source = db.relationship(
+        "DataSource",
+        foreign_keys=[source_id],
+        backref=db.backref("annotations", lazy=True),
+    )
+    type = db.Column(db.Enum("alert", "holiday", "label", name="annotation_type"))
+    __table_args__ = (
+        db.UniqueConstraint(
+            "name",
+            "start",
+            "source_id",
+            "type",
+            name="annotation_name_key",
+        ),
+    )
+
+    @property
+    def duration(self) -> timedelta:
+        return self.end - self.start
+
+    def __repr__(self) -> str:
+        return f"<Annotation {self.id}: {self.name} ({self.type}), start: {self.start} end: {self.end}, source: {self.source}>"
+
+
+class GenericAssetAnnotationRelationship(db.Model):
+    """Links annotations to generic assets."""
+
+    __tablename__ = "annotations_assets"
+
+    id = db.Column(db.Integer(), primary_key=True)
+    generic_asset_id = db.Column(db.Integer, db.ForeignKey("generic_asset.id"))
+    annotation_id = db.Column(db.Integer, db.ForeignKey("annotation.id"))
+
+
+class SensorAnnotationRelationship(db.Model):
+    """Links annotations to sensors."""
+
+    __tablename__ = "annotations_sensors"
+
+    id = db.Column(db.Integer(), primary_key=True)
+    sensor_id = db.Column(db.Integer, db.ForeignKey("sensor.id"))
+    annotation_id = db.Column(db.Integer, db.ForeignKey("annotation.id"))
diff --git a/flexmeasures/data/models/data_sources.py b/flexmeasures/data/models/data_sources.py
@@ -87,11 +87,16 @@ def __str__(self):
 
 
 def get_or_create_source(
-    source: Union[User, str], source_type: Optional[str] = None, flush: bool = True
+    source: Union[User, str],
+    source_type: Optional[str] = None,
+    model: Optional[str] = None,
+    flush: bool = True,
 ) -> DataSource:
     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 is_user(source):
         query = query.filter(DataSource.user == source)
     elif isinstance(source, str):
@@ -100,13 +105,13 @@ def get_or_create_source(
         raise TypeError("source should be of type User or str")
     _source = query.one_or_none()
     if not _source:
-        current_app.logger.info(f"Setting up '{source}' as new data source...")
         if is_user(source):
-            _source = DataSource(user=source)
+            _source = DataSource(user=source, model=model)
         else:
             if source_type is None:
                 raise TypeError("Please specify a source type")
-            _source = DataSource(name=source, type=source_type)
+            _source = DataSource(name=source, model=model, type=source_type)
+        current_app.logger.info(f"Setting up {_source} as new data source...")
         db.session.add(_source)
         if flush:
             # assigns id so that we can reference the new object in the current db session

diff --git a/flexmeasures/data/models/generic_assets.py b/flexmeasures/data/models/generic_assets.py
@@ -47,6 +47,11 @@ class GenericAsset(db.Model, AuthModelMixin):
         foreign_keys=[generic_asset_type_id],
         backref=db.backref("generic_assets", lazy=True),
     )
+    annotations = db.relationship(
+        "Annotation",
+        secondary="annotations_assets",
+        backref=db.backref("assets", lazy="dynamic"),
+    )
 
     __table_args__ = (
         UniqueConstraint(

diff --git a/flexmeasures/data/models/time_series.py b/flexmeasures/data/models/time_series.py
@@ -47,6 +47,11 @@ class Sensor(db.Model, tb.SensorDBMixin, AuthModelMixin):
             "sensors", lazy=True, cascade="all, delete-orphan", passive_deletes=True
         ),
     )
+    annotations = db.relationship(
+        "Annotation",
+        secondary="annotations_sensors",
+        backref=db.backref("sensors", lazy="dynamic"),
+    )
 
     def __init__(
         self,

diff --git a/flexmeasures/utils/flexmeasures_inflection.py b/flexmeasures/utils/flexmeasures_inflection.py
@@ -1,7 +1,10 @@
+from __future__ import annotations
 import re
 
+import inflect
 import inflection
 
+p = inflect.engine()
 
 # Give the inflection module some help for our domain
 inflection.UNCOUNTABLES.add("solar")
@@ -30,10 +33,10 @@ def parameterize(word):
     return inflection.parameterize(word).replace("-", "_")
 
 
-def pluralize(word):
+def pluralize(word, count: str | int | None = None):
     if word.lower().split()[-1] in inflection.UNCOUNTABLES:
         return word
-    return inflection.pluralize(word)
+    return p.plural(word, count)
 
 
 def titleize(word):

diff --git a/requirements/app.in b/requirements/app.in
@@ -10,6 +10,7 @@ pint
 py-moneyed
 iso8601
 xlrd
+workalendar
 inflection
 inflect
 humanize