From c0eed0031946832a35fcd00f8a5d0cbf509a1d78 Mon Sep 17 00:00:00 2001
From: Benjamin Dauvergne <bdauvergne@entrouvert.com>
Date: Tue, 19 Aug 2014 15:59:15 +0200
Subject: [PATCH 08/11] Import django-tenant-schemas commands to adapt them to
 our way of managing tenants

---
 .../multitenant/management/commands/__init__.py    |  166 ++++++++++++++++++++
 .../management/commands/createsuperuser.py         |   11 ++
 .../multitenant/management/commands/migrate.py     |   23 +++
 .../management/commands/migrate_schemas.py         |   91 +++++++++++
 .../management/commands/sync_schemas.py            |   80 ++++++++++
 .../multitenant/management/commands/syncdb.py      |   26 +++
 .../management/commands/tenant_command.py          |   44 ++++++
 7 files changed, 441 insertions(+)
 create mode 100644 entrouvert/djommon/multitenant/management/commands/createsuperuser.py
 create mode 100644 entrouvert/djommon/multitenant/management/commands/migrate.py
 create mode 100644 entrouvert/djommon/multitenant/management/commands/migrate_schemas.py
 create mode 100644 entrouvert/djommon/multitenant/management/commands/sync_schemas.py
 create mode 100644 entrouvert/djommon/multitenant/management/commands/syncdb.py
 create mode 100644 entrouvert/djommon/multitenant/management/commands/tenant_command.py

