initial commit
This commit is contained in:
TheOnlyMace
213
api/app/database.py
Normal file
213
api/app/database.py
Normal file
@@ -0,0 +1,213 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from contextlib import asynccontextmanager
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
import asyncpg
|
||||
|
||||
from app.config import DATABASE_URL
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
_pool: asyncpg.Pool | None = None
|
||||
|
||||
|
||||
async def init_db() -> None:
|
||||
global _pool
|
||||
try:
|
||||
_pool = await asyncpg.create_pool(DATABASE_URL, min_size=1, max_size=8, command_timeout=30)
|
||||
async with _pool.acquire() as conn:
|
||||
await conn.execute("SELECT 1")
|
||||
logger.info("PostgreSQL connected")
|
||||
except Exception as exc:
|
||||
logger.warning("PostgreSQL unavailable, running without DB cache: %s", exc)
|
||||
_pool = None
|
||||
|
||||
|
||||
async def close_db() -> None:
|
||||
global _pool
|
||||
if _pool is not None:
|
||||
await _pool.close()
|
||||
_pool = None
|
||||
|
||||
|
||||
def pool_available() -> bool:
|
||||
return _pool is not None
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def acquire():
|
||||
if _pool is None:
|
||||
yield None
|
||||
return
|
||||
async with _pool.acquire() as conn:
|
||||
yield conn
|
||||
|
||||
|
||||
async def get_cached_json(table: str, key_cols: dict[str, Any], max_age_seconds: int) -> Any | None:
|
||||
async with acquire() as conn:
|
||||
if conn is None:
|
||||
return None
|
||||
if table == "forecast_cache":
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT data, fetched_at FROM forecast_cache
|
||||
WHERE lat = $1 AND lon = $2
|
||||
AND fetched_at > NOW() - ($3 || ' seconds')::interval
|
||||
""",
|
||||
key_cols["lat"],
|
||||
key_cols["lon"],
|
||||
str(max_age_seconds),
|
||||
)
|
||||
elif table == "observations_cache":
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT data, fetched_at FROM observations_cache
|
||||
WHERE lat = $1 AND lon = $2
|
||||
AND fetched_at > NOW() - ($3 || ' seconds')::interval
|
||||
""",
|
||||
key_cols["lat"],
|
||||
key_cols["lon"],
|
||||
str(max_age_seconds),
|
||||
)
|
||||
elif table == "geocode_cache":
|
||||
row = await conn.fetchrow(
|
||||
"""
|
||||
SELECT results, fetched_at FROM geocode_cache
|
||||
WHERE query_hash = $1
|
||||
AND fetched_at > NOW() - ($2 || ' seconds')::interval
|
||||
""",
|
||||
key_cols["query_hash"],
|
||||
str(max_age_seconds),
|
||||
)
|
||||
else:
|
||||
return None
|
||||
if not row:
|
||||
return None
|
||||
data = row["data"] if "data" in row else row["results"]
|
||||
return json.loads(data) if isinstance(data, str) else data
|
||||
|
||||
|
||||
async def set_cached_json(table: str, key_cols: dict[str, Any], payload: Any) -> None:
|
||||
async with acquire() as conn:
|
||||
if conn is None:
|
||||
return
|
||||
payload_json = json.dumps(payload)
|
||||
if table == "forecast_cache":
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO forecast_cache (lat, lon, data, fetched_at)
|
||||
VALUES ($1, $2, $3::jsonb, NOW())
|
||||
ON CONFLICT (lat, lon) DO UPDATE
|
||||
SET data = EXCLUDED.data, fetched_at = NOW()
|
||||
""",
|
||||
key_cols["lat"],
|
||||
key_cols["lon"],
|
||||
payload_json,
|
||||
)
|
||||
elif table == "observations_cache":
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO observations_cache (lat, lon, data, fetched_at)
|
||||
VALUES ($1, $2, $3::jsonb, NOW())
|
||||
ON CONFLICT (lat, lon) DO UPDATE
|
||||
SET data = EXCLUDED.data, fetched_at = NOW()
|
||||
""",
|
||||
key_cols["lat"],
|
||||
key_cols["lon"],
|
||||
payload_json,
|
||||
)
|
||||
elif table == "geocode_cache":
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO geocode_cache (query_hash, query_text, results, fetched_at)
|
||||
VALUES ($1, $2, $3::jsonb, NOW())
|
||||
ON CONFLICT (query_hash) DO UPDATE
|
||||
SET results = EXCLUDED.results, fetched_at = NOW(), query_text = EXCLUDED.query_text
|
||||
""",
|
||||
key_cols["query_hash"],
|
||||
key_cols.get("query_text", ""),
|
||||
payload_json,
|
||||
)
|
||||
elif table == "warnings_cache":
|
||||
await conn.execute(
|
||||
"INSERT INTO warnings_cache (data, fetched_at) VALUES ($1::jsonb, NOW())",
|
||||
payload_json,
|
||||
)
|
||||
|
||||
|
||||
async def upsert_radar_meta(
|
||||
product: str,
|
||||
filename: str,
|
||||
file_size: int | None,
|
||||
modified_at: datetime | None,
|
||||
render_status: str = "pending",
|
||||
render_path: str | None = None,
|
||||
render_bounds: dict | None = None,
|
||||
render_error: str | None = None,
|
||||
) -> None:
|
||||
async with acquire() as conn:
|
||||
if conn is None:
|
||||
return
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO radar_file_meta (
|
||||
product, filename, file_size, modified_at, render_status,
|
||||
render_path, render_bounds, render_error, updated_at
|
||||
) VALUES ($1, $2, $3, $4, $5, $6, $7::jsonb, $8, NOW())
|
||||
ON CONFLICT (filename) DO UPDATE SET
|
||||
file_size = EXCLUDED.file_size,
|
||||
modified_at = EXCLUDED.modified_at,
|
||||
render_status = EXCLUDED.render_status,
|
||||
render_path = EXCLUDED.render_path,
|
||||
render_bounds = EXCLUDED.render_bounds,
|
||||
render_error = EXCLUDED.render_error,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
product,
|
||||
filename,
|
||||
file_size,
|
||||
modified_at,
|
||||
render_status,
|
||||
render_path,
|
||||
json.dumps(render_bounds) if render_bounds else None,
|
||||
render_error,
|
||||
)
|
||||
|
||||
|
||||
async def list_rendered_radar(product: str, limit: int = 24) -> list[dict[str, Any]]:
|
||||
async with acquire() as conn:
|
||||
if conn is None:
|
||||
return []
|
||||
rows = await conn.fetch(
|
||||
"""
|
||||
SELECT filename, render_path, render_bounds, modified_at, updated_at
|
||||
FROM radar_file_meta
|
||||
WHERE product = $1 AND render_status = 'done' AND render_path IS NOT NULL
|
||||
ORDER BY modified_at DESC NULLS LAST
|
||||
LIMIT $2
|
||||
""",
|
||||
product,
|
||||
limit,
|
||||
)
|
||||
result = []
|
||||
for row in rows:
|
||||
bounds = row["render_bounds"]
|
||||
if isinstance(bounds, str):
|
||||
bounds = json.loads(bounds)
|
||||
result.append(
|
||||
{
|
||||
"filename": row["filename"],
|
||||
"image_url": f"/api/radar/rendered/{product}/{row['filename']}.png",
|
||||
"bounds": bounds,
|
||||
"modified_at": row["modified_at"].isoformat() if row["modified_at"] else None,
|
||||
"rendered_at": row["updated_at"].isoformat() if row["updated_at"] else None,
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
|
||||
def utc_now() -> datetime:
|
||||
return datetime.now(timezone.utc)
|
||||
Reference in New Issue
Block a user