Improve multiple database support.

The token models might not be stored in the default database. There might not _be_ a default database.
Intead, the code now relies on Django's routers to determine the actual database to use when creating
transactions. This required moving from decorators to context managers for those transactions.

To test the multiple database scenario a new settings file as added which derives from settings.py
and then defines different databases and the routers needed to access them. The commit is larger than
might be expected because when there are multiple databases the Django tests have to be told which
databases to work on. Rather than copying the various test cases or making multiple database specific
ones the decision was made to add wrappers around the standard Django TestCase classes and programmatically
define the databases for them. This enables all of the same test code to work for both the one database
and the multi database scenarios with minimal maintenance costs. A tox environment that uses the
multi db settings file has been added to ensure both scenarios are always tested.

Author: shaleh

oauth2_provider/models.py

@@ -2,14 +2,15 @@
 import logging
 import time
 import uuid
+from contextlib import suppress
 from datetime import timedelta
 from urllib.parse import parse_qsl, urlparse
 
 from django.apps import apps
 from django.conf import settings
 from django.contrib.auth.hashers import identify_hasher, make_password
 from django.core.exceptions import ImproperlyConfigured
-from django.db import models, transaction
+from django.db import models, router, transaction
 from django.urls import reverse
 from django.utils import timezone
 from django.utils.translation import gettext_lazy as _
@@ -513,16 +514,27 @@ def revoke(self):
         """
         access_token_model = get_access_token_model()
         refresh_token_model = get_refresh_token_model()
-        with transaction.atomic():
+
+        access_token_database = router.db_for_write(access_token_model)
+        refresh_token_database = router.db_for_write(refresh_token_model)
+
+        # This is highly unlikely, but let's warn people just in case it does.
+        if access_token_database != refresh_token_database:
+            logger.warning(
+                "access token and refresh token are in separate databases but a transaction"
+                " is only used for the access token"
+            )
+
+        # Use the access_token_database instead of making the assumption it is in 'default'.
+        with transaction.atomic(using=access_token_database):
             token = refresh_token_model.objects.select_for_update().filter(pk=self.pk, revoked__isnull=True)
             if not token:
                 return
             self = list(token)[0]
 
-            try:
-                access_token_model.objects.get(pk=self.access_token_id).revoke()
-            except access_token_model.DoesNotExist:
-                pass
+            with suppress(access_token_model.DoesNotExist):
+                access_token_model.objects.get(id=self.access_token_id).revoke()
+
             self.access_token = None
             self.revoked = timezone.now()
             self.save()

oauth2_provider/oauth2_validators.py

@@ -15,7 +15,7 @@
 from django.contrib.auth import authenticate, get_user_model
 from django.contrib.auth.hashers import check_password, identify_hasher
 from django.core.exceptions import ObjectDoesNotExist
-from django.db import transaction
+from django.db import router, transaction
 from django.http import HttpRequest
 from django.utils import dateformat, timezone
 from django.utils.crypto import constant_time_compare
@@ -562,8 +562,12 @@ def rotate_refresh_token(self, request):
         """
         return oauth2_settings.ROTATE_REFRESH_TOKEN
 
-    @transaction.atomic
     def save_bearer_token(self, token, request, *args, **kwargs):
