Refactor to pass OAuth and token decode methods consistently

Author: DiamondJoseph

example_configs/external_service/custom.py

@@ -1,3 +1,5 @@
+from typing import Any, Self
+
 import numpy
 
 from tiled.adapters.array import ArrayAdapter
@@ -55,7 +57,7 @@ def __init__(self, base_url, metadata=None):
         self.client = MockClient(base_url)
         self.metadata = metadata
 
-    def with_session_state(self, state):
+    def with_session_state(self, state: dict[str, Any]) -> Self:
         return AuthenticatedAdapter(self.client, state["token"], metadata=self.metadata)
 
 

tiled/authn_database/connection_pool.py

@@ -1,7 +1,7 @@
 from fastapi import Depends
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 
-from ..server.settings import get_settings
+from ..server.settings import Settings, get_settings
 from ..utils import ensure_specified_sql_driver
 
 # A given process probably only has one of these at a time, but we
@@ -31,7 +31,7 @@ async def close_database_connection_pool(database_settings):
         await engine.dispose()
 
 
-async def get_database_engine(settings=Depends(get_settings)):
+async def get_database_engine(settings: Settings = Depends(get_settings)):
     # Special case for single-user mode
     if settings.database_uri is None:
         return None

tiled/server/app.py

@@ -411,18 +411,20 @@ async def unhandled_exception_handler(
             token_decoder, authenticators, oauth2_scheme
         )
         get_current_principal = current_principal_getter(
+            token_decoder,
             authenticators,
             oauth2_scheme,
-            token_decoder,
         )
 
         # And add this authentication_router itself to the app.
         app.include_router(authentication_router, prefix="/api/v1/auth")
+        get_session_state = session_state_getter(token_decoder, oauth2_scheme)
 
     else:
         get_current_principal = get_current_principal_from_api_key
 
-    get_session_state = session_state_getter(token_decoder)
+        def get_session_state():
+            return None
 
     app.include_router(
         get_router(

tiled/server/authentication.py

@@ -270,12 +270,13 @@ async def token_from_request(
 
 def session_state_getter(
     token_decoder: Callable[[str], Awaitable[Optional[dict[str, Any]]]],
+    oauth2_scheme: OAuth2,
 ):
-    async def get_session_state(
-        decoded_access_token: Optional[dict[str, Any]] = Depends(token_decoder)
-    ):
-        if decoded_access_token:
-            return decoded_access_token.get("state")
+    async def get_session_state(access_token: Optional[str] = Depends(oauth2_scheme)):
+        if access_token:
+            decoded_access_token = await token_decoder(access_token)
+            if decoded_access_token:
+                return decoded_access_token.get("state")
 
     return get_session_state
 

tiled/server/dependencies.py

@@ -1,4 +1,4 @@
-from typing import Optional, Tuple, Union
+from typing import Any, Callable, Mapping, Optional, Tuple, Union
 
 from fastapi import Depends, HTTPException, Query, Request, Security
 from starlette.status import HTTP_403_FORBIDDEN, HTTP_404_NOT_FOUND
@@ -13,13 +13,17 @@
 SLICE_REGEX = rf"^{DIM_REGEX}(?:,{DIM_REGEX})*$"
 
 
-def SecureEntryBuilder(get_current_principal, tree, get_session_state):
+def SecureEntryBuilder(
+    tree: Mapping[str, Any],
+    get_current_principal: Callable[..., Optional[str]],
+    get_session_state: Callable[..., Optional[dict[str, Any]]],
+):
     def SecureEntry(scopes, structure_families=None):
         async def inner(
             path: str,
             request: Request,
-            principal: str = Depends(get_current_principal),
-            session_state: dict = Depends(get_session_state),
+            principal: Optional[str] = Depends(get_current_principal),
+            session_state: Optional[dict[str, Any]] = Depends(get_session_state),
         ):
             """
             Obtain a node in the tree from its path.

tiled/server/router.py

@@ -77,7 +77,7 @@ def get_router(
     validation_registry: ValidationRegistry,
 ) -> APIRouter:
     router = APIRouter()
-    SecureEntry = SecureEntryBuilder(get_current_principal, tree, get_session_state)
+    SecureEntry = SecureEntryBuilder(tree, get_current_principal, get_session_state)
 
     @router.get("/", response_model=About)
     async def about(