From 51ed4f89ec0c0bfc44aa9910006c342b29e14016 Mon Sep 17 00:00:00 2001 From: "J. Nick Koston" Date: Sun, 2 Jun 2024 13:04:53 -0500 Subject: [PATCH] Use more efficient chunked_or_all for recorder table managers (#118646) --- .../components/recorder/table_managers/event_data.py | 8 ++++---- .../components/recorder/table_managers/event_types.py | 4 ++-- .../recorder/table_managers/state_attributes.py | 8 ++++---- .../components/recorder/table_managers/states_meta.py | 4 ++-- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/homeassistant/components/recorder/table_managers/event_data.py b/homeassistant/components/recorder/table_managers/event_data.py index 28f02127d42..1d2fa580b3c 100644 --- a/homeassistant/components/recorder/table_managers/event_data.py +++ b/homeassistant/components/recorder/table_managers/event_data.py @@ -2,14 +2,14 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Collection, Iterable import logging from typing import TYPE_CHECKING, cast from sqlalchemy.orm.session import Session from homeassistant.core import Event -from homeassistant.util.collection import chunked +from homeassistant.util.collection import chunked_or_all from homeassistant.util.json import JSON_ENCODE_EXCEPTIONS from ..db_schema import EventData @@ -87,7 +87,7 @@ class EventDataManager(BaseLRUTableManager[EventData]): return results | self._load_from_hashes(missing_hashes, session) def _load_from_hashes( - self, hashes: Iterable[int], session: Session + self, hashes: Collection[int], session: Session ) -> dict[str, int | None]: """Load the shared_datas to data_ids mapping into memory from a list of hashes. @@ -96,7 +96,7 @@ class EventDataManager(BaseLRUTableManager[EventData]): """ results: dict[str, int | None] = {} with session.no_autoflush: - for hashs_chunk in chunked(hashes, self.recorder.max_bind_vars): + for hashs_chunk in chunked_or_all(hashes, self.recorder.max_bind_vars): for data_id, shared_data in execute_stmt_lambda_element( session, get_shared_event_datas(hashs_chunk), orm_rows=False ): diff --git a/homeassistant/components/recorder/table_managers/event_types.py b/homeassistant/components/recorder/table_managers/event_types.py index 29eaf2450ad..266c970fe1f 100644 --- a/homeassistant/components/recorder/table_managers/event_types.py +++ b/homeassistant/components/recorder/table_managers/event_types.py @@ -9,7 +9,7 @@ from lru import LRU from sqlalchemy.orm.session import Session from homeassistant.core import Event -from homeassistant.util.collection import chunked +from homeassistant.util.collection import chunked_or_all from homeassistant.util.event_type import EventType from ..db_schema import EventTypes @@ -88,7 +88,7 @@ class EventTypeManager(BaseLRUTableManager[EventTypes]): return results with session.no_autoflush: - for missing_chunk in chunked(missing, self.recorder.max_bind_vars): + for missing_chunk in chunked_or_all(missing, self.recorder.max_bind_vars): for event_type_id, event_type in execute_stmt_lambda_element( session, find_event_type_ids(missing_chunk), orm_rows=False ): diff --git a/homeassistant/components/recorder/table_managers/state_attributes.py b/homeassistant/components/recorder/table_managers/state_attributes.py index 4a705858d44..5ed67b0504f 100644 --- a/homeassistant/components/recorder/table_managers/state_attributes.py +++ b/homeassistant/components/recorder/table_managers/state_attributes.py @@ -2,14 +2,14 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Collection, Iterable import logging from typing import TYPE_CHECKING, cast from sqlalchemy.orm.session import Session from homeassistant.core import Event, EventStateChangedData -from homeassistant.util.collection import chunked +from homeassistant.util.collection import chunked_or_all from homeassistant.util.json import JSON_ENCODE_EXCEPTIONS from ..db_schema import StateAttributes @@ -98,7 +98,7 @@ class StateAttributesManager(BaseLRUTableManager[StateAttributes]): return results | self._load_from_hashes(missing_hashes, session) def _load_from_hashes( - self, hashes: Iterable[int], session: Session + self, hashes: Collection[int], session: Session ) -> dict[str, int | None]: """Load the shared_attrs to attributes_ids mapping into memory from a list of hashes. @@ -107,7 +107,7 @@ class StateAttributesManager(BaseLRUTableManager[StateAttributes]): """ results: dict[str, int | None] = {} with session.no_autoflush: - for hashs_chunk in chunked(hashes, self.recorder.max_bind_vars): + for hashs_chunk in chunked_or_all(hashes, self.recorder.max_bind_vars): for attributes_id, shared_attrs in execute_stmt_lambda_element( session, get_shared_attributes(hashs_chunk), orm_rows=False ): diff --git a/homeassistant/components/recorder/table_managers/states_meta.py b/homeassistant/components/recorder/table_managers/states_meta.py index 5e5f2f06796..0ea2c7415b9 100644 --- a/homeassistant/components/recorder/table_managers/states_meta.py +++ b/homeassistant/components/recorder/table_managers/states_meta.py @@ -8,7 +8,7 @@ from typing import TYPE_CHECKING, cast from sqlalchemy.orm.session import Session from homeassistant.core import Event, EventStateChangedData -from homeassistant.util.collection import chunked +from homeassistant.util.collection import chunked_or_all from ..db_schema import StatesMeta from ..queries import find_all_states_metadata_ids, find_states_metadata_ids @@ -107,7 +107,7 @@ class StatesMetaManager(BaseLRUTableManager[StatesMeta]): update_cache = from_recorder or not self._did_first_load with session.no_autoflush: - for missing_chunk in chunked(missing, self.recorder.max_bind_vars): + for missing_chunk in chunked_or_all(missing, self.recorder.max_bind_vars): for metadata_id, entity_id in execute_stmt_lambda_element( session, find_states_metadata_ids(missing_chunk) ):