added support for home assistant postgres database and convenience function to get state of entity

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",
+]

utils/database/dataframes/__init__.py

@@ -0,0 +1,7 @@
+"""Definition of dataframes."""
+
+from .sensor_state import SensorStateSchema
+
+__all__ = [
+    "SensorStateSchema",
+]

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]

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

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)

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",
+]

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()

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)

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)

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
+    )

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)

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)

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)

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)

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)

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)

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)

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)

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)

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)