Schema.py

gistfile1.txt

import os
import psycopg2
import psycopg2.extras

from psycopg2 import ProgrammingError
from psycopg2 import OperationalError

from controllers.settings import test_connection as TEST_CONN
from controllers.settings import connection as CONN
from controllers.settings import cursor_composites
from controllers.settings import schema_file_path as SCHEMA_PATH
from controllers.settings import schema_values_path as VALUES_PATH
from controllers.settings import global_debug as GLOBAL_DEBUG

'''
    Method: convert
    Converts psycopg2 RealDictRow & list types to native python dictionaries and lists rescursively.

    Parameters:
        results - The results from the query
'''
def convert(results):
    if type(results) == list:
        results = list(results)
        for i in range(len(results)):
            results[i] = convert(results[i])
    elif type(results) == psycopg2.extras.RealDictRow:
        results = dict(results)
    return results

'''
    Class: DictComposite
    Dictionary composite factory for psycopg2. Use to transform returned psycopg2 composites into dictionaries.

    Example:
        >>> class DictComposite(psycopg2.extras.CompositeCaster):
        ...     def make(self, values):
        ...         return dict(zip(self.attnames, values))

        >>> psycopg2.extras.register_composite('card', cur,
        ...     factory=DictComposite)

        >>> cur.execute("select (8, 'hearts')::card")
        >>> cur.fetchone()[0]
        {'suit': 'hearts', 'value': 8}
'''
class DictComposite(psycopg2.extras.CompositeCaster):
    def make(self, values):
        return dict(zip(self.attnames, values))

'''
    Method: make_composite

        Returns a cursor made for retrieving dictionary composites from the database by instantiating with the DictComposite Dictionary.

    Parameters:

        cur - A cursor instance instantiated from the psycopg2 connection.
'''
def make_composite(cur):
    for ccomposite in cursor_composites:
        psycopg2.extras.register_composite(ccomposite, cur, factory=DictComposite)
    return cur

'''
    Class: Schema
    Base class for working aggressively with the database.

    Attributes:
        name - The name of the database
        user - The user of the database
        port - The port of the database
        path - The path to the database schema
        values_path - The path to default database values
'''
class Schema():
    def __init__(self, testing=True, schema_path=SCHEMA_PATH):
        self.conn = CONN
        if testing is True:
            self.conn = TEST_CONN

        self.name = self.conn.get('database')
        self.user = self.conn.get('user')
        self.port = self.conn.get('port')
        self.path = SCHEMA_PATH
        self.values_path = VALUES_PATH

        # drop and build
        drop_db = raw_input('\n[DANGER] Drop and rebuild the database schema? (y)es or (n)o: ')
        if drop_db in ['Y','y']:
            os.system('dropdb -p %s %s' % (self.port, self.name))
            self.build()
        else:
            self.conn = psycopg2.connect(**self.conn)

    '''
        Method: empty
        Needed to empty the database after it has already been created & setup.

    '''
    def empty(self):
        cur = self.conn.cursor()

        cur.execute('SELECT * FROM truncate(%s);', (self.user,))
        cur.fetchall()
        cur.close()
        self.conn.commit()

        os.system("psql -U %s -E -p %s -d %s --command '\i %s'" % (self.user, self.port, self.name, self.values_path))

        return True

    '''
        Method: build
        Creates the database via system commands. 
        Adds the plpythonu trusted language. (capability)
        Adds the intarray extension. (capability)
        Optionally loads in the schema line-by-line.
    '''
    def build(self):
        # create database
        os.system('createdb -p %s --owner %s %s' % (self.port, self.user, self.name))

        # add python to database
        os.system("psql -U %s -p %s -d %s --command 'create or replace trusted language plpythonu'" % (self.user, self.port, self.name))

        # trust python
        os.system("psql -U %s -p %s -d %s --command \"UPDATE pg_language SET lanpltrusted = true WHERE lanname = 'plpythonu';\"" % (self.user, self.port, self.name))

        # grant all to python
        os.system("psql -U %s -p %s -d %s --command 'grant all on language plpythonu to %s'" % (self.user, self.port, self.name, self.user))

        # add extensions to database
        os.system("psql -U %s -p %s -d %s --command 'CREATE EXTENSION intarray'"  % (self.user, self.port, self.name))

        # open up the connection
        self.conn = psycopg2.connect(**self.conn)

        # test schema if desired
        slow_test_db = raw_input('Test database schema (line-by-line, this is slow)? (y)es or (n)o: ')
        if slow_test_db in ['y', 'Y']: # load the tables
            cur = self.conn.cursor()
            schema = open(self.path).read()
            cur.execute(schema)
            self.conn.commit()
        else: # build the database using psql \i meta-command
            os.system("psql -U %s -E -p %s -d %s --command '\i %s'" % (self.user, self.port, self.name, self.path))