feat: add functions to cleanup the AuthDB tables

Author: HeloiseJoffe

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

@@ -183,6 +183,13 @@ class AuthSettings(ServiceSettingsBase):
     before it must be exchanged for tokens. Default: 5 minutes.
     """
 
+    completed_flow_retention_minutes: int = 60
+    """Retention time in minutes for completed flow.
+
+    The maximum retention time of flow after being completed
+    and before they are deleted. Default: 60 minutes.
+    """
+
     state_key: FernetKey
     """Encryption key used to encrypt/decrypt the state parameter passed to the IAM.
 
@@ -226,6 +233,13 @@ 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.
+
+    The maximum retention time of refresh tokens after being
+    revoked and before they are deleted. Default: 30 days.
+    """
+
     available_properties: set[SecurityProperty] = Field(
         default_factory=SecurityProperty.available_properties
     )

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

@@ -6,7 +6,7 @@
 from itertools import pairwise
 
 from dateutil.rrule import MONTHLY, rrule
-from sqlalchemy import insert, select, text, update
+from sqlalchemy import delete, insert, select, text, update
 from sqlalchemy.exc import IntegrityError, NoResultFound
 from sqlalchemy.ext.asyncio import AsyncConnection
 from uuid_utils import UUID, uuid7
@@ -339,3 +339,47 @@ async def revoke_user_refresh_tokens(self, subject):
             .where(RefreshTokens.sub == 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."""
+        expired_date = str(
+            uuid7_from_datetime(substract_date(minutes=max_validity), randomize=False)
+        )
+        stmt_expired = delete(RefreshTokens).where(
+            RefreshTokens.status == RefreshTokenStatus.CREATED,
+            RefreshTokens.jti < expired_date,
+        )
+        res_expired = await self.conn.execute(stmt_expired)
+
+        revoked_date = str(
+            uuid7_from_datetime(substract_date(days=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
+
+    async def clean_expired_authorization_flows(self, max_retention: int) -> int:
+        """Delete old DONE/ERROR authorization flows."""
+        stmt_auth = delete(AuthorizationFlows).where(
+            AuthorizationFlows.status.in_([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."""
+        stmt_device = delete(DeviceFlows).where(
+            DeviceFlows.status.in_([FlowStatus.DONE, FlowStatus.ERROR]),
+            DeviceFlows.creation_time < substract_date(minutes=max_retention),
+        )
+        res_device = await self.conn.execute(stmt_device)
+
+        return res_device.rowcount

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

@@ -4,6 +4,7 @@
 
 import base64
 import hashlib
+import logging
 import re
 from datetime import datetime, timedelta, timezone
 from typing import cast
@@ -37,6 +38,8 @@
     verify_dirac_refresh_token,
 )
 
+logger = logging.getLogger(__name__)
+
 
 async def get_oidc_token(
     grant_type: GrantType,
@@ -429,6 +432,8 @@ 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
@@ -437,3 +442,32 @@ 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,
+    )

docs/admin/reference/env-variables.md

@@ -42,6 +42,15 @@ Expiration time in seconds for authorization code flow.
 The time window during which the authorization code remains valid
 before it must be exchanged for tokens. Default: 5 minutes.
 
+### `DIRACX_SERVICE_AUTH_COMPLETED_FLOW_RETENTION_MINUTES`
+
+*Optional*, default value: `60`
+
+Retention time in minutes for completed flow.
+
+The maximum retention time of flow after being completed
+and before they are deleted. Default: 60 minutes.
+
 ### `DIRACX_SERVICE_AUTH_STATE_KEY`
 
 **Required**
@@ -96,6 +105,15 @@ 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`
+
+*Optional*, default value: `30`
+
+Retention time in days for revoked refresh tokens.
+
+The maximum retention time of refresh tokens after being
+revoked and before they are deleted. Default: 30 days.
+
 ### `DIRACX_SERVICE_AUTH_AVAILABLE_PROPERTIES`
 
 *Optional*