Allow specifying the target table when importing statistics (#80230)

Allow specifying the table when importing statistics

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:

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
 

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