manyi/venv/Lib/site-packages/django_extensions/management/commands/reset_schema.py

# -*- coding: utf-8 -*-
"""
Recreates the public schema for current database (PostgreSQL only).
Useful for Docker environments where you need to reset database
schema while there are active connections.
"""

from django.core.management import BaseCommand, CommandError
from django.db import connections
from django.conf import settings
from six.moves import input


class Command(BaseCommand):
    """`reset_schema` command implementation."""

    help = "Recreates the public schema for this project."

    def add_arguments(self, parser):
        super().add_arguments(parser)
        parser.add_argument(
            '--noinput', action='store_false',
            dest='interactive', default=True,
            help='Tells Django to NOT prompt the user for input of any kind.'
        )
        parser.add_argument(
            '-R', '--router', action='store', dest='router', default='default',
            help='Use this router-database instead of the one defined in settings.py'
        )
        parser.add_argument(
            '-S', '--schema', action='store', dest='schema', default='public',
            help='Drop this schema instead of "public"'
        )

    def handle(self, *args, **options):
        router = options['router']
        dbinfo = settings.DATABASES.get(router)
        if dbinfo is None:
            raise CommandError("Unknown database router %s" % router)

        engine = dbinfo.get('ENGINE').split('.')[-1]
        if engine not in ('postgresql', 'postgresql_psycopg2', 'postgis'):
            raise CommandError('This command can be used only with PostgreSQL databases.')

        database_name = dbinfo['NAME']

        schema = options['schema']

        if options['interactive']:
            confirm = input("""
You have requested a database schema reset.
This will IRREVERSIBLY DESTROY ALL data
in the "{}" schema of database "{}".
Are you sure you want to do this?

Type 'yes' to continue, or 'no' to cancel: """.format(schema, database_name))
        else:
            confirm = 'yes'

        if confirm != 'yes':
            print("Reset cancelled.")
            return
        with connections[router].cursor() as cursor:
            cursor.execute("DROP SCHEMA {} CASCADE".format(schema))
            cursor.execute("CREATE SCHEMA {}".format(schema))