refactor: convert InsertedJob and TokenPayload from TypedDict to BaseModel

TypedDict fields are not validated by pydantic at runtime, so
UTCDatetime annotations had no effect. Converting to BaseModel
ensures datetime fields are actually validated on construction.

Author: ryuwd

diracx-core/src/diracx/core/models/auth.py

@@ -5,7 +5,7 @@
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
-from ._types import UTCDatetime
+from .types import UTCDatetime
 
 
 class UserInfo(BaseModel):
@@ -56,7 +56,7 @@ class OpenIDConfiguration(TypedDict):
     code_challenge_methods_supported: list[str]
 
 
-class TokenPayload(TypedDict):
+class TokenPayload(BaseModel):
     jti: str
     exp: UTCDatetime
     dirac_policies: dict

diracx-core/src/diracx/core/models/job.py

@@ -9,12 +9,11 @@
 from typing import Literal
 
 from pydantic import BaseModel, Field, field_validator
-from typing_extensions import TypedDict
 
-from ._types import UTCDatetime
+from .types import UTCDatetime
 
 
-class InsertedJob(TypedDict):
+class InsertedJob(BaseModel):
     JobID: int
     Status: str
     MinorStatus: str

…cx-core/src/diracx/core/models/_types.py …acx-core/src/diracx/core/models/types.pydiracx-core/src/diracx/core/models/_types.py renamed to diracx-core/src/diracx/core/models/types.py diracx-core/src/diracx/core/models/_types.py renamed to diracx-core/src/diracx/core/models/types.py

@@ -11,7 +11,7 @@
 from pydantic import BaseModel, ValidationError
 
 import diracx.core.models
-from diracx.core.models._types import UTCDatetime, _validate_utc
+from diracx.core.models.types import UTCDatetime, _validate_utc
 
 
 class SampleModel(BaseModel):
@@ -80,7 +80,7 @@ def _collect_model_classes() -> list[type[BaseModel]]:
     for _importer, modname, _ispkg in pkgutil.walk_packages(
         package.__path__, prefix=package.__name__ + "."
     ):
-        if modname.endswith("._types"):
+        if modname.endswith(".types"):
             continue
         module = importlib.import_module(modname)
         for _name, obj in inspect.getmembers(module, inspect.isclass):

diracx-core/tests/test_utc_datetime.py

@@ -332,14 +332,14 @@ async def exchange_token(
         refresh_exp = uuid7_to_datetime(refresh_jti) + timedelta(
             minutes=refresh_token_expire_minutes
         )
-        refresh_payload = {
-            "jti": str(refresh_jti),
-            "exp": refresh_exp,
+        refresh_payload = RefreshTokenPayload(
+            jti=str(refresh_jti),
+            exp=refresh_exp,
             # legacy_exchange is used to indicate that the original refresh token
             # was obtained from the legacy_exchange endpoint
-            "legacy_exchange": legacy_exchange,
-            "dirac_policies": {},
-        }
+            legacy_exchange=legacy_exchange,
+            dirac_policies={},
+        )
 
     # Generate access token payload
     # For now, the access token is only used to access DIRAC services,
@@ -348,22 +348,22 @@ async def exchange_token(
     access_exp = uuid7_to_datetime(access_jti) + timedelta(
         minutes=settings.access_token_expire_minutes
     )
-    access_payload: AccessTokenPayload = {
-        "sub": sub,
-        "vo": vo,
-        "iss": settings.token_issuer,
-        "dirac_properties": list(properties),
-        "jti": str(access_jti),
-        "preferred_username": preferred_username,
-        "dirac_group": dirac_group,
-        "exp": access_exp,
-        "dirac_policies": {},
-    }
+    access_payload = AccessTokenPayload(
+        sub=sub,
+        vo=vo,
+        iss=settings.token_issuer,
+        dirac_properties=list(properties),
+        jti=str(access_jti),
+        preferred_username=preferred_username,
+        dirac_group=dirac_group,
+        exp=access_exp,
+        dirac_policies={},
+    )
 
     return access_payload, refresh_payload
 
 
-def create_token(payload: TokenPayload, settings: AuthSettings) -> str:
+def create_token(payload: TokenPayload | dict, settings: AuthSettings) -> str:
     """Create a JWT token with the given payload and settings."""
     signing_key = None
     for key in settings.token_keystore.jwks.keys:
@@ -377,9 +377,10 @@ def create_token(payload: TokenPayload, settings: AuthSettings) -> str:
     if not signing_key:
         raise ValueError("No signing key found in JWKS")
 
+    claims = payload.model_dump() if isinstance(payload, TokenPayload) else payload
     return jwt.encode(
         header={"alg": signing_key.get("alg"), "kid": signing_key.get("kid")},
-        claims=cast(Claims, payload),
+        claims=cast(Claims, claims),
         key=settings.token_keystore.jwks,
         algorithms=settings.token_allowed_algorithms,
     )

diracx-logic/src/diracx/logic/auth/token.py

@@ -61,12 +61,12 @@ async def mint_token(
             dirac_refresh_policies[policy_name] = refresh_extra
 
     # Create the access token
-    access_payload["dirac_policies"] = dirac_access_policies
+    access_payload.dirac_policies = dirac_access_policies
     access_token = create_token(access_payload, settings)
 
     # Create the refresh token
     if refresh_payload:
-        refresh_payload["dirac_policies"] = dirac_refresh_policies
+        refresh_payload.dirac_policies = dirac_refresh_policies
         refresh_token = create_token(refresh_payload, settings)
     elif existing_refresh_token:
         refresh_token = existing_refresh_token

diracx-routers/src/diracx/routers/auth/token.py

@@ -35,17 +35,17 @@ def main(token_keystore: str):
     jti = uuid7()
     expires_at = uuid7_to_datetime(jti) + timedelta(seconds=expires_in)
 
-    access_payload: AccessTokenPayload = {
-        "sub": f"{vo}:{sub}",
-        "vo": vo,
-        "iss": settings.token_issuer,
-        "dirac_properties": dirac_properties,
-        "jti": str(jti),
-        "preferred_username": preferred_username,
-        "dirac_group": dirac_group,
-        "exp": expires_at,
-        "dirac_policies": {},
-    }
+    access_payload = AccessTokenPayload(
+        sub=f"{vo}:{sub}",
+        vo=vo,
+        iss=settings.token_issuer,
+        dirac_properties=dirac_properties,
+        jti=str(jti),
+        preferred_username=preferred_username,
+        dirac_group=dirac_group,
+        exp=expires_at,
+        dirac_policies={},
+    )
     token = TokenResponse(
         access_token=create_token(access_payload, settings),
         expires_in=expires_in,