1 from optparse import make_option

     2 

     3 from django.conf import settings

     4 from django.core.management.base import AppCommand, CommandError

     5 from django.core.management.color import no_style

     6 from django.core.management.sql import sql_reset

     7 from django.db import connections, transaction, DEFAULT_DB_ALIAS

     8 

     9 class Command(AppCommand):

    10     option_list = AppCommand.option_list + (

    11         make_option('--noinput', action='store_false', dest='interactive', default=True,

    12             help='Tells Django to NOT prompt the user for input of any kind.'),

    13         make_option('--database', action='store', dest='database',

    14             default=DEFAULT_DB_ALIAS, help='Nominates a database to reset. '

    15                 'Defaults to the "default" database.'),

    16     )

    17     help = "Executes ``sqlreset`` for the given app(s) in the current database."

    18     args = '[appname ...]'

    19 

    20     output_transaction = True

    21 

    22     def handle_app(self, app, **options):

    23         using = options.get('database', DEFAULT_DB_ALIAS)

    24         connection = connections[using]

    25 

    26         app_name = app.__name__.split('.')[-2]

    27         self.style = no_style()

    28 

    29         sql_list = sql_reset(app, self.style, connection)

    30 

    31         if options.get('interactive'):

    32             confirm = raw_input("""

    33 You have requested a database reset.

    34 This will IRREVERSIBLY DESTROY any data for

    35 the "%s" application in the database "%s".

    36 Are you sure you want to do this?

    37 

    38 Type 'yes' to continue, or 'no' to cancel: """ % (app_name, connection.settings_dict['NAME']))

    39         else:

    40             confirm = 'yes'

    41 

    42         if confirm == 'yes':

    43             try:

    44                 cursor = connection.cursor()

    45                 for sql in sql_list:

    46                     cursor.execute(sql)

    47             except Exception, e:

    48                 transaction.rollback_unless_managed()

    49                 raise CommandError("""Error: %s couldn't be reset. Possible reasons:

    50   * The database isn't running or isn't configured correctly.

    51   * At least one of the database tables doesn't exist.

    52   * The SQL was invalid.

    53 Hint: Look at the output of 'django-admin.py sqlreset %s'. That's the SQL this command wasn't able to run.

    54 The full error: %s""" % (app_name, app_name, e))

    55             transaction.commit_unless_managed()

    56         else:

    57             print "Reset cancelled."