1 """

     2 originally from http://www.djangosnippets.org/snippets/828/ by dnordberg

     3 """

     4 

     5 

     6 from django.conf import settings

     7 from django.core.management.base import CommandError, BaseCommand

     8 from django.db import connection

     9 import logging

    10 from optparse import make_option

    11 

    12 class Command(BaseCommand):

    13     option_list = BaseCommand.option_list + (

    14         make_option('--noinput', action='store_false',

    15                     dest='interactive', default=True,

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

    17         make_option('--no-utf8', action='store_true',

    18                     dest='no_utf8_support', default=False,

    19                     help='Tells Django to not create a UTF-8 charset database'),

    20     )

    21     help = "Resets the database for this project."

    22 

    23     def handle(self, *args, **options):

    24         """

    25         Resets the database for this project.

    26     

    27         Note: Transaction wrappers are in reverse as a work around for

    28         autocommit, anybody know how to do this the right way?

    29         """

    30 

    31         if options.get('interactive'):

    32             confirm = raw_input("""

    33 You have requested a database reset.

    34 This will IRREVERSIBLY DESTROY

    35 ALL data in the database "%s".

    36 Are you sure you want to do this?

    37 

    38 Type 'yes' to continue, or 'no' to cancel: """ % (settings.DATABASE_NAME,))

    39         else:

    40             confirm = 'yes'

    41 

    42         if confirm != 'yes':

    43             print "Reset cancelled."

    44             return

    45 

    46         engine = settings.DATABASE_ENGINE

    47     

    48         if engine == 'sqlite3':

    49             import os

    50             try:

    51                 logging.info("Unlinking sqlite3 database")

    52                 os.unlink(settings.DATABASE_NAME)

    53             except OSError:

    54                 pass

    55         elif engine == 'mysql':

    56             import MySQLdb as Database

    57             kwargs = {

    58                 'user': settings.DATABASE_USER,

    59                 'passwd': settings.DATABASE_PASSWORD,

    60             }

    61             if settings.DATABASE_HOST.startswith('/'):

    62                 kwargs['unix_socket'] = settings.DATABASE_HOST

    63             else:

    64                 kwargs['host'] = settings.DATABASE_HOST

    65             if settings.DATABASE_PORT:

    66                 kwargs['port'] = int(settings.DATABASE_PORT)

    67             connection = Database.connect(**kwargs)

    68             drop_query = 'DROP DATABASE IF EXISTS %s' % settings.DATABASE_NAME

    69             utf8_support = options.get('no_utf8_support', False) and '' or 'CHARACTER SET utf8'

    70             create_query = 'CREATE DATABASE %s %s' % (settings.DATABASE_NAME, utf8_support)

    71             logging.info('Executing... "' + drop_query + '"')

    72             connection.query(drop_query)

    73             logging.info('Executing... "' + create_query + '"')

    74             connection.query(create_query)

    75         elif engine == 'postgresql' or engine == 'postgresql_psycopg2':

    76             if engine == 'postgresql':

    77                 import psycopg as Database

    78             elif engine == 'postgresql_psycopg2':

    79                 import psycopg2 as Database

    80             

    81             if settings.DATABASE_NAME == '':

    82                 from django.core.exceptions import ImproperlyConfigured

    83                 raise ImproperlyConfigured, "You need to specify DATABASE_NAME in your Django settings file."

    84             

    85             conn_string = "dbname=%s" % settings.DATABASE_NAME

    86             if settings.DATABASE_USER:

    87                 conn_string += " user=%s" % settings.DATABASE_USER

    88             if settings.DATABASE_PASSWORD:

    89                 conn_string += " password='%s'" % settings.DATABASE_PASSWORD

    90             if settings.DATABASE_HOST:

    91                 conn_string += " host=%s" % settings.DATABASE_HOST

    92             if settings.DATABASE_PORT:

    93                 conn_string += " port=%s" % settings.DATABASE_PORT

    94             connection = Database.connect(conn_string)

    95             connection.set_isolation_level(0) #autocommit false

    96             cursor = connection.cursor()

    97             drop_query = 'DROP DATABASE %s' % settings.DATABASE_NAME

    98             logging.info('Executing... "' + drop_query + '"')

    99     

   100             try:

   101                 cursor.execute(drop_query)

   102             except Database.ProgrammingError, e:

   103                 logging.info("Error: "+str(e))

   104     

   105             # Encoding should be SQL_ASCII (7-bit postgres default) or prefered UTF8 (8-bit)

   106             create_query = ("""

   107 CREATE DATABASE %s

   108     WITH OWNER = %s

   109         ENCODING = 'UTF8'

   110         TABLESPACE = pg_default;

   111 """ % (settings.DATABASE_NAME, settings.DATABASE_USER))

   112             logging.info('Executing... "' + create_query + '"')

   113             cursor.execute(create_query)

   114     

   115         else:

   116             raise CommandError, "Unknown database engine %s", engine

   117     

   118         print "Reset successful."