This commit is contained in:
J. Nick Koston
2023-04-10 17:10:43 -10:00
parent 15d7bc3e0c
commit b780f10de8
3 changed files with 42 additions and 17 deletions

View File

@@ -651,7 +651,8 @@ def _sorted_states_to_dict(
""" """
field_map = _FIELD_MAP field_map = _FIELD_MAP
state_class: Callable[ state_class: Callable[
[Row, dict[str, dict[str, Any]], float | None], State | dict[str, Any] [Row, dict[str, dict[str, Any]], float | None, str, str, float | None],
State | dict[str, Any],
] ]
if compressed_state_format: if compressed_state_format:
state_class = row_to_compressed_state state_class = row_to_compressed_state
@@ -699,7 +700,14 @@ def _sorted_states_to_dict(
or split_entity_id(entity_id)[0] in NEED_ATTRIBUTE_DOMAINS or split_entity_id(entity_id)[0] in NEED_ATTRIBUTE_DOMAINS
): ):
ent_results.extend( ent_results.extend(
state_class(db_state, attr_cache, start_time_ts, entity_id=entity_id) # type: ignore[call-arg] state_class(
db_state,
attr_cache,
start_time_ts,
entity_id,
db_state[state_idx],
db_state[last_updated_ts_idx],
)
for db_state in group for db_state in group
) )
continue continue
@@ -713,7 +721,14 @@ def _sorted_states_to_dict(
continue continue
prev_state = first_state[state_idx] prev_state = first_state[state_idx]
ent_results.append( ent_results.append(
state_class(first_state, attr_cache, start_time_ts, entity_id=entity_id) # type: ignore[call-arg] state_class(
first_state,
attr_cache,
start_time_ts,
entity_id,
prev_state, # type: ignore[arg-type]
first_state[last_updated_ts_idx],
)
) )
# #

View File

@@ -52,16 +52,18 @@ class LazyState(State):
row: Row, row: Row,
attr_cache: dict[str, dict[str, Any]], attr_cache: dict[str, dict[str, Any]],
start_time_ts: float | None, start_time_ts: float | None,
entity_id: str | None = None, entity_id: str,
state: str,
last_updated_ts: float | None,
) -> None: ) -> None:
"""Init the lazy state.""" """Init the lazy state."""
self._row = row self._row = row
self.entity_id = entity_id or self._row.entity_id self.entity_id = entity_id
self.state = self._row.state or "" self.state = state or ""
self._attributes: dict[str, Any] | None = None self._attributes: dict[str, Any] | None = None
self._last_updated_ts: float | None = self._row.last_updated_ts or start_time_ts self._last_updated_ts: float | None = last_updated_ts or start_time_ts
self._last_changed_ts: float | None = ( self._last_changed_ts: float | None = (
getattr(self._row, "last_changed_ts", None) or self._last_updated_ts getattr(self._row, "last_changed_ts", None) or last_updated_ts
) )
self._context: Context | None = None self._context: Context | None = None
self.attr_cache = attr_cache self.attr_cache = attr_cache
@@ -137,14 +139,16 @@ def row_to_compressed_state(
row: Row, row: Row,
attr_cache: dict[str, dict[str, Any]], attr_cache: dict[str, dict[str, Any]],
start_time_ts: float | None, start_time_ts: float | None,
entity_id: str | None = None, entity_id: str,
state: str,
last_updated_ts: float | None,
) -> dict[str, Any]: ) -> dict[str, Any]:
"""Convert a database row to a compressed state schema 31 and later.""" """Convert a database row to a compressed state schema 31 and later."""
comp_state = { comp_state: dict[str, Any] = {
COMPRESSED_STATE_STATE: row.state, COMPRESSED_STATE_STATE: state,
COMPRESSED_STATE_ATTRIBUTES: decode_attributes_from_row(row, attr_cache), COMPRESSED_STATE_ATTRIBUTES: decode_attributes_from_row(row, attr_cache),
} }
row_last_updated_ts: float = row.last_updated_ts or start_time_ts # type: ignore[assignment] row_last_updated_ts: float = last_updated_ts or start_time_ts # type: ignore[assignment]
comp_state[COMPRESSED_STATE_LAST_UPDATED] = row_last_updated_ts comp_state[COMPRESSED_STATE_LAST_UPDATED] = row_last_updated_ts
if ( if (
(row_last_changed_ts := getattr(row, "last_changed_ts", None)) (row_last_changed_ts := getattr(row, "last_changed_ts", None))

View File

@@ -55,14 +55,20 @@ async def _async_get_states(
"""Get states from the database.""" """Get states from the database."""
def _get_states_with_session(): def _get_states_with_session():
if get_instance(hass).schema_version < 31:
klass = LegacyLazyStatePreSchema31
else:
klass = LazyState
with session_scope(hass=hass, read_only=True) as session: with session_scope(hass=hass, read_only=True) as session:
attr_cache = {} attr_cache = {}
pre_31_schema = get_instance(hass).schema_version < 31
return [ return [
klass(row, attr_cache, None) LegacyLazyStatePreSchema31(row, attr_cache, None)
if pre_31_schema
else LazyState(
row,
attr_cache,
None,
row.entity_id,
row.state,
getattr(row, "last_updated_ts", None),
)
for row in legacy._get_rows_with_session( for row in legacy._get_rows_with_session(
hass, hass,
session, session,