Add WS API for removing statistics for a list of statistic_ids (#55078)

* Add WS API for removing statistics for a list of statistic_ids

* Refactor according to code review, enable foreign keys support for sqlite

* Adjust tests

* Move clear_statistics WS API to recorder

* Adjust tests after rebase

* Update docstring

* Update homeassistant/components/recorder/websocket_api.py

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

* Adjust tests after rebase

Co-authored-by: Paulus Schoutsen <paulus@home-assistant.io>

homeassistant/components/recorder/__init__.py

@@ -323,6 +323,12 @@ def _async_register_services(hass, instance):
     )
 
 
+class ClearStatisticsTask(NamedTuple):
+    """Object to store statistics_ids which for which to remove statistics."""
+
+    statistic_ids: list[str]
+
+
 class PurgeTask(NamedTuple):
     """Object to store information about purge task."""
 
@@ -570,6 +576,11 @@ class Recorder(threading.Thread):
         start = statistics.get_start_time()
         self.queue.put(StatisticsTask(start))
 
+    @callback
+    def async_clear_statistics(self, statistic_ids):
+        """Clear statistics for a list of statistic_ids."""
+        self.queue.put(ClearStatisticsTask(statistic_ids))
+
     @callback
     def _async_setup_periodic_tasks(self):
         """Prepare periodic tasks."""
@@ -763,6 +774,9 @@ class Recorder(threading.Thread):
         if isinstance(event, StatisticsTask):
             self._run_statistics(event.start)
             return
+        if isinstance(event, ClearStatisticsTask):
+            statistics.clear_statistics(self, event.statistic_ids)
+            return
         if isinstance(event, WaitTask):
             self._queue_watch.set()
             return

homeassistant/components/recorder/statistics.py

@@ -458,6 +458,14 @@ def _configured_unit(unit: str, units: UnitSystem) -> str:
     return unit
 
 
