From adf4cc6491ac2921a47523c86bcb4ee527e55380 Mon Sep 17 00:00:00 2001 From: tjellicoe Date: Mon, 22 Apr 2024 16:17:25 +0100 Subject: [PATCH] Added collection search extension --- stac_fastapi/api/stac_fastapi/api/config.py | 1 + .../stac_fastapi/extensions/core/__init__.py | 2 + .../core/collectionSearch/__init__.py | 6 + .../core/collectionSearch/collectionSearch.py | 109 ++++++++++++++++++ .../core/collectionSearch/request.py | 28 +++++ stac_fastapi/types/stac_fastapi/types/core.py | 77 ++++++++++++- .../types/stac_fastapi/types/search.py | 92 +++++++++++++++ 7 files changed, 314 insertions(+), 1 deletion(-) create mode 100644 stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/__init__.py create mode 100644 stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/collectionSearch.py create mode 100644 stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/request.py diff --git a/stac_fastapi/api/stac_fastapi/api/config.py b/stac_fastapi/api/stac_fastapi/api/config.py index e6e4d882a..b4e2c8620 100644 --- a/stac_fastapi/api/stac_fastapi/api/config.py +++ b/stac_fastapi/api/stac_fastapi/api/config.py @@ -23,3 +23,4 @@ class AddOns(enum.Enum): """Enumeration of available third party add ons.""" bulk_transaction = "bulk-transaction" + collection_search = "collection-search" diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py index 96317fe4a..be608101c 100644 --- a/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py +++ b/stac_fastapi/extensions/stac_fastapi/extensions/core/__init__.py @@ -6,6 +6,7 @@ from .query import QueryExtension from .sort import SortExtension from .transaction import TransactionExtension +from .collectionSearch import CollectionSearchExtension __all__ = ( "ContextExtension", @@ -16,4 +17,5 @@ "SortExtension", "TokenPaginationExtension", "TransactionExtension", + "CollectionSearchExtension", ) diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/__init__.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/__init__.py new file mode 100644 index 000000000..21f35e8ca --- /dev/null +++ b/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/__init__.py @@ -0,0 +1,6 @@ +"""Filter extension module.""" + + +from .collectionSearch import CollectionSearchExtension + +__all__ = ["CollectionSearchExtension"] diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/collectionSearch.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/collectionSearch.py new file mode 100644 index 000000000..91f4fc6d6 --- /dev/null +++ b/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/collectionSearch.py @@ -0,0 +1,109 @@ +# encoding: utf-8 +"""Collection Search Extension.""" +from enum import Enum +from typing import List, Type, Union + +import attr +from fastapi import APIRouter, FastAPI +from starlette.responses import Response +from stac_pydantic.api.collections import Collections + +from stac_fastapi.api.models import JSONSchemaResponse +from stac_fastapi.api.routes import create_async_endpoint +from stac_fastapi.types.core import AsyncCollectionSearchClient, CollectionSearchClient +from stac_fastapi.types.extension import ApiExtension +from stac_fastapi.types.search import BaseCollectionSearchGetRequest, BaseCollectionSearchPostRequest + +from .request import CollectionSearchExtensionGetRequest, CollectionSearchExtensionPostRequest + +class CollectionSearchConformanceClasses(str, Enum): + """Conformance classes for the Collection Search extension. + + See + https://github.com/stac-api-extensions/collection-search + """ + + CORE = "https://api.stacspec.org/v1.0.0-rc.1/core" + COLLECTION_SEARCH = "https://api.stacspec.org/v1.0.0-rc.1/collection-search" + SIMPLE_QUERY = "http://www.opengis.net/spec/ogcapi-common-2/1.0/conf/simple-query" + +@attr.s +class CollectionSearchExtension(ApiExtension): + """CollectionSearch Extension. + + The collection search extension adds two endpoints which allow searching of + collections via GET and POST: + GET /collection-search + POST /collection-search + + https://github.com/stac-api-extensions/collection-search + + Attributes: + search_get_request_model: Get request model for collection search + search_post_request_model: Post request model for collection search + client: Collection Search endpoint logic + conformance_classes: Conformance classes provided by the extension + """ + + GET = CollectionSearchExtensionGetRequest + POST = CollectionSearchExtensionPostRequest + + collection_search_get_request_model: Type[BaseCollectionSearchGetRequest] = attr.ib( + default=BaseCollectionSearchGetRequest + ) + collection_search_post_request_model: Type[BaseCollectionSearchPostRequest] = attr.ib( + default=BaseCollectionSearchPostRequest + ) + + client: Union[AsyncCollectionSearchClient, CollectionSearchClient] = attr.ib( + factory=CollectionSearchClient + ) + + conformance_classes: List[str] = attr.ib( + default=[ + CollectionSearchConformanceClasses.CORE, + CollectionSearchConformanceClasses.COLLECTION_SEARCH, + CollectionSearchConformanceClasses.SIMPLE_QUERY, + ] + ) + router: APIRouter = attr.ib(factory=APIRouter) + response_class: Type[Response] = attr.ib(default=JSONSchemaResponse) + + + def register(self, app: FastAPI) -> None: + """Register the extension with a FastAPI application. + + Args: + app: target FastAPI application. + + Returns: + None + """ + self.router.prefix = app.state.router_prefix + self.router.add_api_route( + name="Collection Search", + path="/collection-search", + response_model=Collections, + response_class=self.response_class, + response_model_exclude_unset=True, + response_model_exclude_none=True, + methods=["GET"], + endpoint=create_async_endpoint( + self.client.get_collection_search, self.collection_search_get_request_model + ), + ) + + self.router.add_api_route( + name="Collection Search", + path="/collection-search", + response_model=Collections, + response_class=self.response_class, + response_model_exclude_unset=True, + response_model_exclude_none=True, + methods=["POST"], + endpoint=create_async_endpoint( + self.client.post_collection_search, self.collection_search_post_request_model + ), + ) + + app.include_router(self.router, tags=["Collection Search Extension"]) diff --git a/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/request.py b/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/request.py new file mode 100644 index 000000000..f159d6a87 --- /dev/null +++ b/stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/request.py @@ -0,0 +1,28 @@ +"""Collection Search extension request models.""" + +from enum import Enum +from typing import Any, Dict, Optional, List + +import attr +from pydantic import BaseModel, Field +from stac_pydantic.shared import BBox + +from stac_fastapi.types.search import APIRequest, Limit, str2bbox + +from stac_fastapi.types.rfc3339 import DateTimeType, str_to_interval + +@attr.s +class CollectionSearchExtensionGetRequest(APIRequest): + """Collection Search extension GET request model.""" + + bbox: Optional[BBox] = attr.ib(default=None, converter=str2bbox) + datetime: Optional[DateTimeType] = attr.ib(default=None, converter=str_to_interval) + limit: Optional[int] = attr.ib(default=10) + + +class CollectionSearchExtensionPostRequest(BaseModel): + """Collection Search extension POST request model.""" + + bbox: Optional[BBox] + datetime: Optional[DateTimeType] + limit: Optional[Limit] = Field(default=10) diff --git a/stac_fastapi/types/stac_fastapi/types/core.py b/stac_fastapi/types/stac_fastapi/types/core.py index 77078ace3..e8d75b602 100644 --- a/stac_fastapi/types/stac_fastapi/types/core.py +++ b/stac_fastapi/types/stac_fastapi/types/core.py @@ -17,7 +17,7 @@ from stac_fastapi.types.extension import ApiExtension from stac_fastapi.types.requests import get_base_url from stac_fastapi.types.rfc3339 import DateTimeType -from stac_fastapi.types.search import BaseSearchPostRequest +from stac_fastapi.types.search import BaseSearchPostRequest, BaseCollectionSearchPostRequest from stac_fastapi.types.stac import Conformance NumType = Union[float, int] @@ -801,3 +801,78 @@ def get_queryables( "description": "Queryable names for the example STAC API Item Search filter.", "properties": {}, } + +@attr.s +class AsyncCollectionSearchClient(abc.ABC): + """Defines a pattern for implementing the STAC Collection Search extension.""" + + @abc.abstractmethod + async def post_collection_search( + self, search_request: BaseCollectionSearchPostRequest, **kwargs + ) -> stac_types.ItemCollection: + """Cross catalog search (POST) of collections. + + Called with `POST /collection-search`. + + Args: + search_request: search request parameters. + + Returns: + A tuple of (collections, next pagination token if any). + """ + ... + + @abc.abstractmethod + async def get_collection_search( + self, + bbox: Optional[BBox] = None, + datetime: Optional[DateTimeType] = None, + limit: Optional[int] = 10, + **kwargs, + ) -> stac_types.Collections: + """Cross catalog search (GET) of collections. + + Called with `GET /collection-search`. + + Returns: + A tuple of (collections, next pagination token if any). + """ + ... + + +@attr.s +class CollectionSearchClient(abc.ABC): + """Defines a pattern for implementing the STAC Collection Search extension.""" + + @abc.abstractmethod + def post_collection_search( + self, search_request: BaseCollectionSearchPostRequest, **kwargs + ) -> stac_types.Collections: + """Cross catalog search (POST) of collections. + + Called with `POST /collection-search`. + + Args: + search_request: search request parameters. + + Returns: + A tuple of (collections, next pagination token if any). + """ + ... + + @abc.abstractmethod + def get_collection_search( + self, + bbox: Optional[BBox] = None, + datetime: Optional[DateTimeType] = None, + limit: Optional[int] = 10, + **kwargs, + ) -> stac_types.Collections: + """Cross catalog search (GET) of collections. + + Called with `GET /collection-search`. + + Returns: + A tuple of (collections, next pagination token if any). + """ + ... \ No newline at end of file diff --git a/stac_fastapi/types/stac_fastapi/types/search.py b/stac_fastapi/types/stac_fastapi/types/search.py index 0851c1d30..46831c056 100644 --- a/stac_fastapi/types/stac_fastapi/types/search.py +++ b/stac_fastapi/types/stac_fastapi/types/search.py @@ -187,7 +187,99 @@ def validate_bbox(cls, v: Union[str, BBox]) -> BBox: # Validate against WGS84 if xmin < -180 or ymin < -90 or xmax > 180 or ymax > 90: raise ValueError("Bounding box must be within (-180, -90, 180, 90)") + return v + + @validator("datetime", pre=True) + def validate_datetime(cls, v: Union[str, DateTimeType]) -> DateTimeType: + """Parse datetime.""" + if type(v) == str: + v = str_to_interval(v) + return v + + @property + def spatial_filter(self) -> Optional[_GeometryBase]: + """Return a geojson-pydantic object representing the spatial filter for the search + request. + + Check for both because the ``bbox`` and ``intersects`` parameters are + mutually exclusive. + """ + if self.bbox: + return Polygon( + coordinates=[ + [ + [self.bbox[0], self.bbox[3]], + [self.bbox[2], self.bbox[3]], + [self.bbox[2], self.bbox[1]], + [self.bbox[0], self.bbox[1]], + [self.bbox[0], self.bbox[3]], + ] + ] + ) + if self.intersects: + return self.intersects + return + + +@attr.s +class BaseCollectionSearchGetRequest(APIRequest): + """Base arguments for Collection Search GET Request.""" + + bbox: Optional[BBox] = attr.ib(default=None, converter=str2bbox) + datetime: Optional[DateTimeType] = attr.ib(default=None, converter=str_to_interval) + limit: Optional[int] = attr.ib(default=10) + + +class BaseCollectionSearchPostRequest(BaseModel): + """Search model. + + Replace base model in STAC-pydantic as it includes additional fields, not in the core + model. + """ + bbox: Optional[BBox] + datetime: Optional[DateTimeType] + limit: Optional[Limit] = Field(default=10) + + @property + def start_date(self) -> Optional[datetime]: + """Extract the start date from the datetime string.""" + return self.datetime[0] if self.datetime else None + + @property + def end_date(self) -> Optional[datetime]: + """Extract the end date from the datetime string.""" + return self.datetime[1] if self.datetime else None + + @validator("bbox", pre=True) + def validate_bbox(cls, v: Union[str, BBox]) -> BBox: + """Check order of supplied bbox coordinates.""" + if v: + if type(v) == str: + v = str2bbox(v) + # Validate order + if len(v) == 4: + xmin, ymin, xmax, ymax = v + else: + xmin, ymin, min_elev, xmax, ymax, max_elev = v + if max_elev < min_elev: + raise ValueError( + "Maximum elevation must greater than minimum elevation" + ) + + if xmax < xmin: + raise ValueError( + "Maximum longitude must be greater than minimum longitude" + ) + + if ymax < ymin: + raise ValueError( + "Maximum longitude must be greater than minimum longitude" + ) + + # Validate against WGS84 + if xmin < -180 or ymin < -90 or xmax > 180 or ymax > 90: + raise ValueError("Bounding box must be within (-180, -90, 180, 90)") return v @validator("datetime", pre=True)