# Database access

import copy

import re

import sys

import warnings

import exc

import iters

from Proxy import Proxy

import rand

import strings

import util

##### Exceptions

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 ExceptionWithColumns(DbException):

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

        DbException.__init__(self, 'columns: ' + ', '.join(cols), cause)

        self.cols = cols

class DuplicateKeyException(ExceptionWithColumns): pass

class NullValueException(ExceptionWithColumns): pass

class EmptyRowException(DbException): pass

##### Warnings

class DbWarning(UserWarning): pass

##### Input validation

def check_name(name):

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

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

def esc_name(db, name):

    module = util.root_module(db.db)

    if module == 'psycopg2': return name

        # Don't enclose in quotes because this disables case-insensitivity

    elif module == 'MySQLdb': quote = '`'

    else: raise NotImplementedError("Can't escape name for "+module+' database')

    return quote + name.replace(quote, '') + quote

##### Database connections

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 db_config_str(db_config):

    return db_config['engine']+' database '+db_config['database']

def _query_lookup(query, params): return (query, util.cast(tuple, params))

class DbConn:

    def __init__(self, db_config, serializable=True, debug=False):

        self.db_config = db_config

        self.serializable = serializable

        self.debug = debug

        self.__db = None

        self.pkeys = {}

        self.index_cols = {}

        self.query_results = {}

    def __getattr__(self, name):

        if name == '__dict__': raise Exception('getting __dict__')

        if name == 'db': return self._db()

        else: raise AttributeError()

    def __getstate__(self):

        state = copy.copy(self.__dict__) # shallow copy

        state['_DbConn__db'] = None # don't pickle the connection

        return state

    def _db(self):

        if self.__db == None:

            # Process db_config

            db_config = self.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

            # Connect

            self.__db = module.connect(**db_config)

            # Configure connection

            if self.serializable: run_raw_query(self,

                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')

        return self.__db

    class DbCursor(Proxy):

        def __init__(self, outer):

            Proxy.__init__(self, outer.db.cursor())

            self.outer = outer

            self.query_lookup = None

            self.result = []

        def execute(self, query, params=None):

            self.query_lookup = _query_lookup(query, params)

            return_value = self.inner.execute(query, params)

            self.query = get_cur_query(self.inner)

            return return_value

        def fetchone(self):

            row = self.inner.fetchone()

            if row == None: # fetched all rows

                assert self.query_lookup != None

                pass #self.outer.query_results[self.query_lookup] = (self.query,

                    #self.result)

            else: self.result.append(row)

            return row

    class CacheCursor:

        def __init__(self, query, result):

            self.query = query

            self.rowcount = len(result)

            self.iter = iter(result)

        def fetchone(self):

            try: return self.iter.next()

            except StopIteration: return None

    def run_query(self, query, params=None, cacheable=False):

        query_lookup = _query_lookup(query, params)

        try: actual_query, result = self.query_results[query_lookup]

        except KeyError:

            cur = self.DbCursor(self)

            try: cur.execute(query, params)

            except Exception, e:

                _add_cursor_info(e, cur)

                raise

            if self.debug:

                sys.stderr.write(strings.one_line(get_cur_query(cur))+'\n')

            return cur

        else: return self.CacheCursor(actual_query, result)

connect = DbConn

##### Querying

def run_raw_query(db, *args, **kw_args):

    '''For args, see DbConn.run_query()'''

    return db.run_query(*args, **kw_args)

##### Recoverable querying

def with_savepoint(db, func):

    savepoint = 'savepoint_'+str(rand.rand_int()) # must be unique

    run_raw_query(db, 'SAVEPOINT '+savepoint)

    try: return_val = func()

    except:

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

        raise

    else:

        run_raw_query(db, 'RELEASE SAVEPOINT '+savepoint)

        return return_val

def run_query(db, query, params=None, recover=None, cacheable=False):

    if recover == None: recover = False

    def run(): return run_raw_query(db, query, params, cacheable)

    if recover: return with_savepoint(db, run)

    else: return run()

##### Result retrieval

def col_names(cur): return (col[0] for col in cur.description)

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

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

def row(cur):

    row_iter = rows(cur)

    row_ = row_iter.next()

    iters.consume_iter(row_iter) # fetch all rows so result will be cached

    return row_

def next_value(cur): return next_row(cur)[0]

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

def values(cur): return iters.func_iter(lambda: next_value(cur))

def value_or_none(cur):

    try: return value(cur)

    except StopIteration: return None

##### Basic queries

def select(db, table, fields=None, conds=None, limit=None, start=None,

    recover=None, cacheable=True):

    '''@param fields Use None to select all fields in the table'''

    if conds == None: conds = {}

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

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

    check_name(table)

    if fields != None: map(check_name, fields)

    map(check_name, conds.keys())

    def cond(entry):

        col, value = entry

        cond_ = esc_name(db, col)+' '

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

        else: cond_ += '='

        cond_ += ' %s'

        return cond_

    query = 'SELECT '

    if fields == None: query += '*'

    else: query += ', '.join([esc_name(db, field) for field in fields])

    query += ' FROM '+esc_name(db, table)

    missing = True

    if conds != {}:

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

        missing = False

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

    if start != None:

        if start != 0: query += ' OFFSET '+str(start)

        missing = False

    if missing: warnings.warn(DbWarning(

        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))

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

def insert(db, table, row, returning=None, recover=None):

    '''@param returning str|None An inserted column (such as pkey) to return'''

    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'

    if returning != None:

        check_name(returning)

        query += ' RETURNING '+returning

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

def last_insert_id(db):

    module = util.root_module(db.db)

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

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

    else: return None

def truncate(db, table):

    check_name(table)

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

##### Database structure queries

def pkey(db, table, recover=None):

    '''Assumed to be first column in table'''

    check_name(table)

    if table not in db.pkeys:

        db.pkeys[table] = col_names(run_query(db,

            'SELECT * FROM '+table+' LIMIT 0', recover=recover)).next()

    return db.pkeys[table]

def index_cols(db, table, index):

    '''Can also use this for UNIQUE constraints, because a UNIQUE index is

    automatically created. When you don't know whether something is a UNIQUE

    constraint or a UNIQUE index, use this function.'''

    check_name(table)

    check_name(index)

    lookup = (table, index)

    if lookup not in db.index_cols:

        module = util.root_module(db.db)

        if module == 'psycopg2':

            db.index_cols[lookup] = list(values(run_query(db, '''\

SELECT attname

FROM

(

        SELECT attnum, attname

        FROM pg_index

        JOIN pg_class index ON index.oid = indexrelid

        JOIN pg_class table_ ON table_.oid = indrelid

        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)

        WHERE

            table_.relname = %(table)s

            AND index.relname = %(index)s

    UNION

        SELECT attnum, attname

        FROM

        (

            SELECT

                indrelid

                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int

                    AS indkey

            FROM pg_index

            JOIN pg_class index ON index.oid = indexrelid

            JOIN pg_class table_ ON table_.oid = indrelid

            WHERE

                table_.relname = %(table)s

                AND index.relname = %(index)s

        ) s

        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey

) s

ORDER BY attnum

''',

                {'table': table, 'index': index})))

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

            ' database')

    return db.index_cols[lookup]

