feat: switch to collections for config (#12)

* feat: switch to collections for config

* fix: validation script

Author: gadomski

.env.local

@@ -1 +1 @@
-STAC_FASTAPI_GEOPARQUET_HREF = "data/naip.parquet"
+STAC_FASTAPI_COLLECTIONS_HREF = "data/collections.json"

data/collections.json

@@ -0,0 +1,36 @@
+[
+  {
+    "type": "Collection",
+    "stac_version": "1.1.0",
+    "id": "naip",
+    "description": "This collection was generated by rustac v0.12.0 from 10000 items",
+    "license": "other",
+    "extent": {
+      "spatial": {
+        "bbox": [
+          [
+            -109.130426,
+            36.933639,
+            -101.993445,
+            41.067413
+          ]
+        ]
+      },
+      "temporal": {
+        "interval": [
+          [
+            "2019-09-19T00:00:00Z",
+            "2022-08-27T16:00:00Z"
+          ]
+        ]
+      }
+    },
+    "links": [],
+    "assets": {
+      "data": {
+        "href": "./naip.parquet",
+        "type": "application/vnd.apache.parquet"
+      }
+    }
+  }
+]

data/config.toml

@@ -9,7 +9,7 @@ dependencies = [
     "fastapi>=0.115.8",
     "geojson-pydantic>=1.2.0",
     "pydantic>=2.10.4",
-    "rustac==0.7.0b1",
+    "rustac @ git+https://github.com/stac-utils/rustac-py",
     "stac-fastapi-api>=5.0.2",
     "stac-fastapi-extensions>=5.0.2",
     "stac-fastapi-types>=5.0.2",

pyproject.toml

@@ -2,4 +2,4 @@
 
 set -e
 
-STAC_FASTAPI_GEOPARQUET_HREF=data/naip.parquet uv run fastapi dev src/stac_fastapi/geoparquet/main.py
+STAC_FASTAPI_COLLECTIONS_HREF=data/collections.json uv run fastapi dev src/stac_fastapi/geoparquet/main.py

scripts/dev

@@ -0,0 +1,41 @@
+#!/usr/bin/env python3
+
+import asyncio
+import json
+import sys
+from pathlib import Path
+
+import pystac.client
+import pystac.utils
+import rustac
+
+OUTPUT_PATH = Path(__file__).parents[1] / "data" / "collections.json"
+
+
+async def main() -> None:
+    if len(sys.argv) == 1:
+        raise Exception("ERROR: must provide at least one href on the command line")
+
+    collections = []
+    for href in sys.argv[1:]:
+        items = await rustac.read(href)
+        collection = rustac.collection_from_id_and_items(
+            Path(href).stem, items["features"]
+        )
+        collection["links"] = []
+        collection["assets"] = {
+            "data": {
+                "href": pystac.utils.make_relative_href(
+                    str(Path(href).absolute()), OUTPUT_PATH
+                ),
+                "type": "application/vnd.apache.parquet",
+            }
+        }
+        collections.append(collection)
+
+    with open(OUTPUT_PATH, "w") as f:
+        json.dump(collections, f, indent=2)
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

scripts/generate-collections

@@ -1,9 +1,10 @@
+import json
+from collections import defaultdict
 from contextlib import asynccontextmanager
-from pathlib import Path
 from typing import Any, AsyncIterator, TypedDict
 
-import tomllib
-from fastapi import FastAPI
+import pystac.utils
+from fastapi import FastAPI, HTTPException
 from rustac import DuckdbClient
 
 import stac_fastapi.api.models
@@ -12,7 +13,9 @@
 
 from .client import Client
 from .search import SearchGetRequest, SearchPostRequest
-from .settings import Settings, StacFastapiGeoparquetSettings
+from .settings import Settings
+
+GEOPARQUET_MEDIA_TYPE = "application/vnd.apache.parquet"
 
 GetSearchRequestModel = stac_fastapi.api.models.create_request_model(
     model_name="SearchGetRequest",
@@ -40,7 +43,7 @@ class State(TypedDict):
     collections: dict[str, dict[str, Any]]
     """A mapping of collection id to collection."""
 
-    hrefs: dict[str, str]
+    hrefs: dict[str, list[str]]
     """A mapping of collection id to geoparquet href."""
 
 
@@ -57,41 +60,36 @@ def create(
             stac_fastapi_description="A stac-fastapi server backend by stac-geoparquet",
         )
 
-    if settings.stac_fastapi_geoparquet_href.endswith(".toml"):
-        with open(settings.stac_fastapi_geoparquet_href, "rb") as f:
-            data = tomllib.load(f)
-        stac_fastapi_geoparquet_settings = StacFastapiGeoparquetSettings.model_validate(
-            data
-        )
-        config_directory = Path(settings.stac_fastapi_geoparquet_href).parent
-        hrefs = []
-        for href in stac_fastapi_geoparquet_settings.hrefs:
-            if Path(href).is_absolute():
-                hrefs.append(href)
-            else:
-                hrefs.append(str(config_directory.joinpath(href).resolve()))
-    else:
-        hrefs = [settings.stac_fastapi_geoparquet_href]
+    with open(settings.stac_fastapi_collections_href, "rb") as f:
+        collections = json.load(f)
 
     @asynccontextmanager
     async def lifespan(app: FastAPI) -> AsyncIterator[State]:
         client = app.extra["duckdb_client"]
-        collections = dict()
-        href_dict = dict()
-        for href in hrefs:
-            for collection in client.get_collections(href):
-                if collection["id"] in collections:
-                    raise ValueError(
-                        "cannot have two items in the same collection in different "
-                        "geoparquet files"
+        settings: Settings = app.extra["settings"]
+        collections = app.extra["collections"]
+        collection_dict = dict()
+        hrefs = defaultdict(list)
+        for collection in collections:
+            if collection["id"] in collection_dict:
+                raise HTTPException(
+                    500, f"two collections with the same id: {collection.id}"
+                )
+            else:
+                collection_dict[collection["id"]] = collection
+            for key, asset in collection["assets"].items():
+                if asset["type"] == GEOPARQUET_MEDIA_TYPE:
+                    hrefs[collection["id"]].append(
+                        pystac.utils.make_absolute_href(
+                            asset["href"],
+                            settings.stac_fastapi_collections_href,
+                            start_is_dir=False,
+                        )
                     )
-                else:
-                    collections[collection["id"]] = collection
-                href_dict[collection["id"]] = href
         yield {
             "client": client,
-            "collections": collections,
-            "hrefs": href_dict,
+            "collections": collection_dict,
+            "hrefs": hrefs,
         }
 
     api = StacApi(
@@ -102,6 +100,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[State]:
             openapi_url=settings.openapi_url,
             docs_url=settings.docs_url,
             redoc_url=settings.docs_url,
+            settings=settings,
+            collections=collections,
             duckdb_client=duckdb_client,
         ),
         search_get_request_model=GetSearchRequestModel,

src/stac_fastapi/geoparquet/api.py

@@ -162,15 +162,16 @@ async def search(
         **kwargs: Any,
     ) -> ItemCollection:
         client = cast(DuckdbClient, request.state.client)
-        hrefs = cast(dict[str, str], request.state.hrefs)
+        hrefs = cast(dict[str, list[str]], request.state.hrefs)
 
-        hrefs_to_search = set()
+        hrefs_to_search = []
         if search.collections:
             for collection in search.collections:
-                if href := hrefs.get(collection):
-                    hrefs_to_search.add(href)
+                if collection_hrefs := hrefs.get(collection):
+                    hrefs_to_search.extend(collection_hrefs)
         else:
-            hrefs_to_search.update(hrefs.values())
+            for collection_hrefs in hrefs.values():
+                hrefs_to_search.extend(collection_hrefs)
 
         if len(hrefs) > 1:
             raise ValidationError(

src/stac_fastapi/geoparquet/client.py

@@ -1,20 +1,10 @@
-from pydantic import BaseModel
-
 from stac_fastapi.types.config import ApiSettings
 
 
 class Settings(ApiSettings):  # type: ignore
     """stac-fastapi-geoparquet settings"""
 
-    stac_fastapi_geoparquet_href: str
-    """This can either be the href of a single geoparquet file, or the href of a TOML
-    configuration file.
-    """
-
-    stac_fastapi_duckdb_extension_directory: str | None = None
-    """DuckDB extension directory (if none, the default will be used)"""
-
+    stac_fastapi_collections_href: str
+    """A file containing JSON list of collections.
 
-class StacFastapiGeoparquetSettings(BaseModel):
-    hrefs: list[str]
-    """Geoparquet hrefs"""
+    Any parquet assets on the collection will be loaded into the server."""

src/stac_fastapi/geoparquet/settings.py

@@ -3,18 +3,16 @@
 
 import pytest
 from fastapi.testclient import TestClient
-from pytest import FixtureRequest
 
 import stac_fastapi.geoparquet.api
 from stac_fastapi.geoparquet import Settings
 
-GEOPARQUET_FILE = Path(__file__).parents[1] / "data" / "naip.parquet"
-TOML_FILE = Path(__file__).parents[1] / "data" / "config.toml"
+COLLECTIONS_PATH = Path(__file__).parents[1] / "data" / "collections.json"
 
 
-@pytest.fixture(params=[GEOPARQUET_FILE, TOML_FILE])
-async def client(request: FixtureRequest) -> AsyncIterator[TestClient]:
-    settings = Settings(stac_fastapi_geoparquet_href=str(request.param))
+@pytest.fixture
+async def client() -> AsyncIterator[TestClient]:
+    settings = Settings(stac_fastapi_collections_href=str(COLLECTIONS_PATH))
     api = stac_fastapi.geoparquet.api.create(settings)
     with TestClient(api.app) as client:
         yield client

tests/conftest.py

@@ -7,7 +7,7 @@
 import stac_fastapi.geoparquet.api
 from stac_fastapi.geoparquet import Settings
 
-from .conftest import GEOPARQUET_FILE
+from .conftest import COLLECTIONS_PATH
 
 
 @pytest.fixture
@@ -17,7 +17,7 @@ def extension_directory() -> Path:
 
 def test_create(extension_directory: Path) -> None:
     duckdb_client = DuckdbClient(extension_directory=str(extension_directory))
-    settings = Settings(stac_fastapi_geoparquet_href=str(GEOPARQUET_FILE))
+    settings = Settings(stac_fastapi_collections_href=str(COLLECTIONS_PATH))
     api = stac_fastapi.geoparquet.api.create(
         duckdb_client=duckdb_client, settings=settings
     )