first commit

2026-06-19 13:38:17 +08:00
commit f38814e9bd
11 changed files with 915 additions and 0 deletions
@@ -0,0 +1,125 @@
+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]