Add keep_days to recorder.purge_entities (#89726)

homeassistant/components/recorder/services.py

@@ -9,7 +9,10 @@ import voluptuous as vol
 from homeassistant.core import HomeAssistant, ServiceCall, callback
 import homeassistant.helpers.config_validation as cv
 from homeassistant.helpers.entityfilter import generate_filter
-from homeassistant.helpers.service import async_extract_entity_ids
+from homeassistant.helpers.service import (
+    async_extract_entity_ids,
+    async_register_admin_service,
+)
 import homeassistant.util.dt as dt_util
 
 from .const import ATTR_APPLY_FILTER, ATTR_KEEP_DAYS, ATTR_REPACK, DOMAIN
@@ -38,6 +41,7 @@ SERVICE_PURGE_ENTITIES_SCHEMA = vol.Schema(
         vol.Optional(ATTR_ENTITY_GLOBS, default=[]): vol.All(
             cv.ensure_list, [cv.string]
         ),
+        vol.Optional(ATTR_KEEP_DAYS, default=0): cv.positive_int,
     }
 ).extend(cv.ENTITY_SERVICE_FIELDS)
 
@@ -56,8 +60,12 @@ def _async_register_purge_service(hass: HomeAssistant, instance: Recorder) -> No
         purge_before = dt_util.utcnow() - timedelta(days=keep_days)
         instance.queue_task(PurgeTask(purge_before, repack, apply_filter))
 
-    hass.services.async_register(
+    async_register_admin_service(
-        DOMAIN, SERVICE_PURGE, async_handle_purge_service, schema=SERVICE_PURGE_SCHEMA
+        hass,
+        DOMAIN,
+        SERVICE_PURGE,
+        async_handle_purge_service,
+        schema=SERVICE_PURGE_SCHEMA,
     )
 
 
@@ -69,12 +77,14 @@ def _async_register_purge_entities_service(
         """Handle calls to the purge entities service."""
         entity_ids = await async_extract_entity_ids(hass, service)
         domains = service.data.get(ATTR_DOMAINS, [])
+        keep_days = service.data.get(ATTR_KEEP_DAYS, 0)
         entity_globs = service.data.get(ATTR_ENTITY_GLOBS, [])
         entity_filter = generate_filter(domains, list(entity_ids), [], [], entity_globs)
-        purge_before = dt_util.utcnow()
+        purge_before = dt_util.utcnow() - timedelta(days=keep_days)
         instance.queue_task(PurgeEntitiesTask(entity_filter, purge_before))
 
-    hass.services.async_register(
+    async_register_admin_service(
+        hass,
         DOMAIN,
         SERVICE_PURGE_ENTITIES,
         async_handle_purge_entities_service,
@@ -87,7 +97,8 @@ def _async_register_enable_service(hass: HomeAssistant, instance: Recorder) -> N
     async def async_handle_enable_service(service: ServiceCall) -> None:
         instance.set_enable(True)
 
-    hass.services.async_register(
+    async_register_admin_service(
+        hass,
         DOMAIN,
         SERVICE_ENABLE,
         async_handle_enable_service,
@@ -100,7 +111,8 @@ def _async_register_disable_service(hass: HomeAssistant, instance: Recorder) ->
     async def async_handle_disable_service(service: ServiceCall) -> None:
         instance.set_enable(False)
 
-    hass.services.async_register(
+    async_register_admin_service(
+        hass,
         DOMAIN,
         SERVICE_DISABLE,
         async_handle_disable_service,

homeassistant/components/recorder/services.yaml

@@ -51,6 +51,16 @@ purge_entities:
       selector:
         object:
 
+    keep_days:
+      name: Days to keep
+      description: Number of history days to keep in database of matching rows. The default of 0 days will remove all matching rows.
+      default: 0
+      selector:
+        number:
+          min: 0
+          max: 365
+          unit_of_measurement: days
+
 disable:
   name: Disable
   description: Stop the recording of events and state changes

tests/components/recorder/test_purge.py

@@ -4,6 +4,7 @@ import json
 import sqlite3
 from unittest.mock import patch
 
+from freezegun import freeze_time
 import pytest
 from sqlalchemy.exc import DatabaseError, OperationalError
 from sqlalchemy.orm.session import Session
@@ -25,6 +26,7 @@ from homeassistant.components.recorder.db_schema import (
     StatisticsRuns,
     StatisticsShortTerm,
 )
+from homeassistant.components.recorder.history import get_significant_states
 from homeassistant.components.recorder.purge import purge_old_data
 from homeassistant.components.recorder.queries import select_event_type_ids
 from homeassistant.components.recorder.services import (
@@ -2021,3 +2023,70 @@ async def test_purge_old_states_purges_the_state_metadata_ids(
         assert finished
         assert states.count() == 0
         assert states_meta.count() == 0
+
+
+async def test_purge_entities_keep_days(
+    async_setup_recorder_instance: RecorderInstanceGenerator,
+    hass: HomeAssistant,
+) -> None:
+    """Test purging states with an entity filter and keep_days."""
+    instance = await async_setup_recorder_instance(hass, {})
+    await hass.async_block_till_done()
+    await async_wait_recording_done(hass)
+    start = dt_util.utcnow()
+    two_days_ago = start - timedelta(days=2)
+    one_week_ago = start - timedelta(days=7)
+    one_month_ago = start - timedelta(days=30)
+    with freeze_time(one_week_ago):
+        hass.states.async_set("sensor.keep", "initial")
+        hass.states.async_set("sensor.purge", "initial")
+
+    await async_wait_recording_done(hass)
+
+    with freeze_time(two_days_ago):
+        hass.states.async_set("sensor.purge", "two_days_ago")
+
+    await async_wait_recording_done(hass)
+
+    hass.states.async_set("sensor.purge", "now")
+    hass.states.async_set("sensor.keep", "now")
+    await async_recorder_block_till_done(hass)
+
+    states = await instance.async_add_executor_job(
+        get_significant_states, hass, one_month_ago
+    )
+    assert len(states["sensor.keep"]) == 2
+    assert len(states["sensor.purge"]) == 3
+
+    await hass.services.async_call(
+        recorder.DOMAIN,
+        SERVICE_PURGE_ENTITIES,
+        {
+            "entity_id": "sensor.purge",
+            "keep_days": 1,
+        },
+    )
+    await async_recorder_block_till_done(hass)
+    await async_wait_purge_done(hass)
+
+    states = await instance.async_add_executor_job(
+        get_significant_states, hass, one_month_ago
+    )
+    assert len(states["sensor.keep"]) == 2
+    assert len(states["sensor.purge"]) == 1
+
+    await hass.services.async_call(
+        recorder.DOMAIN,
+        SERVICE_PURGE_ENTITIES,
+        {
+            "entity_id": "sensor.purge",
+        },
+    )
+    await async_recorder_block_till_done(hass)
+    await async_wait_purge_done(hass)
+
+    states = await instance.async_add_executor_job(
+        get_significant_states, hass, one_month_ago
+    )
+    assert len(states["sensor.keep"]) == 2
+    assert "sensor.purge" not in states