Reorg utils

Author: alukach

src/stac_auth_proxy/app.py

@@ -56,6 +56,22 @@ def create_app(settings: Optional[Settings] = None) -> FastAPI:
         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

@@ -18,7 +18,7 @@ class ClassInput(BaseModel):
     kwargs: dict[str, str] = Field(default_factory=dict)
 
     def __call__(self, token_dependency):
-        """Dynamically load a class and instantiate it with kwargs."""
+        """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)
@@ -49,6 +49,10 @@ class Settings(BaseSettings):
     openapi_spec_endpoint: Optional[str] = None
 
     collections_filter: Optional[ClassInput] = None
+    collections_filter_endpoints: Optional[EndpointMethods] = {
+        "/collections": ["GET"],
+        "/collections/{collection_id}": ["GET"],
+    }
     items_filter: Optional[ClassInput] = None
 
     model_config = SettingsConfigDict(env_prefix="STAC_AUTH_PROXY_")

src/stac_auth_proxy/filters/template.py

@@ -6,7 +6,7 @@
 from fastapi import Request, Security
 from jinja2 import BaseLoader, Environment
 
-from ..utils import extract_variables
+from ..utils.requests import extract_variables
 
 
 def Template(template_str: str, token_dependency: Callable[..., Any]):

src/stac_auth_proxy/handlers/open_api_spec.py

@@ -5,7 +5,8 @@
 from fastapi import Request, Response
 from fastapi.routing import APIRoute
 
-from ..utils import has_any_security_requirements, safe_headers
+from ..utils.di import has_any_security_requirements
+from ..utils.requests import safe_headers
 from .reverse_proxy import ReverseProxyHandler
 
 logger = logging.getLogger(__name__)

src/stac_auth_proxy/handlers/reverse_proxy.py

@@ -12,7 +12,7 @@
 from starlette.datastructures import MutableHeaders
 from starlette.responses import StreamingResponse
 
-from .. import utils
+from ..utils import filters
 
 logger = logging.getLogger(__name__)
 
@@ -23,6 +23,8 @@ class ReverseProxyHandler:
 
     upstream: str
     client: httpx.AsyncClient = None
+
+    # Filters
     collections_filter: Optional[Callable] = None
     items_filter: Optional[Callable] = None
 
@@ -55,13 +57,13 @@ async def proxy_request(
         path = request.url.path
         query = request.url.query
 
-        # Appliy filters
-        if utils.is_collection_endpoint(path) and collections_filter:
+        # Apply filters
+        if filters.is_collection_endpoint(path) and collections_filter:
             if request.method == "GET" and path == "/collections":
-                query = utils.insert_filter(qs=query, filter=collections_filter)
-        elif utils.is_item_endpoint(path) and self.items_filter:
+                query = filters.insert_filter(qs=query, filter=collections_filter)
+        elif filters.is_item_endpoint(path) and self.items_filter:
             if request.method == "GET":
-                query = utils.insert_filter(qs=query, filter=self.items_filter)
+                query = filters.insert_filter(qs=query, filter=self.items_filter)
 
         # https://github.com/fastapi/fastapi/discussions/7382#discussioncomment-5136466
         rp_req = self.client.build_request(

src/stac_auth_proxy/utils.py

@@ -0,0 +1,13 @@
+from fastapi.dependencies.models import Dependant
+
+
+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(
+        has_any_security_requirements(sub_dep) for sub_dep in dependency.dependencies
+    )

src/stac_auth_proxy/utils/__init__.py

@@ -0,0 +1,32 @@
+"""Utility functions."""
+
+from urllib.parse import parse_qs, urlencode
+
+from cql2 import Expr
+
+
+def insert_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)
+
+    filters = [Expr(f) for f in qs_dict.get("filter", [])]
+    filters.append(filter)
+
+    combined_filter = Expr(" AND ".join(e.to_text() for e in filters))
+    combined_filter.validate()
+
+    qs_dict["filter"] = [combined_filter.to_text()]
+
+    return urlencode(qs_dict, doseq=True)
+
+
+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
+    return path == "/collections"
+
+
+def is_item_endpoint(path: str) -> bool:
+    """Check if the path is an item endpoint."""
+    # TODO: Expand this to cover all cases where an item filter should be applied
+    return path == "/collection/{collection_id}/items"

src/stac_auth_proxy/utils/di.py

@@ -0,0 +1,29 @@
+import re
+
+from urllib.parse import urlparse
+from httpx import Headers
+
+
+def safe_headers(headers: Headers) -> dict[str, str]:
+    """Scrub headers that should not be proxied to the client."""
+    excluded_headers = [
+        "content-length",
+        "content-encoding",
+    ]
+    return {
+        key: value
+        for key, value in headers.items()
+        if key.lower() not in excluded_headers
+    }
+
+
+def extract_variables(url: str) -> dict:
+    """
+    Extract variables from a URL path. Being that we use a catch-all endpoint for the proxy,
+    we can't rely on the path parameters that FastAPI provides.
+    """
+    path = urlparse(url).path
+    # This allows either /items or /bulk_items, with an optional item_id following.
+    pattern = r"^/collections/(?P<collection_id>[^/]+)(?:/(?:items|bulk_items)(?:/(?P<item_id>[^/]+))?)?/?$"
+    match = re.match(pattern, path)
+    return {k: v for k, v in match.groupdict().items() if v} if match else {}