# Database access

import random
import re
import sys

import ex
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): ex.add_msg(e, 'query: '+get_cur_query(cur))

class NameException(Exception): pass

class DbException(ex.ExceptionWithCause):
    def __init__(self, msg, cause=None, cur=None):
        ex.ExceptionWithCause.__init__(self, msg, cause)
        if cur != None: _add_cursor_info(self, cur)

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 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 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]}))

db_engines = {
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
    'PostgreSQL': ('psycopg2', {}),
}

def connect(db_config):
    db_config = db_config.copy() # don't modify input!
    module, mappings = db_engines[db_config.pop('engine')]
    for orig, new in mappings.iteritems(): util.rename_key(db_config, orig, new)
    return __import__(module).connect(**db_config)