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)
+withself.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 @@
 classDOTConfig(AppConfig):
 name="oauth2_provider"
 verbose_name="Django OAuth Toolkit"
+
+defready(self):
+# Import checks to ensure they run.
+from . importchecks# noqa: F401

oauth2_provider/checks.py

@@ -0,0 +1,28 @@
+fromdjango.appsimportapps
+fromdjango.coreimportchecks
+fromdjango.dbimportrouter
+
+from .settingsimportoauth2_settings
+
+
+@checks.register(checks.Tags.database)
+defvalidate_token_configuration(app_configs, **kwargs):
+databases=set(
+router.db_for_write(apps.get_model(model))
+formodelin (
+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.
+iflen(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 @@
 importlogging
 importtime
 importuuid
+fromcontextlibimportsuppress
 fromdatetimeimporttimedelta
 fromurllib.parseimportparse_qsl, urlparse
 
 fromdjango.appsimportapps
 fromdjango.confimportsettings
 fromdjango.contrib.auth.hashersimportidentify_hasher, make_password
 fromdjango.core.exceptionsimportImproperlyConfigured
-fromdjango.dbimportmodels, transaction
+fromdjango.dbimportmodels, router, transaction
 fromdjango.urlsimportreverse
 fromdjango.utilsimporttimezone
 fromdjango.utils.translationimportgettext_lazyas_
@@ -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()
-withtransaction.atomic():
+
+# Use the access_token_database instead of making the assumption it is in 'default'.
+withtransaction.atomic(using=access_token_database):
 token=refresh_token_model.objects.select_for_update().filter(pk=self.pk, revoked__isnull=True)
 ifnottoken:
 return
 self=list(token)[0]
 
-try:
-access_token_model.objects.get(pk=self.access_token_id).revoke()
-exceptaccess_token_model.DoesNotExist:
-pass
+withsuppress(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():
 
 
 defget_id_token_model():
-"""Return the AccessToken model that is active in this project."""
+"""Return the IDToken model that is active in this project."""
 returnapps.get_model(oauth2_settings.ID_TOKEN_MODEL)
 
 

oauth2_provider/oauth2_validators.py

@@ -15,7 +15,7 @@
 fromdjango.contrib.authimportauthenticate, get_user_model
 fromdjango.contrib.auth.hashersimportcheck_password, identify_hasher
 fromdjango.core.exceptionsimportObjectDoesNotExist
-fromdjango.dbimporttransaction
+fromdjango.dbimportrouter, transaction
 fromdjango.httpimportHttpRequest
 fromdjango.utilsimportdateformat, timezone
 fromdjango.utils.cryptoimportconstant_time_compare
@@ -567,11 +567,23 @@ def rotate_refresh_token(self, request):
  """
 returnoauth2_settings.ROTATE_REFRESH_TOKEN
 
-@transaction.atomic
 defsave_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'.
+withtransaction.atomic(using=router.db_for_write(AccessToken)):
+returnself._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
 
 returnrt.application==client
 
-@transaction.atomic
 def_save_id_token(self, jti, request, expires, *args, **kwargs):
 scopes=request.scopeor" ".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'.
+withtransaction.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 @@
+fromdjango.confimportsettings
+fromdjango.testimportTestCaseasDjangoTestCase
+fromdjango.testimportTransactionTestCaseasDjangoTransactionTestCase
+
+
+# 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.
+
+
+defretrieve_current_databases():
+iflen(settings.DATABASES) >1:
+return [namefornameinsettings.DATABASESifname!="default"]
+else:
+return ["default"]
+
+
+classOAuth2ProviderBase:
+@classmethod
+defsetUpClass(cls):
+cls.databases=retrieve_current_databases()
+super().setUpClass()
+
+
+classOAuth2ProviderTestCase(OAuth2ProviderBase, DjangoTestCase):
+"""Place holder to allow overriding behaviors."""
+
+
+classOAuth2ProviderTransactionTestCase(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.
+
+
+classAlphaRouter:
+# 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.
+
+defdb_for_read(self, model, **hints):
+ifmodel._meta.app_labelnotinapps_in_beta:
+return"alpha"
+returnNone
+
+defdb_for_write(self, model, **hints):
+ifmodel._meta.app_labelnotinapps_in_beta:
+return"alpha"
+returnNone
+
+defallow_relation(self, obj1, obj2, **hints):
+ifobj1._state.db=="alpha"andobj2._state.db=="alpha":
+returnTrue
+returnNone
+
+defallow_migrate(self, db, app_label, model_name=None, **hints):
+ifapp_labelnotinapps_in_beta:
+returndb=="alpha"
+returnNone
+
+
+classBetaRouter:
+defdb_for_read(self, model, **hints):
+ifmodel._meta.app_labelinapps_in_beta:
+return"beta"
+returnNone
+
+defdb_for_write(self, model, **hints):
+ifmodel._meta.app_labelinapps_in_beta:
+return"beta"
+returnNone
+
+defallow_relation(self, obj1, obj2, **hints):
+ifobj1._state.db=="beta"andobj2._state.db=="beta":
+returnTrue
+returnNone
+
+defallow_migrate(self, db, app_label, model_name=None, **hints):
+ifapp_labelinapps_in_beta:
+returndb=="beta"
+
+
+classCrossDatabaseRouter:
+# 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.
+defdb_for_read(self, model, **hints):
+ifmodel._meta.model_name=="accesstoken":
+return"beta"
+returnNone
+
+defdb_for_write(self, model, **hints):
+ifmodel._meta.model_name=="accesstoken":
+return"beta"
+returnNone
+
+defallow_relation(self, obj1, obj2, **hints):
+ifobj1._state.db=="beta"andobj2._state.db=="beta":
+returnTrue
+returnNone
+
+defallow_migrate(self, db, app_label, model_name=None, **hints):
+ifmodel_name=="accesstoken":
+returndb=="beta"
+returnNone