hobo/hobo/multitenant/management/commands/__init__.py

# 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

import django
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
from tenant_schemas.utils import get_public_schema_name

from hobo.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().__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.
            obj._original_command = app_name
        else:
            obj._original_command = load_command_class(app_name, obj.COMMAND_NAME)

        # prepend the command's original help with the info about schemata
        # iteration
        obj.help = (
            "Calls {cmd} for all registered schemata. You can use regular "
            "{cmd} options.\n\nOriginal help for {cmd}:\n\n{help}".format(
                cmd=obj.COMMAND_NAME,
                help=getattr(obj._original_command, 'help', 'none'),
            )
        )

        return obj

    def add_arguments(self, parser):
        super().add_arguments(parser)
        parser.add_argument("-d", "--domain", dest="domain")
        # 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)

    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():
                self.execute_command(tenant, self.COMMAND_NAME, *args, **options)


class InteractiveTenantOption:
    def add_arguments(self, parser):
        parser.add_argument("-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'):
            domain = options['domain']
        else:
            if connection.get_tenant().schema_name != 'public':
                return connection.get_tenant()
            displayed_tenants = all_tenants
            while True:
                domain = input("Enter Tenant Domain ('?' to list): ")
                if domain in [t.domain_url for t in all_tenants]:
                    break
                try:
                    domain = displayed_tenants[int(domain) - 1].domain_url
                    break
                except (ValueError, IndexError):
                    pass
                if domain == '?':
                    displayed_tenants = all_tenants
                else:
                    displayed_tenants = [x for x in all_tenants if domain in x.domain_url]
                for i, tenant in enumerate(displayed_tenants):
                    print("[%2d] %s (schema %s)" % (i + 1, tenant.domain_url, tenant.schema_name))

        return TenantMiddleware.get_tenant_by_hostname(domain)


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().__new__(cls, *args, **kwargs)
        obj.command_instance = obj.COMMAND()
        return obj

    def add_arguments(self, parser):
        super().add_arguments(parser)
        self.command_instance.add_arguments(parser)

    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):
    def add_arguments(self, parser):
        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='?',
            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("-s", "--schema", dest="schema_name")

    def handle(self, *args, **options):
        self.domain = options.get('domain')
        self.installed_apps = settings.INSTALLED_APPS
        self.args = args
        self.options = options

        if self.domain:
            self.schema_name = TenantMiddleware.hostname2schema(self.domain)
        else:
            self.schema_name = options.get('schema_name')

        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):
        if int(self.options.get('verbosity', 1)) >= 1:
            self.stdout.write(self.style.NOTICE(output))


def disable_global_logging():
    import logging
    import os
    import sys

    if 'TERM' not in os.environ:
        return

    # try to disable sentry
    if 'sentry_sdk' in sys.modules:
        import sentry_sdk

        sentry_sdk.init()

    # then try to disable journald, syslog logging and mail admins
    root_logger = logging.getLogger()
    for handler in list(root_logger.handlers):
        hdlr_class_name = handler.__class__.__name__
        if hdlr_class_name in ['JournalHandler', 'SysLogHandler', 'AdminEmailHandler']:
            root_logger.handlers.remove(handler)