# Database access

import random

import re

import sys

import exc

import util

def get_cur_query(cur):

    if hasattr(cur, 'query'): return cur.query

    elif hasattr(cur, '_last_executed'): return cur._last_executed

    else: return None

def _add_cursor_info(e, cur): exc.add_msg(e, 'query: '+get_cur_query(cur))

class DbException(exc.ExceptionWithCause):

    def __init__(self, msg, cause=None, cur=None):

        exc.ExceptionWithCause.__init__(self, msg, cause)

        if cur != None: _add_cursor_info(self, cur)

class NameException(DbException): pass

class ExceptionWithColumn(DbException):

    def __init__(self, col, cause=None):

        DbException.__init__(self, 'column: '+col, cause)

        self.col = col

class DuplicateKeyException(ExceptionWithColumn): pass

class NullValueException(ExceptionWithColumn): pass

class EmptyRowException(DbException): pass

def check_name(name):

    if re.search(r'\W', name) != None: raise NameException('Name "'+name

        +'" may contain only alphanumeric characters and _')

def run_query(db, query, params=None):

    cur = db.cursor()

    try: cur.execute(query, params)

    except Exception, e:

        _add_cursor_info(e, cur)

        raise

    return cur

def col(cur, idx): return cur.description[idx][0]

def rows(cur): return iter(lambda: cur.fetchone(), None)

def row(cur): return rows(cur).next()

def value(cur): return row(cur)[0]

def values(cur): return iter(lambda: value(cur), None)

def value_or_none(cur):

    try: return value(cur)

    except StopIteration: return None

def with_savepoint(db, func):

    savepoint = 'savepoint_'+str(random.randint(0, sys.maxint)) # must be unique

    run_query(db, 'SAVEPOINT '+savepoint)

    try: return_val = func()

    except:

        run_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)

        raise

    else:

        run_query(db, 'RELEASE SAVEPOINT '+savepoint)

        return return_val

def select(db, table, fields, conds, limit=None):

    assert limit == None or type(limit) == int

    check_name(table)

    map(check_name, fields)

    map(check_name, conds.keys())

    def cond(entry):

        col, value = entry

        cond_ = col+' '

        if value == None: cond_ += 'IS'

        else: cond_ += '='

        cond_ += ' %s'

        return cond_

    query = 'SELECT '+', '.join(fields)+' FROM '+table

    if conds != {}:

        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))

    if limit != None: query += ' LIMIT '+str(limit)

    return run_query(db, query, conds.values())

def insert(db, table, row):

    check_name(table)

    cols = row.keys()

    map(check_name, cols)

    query = 'INSERT INTO '+table

    if row != {}: query += ' ('+', '.join(cols)+') VALUES ('\

        +', '.join(['%s']*len(cols))+')'

    else: query += ' DEFAULT VALUES'

    return run_query(db, query, row.values())

def last_insert_id(db):

    module = util.root_module(db)

    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))

    elif module == 'MySQLdb': return db.insert_id()

    else: return None

def constraint_cols(db, table, constraint):

    check_name(table)

    check_name(constraint)

    module = util.root_module(db)

    if module == 'psycopg2':

        return list(values(run_query(db, '''\

SELECT attname

FROM pg_constraint

JOIN pg_class ON pg_class.oid = conrelid

JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)

WHERE

    relname = %(table)s

    AND conname = %(constraint)s

ORDER BY attnum

''',

            {'table': table, 'constraint': constraint})))

    else: raise NotImplementedError("Can't list constraint columns for "+module+

        ' database')

def try_insert(db, table, row):

    try: return with_savepoint(db, lambda: insert(db, table, row))

    except Exception, e:

        msg = str(e)

        match = re.search(r'duplicate key value violates unique constraint "'

            +table+'_(\w+)_index"', msg)

        if match: raise DuplicateKeyException(match.group(1), e)

        match = re.search(r'null value in column "(\w+)" violates not-null '

            'constraint', msg)

        if match: raise NullValueException(match.group(1), e)

        raise # no specific exception raised

def pkey(db, cache, table): # Assumed to be first column in table

    check_name(table)

    if table not in cache:

        cache[table] = col(run_query(db, 'SELECT * FROM '+table+' LIMIT 0'), 0)

    return cache[table]

def get(db, table, row, pkey, create=False, row_ct_ref=None):

    try: return value(select(db, table, [pkey], row, 1))

    except StopIteration:

        if not create: raise

        # Insert new row

        try:

            row_ct = try_insert(db, table, row).rowcount

            if row_ct_ref != None and row_ct >= 0: row_ct_ref[0] += row_ct

            return last_insert_id(db)

        except DuplicateKeyException, e:

            return value(select(db, table, [pkey], {e.col: row[e.col]}))

def truncate(db, table):

    check_name(table)

    return run_query(db, 'TRUNCATE '+table+' CASCADE')

def tables(db):

    module = util.root_module(db)

    if module == 'psycopg2':

        return values(run_query(db, "SELECT tablename from pg_tables "

            "WHERE schemaname = 'public' ORDER BY tablename"))

    elif module == 'MySQLdb': return values(run_query(db, 'SHOW TABLES'))

    else: raise NotImplementedError("Can't list tables for "+module+' database')

def empty_db(db):

    for table in tables(db): truncate(db, table)

db_engines = {

    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),

    'PostgreSQL': ('psycopg2', {}),

}

DatabaseErrors_set = set([DbException])

DatabaseErrors = tuple(DatabaseErrors_set)

def _add_module(module):

    DatabaseErrors_set.add(module.DatabaseError)

    global DatabaseErrors

    DatabaseErrors = tuple(DatabaseErrors_set)

def connect(db_config):

    db_config = db_config.copy() # don't modify input!

    module_name, mappings = db_engines[db_config.pop('engine')]

    module = __import__(module_name)

    _add_module(module)

    for orig, new in mappings.iteritems():

        try: util.rename_key(db_config, orig, new)

        except KeyError: pass

    return module.connect(**db_config)