From 389fc515a13d5f5288b2ba746a295067beab9a69 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@koston.org>
Date: Thu, 2 Feb 2023 11:49:33 -0600
Subject: [PATCH] Prevent overly large attributes from being stored in the
 database (#87105)

---
 .../components/recorder/db_schema.py          | 20 ++++++++++----
 tests/components/recorder/test_init.py        | 27 +++++++++++++++++++
 2 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/homeassistant/components/recorder/db_schema.py b/homeassistant/components/recorder/db_schema.py
index d614e3a28a6..dfbb83c495a 100644
--- a/homeassistant/components/recorder/db_schema.py
+++ b/homeassistant/components/recorder/db_schema.py
@@ -72,6 +72,9 @@ TABLE_STATISTICS_META = "statistics_meta"
 TABLE_STATISTICS_RUNS = "statistics_runs"
 TABLE_STATISTICS_SHORT_TERM = "statistics_short_term"
 
+MAX_STATE_ATTRS_BYTES = 16384
+PSQL_DIALECT = SupportedDialect.POSTGRESQL
+
 ALL_TABLES = [
     TABLE_STATES,
     TABLE_STATE_ATTRIBUTES,
@@ -415,13 +418,20 @@ class StateAttributes(Base):  # type: ignore[misc,valid-type]
         exclude_attrs = (
             exclude_attrs_by_domain.get(domain, set()) | ALL_DOMAIN_EXCLUDE_ATTRS
         )
-        if dialect == SupportedDialect.POSTGRESQL:
-            return json_bytes_strip_null(
-                {k: v for k, v in state.attributes.items() if k not in exclude_attrs}
-            )
-        return json_bytes(
+        encoder = json_bytes_strip_null if dialect == PSQL_DIALECT else json_bytes
+        bytes_result = encoder(
             {k: v for k, v in state.attributes.items() if k not in exclude_attrs}
         )
+        if len(bytes_result) > MAX_STATE_ATTRS_BYTES:
+            _LOGGER.warning(
+                "State attributes for %s exceed maximum size of %s bytes. "
+                "This can cause database performance issues; Attributes "
+                "will not be stored",
+                state.entity_id,
+                MAX_STATE_ATTRS_BYTES,
+            )
+            return b"{}"
+        return bytes_result
 
     @staticmethod
     def hash_shared_attrs_bytes(shared_attrs_bytes: bytes) -> int:
diff --git a/tests/components/recorder/test_init.py b/tests/components/recorder/test_init.py
index b8300a18480..cf4a5935933 100644
--- a/tests/components/recorder/test_init.py
+++ b/tests/components/recorder/test_init.py
@@ -702,6 +702,33 @@ def test_saving_state_and_removing_entity(hass, hass_recorder):
         assert states[2].state is None
 
 
+def test_saving_state_with_oversized_attributes(hass_recorder, caplog):
+    """Test saving states is limited to 16KiB of JSON encoded attributes."""
+    hass = hass_recorder()
+    massive_dict = {"a": "b" * 16384}
+    attributes = {"test_attr": 5, "test_attr_10": "nice"}
+    hass.states.set("switch.sane", "on", attributes)
+    hass.states.set("switch.too_big", "on", massive_dict)
+    wait_recording_done(hass)
+    states = []
+
+    with session_scope(hass=hass) as session:
+        for state, state_attributes in session.query(States, StateAttributes).outerjoin(
+            StateAttributes, States.attributes_id == StateAttributes.attributes_id
+        ):
+            native_state = state.to_native()
+            native_state.attributes = state_attributes.to_native()
+            states.append(native_state)
+
+    assert "switch.too_big" in caplog.text
+
+    assert len(states) == 2
+    assert _state_with_context(hass, "switch.sane").as_dict() == states[0].as_dict()
+    assert states[1].state == "on"
+    assert states[1].entity_id == "switch.too_big"
+    assert states[1].attributes == {}
+
+
 def test_recorder_setup_failure(hass):
     """Test some exceptions."""
     recorder_helper.async_initialize_recorder(hass)