<s> from django . conf import settings <EOL> from django . core . management import call_command , get_commands , load_command_class <EOL> from django . core . management . base import BaseCommand , CommandError <EOL> from django . db import connection <EOL> try : <EOL> from django . utils . six . moves import input <EOL> except ImportError : <EOL> input = raw_input <EOL> from django_tenants . utils import get_tenant_model , get_public_schema_name <EOL> class BaseTenantCommand ( BaseCommand ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( cls , * args , ** kwargs ) : <EOL> """<STR_LIT>""" <EOL> obj = super ( BaseTenantCommand , cls ) . __new__ ( cls , * args , ** kwargs ) <EOL> app_name = get_commands ( ) [ obj . COMMAND_NAME ] <EOL> if isinstance ( app_name , BaseCommand ) : <EOL> cmdclass = app_name <EOL> else : <EOL> cmdclass = load_command_class ( app_name , obj . COMMAND_NAME ) <EOL> obj . help = "<STR_LIT>" "<STR_LIT>" % ( obj . COMMAND_NAME , obj . COMMAND_NAME , obj . COMMAND_NAME , <EOL> getattr ( cmdclass , '<STR_LIT>' , '<STR_LIT:none>' ) ) <EOL> return obj <EOL> def add_arguments ( self , parser ) : <EOL> super ( BaseTenantCommand , self ) . add_arguments ( parser ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , dest = "<STR_LIT>" ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , dest = "<STR_LIT>" , action = "<STR_LIT:store_true>" , default = False ) <EOL> def execute_command ( self , tenant , command_name , * args , ** options ) : <EOL> verbosity = int ( options . get ( '<STR_LIT>' ) ) <EOL> if verbosity >= <NUM_LIT:1> : <EOL> print ( ) <EOL> print ( self . style . NOTICE ( "<STR_LIT>" ) <EOL> + self . style . SQL_TABLE ( tenant . schema_name ) <EOL> + self . style . NOTICE ( "<STR_LIT>" % command_name ) ) <EOL> connection . set_tenant ( tenant ) <EOL> call_command ( command_name , * args , ** options ) <EOL> def handle ( self , * args , ** options ) : <EOL> """<STR_LIT>""" <EOL> if options [ '<STR_LIT>' ] : <EOL> connection . set_schema_to_public ( ) <EOL> self . execute_command ( get_tenant_model ( ) . objects . get ( schema_name = options [ '<STR_LIT>' ] ) , self . COMMAND_NAME , <EOL> * args , ** options ) <EOL> else : <EOL> for tenant in get_tenant_model ( ) . objects . all ( ) : <EOL> if not ( options [ '<STR_LIT>' ] and tenant . schema_name == get_public_schema_name ( ) ) : <EOL> self . execute_command ( tenant , self . COMMAND_NAME , * args , ** options ) <EOL> class InteractiveTenantOption ( object ) : <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> super ( InteractiveTenantOption , self ) . __init__ ( * args , ** kwargs ) <EOL> def add_arguments ( self , parser ) : <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , dest = "<STR_LIT>" , help = "<STR_LIT>" ) <EOL> def get_tenant_from_options_or_interactive ( self , ** options ) : <EOL> TenantModel = get_tenant_model ( ) <EOL> all_tenants = TenantModel . objects . all ( ) <EOL> if not all_tenants : <EOL> raise CommandError ( """<STR_LIT>""" ) <EOL> if options . get ( '<STR_LIT>' ) : <EOL> tenant_schema = options [ '<STR_LIT>' ] <EOL> else : <EOL> while True : <EOL> tenant_schema = input ( "<STR_LIT>" ) <EOL> if tenant_schema == '<STR_LIT:?>' : <EOL> print ( '<STR_LIT:\n>' . join ( [ "<STR_LIT:%s>" % t . schema_name for t in all_tenants ] ) ) <EOL> else : <EOL> break <EOL> if tenant_schema not in [ t . schema_name for t in all_tenants ] : <EOL> raise CommandError ( "<STR_LIT>" % ( tenant_schema , ) ) <EOL> return TenantModel . objects . get ( schema_name = tenant_schema ) <EOL> class TenantWrappedCommand ( InteractiveTenantOption , BaseCommand ) : <EOL> """<STR_LIT>""" <EOL> def __new__ ( cls , * args , ** kwargs ) : <EOL> obj = super ( TenantWrappedCommand , cls ) . __new__ ( cls , * args , ** kwargs ) <EOL> obj . command_instance = obj . COMMAND ( ) <EOL> obj . option_list = obj . command_instance . option_list <EOL> return obj <EOL> def handle ( self , * args , ** options ) : <EOL> tenant = self . get_tenant_from_options_or_interactive ( ** options ) <EOL> connection . set_tenant ( tenant ) <EOL> self . command_instance . execute ( * args , ** options ) <EOL> def add_arguments ( self , parser ) : <EOL> super ( TenantWrappedCommand , self ) . add_arguments ( parser ) <EOL> self . command_instance . add_arguments ( parser ) <EOL> class SyncCommon ( BaseCommand ) : <EOL> def add_arguments ( self , parser ) : <EOL> parser . add_argument ( '<STR_LIT>' , action = '<STR_LIT:store_true>' , dest = '<STR_LIT>' , default = False , <EOL> help = '<STR_LIT>' ) <EOL> parser . add_argument ( '<STR_LIT>' , action = '<STR_LIT:store_true>' , dest = '<STR_LIT>' , default = False , <EOL> help = '<STR_LIT>' ) <EOL> parser . add_argument ( "<STR_LIT>" , "<STR_LIT>" , dest = "<STR_LIT>" ) <EOL> parser . add_argument ( '<STR_LIT>' , action = '<STR_LIT:store>' , dest = '<STR_LIT>' , default = None , <EOL> help = '<STR_LIT>' ) <EOL> def handle ( self , * args , ** options ) : <EOL> self . sync_tenant = options . get ( '<STR_LIT>' ) <EOL> self . sync_public = options . get ( '<STR_LIT>' ) <EOL> self . schema_name = options . get ( '<STR_LIT>' ) <EOL> self . executor = options . get ( '<STR_LIT>' ) <EOL> self . installed_apps = settings . INSTALLED_APPS <EOL> self . args = args <EOL> self . options = options <EOL> if self . schema_name : <EOL> if self . sync_public : <EOL> raise CommandError ( "<STR_LIT>" ) <EOL> elif self . schema_name == get_public_schema_name ( ) : <EOL> self . sync_public = True <EOL> else : <EOL> self . sync_tenant = True <EOL> elif not self . sync_public and not self . sync_tenant : <EOL> self . sync_tenant = True <EOL> self . sync_public = True <EOL> if hasattr ( settings , '<STR_LIT>' ) : <EOL> self . tenant_apps = settings . TENANT_APPS <EOL> if hasattr ( settings , '<STR_LIT>' ) : <EOL> self . shared_apps = settings . SHARED_APPS <EOL> def _notice ( self , output ) : <EOL> self . stdout . write ( self . style . NOTICE ( output ) ) </s>
