diff --git a/apps/gooyal_oauth2/models.py b/apps/gooyal_oauth2/models.py index 82b9793..476cdcb 100644 --- a/apps/gooyal_oauth2/models.py +++ b/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) +