126 lines
3.6 KiB
Python
126 lines
3.6 KiB
Python
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]
|