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]