[OAuth] Hand-craft migration to replace upstream models

aleksis/core/migrations/0023_oauth_application_model.py

 # Generated by Django 3.2.8 on 2021-11-04 09:52
 
+# Manual migration to replace Django OAuth Toolkit's models with custom models
+# This migration is special, because it needs to work correctly both when applied
+# during the initial database migration of the whole project, and when updating
+# from an older schema version
+#
+# It strictly has to be applied before Django OAuth Toolkit's initial migration,
+# because when the planner introspects existing models, it will stumble upon the swapped
+# foreign key relations that now point to our custom models (as configured in settings).
+#
+# When run during the initial database migration, this simply works by setting run_before.
+#
+# When run during an update later, Django will refuse to migrate, because run_before will
+# lead to an inconsistent migration history. Django OAuth Toolkit's migrations have already
+# been run in the past, and we cannot move a new migration to before that.
+#
+# Therefore, we treat the two cases explicitly:
+#  – If tables from the oauth2_provider label exist, we…
+#     …assume that its migrations have been run in the past, and we are in an udpate
+#     …forcefully drop the migration from the database to convince Django to let our own
+#      migration run first
+#     …rename the tables to the new app label through raw SQL, passing the original
+#      operations that Django OAuth Toolkit would have done to the planner to spoof them
+#      in the state
+#  – If no tables from the oauth2_provider label exist, we…
+#     …assume that we are in the first ever migration of the project
+#     …introspect the operations that Django OAuth Toolkit would do, and re-create them for
+#      our own app
+#  – In both cases, we then let Django OAuth Toolkit do its migrations like normal; they will
+#    will be no-ops as we swapped all models
+
 from django.apps import apps as global_apps
 from django.conf import settings
-from django.db import connection, migrations, models
-import django.db.models.deletion
-import oauth2_provider.generators
+from django.db import connection, migrations
+from django.utils.module_loading import import_string
+
+MODEL_NAMES = ("Application", "Grant", "AccessToken", "IDToken", "RefreshToken")
+_MIGRATION_NAMES = ["0001_initial", "0002_auto_20190406_1805", "0003_auto_20201211_1314", "0004_auto_20200902_2022"]
+
+
+class Migration(migrations.Migration):
+    """Empty dummy migration to satisfy loader in utility functions."""
+    # This migration will later be replaced by the real one.
+    # We only need a class called Migration here so the loader we use in
+    # the functions below does not complain that this migration does not
+    # have a Migration class (yet).
+    pass
+
+def _get_oauth_migrations():
+    """Get all original migrations from Django OAuth Toolkit."""
+    loader = migrations.loader.MigrationLoader(connection)
+    return [loader.get_migration("oauth2_provider", name) for name in _MIGRATION_NAMES]
 
-def migrate_oauth_applications(apps, schema_editor):
-    db_alias = schema_editor.connection.alias
+def _get_oauth_migration_operations(for_model=None):
+    """Get all original operations from Django Oauth Toolkit and re-create them for the new models."""
+    operations = []
+    for migration in _get_oauth_migrations():
+        for old_operation in migration.operations:
+            # We need to re-create a new Operation class because Operations
+            # are immutably linked to an app
+            op_name, args, kwargs = old_operation.deconstruct()
+            if not "." in op_name:
+                op_name = f"django.db.migrations.{op_name}"
+            op_cls = import_string(op_name)
+            if "model_name" in kwargs:
+                # If we are not interested for this model, skip it
+                if for_model is not None and for_model != "oauth" + kwargs["model_name"].lower():
+                    continue
+                # We are modifying a field. Replace the model name it belongs to
+                kwargs["model_name"] = "OAuth" + kwargs["model_name"]
+            elif "name" in kwargs:
+                # If we are not interested for this model, skip it
+                if for_model is not None and for_model != "oauth" + kwargs["name"].lower():
+                    continue
+                # We are modifying a model. Replace the full name
+                kwargs["name"] = "OAuth" + kwargs["name"]
+            operations.append(op_cls(*args, **kwargs))
 
+    return operations
+
+def _get_original_models():
+    """Get the original models of Django OAuth Toolkit."""
     try:
-        OldApp = apps.get_model("oauth2_provider", "Application")
+        return [global_apps.get_model("oauth2_provider", name) for name in MODEL_NAMES]
     except LookupError:
-        return
-    NewApp = apps.get_model("core", "OAuthApplication")
+        return []
+
+def _get_new_models():
+    """Get the new customised models."""
+    return [global_apps.get_model("core", f"OAuth{name}") for name in MODEL_NAMES]
 
