Set up API endpoint to query DB

README.md

@@ -137,6 +137,19 @@ where
 ```
 
 
+### Run API
+
+```bash
+fastapi dev src/aross_stations_db/api
+```
+
+Example query:
+
+```
+http://127.0.0.1:8000/v1/?start=2023-01-01&end=2023-06-01&polygon=POLYGON%20((-159.32130625160698%2069.56469019745796,%20-159.32130625160698%2068.08208920517862,%20-150.17196253090276%2068.08208920517862,%20-150.17196253090276%2069.56469019745796,%20-159.32130625160698%2069.56469019745796))
+```
+
+
 <!-- prettier-ignore-start -->
 [actions-badge]:            https://github.com/nsidc/aross-stations-db/workflows/CI/badge.svg
 [actions-link]:             https://github.com/nsidc/aross-stations-db/actions

pyproject.toml

@@ -23,9 +23,9 @@ classifiers = [
 dynamic = ["version"]
 dependencies = [
   "loguru",
+  "fastapi ~=0.111.0",
   "pydantic ~=2.0",
   "pydantic-settings",
-  # TODO: Good idea? Will help with API development :)
   "sqlalchemy ~=2.0",
   "geoalchemy2",
   "psycopg[binary,pool]",

src/aross_stations_db/api/__init__.py

@@ -0,0 +1,11 @@
+from fastapi import FastAPI
+
+from aross_stations_db.api.v1.routes import router as v1_router
+
+api = FastAPI(
+    title=(
+        "Rain on snow events detected by"
+        " Automated Surface Observing System (ASOS) stations"
+    ),
+)
+api.include_router(v1_router)

src/aross_stations_db/api/v1/routes.py

@@ -0,0 +1,59 @@
+import datetime as dt
+from typing import Annotated
+
+from fastapi import APIRouter, Depends, Query
+from geoalchemy2 import WKBElement
+from pydantic import BaseModel
+from sqlalchemy import func
+from sqlalchemy.orm import Session
+
+from aross_stations_db.middleware import get_db_session
+from aross_stations_db.tables import Event, Station
+
+router = APIRouter(prefix="/v1", tags=["v1"])
+
+
+class ReturnElement(BaseModel):
+    name: str
+    location: str
+    event_count: int
+
+
+@router.get("/")
+def get(
+    start: Annotated[dt.datetime, Query(description="ISO-format timestamp")],
+    end: Annotated[dt.datetime, Query(description="ISO-format timestamp")],
+    db: Annotated[Session, Depends(get_db_session)],
+    polygon: Annotated[str | None, WKBElement, Query(description="WKT shape")] = None,
+) -> list[ReturnElement]:
+    """Get stations and their events matching query parameters."""
+    query = (
+        db.query(
+            Station,
+            Station.location.ST_AsEWKT(),
+            func.count(),
+        )
+        .join(
+            Event,
+        )
+        .filter(Event.time_start >= start, Event.time_end < end)
+    )
+
+    if polygon:
+        query = query.filter(
+            Station.location.ST_Within(
+                func.ST_SetSRID(func.ST_GeomFromText(polygon), 4326),
+            )
+        )
+
+    query = query.group_by(Station.id)
+
+    results = query.all()
+    return [
+        ReturnElement(
+            name=station.name,
+            location=location,
+            event_count=event_count,
+        )
+        for station, location, event_count in results
+    ]

src/aross_stations_db/cli.py

@@ -1,5 +1,6 @@
 import click
 from loguru import logger
+from sqlalchemy.orm import Session
 
 from aross_stations_db.config import Settings
 from aross_stations_db.db import (
@@ -12,10 +13,6 @@
     get_stations,
 )
 
-# TODO: False-positive. Remove type-ignore.
-#       See: https://github.com/pydantic/pydantic/issues/6713
-config = Settings()  # type:ignore[call-arg]
-
 
 @click.group()
 def cli() -> None:
@@ -25,18 +22,28 @@ def cli() -> None:
 @cli.command
 def init() -> None:
     """Create the database tables."""
-    create_tables(config.db_session)
+    # TODO: False-positive. Remove type-ignore.
+    #       See: https://github.com/pydantic/pydantic/issues/6713
+    config = Settings()  # type:ignore[call-arg]
+
+    with Session(config.db_engine) as db_session:
+        create_tables(db_session)
 
     logger.success("Tables created")
 
 
 @cli.command
 def load() -> None:
     """Load the database tables from files on disk."""
+    # TODO: False-positive. Remove type-ignore.
+    #       See: https://github.com/pydantic/pydantic/issues/6713
+    config = Settings()  # type:ignore[call-arg]
+
     stations = get_stations(config.stations_metadata_filepath)
     events = get_events(config.events_dir)
 
-    load_stations(stations, session=config.db_session)
-    load_events(events, session=config.db_session)
+    with Session(config.db_engine) as db_session:
+        load_stations(stations, session=db_session)
+        load_events(events, session=db_session)
 
     logger.success("Data loaded")

src/aross_stations_db/config.py

@@ -3,8 +3,7 @@
 from dotenv import load_dotenv
 from pydantic import DirectoryPath, FilePath, PostgresDsn, computed_field
 from pydantic_settings import BaseSettings, SettingsConfigDict
-from sqlalchemy import create_engine
-from sqlalchemy.orm import Session
+from sqlalchemy import Engine, create_engine
 
 # Add magic support for `.env` file 🪄
 load_dotenv()
@@ -29,8 +28,8 @@ def events_dir(self) -> DirectoryPath:
     def stations_metadata_filepath(self) -> FilePath:
         return self.DATA_BASEDIR / "metadata" / "aross.asos_stations.metadata.csv"
 
+    # TODO: Remove?
     @computed_field  # type:ignore[misc]
     @cached_property
-    def db_session(self) -> Session:
-        engine = create_engine(str(self.DB_CONNSTR))
-        return Session(engine)
+    def db_engine(self) -> Engine:
+        return create_engine(str(self.DB_CONNSTR))

src/aross_stations_db/db.py

@@ -37,8 +37,8 @@ def load_events(events: Iterator[dict[str, str]], *, session: Session) -> None:
         [
             Event(
                 station_id=event["station_id"],
-                start_timestamp=dt.datetime.fromisoformat(event["start"]),
-                end_timestamp=dt.datetime.fromisoformat(event["end"]),
+                time_start=dt.datetime.fromisoformat(event["start"]),
+                time_end=dt.datetime.fromisoformat(event["end"]),
             )
             for event in events
         ]

src/aross_stations_db/middleware.py

@@ -0,0 +1,26 @@
+from collections.abc import Iterator
+from functools import cache
+from typing import Annotated
+
+from fastapi import Depends
+from sqlalchemy.orm import Session, sessionmaker
+
+from aross_stations_db.config import Settings
+
+
+@cache
+def get_config() -> Settings:
+    # TODO: False-positive. Remove type-ignore.
+    #       See: https://github.com/pydantic/pydantic/issues/6713
+    return Settings()  # type:ignore[call-arg]
+
+
+def get_db_session(
+    config: Annotated[Settings, Depends(get_config)],
+) -> Iterator[Session]:
+    SessionFactory = sessionmaker(config.db_engine)
+    session = SessionFactory()
+    try:
+        yield session
+    finally:
+        session.close()

src/aross_stations_db/tables.py

@@ -7,6 +7,7 @@
     DeclarativeBase,
     Mapped,
     mapped_column,
+    relationship,
 )
 
 
@@ -29,16 +30,15 @@ class Station(Base):
         index=True,
     )
 
+    events = relationship("Event", backref="station")
+
 
 class Event(Base):
     __tablename__ = "event"
 
-    # TODO: Is this a good PK or should it be station_id + start?
-    id: Mapped[int] = mapped_column(primary_key=True)
-
-    station_id: Mapped[str] = mapped_column(ForeignKey("station.id"), index=True)
-    start_timestamp: Mapped[dt.datetime] = mapped_column(index=True)
-    end_timestamp: Mapped[dt.datetime] = mapped_column(index=True)
+    station_id: Mapped[str] = mapped_column(ForeignKey("station.id"), primary_key=True)
+    time_start: Mapped[dt.datetime] = mapped_column(primary_key=True)
+    time_end: Mapped[dt.datetime] = mapped_column(primary_key=True)
 
     # TODO: More fields: duration,RA,UP,FZRA,SOLID,t2m_mean,t2m_min,t2m_max,sog
     #       Don't think we need to keep duration.