hobo/hobo/multitenant/management/commands/tenant_command.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

import argparse
import sys

import django
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.core.management import call_command, get_commands, load_command_class
from django.core.management.base import BaseCommand, CommandError, SystemCheckError, handle_default_options
from django.db import connection, connections
from django.utils.encoding import force_str

from hobo.multitenant.management.commands import InteractiveTenantOption
from hobo.multitenant.middleware import TenantMiddleware


def exception_to_text(e):
    try:
        return str(e)
    except Exception:
        pass

    try:
        return force_str(str(e), errors='ignore')
    except Exception:
        pass

    try:
        return force_str(repr(e), errors='ignore')
    except Exception:
        pass

    return 'Unrepresentable exception'


def run_command_from_argv(command, argv):
    # copied/adapted from Django run_from_argv
    command._called_from_command_line = True
    parser = command.create_parser(argv[0], argv[1])

    options = parser.parse_args(argv[2:])
    cmd_options = vars(options)
    # Move positional args out of options to mimic legacy optparse
    args = cmd_options.pop('args', ())
    handle_default_options(options)
    try:
        command.execute(*args, **cmd_options)
    except Exception as e:
        if options.traceback:
            raise

        # SystemCheckError takes care of its own formatting.
        if isinstance(e, SystemCheckError):
            command.stderr.write(str(e), lambda x: x)
        else:
            command.stderr.write(
                '%s: %s: %s' % (connection.tenant, e.__class__.__name__, exception_to_text(e))
            )
        return e


class Command(InteractiveTenantOption, BaseCommand):
    help = "Wrapper around django commands for use with an individual tenant"
    args = '<other_command>'

    def run_from_argv(self, argv):
        """
        Changes the option_list to use the options from the wrapped command.
        Adds schema parameter to specify 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])

        klass.requires_migrations_checks = False
        klass.requires_system_checks = False

        # Ugly, but works. Delete tenant_command from the argv, parse the schema manually
        # and forward the rest of the arguments to the actual command being wrapped.
        del argv[1]
        args_parser = argparse.ArgumentParser()
        args_parser.add_argument("--all-tenants", help="apply command to all tenants", action='store_true')
        args_parser.add_argument("-d", "--domain", dest="domain_name", help="specify tenant domain name")
        args_parser.add_argument(
            '--force-job',
            dest='force_job',
            action='store_true',
            help='Run command even if DISABLE_CRON_JOBS is set',
        )
        args_namespace, args = args_parser.parse_known_args(argv)
        try:
            command = args[1]
        except IndexError:
            command = '(unknown)'

        # Continue weirdness: parse verbosity option and also leave it in args
        # for subcommand consumption
        verbosity_parser = argparse.ArgumentParser()
        verbosity_parser.add_argument(
            '-v', '--verbosity', action='store', dest='verbosity', default=1, type=int
        )
        args_verbosity, _ = verbosity_parser.parse_known_args(args)

        if (
            args_namespace.all_tenants
            and not args_namespace.force_job
            and getattr(settings, 'DISABLE_CRON_JOBS', False)
        ):
            if args_verbosity.verbosity > 0:
                print('Command %s is ignored because DISABLE_CRON_JOBS is set' % command)
            return

        if args_namespace.all_tenants:
            errors = []
            for tenant in TenantMiddleware.get_tenants():
                connection.set_tenant(tenant)
                if getattr(settings, 'TENANT_DISABLE_CRON_JOBS', False):
                    if args_verbosity.verbosity > 0 or args_namespace.force_job:
                        msg = 'Command %s is ignored on tenant %s because TENANT_DISABLE_CRON_JOBS is set' % (
                            command,
                            tenant.domain_url,
                        )
                        prefix = '* '
                        if args_namespace.force_job:
                            prefix = '* WARNING: '
                        print(prefix + msg)
                    continue
                if args_verbosity.verbosity > 1:
                    print('* Running command %s on tenant %s' % (command, tenant.domain_url))
                error = run_command_from_argv(klass, args)
                if error:
                    errors.append(error)
            try:
                connections.close_all()
            except ImproperlyConfigured:
                # Ignore if connections aren't setup at this point (e.g. no
                # configured settings).
                pass
            if errors:
                self.stderr.write('Command failed on multiple tenants')
                sys.exit(1)
        else:
            tenant = self.get_tenant_from_options_or_interactive(domain=args_namespace.domain_name)
            connection.set_tenant(tenant)
            klass.run_from_argv(args)

    def handle(self, *args, **options):
        tenant = self.get_tenant_from_options_or_interactive(**options)
        connection.set_tenant(tenant)
        call_command(*args, **options)