# Database access

import random

import re

import sys

import ex_util

class NameException(Exception): pass

class DuplicateKeyException(Exception): 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, ex:

        ex_util.add_msg(ex, 'query: '+cur.query)

        raise

    return cur

def row(cur): return cur.fetchone()

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

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):

    for field in fields: check_name(field)

    for key in conds.keys(): check_name(key)

    def cond(entry):

        key, value = entry

        cond_ = key+' '

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

        else: cond_ += '='

        cond_ += ' %s'

        return cond_

    return run_query(db, 'SELECT '+', '.join(fields)+' FROM '+table+' WHERE '

        +' AND '.join(map(cond, conds.iteritems())), conds.values())

def insert(db, table, row, get_id=False):

    check_name(table)

    keys = row.keys()

    for key in keys: check_name(key)

    query = 'INSERT INTO '+table+' ('+', '.join(keys)+') VALUES ('\

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

    if get_id: query += ' RETURNING lastval()'

    cur = run_query(db, query, row.values())

    if get_id: return value(cur)

def insert_ignore(db, table, row, get_id=False):

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

    except Exception, ex:

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

            +table+'_(\w+)_index"', str(ex))

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

        else: raise

def insert_or_get(db, table, row, pkey):

    try: return insert_ignore(db, table, row, True)

    except DuplicateKeyException, ex:

        dup_key = str(ex)

        return value(select(db, table, [pkey], {dup_key: row[dup_key]}))