+        # Use the AccessToken's database instead of making the assumption it is in 'default'.
+        with transaction.atomic(using=router.db_for_write(AccessToken)):
+            return self._save_bearer_token_internals(token, request, *args, **kwargs)
+
+    def _save_bearer_token_internals(self, token, request, *args, **kwargs):
         """
         Save access and refresh token, If refresh token is issued, remove or
         reuse old refresh token as in rfc:`6`
@@ -788,7 +792,6 @@ def validate_refresh_token(self, refresh_token, client, request, *args, **kwargs
 
         return rt.application == client
 
-    @transaction.atomic
     def _save_id_token(self, jti, request, expires, *args, **kwargs):
         scopes = request.scope or " ".join(request.scopes)
 
@@ -889,7 +892,9 @@ def finalize_id_token(self, id_token, token, token_handler, request):
             claims=json.dumps(id_token, default=str),
         )
         jwt_token.make_signed_token(request.client.jwk_key)
-        id_token = self._save_id_token(id_token["jti"], request, expiration_time)
+        # Use the IDToken's database instead of making the assumption it is in 'default'.
+        with transaction.atomic(using=router.db_for_write(IDToken)):
+            id_token = self._save_id_token(id_token["jti"], request, expiration_time)
         # this is needed by django rest framework
         request.access_token = id_token
         request.id_token = id_token

tests/common_testing.py

@@ -0,0 +1,35 @@
+from django.conf import settings
+from django.test import TestCase as DjangoTestCase
+from django.test import TransactionTestCase as DjangoTransactionTestCase
+
+
+class OAuth2ProviderTestCase(DjangoTestCase):
+    """Place holder to allow overriding behaviors."""
+
+
+class OAuth2ProviderTransactionTestCase(DjangoTransactionTestCase):
+    """Place holder to allow overriding behaviors."""
+
+
+if len(settings.DATABASES) > 1:
+    # There are multiple databases defined. When this happens Django tests will not
+    # work unless they are told which database(s) to work with. The multiple
+    # database scenario setup for these tests purposefully defines 'default' as an
+    # empty database in order to catch any assumptions in this package about database
+    # names and in particular to ensure there is no assumption that 'default' is a
+    # valid database.
+    # For any test that would usually use Django's TestCase or TransactionTestCase
+    # using the classes defined here is all that is required.
+    # Any test that uses pytest's django_db need to base in a databases parameter
+    # using this definition of test_database_names.
+    # In test code, anywhere the default database is used the variable
+    # database_for_oauth2_provider must be used in its place. For instance,
+    #     with self.assertNumQueries(1, using=database_for_oauth2_provider):
+    # without the using option this fails because default is used.
+    test_database_names = {name for name in settings.DATABASES if name != "default"}
+    database_for_oauth2_provider = "alpha"
+    OAuth2ProviderTestCase.databases = test_database_names
+    OAuth2ProviderTransactionTestCase.databases = test_database_names
+else:
+    test_database_names = {"default"}
+    database_for_oauth2_provider = "default"

tests/db_router.py

@@ -0,0 +1,51 @@
+apps_in_beta = {"some_other_app", "this_one_too"}
+
+# These are bare minimum routers to fake the scenario where there is actually a
+# decision around where an application's models might live.
+# alpha is where the core Django models are stored including user. To keep things
+# simple this is where the oauth2 provider models are stored as well because they
+# have a foreign key to User.
+
+
+class AlphaRouter:
+    def db_for_read(self, model, **hints):
+        if model._meta.app_label not in apps_in_beta:
+            return "alpha"
+        return None
+
+    def db_for_write(self, model, **hints):
+        if model._meta.app_label not in apps_in_beta:
+            return "alpha"
+        return None
+
+    def allow_relation(self, obj1, obj2, **hints):
+        if obj1._state.db == "alpha" and obj2._state.db == "alpha":
+            return True
+        return None
+
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        if app_label not in apps_in_beta:
+            return db == "alpha"
+        return None
+
+
+class BetaRouter:
+    def db_for_read(self, model, **hints):
+        if model._meta.app_label in apps_in_beta:
+            return "beta"
+        return None
+
+    def db_for_write(self, model, **hints):
+        if model._meta.app_label in apps_in_beta:
+            return "beta"
+        return None
+
+    def allow_relation(self, obj1, obj2, **hints):
+        if obj1._state.db == "beta" and obj2._state.db == "beta":
+            return True
+        return None
+
+    def allow_migrate(self, db, app_label, model_name=None, **hints):
+        if app_label in apps_in_beta:
+            return db == "beta"
+        return None

tests/multi_db_settings.py

@@ -0,0 +1,19 @@
+# Import the test settings and then override DATABASES.
+
+from .settings import *  # noqa: F401, F403
+
+
+DATABASES = {
+    "alpha": {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": ":memory:",
+    },
+    "beta": {
+        "ENGINE": "django.db.backends.sqlite3",
+        "NAME": ":memory:",
+    },
+    # As https://docs.djangoproject.com/en/4.2/topics/db/multi-db/#defining-your-databases
+    # indicates, it is ok to have no default database.
+    "default": {},
+}
+DATABASE_ROUTERS = ["tests.db_router.AlphaRouter", "tests.db_router.BetaRouter"]

tests/test_application_views.py

@@ -1,11 +1,11 @@
 import pytest
 from django.contrib.auth import get_user_model
-from django.test import TestCase
 from django.urls import reverse
 
 from oauth2_provider.models import get_application_model
 from oauth2_provider.views.application import ApplicationRegistration
 
+from .common_testing import OAuth2ProviderTestCase as TestCase
 from .models import SampleApplication
 
 

tests/test_auth_backends.py

@@ -5,14 +5,16 @@
 from django.contrib.auth.models import AnonymousUser
 from django.core.exceptions import SuspiciousOperation
 from django.http import HttpResponse
-from django.test import RequestFactory, TestCase
+from django.test import RequestFactory
 from django.test.utils import modify_settings, override_settings
 from django.utils.timezone import now, timedelta
 
 from oauth2_provider.backends import OAuth2Backend
 from oauth2_provider.middleware import OAuth2ExtraTokenMiddleware, OAuth2TokenMiddleware
 from oauth2_provider.models import get_access_token_model, get_application_model
 
+from .common_testing import OAuth2ProviderTestCase as TestCase
+
 
 UserModel = get_user_model()
 ApplicationModel = get_application_model()

tests/test_authorization_code.py

@@ -7,7 +7,7 @@
 import pytest
 from django.conf import settings
 from django.contrib.auth import get_user_model
-from django.test import RequestFactory, TestCase
+from django.test import RequestFactory
 from django.urls import reverse
 from django.utils import timezone
 from django.utils.crypto import get_random_string
@@ -23,6 +23,7 @@
 from oauth2_provider.views import ProtectedResourceView
 
 from . import presets
+from .common_testing import OAuth2ProviderTestCase as TestCase
 from .utils import get_basic_auth_header
 
 

tests/test_client_credential.py

@@ -4,7 +4,7 @@
 import pytest
 from django.contrib.auth import get_user_model
 from django.core.exceptions import SuspiciousOperation
-from django.test import RequestFactory, TestCase
+from django.test import RequestFactory
 from django.urls import reverse
 from django.views.generic import View
 from oauthlib.oauth2 import BackendApplicationServer
@@ -16,6 +16,7 @@
 from oauth2_provider.views.mixins import OAuthLibMixin
 
 from . import presets
+from .common_testing import OAuth2ProviderTestCase as TestCase
 from .utils import get_basic_auth_header
 
 

tests/test_commands.py

@@ -5,11 +5,11 @@
 from django.contrib.auth.hashers import check_password
 from django.core.management import call_command
 from django.core.management.base import CommandError
-from django.test import TestCase
 
 from oauth2_provider.models import get_application_model
 
 from . import presets
+from .common_testing import OAuth2ProviderTestCase as TestCase
 
 
 Application = get_application_model()