From c1d29e160f83a3363c8bfb5bcd2f31ef66a6c112 Mon Sep 17 00:00:00 2001
From: Thomas NOEL <tnoel@entrouvert.com>
Date: Tue, 15 Jan 2019 14:47:27 +0100
Subject: [PATCH] migrate_schemas: add executor option from
 django-tenants-schemas (#23045)

---
 debian/debian_config_common.py                |  7 ++++
 hobo/multitenant/apps.py                      | 22 ----------
 .../management/commands/__init__.py           |  4 ++
 .../management/commands/migrate.py            | 25 ++++-------
 .../management/commands/migrate_schemas.py    | 42 +++++++++----------
 5 files changed, 37 insertions(+), 63 deletions(-)

diff --git a/debian/debian_config_common.py b/debian/debian_config_common.py
index e9cc69d..affc725 100644
--- a/debian/debian_config_common.py
+++ b/debian/debian_config_common.py
@@ -332,6 +332,13 @@ if PROJECT_NAME != 'wcs':
     if 'authentic2' not in INSTALLED_APPS:
         MELLON_ADAPTER = ('hobo.multitenant.mellon.MellonAdapter',)
 
+    # migrate_schemas in parallel
+    # see https://django-tenant-schemas.readthedocs.io/en/latest/use.html#migrate-schemas-in-parallel
+    # maximum number of processes for migration pool (avoid exhausting the database connection pool)
+    TENANT_PARALLEL_MIGRATION_MAX_PROCESSES = 4
+    # number of migrations to be sent at once to every worker
+    TENANT_PARALLEL_MIGRATION_CHUNKS = 2
+
 
 if 'authentic2' not in INSTALLED_APPS:
     MELLON_DEFAULT_ASSERTION_CONSUMER_BINDING = 'artifact'
diff --git a/hobo/multitenant/apps.py b/hobo/multitenant/apps.py
index 68cc347..c655555 100644
--- a/hobo/multitenant/apps.py
+++ b/hobo/multitenant/apps.py
@@ -138,25 +138,3 @@ class MultitenantAppConfig(AppConfig):
         threading._DummyThread = _DummyThread
         threading._MainThread = _MainThread
         threading._Timer = _Timer
-
-        if (1, 7) <= django.VERSION < (1, 8):
-            class RunPythonOverride(operations.RunPython):
-                def can_migrate(self, app_label, schema_editor):
-                    app = apps.get_app_config(app_label)
-                    for model in app.get_models():
-                        if self.allowed_to_migrate(schema_editor.connection.alias, model):
-                            return True
-                    return False
-
-                def database_forwards(self, app_label, schema_editor, *args, **kwargs):
-                    if not self.can_migrate(app_label, schema_editor):
-                        return
-                    super(RunPythonOverride, self).database_forwards(app_label, schema_editor, *args, **kwargs)
-
-                def database_backwards(self, app_label, schema_editor, *args, **kwargs):
-                    if not self.can_migrate(app_label, schema_editor):
-                        return
-                    super(RunPythonOverride, self).database_backwards(app_label, schema_editor, *args, **kwargs)
-
-            operations.RunPython = RunPythonOverride
-            migrations.RunPython = RunPythonOverride
diff --git a/hobo/multitenant/management/commands/__init__.py b/hobo/multitenant/management/commands/__init__.py
index 139d61c..658d08e 100644
--- a/hobo/multitenant/management/commands/__init__.py
+++ b/hobo/multitenant/management/commands/__init__.py
@@ -162,11 +162,15 @@ class SyncCommon(BaseCommand):
                     help=('Database state will be brought to the state after that '
                           'migration. Use the name "zero" to unapply all migrations.'))
         parser.add_argument("-d", "--domain", dest="domain")
+        parser.add_argument('--executor', action='store', dest='executor', default=None,
+                    help='Executor for running migrations [standard (default)|parallel]')
+
 
     def handle(self, *args, **options):
         self.sync_tenant = options.get('tenant')
         self.sync_public = options.get('shared')
         self.domain = options.get('domain')
+        self.executor = options.get('executor')
         self.installed_apps = settings.INSTALLED_APPS
         self.args = args
         self.options = options
diff --git a/hobo/multitenant/management/commands/migrate.py b/hobo/multitenant/management/commands/migrate.py
index bc7abf3..5cdc5e2 100644
--- a/hobo/multitenant/management/commands/migrate.py
+++ b/hobo/multitenant/management/commands/migrate.py
@@ -4,29 +4,18 @@
 #   License: MIT license
 #   Home-page: http://github.com/bcarneiro/django-tenant-schemas
 import django
-from django.conf import settings
 from django.core.management.base import CommandError, BaseCommand
 from tenant_schemas.utils import django_is_in_test_mode
 
-if django.VERSION < (1, 7, 0):
-    try:
-        from south.management.commands.migrate import Command as MigrateCommand
-    except ImportError:
-        MigrateCommand = BaseCommand
-else:
-    MigrateCommand = BaseCommand
 
