From 2600a93bd178ba38c11019b99f6efb755bb1cd50 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20P=C3=A9ters?= <fpeters@entrouvert.com>
Date: Wed, 3 Jul 2019 07:16:46 +0200
Subject: [PATCH] revert "migrate_schemas: remove migration of public schema"
 (#34504)

This reverts commit 81e955d803e667b374e8d6c976d886394f513ac9.
---
 .../management/commands/__init__.py           | 23 ++++++-
 .../management/commands/migrate_schemas.py    | 66 +++++++++++--------
 2 files changed, 59 insertions(+), 30 deletions(-)

diff --git a/hobo/multitenant/management/commands/__init__.py b/hobo/multitenant/management/commands/__init__.py
index 5ed2b73..139d61c 100644
--- a/hobo/multitenant/management/commands/__init__.py
+++ b/hobo/multitenant/management/commands/__init__.py
@@ -52,6 +52,8 @@ class BaseTenantCommand(BaseCommand):
     def add_arguments(self, parser):
         super(BaseTenantCommand, self).add_arguments(parser)
         parser.add_argument("-d", "--domain", dest="domain")
+        parser.add_argument("-p", "--skip-public", dest="skip_public",
+                            action="store_true", default=False)
         # use the privately held reference to the underlying command to invoke
         # the add_arguments path on this parser instance
         self._original_command.add_arguments(parser)
@@ -80,7 +82,8 @@ class BaseTenantCommand(BaseCommand):
             self.execute_command(TenantMiddleware.get_tenant_by_hostname(options['domain']), self.COMMAND_NAME, *args, **options)
         else:
             for tenant in TenantMiddleware.get_tenants():
-                self.execute_command(tenant, self.COMMAND_NAME, *args, **options)
+                if not (options['skip_public'] and tenant.schema_name == get_public_schema_name()):
+                    self.execute_command(tenant, self.COMMAND_NAME, *args, **options)
 
 
 class InteractiveTenantOption(object):
@@ -149,6 +152,10 @@ class TenantWrappedCommand(InteractiveTenantOption, BaseCommand):
 class SyncCommon(BaseCommand):
 
     def add_arguments(self, parser):
+        parser.add_argument('--tenant', action='store_true', dest='tenant', default=False,
+                    help='Tells Django to populate only tenant applications.')
+        parser.add_argument('--shared', action='store_true', dest='shared', default=False,
+                    help='Tells Django to populate only shared applications.')
         parser.add_argument('--app_label', action='store', dest='app_label', nargs='?',
                     help='App label of an application to synchronize the state.')
         parser.add_argument('--migration_name', action='store', dest='migration_name', nargs='?',
@@ -157,6 +164,8 @@ class SyncCommon(BaseCommand):
         parser.add_argument("-d", "--domain", dest="domain")
 
     def handle(self, *args, **options):
+        self.sync_tenant = options.get('tenant')
+        self.sync_public = options.get('shared')
         self.domain = options.get('domain')
         self.installed_apps = settings.INSTALLED_APPS
         self.args = args
@@ -167,6 +176,18 @@ class SyncCommon(BaseCommand):
         else:
             self.schema_name = options.get('schema_name')
 
+        if self.schema_name:
+            if self.sync_public:
+                raise CommandError("domain should only be used with the --tenant switch.")
+            elif self.schema_name == get_public_schema_name():
+                self.sync_public = True
+            else:
+                self.sync_tenant = True
+        elif not self.sync_public and not self.sync_tenant:
+            # no options set, sync both
+            self.sync_tenant = True
+            self.sync_public = True
+
         if hasattr(settings, 'TENANT_APPS'):
             self.tenant_apps = settings.TENANT_APPS
         if hasattr(settings, 'SHARED_APPS'):
diff --git a/hobo/multitenant/management/commands/migrate_schemas.py b/hobo/multitenant/management/commands/migrate_schemas.py
index 579c1d8..05c26d5 100644
--- a/hobo/multitenant/management/commands/migrate_schemas.py
+++ b/hobo/multitenant/management/commands/migrate_schemas.py
@@ -47,36 +47,44 @@ class MigrateSchemasCommand(SyncCommon):
 
     def handle(self, *args, **options):
         super(MigrateSchemasCommand, self).handle(*args, **options)
-        if self.schema_name:
-            if not schema_exists(self.schema_name):
-                raise RuntimeError('Schema "{}" does not exist'.format(
-                    self.schema_name))
+        self.PUBLIC_SCHEMA_NAME = get_public_schema_name()
+
+        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)
+        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(
+                        self.schema_name))
+                else:
+                    self.run_migrations(self.schema_name, settings.TENANT_APPS)
             else:
-                self.run_migrations(self.schema_name, settings.TENANT_APPS)
-        else:
-            app_labels = []
-            for app in apps.get_app_configs():
-                if app.name in settings.TENANT_APPS:
-                    app_labels.append(app.label)
-            loader = MigrationLoader(connection)
-            loader.load_disk()
-            recorder = MigrationRecorder(connection)
-            applied_public_migrations = set(
-                [(app, migration)
-                 for app, migration in recorder.applied_migrations()
-                 if app in app_labels and (app, migration) in loader.disk_migrations])
-            for tenant in TenantMiddleware.get_tenants():
-                connection.set_schema(tenant.schema_name, include_public=False)
-                applied_migrations = self.get_applied_migrations(app_labels)
-                if options.get('fake') or options.get('migration_name') or options.get('app_label'):
-                    # never skip migrations if explicit migration actions
-                    # are given.
-                    applied_migrations = []
-                if all([x in applied_migrations for x in applied_public_migrations]):
-                    if int(self.options.get('verbosity', 1)) >= 1:
-                        self._notice("=== Skipping migrations of schema %s" % tenant.schema_name)
-                    continue
-                self.run_migrations(tenant.schema_name, settings.TENANT_APPS)
+                app_labels = []
+                for app in apps.get_app_configs():
+                    if app.name in settings.TENANT_APPS:
+                        app_labels.append(app.label)
+                loader = MigrationLoader(connection)
+                loader.load_disk()
+                recorder = MigrationRecorder(connection)
+                applied_public_migrations = set(
+                    [(app, migration)
+                     for app, migration in recorder.applied_migrations()
+                     if app in app_labels and (app, migration) in loader.disk_migrations])
+                for tenant in TenantMiddleware.get_tenants():
+                    connection.set_schema(tenant.schema_name, include_public=False)
+                    applied_migrations = self.get_applied_migrations(app_labels)
+                    if options.get('fake') or options.get('migration_name') or options.get('app_label'):
+                        # never skip migrations if explicit migration actions
+                        # are given.
+                        applied_migrations = []
+                    if all([x in applied_migrations for x in applied_public_migrations]):
+                        if int(self.options.get('verbosity', 1)) >= 1:
+                            self._notice("=== Skipping migrations of schema %s" % tenant.schema_name)
+                        continue
+                    self.run_migrations(tenant.schema_name, settings.TENANT_APPS)
 
     def get_applied_migrations(self, app_labels):
         applied_migrations = []
-- 
2.20.1

