Use more efficient chunked_or_all for recorder table managers (#118646)

This commit is contained in:
J. Nick Koston 2024-06-02 13:04:53 -05:00 committed by GitHub
parent 54a1a4ab41
commit 51ed4f89ec
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 12 additions and 12 deletions

View File

@ -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
):

View File

@ -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
):

View File

@ -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
):

View File

@ -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)
):