clear tokens

apps/gooyal_oauth2/models.py

@@ -1,7 +1,12 @@
+import logging
+import time
+from datetime import timedelta
+
 import requests
+from django.core.exceptions import ImproperlyConfigured
 from django.db.models import JSONField
 from oauth2_provider.models import AbstractApplication, AbstractAccessToken, AbstractGrant, AbstractRefreshToken, \
-    AbstractIDToken, get_access_token_model, get_refresh_token_model
+    AbstractIDToken, get_access_token_model, get_refresh_token_model, get_id_token_model, get_grant_model
 from oauth2_provider.scopes import get_scopes_backend
 from oauth2_provider.settings import oauth2_settings
 import uuid
@@ -14,6 +19,8 @@ from django.utils import timezone
 from .settings import oauth2_settings
 
 
+logger = logging.getLogger(__name__)
+
 
 class Resource(models.Model):
     uuid = models.UUIDField(primary_key=True, editable=False, default=uuid.uuid4, unique=True, db_index=True)
@@ -221,3 +228,74 @@ class AccessToken(AbstractAccessToken):
 class IDToken(AbstractIDToken):
     id = None
     uuid = models.UUIDField(primary_key=True, editable=False, default=uuid.uuid4, unique=True, db_index=True)
+
+
+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("pk", flat=True)[:CLEAR_EXPIRED_TOKENS_BATCH_SIZE]
+            batch_length = flat_queryset.count()
+            queryset.model.objects.filter(pk__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()
+    id_token_model = get_id_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:
+                REFRESH_TOKEN_EXPIRE_SECONDS = timedelta(seconds=REFRESH_TOKEN_EXPIRE_SECONDS)
+            except TypeError:
+                e = "REFRESH_TOKEN_EXPIRE_SECONDS must be either a timedelta or seconds"
+                raise ImproperlyConfigured(e)
+        refresh_expire_at = now - REFRESH_TOKEN_EXPIRE_SECONDS
+
+    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)
+
+        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)
+
+    access_token_query = models.Q(refresh_token__isnull=True, expires__lt=now)
+    access_tokens = access_token_model.objects.filter(access_token_query)
+
+    access_tokens_delete_no = batch_delete(access_tokens, access_token_query)
+    logger.info("%s Expired access tokens deleted", access_tokens_delete_no)
+
+    id_token_query = models.Q(access_token__isnull=True, expires__lt=now)
+    id_tokens = id_token_model.objects.filter(id_token_query)
+
+    id_tokens_delete_no = batch_delete(id_tokens, id_token_query)
+    logger.info("%s Expired ID tokens deleted", id_tokens_delete_no)
+
+    grants_query = models.Q(expires__lt=now)
+    grants = grant_model.objects.filter(grants_query)
+
+    grants_deleted_no = batch_delete(grants, grants_query)
+    logger.info("%s Expired grant tokens deleted", grants_deleted_no)
+