Adding Collection Search Extension

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"

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",
 )

stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/__init__.py

@@ -0,0 +1,6 @@
+"""Filter extension module."""
+
+
+from .collectionSearch import CollectionSearchExtension
+
+__all__ = ["CollectionSearchExtension"]

stac_fastapi/extensions/stac_fastapi/extensions/core/collectionSearch/collectionSearch.py

@@ -0,0 +1,108 @@
+# 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 JSONResponse, 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=JSONResponse)    
+
+    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=None,
+            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=None,
+            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"])

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)

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).
+        """
+        ...

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)