pbp_com/apps/pbp_core/createdb.py

import os
import time

import MySQLdb
import psycopg2
import _mysql_exceptions
from wsgi import *


def create_dbs():
    deadline = time.time() + 60
    while time.time() < deadline:
        try:
            print ("create_dbs: let's go.")
            django_settings = __import__(
                os.environ["DJANGO_SETTINGS_MODULE"], fromlist="DATABASES"
            )
            print ("create_dbs: got settings.")
            databases = django_settings.DATABASES
            for name, db in databases.iteritems():
                host = db["HOST"]
                user = db["USER"]
                password = db["PASSWORD"]
                port = db["PORT"]
                db_name = db["NAME"]
                db_type = db["ENGINE"]
                # see if it is mysql
                if db_type.endswith("mysql"):
                    print "creating database %s on %s" % (db_name, host)
                    db = MySQLdb.connect(
                        user=user, passwd=password, host=host, port=port
                    )
                    cur = db.cursor()
                    print ("Check if database is already there.")
                    cur.execute(
                        """SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA
                                 WHERE SCHEMA_NAME = %s""",
                        (db_name,),
                    )
                    results = cur.fetchone()
                    if not results:
                        print ("Database %s doesn't exist, lets create it." % db_name)
                        sql = """CREATE DATABASE IF NOT EXISTS %s """ % (db_name,)
                        print ("> %s" % sql)
                        cur.execute(sql)
                        print (".....")
                    else:
                        print ("database already exists, moving on to next step.")
                    exit(0)
                # see if it is postgresql
                elif db_type.endswith("postgresql_psycopg2"):
                    print "creating database %s on %s" % (db_name, host)
                    con = psycopg2.connect(
                        host=host,
                        user=user,
                        password=password,
                        port=port,
                        database="postgres",
                    )
                    con.set_isolation_level(0)
                    cur = con.cursor()
                    try:
                        cur.execute("CREATE DATABASE %s" % db_name)
                    except psycopg2.ProgrammingError as detail:
                        print detail
                        print "moving right along..."
                    exit(0)
                else:
                    print (
                        "ERROR: {0} is not supported by this script, you will need to create your database by hand.".format(
                            db_type
                        )
                    )
                    exit(1)
        except psycopg2.OperationalError:
            print "Could not connect to database. Waiting a little bit."
            time.sleep(10)
        except _mysql_exceptions.OperationalError:
            print "Could not connect to database. Waiting a little bit."
            time.sleep(10)

    print "Could not connect to database after 1 minutes. Something is wrong."
    exit(1)


if __name__ == "__main__":
    import sys

    print ("create_dbs start")
    create_dbs()
    print ("create_dbs all done")