diff --git a/entrouvert/djommon/multitenant/management/commands/__init__.py b/entrouvert/djommon/multitenant/management/commands/__init__.py
index e69de29..70e075b 100644
--- a/entrouvert/djommon/multitenant/management/commands/__init__.py
+++ b/entrouvert/djommon/multitenant/management/commands/__init__.py
@@ -0,0 +1,166 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from optparse import make_option
+from django.conf import settings
+from django.core.management import call_command, get_commands, load_command_class
+from django.core.management.base import BaseCommand, CommandError
+from django.db import connection
+try:
+    from django.utils.six.moves import input
+except ImportError:
+    input = raw_input
+from tenant_schemas.utils import get_public_schema_name
+from entrouvert.djommon.multitenant.middleware import TenantMiddleware
+
+
+class BaseTenantCommand(BaseCommand):
+    """
+    Generic command class useful for iterating any existing command
+    over all schemata. The actual command name is expected in the
+    class variable COMMAND_NAME of the subclass.
+    """
+    def __new__(cls, *args, **kwargs):
+        """
+        Sets option_list and help dynamically.
+        """
+        obj = super(BaseTenantCommand, cls).__new__(cls, *args, **kwargs)
+
+        app_name = get_commands()[obj.COMMAND_NAME]
+        if isinstance(app_name, BaseCommand):
+            # If the command is already loaded, use it directly.
+            cmdclass = app_name
+        else:
+            cmdclass = load_command_class(app_name, obj.COMMAND_NAME)
+
+        # inherit the options from the original command
+        obj.option_list = cmdclass.option_list
+        obj.option_list += (
+            make_option("-d", "--domain", dest="domain"),
+        )
+        obj.option_list += (
+            make_option("-p", "--skip-public", dest="skip_public", action="store_true", default=False),
+        )
+
+        # prepend the command's original help with the info about schemata iteration
+        obj.help = "Calls %s for all registered schemata. You can use regular %s options. "\
+                   "Original help for %s: %s" % (obj.COMMAND_NAME, obj.COMMAND_NAME, obj.COMMAND_NAME,
+                                                 getattr(cmdclass, 'help', 'none'))
+        return obj
+
+    def execute_command(self, tenant, command_name, *args, **options):
+        verbosity = int(options.get('verbosity'))
+
+        if verbosity >= 1:
+            print()
+            print(self.style.NOTICE("=== Switching to schema '") \
+                + self.style.SQL_TABLE(tenant.schema_name)\
+                + self.style.NOTICE("' then calling %s:" % command_name))
+
+        connection.set_tenant(tenant)
+
+        # call the original command with the args it knows
+        call_command(command_name, *args, **options)
+
+    def handle(self, *args, **options):
+        """
+        Iterates a command over all registered schemata.
+        """
+        if options['domain']:
+            # only run on a particular schema
+            connection.set_schema_to_public()
+            self.execute_command(TenantMiddleware.get_tenant_by_hostname(options['domain']), self.COMMAND_NAME, *args, **options)
+        else:
+            for tenant in TenantMiddleware.get_tenants():
+                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):
+    def __init__(self, *args, **kwargs):
+        super(InteractiveTenantOption, self).__init__(*args, **kwargs)
+        self.option_list += (
+            make_option("-d", "--domain", dest="domain", help="specify tenant domain"),
+        )
+
+    def get_tenant_from_options_or_interactive(self, **options):
+        all_tenants = list(TenantMiddleware.get_tenants())
+
+        if not all_tenants:
+            raise CommandError("""There are no tenants in the system.
+To learn how create a tenant, see:
+https://django-tenant-schemas.readthedocs.org/en/latest/use.html#creating-a-tenant""")
+
+        if options.get('domain'):
+            tenant_schema = options['domain']
+        else:
+            while True:
+                tenant_schema = input("Enter Tenant Domain ('?' to list schemas): ")
+                if tenant_schema == '?':
+                    print('\n'.join(["%s - %s" % (t.schema_name, t.domain_url,) for t in all_tenants]))
+                else:
+                    break
+
+        if tenant_schema not in [t.schema_name for t in all_tenants]:
+            raise CommandError("Invalid tenant schema, '%s'" % (tenant_schema,))
+
+        return TenantMiddleware.get_tenant_by_hostname(tenant_schema)
+
+
+class TenantWrappedCommand(InteractiveTenantOption, BaseCommand):
+    """
+    Generic command class useful for running any existing command
+    on a particular tenant. The actual command name is expected in the
+    class variable COMMAND_NAME of the subclass.
+    """
+    def __new__(cls, *args, **kwargs):
+        obj = super(TenantWrappedCommand, cls).__new__(cls, *args, **kwargs)
+        obj.command_instance = obj.COMMAND()
+        obj.option_list = obj.command_instance.option_list
+        return obj
+
+    def handle(self, *args, **options):
+        tenant = self.get_tenant_from_options_or_interactive(**options)
+        connection.set_tenant(tenant)
+
+        self.command_instance.execute(*args, **options)
+
+
+class SyncCommon(BaseCommand):
+    option_list = (
+        make_option('--tenant', action='store_true', dest='tenant', default=False,
+                    help='Tells Django to populate only tenant applications.'),
+        make_option('--shared', action='store_true', dest='shared', default=False,
+                    help='Tells Django to populate only shared applications.'),
+        make_option("-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
+        self.options = options
+
+        if self.schema_name:
+            if self.sync_public:
+                raise CommandError("schema 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'):
+            self.shared_apps = settings.SHARED_APPS
+
+    def _notice(self, output):
+        self.stdout.write(self.style.NOTICE(output))
diff --git a/entrouvert/djommon/multitenant/management/commands/createsuperuser.py b/entrouvert/djommon/multitenant/management/commands/createsuperuser.py
new file mode 100644
index 0000000..54fc9f3
--- /dev/null
+++ b/entrouvert/djommon/multitenant/management/commands/createsuperuser.py
@@ -0,0 +1,11 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from entrouvert.djommon.multitenant.management.commands import TenantWrappedCommand
+from django.contrib.auth.management.commands import createsuperuser
+
+
+class Command(TenantWrappedCommand):
+    COMMAND = createsuperuser.Command
diff --git a/entrouvert/djommon/multitenant/management/commands/migrate.py b/entrouvert/djommon/multitenant/management/commands/migrate.py
new file mode 100644
index 0000000..ffa0fb1
--- /dev/null
+++ b/entrouvert/djommon/multitenant/management/commands/migrate.py
@@ -0,0 +1,23 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from django.conf import settings
+from django.core.management.base import CommandError, BaseCommand
+try:
+    from south.management.commands.migrate import Command as MigrateCommand
+except ImportError:
+    MigrateCommand = BaseCommand
+
+
+class Command(MigrateCommand):
+
+    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)
diff --git a/entrouvert/djommon/multitenant/management/commands/migrate_schemas.py b/entrouvert/djommon/multitenant/management/commands/migrate_schemas.py
new file mode 100644
index 0000000..f67bda5
--- /dev/null
+++ b/entrouvert/djommon/multitenant/management/commands/migrate_schemas.py
@@ -0,0 +1,91 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from django.conf import settings
+from django.db import connection
+from south import migration
+from south.migration.base import Migrations
+from south.management.commands.migrate import Command as MigrateCommand
+from entrouvert.djommon.multitenant.middleware import TenantMiddleware
+from entrouvert.djommon.multitenant.management.commands import SyncCommon
+
+
+class Command(SyncCommon):
+    help = "Migrate schemas with South"
+    option_list = MigrateCommand.option_list + SyncCommon.option_list
+
+    def handle(self, *args, **options):
+        super(Command, self).handle(*args, **options)
+
+        if self.sync_public:
+            self.migrate_public_apps()
+        if self.sync_tenant:
+            self.migrate_tenant_apps(self.schema_name)
+
+    def _set_managed_apps(self, included_apps, excluded_apps):
+        """ while sync_schemas works by setting which apps are managed, on south we set which apps should be ignored """
+        ignored_apps = []
+        if excluded_apps:
+            for item in excluded_apps:
+                if item not in included_apps:
+                    ignored_apps.append(item)
+
+        for app in ignored_apps:
+            app_label = app.split('.')[-1]
+            settings.SOUTH_MIGRATION_MODULES[app_label] = 'ignore'
+
+    def _save_south_settings(self):
+        self._old_south_modules = None
+        if hasattr(settings, "SOUTH_MIGRATION_MODULES") and settings.SOUTH_MIGRATION_MODULES is not None:
+            self._old_south_modules = settings.SOUTH_MIGRATION_MODULES.copy()
+        else:
+            settings.SOUTH_MIGRATION_MODULES = dict()
+
+    def _restore_south_settings(self):
+        settings.SOUTH_MIGRATION_MODULES = self._old_south_modules
+
+    def _clear_south_cache(self):
+        for mig in list(migration.all_migrations()):
+            delattr(mig._application, "migrations")
+        Migrations._clear_cache()
+
+    def _migrate_schema(self, tenant):
+        connection.set_tenant(tenant, include_public=False)
+        MigrateCommand().execute(*self.args, **self.options)
+
+    def migrate_tenant_apps(self, schema_name=None):
+        self._save_south_settings()
+
+        apps = self.tenant_apps or self.installed_apps
+        self._set_managed_apps(included_apps=apps, excluded_apps=self.shared_apps)
+
+        if schema_name:
+            self._notice("=== Running migrate for schema: %s" % schema_name)
+            connection.set_schema_to_public()
+            tenant = TenantMiddleware.get_tenant_by_hostname(schema_name)
+            self._migrate_schema(tenant)
+        else:
+            all_tenants = TenantMiddleware.get_tenants()
+            if not all_tenants:
+                self._notice("No tenants found")
+
+            for tenant in all_tenants:
+                Migrations._dependencies_done = False  # very important, the dependencies need to be purged from cache
+                self._notice("=== Running migrate for schema %s" % tenant.schema_name)
+                self._migrate_schema(tenant)
+
+        self._restore_south_settings()
+
+    def migrate_public_apps(self):
+        self._save_south_settings()
+
+        apps = self.shared_apps or self.installed_apps
+        self._set_managed_apps(included_apps=apps, excluded_apps=self.tenant_apps)
+
+        self._notice("=== Running migrate for schema public")
+        MigrateCommand().execute(*self.args, **self.options)
+
+        self._clear_south_cache()
+        self._restore_south_settings()
diff --git a/entrouvert/djommon/multitenant/management/commands/sync_schemas.py b/entrouvert/djommon/multitenant/management/commands/sync_schemas.py
new file mode 100644
index 0000000..6b126eb
--- /dev/null
+++ b/entrouvert/djommon/multitenant/management/commands/sync_schemas.py
@@ -0,0 +1,80 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from django.conf import settings
+from django.contrib.contenttypes.models import ContentType
+from django.db.models import get_apps, get_models
+if "south" in settings.INSTALLED_APPS:
+    from south.management.commands.syncdb import Command as SyncdbCommand
+else:
+    from django.core.management.commands.syncdb import Command as SyncdbCommand
+from django.db import connection
+from entrouvert.djommon.multitenant.middleware import TenantMiddleware
+from entrouvert.djommon.multitenant.management.commands import SyncCommon
+
+
+class Command(SyncCommon):
+    help = "Sync schemas based on TENANT_APPS and SHARED_APPS settings"
+    option_list = SyncdbCommand.option_list + SyncCommon.option_list
+
+    def handle(self, *args, **options):
+        super(Command, self).handle(*args, **options)
+
+        if "south" in settings.INSTALLED_APPS:
+            self.options["migrate"] = False
+
+        # save original settings
+        for model in get_models(include_auto_created=True):
+            setattr(model._meta, 'was_managed', model._meta.managed)
+
+        ContentType.objects.clear_cache()
+
+        if self.sync_public:
+            self.sync_public_apps()
+        if self.sync_tenant:
+            self.sync_tenant_apps(self.schema_name)
+
+        # restore settings
+        for model in get_models(include_auto_created=True):
+            model._meta.managed = model._meta.was_managed
+
+    def _set_managed_apps(self, included_apps):
+        """ sets which apps are managed by syncdb """
+        for model in get_models(include_auto_created=True):
+            model._meta.managed = False
+
+        verbosity = int(self.options.get('verbosity'))
+        for app_model in get_apps():
+            app_name = app_model.__name__.replace('.models', '')
+            if app_name in included_apps:
+                for model in get_models(app_model, include_auto_created=True):
+                    model._meta.managed = model._meta.was_managed
+                    if model._meta.managed and verbosity >= 3:
+                        self._notice("=== Include Model: %s: %s" % (app_name, model.__name__))
+
+    def _sync_tenant(self, tenant):
+        self._notice("=== Running syncdb for schema: %s" % tenant.schema_name)
+        connection.set_tenant(tenant, include_public=False)
+        SyncdbCommand().execute(**self.options)
+
+    def sync_tenant_apps(self, schema_name=None):
+        apps = self.tenant_apps or self.installed_apps
+        self._set_managed_apps(apps)
+        if schema_name:
+            tenant = TenantMiddleware.get_tenant_by_hostname(schema_name)
+            self._sync_tenant(tenant)
+        else:
+            all_tenants = TenantMiddleware.get_tenants()
+            if not all_tenants:
+                self._notice("No tenants found!")
+
+            for tenant in all_tenants:
+                self._sync_tenant(tenant)
+
+    def sync_public_apps(self):
+        apps = self.shared_apps or self.installed_apps
+        self._set_managed_apps(apps)
+        SyncdbCommand().execute(**self.options)
+        self._notice("=== Running syncdb for schema public")
diff --git a/entrouvert/djommon/multitenant/management/commands/syncdb.py b/entrouvert/djommon/multitenant/management/commands/syncdb.py
new file mode 100644
index 0000000..63ba276
--- /dev/null
+++ b/entrouvert/djommon/multitenant/management/commands/syncdb.py
@@ -0,0 +1,26 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from django.core.management.base import CommandError
+from django.conf import settings
+from tenant_schemas.utils import django_is_in_test_mode
+
+try:
+    from south.management.commands import syncdb
+except ImportError:
+    from django.core.management.commands import syncdb
+
+
+class Command(syncdb.Command):
+
+    def handle(self, *args, **options):
+        database = options.get('database', 'default')
+        if (settings.DATABASES[database]['ENGINE'] == 'tenant_schemas.postgresql_backend' and not
+                django_is_in_test_mode()):
+            raise CommandError("syncdb has been disabled, for database '{}'. "
+                               "Use sync_schemas instead. Please read the "
+                               "documentation if you don't know why "
+                               "you shouldn't call syncdb directly!".format(database))
+        super(Command, self).handle(*args, **options)
diff --git a/entrouvert/djommon/multitenant/management/commands/tenant_command.py b/entrouvert/djommon/multitenant/management/commands/tenant_command.py
new file mode 100644
index 0000000..8ddd9d9
--- /dev/null
+++ b/entrouvert/djommon/multitenant/management/commands/tenant_command.py
@@ -0,0 +1,44 @@
+# this file derive from django-tenant-schemas
+#   Author: Bernardo Pires Carneiro
+#   Email: carneiro.be@gmail.com
+#   License: MIT license
+#   Home-page: http://github.com/bcarneiro/django-tenant-schemas
+from optparse import make_option
+from django.core.management.base import BaseCommand, CommandError
+from django.core.management import call_command, get_commands, load_command_class
+from django.db import connection
+from entrouvert.djommon.multitenant.management.commands import InteractiveTenantOption
+
+
+class Command(InteractiveTenantOption, BaseCommand):
+    help = "Wrapper around django commands for use with an individual tenant"
+
+    def run_from_argv(self, argv):
+        """
+        Changes the option_list to use the options from the wrapped command.
+        Adds schema parameter to specifiy which schema will be used when
+        executing the wrapped command.
+        """
+        # load the command object.
+        try:
+            app_name = get_commands()[argv[2]]
+        except KeyError:
+            raise CommandError("Unknown command: %r" % argv[2])
+
+        if isinstance(app_name, BaseCommand):
+            # if the command is already loaded, use it directly.
+            klass = app_name
+        else:
+            klass = load_command_class(app_name, argv[2])
+
+        self.option_list = klass.option_list + (
+            make_option("-d", "--domain", dest="domain", help="specify tenant schema"),
+        )
+
+        super(Command, self).run_from_argv(argv)
+
+    def handle(self, *args, **options):
+        tenant = self.get_tenant_from_options_or_interactive(**options)
+        connection.set_tenant(tenant)
+
+        call_command(*args, **options)
-- 
1.7.10.4

