90 lines
2.5 KiB
Python
90 lines
2.5 KiB
Python
"""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)
|