Fast most_recent_only belief lookup for sensor statuses

documentation/changelog.rst

@@ -23,6 +23,7 @@ Infrastructure / Support
 -------------------------
 * Speed up sensor data queries and free up disk space by reordering the ``timed_belief`` primary key to lead with ``sensor_id`` and dropping the indexes it makes redundant, in a migration that runs online and so needs no maintenance window (though it can take a while on a large database) [see `PR #2378 <https://www.github.com/FlexMeasures/flexmeasures/pull/2378>`_]
 * Look up which data sources recorded for which sensors from a small summary table instead of scanning the beliefs table [see `PR #2382 <https://www.github.com/FlexMeasures/flexmeasures/pull/2382>`_]
+* The asset page no longer slows down with the amount of data a sensor holds while working out how up to date each of its sensors is [see `PR #2463 <https://www.github.com/FlexMeasures/flexmeasures/pull/2463>`_]
 * Shrink the Docker image by excluding dev-only dependencies, pruning stray ``docs``/``examples`` payloads bundled by ``sktime``/``scikit-base`` (issue: https://github.com/sktime/sktime/issues/10891), stripping the symbol tables that the compiled extensions ship with, and dropping the ``sktime``-backed belief-formation extra of ``timely-beliefs``, which FlexMeasures does not use [see `PR #2438 <https://www.github.com/FlexMeasures/flexmeasures/pull/2438>`_, `PR #2439 <https://www.github.com/FlexMeasures/flexmeasures/pull/2439>`_ and `PR #2440 <https://www.github.com/FlexMeasures/flexmeasures/pull/2440>`_]
 
 * The UI's JavaScript modules can now be tested, by running them in a headless browser from pytest, without adding a Node.js toolchain [see `PR #2435 <https://www.github.com/FlexMeasures/flexmeasures/pull/2435>`_]

flexmeasures/data/services/sensors.py

@@ -26,6 +26,7 @@
 from flexmeasures.data.models.audit_log import AssetAuditLog
 from flexmeasures.data.models.automations import Automation
 from flexmeasures.data.models.data_sources import DataSource, DEFAULT_DATASOURCE_TYPES
+from flexmeasures.data.models.parsing_utils import parse_source_arg
 from flexmeasures.data.models.generic_assets import GenericAsset
 from flexmeasures.data.models.planning.devices import INFLEXIBLE_DEVICE_KEYS
 from flexmeasures.data.schemas.generic_assets import SensorsToShowSchema
@@ -449,19 +450,68 @@ def get_sensors(
     return db.session.scalars(sensor_query).all()
 
 
+def _sensor_sources_by_type(
+    sensor: Sensor, staleness_search: dict
+) -> dict[str, list[DataSource]]:
+    """Group the sensor's data sources by source type, honouring the source filters of the staleness search.
+
+    Reading which sources ever recorded for this sensor is a lookup in the ``sensor_data_source`` summary,
+    so it costs a handful of rows rather than a scan of the beliefs table.
+    Only the default source types are considered, since those are the ones a status is reported for.
+
+    The summary is a superset (see :class:`~flexmeasures.data.models.data_sources.SensorDataSource`),
+    so a source may be listed whose beliefs have since been deleted.
+    That only costs a belief query returning nothing; no source that has data can be missing.
+    """
+    sources = sensor.search_data_sources(
+        source_types=DEFAULT_DATASOURCE_TYPES,
+        exclude_source_types=staleness_search.get("exclude_source_types"),
+    )
+    requested_sources = parse_source_arg(staleness_search.get("source"))
+    if requested_sources is not None:
+        requested_source_ids = {source.id for source in requested_sources}
+        sources = [source for source in sources if source.id in requested_source_ids]
+
+    sources_by_type: dict[str, list[DataSource]] = {}
+    for source in sources:
+        sources_by_type.setdefault(source.type, []).append(source)
+    return sources_by_type
+
+
 def _get_sensor_bdfs_by_source_type(
     sensor: Sensor, staleness_search: dict
 ) -> dict[str, BeliefsDataFrame] | None:
     """Get latest event, split by source type for a given sensor with given search parameters.
     We only look for the default data source types!
+
+    Each type is searched for by naming its sources explicitly, rather than by filtering on the type of the source.
+    A type filter cannot be served by an index on the beliefs table,
+    so the "most recent belief" query would walk the sensor's events from the newest backwards,
+    rechecking the type of each belief's source until it found a match;
+    it would read every belief the sensor has whenever this type recorded none of them.
+    Naming the sources instead lets the primary key answer the query directly,
+    and lets a type with no sources at all be skipped without a query.
     """
+    sources_by_type = _sensor_sources_by_type(sensor, staleness_search)
+
+    # The source filters are already applied by the source lookup above,
+    # so passing them on as well would only re-apply them.
+    belief_search = {
+        key: value
+        for key, value in staleness_search.items()
+        if key not in ("source", "exclude_source_types")
+    }
+
     bdfs_by_source = dict()
     for source_type in DEFAULT_DATASOURCE_TYPES:
+        sources = sources_by_type.get(source_type)
+        if not sources:
+            continue
         bdf = TimedBelief.search(
             sensors=sensor,
             most_recent_only=True,
-            source_types=[source_type],
-            **staleness_search,
+            source=sources,
+            **belief_search,
         )
         if not bdf.empty:
             bdfs_by_source[source_type] = bdf

flexmeasures/data/tests/test_sensor_status_queries.py

@@ -0,0 +1,105 @@
+"""Tests for how the sensor status service looks up the most recent beliefs."""
+
+from unittest import mock
+
+import pandas as pd
+from timely_beliefs import BeliefsDataFrame
+
+from flexmeasures.data.models.data_sources import DataSource
+from flexmeasures.data.models.time_series import TimedBelief
+from flexmeasures.data.services.sensors import _get_sensor_bdfs_by_source_type
+from flexmeasures.data.utils import save_to_db
+from flexmeasures.tests.utils import get_test_sensor
+
+
+def add_belief(db, sensor, source_name: str, source_type: str) -> DataSource:
+    """Record one belief for the sensor from a new source of the given type."""
+    source = DataSource(name=source_name, type=source_type)
+    db.session.add(source)
+    db.session.commit()
+    save_to_db(
+        BeliefsDataFrame(
+            [
+                TimedBelief(
+                    sensor=sensor,
+                    source=source,
+                    event_start=pd.Timestamp("2021-03-28 16:00:00+00:00"),
+                    belief_time=pd.Timestamp("2021-03-27 08:00:00+00:00"),
+                    event_value=1.0,
+                )
+            ]
+        )
+    )
+    db.session.commit()
+    return source
+
+
+def test_only_source_types_that_recorded_are_queried(setup_beliefs, db):
+    """A source type that never recorded for this sensor must not cost a belief query.
+
+    Such a query cannot be answered from an index: it would walk the sensor's events,
+    from the newest backwards, checking the type of every belief's source,
+    and read all of them before concluding that this type recorded none.
+    The sensor's sources are known from the summary table, so those types can be skipped entirely.
+    """
+    sensor = get_test_sensor(db)
+
+    with mock.patch.object(
+        TimedBelief, "search", wraps=TimedBelief.search
+    ) as search_spy:
+        bdfs = _get_sensor_bdfs_by_source_type(sensor=sensor, staleness_search={})
+
+    # The fixture records beliefs from one source only, of type "demo script"
+    assert set(bdfs) == {"demo script"}
+    assert search_spy.call_count == 1
+
+    # And that one query names its sources, rather than filtering on the type of the source
+    kwargs = search_spy.call_args.kwargs
+    assert "source_types" not in kwargs
+    assert [source.type for source in kwargs["source"]] == ["demo script"]
+
+
+def test_a_second_source_type_is_picked_up(setup_beliefs, db):
+    """Adding a source of another type must add that type, and only that type, to the results."""
+    sensor = get_test_sensor(db)
+    add_belief(db, sensor, "A reporter", "reporter")
+
+    with mock.patch.object(
+        TimedBelief, "search", wraps=TimedBelief.search
+    ) as search_spy:
+        bdfs = _get_sensor_bdfs_by_source_type(sensor=sensor, staleness_search={})
+
+    assert set(bdfs) == {"demo script", "reporter"}
+    assert search_spy.call_count == 2
+
+
+def test_source_filter_in_the_staleness_search_is_honoured(setup_beliefs, db):
+    """A search restricted to one source must report on that source only.
+
+    The source types are resolved before the beliefs are queried,
+    so this filter has to be applied to that resolution as well.
+    """
+    sensor = get_test_sensor(db)
+    reporter_source = add_belief(db, sensor, "The only source of interest", "reporter")
+
+    bdfs = _get_sensor_bdfs_by_source_type(
+        sensor=sensor, staleness_search=dict(source=[reporter_source])
+    )
+
+    assert set(bdfs) == {"reporter"}
+
+
+def test_excluded_source_types_are_honoured(setup_beliefs, db):
+    """An excluded source type must not be reported on, nor queried for."""
+    sensor = get_test_sensor(db)
+    add_belief(db, sensor, "An excluded reporter", "reporter")
+
+    with mock.patch.object(
+        TimedBelief, "search", wraps=TimedBelief.search
+    ) as search_spy:
+        bdfs = _get_sensor_bdfs_by_source_type(
+            sensor=sensor, staleness_search=dict(exclude_source_types=["reporter"])
+        )
+
+    assert set(bdfs) == {"demo script"}
+    assert search_spy.call_count == 1