gistfile1.txt

#! /usr/bin/env python
"""
Create, prepare and load schema for Socorro PostgreSQL database.
"""

import psycopg2
import psycopg2.extensions
from psycopg2 import ProgrammingError
import sys
import re
import logging

def createDBObject(sql, allowable_errors=None, databaseName='postgres'):
    conn = psycopg2.connect(database=databaseName, user='socorro')
    conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
    cur = conn.cursor()
    try:
        cur.execute(sql)
    except ProgrammingError, e:
        dberr = e.pgerror.strip().split('ERROR:  ')[1]
        if allowable_errors:
            for err in allowable_errors:
                if re.match(err, dberr):
                    logging.warn(dberr)
                else:
                    raise
        else:
            raise
    conn.close()

def main():
    databaseName = 'test'

    createDBObject('DROP DATABASE %s' % databaseName,
                   ['database "test" does not exist'])

    createDBObject('CREATE DATABASE %s' % databaseName,
                   ['database "test" already exists'])

    with open('sql/roles.sql') as f:
        for line in f:
            createDBObject(line, [r'role "\w+" already exists'])

    for lang in ['plpgsql', 'plperl']:
        createDBObject('CREATE LANGUAGE "%s"' % lang,
                       ['language "%s" already exists' % lang])

    createDBObject("GRANT ALL ON database %s TO breakpad_rw" % databaseName)

    with open('sql/schema.sql') as f:
        createDBObject(f.read(), databaseName=databaseName)

if __name__ == '__main__':
    main()