root/GratiaWeb/db.py

import cherrypy, re, types

class IncorrectSchemaError(Exception):
    pass

class SqlTranslator(object):

    def __init__(self):
        self.default_sql = {}
        self.dialect_sql = {}
        self.schema = {}

    def register_sql(self, query_name, sql_str, dialect = None):
        if dialect == None:
            if query_name in self.default_sql:
                raise ValueError("Query name `%s` already registered." % \
                    query_name)
            self.default_sql[query_name] = sql_str
        else:
            if not self.dialect_sql.has_key(dialect):
                self.dialect_sql[dialect] = {}
            if query_name in self.dialect_sql[dialect]:
                raise ValueError("Query name `%s` already registered for " \
                    "dialect `%s`." % (query_name, dialect))
            self.dialect_sql[dialect][query_name] = sql_str

    def register_schema(self, name, schema, version, dialect=None):
        self.schema[dialect] = self.schema.get(dialect, {})
        self.schema[dialect][name] = self.schema[dialect].get(name, {})
        if isinstance(schema, types.StringType):
            schema = [schema]
        self.schema[dialect][name][version] = schema

    def get_schema(self, name, version, dialect=None):
        schema = self.schema.get(dialect, {}).get(name, {}).get(version, None)
        if schema == None:
            schema = self.schema.get(None, {}).get(name, {}).get(version, None)
        if schema == None:
            raise ValueError("Schema name %s, version %s, dialect %s not " \
                             "found!" % (name, version, dialect))
        return schema

    def get_sql(self, query_name, dialect=None):
        if dialect == None:
            return self.default_sql[query_name]
        elif not (dialect in self.dialect_sql):
            return self.default_sql[query_name]
        elif query_name in self.dialect_sql[dialect]:
            return self.dialect_sql[dialect][query_name]
        else:
            return self.default_sql[query_name]

translator = SqlTranslator()

def get_config():
    if cherrypy.request.config:
        return cherrypy.request.config
    else:
        return cherrypy.config


def find_db():
    conn_dict = get_config().get('db.connections', {})
    conn_str = None
    for key, val in conn_dict.items():
        if re.search(key, cherrypy.request.path_info):
            conn_str = val
    if conn_str == None:
        conn_str = get_config().get('db.connection', None)
    if conn_str == None:
        raise ValueError("DB Connection string not provided")
    cherrypy.request.db = db.get_db(conn_str)


class DBHandler(object):

    def __init__(self):
        self.db_conns = {}

    def __getitem__(self, attr_name):
        db = self.find_db(attr_name)
        return db

    def __getattr__(self, attr_name):
        db = self.find_db()
        return getattr(db, attr_name)

    def find_db(self, module=None):
        if module != None:
            prefix = 'db.' + module + '.'
        else:
            prefix = 'db.'
        conn_dict = get_config().get(prefix + 'connections', {})
        conn_str = None
        for key, val in conn_dict.items():
            if re.search(key, cherrypy.request.path_info):
                conn_str = val
        if conn_str == None:
            conn_str = get_config().get(prefix + 'connection', None)
        if conn_str == None:
            try:
                return self.find_db()
            except:
                raise ValueError("DB Connection string not provided")
        if conn_str in self.db_conns:
            return self.db_conns[conn_str]
        self.db_conns[conn_str] = DBConn(conn_str)
        return self.db_conns[conn_str]
 
db = DBHandler()

cherrypy.tools.find_db = cherrypy.Tool('before_handler', db.find_db)

