Honor database assignment from router (#1450)

* 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.

* changelog entry and authors update

* PR review response.

Document multiple database requires in advanced_topics.rst.
Add an ImproperlyConfigured validator to the ready method of the DOTConfig app.
Fix IDToken doc string.
Document the use of _save_bearer_token.
Define LocalIDToken and use it for validating the configuration test.
Questionably, define py39-multi-db-invalid-token-configuration-dj42. This will
consistently cause tox runs to fail until it is worked out how to mark this
as an expected failure.

* move migration

* update migration

* use django checks system

* drop misconfigured db check. Let's find a better way.

* run checks

* maybe a better test definition

* listing tests was breaking things

* No more magic.

* Oops. Debugger.

* Use retrieven_current_databases in django_db marked tests.

* Updates.

Prove the checks work.
Document test requirements.

* fix typo

---------

Co-authored-by: Alan Crosswell <[email protected]>
Co-authored-by: Alan Crosswell <[email protected]>

Author: shaleh

AUTHORS

@@ -102,6 +102,7 @@ Rodney Richardson
 Rustem Saiargaliev
 Rustem Saiargaliev
 Sandro Rodrigues
+Sean 'Shaleh' Perry
 Shaheed Haque
 Shaun Stanworth
 Sayyid Hamid Mahdavi

CHANGELOG.md

@@ -23,6 +23,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Update token to TextField from CharField with 255 character limit and SHA-256 checksum in AbstractAccessToken model. Removing the 255 character limit enables supporting JWT tokens with additional claims
 * Update middleware, validators, and views to use token checksums instead of token for token retrieval and validation.
 * #1446 use generic models pk instead of id.
+* Transactions wrapping writes of the Tokens now rely on Django's database routers to determine the correct
+  database to use instead of assuming that 'default' is the correct one.
 * Bump oauthlib version to 3.2.2 and above
 * Update the OAuth2Validator's invalidate_authorization_code method to return an InvalidGrantError if the associated grant does not exist.
 

docs/advanced_topics.rst

@@ -65,6 +65,17 @@ That's all, now Django OAuth Toolkit will use your model wherever an Application
     is because of the way Django currently implements swappable models.
     See `issue #90 <https://github.com/jazzband/django-oauth-toolkit/issues/90>`_ for details.
 
+Configuring multiple databases
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+There is no requirement that the tokens are stored in the default database or that there is a
+default database provided the database routers can determine the correct Token locations. Because the
+Tokens have foreign keys to the ``User`` model, you likely want to keep the tokens in the same database
+as your User model. It is also important that all of the tokens are stored in the same database.
+This could happen for instance if one of the Tokens is locally overridden and stored in a separate database.
+The reason for this is transactions will only be made for the database where AccessToken is stored
+even when writing to RefreshToken or other tokens.
+
 Multiple Grants
 ~~~~~~~~~~~~~~~
 

docs/contributing.rst

@@ -252,6 +252,26 @@ Open :file:`mycoverage/index.html` in your browser and you can see a coverage su
 
 There's no need to wait for Codecov to complain after you submit your PR.
 
+The tests are generic and written to work with both single database and multiple database configurations. tox will run
+tests both ways. You can see the configurations used in tests/settings.py and tests/multi_db_settings.py.
+
+When there are multiple databases defined, Django tests will not work unless they are told which database(s) to work with.
+For test writers this means any test must either:
+- instead of Django's TestCase or TransactionTestCase use the versions of those
+  classes defined in tests/common_testing.py
+- when using pytest's `django_db` mark, define it like this:
+  `@pytest.mark.django_db(databases=retrieve_current_databases())`
+
+In test code, anywhere the database is referenced the Django router needs to be used exactly like the package's code.
+
+.. code-block:: python
+
+    token_database = router.db_for_write(AccessToken)
+    with self.assertNumQueries(1, using=token_database):
+        # call something using the database
+
+Without the 'using' option, this test fails in the multiple database scenario because 'default' will be used instead.
+
 Code conventions matter
 -----------------------
 

oauth2_provider/apps.py

@@ -4,3 +4,7 @@
 class DOTConfig(AppConfig):
     name = "oauth2_provider"
     verbose_name = "Django OAuth Toolkit"
+
+    def ready(self):
+        # Import checks to ensure they run.
+        from . import checks  # noqa: F401

oauth2_provider/checks.py

@@ -0,0 +1,28 @@
+from django.apps import apps
+from django.core import checks
+from django.db import router
+
+from .settings import oauth2_settings
+
+
+@checks.register(checks.Tags.database)
+def validate_token_configuration(app_configs, **kwargs):
+    databases = set(
+        router.db_for_write(apps.get_model(model))
+        for model in (
+            oauth2_settings.ACCESS_TOKEN_MODEL,
+            oauth2_settings.ID_TOKEN_MODEL,
+            oauth2_settings.REFRESH_TOKEN_MODEL,
+        )
+    )
+
+    # This is highly unlikely, but let's warn people just in case it does.
+    # If the tokens were allowed to be in different databases this would require all
+    # writes to have a transaction around each database. Instead, let's enforce that
+    # they all live together in one database.
+    # The tokens are not required to live in the default database provided the Django
+    # routers know the correct database for them.
+    if len(databases) > 1:
+        return [checks.Error("The token models are expected to be stored in the same database.")]
+
+    return []

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 _
@@ -512,17 +513,19 @@ def revoke(self):
         Mark this refresh token revoked and revoke related access token
         """
         access_token_model = get_access_token_model()
+        access_token_database = router.db_for_write(access_token_model)
         refresh_token_model = get_refresh_token_model()
-        with transaction.atomic():
+
+        # 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()
@@ -655,7 +658,7 @@ def get_access_token_model():
 
 
 def get_id_token_model():
-    """Return the AccessToken model that is active in this project."""
+    """Return the IDToken model that is active in this project."""
     return apps.get_model(oauth2_settings.ID_TOKEN_MODEL)
 
 

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
@@ -567,11 +567,23 @@ def rotate_refresh_token(self, request):
         """
         return oauth2_settings.ROTATE_REFRESH_TOKEN
 
-    @transaction.atomic
     def save_bearer_token(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`
+        Save access and refresh token.
+
+        Override _save_bearer_token and not this function when adding custom logic
+        for the storing of these token. This allows the transaction logic to be
+        separate from the token handling.
+        """
+        # 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(token, request, *args, **kwargs)
+
+    def _save_bearer_token(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`.
 
         @see: https://rfc-editor.org/rfc/rfc6749.html#section-6
         """
@@ -793,7 +805,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)
 
@@ -894,7 +905,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,33 @@
+from django.conf import settings
+from django.test import TestCase as DjangoTestCase
+from django.test import TransactionTestCase as DjangoTransactionTestCase
+
+
+# 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.
+#
+# When there are multiple databases defined, Django tests will not work unless they are
+# told which database(s) to work with.
+
+
+def retrieve_current_databases():
+    if len(settings.DATABASES) > 1:
+        return [name for name in settings.DATABASES if name != "default"]
+    else:
+        return ["default"]
+
+
+class OAuth2ProviderBase:
+    @classmethod
+    def setUpClass(cls):
+        cls.databases = retrieve_current_databases()
+        super().setUpClass()
+
+
+class OAuth2ProviderTestCase(OAuth2ProviderBase, DjangoTestCase):
+    """Place holder to allow overriding behaviors."""
+
+
+class OAuth2ProviderTransactionTestCase(OAuth2ProviderBase, DjangoTransactionTestCase):
+    """Place holder to allow overriding behaviors."""

tests/db_router.py

@@ -0,0 +1,76 @@
+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.
+
+
+class AlphaRouter:
+    # 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.
+
+    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"
+
+
+class CrossDatabaseRouter:
+    # 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.
+    def db_for_read(self, model, **hints):
+        if model._meta.model_name == "accesstoken":
+            return "beta"
+        return None
+
+    def db_for_write(self, model, **hints):
+        if model._meta.model_name == "accesstoken":
+            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 model_name == "accesstoken":
+            return db == "beta"
+        return None