refactor: split refresh tokens cleanup function, harmonize units and handle pending/ready flows

Author: HeloiseJoffe

diracx-core/src/diracx/core/settings.py

@@ -32,6 +32,7 @@
     SecretStr,
     TypeAdapter,
     UrlConstraints,
+    model_validator,
 )
 from pydantic_settings import BaseSettings, SettingsConfigDict
 
@@ -151,6 +152,21 @@ def create(cls) -> Self:
 class AuthSettings(ServiceSettingsBase):
     """Settings for the authentication service."""
 
+    @model_validator(mode="after")
+    def check_retention_greater_than_expiration(self) -> Self:
+        """Ensure retention times are bigger than expiration times to avoid deleting valid flows."""
+        if self.completed_flow_retention_minutes <= (
+            self.device_flow_expiration_seconds / 60
+        ) or self.completed_flow_retention_minutes <= (
+            self.authorization_flow_expiration_seconds / 60
+        ):
+            raise ValueError(
+                f"completed_flow_retention_minutes ({self.completed_flow_retention_minutes} minutes) must be bigger"
+                f" than device_flow_expiration_seconds ({self.device_flow_expiration_seconds / 60} minutes) and"
+                f" authorization_flow_expiration_seconds: ({self.authorization_flow_expiration_seconds / 60} minutes)"
+            )
+        return self
+
     model_config = SettingsConfigDict(
         env_prefix="DIRACX_SERVICE_AUTH_", use_attribute_docstrings=True
     )
@@ -233,11 +249,11 @@ class AuthSettings(ServiceSettingsBase):
     through a new authentication flow. Default: 60 minutes.
     """
 
-    revoked_refresh_token_retention_days: int = 30
-    """Retention time in days for revoked refresh tokens.
+    revoked_refresh_token_retention_minutes: int = 43200
+    """Retention time in minutes for revoked refresh tokens.
 
     The maximum retention time of refresh tokens after being
-    revoked and before they are deleted. Default: 30 days.
+    revoked and before they are deleted. Default: 43200 minutes (30 days).
     """
 
     available_properties: set[SecurityProperty] = Field(

diracx-db/src/diracx/db/sql/auth/db.py

@@ -340,10 +340,11 @@ async def revoke_user_refresh_tokens(self, subject):
             .values(status=RefreshTokenStatus.REVOKED)
         )
 
-    async def clean_expired_refresh_token(
-        self, max_validity: int, max_retention: int
-    ) -> tuple[int, int]:
-        """Delete expired and old revoked refresh tokens."""
+    async def clean_expired_refresh_tokens(self, max_validity: int) -> int:
+        """Delete expired refresh tokens.
+
+        max_validity: Maximum validity time in minutes for refresh tokens.
+        """
         expired_date = str(
             uuid7_from_datetime(substract_date(minutes=max_validity), randomize=False)
         )
@@ -353,31 +354,60 @@ async def clean_expired_refresh_token(
         )
         res_expired = await self.conn.execute(stmt_expired)
 
+        return res_expired.rowcount
+
+    async def clean_revoked_refresh_tokens(self, max_retention: int) -> int:
+        """Delete old revoked refresh tokens.
+
+        max_retention: Maximum retention time in minutes for revoked refresh tokens.
+        """
         revoked_date = str(
-            uuid7_from_datetime(substract_date(days=max_retention), randomize=False)
+            uuid7_from_datetime(substract_date(minutes=max_retention), randomize=False)
         )
         stmt_revoked = delete(RefreshTokens).where(
             RefreshTokens.status == RefreshTokenStatus.REVOKED,
             RefreshTokens.jti < revoked_date,
         )
         res_revoked = await self.conn.execute(stmt_revoked)
 
-        return res_expired.rowcount, res_revoked.rowcount
+        return res_revoked.rowcount
 
     async def clean_expired_authorization_flows(self, max_retention: int) -> int:
-        """Delete old DONE/ERROR authorization flows."""
+        """Delete old authorization flows.
+
+        max_retention: Maximum retention time in minutes for expired authorization flows.
+        Must be bigger than authorization_flow_expiration_seconds.
+        """
         stmt_auth = delete(AuthorizationFlows).where(
-            AuthorizationFlows.status.in_([FlowStatus.DONE, FlowStatus.ERROR]),
+            AuthorizationFlows.status.in_(
+                [
+                    FlowStatus.PENDING,
+                    FlowStatus.READY,
+                    FlowStatus.DONE,
+                    FlowStatus.ERROR,
+                ]
+            ),
             AuthorizationFlows.creation_time < substract_date(minutes=max_retention),
         )
         res_auth = await self.conn.execute(stmt_auth)
 
         return res_auth.rowcount
 
     async def clean_expired_device_flows(self, max_retention: int) -> int:
