Initial pass at DI tooling

Author: alukach

pyproject.toml

@@ -43,6 +43,7 @@ requires = ["hatchling>=1.12.0"]
 dev = [
   "jwcrypto>=1.5.6",
   "pre-commit>=3.5.0",
+  "pytest-asyncio>=0.25.1",
   "pytest-cov>=5.0.0",
   "pytest>=8.3.3",
 ]

src/stac_auth_proxy/app.py

@@ -40,38 +40,14 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
 
     proxy_handler = ReverseProxyHandler(
         upstream=str(settings.upstream_url),
-        collections_filter=(
-            settings.collections_filter(auth_scheme.maybe_validated_user)
-            if settings.collections_filter
-            else None
-        ),
-        items_filter=(
-            settings.items_filter(auth_scheme.maybe_validated_user)
-            if settings.items_filter
-            else None
-        ),
+        auth_dependency=auth_scheme.maybe_validated_user,
+        collections_filter=settings.collections_filter,
+        items_filter=settings.items_filter,
     )
     openapi_handler = build_openapi_spec_handler(
         proxy=proxy_handler,
         oidc_config_url=str(settings.oidc_discovery_url),
     )
-
-    # TODO: How can we inject the collections_filter into only the endpoints that need it?
-    # for endpoint, methods in settings.collections_filter_endpoints.items():
-    #     app.add_api_route(
-    #         endpoint,
-    #         partial(
-    #             proxy_handler.stream, collections_filter=settings.collections_filter
-    #         ),
-    #         methods=methods,
-    #         dependencies=[Security(auth_scheme.maybe_validated_user)],
-    #     )
-
-    # skip = [
-    #     settings.openapi_spec_endpoint,
-    #     *settings.collections_filter_endpoints,
-    # ]
-
     # Endpoints that are explicitely marked private
     for path, methods in settings.private_endpoints.items():
         app.add_api_route(

src/stac_auth_proxy/config.py

@@ -17,12 +17,12 @@ class ClassInput(BaseModel):
     args: Sequence[str] = Field(default_factory=list)
     kwargs: dict[str, str] = Field(default_factory=dict)
 
-    def __call__(self, token_dependency):
+    def __call__(self):
         """Dynamically load a class and instantiate it with args & kwargs."""
         module_path, class_name = self.cls.rsplit(".", 1)
         module = importlib.import_module(module_path)
         cls = getattr(module, class_name)
-        return cls(*self.args, **self.kwargs, token_dependency=token_dependency)
+        return cls(*self.args, **self.kwargs)
 
 
 class Settings(BaseSettings):

src/stac_auth_proxy/filters/template.py

@@ -1,21 +1,21 @@
 """Generate CQL2 filter expressions via Jinja2 templating."""
 
-from typing import Annotated, Any, Callable
+from typing import Annotated, Any
 
 from cql2 import Expr
-from fastapi import Request, Security
+from fastapi import Request
 from jinja2 import BaseLoader, Environment
 
 from ..utils.requests import extract_variables
 
 
-def Template(template_str: str, token_dependency: Callable[..., Any]):
+def Template(template_str: str):
     """Generate CQL2 filter expressions via Jinja2 templating."""
     env = Environment(loader=BaseLoader).from_string(template_str)
 
     async def dependency(
         request: Request,
-        auth_token: Annotated[dict[str, Any], Security(token_dependency)],
+        auth_token: Annotated[dict[str, Any], ...],
     ) -> Expr:
         """Render a CQL2 filter expression with the request and auth token."""
         # TODO: How to handle the case where auth_token is null?

src/stac_auth_proxy/handlers/reverse_proxy.py

@@ -12,7 +12,7 @@
 from starlette.datastructures import MutableHeaders
 from starlette.responses import StreamingResponse
 
-from ..utils import filters
+from ..utils import di, filters
 
 logger = logging.getLogger(__name__)
 
@@ -22,6 +22,8 @@ class ReverseProxyHandler:
     """Reverse proxy functionality."""
 
     upstream: str
+    auth_dependency: Callable
+
     client: httpx.AsyncClient = None
 
     # Filters
@@ -34,21 +36,21 @@ def __post_init__(self):
             base_url=self.upstream,
             timeout=httpx.Timeout(timeout=15.0),
         )
+        self.collections_filter = (
+            self.collections_filter() if self.collections_filter else None
+        )
+        self.items_filter = self.items_filter() if self.items_filter else None
 
