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

Allow specifying the table when importing statistics
This commit is contained in:
Erik Montnemery 2022-10-14 08:32:19 +02:00 committed by GitHub
parent d9af274da3
commit 2e261d5dc2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 20 additions and 8 deletions

View File

@ -61,7 +61,9 @@ from .db_schema import (
Events, Events,
StateAttributes, StateAttributes,
States, States,
Statistics,
StatisticsRuns, StatisticsRuns,
StatisticsShortTerm,
) )
from .executor import DBInterruptibleThreadPoolExecutor from .executor import DBInterruptibleThreadPoolExecutor
from .models import ( from .models import (
@ -534,10 +536,13 @@ class Recorder(threading.Thread):
@callback @callback
def async_import_statistics( def async_import_statistics(
self, metadata: StatisticMetaData, stats: Iterable[StatisticData] self,
metadata: StatisticMetaData,
stats: Iterable[StatisticData],
table: type[Statistics | StatisticsShortTerm],
) -> None: ) -> None:
"""Schedule import of statistics.""" """Schedule import of statistics."""
self.queue_task(ImportStatisticsTask(metadata, stats)) self.queue_task(ImportStatisticsTask(metadata, stats, table))
@callback @callback
def _async_setup_periodic_tasks(self) -> None: def _async_setup_periodic_tasks(self) -> None:

View File

@ -1461,7 +1461,7 @@ def _async_import_statistics(
statistic["last_reset"] = dt_util.as_utc(last_reset) statistic["last_reset"] = dt_util.as_utc(last_reset)
# Insert job in recorder's queue # Insert job in recorder's queue
get_instance(hass).async_import_statistics(metadata, statistics) get_instance(hass).async_import_statistics(metadata, statistics, Statistics)
@callback @callback
@ -1551,6 +1551,7 @@ def import_statistics(
instance: Recorder, instance: Recorder,
metadata: StatisticMetaData, metadata: StatisticMetaData,
statistics: Iterable[StatisticData], statistics: Iterable[StatisticData],
table: type[Statistics | StatisticsShortTerm],
) -> bool: ) -> bool:
"""Process an import_statistics job.""" """Process an import_statistics job."""
@ -1564,11 +1565,11 @@ def import_statistics(
metadata_id = _update_or_add_metadata(session, metadata, old_metadata_dict) metadata_id = _update_or_add_metadata(session, metadata, old_metadata_dict)
for stat in statistics: for stat in statistics:
if stat_id := _statistics_exists( 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: else:
_insert_statistics(session, Statistics, metadata_id, stat) _insert_statistics(session, table, metadata_id, stat)
return True return True

View File

@ -14,6 +14,7 @@ from homeassistant.helpers.typing import UndefinedType
from . import purge, statistics from . import purge, statistics
from .const import DOMAIN, EXCLUDE_ATTRIBUTES from .const import DOMAIN, EXCLUDE_ATTRIBUTES
from .db_schema import Statistics, StatisticsShortTerm
from .models import StatisticData, StatisticMetaData from .models import StatisticData, StatisticMetaData
from .util import periodic_db_cleanups from .util import periodic_db_cleanups
@ -147,13 +148,18 @@ class ImportStatisticsTask(RecorderTask):
metadata: StatisticMetaData metadata: StatisticMetaData
statistics: Iterable[StatisticData] statistics: Iterable[StatisticData]
table: type[Statistics | StatisticsShortTerm]
def run(self, instance: Recorder) -> None: def run(self, instance: Recorder) -> None:
"""Run statistics task.""" """Run statistics task."""
if statistics.import_statistics(instance, self.metadata, self.statistics): if statistics.import_statistics(
instance, self.metadata, self.statistics, self.table
):
return return
# Schedule a new statistics task if this one didn't finish # 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 @dataclass