Simplify recorder RecorderRunsManager (#131785)

This commit is contained in:
Erik Montnemery 2024-12-01 18:26:29 +01:00 committed by GitHub
parent c54eed3607
commit cf0ee63507
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 15 additions and 90 deletions

View File

@ -2,8 +2,6 @@
from __future__ import annotations
import bisect
from dataclasses import dataclass
from datetime import datetime
from sqlalchemy.orm.session import Session
@ -11,34 +9,6 @@ from sqlalchemy.orm.session import Session
import homeassistant.util.dt as dt_util
from ..db_schema import RecorderRuns
from ..models import process_timestamp
def _find_recorder_run_for_start_time(
run_history: _RecorderRunsHistory, start: datetime
) -> RecorderRuns | None:
"""Find the recorder run for a start time in _RecorderRunsHistory."""
run_timestamps = run_history.run_timestamps
runs_by_timestamp = run_history.runs_by_timestamp
# bisect_left tells us were we would insert
# a value in the list of runs after the start timestamp.
#
# The run before that (idx-1) is when the run started
#
# If idx is 0, history never ran before the start timestamp
#
if idx := bisect.bisect_left(run_timestamps, start.timestamp()):
return runs_by_timestamp[run_timestamps[idx - 1]]
return None
@dataclass(frozen=True)
class _RecorderRunsHistory:
"""Bisectable history of RecorderRuns."""
run_timestamps: list[int]
runs_by_timestamp: dict[int, RecorderRuns]
class RecorderRunsManager:
@ -48,7 +18,7 @@ class RecorderRunsManager:
"""Track recorder run history."""
self._recording_start = dt_util.utcnow()
self._current_run_info: RecorderRuns | None = None
self._run_history = _RecorderRunsHistory([], {})
self._first_run: RecorderRuns | None = None
@property
def recording_start(self) -> datetime:
@ -58,9 +28,7 @@ class RecorderRunsManager:
@property
def first(self) -> RecorderRuns:
"""Get the first run."""
if runs_by_timestamp := self._run_history.runs_by_timestamp:
return next(iter(runs_by_timestamp.values()))
return self.current
return self._first_run or self.current
@property
def current(self) -> RecorderRuns:
@ -78,15 +46,6 @@ class RecorderRunsManager:
"""Return if a run is active."""
return self._current_run_info is not None
def get(self, start: datetime) -> RecorderRuns | None:
"""Return the recorder run that started before or at start.
If the first run started after the start, return None
"""
if start >= self.recording_start:
return self.current
return _find_recorder_run_for_start_time(self._run_history, start)
def start(self, session: Session) -> None:
"""Start a new run.
@ -122,31 +81,17 @@ class RecorderRunsManager:
Must run in the recorder thread.
"""
run_timestamps: list[int] = []
runs_by_timestamp: dict[int, RecorderRuns] = {}
for run in session.query(RecorderRuns).order_by(RecorderRuns.start.asc()).all():
if (
run := session.query(RecorderRuns)
.order_by(RecorderRuns.start.asc())
.first()
):
session.expunge(run)
if run_dt := process_timestamp(run.start):
# Not sure if this is correct or runs_by_timestamp annotation should be changed
timestamp = int(run_dt.timestamp())
run_timestamps.append(timestamp)
runs_by_timestamp[timestamp] = run
#
# self._run_history is accessed in get()
# which is allowed to be called from any thread
#
# We use a dataclass to ensure that when we update
# run_timestamps and runs_by_timestamp
# are never out of sync with each other.
#
self._run_history = _RecorderRunsHistory(run_timestamps, runs_by_timestamp)
self._first_run = run
def clear(self) -> None:
"""Clear the current run after ending it.
Must run in the recorder thread.
"""
if self._current_run_info:
self._current_run_info = None
self._current_run_info = None

View File

@ -21,6 +21,11 @@ async def test_run_history(recorder_mock: Recorder, hass: HomeAssistant) -> None
two_days_ago = now - timedelta(days=2)
one_day_ago = now - timedelta(days=1)
# Test that the first run falls back to the current run
assert process_timestamp(
instance.recorder_runs_manager.first.start
) == process_timestamp(instance.recorder_runs_manager.current.start)
with instance.get_session() as session:
session.add(RecorderRuns(start=three_days_ago, created=three_days_ago))
session.add(RecorderRuns(start=two_days_ago, created=two_days_ago))
@ -29,32 +34,7 @@ async def test_run_history(recorder_mock: Recorder, hass: HomeAssistant) -> None
instance.recorder_runs_manager.load_from_db(session)
assert (
process_timestamp(
instance.recorder_runs_manager.get(
three_days_ago + timedelta(microseconds=1)
).start
)
== three_days_ago
)
assert (
process_timestamp(
instance.recorder_runs_manager.get(
two_days_ago + timedelta(microseconds=1)
).start
)
== two_days_ago
)
assert (
process_timestamp(
instance.recorder_runs_manager.get(
one_day_ago + timedelta(microseconds=1)
).start
)
== one_day_ago
)
assert (
process_timestamp(instance.recorder_runs_manager.get(now).start)
== instance.recorder_runs_manager.recording_start
process_timestamp(instance.recorder_runs_manager.first.start) == three_days_ago
)