From 753ab08b5ec06c7e3639be2eda366e25159fd553 Mon Sep 17 00:00:00 2001 From: G Johansson Date: Sat, 22 Jun 2024 19:30:28 +0200 Subject: [PATCH] Add capability to exclude all attributes from recording (#119725) --- .../components/recorder/db_schema.py | 24 ++++++- tests/components/recorder/db_schema_42.py | 24 ++++++- tests/components/recorder/test_init.py | 65 +++++++++++++++++++ 3 files changed, 111 insertions(+), 2 deletions(-) diff --git a/homeassistant/components/recorder/db_schema.py b/homeassistant/components/recorder/db_schema.py index 186b873047b..ce463067824 100644 --- a/homeassistant/components/recorder/db_schema.py +++ b/homeassistant/components/recorder/db_schema.py @@ -35,7 +35,12 @@ from sqlalchemy.ext.compiler import compiles from sqlalchemy.orm import DeclarativeBase, Mapped, aliased, mapped_column, relationship from sqlalchemy.types import TypeDecorator +from homeassistant.components.sensor import ATTR_STATE_CLASS from homeassistant.const import ( + ATTR_DEVICE_CLASS, + ATTR_FRIENDLY_NAME, + ATTR_UNIT_OF_MEASUREMENT, + MATCH_ALL, MAX_LENGTH_EVENT_EVENT_TYPE, MAX_LENGTH_STATE_ENTITY_ID, MAX_LENGTH_STATE_STATE, @@ -584,10 +589,27 @@ class StateAttributes(Base): if (state := event.data["new_state"]) is None: return b"{}" if state_info := state.state_info: + unrecorded_attributes = state_info["unrecorded_attributes"] exclude_attrs = { *ALL_DOMAIN_EXCLUDE_ATTRS, - *state_info["unrecorded_attributes"], + *unrecorded_attributes, } + if MATCH_ALL in unrecorded_attributes: + # Don't exclude device class, state class, unit of measurement + # or friendly name when using the MATCH_ALL exclude constant + _exclude_attributes = { + k: v + for k, v in state.attributes.items() + if k + not in ( + ATTR_DEVICE_CLASS, + ATTR_STATE_CLASS, + ATTR_UNIT_OF_MEASUREMENT, + ATTR_FRIENDLY_NAME, + ) + } + exclude_attrs.update(_exclude_attributes) + else: exclude_attrs = ALL_DOMAIN_EXCLUDE_ATTRS encoder = json_bytes_strip_null if dialect == PSQL_DIALECT else json_bytes diff --git a/tests/components/recorder/db_schema_42.py b/tests/components/recorder/db_schema_42.py index b8e49aef592..c0dfc70571d 100644 --- a/tests/components/recorder/db_schema_42.py +++ b/tests/components/recorder/db_schema_42.py @@ -54,7 +54,12 @@ from homeassistant.components.recorder.models import ( ulid_to_bytes_or_none, uuid_hex_to_bytes_or_none, ) +from homeassistant.components.sensor import ATTR_STATE_CLASS from homeassistant.const import ( + ATTR_DEVICE_CLASS, + ATTR_FRIENDLY_NAME, + ATTR_UNIT_OF_MEASUREMENT, + MATCH_ALL, MAX_LENGTH_EVENT_EVENT_TYPE, MAX_LENGTH_STATE_ENTITY_ID, MAX_LENGTH_STATE_STATE, @@ -577,10 +582,27 @@ class StateAttributes(Base): if state is None: return b"{}" if state_info := state.state_info: + unrecorded_attributes = state_info["unrecorded_attributes"] exclude_attrs = { *ALL_DOMAIN_EXCLUDE_ATTRS, - *state_info["unrecorded_attributes"], + *unrecorded_attributes, } + if MATCH_ALL in unrecorded_attributes: + # Don't exclude device class, state class, unit of measurement + # or friendly name when using the MATCH_ALL exclude constant + _exclude_attributes = { + k: v + for k, v in state.attributes.items() + if k + not in ( + ATTR_DEVICE_CLASS, + ATTR_STATE_CLASS, + ATTR_UNIT_OF_MEASUREMENT, + ATTR_FRIENDLY_NAME, + ) + } + exclude_attrs.update(_exclude_attributes) + else: exclude_attrs = ALL_DOMAIN_EXCLUDE_ATTRS encoder = json_bytes_strip_null if dialect == PSQL_DIALECT else json_bytes diff --git a/tests/components/recorder/test_init.py b/tests/components/recorder/test_init.py index 300d338fcb3..52947ce0c19 100644 --- a/tests/components/recorder/test_init.py +++ b/tests/components/recorder/test_init.py @@ -2420,6 +2420,71 @@ async def test_excluding_attributes_by_integration( assert state.as_dict() == expected.as_dict() +async def test_excluding_all_attributes_by_integration( + hass: HomeAssistant, + entity_registry: er.EntityRegistry, + setup_recorder: None, +) -> None: + """Test that an entity can exclude all attributes from being recorded using MATCH_ALL.""" + state = "restoring_from_db" + attributes = { + "test_attr": 5, + "excluded_component": 10, + "excluded_integration": 20, + "device_class": "test", + "state_class": "test", + "friendly_name": "Test entity", + "unit_of_measurement": "mm", + } + mock_platform( + hass, + "fake_integration.recorder", + Mock(exclude_attributes=lambda hass: {"excluded"}), + ) + hass.config.components.add("fake_integration") + hass.bus.async_fire(EVENT_COMPONENT_LOADED, {"component": "fake_integration"}) + await hass.async_block_till_done() + + class EntityWithExcludedAttributes(MockEntity): + _unrecorded_attributes = frozenset({MATCH_ALL}) + + entity_id = "test.fake_integration_recorder" + entity_platform = MockEntityPlatform(hass, platform_name="fake_integration") + entity = EntityWithExcludedAttributes( + entity_id=entity_id, + extra_state_attributes=attributes, + ) + await entity_platform.async_add_entities([entity]) + await hass.async_block_till_done() + + await async_wait_recording_done(hass) + + with session_scope(hass=hass, read_only=True) as session: + db_states = [] + for db_state, db_state_attributes, states_meta in ( + session.query(States, StateAttributes, StatesMeta) + .outerjoin( + StateAttributes, States.attributes_id == StateAttributes.attributes_id + ) + .outerjoin(StatesMeta, States.metadata_id == StatesMeta.metadata_id) + ): + db_state.entity_id = states_meta.entity_id + db_states.append(db_state) + state = db_state.to_native() + state.attributes = db_state_attributes.to_native() + assert len(db_states) == 1 + assert db_states[0].event_id is None + + expected = _state_with_context(hass, entity_id) + expected.attributes = { + "device_class": "test", + "state_class": "test", + "friendly_name": "Test entity", + "unit_of_measurement": "mm", + } + assert state.as_dict() == expected.as_dict() + + async def test_lru_increases_with_many_entities( small_cache_size: None, hass: HomeAssistant, setup_recorder: None ) -> None: