diff --git a/utils/database/__init__.py b/utils/database/__init__.py new file mode 100644 index 0000000..608a336 --- /dev/null +++ b/utils/database/__init__.py @@ -0,0 +1,9 @@ +"""Shared utilities for Home Assistant data analysis.""" + +from .get_db_session import get_db_session +from .get_state import get_state + +__all__ = [ + "get_db_session", + "get_state", +] diff --git a/utils/database/dataframes/__init__.py b/utils/database/dataframes/__init__.py new file mode 100644 index 0000000..0310ef4 --- /dev/null +++ b/utils/database/dataframes/__init__.py @@ -0,0 +1,7 @@ +"""Definition of dataframes.""" + +from .sensor_state import SensorStateSchema + +__all__ = [ + "SensorStateSchema", +] diff --git a/utils/database/dataframes/sensor_state.py b/utils/database/dataframes/sensor_state.py new file mode 100644 index 0000000..bdd8018 --- /dev/null +++ b/utils/database/dataframes/sensor_state.py @@ -0,0 +1,12 @@ +"""Definition of SensorStateSchema dataframe schema.""" + +from typing import Optional +import pandera.pandas as pa +import pandas as pd + + +class SensorStateSchema(pa.DataFrameModel): + """Schema for sensor state dataframe.""" + + time: pd.DatetimeTZDtype = pa.Field(dtype_kwargs={"tz": "UTC"}) + state: Optional[str] diff --git a/utils/database/get_db_session.py b/utils/database/get_db_session.py new file mode 100644 index 0000000..9b7e2fc --- /dev/null +++ b/utils/database/get_db_session.py @@ -0,0 +1,61 @@ +"""Database connection utilities for Home Assistant PostgreSQL database.""" + +import os +from functools import lru_cache + +from dotenv import load_dotenv +from sqlalchemy import create_engine, text +from sqlalchemy.orm import sessionmaker + +# Load environment variables from .env file +load_dotenv() + + +@lru_cache(maxsize=1) +def _get_engine(echo: bool = False): + """Create and cache the database engine (singleton).""" + required_vars = [ + "POSTGRES_HOST", + "POSTGRES_PORT", + "POSTGRES_DB", + "POSTGRES_USER", + "POSTGRES_PASSWORD", + ] + missing_vars = [var for var in required_vars if not os.getenv(var)] + + if missing_vars: + raise ValueError(f"Missing required environment variables: {', '.join(missing_vars)}") + + connection_string = ( + f"postgresql://{os.getenv('POSTGRES_USER')}:{os.getenv('POSTGRES_PASSWORD')}" + f"@{os.getenv('POSTGRES_HOST')}:{os.getenv('POSTGRES_PORT')}/{os.getenv('POSTGRES_DB')}" + ) + + return create_engine(connection_string, echo=echo) + + +def get_db_session(echo: bool = False) -> sessionmaker: + """ + Create a SQLAlchemy sessionmaker for the Home Assistant PostgreSQL database. + + Args: + echo: If True, log all SQL statements (default: False) + + Returns: + sessionmaker: Database session maker object + + Raises: + ValueError: If required environment variables are not set + """ + + # Create the SQLAlchemy engine + engine = _get_engine(echo=echo) + + # Create a session maker + session = sessionmaker(bind=engine) + + # Test connection to database + with session.begin() as s: + s.execute(text("SELECT 1")) + + return session diff --git a/utils/database/get_state.py b/utils/database/get_state.py new file mode 100644 index 0000000..8ed393d --- /dev/null +++ b/utils/database/get_state.py @@ -0,0 +1,89 @@ +"""Definition of get_state function.""" + +from datetime import datetime, UTC +from pandera.typing import DataFrame +import pandas as pd +import pandera.pandas as pa +import sqlalchemy as sa + +from utils.database.tables import ( + States, + StatesMeta, +) +from utils.database.dataframes import SensorStateSchema +from utils.database import get_db_session + + +@pa.check_types +def get_state( + entity_id: str, + limit: int | None = None, + start_time: datetime | None = None, + end_time: datetime | None = None, +) -> DataFrame[SensorStateSchema]: + """ + Retrieve sensor state data. + + Args: + entity_id: The entity ID of the sensor. + limit: Number of recent records to retrieve (default is None for all records). + start_time: Start of time range (inclusive). If None, no lower bound. + end_time: End of time range (inclusive). If None, no upper bound. + + Returns: + DataFrame[SensorStateSchema]: The retrieved sensor state. + """ + + # Create a new database session + session = get_db_session() + + # Prepare the base statement + stmt = ( + sa.select( + States.state, + sa.func.to_timestamp(States.last_updated_ts).label("time"), + ) + .join( + StatesMeta, + States.metadata_id == StatesMeta.metadata_id, + ) + .where(StatesMeta.entity_id == entity_id) + ) + + # Add time range filters if provided + if start_time is not None: + start_ts = start_time.astimezone(UTC).timestamp() + stmt = stmt.where(States.last_updated_ts >= start_ts) + + if end_time is not None: + end_ts = end_time.astimezone(UTC).timestamp() + stmt = stmt.where(States.last_updated_ts <= end_ts) + + # Order by time descending and apply limit if specified + stmt = stmt.order_by(States.last_updated_ts.desc()) + + if limit is not None: + stmt = stmt.limit(limit) + + # Execute query + with session.begin() as s: + df = pd.read_sql(stmt, s.connection()) + + # Handle empty results + if len(df) == 0: + # Create empty DataFrame with correct schema + empty_df = pd.DataFrame( + { + "time": pd.Series([], dtype="datetime64[ns, UTC]"), + "state": pd.Series([], dtype="object"), + } + ) + return DataFrame[SensorStateSchema](empty_df) + + # Reorder columns to match schema definition + df = df[["time", "state"]] + + # Convert time column to nanosecond precision (PostgreSQL returns microseconds) + df["time"] = df["time"].dt.as_unit("ns") + + return DataFrame[SensorStateSchema](df) diff --git a/utils/database/tables/__init__.py b/utils/database/tables/__init__.py new file mode 100644 index 0000000..56d9441 --- /dev/null +++ b/utils/database/tables/__init__.py @@ -0,0 +1,31 @@ +"""Database table definitions for Home Assistant PostgreSQL database.""" + +from .event_data import EventData +from .event_types import EventTypes +from .events import Events +from .migration_changes import MigrationChanges +from .recorder_runs import RecorderRuns +from .schema_changes import SchemaChanges +from .state_attributes import StateAttributes +from .states_meta import StatesMeta +from .states import States +from .statistics_meta import StatisticsMeta +from .statistics_runs import StatisticsRuns +from .statistics_short_term import StatisticsShortTerm +from .statistics import Statistics + +__all__ = [ + "EventData", + "EventTypes", + "Events", + "MigrationChanges", + "RecorderRuns", + "SchemaChanges", + "StateAttributes", + "StatesMeta", + "States", + "StatisticsMeta", + "StatisticsRuns", + "StatisticsShortTerm", + "Statistics", +] diff --git a/utils/database/tables/base.py b/utils/database/tables/base.py new file mode 100644 index 0000000..0b9d5bb --- /dev/null +++ b/utils/database/tables/base.py @@ -0,0 +1,5 @@ +"""Base declarative class for SQLAlchemy ORM models.""" + +from sqlalchemy.orm import declarative_base + +Base = declarative_base() diff --git a/utils/database/tables/event_data.py b/utils/database/tables/event_data.py new file mode 100644 index 0000000..90bcd20 --- /dev/null +++ b/utils/database/tables/event_data.py @@ -0,0 +1,16 @@ +"""SQLAlchemy model for the 'event_data' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, Text + +from .base import Base + + +class EventData(Base): + """SQLAlchemy model for the 'event_data' table.""" + + __tablename__ = "event_data" + __table_args__ = {"schema": "public"} + + data_id = Column(BigInteger, primary_key=True, autoincrement=True) + hash = Column(BigInteger, nullable=True) + shared_data = Column(Text, nullable=True) diff --git a/utils/database/tables/event_types.py b/utils/database/tables/event_types.py new file mode 100644 index 0000000..40d1cf7 --- /dev/null +++ b/utils/database/tables/event_types.py @@ -0,0 +1,15 @@ +"""SQLAlchemy model for the 'event_types' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, String + +from .base import Base + + +class EventTypes(Base): + """SQLAlchemy model for the 'event_types' table.""" + + __tablename__ = "event_types" + __table_args__ = {"schema": "public"} + + event_type_id = Column(BigInteger, primary_key=True, autoincrement=True) + event_type = Column(String(64), nullable=True) diff --git a/utils/database/tables/events.py b/utils/database/tables/events.py new file mode 100644 index 0000000..9a0706c --- /dev/null +++ b/utils/database/tables/events.py @@ -0,0 +1,39 @@ +"""SQLAlchemy model for the 'events' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import ( + Column, + BigInteger, + CHAR, + SmallInteger, + TIMESTAMP, + Float, + LargeBinary, + ForeignKey, +) + +from .base import Base + + +class Events(Base): + """SQLAlchemy model for the 'events' table.""" + + __tablename__ = "events" + __table_args__ = {"schema": "public"} + + event_id = Column(BigInteger, primary_key=True, autoincrement=True) + event_type = Column(CHAR(1), nullable=True) + event_data = Column(CHAR(1), nullable=True) + origin = Column(CHAR(1), nullable=True) + origin_idx = Column(SmallInteger, nullable=True) + time_fired = Column(TIMESTAMP(timezone=True), nullable=True) + time_fired_ts = Column(Float, nullable=True) + context_id = Column(CHAR(1), nullable=True) + context_user_id = Column(CHAR(1), nullable=True) + context_parent_id = Column(CHAR(1), nullable=True) + data_id = Column(BigInteger, ForeignKey("public.event_data.data_id"), nullable=True) + context_id_bin = Column(LargeBinary, nullable=True) + context_user_id_bin = Column(LargeBinary, nullable=True) + context_parent_id_bin = Column(LargeBinary, nullable=True) + event_type_id = Column( + BigInteger, ForeignKey("public.event_types.event_type_id"), nullable=True + ) diff --git a/utils/database/tables/migration_changes.py b/utils/database/tables/migration_changes.py new file mode 100644 index 0000000..a8a6e5d --- /dev/null +++ b/utils/database/tables/migration_changes.py @@ -0,0 +1,15 @@ +"""SQLAlchemy model for the 'migration_changes' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, String, SmallInteger + +from .base import Base + + +class MigrationChanges(Base): + """SQLAlchemy model for the 'migration_changes' table.""" + + __tablename__ = "migration_changes" + __table_args__ = {"schema": "public"} + + migration_id = Column(String(255), primary_key=True) + version = Column(SmallInteger, nullable=False) diff --git a/utils/database/tables/recorder_runs.py b/utils/database/tables/recorder_runs.py new file mode 100644 index 0000000..c1af572 --- /dev/null +++ b/utils/database/tables/recorder_runs.py @@ -0,0 +1,18 @@ +"""SQLAlchemy model for the 'recorder_runs' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, TIMESTAMP, Boolean + +from .base import Base + + +class RecorderRuns(Base): + """SQLAlchemy model for the 'recorder_runs' table.""" + + __tablename__ = "recorder_runs" + __table_args__ = {"schema": "public"} + + run_id = Column(BigInteger, primary_key=True, autoincrement=True) + start = Column(TIMESTAMP(timezone=True), nullable=False) + end = Column(TIMESTAMP(timezone=True), nullable=True) + closed_incorrect = Column(Boolean, nullable=False) + created = Column(TIMESTAMP(timezone=True), nullable=False) diff --git a/utils/database/tables/schema_changes.py b/utils/database/tables/schema_changes.py new file mode 100644 index 0000000..854fdb6 --- /dev/null +++ b/utils/database/tables/schema_changes.py @@ -0,0 +1,16 @@ +"""SQLAlchemy model for the 'schema_changes' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, Integer, TIMESTAMP + +from .base import Base + + +class SchemaChanges(Base): + """SQLAlchemy model for the 'schema_changes' table.""" + + __tablename__ = "schema_changes" + __table_args__ = {"schema": "public"} + + change_id = Column(BigInteger, primary_key=True, autoincrement=True) + schema_version = Column(Integer, nullable=True) + changed = Column(TIMESTAMP(timezone=True), nullable=False) diff --git a/utils/database/tables/state_attributes.py b/utils/database/tables/state_attributes.py new file mode 100644 index 0000000..bd79455 --- /dev/null +++ b/utils/database/tables/state_attributes.py @@ -0,0 +1,16 @@ +"""SQLAlchemy model for the 'state_attributes' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, Text + +from .base import Base + + +class StateAttributes(Base): + """SQLAlchemy model for the 'state_attributes' table.""" + + __tablename__ = "state_attributes" + __table_args__ = {"schema": "public"} + + attributes_id = Column(BigInteger, primary_key=True, autoincrement=True) + hash = Column(BigInteger, nullable=True) + shared_attrs = Column(Text, nullable=True) diff --git a/utils/database/tables/states.py b/utils/database/tables/states.py new file mode 100644 index 0000000..f9c0ced --- /dev/null +++ b/utils/database/tables/states.py @@ -0,0 +1,45 @@ +"""SQLAlchemy model for the 'states' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import ( + Column, + BigInteger, + CHAR, + String, + SmallInteger, + TIMESTAMP, + Float, + LargeBinary, + ForeignKey, +) + +from .base import Base + + +class States(Base): + """SQLAlchemy model for the 'states' table.""" + + __tablename__ = "states" + __table_args__ = {"schema": "public"} + + state_id = Column(BigInteger, primary_key=True, autoincrement=True) + entity_id = Column(CHAR(1), nullable=True) + state = Column(String(255), nullable=True) + attributes = Column(CHAR(1), nullable=True) + event_id = Column(SmallInteger, nullable=True) + last_changed = Column(TIMESTAMP(timezone=True), nullable=True) + last_changed_ts = Column(Float, nullable=True) + last_reported_ts = Column(Float, nullable=True) + last_updated = Column(TIMESTAMP(timezone=True), nullable=True) + last_updated_ts = Column(Float, nullable=True) + old_state_id = Column(BigInteger, ForeignKey("public.states.state_id"), nullable=True) + attributes_id = Column( + BigInteger, ForeignKey("public.state_attributes.attributes_id"), nullable=True + ) + context_id = Column(CHAR(1), nullable=True) + context_user_id = Column(CHAR(1), nullable=True) + context_parent_id = Column(CHAR(1), nullable=True) + origin_idx = Column(SmallInteger, nullable=True) + context_id_bin = Column(LargeBinary, nullable=True) + context_user_id_bin = Column(LargeBinary, nullable=True) + context_parent_id_bin = Column(LargeBinary, nullable=True) + metadata_id = Column(BigInteger, ForeignKey("public.states_meta.metadata_id"), nullable=True) diff --git a/utils/database/tables/states_meta.py b/utils/database/tables/states_meta.py new file mode 100644 index 0000000..dbf995b --- /dev/null +++ b/utils/database/tables/states_meta.py @@ -0,0 +1,15 @@ +"""SQLAlchemy model for the 'states_meta' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, String + +from .base import Base + + +class StatesMeta(Base): + """SQLAlchemy model for the 'states_meta' table.""" + + __tablename__ = "states_meta" + __table_args__ = {"schema": "public"} + + metadata_id = Column(BigInteger, primary_key=True, autoincrement=True) + entity_id = Column(String(255), nullable=True) diff --git a/utils/database/tables/statistics.py b/utils/database/tables/statistics.py new file mode 100644 index 0000000..96e6d07 --- /dev/null +++ b/utils/database/tables/statistics.py @@ -0,0 +1,29 @@ +"""SQLAlchemy model for the 'statistics' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, TIMESTAMP, Float, ForeignKey + +from .base import Base + + +class Statistics(Base): + """SQLAlchemy model for the 'statistics' table.""" + + __tablename__ = "statistics" + __table_args__ = {"schema": "public"} + + id = Column(BigInteger, primary_key=True, autoincrement=True) + created = Column(TIMESTAMP(timezone=True), nullable=True) + created_ts = Column(Float, nullable=True) + metadata_id = Column( + BigInteger, ForeignKey("public.statistics_meta.id", ondelete="CASCADE"), nullable=True + ) + start = Column(TIMESTAMP(timezone=True), nullable=True) + start_ts = Column(Float, nullable=True) + mean = Column(Float, nullable=True) + mean_weight = Column(Float, nullable=True) + min = Column(Float, nullable=True) + max = Column(Float, nullable=True) + last_reset = Column(TIMESTAMP(timezone=True), nullable=True) + last_reset_ts = Column(Float, nullable=True) + state = Column(Float, nullable=True) + sum = Column(Float, nullable=True) diff --git a/utils/database/tables/statistics_meta.py b/utils/database/tables/statistics_meta.py new file mode 100644 index 0000000..0f91b24 --- /dev/null +++ b/utils/database/tables/statistics_meta.py @@ -0,0 +1,22 @@ +"""SQLAlchemy model for the 'statistics_meta' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, String, Boolean, SmallInteger + +from .base import Base + + +class StatisticsMeta(Base): + """SQLAlchemy model for the 'statistics_meta' table.""" + + __tablename__ = "statistics_meta" + __table_args__ = {"schema": "public"} + + id = Column(BigInteger, primary_key=True, autoincrement=True) + statistic_id = Column(String(255), nullable=True) + source = Column(String(32), nullable=True) + unit_of_measurement = Column(String(255), nullable=True) + unit_class = Column(String(255), nullable=True) + has_mean = Column(Boolean, nullable=True) + has_sum = Column(Boolean, nullable=True) + name = Column(String(255), nullable=True) + mean_type = Column(SmallInteger, nullable=False) diff --git a/utils/database/tables/statistics_runs.py b/utils/database/tables/statistics_runs.py new file mode 100644 index 0000000..8830179 --- /dev/null +++ b/utils/database/tables/statistics_runs.py @@ -0,0 +1,15 @@ +"""SQLAlchemy model for the 'statistics_runs' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, TIMESTAMP + +from .base import Base + + +class StatisticsRuns(Base): + """SQLAlchemy model for the 'statistics_runs' table.""" + + __tablename__ = "statistics_runs" + __table_args__ = {"schema": "public"} + + run_id = Column(BigInteger, primary_key=True, autoincrement=True) + start = Column(TIMESTAMP(timezone=True), nullable=False) diff --git a/utils/database/tables/statistics_short_term.py b/utils/database/tables/statistics_short_term.py new file mode 100644 index 0000000..6f41a56 --- /dev/null +++ b/utils/database/tables/statistics_short_term.py @@ -0,0 +1,29 @@ +"""SQLAlchemy model for the 'statistics_short_term' table in the Home Assistant PostgreSQL database.""" + +from sqlalchemy import Column, BigInteger, TIMESTAMP, Float, ForeignKey + +from .base import Base + + +class StatisticsShortTerm(Base): + """SQLAlchemy model for the 'statistics_short_term' table.""" + + __tablename__ = "statistics_short_term" + __table_args__ = {"schema": "public"} + + id = Column(BigInteger, primary_key=True, autoincrement=True) + created = Column(TIMESTAMP(timezone=True), nullable=True) + created_ts = Column(Float, nullable=True) + metadata_id = Column( + BigInteger, ForeignKey("public.statistics_meta.id", ondelete="CASCADE"), nullable=True + ) + start = Column(TIMESTAMP(timezone=True), nullable=True) + start_ts = Column(Float, nullable=True) + mean = Column(Float, nullable=True) + mean_weight = Column(Float, nullable=True) + min = Column(Float, nullable=True) + max = Column(Float, nullable=True) + last_reset = Column(TIMESTAMP(timezone=True), nullable=True) + last_reset_ts = Column(Float, nullable=True) + state = Column(Float, nullable=True) + sum = Column(Float, nullable=True)