catalyst/zipline/db.py

import atexit
import pymongo
import zipline.util as qutil

class MongoOptions(object):
    
    def __init__(self, host, port, dbname, user, password):
        self.mongodb_host       = host
        self.mongodb_port       = port
        self.mongodb_dbname     = dbname
        self.mongodb_user       = user
        self.mongodb_password   = password

class NoDatabase(Exception):
    def __repr__(self):
        return 'The database has not been set up yet.'

def setup_db(credentials):
    """
    Setup the database. Has global side effects.
    """
    qutil.LOGGER.info(dir(DbConnection))
    if not DbConnection.initd:
        connector = connect_db(credentials)
        DbConnection.set(*connector)

def connect_db(options):
    """
    Connect to pymongo, return a connection and database instance
    as a tuple.
    """

    connection = pymongo.Connection(options.mongodb_host, options.mongodb_port)

    db = connection[options.mongodb_dbname]
    db.authenticate(options.mongodb_user, options.mongodb_password)

    def _gc_connection(): # pragma: no cover
        connection.close()

    atexit.register(_gc_connection)
    return connection, db

class DbConnection(object):
    """
    Hold the shared state of the database connection.
    """

    initd       = False
    __shared    = {}
    
    def __init__(self):
        self.__dict__ = self.__shared

    @staticmethod
    def set(conn, db):
        DbConnection.__shared['conn'] = conn
        DbConnection.__shared['db'] = db
        DbConnection.initd = True

    @staticmethod
    def get():
        return (
            DbConnection.__shared['conn'],
            DbConnection.__shared['db']
        )

    def __getattr__(self, key):
        if not DbConnection.__shared.get('initd'):
            raise NoDatabase()
        else:
            return DbConnection.__shared.get(key)

    def destory(self): # pragma: no cover
        DbConnection.__shared['initd'] = False
        self.conn.close()