-        """Delete old DONE/ERROR device flows."""
+        """Delete old device flows.
+
+        max_retention: Maximum retention time in minutes for expired device flows.
+        Must be bigger than device_flow_expiration_seconds.
+        """
         stmt_device = delete(DeviceFlows).where(
-            DeviceFlows.status.in_([FlowStatus.DONE, FlowStatus.ERROR]),
+            DeviceFlows.status.in_(
+                [
+                    FlowStatus.PENDING,
+                    FlowStatus.READY,
+                    FlowStatus.DONE,
+                    FlowStatus.ERROR,
+                ]
+            ),
             DeviceFlows.creation_time < substract_date(minutes=max_retention),
         )
         res_device = await self.conn.execute(stmt_device)

diracx-db/tests/auth/test_authorization_flow.py

@@ -80,32 +80,39 @@ async def test_insert(auth_db: AuthDB):
 async def test_clean_authorization_flows(auth_db: AuthDB):
     # Insert two authorization flows
     async with auth_db as auth_db:
-        uuid1 = await auth_db.insert_authorization_flow(
+        _ = await auth_db.insert_authorization_flow(
             "client_id", "scope", "code_challenge", "S256", "redirect_uri"
         )
         uuid2 = await auth_db.insert_authorization_flow(
             "client_id2", "scope2", "code_challenge2", "S256", "redirect_uri2"
         )
+        uuid3 = await auth_db.insert_authorization_flow(
+            "client_id3", "scope3", "code_challenge3", "S256", "redirect_uri3"
+        )
+        uuid4 = await auth_db.insert_authorization_flow(
+            "client_id4", "scope4", "code_challenge4", "S256", "redirect_uri4"
+        )
 
     id_token = {"sub": "myIdToken"}
 
     async with auth_db as auth_db:
-        code1, _ = await auth_db.authorization_flow_insert_id_token(uuid1, id_token, 1)
-        code2, _ = await auth_db.authorization_flow_insert_id_token(uuid2, id_token, 1)
+        _, _ = await auth_db.authorization_flow_insert_id_token(uuid2, id_token, 1)
+        code3, _ = await auth_db.authorization_flow_insert_id_token(uuid3, id_token, 1)
+        code4, _ = await auth_db.authorization_flow_insert_id_token(uuid4, id_token, 1)
 
     async with auth_db as auth_db:
-        await auth_db.update_authorization_flow_status(code1, FlowStatus.DONE)
-        await auth_db.update_authorization_flow_status(code2, FlowStatus.ERROR)
+        await auth_db.update_authorization_flow_status(code3, FlowStatus.DONE)
+        await auth_db.update_authorization_flow_status(code4, FlowStatus.ERROR)
 
     # Check the number of deleted authorization flow (should be 0)
     async with auth_db as auth_db:
         deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=30)
         assert deleted_auth == 0
 
-    # Check the number of deleted authorization flow (should be 2)
+    # Check the number of deleted authorization flow (should be 4: 1 PENDING, 1 READY, 1 DONE, 1 ERROR)
     async with auth_db as auth_db:
         deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=0)
-        assert deleted_auth == 2
+        assert deleted_auth == 4
 
     # Check the number of deleted authorization flow (should be 0 because there is nothing left to delete)
     async with auth_db as auth_db:

diracx-db/tests/auth/test_device_flow.py

@@ -144,26 +144,28 @@ async def test_device_flow_insert_id_token(auth_db: AuthDB):
 async def test_clean_device_flows(auth_db: AuthDB):
     # Insert two device flows
     async with auth_db as auth_db:
-        user_code1, device_code1 = await auth_db.insert_device_flow(
-            "client_id", "scope"
-        )
-        user_code2, device_code2 = await auth_db.insert_device_flow(
-            "client_id", "scope"
-        )
+        _, _ = await auth_db.insert_device_flow("client_id1", "scope1")
+        user_code2, _ = await auth_db.insert_device_flow("client_id2", "scope2")
+        _, device_code3 = await auth_db.insert_device_flow("client_id3", "scope3")
+        _, device_code4 = await auth_db.insert_device_flow("client_id4", "scope4")
 
     async with auth_db as auth_db:
-        await auth_db.update_device_flow_status(device_code1, FlowStatus.DONE)
-        await auth_db.update_device_flow_status(device_code2, FlowStatus.ERROR)
+        await auth_db.device_flow_validate_user_code(user_code2, MAX_VALIDITY)
+        await auth_db.device_flow_insert_id_token(
+            user_code2, {"sub": "myIdToken"}, MAX_VALIDITY
+        )
+        await auth_db.update_device_flow_status(device_code3, FlowStatus.DONE)
+        await auth_db.update_device_flow_status(device_code4, FlowStatus.ERROR)
 
     # Check the number of deleted device flows (should be 0)
     async with auth_db as auth_db:
         deleted_device = await auth_db.clean_expired_device_flows(max_retention=30)
         assert deleted_device == 0
 
