Add Estimated Database Size to the recorder system health

homeassistant/components/recorder/const.py

@@ -4,6 +4,7 @@
 import json
 from typing import Final
 
+from homeassistant.backports.enum import StrEnum
 from homeassistant.const import ATTR_ATTRIBUTION, ATTR_RESTORED, ATTR_SUPPORTED_FEATURES
 from homeassistant.helpers.json import JSONEncoder
 
@@ -37,3 +38,11 @@
 
 
 EXCLUDE_ATTRIBUTES = f"{DOMAIN}_exclude_attributes_by_domain"
+
+
+class SupportedDialect(StrEnum):
+    """Supported dialects."""
+
+    SQLITE = "sqlite"
+    MYSQL = "mysql"
+    POSTGRESQL = "postgresql"

homeassistant/components/recorder/core.py

@@ -3,6 +3,7 @@
 
 import asyncio
 from collections.abc import Callable, Iterable
+import contextlib
 from datetime import datetime, timedelta
 import logging
 import queue
@@ -41,6 +42,7 @@
     KEEPALIVE_TIME,
     MAX_QUEUE_BACKLOG,
     SQLITE_URL_PREFIX,
+    SupportedDialect,
 )
 from .executor import DBInterruptibleThreadPoolExecutor
 from .models import (
@@ -193,6 +195,13 @@ def backlog(self) -> int:
         """Return the number of items in the recorder backlog."""
         return self._queue.qsize()
 
+    @property
+    def dialect_name(self) -> SupportedDialect | None:
+        """Return the dialect the recorder uses."""
+        with contextlib.suppress(ValueError):
+            return SupportedDialect(self.engine.dialect.name) if self.engine else None
+        return None
+
     @property
     def _using_file_sqlite(self) -> bool:
         """Short version to check if we are using sqlite3 as a file."""
@@ -459,11 +468,6 @@ def async_external_statistics(
         """Schedule external statistics."""
         self.queue_task(ExternalStatisticsTask(metadata, stats))
 
-    @callback
-    def using_sqlite(self) -> bool:
-        """Return if recorder uses sqlite as the engine."""
-        return bool(self.engine and self.engine.dialect.name == "sqlite")
-
     @callback
     def _async_setup_periodic_tasks(self) -> None:
         """Prepare periodic tasks."""
@@ -473,7 +477,7 @@ def _async_setup_periodic_tasks(self) -> None:
 
         # If the db is using a socket connection, we need to keep alive
         # to prevent errors from unexpected disconnects
-        if not self.using_sqlite():
+        if self.dialect_name != SupportedDialect.SQLITE:
             self._keep_alive_listener = async_track_time_interval(
                 self.hass, self._async_keep_alive, timedelta(seconds=KEEPALIVE_TIME)
             )
@@ -939,7 +943,7 @@ def block_till_done(self) -> None:
 
     async def lock_database(self) -> bool:
         """Lock database so it can be backed up safely."""
-        if not self.using_sqlite():
+        if self.dialect_name != SupportedDialect.SQLITE:
             _LOGGER.debug(
                 "Not a SQLite database or not connected, locking not necessary"
             )
@@ -968,7 +972,7 @@ def unlock_database(self) -> bool:
 
         Returns true if database lock has been held throughout the process.
         """
-        if not self.using_sqlite():
+        if self.dialect_name != SupportedDialect.SQLITE:
             _LOGGER.debug(
                 "Not a SQLite database or not connected, unlocking not necessary"
             )

homeassistant/components/recorder/migration.py

@@ -21,6 +21,7 @@
 
 from homeassistant.core import HomeAssistant
 
+from .const import SupportedDialect
 from .models import (
     SCHEMA_VERSION,
     TABLE_STATES,
@@ -263,7 +264,7 @@ def _modify_columns(
     columns_def: list[str],
 ) -> None:
     """Modify columns in a table."""
-    if engine.dialect.name == "sqlite":
+    if engine.dialect.name == SupportedDialect.SQLITE:
         _LOGGER.debug(
             "Skipping to modify columns %s in table %s; "
             "Modifying column length in SQLite is unnecessary, "
@@ -281,7 +282,7 @@ def _modify_columns(
         table_name,
     )
 
-    if engine.dialect.name == "postgresql":
+    if engine.dialect.name == SupportedDialect.POSTGRESQL:
         columns_def = [
             "ALTER {column} TYPE {type}".format(
                 **dict(zip(["column", "type"], col_def.split(" ", 1)))
@@ -408,7 +409,7 @@ def _apply_update(  # noqa: C901
 ) -> None:
     """Perform operations to bring schema up to date."""
     dialect = engine.dialect.name
-    big_int = "INTEGER(20)" if dialect == "mysql" else "INTEGER"
+    big_int = "INTEGER(20)" if dialect == SupportedDialect.MYSQL else "INTEGER"
 
     if new_version == 1:
         _create_index(session_maker, "events", "ix_events_time_fired")
@@ -487,11 +488,11 @@ def _apply_update(  # noqa: C901
         _create_index(session_maker, "states", "ix_states_old_state_id")
         _update_states_table_with_foreign_key_options(session_maker, engine)
     elif new_version == 12:
-        if engine.dialect.name == "mysql":
+        if engine.dialect.name == SupportedDialect.MYSQL:
             _modify_columns(session_maker, engine, "events", ["event_data LONGTEXT"])
             _modify_columns(session_maker, engine, "states", ["attributes LONGTEXT"])
     elif new_version == 13:
-        if engine.dialect.name == "mysql":
+        if engine.dialect.name == SupportedDialect.MYSQL:
             _modify_columns(
                 session_maker,
                 engine,
@@ -545,7 +546,7 @@ def _apply_update(  # noqa: C901
             session.add(StatisticsRuns(start=get_start_time()))
     elif new_version == 20:
         # This changed the precision of statistics from float to double
-        if engine.dialect.name in ["mysql", "postgresql"]:
+        if engine.dialect.name in [SupportedDialect.MYSQL, SupportedDialect.POSTGRESQL]:
             _modify_columns(
                 session_maker,
                 engine,
@@ -560,7 +561,7 @@ def _apply_update(  # noqa: C901
             )
     elif new_version == 21:
         # Try to change the character set of the statistic_meta table
-        if engine.dialect.name == "mysql":
+        if engine.dialect.name == SupportedDialect.MYSQL:
             for table in ("events", "states", "statistics_meta"):
                 _LOGGER.warning(
                     "Updating character set and collation of table %s to utf8mb4. "

homeassistant/components/recorder/purge.py

@@ -13,7 +13,7 @@
 
 from homeassistant.const import EVENT_STATE_CHANGED
 
-from .const import MAX_ROWS_TO_PURGE
+from .const import MAX_ROWS_TO_PURGE, SupportedDialect
 from .models import (
     EventData,
     Events,
@@ -45,7 +45,7 @@ def purge_old_data(
         "Purging states and events before target %s",
         purge_before.isoformat(sep=" ", timespec="seconds"),
     )
-    using_sqlite = instance.using_sqlite()
+    using_sqlite = instance.dialect_name == SupportedDialect.SQLITE
 
     with session_scope(session=instance.get_session()) as session:
         # Purge a max of MAX_ROWS_TO_PURGE, based on the oldest states or events record
@@ -425,7 +425,7 @@ def _purge_old_recorder_runs(
 def _purge_filtered_data(instance: Recorder, session: Session) -> bool:
     """Remove filtered states and events that shouldn't be in the database."""
     _LOGGER.debug("Cleanup filtered data")
-    using_sqlite = instance.using_sqlite()
+    using_sqlite = instance.dialect_name == SupportedDialect.SQLITE
 
     # Check if excluded entity_ids are in database
     excluded_entity_ids: list[str] = [
@@ -484,7 +484,7 @@ def _purge_filtered_events(
     instance: Recorder, session: Session, excluded_event_types: list[str]
 ) -> None:
     """Remove filtered events and linked states."""
-    using_sqlite = instance.using_sqlite()
+    using_sqlite = instance.dialect_name == SupportedDialect.SQLITE
     event_ids, data_ids = zip(
         *(
             session.query(Events.event_id, Events.data_id)
@@ -514,7 +514,7 @@ def _purge_filtered_events(
 @retryable_database_job("purge")
 def purge_entity_data(instance: Recorder, entity_filter: Callable[[str], bool]) -> bool:
     """Purge states and events of specified entities."""
-    using_sqlite = instance.using_sqlite()
+    using_sqlite = instance.dialect_name == SupportedDialect.SQLITE
     with session_scope(session=instance.get_session()) as session:
         selected_entity_ids: list[str] = [
             entity_id

homeassistant/components/recorder/repack.py

@@ -6,6 +6,8 @@
 
 from sqlalchemy import text
 
+from .const import SupportedDialect
+
 if TYPE_CHECKING:
     from . import Recorder
 
@@ -18,15 +20,15 @@ def repack_database(instance: Recorder) -> None:
     dialect_name = instance.engine.dialect.name
 
     # Execute sqlite command to free up space on disk
-    if dialect_name == "sqlite":
+    if dialect_name == SupportedDialect.SQLITE:
         _LOGGER.debug("Vacuuming SQL DB to free space")
         with instance.engine.connect() as conn:
             conn.execute(text("VACUUM"))
             conn.commit()
         return
 
     # Execute postgresql vacuum command to free up space on disk
-    if dialect_name == "postgresql":
+    if dialect_name == SupportedDialect.POSTGRESQL:
         _LOGGER.debug("Vacuuming SQL DB to free space")
         with instance.engine.connect().execution_options(
             isolation_level="AUTOCOMMIT"
@@ -36,7 +38,7 @@ def repack_database(instance: Recorder) -> None:
         return
 
     # Optimize mysql / mariadb tables to free up space on disk
-    if dialect_name == "mysql":
+    if dialect_name == SupportedDialect.MYSQL:
         _LOGGER.debug("Optimizing SQL DB to free space")
         with instance.engine.connect() as conn:
             conn.execute(text("OPTIMIZE TABLE states, events, recorder_runs"))

homeassistant/components/recorder/statistics.py

@@ -38,7 +38,7 @@
 from homeassistant.util.unit_system import UnitSystem
 import homeassistant.util.volume as volume_util
 
-from .const import DATA_INSTANCE, DOMAIN, MAX_ROWS_TO_PURGE
+from .const import DATA_INSTANCE, DOMAIN, MAX_ROWS_TO_PURGE, SupportedDialect
 from .models import (
     StatisticData,
     StatisticMetaData,
@@ -1342,10 +1342,13 @@ def _filter_unique_constraint_integrity_error(err: Exception) -> bool:
         dialect_name = instance.engine.dialect.name
 
         ignore = False
-        if dialect_name == "sqlite" and "UNIQUE constraint failed" in str(err):
+        if (
+            dialect_name == SupportedDialect.SQLITE
+            and "UNIQUE constraint failed" in str(err)
+        ):
             ignore = True
         if (
-            dialect_name == "postgresql"
+            dialect_name == SupportedDialect.POSTGRESQL
             and hasattr(err.orig, "pgcode")
             and err.orig.pgcode == "23505"
         ):

homeassistant/components/recorder/strings.json

@@ -2,7 +2,8 @@
   "system_health": {
     "info": {
       "oldest_recorder_run": "Oldest Run Start Time",
-      "current_recorder_run": "Current Run Start Time"
+      "current_recorder_run": "Current Run Start Time",
+      "estimated_db_size": "Estimated Database Size (MiB)"
     }
   }
 }

homeassistant/components/recorder/system_health/__init__.py

@@ -0,0 +1,60 @@
+"""Provide info to system health."""
+from __future__ import annotations
+
+from typing import Any
+
+from yarl import URL
+
+from homeassistant.components import system_health
+from homeassistant.components.recorder.core import Recorder
+from homeassistant.components.recorder.util import session_scope
+from homeassistant.core import HomeAssistant, callback
+
+from .. import get_instance
+from ..const import SupportedDialect
+from .mysql import db_size_bytes as mysql_db_size_bytes
+from .postgresql import db_size_bytes as postgresql_db_size_bytes
+from .sqlite import db_size_bytes as sqlite_db_size_bytes
+
+DIALECT_TO_GET_SIZE = {
+    SupportedDialect.SQLITE: sqlite_db_size_bytes,
+    SupportedDialect.MYSQL: mysql_db_size_bytes,
+    SupportedDialect.POSTGRESQL: postgresql_db_size_bytes,
+}
+
+
+@callback
+def async_register(
+    hass: HomeAssistant, register: system_health.SystemHealthRegistration
+) -> None:
+    """Register system health callbacks."""
+    register.async_register_info(system_health_info)
+
+
+def _get_db_stats(instance: Recorder, database_name: str) -> dict[str, Any]:
+    """Get the stats about the database."""
+    db_stats: dict[str, Any] = {}
+    with session_scope(session=instance.get_session()) as session:
+        if (
+            (dialect_name := instance.dialect_name)
+            and (get_size := DIALECT_TO_GET_SIZE.get(dialect_name))
+            and (db_bytes := get_size(session, database_name))
+        ):
+            db_stats["estimated_db_size"] = f"{db_bytes/1024/1024:.2f} MiB"
+    return db_stats
+
+
+async def system_health_info(hass: HomeAssistant) -> dict[str, Any]:
+    """Get info for the info page."""
+    instance = get_instance(hass)
+    run_history = instance.run_history
+    database_name = URL(instance.db_url).path.lstrip("/")
+    db_stats: dict[str, Any] = {}
+    if instance.async_db_ready.done():
+        db_stats = await instance.async_add_executor_job(
+            _get_db_stats, instance, database_name
+        )
+    return {
+        "oldest_recorder_run": run_history.first.start,
+        "current_recorder_run": run_history.current.start,
+    } | db_stats

homeassistant/components/recorder/system_health/mysql.py

@@ -0,0 +1,19 @@
+"""Provide info to system health for mysql."""
+from __future__ import annotations
+
+from sqlalchemy import text
+from sqlalchemy.orm.session import Session
+
+
+def db_size_bytes(session: Session, database_name: str) -> float:
+    """Get the mysql database size."""
+    return float(
+        session.execute(
+            text(
+                "SELECT ROUND(SUM(DATA_LENGTH + INDEX_LENGTH), 2) "
+                "FROM information_schema.TABLES WHERE "
+                "TABLE_SCHEMA=:database_name"
+            ),
+            {"database_name": database_name},
+        ).first()[0]
+    )

homeassistant/components/recorder/system_health/postgresql.py

@@ -0,0 +1,15 @@
+"""Provide info to system health for postgresql."""
+from __future__ import annotations
+
+from sqlalchemy import text
+from sqlalchemy.orm.session import Session
+
+
+def db_size_bytes(session: Session, database_name: str) -> float:
+    """Get the mysql database size."""
+    return float(
+        session.execute(
+            text("select pg_database_size(:database_name);"),
+            {"database_name": database_name},
+        ).first()[0]
+    )