add query/sort/fields extension support for item_collection and refactor extensions mapping

vincentsarago · vincentsarago · commit c009e32785cb · 2025-01-23T14:57:34.000+01:00
diff --git a/stac_fastapi/pgstac/app.py b/stac_fastapi/pgstac/app.py
@@ -23,22 +23,29 @@
 from stac_fastapi.api.openapi import update_openapi
 from stac_fastapi.extensions.core import (
     FieldsExtension,
-    FilterExtension,
     FreeTextExtension,
     OffsetPaginationExtension,
     SortExtension,
     TokenPaginationExtension,
     TransactionExtension,
 )
 from stac_fastapi.extensions.core.collection_search import CollectionSearchExtension
+from stac_fastapi.extensions.core.collection_search.request import (
+    BaseCollectionSearchGetRequest,
+)
 from stac_fastapi.extensions.third_party import BulkTransactionExtension
 from starlette.middleware import Middleware
 
 from stac_fastapi.pgstac.config import Settings
 from stac_fastapi.pgstac.core import CoreCrudClient
 from stac_fastapi.pgstac.db import close_db_connection, connect_to_db
 from stac_fastapi.pgstac.extensions import QueryExtension
-from stac_fastapi.pgstac.extensions.filter import FiltersClient
+from stac_fastapi.pgstac.extensions.filter import (
+    CollectionSearchFilterExtension,
+    FiltersClient,
+    ItemCollectionFilterExtension,
+    SearchFilterExtension,
+)
 from stac_fastapi.pgstac.transactions import BulkTransactionsClient, TransactionsClient
 from stac_fastapi.pgstac.types.search import PgstacSearch
 
@@ -59,23 +66,48 @@
     "query": QueryExtension(),
     "sort": SortExtension(),
     "fields": FieldsExtension(),
-    "filter": FilterExtension(client=FiltersClient()),
+    "filter": SearchFilterExtension(client=FiltersClient()),
     "pagination": TokenPaginationExtension(),
 }
 
 # collection_search extensions
 cs_extensions_map = {
-    "query": QueryExtension(),
-    "sort": SortExtension(),
-    "fields": FieldsExtension(),
-    "filter": FilterExtension(client=FiltersClient()),
-    "free_text": FreeTextExtension(),
+    "query": QueryExtension(
+        conformance_classes=[
+            "https://api.stacspec.org/v1.0.0-rc.1/collection-search#query"
+        ]
+    ),
+    "sort": SortExtension(
+        conformance_classes=[
+            "https://api.stacspec.org/v1.0.0-rc.1/collection-search#sort"
+        ]
+    ),
+    "fields": FieldsExtension(
+        conformance_classes=[
+            "https://api.stacspec.org/v1.0.0-rc.1/collection-search#fields"
+        ]
+    ),
+    "filter": CollectionSearchFilterExtension(client=FiltersClient()),
+    "free_text": FreeTextExtension(
+        conformance_classes=[
+            "https://api.stacspec.org/v1.0.0-rc.1/collection-search#free-text",
+        ],
+    ),
     "pagination": OffsetPaginationExtension(),
 }
 
 # item_collection extensions
 itm_col_extensions_map = {
-    "filter": FilterExtension(client=FiltersClient()),
+    "query": QueryExtension(
+        conformance_classes=["https://api.stacspec.org/v1.0.0/ogcapi-features#query"],
+    ),
+    "sort": SortExtension(
+        conformance_classes=["https://api.stacspec.org/v1.0.0/ogcapi-features#sort"],
+    ),
+    "fields": FieldsExtension(
+        conformance_classes=["https://api.stacspec.org/v1.0.0/ogcapi-features#fields"],
+    ),
+    "filter": ItemCollectionFilterExtension(client=FiltersClient()),
     "pagination": TokenPaginationExtension(),
 }
 
@@ -123,17 +155,31 @@
         extensions=itm_col_extensions,
         request_type="GET",
     )
+    application_extensions.extend(itm_col_extensions)
 
 # /collections model
 collections_get_request_model = EmptyRequest
 if "collection_search" in enabled_extensions:
-    cs_extensions = [
-        extension
-        for key, extension in cs_extensions_map.items()
-        if key in enabled_extensions
+    mixins = []
+    mixing_conformances = [
+        "https://api.stacspec.org/v1.0.0-rc.1/collection-search",
+        "http://www.opengis.net/spec/ogcapi-common-2/1.0/conf/simple-query",
     ]
-    collection_search_extension = CollectionSearchExtension.from_extensions(cs_extensions)
-    collections_get_request_model = collection_search_extension.GET
+    for key, extension in cs_extensions_map.items():
+        if key not in enabled_extensions:
+            continue
+        mixins.append(extension.GET)
+        mixing_conformances.extend(extension.conformance_classes)
+
+    collections_get_request_model = create_request_model(
+        model_name="CollectionsGetRequest",
+        base_model=BaseCollectionSearchGetRequest,
+        mixins=mixins,
+        request_type="GET",
+    )
+    collection_search_extension = CollectionSearchExtension(
+        GET=collections_get_request_model, conformance_classes=mixing_conformances
+    )
     application_extensions.append(collection_search_extension)
 
 
diff --git a/stac_fastapi/pgstac/core.py b/stac_fastapi/pgstac/core.py
@@ -343,9 +343,12 @@ async def item_collection(
         datetime: Optional[str] = None,
         limit: Optional[int] = None,
         # Extensions
-        token: Optional[str] = None,
+        query: Optional[str] = None,
+        fields: Optional[List[str]] = None,
+        sortby: Optional[str] = None,
         filter_expr: Optional[str] = None,
         filter_lang: Optional[str] = None,
+        token: Optional[str] = None,
         **kwargs,
     ) -> ItemCollection:
         """Get all items from a specific collection.
@@ -369,12 +372,15 @@ async def item_collection(
             "datetime": datetime,
             "limit": limit,
             "token": token,
+            "query": orjson.loads(unquote_plus(query)) if query else query,
         }
 
         clean = self._clean_search_args(
             base_args=base_args,
             filter_query=filter_expr,
             filter_lang=filter_lang,
+            fields=fields,
+            sortby=sortby,
         )
 
         search_request = self.pgstac_search_model(**clean)
@@ -450,11 +456,11 @@ async def get_search(
         limit: Optional[int] = None,
         # Extensions
         query: Optional[str] = None,
-        token: Optional[str] = None,
         fields: Optional[List[str]] = None,
         sortby: Optional[str] = None,
         filter_expr: Optional[str] = None,
         filter_lang: Optional[str] = None,
+        token: Optional[str] = None,
         **kwargs,
     ) -> ItemCollection:
         """Cross catalog search (GET).
diff --git a/stac_fastapi/pgstac/extensions/filter.py b/stac_fastapi/pgstac/extensions/filter.py
@@ -1,11 +1,24 @@
 """Get Queryables."""
 
-from typing import Any, Dict, Optional
+from typing import Any, Dict, List, Optional, Type
 
+import attr
 from buildpg import render
-from fastapi import Request
+from fastapi import APIRouter, FastAPI, Request
+from stac_fastapi.api.models import CollectionUri, EmptyRequest, JSONSchemaResponse
+from stac_fastapi.api.routes import create_async_endpoint
+from stac_fastapi.extensions.core import FilterExtension
+from stac_fastapi.extensions.core.collection_search.collection_search import (
+    ConformanceClasses as CollectionSearchConformanceClasses,
+)
 from stac_fastapi.extensions.core.filter.client import AsyncBaseFiltersClient
+from stac_fastapi.extensions.core.filter.filter import FilterConformanceClasses
+from stac_fastapi.extensions.core.filter.request import (
+    FilterExtensionGetRequest,
+    FilterExtensionPostRequest,
+)
 from stac_fastapi.types.errors import NotFoundError
+from starlette.responses import Response
 
 
 class FiltersClient(AsyncBaseFiltersClient):
@@ -40,3 +53,124 @@ async def get_queryables(
 
             queryables["$id"] = str(request.url)
             return queryables
+
+
+@attr.s
+class SearchFilterExtension(FilterExtension):
+    """Item Search Filter Extension."""
+
+    GET = FilterExtensionGetRequest
+    POST = FilterExtensionPostRequest
+
+    client: FiltersClient = attr.ib(factory=FiltersClient)
+    conformance_classes: List[str] = attr.ib(
+        default=[
+            FilterConformanceClasses.FILTER,
+            FilterConformanceClasses.ITEM_SEARCH_FILTER,
+            FilterConformanceClasses.BASIC_CQL2,
+            FilterConformanceClasses.CQL2_JSON,
+            FilterConformanceClasses.CQL2_TEXT,
+        ]
+    )
+    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="Queryables",
+            path="/queryables",
+            methods=["GET"],
+            responses={
+                200: {
+                    "content": {
+                        "application/schema+json": {},
+                    },
+                    # TODO: add output model in stac-pydantic
+                },
+            },
+            response_class=self.response_class,
+            endpoint=create_async_endpoint(self.client.get_queryables, EmptyRequest),
+        )
+        app.include_router(self.router, tags=["Filter Extension"])
+
+
+@attr.s
+class ItemCollectionFilterExtension(FilterExtension):
+    """Item Collection Filter Extension."""
+
+    GET = FilterExtensionGetRequest
+    POST = FilterExtensionPostRequest
+
+    client: FiltersClient = attr.ib(factory=FiltersClient)
+    conformance_classes: List[str] = attr.ib(
+        default=[
+            FilterConformanceClasses.FILTER,
+            FilterConformanceClasses.FEATURES_FILTER,
+            FilterConformanceClasses.BASIC_CQL2,
+            FilterConformanceClasses.CQL2_JSON,
+            FilterConformanceClasses.CQL2_TEXT,
+        ]
+    )
+    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.add_api_route(
+            name="Collection Queryables",
+            path="/collections/{collection_id}/queryables",
+            methods=["GET"],
+            responses={
+                200: {
+                    "content": {
+                        "application/schema+json": {},
+                    },
+                    # TODO: add output model in stac-pydantic
+                },
+            },
+            response_class=self.response_class,
+            endpoint=create_async_endpoint(self.client.get_queryables, CollectionUri),
+        )
+        app.include_router(self.router, tags=["Filter Extension"])
+
+
+@attr.s
+class CollectionSearchFilterExtension(FilterExtension):
+    """Collection Search Filter Extension."""
+
+    GET = FilterExtensionGetRequest
+    POST = FilterExtensionPostRequest
+
+    client: FiltersClient = attr.ib(factory=FiltersClient)
+    conformance_classes: List[str] = attr.ib(
+        default=[CollectionSearchConformanceClasses.FILTER]
+    )
+    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
+        """
+        pass