+def clear_statistics(instance: Recorder, statistic_ids: list[str]) -> None:
+    """Clear statistics for a list of statistic_ids."""
+    with session_scope(session=instance.get_session()) as session:  # type: ignore
+        session.query(StatisticsMeta).filter(
+            StatisticsMeta.statistic_id.in_(statistic_ids)
+        ).delete(synchronize_session=False)
+
+
 def list_statistic_ids(
     hass: HomeAssistant,
     statistic_type: Literal["mean"] | Literal["sum"] | None = None,

homeassistant/components/recorder/util.py

@@ -284,6 +284,9 @@ def setup_connection_for_dialect(dialect_name, dbapi_connection, first_connectio
         # approximately 8MiB of memory
         execute_on_connection(dbapi_connection, "PRAGMA cache_size = -8192")
 
+        # enable support for foreign keys
+        execute_on_connection(dbapi_connection, "PRAGMA foreign_keys=ON")
+
     if dialect_name == "mysql":
         execute_on_connection(dbapi_connection, "SET session wait_timeout=28800")
 

homeassistant/components/recorder/websocket_api.py

@@ -4,6 +4,7 @@ import voluptuous as vol
 from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
 
+from .const import DATA_INSTANCE
 from .statistics import validate_statistics
 
 
@@ -11,6 +12,7 @@ from .statistics import validate_statistics
 def async_setup(hass: HomeAssistant) -> None:
     """Set up the recorder websocket API."""
     websocket_api.async_register_command(hass, ws_validate_statistics)
+    websocket_api.async_register_command(hass, ws_clear_statistics)
 
 
 @websocket_api.websocket_command(
@@ -28,3 +30,23 @@ async def ws_validate_statistics(
         hass,
     )
     connection.send_result(msg["id"], statistic_ids)
+
+
+@websocket_api.require_admin
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "recorder/clear_statistics",
+        vol.Required("statistic_ids"): [str],
+    }
+)
+@callback
+def ws_clear_statistics(
+    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
+) -> None:
+    """Clear statistics for a list of statistic_ids.
+
+    Note: The WS call posts a job to the recorder's queue and then returns, it doesn't
+    wait until the job is completed.
+    """
+    hass.data[DATA_INSTANCE].async_clear_statistics(msg["statistic_ids"])
+    connection.send_result(msg["id"])

tests/components/recorder/test_util.py

@@ -149,15 +149,17 @@ def test_setup_connection_for_dialect_sqlite():
 
     util.setup_connection_for_dialect("sqlite", dbapi_connection, True)
 
-    assert len(execute_mock.call_args_list) == 2
+    assert len(execute_mock.call_args_list) == 3
     assert execute_mock.call_args_list[0][0][0] == "PRAGMA journal_mode=WAL"
     assert execute_mock.call_args_list[1][0][0] == "PRAGMA cache_size = -8192"
+    assert execute_mock.call_args_list[2][0][0] == "PRAGMA foreign_keys=ON"
 
     execute_mock.reset_mock()
     util.setup_connection_for_dialect("sqlite", dbapi_connection, False)
 
-    assert len(execute_mock.call_args_list) == 1
+    assert len(execute_mock.call_args_list) == 2
     assert execute_mock.call_args_list[0][0][0] == "PRAGMA cache_size = -8192"
+    assert execute_mock.call_args_list[1][0][0] == "PRAGMA foreign_keys=ON"
 
 
 def test_basic_sanity_check(hass_recorder):

tests/components/recorder/test_websocket_api.py

@@ -3,6 +3,7 @@
 from datetime import timedelta
 
 import pytest
+from pytest import approx
 
 from homeassistant.components.recorder.const import DATA_INSTANCE
 from homeassistant.components.recorder.models import StatisticsMeta
@@ -11,6 +12,8 @@ from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 from homeassistant.util.unit_system import IMPERIAL_SYSTEM, METRIC_SYSTEM
 
+from .common import trigger_db_commit
+
 from tests.common import init_recorder_component
 
 BATTERY_SENSOR_ATTRIBUTES = {
@@ -240,3 +243,137 @@ async def test_validate_statistics_unsupported_device_class(
     # Remove the state - empty response
     hass.states.async_remove("sensor.test")
     await assert_validation_result(client, {})
+
+
+async def test_clear_statistics(hass, hass_ws_client):
+    """Test removing statistics."""
+    now = dt_util.utcnow()
+
+    units = METRIC_SYSTEM
+    attributes = POWER_SENSOR_ATTRIBUTES
+    state = 10
+    value = 10000
+
+    hass.config.units = units
+    await hass.async_add_executor_job(init_recorder_component, hass)
+    await async_setup_component(hass, "history", {})
+    await async_setup_component(hass, "sensor", {})
+    await hass.async_add_executor_job(hass.data[DATA_INSTANCE].block_till_done)
+    hass.states.async_set("sensor.test1", state, attributes=attributes)
+    hass.states.async_set("sensor.test2", state * 2, attributes=attributes)
+    hass.states.async_set("sensor.test3", state * 3, attributes=attributes)
+    await hass.async_block_till_done()
+
+    await hass.async_add_executor_job(trigger_db_commit, hass)
+    await hass.async_block_till_done()
+
+    hass.data[DATA_INSTANCE].do_adhoc_statistics(start=now)
+    await hass.async_add_executor_job(hass.data[DATA_INSTANCE].block_till_done)
+
+    client = await hass_ws_client()
+    await client.send_json(
+        {
+            "id": 1,
+            "type": "history/statistics_during_period",
+            "start_time": now.isoformat(),
+            "period": "5minute",
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    expected_response = {
+        "sensor.test1": [
+            {
+                "statistic_id": "sensor.test1",
+                "start": now.isoformat(),
+                "end": (now + timedelta(minutes=5)).isoformat(),
+                "mean": approx(value),
+                "min": approx(value),
+                "max": approx(value),
+                "last_reset": None,
+                "state": None,
+                "sum": None,
+                "sum_decrease": None,
+                "sum_increase": None,
+            }
+        ],
+        "sensor.test2": [
+            {
+                "statistic_id": "sensor.test2",
+                "start": now.isoformat(),
+                "end": (now + timedelta(minutes=5)).isoformat(),
+                "mean": approx(value * 2),
+                "min": approx(value * 2),
+                "max": approx(value * 2),
+                "last_reset": None,
+                "state": None,
+                "sum": None,
+                "sum_decrease": None,
+                "sum_increase": None,
+            }
+        ],
+        "sensor.test3": [
+            {
+                "statistic_id": "sensor.test3",
+                "start": now.isoformat(),
+                "end": (now + timedelta(minutes=5)).isoformat(),
+                "mean": approx(value * 3),
+                "min": approx(value * 3),
+                "max": approx(value * 3),
+                "last_reset": None,
+                "state": None,
+                "sum": None,
+                "sum_decrease": None,
+                "sum_increase": None,
+            }
+        ],
+    }
+    assert response["result"] == expected_response
+
+    await client.send_json(
+        {
+            "id": 2,
+            "type": "recorder/clear_statistics",
+            "statistic_ids": ["sensor.test"],
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    await hass.async_add_executor_job(hass.data[DATA_INSTANCE].block_till_done)
+
+    client = await hass_ws_client()
+    await client.send_json(
+        {
+            "id": 3,
+            "type": "history/statistics_during_period",
+            "start_time": now.isoformat(),
+            "period": "5minute",
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == expected_response
+
+    await client.send_json(
+        {
+            "id": 4,
+            "type": "recorder/clear_statistics",
+            "statistic_ids": ["sensor.test1", "sensor.test3"],
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    await hass.async_add_executor_job(hass.data[DATA_INSTANCE].block_till_done)
+
+    client = await hass_ws_client()
+    await client.send_json(
+        {
+            "id": 5,
+            "type": "history/statistics_during_period",
+            "start_time": now.isoformat(),
+            "period": "5minute",
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == {"sensor.test2": expected_response["sensor.test2"]}