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.

Author: shaleh

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
 ~~~~~~~~~~~~~~~
 

oauth2_provider/apps.py

@@ -1,6 +1,32 @@
-from django.apps import AppConfig
+from django.apps import AppConfig, apps
+from django.core.exceptions import ImproperlyConfigured
+from django.db import router
 
 
 class DOTConfig(AppConfig):
     name = "oauth2_provider"
     verbose_name = "Django OAuth Toolkit"
+
+    def _validate_token_configuration(self):
+        from .settings import oauth2_settings
+
+        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:
+            raise ImproperlyConfigured("the token models are expected to be stored in the same database.")
+
+    def ready(self):
+        self._validate_token_configuration()

oauth2_provider/models.py

@@ -513,17 +513,8 @@ def revoke(self):
         Mark this refresh token revoked and revoke related access token
         """
         access_token_model = get_access_token_model()
-        refresh_token_model = get_refresh_token_model()
-
         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"
-            )
+        refresh_token_model = get_refresh_token_model()
 
         # Use the access_token_database instead of making the assumption it is in 'default'.
         with transaction.atomic(using=access_token_database):
@@ -667,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

@@ -563,14 +563,22 @@ def rotate_refresh_token(self, request):
         return oauth2_settings.ROTATE_REFRESH_TOKEN
 
     def save_bearer_token(self, token, request, *args, **kwargs):
+        """
+        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_internals(token, request, *args, **kwargs)
+            return self._save_bearer_token(token, request, *args, **kwargs)
 
-    def _save_bearer_token_internals(self, 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`
+        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
         """

tests/db_router.py

@@ -1,13 +1,14 @@
-apps_in_beta = {"some_other_app", "this_one_too"}
+apps_in_beta = {"tests", "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:
+    # 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"

tests/migrations/0006_add_localidtoken.py

@@ -0,0 +1,34 @@
+# Generated by Django 3.2.25 on 2024-08-08 22:47
+
+from django.conf import settings
+from django.db import migrations, models
+import django.db.models.deletion
+import uuid
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        migrations.swappable_dependency(settings.OAUTH2_PROVIDER_APPLICATION_MODEL),
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ('tests', '0005_basetestapplication_allowed_origins_and_more'),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name='LocalIDToken',
+            fields=[
+                ('id', models.BigAutoField(primary_key=True, serialize=False)),
+                ('jti', models.UUIDField(default=uuid.uuid4, editable=False, unique=True, verbose_name='JWT Token ID')),
+                ('expires', models.DateTimeField()),
+                ('scope', models.TextField(blank=True)),
+                ('created', models.DateTimeField(auto_now_add=True)),
+                ('updated', models.DateTimeField(auto_now=True)),
+                ('application', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, to=settings.OAUTH2_PROVIDER_APPLICATION_MODEL)),
+                ('user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='tests_localidtoken', to=settings.AUTH_USER_MODEL)),
+            ],
+            options={
+                'abstract': False,
+            },
+        ),
+    ]

tests/models.py

@@ -4,6 +4,7 @@
     AbstractAccessToken,
     AbstractApplication,
     AbstractGrant,
+    AbstractIDToken,
     AbstractRefreshToken,
 )
 from oauth2_provider.settings import oauth2_settings
@@ -54,3 +55,9 @@ class SampleRefreshToken(AbstractRefreshToken):
 
 class SampleGrant(AbstractGrant):
     custom_field = models.CharField(max_length=255)
+
+
+class LocalIDToken(AbstractIDToken):
+    """Exists to be improperly configured for multiple databases."""
+
+    # The other token types will be in 'alpha' database.

tests/multi_db_settings_invalid_token_configuration.py

@@ -0,0 +1,8 @@
+from .multi_db_settings import *  # noqa: F401, F403
+
+
+OAUTH2_PROVIDER = {
+    # The other two tokens will be in alpha. This will cause a failure when the
+    # app's ready method is called.
+    "ID_TOKEN_MODEL": "tests.LocalIDToken",
+}

tox.ini

@@ -10,6 +10,7 @@ envlist =
     py{310,311,312}-dj51,
     py{310,311,312}-djmain,
     py39-multi-db-dj-42,
+    py39-multi-db-invalid-token-configuration-dj42,
 
 [gh-actions]
 python =
@@ -114,6 +115,13 @@ setenv =
     PYTHONPATH = {toxinidir}
     PYTHONWARNINGS = all
 
+[testenv:py39-multi-db-invalid-token-configuration-dj42]
+setenv =
+    DJANGO_SETTINGS_MODULE = tests.multi_db_settings_invalid_token_configuration
+    PYTHONPATH = {toxinidir}
+    PYTHONWARNINGS = all
+    ignore_errors = true
+
 [testenv:migrate_swapped]
 setenv =
     DJANGO_SETTINGS_MODULE = tests.settings_swapped