weather-data-api/app/database.py

from datetime import datetime

import asyncpg
from app.config import settings

_pool: asyncpg.Pool | None = None


async def get_pool() -> asyncpg.Pool:
    """Return the global connection pool, creating it if necessary."""
    global _pool
    if _pool is None:
        _pool = await asyncpg.create_pool(
            settings.database_url,
            min_size=2,
            max_size=10,
        )
    return _pool


async def close_pool() -> None:
    """Close the connection pool."""
    global _pool
    if _pool is not None:
        await _pool.close()
        _pool = None


async def init_db() -> None:
    """Create the weather_data table if it doesn't exist."""
    pool = await get_pool()
    async with pool.acquire() as conn:
        await conn.execute("""
            CREATE TABLE IF NOT EXISTS weather_data (
                id          BIGSERIAL PRIMARY KEY,
                station_name TEXT NOT NULL,
                temperature  DOUBLE PRECISION,
                pressure     DOUBLE PRECISION,
                relative_humidity DOUBLE PRECISION,
                received_at  TIMESTAMPTZ NOT NULL DEFAULT NOW()
            );
        """)
        await conn.execute("""
            CREATE INDEX IF NOT EXISTS idx_weather_station_name
                ON weather_data (station_name, received_at DESC);
        """)


async def insert_weather_data(
    station_name: str,
    temperature: float | None,
    pressure: float | None,
    relative_humidity: float | None,
) -> dict:
    """Insert a weather record and return it."""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            INSERT INTO weather_data (station_name, temperature, pressure, relative_humidity)
            VALUES ($1, $2, $3, $4)
            RETURNING id, station_name, temperature, pressure, relative_humidity, received_at
            """,
            station_name,
            temperature,
            pressure,
            relative_humidity,
        )
        return dict(row)


async def get_latest_by_station(station_name: str) -> dict | None:
    """Return the most recent record for a given station, or None."""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            """
            SELECT id, station_name, temperature, pressure, relative_humidity, received_at
            FROM weather_data
            WHERE station_name = $1
            ORDER BY received_at DESC
            LIMIT 1
            """,
            station_name,
        )
        return dict(row) if row else None


async def get_history_by_station(
    station_name: str,
    start_time: datetime | None = None,
    end_time: datetime | None = None,
) -> list[dict]:
    """Return all records for a station within the given time range.

    Returns an empty list if no records match.
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        conditions = ["station_name = $1"]
        params: list = [station_name]
        idx = 2

        if start_time is not None:
            conditions.append(f"received_at >= ${idx}")
            params.append(start_time)
            idx += 1

        if end_time is not None:
            conditions.append(f"received_at <= ${idx}")
            params.append(end_time)
            idx += 1

        where_clause = " AND ".join(conditions)

        rows = await conn.fetch(
            f"""
            SELECT id, station_name, temperature, pressure, relative_humidity, received_at
            FROM weather_data
            WHERE {where_clause}
            ORDER BY received_at DESC
            """,
            *params,
        )
        return [dict(r) for r in rows]


async def list_all_stations() -> list[dict]:
    """Return all distinct station names with observation counts and latest timestamps.

    Derives stations from weather_data since there is no separate stations table.
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        rows = await conn.fetch("""
            SELECT
                station_name AS name,
                COUNT(*)::int AS observation_count,
                MAX(received_at) AS latest_observation_at
            FROM weather_data
            GROUP BY station_name
            ORDER BY station_name
        """)
        return [dict(r) for r in rows]


async def get_station_exists(station_name: str) -> bool:
    """Check whether a station has any observations."""
    pool = await get_pool()
    async with pool.acquire() as conn:
        row = await conn.fetchrow(
            "SELECT 1 FROM weather_data WHERE station_name = $1 LIMIT 1",
            station_name,
        )
        return row is not None


async def get_history_by_station_paginated(
    station_name: str,
    start_time: datetime | None = None,
    end_time: datetime | None = None,
    limit: int | None = None,
    offset: int = 0,
) -> list[dict]:
    """Return observations for a station within the given time range, with pagination.

    Returns an empty list if no records match.
    """
    pool = await get_pool()
    async with pool.acquire() as conn:
        conditions = ["station_name = $1"]
        params: list = [station_name]
        idx = 2

        if start_time is not None:
            conditions.append(f"received_at >= ${idx}")
            params.append(start_time)
            idx += 1

        if end_time is not None:
            conditions.append(f"received_at <= ${idx}")
            params.append(end_time)
            idx += 1

        where_clause = " AND ".join(conditions)

        query = f"""
            SELECT id, station_name, temperature, pressure, relative_humidity, received_at
            FROM weather_data
            WHERE {where_clause}
            ORDER BY received_at DESC
        """

        if limit is not None:
            query += f" LIMIT ${idx}"
            params.append(limit)
            idx += 1

        query += f" OFFSET ${idx}"
        params.append(offset)

        rows = await conn.fetch(query, *params)
        return [dict(r) for r in rows]