Aggregation Extension

CHANGES.md

@@ -1,5 +1,11 @@
 # Changelog
 
+## [Unreleased]
+
+### Added
+
+* Add base support for the Aggregation extension [#684](https://github.com/stac-utils/stac-fastapi/pull/684)
+
 ## [3.0.0a0] - 2024-05-06
 
 ### Added

stac_fastapi/api/stac_fastapi/api/config.py

@@ -18,6 +18,7 @@ class ApiExtensions(enum.Enum):
     query = "query"
     sort = "sort"
     transaction = "transaction"
+    aggregation = "aggregation"
 
 
 class AddOns(enum.Enum):

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

@@ -1,5 +1,6 @@
 """stac_api.extensions.core module."""
 
+from .aggregation import AggregationExtension
 from .context import ContextExtension
 from .fields import FieldsExtension
 from .filter import FilterExtension
@@ -9,6 +10,7 @@
 from .transaction import TransactionExtension
 
 __all__ = (
+    "AggregationExtension",
     "ContextExtension",
     "FieldsExtension",
     "FilterExtension",

stac_fastapi/extensions/stac_fastapi/extensions/core/aggregation/__init__.py

@@ -0,0 +1,5 @@
+"""Aggregation extension module."""
+
+from .aggregation import AggregationExtension
+
+__all__ = ["AggregationExtension"]

stac_fastapi/extensions/stac_fastapi/extensions/core/aggregation/aggregation.py

@@ -0,0 +1,95 @@
+"""Aggregation Extension."""
+from enum import Enum
+from typing import List, Union
+
+import attr
+from fastapi import APIRouter, FastAPI
+
+from stac_fastapi.api.models import CollectionUri, EmptyRequest
+from stac_fastapi.api.routes import create_async_endpoint
+from stac_fastapi.types.core import AsyncBaseAggregationClient, BaseAggregationClient
+from stac_fastapi.types.extension import ApiExtension
+
+from .request import AggregationExtensionGetRequest, AggregationExtensionPostRequest
+
+
+class AggregationConformanceClasses(str, Enum):
+    """Conformance classes for the Aggregation extension.
+
+    See
+    https://github.com/stac-api-extensions/aggregation
+    """
+
+    AGGREGATION = "https://api.stacspec.org/v0.3.0/aggregation"
+
+
+@attr.s
+class AggregationExtension(ApiExtension):
+    """Aggregation Extension.
+
+    The purpose of the Aggregation Extension is to provide an endpoint similar to
+    the Search endpoint (/search), but which will provide aggregated information
+    on matching Items rather than the Items themselves. This is highly influenced
+    by the Elasticsearch and OpenSearch aggregation endpoint, but with a more
+    regular structure for responses.
+
+    The Aggregation extension adds several endpoints which allow the retrieval of
+    available aggregation fields and aggregation buckets based on a seearch query:
+        GET /aggregation
+        GET /aggregate
+        GET /collections/{collection_id}/aggregations
+        GET /collections/{collection_id}/aggregate
+
+    https://github.com/stac-api-extensions/aggregation/blob/main/README.md
+
+    Attributes:
+        conformance_classes: Conformance classes provided by the extension
+    """
+
+    GET = AggregationExtensionGetRequest
+    POST = AggregationExtensionPostRequest
+
+    client: Union[AsyncBaseAggregationClient, BaseAggregationClient] = attr.ib(
+        factory=BaseAggregationClient
+    )
+
+    conformance_classes: List[str] = attr.ib(
+        default=[AggregationConformanceClasses.AGGREGATION]
+    )
+    router: APIRouter = attr.ib(factory=APIRouter)
+
+    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="Aggregations",
+            path="/aggregations",
+            methods=["GET", "POST"],
+            endpoint=create_async_endpoint(self.client.get_aggregations, EmptyRequest),
+        )
+        self.router.add_api_route(
+            name="Collection Aggregations",
+            path="/collections/{collection_id}/aggregations",
+            methods=["GET", "POST"],
+            endpoint=create_async_endpoint(self.client.get_aggregations, CollectionUri),
+        )
+        self.router.add_api_route(
+            name="Aggregate",
+            path="/aggregate",
+            methods=["GET", "POST"],
+            endpoint=create_async_endpoint(self.client.aggregate, self.GET),
+        )
+        self.router.add_api_route(
+            name="Collection Aggregate",
+            path="/collections/{collection_id}/aggregate",
+            methods=["GET", "POST"],
+            endpoint=create_async_endpoint(self.client.aggregate, self.GET),
+        )
+        app.include_router(self.router, tags=["Aggregation Extension"])

stac_fastapi/extensions/stac_fastapi/extensions/core/aggregation/request.py

@@ -0,0 +1,31 @@
+"""Request model for the Aggregation extension."""
+
+from typing import Optional
+
+import attr
+from pydantic import BaseModel
+
+from stac_fastapi.types.search import APIRequest
+
+
+@attr.s
+class AggregationExtensionGetRequest(APIRequest):
+    """Aggregation Extension GET request model."""
+
+    aggregations: Optional[str] = attr.ib(default=None)
+    bbox: Optional[str] = attr.ib(default=None)
+    intersects: Optional[str] = attr.ib(default=None)
+    ids: Optional[str] = attr.ib(default=None)
+    datetime: Optional[str] = attr.ib(default=None)
+    collections: Optional[str] = attr.ib(default=None)
+
+
+class AggregationExtensionPostRequest(BaseModel):
+    """Aggregation Extension POST request model."""
+
+    aggregations: Optional[str] = attr.ib(default=None)
+    bbox: Optional[str] = attr.ib(default=None)
+    intersects: Optional[str] = attr.ib(default=None)
+    ids: Optional[str] = attr.ib(default=None)
+    datetime: Optional[str] = attr.ib(default=None)
+    collections: Optional[str] = attr.ib(default=None)

stac_fastapi/extensions/tests/test_aggregation.py

@@ -0,0 +1,107 @@
+from typing import Iterator
+
+import pytest
+from starlette.testclient import TestClient
+
+from stac_fastapi.api.app import StacApi
+from stac_fastapi.extensions.core import AggregationExtension
+from stac_fastapi.types.config import ApiSettings
+from stac_fastapi.types.core import BaseAggregationClient, BaseCoreClient
+
+
+class DummyCoreClient(BaseCoreClient):
+    def all_collections(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def get_collection(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def get_item(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def get_search(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def post_search(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def item_collection(self, *args, **kwargs):
+        raise NotImplementedError
+
+
+class DummyAggregationsClient(BaseAggregationClient):
+    """Dummy client returning parts of the request, rather than proper STAC items."""
+
+    def get_aggregations(self, *args, **kwargs):
+        return {
+            "aggregations": [{"name": "total_count", "data_type": "integer"}],
+            "links": [
+                {
+                    "rel": "root",
+                    "type": "application/json",
+                    "href": "https://example.org",
+                },
+                {
+                    "rel": "self",
+                    "type": "application/json",
+                    "href": "https://example.org/aggregations",
+                },
+            ],
+        }
+
+    def aggregate(self, *args, **kwargs):
+        return {
+            "aggregations": [],
+            "links": [
+                {
+                    "rel": "self",
+                    "type": "application/json",
+                    "href": "https://example.org/aggregate",
+                },
+                {
+                    "rel": "root",
+                    "type": "application/json",
+                    "href": "https://example.org",
+                },
+            ],
+        }
+
+
+def test_get_aggregations(client: TestClient) -> None:
+    response = client.get("/aggregations")
+    assert response.is_success, response.text
+    assert response.json()["aggregations"] == [
+        {"name": "total_count", "data_type": "integer"}
+    ]
+
+
+def test_aggregate(client: TestClient) -> None:
+    response = client.get("/aggregate")
+    assert response.is_success, response.text
+    assert response.json()["aggregations"] == []
+
+
+@pytest.fixture
+def client(
+    core_client: DummyCoreClient, aggregations_client: DummyAggregationsClient
+) -> Iterator[TestClient]:
+    settings = ApiSettings()
+    api = StacApi(
+        settings=settings,
+        client=core_client,
+        extensions=[
+            AggregationExtension(client=aggregations_client),
+        ],
+    )
+    with TestClient(api.app) as client:
+        yield client
+
+
+@pytest.fixture
+def core_client() -> DummyCoreClient:
+    return DummyCoreClient()
+
+
+@pytest.fixture
+def aggregations_client() -> DummyAggregationsClient:
+    return DummyAggregationsClient()