fix: move cleanup function and minor log fixes

Author: HeloiseJoffe

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

@@ -2,6 +2,8 @@
 
 from __future__ import annotations
 
+import logging
+
 from uuid_utils import UUID
 
 from diracx.core.exceptions import InvalidCredentialsError
@@ -10,6 +12,8 @@
 from diracx.db.sql import AuthDB
 from diracx.logic.auth.utils import verify_dirac_refresh_token
 
+logger = logging.getLogger(__name__)
+
 
 async def get_refresh_tokens(
     auth_db: AuthDB,
@@ -57,3 +61,24 @@ async def revoke_refresh_token_by_refresh_token(
     # Decode and verify the refresh token
     jti, _, _ = await verify_dirac_refresh_token(token, settings)
     return await revoke_refresh_token_by_jti(auth_db=auth_db, subject=subject, jti=jti)
+
+
+async def cleanup_expired_data(auth_db: AuthDB, settings: AuthSettings) -> None:
+    """Remove expired data from the auth database."""
+    expired_tokens, revoked_tokens = await auth_db.clean_expired_refresh_token(
+        max_validity=settings.refresh_token_expire_minutes,
+        max_retention=settings.revoked_refresh_token_retention_days,
+    )
+    logger.info(
+        f"Deleted {expired_tokens} expired and {revoked_tokens} revoked refresh tokens"
+    )
+
+    auth = await auth_db.clean_expired_authorization_flows(
+        max_retention=settings.completed_flow_retention_minutes,
+    )
+    logger.info(f"Deleted {auth} expired authorization flows")
+
+    device = await auth_db.clean_expired_device_flows(
+        max_retention=settings.completed_flow_retention_minutes,
+    )
+    logger.info(f"Deleted {device} expired device flows")

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

@@ -4,7 +4,6 @@
 
 import base64
 import hashlib
-import logging
 import re
 from datetime import datetime, timedelta, timezone
 from typing import cast
@@ -38,8 +37,6 @@
     verify_dirac_refresh_token,
 )
 
-logger = logging.getLogger(__name__)
-
 
 async def get_oidc_token(
     grant_type: GrantType,
@@ -432,8 +429,6 @@ async def get_authorization_flow(auth_db: AuthDB, code: str, max_validity: int):
     """Get the authorization flow from the DB and check few parameters before returning it."""
     res = await auth_db.get_authorization_flow(code, max_validity)
 
-    print(f"Flow status : {res['Status']}")
-
     if res["Status"] == FlowStatus.READY:
         await auth_db.update_authorization_flow_status(code, FlowStatus.DONE)
         return res
@@ -442,32 +437,3 @@ async def get_authorization_flow(auth_db: AuthDB, code: str, max_validity: int):
         raise AuthorizationError("Code was already used")
 
     raise AuthorizationError("Bad state in authorization flow")
-
-
-async def cleanup_expired_data(auth_db: AuthDB, settings: AuthSettings) -> None:
-    """Remove expired data from the auth database."""
-    expired_tokens, revoked_tokens = await auth_db.clean_expired_refresh_token(
-        max_validity=settings.refresh_token_expire_minutes,
-        max_retention=settings.revoked_refresh_token_retention_days,
-    )
-    logger.info(
-        "Deleted %d expired and %d revoked refresh tokens",
-        expired_tokens,
-        revoked_tokens,
-    )
-
-    auth = await auth_db.clean_expired_authorization_flows(
-        max_retention=settings.completed_flow_retention_minutes,
-    )
-    logger.info(
-        "Deleted %d expired authorization flows",
-        auth,
-    )
-
-    device = await auth_db.clean_expired_device_flows(
-        max_retention=settings.completed_flow_retention_minutes,
-    )
-    logger.info(
-        "Deleted %d expired device flows",
-        device,
-    )