From 980b8a91e62c449fab558318573fa756818875a6 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Sat, 14 Dec 2024 14:21:19 +0100 Subject: [PATCH] Revert "Simplify recorder RecorderRunsManager" (#133201) Revert "Simplify recorder RecorderRunsManager (#131785)" This reverts commit cf0ee635077114961f6e508be56ce7620c718c18. --- .../recorder/table_managers/recorder_runs.py | 73 ++++++++++++++++--- .../table_managers/test_recorder_runs.py | 32 ++++++-- 2 files changed, 90 insertions(+), 15 deletions(-) diff --git a/homeassistant/components/recorder/table_managers/recorder_runs.py b/homeassistant/components/recorder/table_managers/recorder_runs.py index 4ca0aa18b88..b0b9818118b 100644 --- a/homeassistant/components/recorder/table_managers/recorder_runs.py +++ b/homeassistant/components/recorder/table_managers/recorder_runs.py @@ -2,6 +2,8 @@ from __future__ import annotations +import bisect +from dataclasses import dataclass from datetime import datetime from sqlalchemy.orm.session import Session @@ -9,6 +11,34 @@ 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: @@ -18,7 +48,7 @@ class RecorderRunsManager: """Track recorder run history.""" self._recording_start = dt_util.utcnow() self._current_run_info: RecorderRuns | None = None - self._first_run: RecorderRuns | None = None + self._run_history = _RecorderRunsHistory([], {}) @property def recording_start(self) -> datetime: @@ -28,7 +58,9 @@ class RecorderRunsManager: @property def first(self) -> RecorderRuns: """Get the first run.""" - return self._first_run or self.current + if runs_by_timestamp := self._run_history.runs_by_timestamp: + return next(iter(runs_by_timestamp.values())) + return self.current @property def current(self) -> RecorderRuns: @@ -46,6 +78,15 @@ 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. @@ -81,17 +122,31 @@ class RecorderRunsManager: Must run in the recorder thread. """ - if ( - run := session.query(RecorderRuns) - .order_by(RecorderRuns.start.asc()) - .first() - ): + run_timestamps: list[int] = [] + runs_by_timestamp: dict[int, RecorderRuns] = {} + + for run in session.query(RecorderRuns).order_by(RecorderRuns.start.asc()).all(): session.expunge(run) - self._first_run = 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) def clear(self) -> None: """Clear the current run after ending it. Must run in the recorder thread. """ - self._current_run_info = None + if self._current_run_info: + self._current_run_info = None diff --git a/tests/components/recorder/table_managers/test_recorder_runs.py b/tests/components/recorder/table_managers/test_recorder_runs.py index e79def01bad..41f3a8fef4d 100644 --- a/tests/components/recorder/table_managers/test_recorder_runs.py +++ b/tests/components/recorder/table_managers/test_recorder_runs.py @@ -21,11 +21,6 @@ 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)) @@ -34,7 +29,32 @@ 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.first.start) == three_days_ago + 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 )