""" originally from http://www.djangosnippets.org/snippets/828/ by dnordberg """ from django.conf import settings from django.core.management.base import CommandError, BaseCommand from django.db import connection import logging from optparse import make_option class Command(BaseCommand): option_list = BaseCommand.option_list + ( make_option('--noinput', action='store_false', dest='interactive', default=True, help='Tells Django to NOT prompt the user for input of any kind.'), make_option('--no-utf8', action='store_true', dest='no_utf8_support', default=False, help='Tells Django to not create a UTF-8 charset database'), ) help = "Resets the database for this project." def handle(self, *args, **options): """ Resets the database for this project. Note: Transaction wrappers are in reverse as a work around for autocommit, anybody know how to do this the right way? """ if options.get('interactive'): confirm = raw_input(""" You have requested a database reset. This will IRREVERSIBLY DESTROY ALL data in the database "%s". Are you sure you want to do this? Type 'yes' to continue, or 'no' to cancel: """ % (settings.DATABASE_NAME,)) else: confirm = 'yes' if confirm != 'yes': print "Reset cancelled." return engine = settings.DATABASE_ENGINE if engine == 'sqlite3': import os try: logging.info("Unlinking sqlite3 database") os.unlink(settings.DATABASE_NAME) except OSError: pass elif engine == 'mysql': import MySQLdb as Database kwargs = { 'user': settings.DATABASE_USER, 'passwd': settings.DATABASE_PASSWORD, } if settings.DATABASE_HOST.startswith('/'): kwargs['unix_socket'] = settings.DATABASE_HOST else: kwargs['host'] = settings.DATABASE_HOST if settings.DATABASE_PORT: kwargs['port'] = int(settings.DATABASE_PORT) connection = Database.connect(**kwargs) drop_query = 'DROP DATABASE IF EXISTS %s' % settings.DATABASE_NAME utf8_support = options.get('no_utf8_support', False) and '' or 'CHARACTER SET utf8' create_query = 'CREATE DATABASE %s %s' % (settings.DATABASE_NAME, utf8_support) logging.info('Executing... "' + drop_query + '"') connection.query(drop_query) logging.info('Executing... "' + create_query + '"') connection.query(create_query) elif engine == 'postgresql' or engine == 'postgresql_psycopg2': if engine == 'postgresql': import psycopg as Database elif engine == 'postgresql_psycopg2': import psycopg2 as Database if settings.DATABASE_NAME == '': from django.core.exceptions import ImproperlyConfigured raise ImproperlyConfigured, "You need to specify DATABASE_NAME in your Django settings file." conn_string = "dbname=%s" % settings.DATABASE_NAME if settings.DATABASE_USER: conn_string += " user=%s" % settings.DATABASE_USER if settings.DATABASE_PASSWORD: conn_string += " password='%s'" % settings.DATABASE_PASSWORD if settings.DATABASE_HOST: conn_string += " host=%s" % settings.DATABASE_HOST if settings.DATABASE_PORT: conn_string += " port=%s" % settings.DATABASE_PORT connection = Database.connect(conn_string) connection.set_isolation_level(0) #autocommit false cursor = connection.cursor() drop_query = 'DROP DATABASE %s' % settings.DATABASE_NAME logging.info('Executing... "' + drop_query + '"') try: cursor.execute(drop_query) except Database.ProgrammingError, e: logging.info("Error: "+str(e)) # Encoding should be SQL_ASCII (7-bit postgres default) or prefered UTF8 (8-bit) create_query = (""" CREATE DATABASE %s WITH OWNER = %s ENCODING = 'UTF8' TABLESPACE = pg_default; """ % (settings.DATABASE_NAME, settings.DATABASE_USER)) logging.info('Executing... "' + create_query + '"') cursor.execute(create_query) else: raise CommandError, "Unknown database engine %s", engine print "Reset successful."