From 753deee48c80663fe9e498c47b079abf2d9110fa Mon Sep 17 00:00:00 2001 From: Paulus Schoutsen Date: Thu, 10 Sep 2026 13:44:36 -0400 Subject: [PATCH] Stream usage prediction events instead of loading them into memory (#181770) Co-authored-by: Claude Fable 5.1 --- .../usage_prediction/common_control.py | 172 +++++++++--------- .../usage_prediction/test_common_control.py | 81 +++++++++ 2 files changed, 162 insertions(+), 91 deletions(-) diff --git a/homeassistant/components/usage_prediction/common_control.py b/homeassistant/components/usage_prediction/common_control.py index 15d17dceaaef..d3f1da724035 100644 --- a/homeassistant/components/usage_prediction/common_control.py +++ b/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 diff --git a/tests/components/usage_prediction/test_common_control.py b/tests/components/usage_prediction/test_common_control.py index b9217fca004f..244fd7163747 100644 --- a/tests/components/usage_prediction/test_common_control.py +++ b/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()