Stream usage prediction events instead of loading them into memory (#181770)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>

homeassistant/components/usage_prediction/common_control.py

@@ -1,15 +1,12 @@
 """Code to generate common control usage patterns."""
 
 from collections import Counter
-from collections.abc import Callable, Sequence
 from datetime import datetime, timedelta
 from functools import cache
 import logging
 from typing import Any, Literal, cast
 
 from sqlalchemy import select
-from sqlalchemy.engine.row import Row
-from sqlalchemy.orm import Session
 
 from homeassistant.components.recorder import get_instance
 from homeassistant.components.recorder.db_schema import EventData, Events, EventTypes
@@ -30,6 +27,9 @@ TIME_CATEGORIES = ["morning", "afternoon", "evening", "night"]
 
 RESULTS_TO_INCLUDE = 8
 
+# Rows fetched per round trip while streaming the events query
+QUERY_YIELD_PER = 4096
+
 # List of domains for which we want to track usage
 ALLOWED_DOMAINS = {
     # Entity platforms
@@ -94,83 +94,16 @@ async def async_predict_common_control(
     recorder = get_instance(hass)
     ent_reg = er.async_get(hass)
 
-    # Execute the database operation in the recorder's executor
-    data = await recorder.async_add_executor_job(
-        _fetch_with_session, hass, _fetch_and_process_data, ent_reg, user_id
-    )
-    # Prepare a dictionary to track results
-    results: dict[str, Counter[str]] = {
-        time_cat: Counter() for time_cat in TIME_CATEGORIES
+    allowed_entities = {
+        entity_id
+        for entity_id in hass.states.async_entity_ids(ALLOWED_DOMAINS)
+        if not ((entry := ent_reg.async_get(entity_id)) and entry.hidden)
     }
 
-    allowed_entities = set(hass.states.async_entity_ids(ALLOWED_DOMAINS))
-    hidden_entities: set[str] = set()
-
-    # Keep track of contexts that we processed so that we will only process
-    # the first service call in a context, and not subsequent calls.
-    context_processed: set[bytes] = set()
-    # Execute the query
-    context_id: bytes
-    time_fired_ts: float
-    shared_data: str | None
-    local_time_zone = dt_util.get_default_time_zone()
-    for context_id, time_fired_ts, shared_data in data:
-        # Skip if we have already processed an event that was part of this context
-        if context_id in context_processed:
-            continue
-
-        # Mark this context as processed
-        context_processed.add(context_id)
-
-        # Parse the event data
-        if not time_fired_ts or not shared_data:
-            continue
-
-        try:
-            event_data = json_loads_object(shared_data)
-        except (ValueError, TypeError) as err:
-            _LOGGER.debug("Failed to parse event data: %s", err)
-            continue
-
-        # Empty event data, skipping
-        if not event_data:
-            continue
-
-        service_data = cast(dict[str, Any] | None, event_data.get("service_data"))
-
-        # No service data found, skipping
-        if not service_data:
-            continue
-
-        entity_ids: str | list[str] | None = service_data.get("entity_id")
-
-        # No entity IDs found, skip this event
-        if entity_ids is None:
-            continue
-
-        if not isinstance(entity_ids, list):
-            entity_ids = [entity_ids]
-
-        # Convert to local time for time category determination
-        period = time_category(
-            datetime.fromtimestamp(time_fired_ts, local_time_zone).hour
-        )
-        period_results = results[period]
-
-        # Count entity usage
-        for entity_id in entity_ids:
-            if entity_id not in allowed_entities or entity_id in hidden_entities:
-                continue
-
-            if (
-                entity_id not in period_results
-                and (entry := ent_reg.async_get(entity_id))
-                and entry.hidden
-            ):
-                hidden_entities.add(entity_id)
-                continue
-
-            period_results[entity_id] += 1
+    # Execute the database operation in the recorder's executor
+    results = await recorder.async_add_executor_job(
+        _fetch_and_process_data, hass, user_id, allowed_entities
+    )
 
     return EntityUsagePredictions(
         morning=[
@@ -190,9 +123,13 @@ async def async_predict_common_control(
 
 
 def _fetch_and_process_data(
-    session: Session, ent_reg: er.EntityRegistry, user_id: str
-) -> Sequence[Row[tuple[bytes | None, float | None, str | None]]]:
-    """Fetch and process service call events from the database."""
+    hass: HomeAssistant, user_id: str, allowed_entities: set[str]
+) -> dict[str, Counter[str]]:
+    """Count service call events per entity and time category.
+
+    Rows are streamed and reduced here so only the counters cross the
+    executor boundary; a busy user can have millions of matching events.
+    """
     thirty_days_ago_ts = (dt_util.utcnow() - timedelta(days=30)).timestamp()
     user_id_bytes = uuid_hex_to_bytes_or_none(user_id)
     if not user_id_bytes:
@@ -213,16 +150,69 @@ def _fetch_and_process_data(
         .where(EventTypes.event_type == "call_service")
         .order_by(Events.time_fired_ts)
     )
-    return session.connection().execute(query).all()
 
+    # Prepare a dictionary to track results
+    results: dict[str, Counter[str]] = {
+        time_cat: Counter() for time_cat in TIME_CATEGORIES
+    }
 
-def _fetch_with_session(
-    hass: HomeAssistant,
-    fetch_func: Callable[
-        [Session], Sequence[Row[tuple[bytes | None, float | None, str | None]]]
-    ],
-    *args: object,
-) -> Sequence[Row[tuple[bytes | None, float | None, str | None]]]:
-    """Execute a fetch function with a database session."""
+    # Keep track of contexts that we processed so that we will only process
+    # the first service call in a context, and not subsequent calls.
+    context_processed: set[bytes] = set()
+    local_time_zone = dt_util.get_default_time_zone()
+
+    context_id: bytes
+    time_fired_ts: float
+    shared_data: str | None
     with session_scope(hass=hass, read_only=True) as session:
-        return fetch_func(session, *args)
+        rows = session.connection().execute(query).yield_per(QUERY_YIELD_PER)
+        for context_id, time_fired_ts, shared_data in rows:
+            # Skip if we have already processed an event that was part of this context
+            if context_id in context_processed:
+                continue
+
+            # Mark this context as processed
+            context_processed.add(context_id)
+
+            # Parse the event data
+            if not time_fired_ts or not shared_data:
+                continue
+
+            try:
+                event_data = json_loads_object(shared_data)
+            except (ValueError, TypeError) as err:
+                _LOGGER.debug("Failed to parse event data: %s", err)
+                continue
+
+            # Empty event data, skipping
+            if not event_data:
+                continue
+
+            service_data = cast(dict[str, Any] | None, event_data.get("service_data"))
+
+            # No service data found, skipping
+            if not service_data:
+                continue
+
+            entity_ids: str | list[str] | None = service_data.get("entity_id")
+
+            # No entity IDs found, skip this event
+            if entity_ids is None:
+                continue
+
+            if not isinstance(entity_ids, list):
+                entity_ids = [entity_ids]
+
+            # Convert to local time for time category determination
+            period_results = results[
+                time_category(
+                    datetime.fromtimestamp(time_fired_ts, local_time_zone).hour
+                )
+            ]
+
+            # Count entity usage
+            for entity_id in entity_ids:
+                if entity_id in allowed_entities:
+                    period_results[entity_id] += 1
+
+    return results

tests/components/usage_prediction/test_common_control.py

@@ -1,5 +1,6 @@
 """Test the common control usage prediction."""
 
+from typing import Any
 from unittest.mock import patch
 import uuid
 
@@ -408,3 +409,83 @@ async def test_different_users_separated(hass: HomeAssistant) -> None:
         evening=[],
         night=["light.user2_light"],
     )
+
+
+@pytest.mark.usefixtures("recorder_mock")
+@pytest.mark.parametrize(
+    "event_data",
+    [
+        pytest.param({}, id="no_event_data"),
+        pytest.param({"domain": "light", "service": "turn_on"}, id="no_service_data"),
+        pytest.param(
+            {
+                "domain": "light",
+                "service": "turn_on",
+                "service_data": {"brightness": 255},
+            },
+            id="no_entity_id",
+        ),
+    ],
+)
+async def test_events_without_entity_ids_ignored(
+    hass: HomeAssistant, event_data: dict[str, Any]
+) -> None:
+    """Test that service call events without entity IDs are skipped."""
+    user_id = str(uuid.uuid4())
+
+    hass.states.async_set("light.kitchen", "off")
+
+    with freeze_time("2023-07-01 10:00:00"):
+        hass.bus.async_fire(
+            EVENT_CALL_SERVICE, event_data, context=Context(user_id=user_id)
+        )
+        await hass.async_block_till_done()
+
+    await async_wait_recording_done(hass)
+
+    with freeze_time("2023-07-02 10:00:00"):
+        results = await async_predict_common_control(hass, user_id)
+
+    assert results == EntityUsagePredictions()
+
+
+@pytest.mark.usefixtures("recorder_mock")
+@pytest.mark.parametrize(
+    "json_loads_kwargs",
+    [
+        pytest.param({"side_effect": ValueError("bad json")}, id="invalid_json"),
+        pytest.param({"return_value": {}}, id="empty_object"),
+    ],
+)
+async def test_unusable_event_data_ignored(
+    hass: HomeAssistant, json_loads_kwargs: dict[str, Any]
+) -> None:
+    """Test that events whose stored data cannot be used are skipped."""
+    user_id = str(uuid.uuid4())
+
+    hass.states.async_set("light.kitchen", "off")
+
+    with freeze_time("2023-07-01 10:00:00"):
+        hass.bus.async_fire(
+            EVENT_CALL_SERVICE,
+            {
+                "domain": "light",
+                "service": "turn_on",
+                "service_data": {"entity_id": "light.kitchen"},
+            },
+            context=Context(user_id=user_id),
+        )
+        await hass.async_block_till_done()
+
+    await async_wait_recording_done(hass)
+
+    with (
+        freeze_time("2023-07-02 10:00:00"),
+        patch(
+            "homeassistant.components.usage_prediction.common_control.json_loads_object",
+            **json_loads_kwargs,
+        ),
+    ):
+        results = await async_predict_common_control(hass, user_id)
+
+    assert results == EntityUsagePredictions()