def set_schema_to_public(self):
     """
     Instructs to stay in the common 'public' schema.
     """
     self.tenant = FakeTenant(schema_name=get_public_schema_name())
     self.schema_name = get_public_schema_name()
     self.set_settings_schema(self.schema_name)

    def setUpClass(cls):
        super(SharedAuthTest, cls).setUpClass()
        settings.SHARED_APPS = ('tenant_schemas',
                                'django.contrib.auth',
                                'django.contrib.contenttypes', )
        settings.TENANT_APPS = ('dts_test_app', )
        settings.INSTALLED_APPS = settings.SHARED_APPS + settings.TENANT_APPS
        cls.sync_shared()
        Tenant(domain_url='test.com', schema_name=get_public_schema_name()).save()

        # Create a tenant
        cls.tenant = Tenant(domain_url='tenant.test.com', schema_name='tenant')
        cls.tenant.save()

        # Create some users
        with schema_context(get_public_schema_name()):  # this could actually also be executed inside a tenant
            cls.user1 = User(username='arbitrary-1', email="[email protected]")
            cls.user1.save()
            cls.user2 = User(username='arbitrary-2', email="[email protected]")
            cls.user2.save()

        # Create instances on the tenant that point to the users on public
        with tenant_context(cls.tenant):
            cls.d1 = ModelWithFkToPublicUser(user=cls.user1)
            cls.d1.save()
            cls.d2 = ModelWithFkToPublicUser(user=cls.user2)
            cls.d2.save()

 def sync_shared(cls):
     if django.VERSION >= (1, 7, 0):
         call_command('migrate_schemas',
                      schema_name=get_public_schema_name(),
                      interactive=False,
                      verbosity=cls.get_verbosity())
     else:
         call_command('sync_schemas',
                      schema_name=get_public_schema_name(),
                      tenant=False,
                      public=True,
                      interactive=False,
                      migrate_all=True,
                      verbosity=cls.get_verbosity())

    def setUpClass(cls):
        settings.TENANT_MODEL = 'tenant_schemas.Tenant'
        settings.SHARED_APPS = ('tenant_schemas', )
        settings.TENANT_APPS = ('dts_test_app',
                                'django.contrib.contenttypes',
                                'django.contrib.auth', )
        settings.INSTALLED_APPS = settings.SHARED_APPS + settings.TENANT_APPS

        # Django calls syncdb by default for the test database, but we want
        # a blank public schema for this set of tests.
        connection.set_schema_to_public()
        cursor = connection.cursor()
        cursor.execute('DROP SCHEMA %s CASCADE; CREATE SCHEMA %s;'
                       % (get_public_schema_name(), get_public_schema_name(), ))
        super(BaseTestCase, cls).setUpClass()

    def save(self, verbosity=1, *args, **kwargs):
        is_new = self.pk is None

        if is_new and connection.schema_name != get_public_schema_name():
            raise Exception("Can't create tenant outside the public schema. Current schema is %s."
                            % connection.schema_name)
        elif not is_new and connection.schema_name not in (self.schema_name, get_public_schema_name()):
            raise Exception("Can't update tenant outside it's own schema or the public schema. Current schema is %s."
                            % connection.schema_name)

        super(TenantMixin, self).save(*args, **kwargs)

        if is_new and self.auto_create_schema:
            self.create_schema(check_if_exists=True, verbosity=verbosity)
            post_schema_sync.send(sender=TenantMixin, tenant=self)

    def _cursor(self):
        """
        Here it happens. We hope every Django db operation using PostgreSQL
        must go through this to get the cursor handle. We change the path.
        """
        cursor = super(DatabaseWrapper, self)._cursor()

        # Actual search_path modification for the cursor. Database will
        # search schemata from left to right when looking for the object
        # (table, index, sequence, etc.).
        if not self.schema_name:
            raise ImproperlyConfigured("Database schema not set. Did you forget "
                                       "to call set_schema() or set_tenant()?")
        _check_identifier(self.schema_name)
        public_schema_name = get_public_schema_name()
        search_paths = []

        if self.schema_name == public_schema_name:
            search_paths = [public_schema_name]
        elif self.include_public_schema:
            search_paths = [self.schema_name, public_schema_name]
        else:
            search_paths = [self.schema_name]

        search_paths.extend(EXTRA_SEARCH_PATHS)
        cursor.execute('SET search_path = {}'.format(','.join(search_paths)))
        return cursor

    def handle_noargs(self, **options):
        # todo: awful lot of duplication from sync_schemas
        sync_tenant = options.get('tenant')
        sync_public = options.get('shared')
        schema_name = options.get('schema_name')
        self.installed_apps = settings.INSTALLED_APPS
        self.options = options

        if sync_public and schema_name:
            raise CommandError("schema should only be used with the --tenant switch.")
        if not hasattr(settings, 'TENANT_APPS') and sync_tenant:
            raise CommandError("No setting found for TENANT_APPS")
        if not hasattr(settings, 'SHARED_APPS') and sync_public:
            raise CommandError("No setting found for SHARED_APPS")

        if not sync_public and not sync_tenant and not schema_name:
            # no options set, sync both
            sync_tenant = True
            sync_public = True

        if schema_name:
            if schema_name == get_public_schema_name():
                sync_public = True
            else:
                sync_tenant = True

        if hasattr(settings, 'TENANT_APPS'):
            self.tenant_apps = settings.TENANT_APPS
        if hasattr(settings, 'SHARED_APPS'):
            self.shared_apps = settings.SHARED_APPS

        if sync_public:
            self.migrate_public_apps()
        if sync_tenant:
            self.migrate_tenant_apps(schema_name)

    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)

        syncdb_command = MigrateCommand()
        if schema_name:
            print self.style.NOTICE("=== Running migrate for schema: %s" % schema_name)
            connection.set_schema_to_public()
            sync_tenant = get_tenant_model().objects.filter(schema_name=schema_name).get()
            connection.set_tenant(sync_tenant, include_public=False)
            syncdb_command.execute(**self.options)
        else:
            public_schema_name = get_public_schema_name()
            tenant_schemas_count = get_tenant_model().objects.exclude(schema_name=public_schema_name).count()
            if not tenant_schemas_count:
                print self.style.NOTICE("No tenants found")

            for tenant_schema in get_tenant_model().objects.exclude(schema_name=public_schema_name).all():
                Migrations._dependencies_done = False  # very important, the dependencies need to be purged from cache
                print self.style.NOTICE("=== Running migrate for schema %s" % tenant_schema.schema_name)
                connection.set_tenant(tenant_schema, include_public=False)
                syncdb_command.execute(**self.options)

        self._restore_south_settings()

    def process_request(self, request):
        """
        Resets to public schema

        Some nasty weird bugs happened at the production environment without this call.
        connection.pg_thread.schema_name would already be set and then terrible errors
        would occur. Any idea why? My theory is django implements connection as some sort
        of threading local variable.
        """
        connection.set_schema_to_public()
        request.tenant = self.set_tenant(request.get_host())

        # do we have tenant-specific URLs?
        if (
            hasattr(settings, "PUBLIC_SCHEMA_URL_TOKEN")
            and request.tenant.schema_name == get_public_schema_name()
            and request.path_info[-1] == "/"
        ):
            # we are not at the public schema, manually alter routing to schema-dependent urls
            request.path_info = settings.PUBLIC_SCHEMA_URL_TOKEN + request.path_info

        if SET_TENANT_SITE_DYNAMICALLY and hasattr(request.tenant, "site") and request.tenant.site:
            SITE_THREAD_LOCAL.SITE_ID = request.tenant.site_id
            # dynamically set the site
        else:
            SITE_THREAD_LOCAL.SITE_ID = DEFAULT_SITE_ID

    def allow_migrate(self, db, app_label, model_name=None, **hints):
        # the imports below need to be done here else django <1.5 goes crazy
        # https://code.djangoproject.com/ticket/20704
        from django.db import connection
        from tenant_schemas.utils import get_public_schema_name, app_labels

        if isinstance(app_label, ModelBase):
            # In django <1.7 the `app_label` parameter is actually `model`
            app_label = app_label._meta.app_label

        if connection.schema_name == get_public_schema_name():
            if app_label not in app_labels(settings.SHARED_APPS):
                return False
        else:
            if app_label not in app_labels(settings.TENANT_APPS):
                return False

        if model_name:
            model = apps.get_model(app_label=app_label, model_name=model_name)
            if hasattr(model, '__schema_name__') and model.__schema_name__:

                conn_schema = connection.tenant.schema_name.strip()
                if conn_schema == 'public' and conn_schema != model.__schema_name__ or \
                   conn_schema != 'public' and model.__schema_name__ == 'public':
                    return False

        return None

    def allow_migrate(self, db, app_label, model_name=None, **hints):
        # the imports below need to be done here else django <1.5 goes crazy
        # https://code.djangoproject.com/ticket/20704
        from django.db import connection
        from tenant_schemas.utils import get_public_schema_name, app_labels
        from tenant_schemas.postgresql_backend.base import DatabaseWrapper as TenantDbWrapper

        db_engine = settings.DATABASES[db]["ENGINE"]
        if not (
            db_engine == "tenant_schemas.postgresql_backend"
            or issubclass(getattr(load_backend(db_engine), "DatabaseWrapper"), TenantDbWrapper)
        ):
            return None

        if isinstance(app_label, ModelBase):
            # In django <1.7 the `app_label` parameter is actually `model`
            app_label = app_label._meta.app_label

        if connection.schema_name == get_public_schema_name():
            if app_label not in app_labels(settings.SHARED_APPS):
                return False
        else:
            if app_label not in app_labels(settings.TENANT_APPS):
                return False

        return None

    def handle(self, *args, **options):
        self.sync_tenant = options.get('tenant')
        self.sync_public = options.get('shared')
        self.schema_name = options.get('schema_name')
        self.executor = options.get('executor')
        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 handle(self, *args, **options):
        self.non_tenant_schemas = settings.PG_EXTRA_SEARCH_PATHS + ['public']
        self.sync_tenant = options.get('tenant')
        self.sync_public = options.get('shared')
        self.schema_name = options.get('schema_name')
        self.args = args
        self.options = options
        self.PUBLIC_SCHEMA_NAME = get_public_schema_name()

        if self.schema_name:
            if self.sync_public:
                raise CommandError("schema should only be used "
                                   "with the --tenant switch.")
            elif self.schema_name == self.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 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):
                # Make sure the tenant exists and the schema belongs to
                # a tenant; We don't want to sync to extensions schema by
                # mistake
                if not schema_exists(self.schema_name):
                    raise RuntimeError('Schema "{}" does not exist'.format(
                        self.schema_name))
                elif self.schema_name in self.non_tenant_schemas:
                    raise RuntimeError(
                        'Schema "{}" does not belong to any tenant'.format(
                            self.schema_name))
                else:
                    self.run_migrations(self.schema_name, settings.TENANT_APPS)
            else:
                all_tenants = get_tenant_model().objects.exclude(
                    schema_name=get_public_schema_name())
                for tenant in all_tenants:
                    self.run_migrations(tenant.schema_name, settings.TENANT_APPS)

    def run_migrations(self, tenants):
        public_schema_name = get_public_schema_name()

        if public_schema_name in tenants:
            run_migrations(self.args, self.options, self.codename, public_schema_name)
            tenants.pop(tenants.index(public_schema_name))

        self.run_tenant_migrations(tenants)

    def get_tenant(self, model, hostname, request):
        try:
            return super(DefaultTenantMiddleware, self).get_tenant(
                model, hostname, request)
        except model.DoesNotExist:
            schema_name = self.DEFAULT_SCHEMA_NAME
            if not schema_name:
                schema_name = get_public_schema_name()

            return model.objects.get(schema_name=schema_name)

 def setUpClass(cls):
     super(RoutesTestCase, cls).setUpClass()
     settings.SHARED_APPS = ('tenant_schemas', )
     settings.TENANT_APPS = ('dts_test_app',
                             'django.contrib.contenttypes',
                             'django.contrib.auth', )
     settings.INSTALLED_APPS = settings.SHARED_APPS + settings.TENANT_APPS
     cls.sync_shared()
     cls.public_tenant = Tenant(domain_url='test.com', schema_name=get_public_schema_name())
     cls.public_tenant.save(verbosity=BaseTestCase.get_verbosity())

    def _cursor(self, name=None):
        """
        Here it happens. We hope every Django db operation using PostgreSQL
        must go through this to get the cursor handle. We change the path.
        """
        if name:
            # Only supported and required by Django 1.11 (server-side cursor)
            cursor = super(DatabaseWrapper, self)._cursor(name=name)
        else:
            cursor = super(DatabaseWrapper, self)._cursor()

        # optionally limit the number of executions - under load, the execution
        # of `set search_path` can be quite time consuming
        if (not get_limit_set_calls()) or not self.search_path_set:
            # Actual search_path modification for the cursor. Database will
            # search schemata from left to right when looking for the object
            # (table, index, sequence, etc.).
            if not self.schema_name:
                raise ImproperlyConfigured("Database schema not set. Did you forget "
                                           "to call set_schema() or set_tenant()?")
            _check_schema_name(self.schema_name)
            public_schema_name = get_public_schema_name()
            search_paths = []

            if self.schema_name == public_schema_name:
                search_paths = [public_schema_name]
            elif self.include_public_schema:
                search_paths = [self.schema_name, public_schema_name]
            else:
                search_paths = [self.schema_name]

            search_paths.extend(EXTRA_SEARCH_PATHS)

            if name:
                # Named cursor can only be used once
                cursor_for_search_path = self.connection.cursor()
            else:
                # Reuse
                cursor_for_search_path = cursor

            # In the event that an error already happened in this transaction and we are going
            # to rollback we should just ignore database error when setting the search_path
            # if the next instruction is not a rollback it will just fail also, so
            # we do not have to worry that it's not the good one
            try:
                cursor_for_search_path.execute('SET search_path = {0}'.format(','.join(search_paths)))
            except (django.db.utils.DatabaseError, psycopg2.InternalError):
                self.search_path_set = False
            else:
                self.search_path_set = True

            if name:
                cursor_for_search_path.close()

        return cursor

    def handle(self, *args, **options):
        super(MigrateSchemasCommand, self).handle(*args, **options)
        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:
                all_tenants = get_tenant_model().objects.exclude(schema_name=get_public_schema_name())
                for tenant in all_tenants:
                    self.run_migrations(tenant.schema_name, settings.TENANT_APPS)

 def set_tenant(self, tenant, include_public=True):
     """
     Main API method to current database schema,
     but it does not actually modify the db connection.
     """
     self.tenant = tenant
     if tenant is None:
         self.schema_name = get_public_schema_name()
     else:
         self.schema_name = tenant.schema_name
         self.include_public_schema = include_public

    def sync_tenant_apps(self, schema_name=None):
        if schema_name:
            tenant = get_tenant_model().objects.filter(schema_name=schema_name).get()
            self._sync_tenant(tenant)
        else:
            all_tenants = get_tenant_model().objects.exclude(schema_name=get_public_schema_name())
            if not all_tenants:
                self._notice("No tenants found!")

            for tenant in all_tenants:
                self._sync_tenant(tenant)