-    # Check the number of deleted device flows (should be 2)
+    # Check the number of deleted device flows (should be 4: 1 PENDING, 1 READY, 1 DONE, 1 ERROR)
     async with auth_db as auth_db:
         deleted_device = await auth_db.clean_expired_device_flows(max_retention=0)
-        assert deleted_device == 2
+        assert deleted_device == 4
 
     # Check the number of deleted device flow (should be 0 because there is nothing left to delete)
     async with auth_db as auth_db:

diracx-db/tests/auth/test_refresh_token.py

@@ -278,18 +278,20 @@ async def test_clean_refresh_tokens(auth_db: AuthDB):
 
     # Check the number of deleted refresh tokens (should be 0)
     async with auth_db as auth_db:
-        deleted_expired, deleted_revoked = await auth_db.clean_expired_refresh_token(
-            max_validity=10, max_retention=30
-        )
+        deleted_expired = await auth_db.clean_expired_refresh_tokens(max_validity=10)
     assert deleted_expired == 0
+
+    async with auth_db as auth_db:
+        deleted_revoked = await auth_db.clean_revoked_refresh_tokens(max_retention=30)
     assert deleted_revoked == 0
 
     # Check the number of deleted refresh tokens (should be 1 of each)
     async with auth_db as auth_db:
-        deleted_expired, deleted_revoked = await auth_db.clean_expired_refresh_token(
-            max_validity=0, max_retention=0
-        )
+        deleted_expired = await auth_db.clean_expired_refresh_tokens(max_validity=0)
     assert deleted_expired == 1
+
+    async with auth_db as auth_db:
+        deleted_revoked = await auth_db.clean_revoked_refresh_tokens(max_retention=0)
     assert deleted_revoked == 1
 
     # Get all refresh tokens (Admin)

diracx-logic/src/diracx/logic/__main__.py

@@ -93,12 +93,15 @@ async def delete_jwk(args):
 async def cleanup_authdb(args):
     """Delete expired tokens and flows from the AuthDB."""
     logger.info("Deleting expired tokens and flows")
+    import os
+
     from diracx.core.settings import AuthSettings
     from diracx.db.sql import AuthDB
     from diracx.logic.auth.management import cleanup_expired_data
 
     settings = AuthSettings()
-    db = AuthDB(args.db_url)
+    db_url = os.environ["DIRACX_DB_URL_AUTHDB"]
+    db = AuthDB(db_url)
     async with db.engine_context():
         async with db:
             await cleanup_expired_data(db, settings)
@@ -137,9 +140,6 @@ def parse_args():
     cleanup_authdb_parser = subparsers.add_parser(
         "cleanup-authdb", help="Delete expired tokens and flows from the AuthDB"
     )
-    cleanup_authdb_parser.add_argument(
-        "--db-url", required=True, help="URL to the AuthDB"
-    )
     cleanup_authdb_parser.set_defaults(func=cleanup_authdb)
 
     args = parser.parse_args()

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

@@ -65,20 +65,22 @@ async def revoke_refresh_token_by_refresh_token(
 
 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(
+    expired_tokens = await auth_db.clean_expired_refresh_tokens(
         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"
+    logger.info("Deleted %d expired refresh tokens", expired_tokens)
+
+    revoked_tokens = await auth_db.clean_revoked_refresh_tokens(
+        max_retention=settings.revoked_refresh_token_retention_minutes,
     )
+    logger.info("Deleted %d revoked refresh tokens", revoked_tokens)
 
     auth = await auth_db.clean_expired_authorization_flows(
         max_retention=settings.completed_flow_retention_minutes,
     )
-    logger.info(f"Deleted {auth} expired authorization flows")
+    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(f"Deleted {device} expired device flows")
+    logger.info("Deleted %d expired device flows", device)

docs/admin/reference/env-variables.md

@@ -105,14 +105,14 @@ Expiration time in minutes for refresh tokens.
 The maximum lifetime of refresh tokens before they must be re-issued
 through a new authentication flow. Default: 60 minutes.
 
-### `DIRACX_SERVICE_AUTH_REVOKED_REFRESH_TOKEN_RETENTION_DAYS`
+### `DIRACX_SERVICE_AUTH_REVOKED_REFRESH_TOKEN_RETENTION_MINUTES`
 
-*Optional*, default value: `30`
+*Optional*, default value: `43200`
 
-Retention time in days for revoked refresh tokens.
+Retention time in minutes for revoked refresh tokens.
 
 The maximum retention time of refresh tokens after being
-revoked and before they are deleted. Default: 30 days.
+revoked and before they are deleted. Default: 43200 minutes (30 days).
 
 ### `DIRACX_SERVICE_AUTH_AVAILABLE_PROPERTIES`