FlightPlot/flight_map/storage.py

import json
import math
import sqlite3
import uuid
from collections.abc import Iterator
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path

from flight_map.geo import Bounds


@dataclass(frozen=True)
class Observation:
    icao24: str
    callsign: str | None
    timestamp: float
    longitude: float
    latitude: float
    baro_altitude: float | None
    geo_altitude: float | None
    on_ground: bool = False

    def __post_init__(self) -> None:
        if len(self.icao24) != 6 or any(c not in "0123456789abcdef" for c in self.icao24):
            raise ValueError("Invalid ICAO24 aircraft identifier.")
        if not all(math.isfinite(v) for v in (self.timestamp, self.longitude, self.latitude)):
            raise ValueError("Observation contains non-finite coordinates or timestamp.")
        if not (-180 <= self.longitude <= 180 and -90 <= self.latitude <= 90):
            raise ValueError("Observation coordinates are outside geographic limits.")
        for altitude in (self.baro_altitude, self.geo_altitude):
            if altitude is not None and not math.isfinite(altitude):
                raise ValueError("Altitude must be finite or missing.")


@dataclass(frozen=True)
class Session:
    id: str
    bounds: Bounds
    collection_bounds: Bounds
    interval: float
    started: float
    ended: float
    status: str
    source: str


@contextmanager
def database(
    path: Path, *, readonly: bool = False, must_exist: bool = False
) -> Iterator[sqlite3.Connection]:
    if readonly or must_exist:
        mode = "ro" if readonly else "rw"
        connection = sqlite3.connect(path.resolve().as_uri() + f"?mode={mode}", uri=True)
    else:
        connection = sqlite3.connect(path)
    connection.row_factory = sqlite3.Row
    connection.execute("PRAGMA foreign_keys = ON")
    try:
        yield connection
    finally:
        connection.close()


def initialize(connection: sqlite3.Connection) -> None:
    version = connection.execute("PRAGMA user_version").fetchone()[0]
    if version not in (0, 1):
        raise ValueError(f"Unsupported flight database version: {version}")
    connection.executescript("""
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY, bounds TEXT NOT NULL, collection_bounds TEXT NOT NULL,
            interval REAL NOT NULL, started REAL NOT NULL, ended REAL NOT NULL,
            status TEXT NOT NULL, source TEXT NOT NULL
        );
        CREATE TABLE IF NOT EXISTS observations (
            session_id TEXT NOT NULL REFERENCES sessions(id),
            icao24 TEXT NOT NULL, callsign TEXT, timestamp REAL NOT NULL,
            longitude REAL NOT NULL, latitude REAL NOT NULL,
            baro_altitude REAL, geo_altitude REAL, on_ground INTEGER NOT NULL,
            PRIMARY KEY (session_id, icao24, timestamp)
        );
        CREATE INDEX IF NOT EXISTS observation_time
            ON observations(session_id, timestamp);
        PRAGMA user_version = 1;
    """)


def create_session(
    connection: sqlite3.Connection, bounds: Bounds, margin: float, interval: float, now: float
) -> Session:
    if not math.isfinite(interval) or interval < 10:
        raise ValueError("Recording interval must be at least 10 seconds.")
    session = Session(
        uuid.uuid4().hex,
        bounds,
        bounds.expand(margin),
        interval,
        now,
        now,
        "recording",
        "OpenSky Network",
    )
    with connection:
        connection.execute(
            "INSERT INTO sessions VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
            (
                session.id,
                json.dumps(bounds.values),
                json.dumps(session.collection_bounds.values),
                interval,
                now,
                now,
                session.status,
                session.source,
            ),
        )
    return session


def get_session(connection: sqlite3.Connection, session_id: str) -> Session:
    row = connection.execute("SELECT * FROM sessions WHERE id = ?", (session_id,)).fetchone()
    if row is None:
        raise ValueError(f"Recording session not found: {session_id}")
    return Session(
        row["id"],
        Bounds(*json.loads(row["bounds"])),
        Bounds(*json.loads(row["collection_bounds"])),
        row["interval"],
        row["started"],
        row["ended"],
        row["status"],
        row["source"],
    )


def checkpoint(connection: sqlite3.Connection, session_id: str, now: float, status: str) -> None:
    with connection:
        connection.execute(
            "UPDATE sessions SET ended = ?, status = ? WHERE id = ?", (now, status, session_id)
        )


def purge_session(connection: sqlite3.Connection, session_id: str) -> int:
    """Delete one session and its observations in a single write transaction."""
    with connection:
        connection.execute("BEGIN IMMEDIATE")
        exists = connection.execute("SELECT 1 FROM sessions WHERE id = ?", (session_id,)).fetchone()
        if exists is None:
            raise ValueError(f"Recording session not found: {session_id}")
        deleted = connection.execute(
            "DELETE FROM observations WHERE session_id = ?", (session_id,)
        ).rowcount
        connection.execute("DELETE FROM sessions WHERE id = ?", (session_id,))
    return deleted


def save_observations(
    connection: sqlite3.Connection, session: Session, observations: list[Observation], now: float
) -> None:
    with connection:
        for item in observations:
            if (
                item.timestamp < session.started
                or item.timestamp > now
                or not session.collection_bounds.contains(item.longitude, item.latitude)
            ):
                continue
            connection.execute(
                "INSERT OR IGNORE INTO observations VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)",
                (
                    session.id,
                    item.icao24,
                    item.callsign,
                    item.timestamp,
                    item.longitude,
                    item.latitude,
                    item.baro_altitude,
                    item.geo_altitude,
                    item.on_ground,
                ),
            )
        connection.execute("UPDATE sessions SET ended = ? WHERE id = ?", (now, session.id))


def load_observations(
    connection: sqlite3.Connection, session_id: str, start: float, end: float
) -> list[Observation]:
    return [
        Observation(
            row["icao24"],
            row["callsign"],
            row["timestamp"],
            row["longitude"],
            row["latitude"],
            row["baro_altitude"],
            row["geo_altitude"],
            bool(row["on_ground"]),
        )
        for row in connection.execute(
            """SELECT * FROM observations WHERE session_id = ? AND timestamp BETWEEN ? AND ?
               ORDER BY icao24, timestamp""",
            (session_id, start, end),
        )
    ]