Feature/refactor extension maps (#190)

* refactor extension map and add filter extension for item_collection * update changelog * transaction -> application * update conftests
stac-utils · Jan 22, 2025 · df4c12a · df4c12a
1 parent af98395
commit df4c12a
Show file tree

Hide file tree

Showing 4 changed files with 98 additions and 64 deletions.
diff --git a/CHANGES.md b/CHANGES.md
@@ -4,8 +4,7 @@
 
 ### Changed
 
-- Handle `next` and `dev` tokens now returned as links from pgstac>=0.9.0 (author @zstatmanweil, <https://github.com/stac-utils/stac-fastapi-pgstac/pull/140>)
-- Add collection search extension ([#139](https://github.com/stac-utils/stac-fastapi-pgstac/pull/139))
+- handle `next` and `dev` tokens now returned as links from pgstac>=0.9.0 (author @zstatmanweil, <https://github.com/stac-utils/stac-fastapi-pgstac/pull/140>)
 - keep `/search` and `/collections` extensions separate ([#158](https://github.com/stac-utils/stac-fastapi-pgstac/pull/158))
 - update `pypgstac` requirement to `>=0.8,<0.10`
 - set `pypgstac==0.9.*` for test requirements
@@ -15,12 +14,17 @@
 - changed `datetime` input type to `string` in GET endpoint methods
 - renamed `filter` to `filter_expr` input attributes in GET endpoint methods
 - delete `utils.format_datetime_range` function
+
+### Added
+
+- add [collection search extension](https://github.com/stac-api-extensions/collection-search) support ([#139](https://github.com/stac-utils/stac-fastapi-pgstac/pull/139))
 - add [free-text extension](https://github.com/stac-api-extensions/freetext-search) to collection search extensions ([#162](https://github.com/stac-utils/stac-fastapi-pgstac/pull/162))
+- add [filter extension](https://github.com/stac-api-extensions/filter) support to Item Collection endpoint
 
 ### Fixed
 
-- Fix Docker compose file, so example data can be loaded into database (author @zstatmanweil, <https://github.com/stac-utils/stac-fastapi-pgstac/pull/142>)
-- Fix `filter` extension implementation in `CoreCrudClient`
+- fix Docker compose file, so example data can be loaded into database (author @zstatmanweil, <https://github.com/stac-utils/stac-fastapi-pgstac/pull/142>)
+- fix `filter` extension implementation in `CoreCrudClient`
 
 ## [3.0.1] - 2024-11-14
 

diff --git a/stac_fastapi/pgstac/app.py b/stac_fastapi/pgstac/app.py
@@ -43,22 +43,28 @@
 from stac_fastapi.pgstac.types.search import PgstacSearch
 
 settings = Settings()
-extensions_map = {
+
+# application extensions
+application_extensions_map = {
     "transaction": TransactionExtension(
         client=TransactionsClient(),
         settings=settings,
         response_class=ORJSONResponse,
     ),
+    "bulk_transactions": BulkTransactionExtension(client=BulkTransactionsClient()),
+}
+
+# search extensions
+search_extensions_map = {
     "query": QueryExtension(),
     "sort": SortExtension(),
     "fields": FieldsExtension(),
-    "pagination": TokenPaginationExtension(),
     "filter": FilterExtension(client=FiltersClient()),
-    "bulk_transactions": BulkTransactionExtension(client=BulkTransactionsClient()),
+    "pagination": TokenPaginationExtension(),
 }
 
-# some extensions are supported in combination with the collection search extension
-collection_extensions_map = {
+# collection_search extensions
+cs_extensions_map = {
     "query": QueryExtension(),
     "sort": SortExtension(),
     "fields": FieldsExtension(),
@@ -67,44 +73,68 @@
     "pagination": OffsetPaginationExtension(),
 }
 
+# item_collection extensions
+itm_col_extensions_map = {
+    "filter": FilterExtension(client=FiltersClient()),
+    "pagination": TokenPaginationExtension(),
+}
+
+known_extensions = {
+    *application_extensions_map.keys(),
+    *search_extensions_map.keys(),
+    *cs_extensions_map.keys(),
+    *itm_col_extensions_map.keys(),
+    "collection_search",
+}
+
 enabled_extensions = (
     os.environ["ENABLED_EXTENSIONS"].split(",")
     if "ENABLED_EXTENSIONS" in os.environ
-    else list(extensions_map.keys()) + ["collection_search"]
+    else known_extensions
 )
-extensions = [
-    extension for key, extension in extensions_map.items() if key in enabled_extensions
+
+application_extensions = [
+    extension
+    for key, extension in application_extensions_map.items()
+    if key in enabled_extensions
 ]
 
-items_get_request_model = (
-    create_request_model(
+# /search models
+search_extensions = [
+    extension
+    for key, extension in search_extensions_map.items()
+    if key in enabled_extensions
+]
+post_request_model = create_post_request_model(search_extensions, base_model=PgstacSearch)
+get_request_model = create_get_request_model(search_extensions)
+application_extensions.extend(search_extensions)
+
+# /collections/{collectionId}/items model
+items_get_request_model = ItemCollectionUri
+itm_col_extensions = [
+    extension
+    for key, extension in itm_col_extensions_map.items()
+    if key in enabled_extensions
+]
+if itm_col_extensions:
+    items_get_request_model = create_request_model(
         model_name="ItemCollectionUri",
         base_model=ItemCollectionUri,
-        mixins=[TokenPaginationExtension().GET],
+        extensions=itm_col_extensions,
         request_type="GET",
     )
-    if any(isinstance(ext, TokenPaginationExtension) for ext in extensions)
-    else ItemCollectionUri
-)
-
-collection_search_extension = (
-    CollectionSearchExtension.from_extensions(
-        [
-            extension
-            for key, extension in collection_extensions_map.items()
-            if key in enabled_extensions
-        ]
-    )
-    if "collection_search" in enabled_extensions
-    else None
-)
-
-collections_get_request_model = (
-    collection_search_extension.GET if collection_search_extension else EmptyRequest
-)
 
-post_request_model = create_post_request_model(extensions, base_model=PgstacSearch)
-get_request_model = create_get_request_model(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
+    ]
+    collection_search_extension = CollectionSearchExtension.from_extensions(cs_extensions)
+    collections_get_request_model = collection_search_extension.GET
+    application_extensions.append(collection_search_extension)
 
 
 @asynccontextmanager
@@ -127,9 +157,7 @@ async def lifespan(app: FastAPI):
 api = StacApi(
     app=update_openapi(fastapp),
     settings=settings,
-    extensions=extensions + [collection_search_extension]
-    if collection_search_extension
-    else extensions,
+    extensions=application_extensions,
     client=CoreCrudClient(pgstac_search_model=post_request_model),
     response_class=ORJSONResponse,
     items_get_request_model=items_get_request_model,

diff --git a/stac_fastapi/pgstac/core.py b/stac_fastapi/pgstac/core.py
@@ -342,7 +342,10 @@ async def item_collection(
         bbox: Optional[BBox] = None,
         datetime: Optional[str] = None,
         limit: Optional[int] = None,
+        # Extensions
         token: Optional[str] = None,
+        filter_expr: Optional[str] = None,
+        filter_lang: Optional[str] = None,
         **kwargs,
     ) -> ItemCollection:
         """Get all items from a specific collection.
@@ -368,21 +371,11 @@ async def item_collection(
             "token": token,
         }
 
-        if self.extension_is_enabled("FilterExtension"):
-            filter_lang = kwargs.get("filter_lang", None)
-            filter_query = kwargs.get("filter_expr", None)
-            if filter_query:
-                if filter_lang == "cql2-text":
-                    filter_query = to_cql2(parse_cql2_text(filter_query))
-                    filter_lang = "cql2-json"
-
-                base_args["filter"] = orjson.loads(filter_query)
-                base_args["filter-lang"] = filter_lang
-
-        clean = {}
-        for k, v in base_args.items():
-            if v is not None and v != []:
-                clean[k] = v
+        clean = self._clean_search_args(
+            base_args=base_args,
+            filter_query=filter_expr,
+            filter_lang=filter_lang,
+        )
 
         search_request = self.pgstac_search_model(**clean)
         item_collection = await self._search_base(search_request, request=request)

diff --git a/tests/conftest.py b/tests/conftest.py
@@ -134,15 +134,19 @@ def api_client(request, database):
         )
     )
 
-    extensions = [
+    application_extensions = [
         TransactionExtension(client=TransactionsClient(), settings=api_settings),
+        BulkTransactionExtension(client=BulkTransactionsClient()),
+    ]
+
+    search_extensions = [
         QueryExtension(),
         SortExtension(),
         FieldsExtension(),
-        TokenPaginationExtension(),
         FilterExtension(client=FiltersClient()),
-        BulkTransactionExtension(client=BulkTransactionsClient()),
+        TokenPaginationExtension(),
     ]
+    application_extensions.extend(search_extensions)
 
     collection_extensions = [
         QueryExtension(),
@@ -155,26 +159,31 @@ def api_client(request, database):
     collection_search_extension = CollectionSearchExtension.from_extensions(
         collection_extensions
     )
+    application_extensions.append(collection_search_extension)
+
+    item_collection_extensions = [
+        FilterExtension(client=FiltersClient()),
+        TokenPaginationExtension(),
+    ]
+    # NOTE: we don't need to add the extensions to application_extensions
+    # because they are already in it
 
     items_get_request_model = create_request_model(
         model_name="ItemCollectionUri",
         base_model=ItemCollectionUri,
-        mixins=[
-            TokenPaginationExtension().GET,
-            FilterExtension(client=FiltersClient()).GET,
-        ],
+        extensions=item_collection_extensions,
         request_type="GET",
     )
-    search_get_request_model = create_get_request_model(extensions)
+    search_get_request_model = create_get_request_model(search_extensions)
     search_post_request_model = create_post_request_model(
-        extensions, base_model=PgstacSearch
+        search_extensions, base_model=PgstacSearch
     )
 
     collections_get_request_model = collection_search_extension.GET
 
     api = StacApi(
         settings=api_settings,
-        extensions=extensions + [collection_search_extension],
+        extensions=application_extensions,
         client=CoreCrudClient(pgstac_search_model=search_post_request_model),
         items_get_request_model=items_get_request_model,
         search_get_request_model=search_get_request_model,