From 56a6244d903bba456c36220fc864ef8c6332f41a Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Wed, 19 Apr 2023 15:54:34 -1000
Subject: [PATCH] Remove legacy context lookup implementation from logbook
 (#91710)

* Remove legacy context lookup implemention from logbook

This object can now be replaced with a simple dict

* Remove legacy context lookup implemention from logbook

This object can now be replaced with a simple dict

* scope

* fix order issue
---
 homeassistant/components/logbook/processor.py | 45 ++++++-------------
 .../components/logbook/websocket_api.py       |  3 +-
 tests/components/logbook/common.py            |  2 +-
 3 files changed, 15 insertions(+), 35 deletions(-)

diff --git a/homeassistant/components/logbook/processor.py b/homeassistant/components/logbook/processor.py
index 0d3e38d5496..6df07684330 100644
--- a/homeassistant/components/logbook/processor.py
+++ b/homeassistant/components/logbook/processor.py
@@ -66,7 +66,7 @@ from .queries.common import PSEUDO_EVENT_STATE_CHANGED
 class LogbookRun:
     """A logbook run which may be a long running event stream or single request."""
 
-    context_lookup: ContextLookup
+    context_lookup: dict[bytes | None, Row | EventAsRow | None]
     external_events: dict[
         str, tuple[str, Callable[[LazyEventPartialState], dict[str, Any]]]
     ]
@@ -74,6 +74,7 @@ class LogbookRun:
     entity_name_cache: EntityNameCache
     include_entity_name: bool
     format_time: Callable[[Row | EventAsRow], Any]
+    memoize_new_contexts: bool = True
 
 
 class EventProcessor:
@@ -105,7 +106,7 @@ class EventProcessor:
             _row_time_fired_timestamp if timestamp else _row_time_fired_isoformat
         )
         self.logbook_run = LogbookRun(
-            context_lookup=ContextLookup(hass),
+            context_lookup={None: None},
             external_events=logbook_config.external_events,
             event_cache=EventCache({}),
             entity_name_cache=EntityNameCache(self.hass),
@@ -126,6 +127,7 @@ class EventProcessor:
         """
         self.logbook_run.event_cache.clear()
         self.logbook_run.context_lookup.clear()
+        self.logbook_run.memoize_new_contexts = False
 
     def get_events(
         self,
@@ -207,10 +209,14 @@ def _humanify(
     entity_name_cache = logbook_run.entity_name_cache
     include_entity_name = logbook_run.include_entity_name
     format_time = logbook_run.format_time
+    memoize_new_contexts = logbook_run.memoize_new_contexts
+    memoize_context = context_lookup.setdefault
 
     # Process rows
     for row in rows:
-        context_id_bin = context_lookup.memorize(row)
+        context_id_bin: bytes = row.context_id_bin
+        if memoize_new_contexts:
+            memoize_context(context_id_bin, row)
         if row.context_only:
             continue
         event_type = row.event_type
@@ -269,33 +275,6 @@ def _humanify(
             yield data
 
 
-class ContextLookup:
-    """A lookup class for context origins."""
-
-    def __init__(self, hass: HomeAssistant) -> None:
-        """Memorize context origin."""
-        self.hass = hass
-        self._memorize_new = True
-        self._lookup: dict[bytes | None, Row | EventAsRow | None] = {None: None}
-
-    def memorize(self, row: Row | EventAsRow) -> bytes | None:
-        """Memorize a context from the database."""
-        if self._memorize_new:
-            context_id_bin: bytes = row.context_id_bin
-            self._lookup.setdefault(context_id_bin, row)
-            return context_id_bin
-        return None
-
-    def clear(self) -> None:
-        """Clear the context origins and stop recording new ones."""
-        self._lookup.clear()
-        self._memorize_new = False
-
-    def get(self, context_id_bin: bytes) -> Row | EventAsRow | None:
-        """Get the context origin."""
-        return self._lookup.get(context_id_bin)
-
-
 class ContextAugmenter:
     """Augment data with context trace."""
 
@@ -311,8 +290,10 @@ class ContextAugmenter:
         self, context_id_bin: bytes | None, row: Row | EventAsRow
     ) -> Row | EventAsRow | None:
         """Get the context row from the id or row context."""
-        if context_id_bin:
-            return self.context_lookup.get(context_id_bin)
+        if context_id_bin is not None and (
+            context_row := self.context_lookup.get(context_id_bin)
+        ):
+            return context_row
         if (context := getattr(row, "context", None)) is not None and (
             origin_event := context.origin_event
         ) is not None:
diff --git a/homeassistant/components/logbook/websocket_api.py b/homeassistant/components/logbook/websocket_api.py
index cbcf0e413f6..c4e6b9814f4 100644
--- a/homeassistant/components/logbook/websocket_api.py
+++ b/homeassistant/components/logbook/websocket_api.py
@@ -221,8 +221,6 @@ async def _async_events_consumer(
     event_processor: EventProcessor,
 ) -> None:
     """Stream events from the queue."""
-    event_processor.switch_to_live()
-
     while True:
         events: list[Event] = [await stream_queue.get()]
         # If the event is older than the last db
@@ -430,6 +428,7 @@ async def ws_event_stream(
         event_processor,
         partial=False,
     )
+    event_processor.switch_to_live()
 
 
 def _ws_formatted_get_events(
diff --git a/tests/components/logbook/common.py b/tests/components/logbook/common.py
index b8d00681112..75e4cd40392 100644
--- a/tests/components/logbook/common.py
+++ b/tests/components/logbook/common.py
@@ -64,7 +64,7 @@ def mock_humanify(hass_, rows):
     entity_name_cache = processor.EntityNameCache(hass_)
     ent_reg = er.async_get(hass_)
     event_cache = processor.EventCache({})
-    context_lookup = processor.ContextLookup(hass_)
+    context_lookup = {}
     logbook_config = hass_.data.get(logbook.DOMAIN, LogbookConfig({}, None, None))
     external_events = logbook_config.external_events
     logbook_run = processor.LogbookRun(