From 26298dac1c8233f6866ac4a6dad01d5cc21df0e0 Mon Sep 17 00:00:00 2001 From: Matt Fisher Date: Mon, 24 Jun 2024 19:02:34 -0600 Subject: [PATCH] Set up API endpoint to query DB --- README.md | 13 ++++++ pyproject.toml | 2 +- src/aross_stations_db/api/__init__.py | 11 +++++ src/aross_stations_db/api/v1/__init__.py | 0 src/aross_stations_db/api/v1/routes.py | 59 ++++++++++++++++++++++++ src/aross_stations_db/cli.py | 21 ++++++--- src/aross_stations_db/config.py | 9 ++-- src/aross_stations_db/db.py | 4 +- src/aross_stations_db/middleware.py | 26 +++++++++++ src/aross_stations_db/tables.py | 12 ++--- 10 files changed, 136 insertions(+), 21 deletions(-) create mode 100644 src/aross_stations_db/api/__init__.py create mode 100644 src/aross_stations_db/api/v1/__init__.py create mode 100644 src/aross_stations_db/api/v1/routes.py create mode 100644 src/aross_stations_db/middleware.py diff --git a/README.md b/README.md index 1646ec5..79335ea 100644 --- a/README.md +++ b/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)) +``` + + [actions-badge]: https://github.com/nsidc/aross-stations-db/workflows/CI/badge.svg [actions-link]: https://github.com/nsidc/aross-stations-db/actions diff --git a/pyproject.toml b/pyproject.toml index ebac266..776ca7a 100644 --- a/pyproject.toml +++ b/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]", diff --git a/src/aross_stations_db/api/__init__.py b/src/aross_stations_db/api/__init__.py new file mode 100644 index 0000000..0940148 --- /dev/null +++ b/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) diff --git a/src/aross_stations_db/api/v1/__init__.py b/src/aross_stations_db/api/v1/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/aross_stations_db/api/v1/routes.py b/src/aross_stations_db/api/v1/routes.py new file mode 100644 index 0000000..2f2ca92 --- /dev/null +++ b/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 + ] diff --git a/src/aross_stations_db/cli.py b/src/aross_stations_db/cli.py index f24e19c..4996b7c 100644 --- a/src/aross_stations_db/cli.py +++ b/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,7 +22,12 @@ 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") @@ -33,10 +35,15 @@ def init() -> None: @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") diff --git a/src/aross_stations_db/config.py b/src/aross_stations_db/config.py index 8aa4f05..ee75d02 100644 --- a/src/aross_stations_db/config.py +++ b/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)) diff --git a/src/aross_stations_db/db.py b/src/aross_stations_db/db.py index ddd5fd7..16a990e 100644 --- a/src/aross_stations_db/db.py +++ b/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 ] diff --git a/src/aross_stations_db/middleware.py b/src/aross_stations_db/middleware.py new file mode 100644 index 0000000..0b399a5 --- /dev/null +++ b/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() diff --git a/src/aross_stations_db/tables.py b/src/aross_stations_db/tables.py index 6e12eee..ab88cc7 100644 --- a/src/aross_stations_db/tables.py +++ b/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.