Add API endpoint get_statistics_metadata (#68471)

Co-authored-by: Erik Montnemery <erik@montnemery.com>

homeassistant/components/history/__init__.py

@@ -171,6 +171,7 @@ async def ws_get_list_statistic_ids(
     statistic_ids = await get_instance(hass).async_add_executor_job(
         list_statistic_ids,
         hass,
+        None,
         msg.get("statistic_type"),
     )
     connection.send_result(msg["id"], statistic_ids)

homeassistant/components/recorder/statistics.py

@@ -718,21 +718,22 @@ def update_statistics_metadata(
 
 def list_statistic_ids(
     hass: HomeAssistant,
+    statistic_ids: list[str] | tuple[str] | None = None,
     statistic_type: Literal["mean"] | Literal["sum"] | None = None,
 ) -> list[dict | None]:
-    """Return all statistic_ids and unit of measurement.
+    """Return all statistic_ids (or filtered one) and unit of measurement.
 
     Queries the database for existing statistic_ids, as well as integrations with
     a recorder platform for statistic_ids which will be added in the next statistics
     period.
     """
     units = hass.config.units
-    statistic_ids = {}
+    result = {}
 
     # Query the database
     with session_scope(hass=hass) as session:
         metadata = get_metadata_with_session(
-            hass, session, statistic_type=statistic_type
+            hass, session, statistic_type=statistic_type, statistic_ids=statistic_ids
         )
 
         for _, meta in metadata.values():
@@ -741,7 +742,7 @@ def list_statistic_ids(
                 unit = _configured_unit(unit, units)
             meta["unit_of_measurement"] = unit
 
-        statistic_ids = {
+        result = {
             meta["statistic_id"]: {
                 "name": meta["name"],
                 "source": meta["source"],
@@ -754,7 +755,9 @@ def list_statistic_ids(
     for platform in hass.data[DOMAIN].values():
         if not hasattr(platform, "list_statistic_ids"):
             continue
-        platform_statistic_ids = platform.list_statistic_ids(hass, statistic_type)
+        platform_statistic_ids = platform.list_statistic_ids(
+            hass, statistic_ids=statistic_ids, statistic_type=statistic_type
+        )
 
         for statistic_id, info in platform_statistic_ids.items():
             if (unit := info["unit_of_measurement"]) is not None:
@@ -763,7 +766,7 @@ def list_statistic_ids(
             platform_statistic_ids[statistic_id]["unit_of_measurement"] = unit
 
         for key, value in platform_statistic_ids.items():
-            statistic_ids.setdefault(key, value)
+            result.setdefault(key, value)
 
     # Return a list of statistic_id + metadata
     return [
@@ -773,7 +776,7 @@ def list_statistic_ids(
             "source": info["source"],
             "unit_of_measurement": info["unit_of_measurement"],
         }
-        for _id, info in statistic_ids.items()
+        for _id, info in result.items()
     ]
 
 

homeassistant/components/recorder/websocket_api.py

@@ -10,7 +10,7 @@ from homeassistant.components import websocket_api
 from homeassistant.core import HomeAssistant, callback
 
 from .const import DATA_INSTANCE, MAX_QUEUE_BACKLOG
-from .statistics import validate_statistics
+from .statistics import list_statistic_ids, validate_statistics
 from .util import async_migration_in_progress
 
 if TYPE_CHECKING:
@@ -24,6 +24,7 @@ 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.async_register_command(hass, ws_get_statistics_metadata)
     websocket_api.async_register_command(hass, ws_update_statistics_metadata)
     websocket_api.async_register_command(hass, ws_info)
     websocket_api.async_register_command(hass, ws_backup_start)
@@ -68,6 +69,23 @@ def ws_clear_statistics(
     connection.send_result(msg["id"])
 
 
+@websocket_api.websocket_command(
+    {
+        vol.Required("type"): "recorder/get_statistics_metadata",
+        vol.Optional("statistic_ids"): [str],
+    }
+)
+@websocket_api.async_response
+async def ws_get_statistics_metadata(
+    hass: HomeAssistant, connection: websocket_api.ActiveConnection, msg: dict
+) -> None:
+    """Get metadata for a list of statistic_ids."""
+    statistic_ids = await hass.async_add_executor_job(
+        list_statistic_ids, hass, msg.get("statistic_ids")
+    )
+    connection.send_result(msg["id"], statistic_ids)
+
+
 @websocket_api.require_admin
 @websocket_api.websocket_command(
     {

homeassistant/components/sensor/recorder.py

@@ -596,11 +596,15 @@ def _compile_statistics(  # noqa: C901
     return result
 
 
-def list_statistic_ids(hass: HomeAssistant, statistic_type: str | None = None) -> dict:
+def list_statistic_ids(
-    """Return statistic_ids and meta data."""
+    hass: HomeAssistant,
+    statistic_ids: list[str] | tuple[str] | None = None,
+    statistic_type: str | None = None,
+) -> dict:
+    """Return all or filtered statistic_ids and meta data."""
     entities = _get_sensor_states(hass)
 
-    statistic_ids = {}
+    result = {}
 
     for state in entities:
         state_class = state.attributes[ATTR_STATE_CLASS]
@@ -611,6 +615,9 @@ def list_statistic_ids(hass: HomeAssistant, statistic_type: str | None = None) -
         if statistic_type is not None and statistic_type not in provided_statistics:
             continue
 
+        if statistic_ids is not None and state.entity_id not in statistic_ids:
+            continue
+
         if (
             "sum" in provided_statistics
             and ATTR_LAST_RESET not in state.attributes
@@ -619,7 +626,7 @@ def list_statistic_ids(hass: HomeAssistant, statistic_type: str | None = None) -
             continue
 
         if device_class not in UNIT_CONVERSIONS:
-            statistic_ids[state.entity_id] = {
+            result[state.entity_id] = {
                 "source": RECORDER_DOMAIN,
                 "unit_of_measurement": native_unit,
             }
@@ -629,12 +636,12 @@ def list_statistic_ids(hass: HomeAssistant, statistic_type: str | None = None) -
             continue
 
         statistics_unit = DEVICE_CLASS_UNITS[device_class]
-        statistic_ids[state.entity_id] = {
+        result[state.entity_id] = {
             "source": RECORDER_DOMAIN,
             "unit_of_measurement": statistics_unit,
         }
 
-    return statistic_ids
+    return result
 
 
 def validate_statistics(

tests/components/recorder/test_websocket_api.py

@@ -9,6 +9,7 @@ from pytest import approx
 
 from homeassistant.components import recorder
 from homeassistant.components.recorder.const import DATA_INSTANCE
+from homeassistant.components.recorder.statistics import async_add_external_statistics
 from homeassistant.setup import async_setup_component
 import homeassistant.util.dt as dt_util
 from homeassistant.util.unit_system import METRIC_SYSTEM
@@ -35,6 +36,16 @@ TEMPERATURE_SENSOR_ATTRIBUTES = {
     "state_class": "measurement",
     "unit_of_measurement": "°C",
 }
+ENERGY_SENSOR_ATTRIBUTES = {
+    "device_class": "energy",
+    "state_class": "total",
+    "unit_of_measurement": "kWh",
+}
+GAS_SENSOR_ATTRIBUTES = {
+    "device_class": "gas",
+    "state_class": "total",
+    "unit_of_measurement": "m³",
+}
 
 
 async def test_validate_statistics(hass, hass_ws_client):
@@ -421,3 +432,125 @@ async def test_backup_end_without_start(
     response = await client.receive_json()
     assert not response["success"]
     assert response["error"]["code"] == "database_unlock_failed"
+
+
+@pytest.mark.parametrize(
+    "units, attributes, unit",
+    [
+        (METRIC_SYSTEM, GAS_SENSOR_ATTRIBUTES, "m³"),
+        (METRIC_SYSTEM, ENERGY_SENSOR_ATTRIBUTES, "kWh"),
+    ],
+)
+async def test_get_statistics_metadata(hass, hass_ws_client, units, attributes, unit):
+    """Test get_statistics_metadata."""
+    now = dt_util.utcnow()
+
+    hass.config.units = units
+    await hass.async_add_executor_job(init_recorder_component, hass)
+    await async_setup_component(hass, "history", {"history": {}})
+    await async_setup_component(hass, "sensor", {})
+    await async_init_recorder_component(hass)
+    await hass.async_add_executor_job(hass.data[recorder.DATA_INSTANCE].block_till_done)
+
+    client = await hass_ws_client()
+    await client.send_json({"id": 1, "type": "recorder/get_statistics_metadata"})
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == []
+
+    period1 = dt_util.as_utc(dt_util.parse_datetime("2021-09-01 00:00:00"))
+    period2 = dt_util.as_utc(dt_util.parse_datetime("2021-09-30 23:00:00"))
+    period3 = dt_util.as_utc(dt_util.parse_datetime("2021-10-01 00:00:00"))
+    period4 = dt_util.as_utc(dt_util.parse_datetime("2021-10-31 23:00:00"))
+    external_energy_statistics_1 = (
+        {
+            "start": period1,
+            "last_reset": None,
+            "state": 0,
+            "sum": 2,
+        },
+        {
+            "start": period2,
+            "last_reset": None,
+            "state": 1,
+            "sum": 3,
+        },
+        {
+            "start": period3,
+            "last_reset": None,
+            "state": 2,
+            "sum": 5,
+        },
+        {
+            "start": period4,
+            "last_reset": None,
+            "state": 3,
+            "sum": 8,
+        },
+    )
+    external_energy_metadata_1 = {
+        "has_mean": False,
+        "has_sum": True,
+        "name": "Total imported energy",
+        "source": "test",
+        "statistic_id": "test:total_gas",
+        "unit_of_measurement": unit,
+    }
+
+    async_add_external_statistics(
+        hass, external_energy_metadata_1, external_energy_statistics_1
+    )
+
+    hass.states.async_set("sensor.test", 10, 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.states.async_set("sensor.test2", 10, attributes=attributes)
+    await hass.async_block_till_done()
+
+    await hass.async_add_executor_job(trigger_db_commit, hass)
+    await hass.async_block_till_done()
+
+    await client.send_json(
+        {
+            "id": 2,
+            "type": "recorder/get_statistics_metadata",
+            "statistic_ids": ["sensor.test"],
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == [
+        {
+            "statistic_id": "sensor.test",
+            "name": None,
+            "source": "recorder",
+            "unit_of_measurement": unit,
+        }
+    ]
+
+    hass.data[recorder.DATA_INSTANCE].do_adhoc_statistics(start=now)
+    await hass.async_add_executor_job(hass.data[recorder.DATA_INSTANCE].block_till_done)
+    # Remove the state, statistics will now be fetched from the database
+    hass.states.async_remove("sensor.test")
+    await hass.async_block_till_done()
+
+    await client.send_json(
+        {
+            "id": 3,
+            "type": "recorder/get_statistics_metadata",
+            "statistic_ids": ["sensor.test"],
+        }
+    )
+    response = await client.receive_json()
+    assert response["success"]
+    assert response["result"] == [
+        {
+            "statistic_id": "sensor.test",
+            "name": None,
+            "source": "recorder",
+            "unit_of_measurement": unit,
+        }
+    ]