140 lines
5.3 KiB
Python
140 lines
5.3 KiB
Python
from datetime import datetime, timezone
|
|
from typing import List, Optional, Union
|
|
from sqlmodel import select, func
|
|
from sqlmodel.ext.asyncio.session import AsyncSession
|
|
from app.models import SleepSession
|
|
|
|
def ensure_aware(dt: datetime) -> datetime:
|
|
if dt.tzinfo is None:
|
|
return dt.replace(tzinfo=timezone.utc)
|
|
return dt
|
|
|
|
async def get_sessions(
|
|
session: AsyncSession,
|
|
from_date: Optional[datetime] = None,
|
|
to_date: Optional[datetime] = None,
|
|
limit: int = 200,
|
|
) -> tuple[List[SleepSession], int]:
|
|
if to_date and to_date.hour == 0 and to_date.minute == 0 and to_date.second == 0 and to_date.microsecond == 0:
|
|
to_date = to_date.replace(hour=23, minute=59, second=59, microsecond=999999)
|
|
|
|
# If we have an aware to_date, but the DB is naive,
|
|
# we should make it naive for comparison to avoid errors in some drivers.
|
|
query_to_date = to_date.replace(tzinfo=None) if to_date and to_date.tzinfo else to_date
|
|
query_from_date = from_date.replace(tzinfo=None) if from_date and from_date.tzinfo else from_date
|
|
|
|
count_statement = select(func.count()).select_from(SleepSession)
|
|
if query_from_date:
|
|
count_statement = count_statement.where(SleepSession.start_time >= query_from_date)
|
|
if query_to_date:
|
|
count_statement = count_statement.where(SleepSession.start_time <= query_to_date)
|
|
count_result = await session.exec(count_statement)
|
|
total = count_result.one() or 0
|
|
|
|
statement = select(SleepSession).order_by(SleepSession.start_time.desc())
|
|
if query_from_date:
|
|
statement = statement.where(SleepSession.start_time >= query_from_date)
|
|
if query_to_date:
|
|
statement = statement.where(SleepSession.start_time <= query_to_date)
|
|
statement = statement.limit(limit)
|
|
|
|
result = await session.exec(statement)
|
|
sessions = result.all()
|
|
return list(sessions), total
|
|
|
|
async def get_active_session(session: AsyncSession) -> Optional[SleepSession]:
|
|
statement = select(SleepSession).where(SleepSession.end_time == None)
|
|
result = await session.exec(statement)
|
|
return result.first()
|
|
|
|
async def create_sleep_session(session: AsyncSession, start_time: datetime) -> SleepSession:
|
|
new_session = SleepSession(start_time=start_time)
|
|
session.add(new_session)
|
|
await session.commit()
|
|
await session.refresh(new_session)
|
|
return new_session
|
|
|
|
async def end_sleep_session(
|
|
session: AsyncSession, session_id: int, end_time: datetime
|
|
) -> Union[SleepSession, bool, str, None]:
|
|
statement = select(SleepSession).where(SleepSession.id == session_id)
|
|
result = await session.exec(statement)
|
|
sleep_session = result.first()
|
|
|
|
if not sleep_session:
|
|
return None # Not found
|
|
|
|
if sleep_session.end_time is not None:
|
|
return False # Already ended
|
|
|
|
# Normalize to aware for comparison
|
|
end_time_aware = end_time.replace(tzinfo=timezone.utc) if end_time.tzinfo is None else end_time
|
|
start_time_aware = sleep_session.start_time.replace(tzinfo=timezone.utc) if sleep_session.start_time.tzinfo is None else sleep_session.start_time
|
|
|
|
if end_time_aware < start_time_aware:
|
|
return "invalid" # End time before start time
|
|
|
|
# To store back to naive DB, we strip tzinfo if it was aware
|
|
sleep_session.end_time = end_time.replace(tzinfo=None) if end_time.tzinfo else end_time
|
|
session.add(sleep_session)
|
|
await session.commit()
|
|
await session.refresh(sleep_session)
|
|
return sleep_session
|
|
|
|
async def delete_session(session: AsyncSession, session_id: int) -> bool:
|
|
statement = select(SleepSession).where(SleepSession.id == session_id)
|
|
result = await session.exec(statement)
|
|
sleep_session = result.first()
|
|
|
|
if not sleep_session:
|
|
return False
|
|
|
|
await session.delete(sleep_session)
|
|
await session.commit()
|
|
return True
|
|
|
|
async def update_session(
|
|
session: AsyncSession, session_id: int, update_data: dict
|
|
) -> Optional[SleepSession]:
|
|
statement = select(SleepSession).where(SleepSession.id == session_id)
|
|
result = await session.exec(statement)
|
|
sleep_session = result.first()
|
|
|
|
if not sleep_session:
|
|
return None
|
|
|
|
for key, value in update_data.items():
|
|
if hasattr(sleep_session, key):
|
|
setattr(sleep_session, key, value)
|
|
|
|
session.add(sleep_session)
|
|
await session.commit()
|
|
await session.refresh(sleep_session)
|
|
return sleep_session
|
|
|
|
|
|
async def get_sleep_counts_by_start_date(
|
|
session: AsyncSession,
|
|
from_date: Optional[datetime] = None,
|
|
to_date: Optional[datetime] = None,
|
|
) -> List[tuple]:
|
|
if to_date and to_date.hour == 0 and to_date.minute == 0 and to_date.second == 0 and to_date.microsecond == 0:
|
|
to_date = to_date.replace(hour=23, minute=59, second=59, microsecond=999999)
|
|
|
|
query_to_date = to_date.replace(tzinfo=None) if to_date and to_date.tzinfo else to_date
|
|
query_from_date = from_date.replace(tzinfo=None) if from_date and from_date.tzinfo else from_date
|
|
|
|
statement = (
|
|
select(func.date(SleepSession.start_time), func.count())
|
|
.group_by(func.date(SleepSession.start_time))
|
|
.order_by(func.date(SleepSession.start_time).desc())
|
|
)
|
|
|
|
if query_from_date:
|
|
statement = statement.where(SleepSession.start_time >= query_from_date)
|
|
if query_to_date:
|
|
statement = statement.where(SleepSession.start_time <= query_to_date)
|
|
|
|
result = await session.exec(statement)
|
|
return list(result.all())
|