-        # Update annotations to support FastAPI's dependency injection
-        for endpoint in [self.proxy_request, self.stream]:
-            endpoint.__annotations__["collections_filter"] = Annotated[
+        # Inject auth dependency into filters
+        for endpoint in [self.collections_filter, self.items_filter]:
+            if not endpoint:
+                continue
+            endpoint.__annotations__["auth_token"] = Annotated[
                 Optional[Expr],
-                Depends(self.collections_filter or (lambda: None)),
+                Depends(self.auth_dependency),
             ]
 
-    async def proxy_request(
-        self,
-        request: Request,
-        *,
-        collections_filter: Annotated[Optional[Expr], Depends(...)] = None,
-        stream=False,
-    ) -> httpx.Response:
+    async def proxy_request(self, request: Request, *, stream=False) -> httpx.Response:
         """Proxy a request to the upstream STAC API."""
         headers = MutableHeaders(request.headers)
         headers.setdefault("X-Forwarded-For", request.client.host)
@@ -58,12 +60,23 @@ async def proxy_request(
         query = request.url.query
 
         # Apply filters
-        if filters.is_collection_endpoint(path) and collections_filter:
-            if request.method == "GET" and path == "/collections":
+        if filters.is_collection_endpoint(path) and self.collections_filter:
+            collections_filter = await di.call_with_injected_dependencies(
+                func=self.collections_filter,
+                request=request,
+            )
+            if request.method == "GET":
                 query = filters.insert_filter(qs=query, filter=collections_filter)
-        elif filters.is_item_endpoint(path) and self.items_filter:
+            else:
+                # TODO: Augment body
+                ...
+
+        if filters.is_item_endpoint(path) and self.items_filter:
             if request.method == "GET":
                 query = filters.insert_filter(qs=query, filter=self.items_filter)
+            else:
+                # TODO: Augment body
+                ...
 
         # https://github.com/fastapi/fastapi/discussions/7382#discussioncomment-5136466
         rp_req = self.client.build_request(
@@ -87,15 +100,11 @@ async def proxy_request(
         rp_resp.headers["X-Upstream-Time"] = f"{proxy_time:.3f}"
         return rp_resp
 
-    async def stream(
-        self,
-        request: Request,
-        collections_filter: Annotated[Optional[Expr], Depends(...)],
-    ) -> StreamingResponse:
+    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,
+            # collections_filter=collections_filter,
             stream=True,
         )
         return StreamingResponse(

src/stac_auth_proxy/utils/__init__.py

@@ -0,0 +1 @@
+"""Utils module for stac_auth_proxy."""

src/stac_auth_proxy/utils/di.py

@@ -1,13 +1,50 @@
+"""Dependency injection utilities for FastAPI."""
+
+import asyncio
+
+from fastapi import Request
 from fastapi.dependencies.models import Dependant
+from fastapi.dependencies.utils import (
+    get_parameterless_sub_dependant,
+    solve_dependencies,
+)
+from fastapi.params import Depends
 
 
 def has_any_security_requirements(dependency: Dependant) -> bool:
     """
     Recursively check if any dependency within the hierarchy has a non-empty
     security_requirements list.
     """
-    if dependency.security_requirements:
-        return True
-    return any(
+    return dependency.security_requirements or any(
         has_any_security_requirements(sub_dep) for sub_dep in dependency.dependencies
     )
+
+
+async def call_with_injected_dependencies(func, request: Request):
+    """
+    Manually solves and injects dependencies for `func` using FastAPI's internal
+    dependency injection machinery.
+    """
+    dependant = get_parameterless_sub_dependant(
+        depends=Depends(dependency=func),
+        path=request.url.path,
+    )
+
+    solved = await solve_dependencies(
+        request=request,
+        dependant=dependant,
+        # response=response,
+        # body=request.body,
+        body=None,
+        async_exit_stack=None,
+        embed_body_fields=False,
+    )
+
+    if solved.errors:
+        raise RuntimeError(f"Dependency resolution error: {solved.errors}")
+
+    results = func(**solved.values)
+    if asyncio.iscoroutine(results):
+        return await results
+    return results

src/stac_auth_proxy/utils/requests.py

@@ -1,6 +1,8 @@
-import re
+"""Utility functions for working with HTTP requests."""
 
+import re
 from urllib.parse import urlparse
+
 from httpx import Headers
 
 

tests/test_di.py

@@ -0,0 +1,56 @@
+"""Tests for the dependency injection utility."""
+
+from typing import Annotated
+
+import pytest
+from fastapi import Depends, Request
+
+from stac_auth_proxy.utils.di import call_with_injected_dependencies
+
+
+async def get_db_connection():
+    """Mock asynchronous function to get a DB connection."""
+    # pretend you open a DB connection or retrieve a session
+    return "some_db_connection"
+
+
+def get_special_value():
+    """Mock synchronous function to get a special value."""
+    return 42
+
+
+async def async_func_with_dependencies(
+    db_conn: Annotated[str, Depends(get_db_connection)],
+    special_value: Annotated[int, Depends(get_special_value)],
+):
+    """Mock asynchronous dependency."""
+    return (db_conn, special_value)
+
+
+def sync_func_with_dependencies(
+    db_conn: Annotated[str, Depends(get_db_connection)],
+    special_value: Annotated[int, Depends(get_special_value)],
+):
+    """Mock synchronous dependency."""
+    return (db_conn, special_value)
+
+
+@pytest.mark.parametrize(
+    "func",
+    [async_func_with_dependencies, sync_func_with_dependencies],
+)
+@pytest.mark.asyncio
+async def test_di(func):
+    """Test dependency injection."""
+    request = Request(
+        scope={
+            "type": "http",
+            "method": "GET",
+            "path": "/test",
+            "headers": [],
+            "query_string": b"",
+        }
+    )
+
+    result = await call_with_injected_dependencies(func, request=request)
+    assert result == ("some_db_connection", 42)

tests/test_filters_jinja2.py

@@ -2,7 +2,6 @@
 
 from urllib.parse import parse_qs
 
-import httpx
 import pytest
 from fastapi.testclient import TestClient
 from utils import AppFactory