from __future__ import annotations from datetime import date, datetime, time, timedelta, timezone from pathlib import Path from typing import Any from fastapi import APIRouter, HTTPException, Query from fastapi.responses import Response from sqlmodel import col, delete, func, select from app.api.deps import CurrentUser, SessionDep from app.models import ( DeleteResult, MqttRawDaysPublic, MqttRawDaySummary, TerminalMessageRaw, TerminalRawArchiveManifest, ) router = APIRouter(prefix="/mqtt-raw", tags=["mqtt_raw"]) UTC_PLUS_8 = timezone(timedelta(hours=8)) def _normalize_day(value: Any) -> date: if isinstance(value, date) and not isinstance(value, datetime): return value if isinstance(value, datetime): return value.date() if isinstance(value, str): return date.fromisoformat(value) raise ValueError(f"Unsupported day value: {value!r}") def _day_bounds(day: date) -> tuple[datetime, datetime]: start_local = datetime.combine(day, time.min, tzinfo=UTC_PLUS_8) end_local = start_local + timedelta(days=1) return start_local.astimezone(timezone.utc), end_local.astimezone(timezone.utc) def _local_day_expr(session: SessionDep, field: Any) -> Any: bind = session.get_bind() dialect_name = bind.dialect.name if bind is not None else "" if dialect_name == "sqlite": return func.date(func.datetime(field, "+8 hours")) if dialect_name in {"postgresql", "postgres"}: return func.date(func.timezone("Asia/Shanghai", field)) return func.date(field) def _build_window( *, day: date, start_seconds: int | None, end_seconds: int | None ) -> tuple[datetime, datetime, int, int]: start_of_day_local = datetime.combine(day, time.min, tzinfo=UTC_PLUS_8) start_sec = 0 if start_seconds is None else start_seconds end_sec = 24 * 3600 if end_seconds is None else end_seconds if not 0 <= start_sec <= 24 * 3600: raise HTTPException(status_code=400, detail="start_seconds must be 0..86400") if not 0 <= end_sec <= 24 * 3600: raise HTTPException(status_code=400, detail="end_seconds must be 0..86400") if start_sec >= end_sec: raise HTTPException( status_code=400, detail="start_seconds must be smaller than end_seconds" ) start = (start_of_day_local + timedelta(seconds=start_sec)).astimezone(timezone.utc) end = (start_of_day_local + timedelta(seconds=end_sec)).astimezone(timezone.utc) return start, end, start_sec, end_sec def _unlink_file(path_str: str) -> None: path = Path(path_str) try: path.unlink(missing_ok=True) except OSError: return @router.get( "/days", response_model=MqttRawDaysPublic, operation_id="mqtt_raw_list_days", ) def list_mqtt_raw_days( session: SessionDep, _current_user: CurrentUser, skip: int = Query(default=0, ge=0), limit: int = Query(default=30, ge=1, le=200), ) -> MqttRawDaysPublic: days: dict[date, MqttRawDaySummary] = {} raw_day_expr = _local_day_expr(session, TerminalMessageRaw.received_at) raw_statement = ( select( raw_day_expr, func.count(), func.min(TerminalMessageRaw.received_at), func.max(TerminalMessageRaw.received_at), ) .group_by(raw_day_expr) .order_by(raw_day_expr.desc()) ) for day_value, message_count, first_at, last_at in session.exec( raw_statement ).all(): day_key = _normalize_day(day_value) days[day_key] = MqttRawDaySummary( day=day_key, message_count=int(message_count or 0), first_received_at=first_at, last_received_at=last_at, ) archive_day_expr = _local_day_expr(session, TerminalRawArchiveManifest.range_start) archive_statement = ( select( archive_day_expr, func.sum(TerminalRawArchiveManifest.record_count), func.min(TerminalRawArchiveManifest.range_start), func.max(TerminalRawArchiveManifest.range_end), ) .group_by(archive_day_expr) .order_by(archive_day_expr.desc()) ) for day_value, message_count, first_at, last_at in session.exec( archive_statement ).all(): day_key = _normalize_day(day_value) existing = days.get(day_key) archive_count = int(message_count or 0) if existing is None: days[day_key] = MqttRawDaySummary( day=day_key, message_count=archive_count, first_received_at=first_at, last_received_at=last_at, ) continue existing.message_count += archive_count if first_at and ( existing.first_received_at is None or first_at < existing.first_received_at ): existing.first_received_at = first_at if last_at and ( existing.last_received_at is None or last_at > existing.last_received_at ): existing.last_received_at = last_at ordered = sorted(days.values(), key=lambda item: item.day, reverse=True) paged = ordered[skip : skip + limit] return MqttRawDaysPublic(data=paged, count=len(ordered)) @router.get( "/download", operation_id="mqtt_raw_download", responses={ 200: { "content": { "application/octet-stream": { "schema": {"type": "string", "format": "binary"} } } } }, ) def download_mqtt_raw( session: SessionDep, _current_user: CurrentUser, day: date, start_seconds: int | None = Query(default=None), end_seconds: int | None = Query(default=None), ) -> Response: range_start, range_end, start_sec, end_sec = _build_window( day=day, start_seconds=start_seconds, end_seconds=end_seconds, ) chunks: list[tuple[datetime, bytes]] = [] raw_statement = ( select(TerminalMessageRaw) .where(TerminalMessageRaw.received_at >= range_start) .where(TerminalMessageRaw.received_at < range_end) .order_by(col(TerminalMessageRaw.received_at), col(TerminalMessageRaw.id)) ) raw_rows = session.exec(raw_statement).all() for row in raw_rows: chunks.append((row.received_at, row.payload_bytes)) manifest_statement = ( select(TerminalRawArchiveManifest) .where(TerminalRawArchiveManifest.range_start >= range_start) .where(TerminalRawArchiveManifest.range_end <= range_end) .order_by( col(TerminalRawArchiveManifest.range_start), col(TerminalRawArchiveManifest.range_end), ) ) manifests = session.exec(manifest_statement).all() for manifest in manifests: path = Path(manifest.file_path) if not path.exists() or not path.is_file(): continue chunks.append((manifest.range_start, path.read_bytes())) chunks.sort(key=lambda item: item[0]) payload = b"".join(chunk for _, chunk in chunks) name_suffix = ( day.isoformat() if start_sec == 0 and end_sec == 24 * 3600 else f"{day.isoformat()}_{start_sec}_{end_sec}" ) filename = f"mqtt_raw_{name_suffix}.txt" return Response( content=payload, media_type="application/octet-stream", headers={"Content-Disposition": f'attachment; filename="{filename}"'}, ) @router.delete( "/day/{day}", response_model=DeleteResult, operation_id="mqtt_raw_delete_day", ) def delete_mqtt_raw_by_day( session: SessionDep, _current_user: CurrentUser, day: date, ) -> DeleteResult: range_start, range_end = _day_bounds(day) raw_statement = ( select(TerminalMessageRaw.id) .where(TerminalMessageRaw.received_at >= range_start) .where(TerminalMessageRaw.received_at < range_end) ) raw_ids = list(session.exec(raw_statement).all()) if raw_ids: session.exec( delete(TerminalMessageRaw).where(col(TerminalMessageRaw.id).in_(raw_ids)) ) manifest_statement = ( select(TerminalRawArchiveManifest) .where(TerminalRawArchiveManifest.range_start >= range_start) .where(TerminalRawArchiveManifest.range_end <= range_end) ) manifests = session.exec(manifest_statement).all() for manifest in manifests: _unlink_file(manifest.file_path) session.delete(manifest) session.commit() deleted_count = len(raw_ids) + len(manifests) return DeleteResult( deleted_count=deleted_count, message=f"Deleted {deleted_count} record(s) for {day.isoformat()}", ) @router.delete( "/all", response_model=DeleteResult, operation_id="mqtt_raw_delete_all", ) def delete_all_mqtt_raw( session: SessionDep, _current_user: CurrentUser, confirm: bool = Query(default=False), ) -> DeleteResult: if not confirm: raise HTTPException( status_code=400, detail="Set confirm=true to delete all saved MQTT raw data", ) raw_ids = list(session.exec(select(TerminalMessageRaw.id)).all()) if raw_ids: session.exec( delete(TerminalMessageRaw).where(col(TerminalMessageRaw.id).in_(raw_ids)) ) manifests = session.exec(select(TerminalRawArchiveManifest)).all() for manifest in manifests: _unlink_file(manifest.file_path) session.delete(manifest) session.commit() deleted_count = len(raw_ids) + len(manifests) return DeleteResult( deleted_count=deleted_count, message=f"Deleted {deleted_count} record(s) from MQTT raw storage", )