-class Command(MigrateCommand):
+class Command(BaseCommand):
 
     def handle(self, *args, **options):
         database = options.get('database', 'default')
-        if (settings.DATABASES[database]['ENGINE'] == 'tenant_schemas.postgresql_backend' or
-                MigrateCommand is BaseCommand):
-            raise CommandError("migrate has been disabled, for database '{}'. Use migrate_schemas "
-                               "instead. Please read the documentation if you don't know why you "
-                               "shouldn't call migrate directly!".format(database))
-        super(Command, self).handle(*args, **options)
+        raise CommandError("migrate has been disabled, for database '{}'. Use migrate_schemas "
+                           "instead. Please read the documentation if you don't know why you "
+                           "shouldn't call migrate directly!".format(database))
 
-if django.VERSION >= (1, 7, 0) and django_is_in_test_mode():
-    from .migrate_schemas import MigrateSchemasCommand
-    Command = MigrateSchemasCommand
+
+if django_is_in_test_mode():
+    from .migrate_schemas import Command
diff --git a/hobo/multitenant/management/commands/migrate_schemas.py b/hobo/multitenant/management/commands/migrate_schemas.py
index f1b9a63..2e23c2f 100644
--- a/hobo/multitenant/management/commands/migrate_schemas.py
+++ b/hobo/multitenant/management/commands/migrate_schemas.py
@@ -1,17 +1,20 @@
 import django
-from optparse import NO_DEFAULT
 
 from django.core.management.commands.migrate import Command as MigrateCommand
-from django.db.migrations.recorder import MigrationRecorder
-from django.db import connection
-from django.conf import settings
 
+from tenant_schemas.migration_executors import get_executor
 from tenant_schemas.utils import get_public_schema_name, schema_exists
-from hobo.multitenant.middleware import TenantMiddleware, TenantNotFound
+from hobo.multitenant.middleware import TenantMiddleware
 from hobo.multitenant.management.commands import SyncCommon
 
+if django.VERSION >= (1, 9, 0):
+    from django.db.migrations.exceptions import MigrationSchemaMissing
+else:
+    class MigrationSchemaMissing(django.db.utils.DatabaseError):
+        pass
 
-class MigrateSchemasCommand(SyncCommon):
+
+class Command(SyncCommon):
     help = "Updates database schema. Manages both apps with migrations and those without."
 
     def __init__(self, stdout=None, stderr=None, no_color=False):
@@ -23,40 +26,33 @@ class MigrateSchemasCommand(SyncCommon):
         super(Command, self).__init__(stdout, stderr, no_color)
 
     def add_arguments(self, parser):
-        super(MigrateSchemasCommand, self).add_arguments(parser)
+        super(Command, self).add_arguments(parser)
         command = MigrateCommand()
         command.add_arguments(parser)
         parser.set_defaults(verbosity=0)
 
     def handle(self, *args, **options):
-        super(MigrateSchemasCommand, self).handle(*args, **options)
+        super(Command, self).handle(*args, **options)
         self.PUBLIC_SCHEMA_NAME = get_public_schema_name()
 
+        executor = get_executor(codename=self.executor)(self.args, self.options)
+
         if self.sync_public and not self.schema_name:
             self.schema_name = self.PUBLIC_SCHEMA_NAME
 
         if self.sync_public:
-            self.run_migrations(self.schema_name, settings.SHARED_APPS)
+            executor.run_migrations(tenants=[self.schema_name])
         if self.sync_tenant:
             if self.schema_name and self.schema_name != self.PUBLIC_SCHEMA_NAME:
                 if not schema_exists(self.schema_name):
-                    raise RuntimeError('Schema "{}" does not exist'.format(
+                    raise MigrationSchemaMissing('Schema "{}" does not exist'.format(
                         self.schema_name))
                 else:
-                    self.run_migrations(self.schema_name, settings.TENANT_APPS)
+                    tenants = [self.schema_name]
             else:
-                for tenant in TenantMiddleware.get_tenants():
-                    self.run_migrations(tenant.schema_name, settings.TENANT_APPS)
-
-    def run_migrations(self, schema_name, included_apps):
-        if int(self.options.get('verbosity', 1)) >= 1:
-            self._notice("=== Running migrate for schema %s" % schema_name)
-        connection.set_schema(schema_name, include_public=False)
-        command = MigrateCommand()
-        command.execute(*self.args, **self.options)
-        connection.set_schema_to_public()
+                tenants = [tenant.schema_name for tenant in TenantMiddleware.get_tenants()
+                           if tenant.schema_name != get_public_schema_name()]
+            executor.run_migrations(tenants=tenants)
 
     def _notice(self, output):
         self.stdout.write(self.style.NOTICE(output))
-
-Command = MigrateSchemasCommand
-- 
2.20.1

