diff --git a/homeassistant/components/recorder/core.py b/homeassistant/components/recorder/core.py index d5e095d8104..833064cee63 100644 --- a/homeassistant/components/recorder/core.py +++ b/homeassistant/components/recorder/core.py @@ -61,7 +61,9 @@ from .db_schema import ( Events, StateAttributes, States, + Statistics, StatisticsRuns, + StatisticsShortTerm, ) from .executor import DBInterruptibleThreadPoolExecutor from .models import ( @@ -534,10 +536,13 @@ class Recorder(threading.Thread): @callback def async_import_statistics( - self, metadata: StatisticMetaData, stats: Iterable[StatisticData] + self, + metadata: StatisticMetaData, + stats: Iterable[StatisticData], + table: type[Statistics | StatisticsShortTerm], ) -> None: """Schedule import of statistics.""" - self.queue_task(ImportStatisticsTask(metadata, stats)) + self.queue_task(ImportStatisticsTask(metadata, stats, table)) @callback def _async_setup_periodic_tasks(self) -> None: diff --git a/homeassistant/components/recorder/statistics.py b/homeassistant/components/recorder/statistics.py index 3ff438f50da..773676b07b0 100644 --- a/homeassistant/components/recorder/statistics.py +++ b/homeassistant/components/recorder/statistics.py @@ -1461,7 +1461,7 @@ def _async_import_statistics( statistic["last_reset"] = dt_util.as_utc(last_reset) # Insert job in recorder's queue - get_instance(hass).async_import_statistics(metadata, statistics) + get_instance(hass).async_import_statistics(metadata, statistics, Statistics) @callback @@ -1551,6 +1551,7 @@ def import_statistics( instance: Recorder, metadata: StatisticMetaData, statistics: Iterable[StatisticData], + table: type[Statistics | StatisticsShortTerm], ) -> bool: """Process an import_statistics job.""" @@ -1564,11 +1565,11 @@ def import_statistics( metadata_id = _update_or_add_metadata(session, metadata, old_metadata_dict) for stat in statistics: if stat_id := _statistics_exists( - session, Statistics, metadata_id, stat["start"] + session, table, metadata_id, stat["start"] ): - _update_statistics(session, Statistics, stat_id, stat) + _update_statistics(session, table, stat_id, stat) else: - _insert_statistics(session, Statistics, metadata_id, stat) + _insert_statistics(session, table, metadata_id, stat) return True diff --git a/homeassistant/components/recorder/tasks.py b/homeassistant/components/recorder/tasks.py index 4fa3a3cc40c..1b8e03ebf17 100644 --- a/homeassistant/components/recorder/tasks.py +++ b/homeassistant/components/recorder/tasks.py @@ -14,6 +14,7 @@ from homeassistant.helpers.typing import UndefinedType from . import purge, statistics from .const import DOMAIN, EXCLUDE_ATTRIBUTES +from .db_schema import Statistics, StatisticsShortTerm from .models import StatisticData, StatisticMetaData from .util import periodic_db_cleanups @@ -147,13 +148,18 @@ class ImportStatisticsTask(RecorderTask): metadata: StatisticMetaData statistics: Iterable[StatisticData] + table: type[Statistics | StatisticsShortTerm] def run(self, instance: Recorder) -> None: """Run statistics task.""" - if statistics.import_statistics(instance, self.metadata, self.statistics): + if statistics.import_statistics( + instance, self.metadata, self.statistics, self.table + ): return # Schedule a new statistics task if this one didn't finish - instance.queue_task(ImportStatisticsTask(self.metadata, self.statistics)) + instance.queue_task( + ImportStatisticsTask(self.metadata, self.statistics, self.table) + ) @dataclass