starred/core
1
0
Fork 0
mirror of https://github.com/home-assistant/core.git synced 2026-09-25 17:04:04 -04:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Paulus Schoutsen
2026-09-10 19:44:36 +02:00
committed by GitHub
co-authored by Claude Fable 5.1
parent 7e94b9aa05
commit 432ebadc17
2 changed files with 162 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files
homeassistant/components/usage_prediction/common_control.py
+81 -91
View File
@@ -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
+81
View File
@@ -1,5 +1,6 @@
"""Test the common control usage prediction."""
from typing import Any
from unittest.mock import patch
import uuid
@@ -409,3 +410,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()