class DBConn(object):

    conn_re = re.compile('(\w*)://(.*)')

    def __init__(self, conn_str):
        if conn_str == None:
            raise ValueError("`db.connection` config parameter not passed!")
        m = self.conn_re.match(conn_str)
        if not m:
            raise ValueError("Unable to parse connection string `%s`" % \
                conn_str)
        proto, conn_str = m.groups()
        self.dialect = proto
        self.conn_str = conn_str
        if proto == 'mysql':
            mysql_re = re.compile('(.*)/(.*)(\?.*)')
            m = mysql_re.match(conn_str)
            if not m:
                raise ValueError("String `%s` does not match format for MySQL"\
                    % conn_str)
            host, dbname, params = m.groups()
            self.execute = self._mysql_execute
            self.connect = self._mysql_connect
            self.module = __import__('MySQLdb')
        if proto == 'sqlite':
            self.execute = self._sqlite_execute
            self.connect = self._sqlite_connect
            try:
                import pysqlite2.dbapi2 as sqlite
            except ImportError:
                try:
                    import sqlite3 as sqlite
                except ImportError:
                    import sqlite
            self.module = sqlite
            self._get_conn = self._sqlite_get_conn
        if proto == 'oracle':
            raise NotImplementedError()
        self.connect(conn_str)
        self._schema_vers = {}
        self.execute_sql = self.execute

    def register_sql(self, query_name, sql_str, dialect=None):
        translator.register_sql(query_name, sql_str, dialect)

    def _mysql_connect(self, conn_str):
        mysql_re = re.compile('(.*)/(.*)(\?.*)')
        m = mysql_re.match(conn_str)
        if not m:
            raise ValueError("String `%s` does not match format for MySQL"\
                             % conn_str)
        host, dbname, params = m.groups()
        conn_args = {}
        if host.find(':') >= 0:
            host, port = host.split(':')
            conn_args['port'] = int(port)
        conn_args['host'] = host
        conn_args['db'] = dbname
        if len(params) > 0:
            params = params[1:]
            kws = params.split('&')
            for kw in kws:
                key, val = kw.split('=')
                conn_args[key] = val
        if 'password' in conn_args:
            conn_args['passwd'] = conn_args['password']
            del conn_args['password']
        self._conn = self.module.connect(**conn_args)

    def _sqlite_connect(self, conn_str):
        cherrypy.thread_data._conn = self.module.connect(conn_str)

    def _sqlite_get_conn(self):
        if hasattr(cherrypy.thread_data, '_conn'):
            return cherrypy.thread_data._conn
        self.connect(self.conn_str)
        return cherrypy.thread_data._conn

    def _get_conn(self):
        return self._conn

    def _get_cursor(self):
        return self._get_conn().cursor()

    def execute(self, query_name, bind_vars):
        raise NotImplementedError()

    def _mysql_execute( self, query_name, sql_vars ):
        my_string = translator.get_sql(query_name, self.dialect)
        sql_vars = dict( sql_vars )
        placement_dict = {}
        for var_name in sql_vars.keys():
            var_string = ':' + var_name
            placement = my_string.find( var_string )
            var_string_len = len(var_string)
            while placement >= 0:
                placement_dict[placement] = var_name
                my_string = my_string[:placement] + '%s' + \
                    my_string[placement+var_string_len:]
                placement = my_string.find( var_string )
        places = placement_dict.keys(); places.sort()
        my_tuple = ()
        for place in places:
            my_tuple += (sql_vars[placement_dict[place]],)
        curs = self._get_cursor()
        curs.arraysize = 500
        curs.execute( my_string, my_tuple )
        return curs

    def _sqlite_execute( self, query_name, sql_vars ):
        my_string = translator.get_sql(query_name, self.dialect)
        sql_vars = dict( sql_vars )
        placement_dict = {}
        for var_name in sql_vars.keys():
            var_string = ':' + var_name
            placement = my_string.find( var_string )
            var_string_len = len(var_string)
            while placement >= 0:
                placement_dict[placement] = var_name
                my_string = my_string[:placement] + '?' + \
                    my_string[placement+var_string_len:]
                placement = my_string.find( var_string )
        places = placement_dict.keys(); places.sort()
        my_tuple = ()
        for place in places:
            my_tuple += (sql_vars[placement_dict[place]],)
        curs = self._get_cursor()
        curs.arraysize = 500
        #print my_string, my_tuple
        curs.execute( my_string, my_tuple )
        return curs

    def verify_schema(self, name, version, create=False):
        if name in self._schema_vers and self._schema_vers[name] == version:
            return True
        try:
            curs = db.execute("get_schema", {'name': name})
        except Exception, e:
            if str(e).find('no such table') >= 0:
                self.create_schema('schema_verify', '0.1', dialect=self.dialect)
                curs = db.execute("get_schema", {'name': name})
            else:
                raise
        rows = curs.fetchall()
        has_schema = False
        has_vers = ''
        for row in rows:
            if row[0] == version:
                has_schema = True
                break
        if not has_schema:
            if not create:
                if has_vers != '':
                    raise IncorrectSchemaError("Schema for %s is currently %s;"\
                        " required version is %s." % (name, has_vers, version))
                else:
                    raise IncorrectSchemaError("No schema for module %s found!"\
                        % name)
            else:
                self.create_schema(name, version, dialect=self.dialect)
        self._schema_vers[name] = version
        return True

    def create_schema(self, name, version, dialect=None):
        schema = self.get_schema(name, version, dialect)
        curs = self._get_cursor()
        for statement in schema:
            print statement
            curs.execute(statement)
        self.execute_sql('update_schema', {'name': name, 'version': version})
        self.commit()

    def register_schema(self, name, schema, version, dialect=None):
        translator.register_schema(name, schema, version, dialect)

    def get_schema(self, name, version, dialect=None):
        return translator.get_schema(name, version, dialect)

    def commit(self):
        self._get_conn().commit()

verify_schema = ["""
CREATE TABLE Schema (
  NAME varchar(255),
  VERSION varchar(255)
)
""", """
CREATE UNIQUE INDEX Schema_unique on Schema (NAME, VERSION)
"""]
translator.register_schema('schema_verify', verify_schema, '0.1', 'sqlite')

update_schema_sql = """
INSERT OR REPLACE INTO Schema(NAME, VERSION) VALUES
    (:name, :version)
"""
translator.register_sql('update_schema', update_schema_sql, 'sqlite')

get_schema_sql = """
SELECT VERSION from Schema where NAME=:name
"""
translator.register_sql('get_schema', get_schema_sql, 'sqlite')