Add ability to exclude attributes from being recorded by entity domain (#68824)

pull/68875/head
J. Nick Koston 2022-03-29 17:13:08 -10:00 committed by GitHub
parent ced68c1b80
commit f5a13fc51b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
5 changed files with 76 additions and 4 deletions

View File

@ -0,0 +1,10 @@
"""Integration platform for recorder."""
from __future__ import annotations
from homeassistant.core import HomeAssistant, callback
@callback
def exclude_attributes(hass: HomeAssistant) -> set[str]:
"""Exclude access_token and entity_picture from being recorded in the database."""
return {"access_token", "entity_picture"}

View File

@ -96,6 +96,7 @@ _LOGGER = logging.getLogger(__name__)
T = TypeVar("T")
EXCLUDE_ATTRIBUTES = f"{DOMAIN}_exclude_attributes_by_domain"
SERVICE_PURGE = "purge"
SERVICE_PURGE_ENTITIES = "purge_entities"
@ -274,6 +275,8 @@ def run_information_with_session(
async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
"""Set up the recorder."""
hass.data[DOMAIN] = {}
exclude_attributes_by_domain: dict[str, set[str]] = {}
hass.data[EXCLUDE_ATTRIBUTES] = exclude_attributes_by_domain
conf = config[DOMAIN]
entity_filter = convert_include_exclude_filter(conf)
auto_purge = conf[CONF_AUTO_PURGE]
@ -301,6 +304,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
db_retry_wait=db_retry_wait,
entity_filter=entity_filter,
exclude_t=exclude_t,
exclude_attributes_by_domain=exclude_attributes_by_domain,
)
instance.async_initialize()
instance.async_register()
@ -317,6 +321,8 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
async def _process_recorder_platform(hass, domain, platform):
"""Process a recorder platform."""
hass.data[DOMAIN][domain] = platform
if hasattr(platform, "exclude_attributes"):
hass.data[EXCLUDE_ATTRIBUTES][domain] = platform.exclude_attributes(hass)
@callback
@ -565,6 +571,7 @@ class Recorder(threading.Thread):
db_retry_wait: int,
entity_filter: Callable[[str], bool],
exclude_t: list[str],
exclude_attributes_by_domain: dict[str, set[str]],
) -> None:
"""Initialize the recorder."""
threading.Thread.__init__(self, name="Recorder")
@ -605,6 +612,7 @@ class Recorder(threading.Thread):
self._db_supports_row_number = True
self._database_lock_task: DatabaseLockTask | None = None
self._db_executor: DBInterruptibleThreadPoolExecutor | None = None
self._exclude_attributes_by_domain = exclude_attributes_by_domain
self.enabled = True
@ -1038,7 +1046,9 @@ class Recorder(threading.Thread):
if event.event_type == EVENT_STATE_CHANGED:
try:
dbstate = States.from_event(event)
shared_attrs = StateAttributes.shared_attrs_from_event(event)
shared_attrs = StateAttributes.shared_attrs_from_event(
event, self._exclude_attributes_by_domain
)
except (TypeError, ValueError) as ex:
_LOGGER.warning(
"State is not JSON serializable: %s: %s",

View File

@ -34,7 +34,7 @@ from homeassistant.const import (
MAX_LENGTH_STATE_ENTITY_ID,
MAX_LENGTH_STATE_STATE,
)
from homeassistant.core import Context, Event, EventOrigin, State
from homeassistant.core import Context, Event, EventOrigin, State, split_entity_id
from homeassistant.helpers.typing import UNDEFINED, UndefinedType
import homeassistant.util.dt as dt_util
@ -260,11 +260,20 @@ class StateAttributes(Base): # type: ignore[misc,valid-type]
return dbstate
@staticmethod
def shared_attrs_from_event(event: Event) -> str:
def shared_attrs_from_event(
event: Event, exclude_attrs_by_domain: dict[str, set[str]]
) -> str:
"""Create shared_attrs from a state_changed event."""
state: State | None = event.data.get("new_state")
# None state means the state was removed from the state machine
return "{}" if state is None else JSON_DUMP(state.attributes)
if state is None:
return "{}"
domain = split_entity_id(state.entity_id)[0]
if exclude_attrs := exclude_attrs_by_domain.get(domain):
return JSON_DUMP(
{k: v for k, v in state.attributes.items() if k not in exclude_attrs}
)
return JSON_DUMP(state.attributes)
@staticmethod
def hash_shared_attrs(shared_attrs: str) -> int:

View File

@ -0,0 +1,42 @@
"""The tests for camera recorder."""
from __future__ import annotations
from datetime import timedelta
from homeassistant.components import camera
from homeassistant.components.recorder.models import StateAttributes, States
from homeassistant.components.recorder.util import session_scope
from homeassistant.core import State
from homeassistant.setup import async_setup_component
from homeassistant.util import dt as dt_util
from tests.common import async_fire_time_changed, async_init_recorder_component
from tests.components.recorder.common import async_wait_recording_done_without_instance
async def test_exclude_attributes(hass):
"""Test camera registered attributes to be excluded."""
await async_init_recorder_component(hass)
await async_setup_component(
hass, camera.DOMAIN, {camera.DOMAIN: {"platform": "demo"}}
)
await hass.async_block_till_done()
async_fire_time_changed(hass, dt_util.utcnow() + timedelta(minutes=5))
await hass.async_block_till_done()
await async_wait_recording_done_without_instance(hass)
def _fetch_camera_states() -> list[State]:
with session_scope(hass=hass) as session:
native_states = []
for db_state, db_state_attributes in session.query(States, StateAttributes):
state = db_state.to_native()
state.attributes = db_state_attributes.to_native()
native_states.append(state)
return native_states
states: list[State] = await hass.async_add_executor_job(_fetch_camera_states)
assert len(states) > 1
for state in states:
assert "access_token" not in state.attributes
assert "entity_picture" not in state.attributes
assert "friendly_name" in state.attributes

View File

@ -77,6 +77,7 @@ def _default_recorder(hass):
db_retry_wait=3,
entity_filter=CONFIG_SCHEMA({DOMAIN: {}}),
exclude_t=[],
exclude_attributes_by_domain={},
)