Minor adjustment of recorder helper (#138941)

This commit is contained in:
Erik Montnemery 2025-02-20 17:28:39 +01:00 committed by GitHub
parent 66f293c8f3
commit ff4f4111d0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 22 additions and 13 deletions

View File

@ -43,6 +43,7 @@ from homeassistant.helpers.event import (
async_track_time_interval, async_track_time_interval,
async_track_utc_time_change, async_track_utc_time_change,
) )
from homeassistant.helpers.recorder import DATA_RECORDER
from homeassistant.helpers.start import async_at_started from homeassistant.helpers.start import async_at_started
from homeassistant.helpers.typing import UNDEFINED, UndefinedType from homeassistant.helpers.typing import UNDEFINED, UndefinedType
from homeassistant.util import dt as dt_util from homeassistant.util import dt as dt_util
@ -183,7 +184,7 @@ class Recorder(threading.Thread):
self.db_retry_wait = db_retry_wait self.db_retry_wait = db_retry_wait
self.database_engine: DatabaseEngine | None = None self.database_engine: DatabaseEngine | None = None
# Database connection is ready, but non-live migration may be in progress # Database connection is ready, but non-live migration may be in progress
db_connected: asyncio.Future[bool] = hass.data[DOMAIN].db_connected db_connected: asyncio.Future[bool] = hass.data[DATA_RECORDER].db_connected
self.async_db_connected: asyncio.Future[bool] = db_connected self.async_db_connected: asyncio.Future[bool] = db_connected
# Database is ready to use but live migration may be in progress # Database is ready to use but live migration may be in progress
self.async_db_ready: asyncio.Future[bool] = hass.loop.create_future() self.async_db_ready: asyncio.Future[bool] = hass.loop.create_future()

View File

