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

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)