Simplify service registration in recorder (#146237)

homeassistant/components/recorder/__init__.py

@@ -45,7 +45,7 @@ from .const import (  # noqa: F401
     SupportedDialect,
 )
 from .core import Recorder
-from .services import async_register_services
+from .services import async_setup_services
 from .tasks import AddRecorderPlatformTask
 from .util import get_instance
 
@@ -174,7 +174,7 @@ async def async_setup(hass: HomeAssistant, config: ConfigType) -> bool:
     instance.async_initialize()
     instance.async_register()
     instance.start()
-    async_register_services(hass, instance)
+    async_setup_services(hass)
     websocket_api.async_setup(hass)
 
     await _async_setup_integration_platform(hass, instance)

homeassistant/components/recorder/services.py

@@ -17,6 +17,7 @@ from homeassistant.core import (
 )
 from homeassistant.helpers import config_validation as cv
 from homeassistant.helpers.entityfilter import generate_filter
+from homeassistant.helpers.recorder import DATA_INSTANCE
 from homeassistant.helpers.service import (
     async_extract_entity_ids,
     async_register_admin_service,
@@ -25,7 +26,6 @@ from homeassistant.util import dt as dt_util
 from homeassistant.util.json import JsonArrayType, JsonObjectType
 
 from .const import ATTR_APPLY_FILTER, ATTR_KEEP_DAYS, ATTR_REPACK, DOMAIN
-from .core import Recorder
 from .statistics import statistics_during_period
 from .tasks import PurgeEntitiesTask, PurgeTask
 
@@ -87,10 +87,10 @@ SERVICE_GET_STATISTICS_SCHEMA = vol.Schema(
 )
 
 
-@callback
+async def _async_handle_purge_service(service: ServiceCall) -> None:
-def _async_register_purge_service(hass: HomeAssistant, instance: Recorder) -> None:
-    async def async_handle_purge_service(service: ServiceCall) -> None:
     """Handle calls to the purge service."""
+    hass = service.hass
+    instance = hass.data[DATA_INSTANCE]
     kwargs = service.data
     keep_days = kwargs.get(ATTR_KEEP_DAYS, instance.keep_days)
     repack = cast(bool, kwargs[ATTR_REPACK])
@@ -98,79 +98,35 @@ 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))
 
-    async_register_admin_service(
-        hass,
-        DOMAIN,
-        SERVICE_PURGE,
-        async_handle_purge_service,
-        schema=SERVICE_PURGE_SCHEMA,
-    )
 
-
+async def _async_handle_purge_entities_service(service: ServiceCall) -> None:
-@callback
-def _async_register_purge_entities_service(
-    hass: HomeAssistant, instance: Recorder
-) -> None:
-    async def async_handle_purge_entities_service(service: ServiceCall) -> None:
     """Handle calls to the purge entities service."""
+    hass = service.hass
     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() - timedelta(days=keep_days)
-        instance.queue_task(PurgeEntitiesTask(entity_filter, purge_before))
+    hass.data[DATA_INSTANCE].queue_task(PurgeEntitiesTask(entity_filter, purge_before))
-
-    async_register_admin_service(
-        hass,
-        DOMAIN,
-        SERVICE_PURGE_ENTITIES,
-        async_handle_purge_entities_service,
-        schema=SERVICE_PURGE_ENTITIES_SCHEMA,
-    )
 
 
-@callback
+async def _async_handle_enable_service(service: ServiceCall) -> None:
-def _async_register_enable_service(hass: HomeAssistant, instance: Recorder) -> None:
+    service.hass.data[DATA_INSTANCE].set_enable(True)
-    async def async_handle_enable_service(service: ServiceCall) -> None:
-        instance.set_enable(True)
-
-    async_register_admin_service(
-        hass,
-        DOMAIN,
-        SERVICE_ENABLE,
-        async_handle_enable_service,
-        schema=SERVICE_ENABLE_SCHEMA,
-    )
 
 
-@callback
+async def _async_handle_disable_service(service: ServiceCall) -> None:
-def _async_register_disable_service(hass: HomeAssistant, instance: Recorder) -> None:
+    service.hass.data[DATA_INSTANCE].set_enable(False)
-    async def async_handle_disable_service(service: ServiceCall) -> None:
-        instance.set_enable(False)
-
-    async_register_admin_service(
-        hass,
-        DOMAIN,
-        SERVICE_DISABLE,
-        async_handle_disable_service,
-        schema=SERVICE_DISABLE_SCHEMA,
-    )
 
 
-@callback
+async def _async_handle_get_statistics_service(
-def _async_register_get_statistics_service(
-    hass: HomeAssistant, instance: Recorder
-) -> None:
-    async def async_handle_get_statistics_service(
     service: ServiceCall,
 ) -> ServiceResponse:
     """Handle calls to the get_statistics service."""
+    hass = service.hass
     start_time = dt_util.as_utc(service.data["start_time"])
     end_time = (
-            dt_util.as_utc(service.data["end_time"])
+        dt_util.as_utc(service.data["end_time"]) if "end_time" in service.data else None
-            if "end_time" in service.data
-            else None
     )
 
     statistic_ids = service.data["statistic_ids"]
@@ -178,7 +134,7 @@ def _async_register_get_statistics_service(
     period = service.data["period"]
     units = service.data.get("units")
 
-        result = await instance.async_add_executor_job(
+    result = await hass.data[DATA_INSTANCE].async_add_executor_job(
         statistics_during_period,
         hass,
         start_time,
@@ -221,21 +177,47 @@ def _async_register_get_statistics_service(
 
     return {"statistics": formatted_result}
 
+
+@callback
+def async_setup_services(hass: HomeAssistant) -> None:
+    """Register recorder services."""
+    async_register_admin_service(
+        hass,
+        DOMAIN,
+        SERVICE_PURGE,
+        _async_handle_purge_service,
+        schema=SERVICE_PURGE_SCHEMA,
+    )
+
+    async_register_admin_service(
+        hass,
+        DOMAIN,
+        SERVICE_PURGE_ENTITIES,
+        _async_handle_purge_entities_service,
+        schema=SERVICE_PURGE_ENTITIES_SCHEMA,
+    )
+
+    async_register_admin_service(
+        hass,
+        DOMAIN,
+        SERVICE_ENABLE,
+        _async_handle_enable_service,
+        schema=SERVICE_ENABLE_SCHEMA,
+    )
+
+    async_register_admin_service(
+        hass,
+        DOMAIN,
+        SERVICE_DISABLE,
+        _async_handle_disable_service,
+        schema=SERVICE_DISABLE_SCHEMA,
+    )
+
     async_register_admin_service(
         hass,
         DOMAIN,
         SERVICE_GET_STATISTICS,
-        async_handle_get_statistics_service,
+        _async_handle_get_statistics_service,
         schema=SERVICE_GET_STATISTICS_SCHEMA,
         supports_response=SupportsResponse.ONLY,
     )
-
-
-@callback
-def async_register_services(hass: HomeAssistant, instance: Recorder) -> None:
-    """Register recorder services."""
-    _async_register_purge_service(hass, instance)
-    _async_register_purge_entities_service(hass, instance)
-    _async_register_enable_service(hass, instance)
-    _async_register_disable_service(hass, instance)
-    _async_register_get_statistics_service(hass, instance)