def constraint_cols(db, table, constraint):

    check_name(table)

    check_name(constraint)

    module = util.root_module(db.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 tables(db):

    module = util.root_module(db.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')

##### Database management

def empty_db(db):

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

##### Heuristic queries

def try_insert(db, table, row, returning=None):

    '''Recovers from errors'''

    try: return insert(db, table, row, returning, recover=True)

    except Exception, e:

        msg = str(e)

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

            r'"(([^\W_]+)_[^"]+)"', msg)

        if match:

            constraint, table = match.groups()

            try: cols = index_cols(db, table, constraint)

            except NotImplementedError: raise e

            else: raise DuplicateKeyException(cols, 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 put(db, table, row, pkey, row_ct_ref=None):

    '''Recovers from errors.

    Only works under PostgreSQL (uses `INSERT ... RETURNING`)'''

    try:

        cur = try_insert(db, table, row, pkey)

        if row_ct_ref != None and cur.rowcount >= 0:

            row_ct_ref[0] += cur.rowcount

        return value(cur)

    except DuplicateKeyException, e:

        return value(select(db, table, [pkey],

            util.dict_subset_right_join(row, e.cols), recover=True))

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

    '''Recovers from errors'''

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

    except StopIteration:

        if not create: raise

        return put(db, table, row, pkey, row_ct_ref) # insert new row