-    if connection.instrospection.table_names() & set(OldApp._meta.db_table):
-        NewApp.objects.using(db_alias).bulk_create(
-            [NewApp(**old_app) for old_app in OldApp.objects.values()]
-        )
+def _original_tables_exist():
+    """Check if original Django OAuth Toolkit tables exist."""
+    original_tables = set([model._meta.db_table for model in _get_original_models()])
+    return bool(set(connection.introspection.table_names()) & original_tables)
 
 
 class Migration(migrations.Migration):
@@ -28,35 +103,26 @@ class Migration(migrations.Migration):
         ('core', '0022_public_favicon'),
     ]
 
-    run_before = []
-
-    operations = [
-        migrations.CreateModel(
-            name='OAuthApplication',
-            fields=[
-                ('id', models.BigAutoField(primary_key=True, serialize=False)),
-                ('client_id', models.CharField(db_index=True, default=oauth2_provider.generators.generate_client_id, max_length=100, unique=True)),
-                ('redirect_uris', models.TextField(blank=True, help_text='Allowed URIs list, space separated')),
-                ('client_type', models.CharField(choices=[('confidential', 'Confidential'), ('public', 'Public')], max_length=32)),
-                ('authorization_grant_type', models.CharField(choices=[('authorization-code', 'Authorization code'), ('implicit', 'Implicit'), ('password', 'Resource owner password-based'), ('client-credentials', 'Client credentials'), ('openid-hybrid', 'OpenID connect hybrid')], max_length=32)),
-                ('client_secret', models.CharField(blank=True, db_index=True, default=oauth2_provider.generators.generate_client_secret, max_length=255)),
-                ('name', models.CharField(blank=True, max_length=255)),
-                ('skip_authorization', models.BooleanField(default=False)),
-                ('created', models.DateTimeField(auto_now_add=True)),
-                ('updated', models.DateTimeField(auto_now=True)),
-                ('algorithm', models.CharField(blank=True, choices=[('', 'No OIDC support'), ('RS256', 'RSA with SHA-2 256'), ('HS256', 'HMAC with SHA-2 256')], default='', max_length=5)),
-                ('user', models.ForeignKey(blank=True, null=True, on_delete=django.db.models.deletion.CASCADE, related_name='core_oauthapplication', to=settings.AUTH_USER_MODEL)),
-            ],
-            options={
-                'abstract': False,
-            },
-        ),
-    ]
+    run_before = [("oauth2_provider", _MIGRATION_NAMES[0])]
 
-    if global_apps.is_installed("oauth2_provider"):
-        operations += [
-            migrations.RunPython(migrate_oauth_applications),
-            migrations.RunSQL("DROP TABLE IF EXISTS oauth2_provider_application;"),
-        ]
+    operations = []
+    if _original_tables_exist():
+        # If original tables existed, we nned to copy data and remove them afterwards
+        for OldModel, NewModel in zip(_get_original_models(), _get_new_models()):
+            operations.append(migrations.RunSQL(
+                f"ALTER TABLE IF EXISTS {OldModel._meta.db_table} RENAME TO {NewModel._meta.db_table};",
+                reverse_sql=f"ALTER TABLE IF EXISTS {NewModel._meta.db_table} RENAME TO {OldModel._meta.db_table};",
+                # We copy the original migrations here to trick the planner into believing the migrations had run normally
+                state_operations=_get_oauth_migration_operations(NewModel._meta.model_name)
+            ))
     else:
-        run_before.append(('oauth2_provider', '0001_initial'))
+        # We need to copy Django OAuth Toolkit's migrations here to re-create
+        # all models, because we cannot auto-generate them due to the big swappable
+        # model circular dependency rabbit hole
+        operations += _get_oauth_migration_operations()
+
+
+if _original_tables_exist():
+    # Fake-unapply migrations so the planner allows run_before
+    for name in _MIGRATION_NAMES:
+        connection.cursor().execute(f"DELETE FROM django_migrations WHERE app='oauth2_provider' AND name='{name}';")

aleksis/core/migrations/0024_oauth_grant_types_optional.py

-# Generated by Django 3.2.8 on 2021-11-04 09:58
+# Generated by Django 3.2.8 on 2021-11-04 19:31
 
 from django.db import migrations, models
 
@@ -7,6 +7,7 @@ class Migration(migrations.Migration):
 
     dependencies = [
         ('core', '0023_oauth_application_model'),
+        ('oauth2_provider', '0004_auto_20200902_2022'),
     ]
 
     operations = [