Batch tokens deletion in cleartokens command

AUTHORS

@@ -24,6 +24,7 @@ Bas van Oostveen
 Dave Burkholder
 David Fischer
 David Smith
+Dawid Wolski
 Diego Garcia
 Dulmandakh Sukhbaatar
 Dylan Giesler

CHANGELOG.md

@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   -->
 
 ## [Unreleased]
+### Added
+* #651 Batch expired token deletions in `cleartokens` management command
 
 ### Added
 

docs/management_commands.rst

@@ -16,5 +16,8 @@ If ``cleartokens`` runs daily the maximum delay before a refresh token is
 removed is ``REFRESH_TOKEN_EXPIRE_SECONDS`` + 1 day. This is normally not a
 problem since refresh tokens are long lived.
 
+To prevent the CPU and RAM high peaks during deletion process use ``CLEAR_EXPIRED_TOKENS_BATCH_SIZE`` and
+``CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL`` settings to adjust the process speed.
+
 Note: Refresh tokens need to expire before AccessTokens can be removed from the
 database. Using ``cleartokens`` without ``REFRESH_TOKEN_EXPIRE_SECONDS`` has limited effect.

docs/settings.rst

@@ -337,6 +337,18 @@ Default: ``["client_secret_post", "client_secret_basic"]``
 
 The authentication methods that are advertised to be supported by this server.
 
+CLEAR_EXPIRED_TOKENS_BATCH_SIZE
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Default: ``10000``
+
+The size of delete batches used by ``cleartokens`` management command.
+
+CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Default: ``0.1``
+
+Time of sleep in seconds used by ``cleartokens`` management command between batch deletions.
+
 
 Settings imported from Django project
 --------------------------

oauth2_provider/models.py

@@ -1,4 +1,5 @@
 import logging
+import time
 import uuid
 from datetime import timedelta
 from urllib.parse import parse_qsl, urlparse
@@ -621,12 +622,31 @@ def get_refresh_token_admin_class():
 
 
 def clear_expired():
+    def batch_delete(queryset, query):
+        CLEAR_EXPIRED_TOKENS_BATCH_SIZE = oauth2_settings.CLEAR_EXPIRED_TOKENS_BATCH_SIZE
+        CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL = oauth2_settings.CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL
+        current_no = start_no = queryset.count()
+
+        while current_no:
+            flat_queryset = queryset.values_list("id", flat=True)[:CLEAR_EXPIRED_TOKENS_BATCH_SIZE]
+            batch_length = flat_queryset.count()
+            queryset.model.objects.filter(id__in=list(flat_queryset)).delete()
+            logger.debug(f"{batch_length} tokens deleted, {current_no-batch_length} left")
+            queryset = queryset.model.objects.filter(query)
+            time.sleep(CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL)
+            current_no = queryset.count()
+
+        stop_no = queryset.model.objects.filter(query).count()
+        deleted = start_no - stop_no
+        return deleted
+
     now = timezone.now()
     refresh_expire_at = None
     access_token_model = get_access_token_model()
     refresh_token_model = get_refresh_token_model()
     grant_model = get_grant_model()
     REFRESH_TOKEN_EXPIRE_SECONDS = oauth2_settings.REFRESH_TOKEN_EXPIRE_SECONDS
+
     if REFRESH_TOKEN_EXPIRE_SECONDS:
         if not isinstance(REFRESH_TOKEN_EXPIRE_SECONDS, timedelta):
             try:
@@ -636,31 +656,32 @@ def clear_expired():
                 raise ImproperlyConfigured(e)
         refresh_expire_at = now - REFRESH_TOKEN_EXPIRE_SECONDS
 
-    with transaction.atomic():
-        if refresh_expire_at:
-            revoked = refresh_token_model.objects.filter(
-                revoked__lt=refresh_expire_at,
-            )
-            expired = refresh_token_model.objects.filter(
-                access_token__expires__lt=refresh_expire_at,
-            )
+    if refresh_expire_at:
+        revoked_query = models.Q(revoked__lt=refresh_expire_at)
+        revoked = refresh_token_model.objects.filter(revoked_query)
+
+        revoked_deleted_no = batch_delete(revoked, revoked_query)
+        logger.info("%s Revoked refresh tokens deleted", revoked_deleted_no)
+
+        expired_query = models.Q(access_token__expires__lt=refresh_expire_at)
+        expired = refresh_token_model.objects.filter(expired_query)
 
-            logger.info("%s Revoked refresh tokens to be deleted", revoked.count())
-            logger.info("%s Expired refresh tokens to be deleted", expired.count())
+        expired_deleted_no = batch_delete(expired, expired_query)
+        logger.info("%s Expired refresh tokens deleted", expired_deleted_no)
+    else:
+        logger.info("refresh_expire_at is %s. No refresh tokens deleted.", refresh_expire_at)
 
-            revoked.delete()
-            expired.delete()
-        else:
-            logger.info("refresh_expire_at is %s. No refresh tokens deleted.", refresh_expire_at)
+    access_token_query = models.Q(refresh_token__isnull=True, expires__lt=now)
+    access_tokens = access_token_model.objects.filter(access_token_query)
 
-        access_tokens = access_token_model.objects.filter(refresh_token__isnull=True, expires__lt=now)
-        grants = grant_model.objects.filter(expires__lt=now)
+    access_tokens_delete_no = batch_delete(access_tokens, access_token_query)
+    logger.info("%s Expired access tokens deleted", access_tokens_delete_no)
 
-        logger.info("%s Expired access tokens to be deleted", access_tokens.count())
-        logger.info("%s Expired grant tokens to be deleted", grants.count())
+    grants_query = models.Q(expires__lt=now)
+    grants = grant_model.objects.filter(grants_query)
 
-        access_tokens.delete()
-        grants.delete()
+    grants_deleted_no = batch_delete(grants, grants_query)
+    logger.info("%s Expired grant tokens deleted", grants_deleted_no)
 
 
 def redirect_to_uri_allowed(uri, allowed_uris):

oauth2_provider/settings.py

@@ -101,6 +101,8 @@
     # Whether to re-create OAuthlibCore on every request.
     # Should only be required in testing.
     "ALWAYS_RELOAD_OAUTHLIB_CORE": False,
+    "CLEAR_EXPIRED_TOKENS_BATCH_SIZE": 10000,
+    "CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL": 0.1,
 }
 
 # List of settings that cannot be empty

tests/settings.py

@@ -156,3 +156,6 @@
 OAUTH2_PROVIDER_APPLICATION_MODEL = "oauth2_provider.Application"
 OAUTH2_PROVIDER_REFRESH_TOKEN_MODEL = "oauth2_provider.RefreshToken"
 OAUTH2_PROVIDER_ID_TOKEN_MODEL = "oauth2_provider.IDToken"
+
+CLEAR_EXPIRED_TOKENS_BATCH_SIZE = 1
+CLEAR_EXPIRED_TOKENS_BATCH_INTERVAL = 0