From fbcc5318c575d4cbabb1e96472636503e0fc8a75 Mon Sep 17 00:00:00 2001 From: Erik Montnemery Date: Wed, 20 Sep 2023 18:09:12 +0200 Subject: [PATCH] Move attributes to be excluded from recording to entity classes (#100239) Co-authored-by: J. Nick Koston --- .../components/automation/__init__.py | 3 ++ .../components/automation/recorder.py | 12 ----- .../components/recorder/db_schema.py | 2 + homeassistant/core.py | 5 ++ homeassistant/helpers/entity.py | 46 +++++++++++++++++-- 5 files changed, 52 insertions(+), 16 deletions(-) delete mode 100644 homeassistant/components/automation/recorder.py diff --git a/homeassistant/components/automation/__init__.py b/homeassistant/components/automation/__init__.py index f4db7831235..fd6a70cce46 100644 --- a/homeassistant/components/automation/__init__.py +++ b/homeassistant/components/automation/__init__.py @@ -314,6 +314,9 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool: class BaseAutomationEntity(ToggleEntity, ABC): """Base class for automation entities.""" + _entity_component_unrecorded_attributes = frozenset( + (ATTR_LAST_TRIGGERED, ATTR_MODE, ATTR_CUR, ATTR_MAX, CONF_ID) + ) raw_config: ConfigType | None @property diff --git a/homeassistant/components/automation/recorder.py b/homeassistant/components/automation/recorder.py deleted file mode 100644 index 3083d271d1f..00000000000 --- a/homeassistant/components/automation/recorder.py +++ /dev/null @@ -1,12 +0,0 @@ -"""Integration platform for recorder.""" -from __future__ import annotations - -from homeassistant.core import HomeAssistant, callback - -from . import ATTR_CUR, ATTR_LAST_TRIGGERED, ATTR_MAX, ATTR_MODE, CONF_ID - - -@callback -def exclude_attributes(hass: HomeAssistant) -> set[str]: - """Exclude extra attributes from being recorded in the database.""" - return {ATTR_LAST_TRIGGERED, ATTR_MODE, ATTR_CUR, ATTR_MAX, CONF_ID} diff --git a/homeassistant/components/recorder/db_schema.py b/homeassistant/components/recorder/db_schema.py index e25c6d6dd5f..e992a683cb1 100644 --- a/homeassistant/components/recorder/db_schema.py +++ b/homeassistant/components/recorder/db_schema.py @@ -576,6 +576,8 @@ class StateAttributes(Base): integration_attrs := exclude_attrs_by_domain.get(entity_info["domain"]) ): exclude_attrs |= integration_attrs + if state_info := state.state_info: + exclude_attrs |= state_info["unrecorded_attributes"] 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} diff --git a/homeassistant/core.py b/homeassistant/core.py index a43fa1997c6..a50d43c1344 100644 --- a/homeassistant/core.py +++ b/homeassistant/core.py @@ -95,6 +95,7 @@ if TYPE_CHECKING: from .auth import AuthManager from .components.http import ApiConfig, HomeAssistantHTTP from .config_entries import ConfigEntries + from .helpers.entity import StateInfo STAGE_1_SHUTDOWN_TIMEOUT = 100 @@ -1249,6 +1250,7 @@ class State: last_updated: datetime.datetime | None = None, context: Context | None = None, validate_entity_id: bool | None = True, + state_info: StateInfo | None = None, ) -> None: """Initialize a new state.""" state = str(state) @@ -1267,6 +1269,7 @@ class State: self.last_updated = last_updated or dt_util.utcnow() self.last_changed = last_changed or self.last_updated self.context = context or Context() + self.state_info = state_info self.domain, self.object_id = split_entity_id(self.entity_id) self._as_dict: ReadOnlyDict[str, Collection[Any]] | None = None @@ -1637,6 +1640,7 @@ class StateMachine: attributes: Mapping[str, Any] | None = None, force_update: bool = False, context: Context | None = None, + state_info: StateInfo | None = None, ) -> None: """Set the state of an entity, add entity if it does not exist. @@ -1688,6 +1692,7 @@ class StateMachine: now, context, old_state is None, + state_info, ) if old_state is not None: old_state.expire() diff --git a/homeassistant/helpers/entity.py b/homeassistant/helpers/entity.py index 5ed16408388..9b16b0c24fd 100644 --- a/homeassistant/helpers/entity.py +++ b/homeassistant/helpers/entity.py @@ -201,6 +201,12 @@ class EntityInfo(TypedDict): config_entry: NotRequired[str] +class StateInfo(TypedDict): + """State info.""" + + unrecorded_attributes: frozenset[str] + + class EntityPlatformState(Enum): """The platform state of an entity.""" @@ -297,6 +303,22 @@ class Entity(ABC): # If entity is added to an entity platform _platform_state = EntityPlatformState.NOT_ADDED + # Attributes to exclude from recording, only set by base components, e.g. light + _entity_component_unrecorded_attributes: frozenset[str] = frozenset() + # Additional integration specific attributes to exclude from recording, set by + # platforms, e.g. a derived class in hue.light + _unrecorded_attributes: frozenset[str] = frozenset() + # Union of _entity_component_unrecorded_attributes and _unrecorded_attributes, + # set automatically by __init_subclass__ + __combined_unrecorded_attributes: frozenset[str] = ( + _entity_component_unrecorded_attributes | _unrecorded_attributes + ) + + # StateInfo. Set by EntityPlatform by calling async_internal_added_to_hass + # While not purely typed, it makes typehinting more useful for us + # and removes the need for constant None checks or asserts. + _state_info: StateInfo = None # type: ignore[assignment] + # Entity Properties _attr_assumed_state: bool = False _attr_attribution: str | None = None @@ -321,6 +343,13 @@ class Entity(ABC): _attr_unique_id: str | None = None _attr_unit_of_measurement: str | None + def __init_subclass__(cls, **kwargs: Any) -> None: + """Initialize an Entity subclass.""" + super().__init_subclass__(**kwargs) + cls.__combined_unrecorded_attributes = ( + cls._entity_component_unrecorded_attributes | cls._unrecorded_attributes + ) + @property def should_poll(self) -> bool: """Return True if entity has to be polled for state. @@ -875,7 +904,12 @@ class Entity(ABC): try: hass.states.async_set( - entity_id, state, attr, self.force_update, self._context + entity_id, + state, + attr, + self.force_update, + self._context, + self._state_info, ) except InvalidStateError: _LOGGER.exception("Failed to set state, fall back to %s", STATE_UNKNOWN) @@ -1081,15 +1115,19 @@ class Entity(ABC): Not to be extended by integrations. """ - info: EntityInfo = { + entity_info: EntityInfo = { "domain": self.platform.platform_name, "custom_component": "custom_components" in type(self).__module__, } if self.platform.config_entry: - info["config_entry"] = self.platform.config_entry.entry_id + entity_info["config_entry"] = self.platform.config_entry.entry_id - entity_sources(self.hass)[self.entity_id] = info + entity_sources(self.hass)[self.entity_id] = entity_info + + self._state_info = { + "unrecorded_attributes": self.__combined_unrecorded_attributes + } if self.registry_entry is not None: # This is an assert as it should never happen, but helps in tests