Feat: authdb tables cleanup (#815)

* feat: add functions to cleanup the AuthDB tables

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
     )
@@ -183,6 +199,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 +249,13 @@ class AuthSettings(ServiceSettingsBase):
     through a new authentication flow. Default: 60 minutes.
     """
 
+    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: 43200 minutes (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,61 @@ async def revoke_user_refresh_tokens(self, subject):
             .where(RefreshTokens.sub == subject)
             .values(status=RefreshTokenStatus.REVOKED)
         )
+
+    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)
+        )
+        stmt_expired = delete(RefreshTokens).where(
+            RefreshTokens.status == RefreshTokenStatus.CREATED,
+            RefreshTokens.jti < expired_date,
+        )
+        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(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_revoked.rowcount
+
+    async def clean_expired_authorization_flows(self, max_retention: int) -> int:
+        """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.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 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.creation_time < substract_date(minutes=max_retention),
+        )
+        res_device = await self.conn.execute(stmt_device)
+
+        return res_device.rowcount

diracx-db/tests/auth/test_authorization_flow.py

@@ -5,6 +5,7 @@
 
 from diracx.core.exceptions import AuthorizationError
 from diracx.db.sql.auth.db import AuthDB
+from diracx.db.sql.auth.schema import FlowStatus
 
 MAX_VALIDITY = 2
 EXPIRED = 0
@@ -74,3 +75,46 @@ async def test_insert(auth_db: AuthDB):
         )
 
     assert uuid1 != uuid2
+
+
+async def test_clean_authorization_flows(auth_db: AuthDB):
+    # Insert two authorization flows
+    async with auth_db as auth_db:
+        _ = 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:
+        _, _ = 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(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 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 == 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:
+        deleted_auth = await auth_db.clean_expired_authorization_flows(max_retention=0)
+        assert deleted_auth == 0

diracx-db/tests/auth/test_device_flow.py

@@ -8,7 +8,7 @@
 
 from diracx.core.exceptions import AuthorizationError
 from diracx.db.sql.auth.db import AuthDB
-from diracx.db.sql.auth.schema import USER_CODE_LENGTH
+from diracx.db.sql.auth.schema import USER_CODE_LENGTH, FlowStatus
 from diracx.db.sql.utils.functions import substract_date
 
 MAX_VALIDITY = 2
@@ -139,3 +139,35 @@ async def test_device_flow_insert_id_token(auth_db: AuthDB):
     async with auth_db as auth_db:
         res = await auth_db.get_device_flow(device_code)
         assert res["IDToken"] == id_token
+
+
+async def test_clean_device_flows(auth_db: AuthDB):
+    # Insert two device flows
+    async with auth_db as auth_db:
+        _, _ = 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.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 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 == 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:
+        deleted_device = await auth_db.clean_expired_device_flows(max_retention=0)
+        assert deleted_device == 0

diracx-db/tests/auth/test_refresh_token.py

@@ -257,3 +257,46 @@ async def test_get_refresh_tokens(auth_db: AuthDB):
 
     # Check the number of retrieved refresh tokens (should be 3 refresh tokens)
     assert len(refresh_tokens) == 2
+
+
+async def test_clean_refresh_tokens(auth_db: AuthDB):
+    # Insert two refresh tokens
+    jtis = []
+    async with auth_db as auth_db:
+        for _ in range(2):
+            jti = uuid7()
+            await auth_db.insert_refresh_token(
+                jti,
+                "subject",
+                "scope",
+            )
+            jtis.append(jti)
+
+    # Revoke one of the refresh token
+    async with auth_db as auth_db:
+        await auth_db.revoke_refresh_token(jtis[0])
+
+    # Check the number of deleted refresh tokens (should be 0)
+    async with auth_db as auth_db:
+        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 = 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)
+    async with auth_db as auth_db:
+        refresh_tokens = await auth_db.get_user_refresh_tokens()
+
+    # Check the number of retrieved refresh tokens (should be 0)
+    assert len(refresh_tokens) == 0

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

@@ -90,6 +90,23 @@ async def delete_jwk(args):
     save_jwks(path, jwks)
 
 
+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_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)
+
+
 def parse_args():
     parser = argparse.ArgumentParser()
     subparsers = parser.add_subparsers(dest="command", required=True)
@@ -120,8 +137,13 @@ def parse_args():
     )
     delete_jwk_parser.set_defaults(func=delete_jwk)
 
+    cleanup_authdb_parser = subparsers.add_parser(
+        "cleanup-authdb", help="Delete expired tokens and flows from the AuthDB"
+    )
+    cleanup_authdb_parser.set_defaults(func=cleanup_authdb)
+
     args = parser.parse_args()
-    logger.setLevel(logging.INFO)
+    logging.basicConfig(level=logging.INFO)
     asyncio.run(args.func(args))
 
 

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,26 @@ 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 = await auth_db.clean_expired_refresh_tokens(
+        max_validity=settings.refresh_token_expire_minutes,
+    )
+    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("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_MINUTES`
+
+*Optional*, default value: `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: 43200 minutes (30 days).
+
 ### `DIRACX_SERVICE_AUTH_AVAILABLE_PROPERTIES`
 
 *Optional*