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
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:
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
):
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
)
continue
@@ -713,7 +721,14 @@ def _sorted_states_to_dict(
continue
prev_state = first_state[state_idx]
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,
attr_cache: dict[str, dict[str, Any]],
start_time_ts: float | None,
entity_id: str | None = None,
entity_id: str,
state: str,
last_updated_ts: float | None,
) -> None:
"""Init the lazy state."""
self._row = row
self.entity_id = entity_id or self._row.entity_id
self.state = self._row.state or ""
self.entity_id = entity_id
self.state = state or ""
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 = (
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.attr_cache = attr_cache
@@ -137,14 +139,16 @@ def row_to_compressed_state(
row: Row,
attr_cache: dict[str, dict[str, Any]],
start_time_ts: float | None,
entity_id: str | None = None,
entity_id: str,
state: str,
last_updated_ts: float | None,
) -> dict[str, Any]:
"""Convert a database row to a compressed state schema 31 and later."""
comp_state = {
COMPRESSED_STATE_STATE: row.state,
comp_state: dict[str, Any] = {
COMPRESSED_STATE_STATE: state,
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
if (
(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."""
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:
attr_cache = {}
pre_31_schema = get_instance(hass).schema_version < 31
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(
hass,
session,