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]