CQL2 Middleware: in progress

Author: alukach

src/stac_auth_proxy/app.py

@@ -56,12 +56,7 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
         )
 
     # Tooling
-    proxy_handler = ReverseProxyHandler(
-        upstream=str(settings.upstream_url),
-        # TODO: Refactor filter tooling into middleare
-        collections_filter=settings.collections_filter,
-        items_filter=settings.items_filter,
-    )
+    proxy_handler = ReverseProxyHandler(upstream=str(settings.upstream_url))
     # Catchall for remainder of the endpoints
     app.add_api_route(
         "/{path:path}",

src/stac_auth_proxy/handlers/reverse_proxy.py

@@ -1,19 +1,15 @@
 """Tooling to manage the reverse proxying of requests to an upstream STAC API."""
 
-import json
 import logging
 import time
-from dataclasses import dataclass
-from typing import Annotated, Callable, Optional
+from dataclasses import dataclass, field
 
 import httpx
-from cql2 import Expr
-from fastapi import Depends, Request
+from fastapi import Request
 from starlette.background import BackgroundTask
 from starlette.datastructures import MutableHeaders
 from starlette.responses import StreamingResponse
 
-from ..utils import di, filters
 
 logger = logging.getLogger(__name__)
 
@@ -24,79 +20,35 @@ class ReverseProxyHandler:
 
     upstream: str
     client: httpx.AsyncClient = None
-
-    # Filters
-    collections_filter: Optional[Callable] = None
-    items_filter: Optional[Callable] = None
+    timeout: httpx.Timeout = field(default_factory=lambda: httpx.Timeout(timeout=15.0))
 
     def __post_init__(self):
         """Initialize the HTTP client."""
         self.client = self.client or httpx.AsyncClient(
             base_url=self.upstream,
-            timeout=httpx.Timeout(timeout=15.0),
-        )
-        self.collections_filter = (
-            self.collections_filter() if self.collections_filter else None
+            timeout=self.timeout,
         )
-        self.items_filter = self.items_filter() if self.items_filter else None
 
-    def _get_filter(self, path: str) -> Optional[Callable[..., Expr]]:
-        """Get the CQL2 filter builder for the given path."""
-        endpoint_filters = [
-            # TODO: Use collections_filter_endpoints & items_filter_endpoints
-            (filters.is_collection_endpoint, self.collections_filter),
-            (filters.is_item_endpoint, self.items_filter),
-            (filters.is_search_endpoint, self.items_filter),
-        ]
-        for check, builder in endpoint_filters:
-            if check(path):
-                return builder
-        return None
-
-    async def proxy_request(self, request: Request, *, stream=False) -> httpx.Response:
+    async def proxy_request(self, request: Request) -> httpx.Response:
         """Proxy a request to the upstream STAC API."""
         headers = MutableHeaders(request.headers)
         headers.setdefault("X-Forwarded-For", request.client.host)
         headers.setdefault("X-Forwarded-Host", request.url.hostname)
 
-        path = request.url.path
-        query = request.url.query
-        # TODO: Should we only do this conditionally based on stream?
-        body = (await request.body()).decode()
-
-        # Apply filter if applicable
-        filter_builder = self._get_filter(path)
-        if filter_builder:
-            cql_filter = await di.call_with_injected_dependencies(
-                func=filter_builder,
-                request=request,
-            )
-            cql_filter.validate()
-
-            if request.method == "GET":
-                query = filters.insert_filter(qs=query, filter=cql_filter)
-            elif request.method in ["POST", "PUT"]:
-                body_dict = json.loads(body)
-                body_filter = body_dict.get("filter")
-                if body_filter:
-                    cql_filter = cql_filter + Expr(body_filter)
-                body_dict["filter"] = cql_filter.to_json()
-                body = json.dumps(body_dict)
-
         # https://github.com/fastapi/fastapi/discussions/7382#discussioncomment-5136466
         rp_req = self.client.build_request(
             request.method,
             url=httpx.URL(
-                path=path,
-                query=query.encode("utf-8"),
+                path=request.url.path,
+                query=request.url.query.encode("utf-8"),
             ),
             headers=headers,
-            content=body,
+            content=request.stream(),
         )
         logger.debug(f"Proxying request to {rp_req.url}")
 
         start_time = time.perf_counter()
-        rp_resp = await self.client.send(rp_req, stream=stream)
+        rp_resp = await self.client.send(rp_req, stream=True)
         proxy_time = time.perf_counter() - start_time
 
         logger.debug(
@@ -107,11 +59,7 @@ async def proxy_request(self, request: Request, *, stream=False) -> httpx.Respon
 
     async def stream(self, request: Request) -> StreamingResponse:
         """Transparently proxy a request to the upstream STAC API."""
-        rp_resp = await self.proxy_request(
-            request,
-            # collections_filter=collections_filter,
-            stream=True,
-        )
+        rp_resp = await self.proxy_request(request)
         return StreamingResponse(
             rp_resp.aiter_raw(),
             status_code=rp_resp.status_code,

src/stac_auth_proxy/middleware/Cql2FilterMiddleware.py

@@ -0,0 +1,105 @@
+import json
+from dataclasses import dataclass
+from typing import Annotated, Callable, Optional
+
+from cql2 import Expr
+from starlette.types import ASGIApp, Message, Receive, Scope, Send
+from starlette.requests import Request
+
+from ..config import EndpointMethods
+from ..utils import di, filters, requests
+
+
+FILTER_STATE_KEY = "cql2_filter"
+
+
+@dataclass(frozen=True)
+class BuildCql2FilterMiddleware:
+    """Middleware to build the Cql2Filter."""
+
+    app: ASGIApp
+
+    # Filters
+    collections_filter: Optional[Callable] = None
+    items_filter: Optional[Callable] = None
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope["type"] != "http":
+            return await self.app(scope, receive, send)
+
+        request = Request(scope)
+        filter_builder = self._get_filter(request.url.path)
+        if filter_builder:
+            cql2_filter = await di.call_with_injected_dependencies(
+                func=filter_builder,
+                request=request,
+            )
+            cql2_filter.validate()
+            scope["state"][FILTER_STATE_KEY] = cql2_filter
+
+        return await self.app(scope, receive, send)
+
+    def _get_filter(self, path: str) -> Optional[Callable[..., Expr]]:
+        """Get the CQL2 filter builder for the given path."""
+        endpoint_filters = [
+            # TODO: Use collections_filter_endpoints & items_filter_endpoints
+            (filters.is_collection_endpoint, self.collections_filter),
+            (filters.is_item_endpoint, self.items_filter),
+            (filters.is_search_endpoint, self.items_filter),
+        ]
+        for check, builder in endpoint_filters:
+            if check(path):
+                return builder
+        return None
+
+
+@dataclass(frozen=True)
+class ApplyCql2FilterMiddleware:
+    """Middleware to add the OpenAPI spec to the response."""
+
+    app: ASGIApp
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        """Add the Cql2Filter to the request."""
+        request = Request(scope)
+        cql2_filter = request.state.get(FILTER_STATE_KEY)
+
+        if scope["type"] != "http" or not cql2_filter:
+            return await self.app(scope, receive, send)
+
+        # Apply filter if applicable
+
+        total_body = b""
+
+        async def receive_with_filter(message: Message):
+            query = request.url.query
+
+            # TODO: How do we handle querystrings in middleware?
+            if request.method == "GET":
+                query = filters.insert_qs_filter(qs=query, filter=cql2_filter)
+
+            if message["type"] == "http.response.body":
+                nonlocal total_body
+                total_body += message["body"]
+                if message["more_body"]:
+                    return await receive({**message, "body": b""})
+
+                # TODO: Only on search, not on create or update...
+                if request.method in ["POST", "PUT"]:
+                    return await receive(
+                        {
+                            "type": "http.response.body",
+                            "body": requests.dict_to_bytes(
+                                filters.append_body_filter(
+                                    json.loads(total_body), cql2_filter
+                                )
+                            ),
+                            "more_body": False,
+                        }
+                    )
+
+                return await receive(message)
+
+            await receive(message)
+
+        return await self.app(scope, receive_with_filter, send)

src/stac_auth_proxy/utils/filters.py

@@ -6,7 +6,7 @@
 from cql2 import Expr
 
 
-def insert_filter(qs: str, filter: Expr) -> str:
+def append_qs_filter(qs: str, filter: Expr) -> str:
     """Insert a filter expression into a query string. If a filter already exists, combine them."""
     qs_dict = parse_qs(qs)
 
@@ -19,6 +19,16 @@ def insert_filter(qs: str, filter: Expr) -> str:
     return urlencode(qs_dict, doseq=True)
 
 
+def append_body_filter(body: dict, filter: Expr) -> dict:
+    """Insert a filter expression into a request body. If a filter already exists, combine them."""
+    cur_filter = body.get("filter")
+    if cur_filter:
+        filter = filter + Expr(cur_filter)
+    body["filter"] = filter.to_json()
+    body["filter-lang"] = "cql2-json"
+    return body
+
+
 def is_collection_endpoint(path: str) -> bool:
     """Check if the path is a collection endpoint."""
     # TODO: Expand this to cover all cases where a collection filter should be applied