From 99d282a9e5bd577776afe79d5f4af718519fdc2a Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Mon, 30 Sep 2024 02:01:41 -0500
Subject: [PATCH] Fix removing nulls when encoding events for PostgreSQL
 (#127053)

---
 .../components/recorder/db_schema.py          |  5 ++--
 tests/components/recorder/test_models.py      | 28 +++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/homeassistant/components/recorder/db_schema.py b/homeassistant/components/recorder/db_schema.py
index 6ba9d971f2c331..7e8343321c3c5d 100644
--- a/homeassistant/components/recorder/db_schema.py
+++ b/homeassistant/components/recorder/db_schema.py
@@ -375,9 +375,8 @@ def shared_data_bytes_from_event(
         event: Event, dialect: SupportedDialect | None
     ) -> bytes:
         """Create shared_data from an event."""
-        if dialect == SupportedDialect.POSTGRESQL:
-            bytes_result = json_bytes_strip_null(event.data)
-        bytes_result = json_bytes(event.data)
+        encoder = json_bytes_strip_null if dialect == PSQL_DIALECT else json_bytes
+        bytes_result = encoder(event.data)
         if len(bytes_result) > MAX_EVENT_DATA_BYTES:
             _LOGGER.warning(
                 "Event data for %s exceed maximum size of %s bytes. "
diff --git a/tests/components/recorder/test_models.py b/tests/components/recorder/test_models.py
index c8ab64c7d89246..9078b2e861c157 100644
--- a/tests/components/recorder/test_models.py
+++ b/tests/components/recorder/test_models.py
@@ -21,6 +21,7 @@
 import homeassistant.core as ha
 from homeassistant.exceptions import InvalidEntityFormatError
 from homeassistant.util import dt as dt_util
+from homeassistant.util.json import json_loads
 
 
 def test_from_event_to_db_event() -> None:
@@ -41,6 +42,18 @@ def test_from_event_to_db_event() -> None:
     assert event.as_dict() == db_event.to_native().as_dict()
 
 
+def test_from_event_to_db_event_with_null() -> None:
+    """Test converting event to EventData with a null with PostgreSQL."""
+    event = ha.Event(
+        "test_event",
+        {"some_data": "withnull\0terminator"},
+    )
+    dialect = SupportedDialect.POSTGRESQL
+    event_data = EventData.shared_data_bytes_from_event(event, dialect)
+    decoded = json_loads(event_data)
+    assert decoded["some_data"] == "withnull"
+
+
 def test_from_event_to_db_state() -> None:
     """Test converting event to db state."""
     state = ha.State(
@@ -78,6 +91,21 @@ def test_from_event_to_db_state_attributes() -> None:
     assert db_attrs.to_native() == attrs
 
 
+def test_from_event_to_db_state_attributes_with_null() -> None:
+    """Test converting a state to StateAttributes with a null with PostgreSQL."""
+    attrs = {"this_attr": "withnull\0terminator"}
+    state = ha.State("sensor.temperature", "18", attrs)
+    event = ha.Event(
+        EVENT_STATE_CHANGED,
+        {"entity_id": "sensor.temperature", "old_state": None, "new_state": state},
+        context=state.context,
+    )
+    dialect = SupportedDialect.POSTGRESQL
+    shared_attrs = StateAttributes.shared_attrs_bytes_from_event(event, dialect)
+    decoded = json_loads(shared_attrs)
+    assert decoded["this_attr"] == "withnull"
+
+
 def test_repr() -> None:
     """Test converting event to db state repr."""
     attrs = {"this_attr": True}