@ -24,6 +24,7 @@ import voluptuous as vol
from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT from homeassistant.const import ATTR_UNIT_OF_MEASUREMENT
from homeassistant.core import HomeAssistant, callback, valid_entity_id from homeassistant.core import HomeAssistant, callback, valid_entity_id
from homeassistant.exceptions import HomeAssistantError from homeassistant.exceptions import HomeAssistantError
from homeassistant.helpers.recorder import DATA_RECORDER
from homeassistant.helpers.singleton import singleton from homeassistant.helpers.singleton import singleton
from homeassistant.helpers.typing import UNDEFINED, UndefinedType from homeassistant.helpers.typing import UNDEFINED, UndefinedType
from homeassistant.util import dt as dt_util from homeassistant.util import dt as dt_util
@ -561,7 +562,9 @@ def _compile_statistics(
platform_stats: list[StatisticResult] = [] platform_stats: list[StatisticResult] = []
current_metadata: dict[str, tuple[int, StatisticMetaData]] = {} current_metadata: dict[str, tuple[int, StatisticMetaData]] = {}
# Collect statistics from all platforms implementing support # Collect statistics from all platforms implementing support
for domain, platform in instance.hass.data[DOMAIN].recorder_platforms.items(): for domain, platform in instance.hass.data[
DATA_RECORDER
].recorder_platforms.items():
if not ( if not (
platform_compile_statistics := getattr( platform_compile_statistics := getattr(
platform, INTEGRATION_PLATFORM_COMPILE_STATISTICS, None platform, INTEGRATION_PLATFORM_COMPILE_STATISTICS, None
@ -599,7 +602,7 @@ def _compile_statistics(
if start.minute == 50: if start.minute == 50:
# Once every hour, update issues # Once every hour, update issues
for platform in instance.hass.data[DOMAIN].recorder_platforms.values(): for platform in instance.hass.data[DATA_RECORDER].recorder_platforms.values():
if not ( if not (
platform_update_issues := getattr( platform_update_issues := getattr(
platform, INTEGRATION_PLATFORM_UPDATE_STATISTICS_ISSUES, None platform, INTEGRATION_PLATFORM_UPDATE_STATISTICS_ISSUES, None
@ -882,7 +885,7 @@ def list_statistic_ids(
# the integrations for the missing ones. # the integrations for the missing ones.
# #
# Query all integrations with a registered recorder platform # Query all integrations with a registered recorder platform
for platform in hass.data[DOMAIN].recorder_platforms.values(): for platform in hass.data[DATA_RECORDER].recorder_platforms.values():
if not ( if not (
platform_list_statistic_ids := getattr( platform_list_statistic_ids := getattr(
platform, INTEGRATION_PLATFORM_LIST_STATISTIC_IDS, None platform, INTEGRATION_PLATFORM_LIST_STATISTIC_IDS, None
@ -2232,7 +2235,7 @@ def _sorted_statistics_to_dict(
def validate_statistics(hass: HomeAssistant) -> dict[str, list[ValidationIssue]]: def validate_statistics(hass: HomeAssistant) -> dict[str, list[ValidationIssue]]:
"""Validate statistics.""" """Validate statistics."""
platform_validation: dict[str, list[ValidationIssue]] = {} platform_validation: dict[str, list[ValidationIssue]] = {}
for platform in hass.data[DOMAIN].recorder_platforms.values(): for platform in hass.data[DATA_RECORDER].recorder_platforms.values():
if platform_validate_statistics := getattr( if platform_validate_statistics := getattr(
platform, INTEGRATION_PLATFORM_VALIDATE_STATISTICS, None platform, INTEGRATION_PLATFORM_VALIDATE_STATISTICS, None
): ):
@ -2243,7 +2246,7 @@ def validate_statistics(hass: HomeAssistant) -> dict[str, list[ValidationIssue]]
def update_statistics_issues(hass: HomeAssistant) -> None: def update_statistics_issues(hass: HomeAssistant) -> None:
"""Update statistics issues.""" """Update statistics issues."""
with session_scope(hass=hass, read_only=True) as session: with session_scope(hass=hass, read_only=True) as session:
for platform in hass.data[DOMAIN].recorder_platforms.values(): for platform in hass.data[DATA_RECORDER].recorder_platforms.values():
if platform_update_statistics_issues := getattr( if platform_update_statistics_issues := getattr(
platform, INTEGRATION_PLATFORM_UPDATE_STATISTICS_ISSUES, None platform, INTEGRATION_PLATFORM_UPDATE_STATISTICS_ISSUES, None
): ):

View File

@ -11,11 +11,11 @@ import logging
import threading import threading
from typing import TYPE_CHECKING, Any from typing import TYPE_CHECKING, Any
from homeassistant.helpers.recorder import DATA_RECORDER
from homeassistant.helpers.typing import UndefinedType from homeassistant.helpers.typing import UndefinedType
from homeassistant.util.event_type import EventType from homeassistant.util.event_type import EventType
from . import entity_registry, purge, statistics from . import entity_registry, purge, statistics
from .const import DOMAIN
from .db_schema import Statistics, StatisticsShortTerm from .db_schema import Statistics, StatisticsShortTerm
from .models import StatisticData, StatisticMetaData from .models import StatisticData, StatisticMetaData
from .util import periodic_db_cleanups, session_scope from .util import periodic_db_cleanups, session_scope
@ -308,7 +308,7 @@ class AddRecorderPlatformTask(RecorderTask):
hass = instance.hass hass = instance.hass
domain = self.domain domain = self.domain
platform = self.platform platform = self.platform
platforms: dict[str, Any] = hass.data[DOMAIN].recorder_platforms platforms: dict[str, Any] = hass.data[DATA_RECORDER].recorder_platforms
platforms[domain] = platform platforms[domain] = platform

View File

@ -20,7 +20,7 @@ if TYPE_CHECKING:
_LOGGER = logging.getLogger(__name__) _LOGGER = logging.getLogger(__name__)
DOMAIN: HassKey[RecorderData] = HassKey("recorder") DATA_RECORDER: HassKey[RecorderData] = HassKey("recorder")
DATA_INSTANCE: HassKey[Recorder] = HassKey("recorder_instance") DATA_INSTANCE: HassKey[Recorder] = HassKey("recorder_instance")
@ -52,11 +52,16 @@ def async_migration_is_live(hass: HomeAssistant) -> bool:
@callback @callback
def async_initialize_recorder(hass: HomeAssistant) -> None: def async_initialize_recorder(hass: HomeAssistant) -> None:
"""Initialize recorder data.""" """Initialize recorder data.
This creates the RecorderData instance stored in hass.data[DATA_RECORDER] and
registers the basic recorder websocket API which is used by frontend to determine
if the recorder is migrating the database.
"""
# pylint: disable-next=import-outside-toplevel # pylint: disable-next=import-outside-toplevel
from homeassistant.components.recorder.basic_websocket_api import async_setup from homeassistant.components.recorder.basic_websocket_api import async_setup
hass.data[DOMAIN] = RecorderData() hass.data[DATA_RECORDER] = RecorderData()
async_setup(hass) async_setup(hass)

View File

@ -82,7 +82,7 @@ async def async_block_recorder(hass: HomeAssistant, seconds: float) -> None:
async def async_wait_recorder(hass: HomeAssistant) -> bool: async def async_wait_recorder(hass: HomeAssistant) -> bool:
"""Wait for recorder to initialize and return connection status.""" """Wait for recorder to initialize and return connection status."""
return await hass.data[recorder_helper.DOMAIN].db_connected return await hass.data[recorder_helper.DATA_RECORDER].db_connected
def get_start_time(start: datetime) -> datetime: def get_start_time(start: datetime) -> datetime:

View File

@ -1557,7 +1557,7 @@ async def _async_init_recorder_component(
assert (recorder.DOMAIN in hass.config.components) == expected_setup_result assert (recorder.DOMAIN in hass.config.components) == expected_setup_result
else: else:
# Wait for recorder to connect to the database # Wait for recorder to connect to the database
await hass.data[recorder_helper.DOMAIN].db_connected await hass.data[recorder_helper.DATA_RECORDER].db_connected
_LOGGER.info( _LOGGER.info(
"Test recorder successfully started, database location: %s", "Test recorder successfully started, database location: %s",
config[recorder.CONF_DB_URL], config[recorder.CONF_DB_URL],