refactor extension map and add filter extension for item_collection

Author: vincentsarago

stac_fastapi/pgstac/app.py

@@ -43,22 +43,28 @@
 from stac_fastapi.pgstac.types.search import PgstacSearch
 
 settings = Settings()
-extensions_map = {
+
+# transaction extensions
+trans_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 = {
+    *trans_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 trans_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,

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)