# Database access

import copy

import operator

import re

import warnings

import exc

import dicts

import iters

import lists

from Proxy import Proxy

import rand

import sql_gen

import strings

import util

##### Exceptions

def get_cur_query(cur, input_query=None, input_params=None):

    raw_query = None

    if hasattr(cur, 'query'): raw_query = cur.query

    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed

    if raw_query != None: return raw_query

    else: return repr(input_query)+' % '+repr(input_params)

def _add_cursor_info(e, *args, **kw_args):

    '''For params, see get_cur_query()'''

    exc.add_msg(e, 'query: '+str(get_cur_query(*args, **kw_args)))

class DbException(exc.ExceptionWithCause):

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

        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)

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

class ExceptionWithName(DbException):

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

        DbException.__init__(self, 'for name: '+str(name), cause)

        self.name = name

class ExceptionWithNameValue(DbException):

    def __init__(self, name, value, cause=None):

        DbException.__init__(self,

            'for name: '+str(name)+'; value: '+strings.ustr(value), cause)

        self.name = name

        self.value = value

class ConstraintException(DbException):

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

        DbException.__init__(self, 'Violated constraint '+name+ ' on columns: '

            +(', '.join(cols)), cause)

        self.name = name

        self.cols = cols

class NameException(DbException): pass

class DuplicateKeyException(ConstraintException): pass

class NullValueException(ConstraintException): pass

class FunctionValueException(ExceptionWithNameValue): pass

class DuplicateTableException(ExceptionWithName): pass

class DuplicateFunctionException(ExceptionWithName): pass

class EmptyRowException(DbException): pass

##### Warnings

class DbWarning(UserWarning): pass

##### Result retrieval

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

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

def consume_rows(cur):

    '''Used to fetch all rows so result will be cached'''

    iters.consume_iter(rows(cur))

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

def row(cur):

    row_ = next_row(cur)

    consume_rows(cur)

    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

##### Input validation

def clean_name(name): return re.sub(r'\W', r'', name).lower()

def check_name(name):

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

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

def esc_name_by_module(module, name, ignore_case=False):

    if module == 'psycopg2':

        if ignore_case:

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

            check_name(name)

            return name

        else: quote = '"'

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

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

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

def esc_name_by_engine(engine, name, **kw_args):

    return esc_name_by_module(db_engines[engine][0], name, **kw_args)

def esc_name(db, name, **kw_args):

    return esc_name_by_module(util.root_module(db.db), name, **kw_args)

def qual_name(db, schema, table):

    def esc_name_(name): return esc_name(db, name)

    table = esc_name_(table)

    if schema != None: return esc_name_(schema)+'.'+table

    else: return table

##### Database connections

db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']

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, dicts.make_hashable(params))

log_debug_none = lambda msg: None

class DbConn:

    def __init__(self, db_config, serializable=True, autocommit=False,

        caching=True, log_debug=log_debug_none):

        self.db_config = db_config

        self.serializable = serializable

        self.autocommit = autocommit

        self.caching = caching

        self.log_debug = log_debug

        self.debug = log_debug != log_debug_none

        self.__db = None

        self.query_results = {}

        self._savepoint = 0

    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['log_debug'] = None # don't pickle the debug callback

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

        return state

    def connected(self): return self.__db != None

    def _db(self):

        if self.__db == None:

            # Process db_config

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

            schemas = db_config.pop('schemas', None)

            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 and not self.autocommit: run_raw_query(self,

                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')

            if schemas != None:

                schemas_ = ''.join((esc_name(self, s)+', '

                    for s in schemas.split(',')))

                run_raw_query(self, "SELECT set_config('search_path', \

%s || current_setting('search_path'), false)", [schemas_])

        return self.__db

    class DbCursor(Proxy):

        def __init__(self, outer):

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

            self.outer = outer

            self.query_results = outer.query_results

            self.query_lookup = None

            self.result = []

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

            self._is_insert = query.upper().find('INSERT') >= 0

            self.query_lookup = _query_lookup(query, params)

            try:

                try:

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

                    self.outer.do_autocommit()

                finally: self.query = get_cur_query(self.inner)

            except Exception, e:

                _add_cursor_info(e, self, query, params)

                self.result = e # cache the exception as the result

                self._cache_result()

                raise

            # Fetch all rows so result will be cached

            if self.rowcount == 0 and not self._is_insert: consume_rows(self)

            return return_value

        def fetchone(self):

            row = self.inner.fetchone()

            if row != None: self.result.append(row)

            # otherwise, fetched all rows

            else: self._cache_result()

            return row

        def _cache_result(self):

            # For inserts, only cache exceptions since inserts are not

            # idempotent, but an invalid insert will always be invalid

            if self.query_results != None and (not self._is_insert

                or isinstance(self.result, Exception)):

                assert self.query_lookup != None

                self.query_results[self.query_lookup] = self.CacheCursor(

                    util.dict_subset(dicts.AttrsDictView(self),

                    ['query', 'result', 'rowcount', 'description']))

        class CacheCursor:

            def __init__(self, cached_result): self.__dict__ = cached_result

            def execute(self, *args, **kw_args):

                if isinstance(self.result, Exception): raise self.result

                # otherwise, result is a rows list

                self.iter = iter(self.result)

            def fetchone(self):

                try: return self.iter.next()

                except StopIteration: return None

    def esc_value(self, value):

        module = util.root_module(self.db)

        if module == 'psycopg2': return self.db.cursor().mogrify('%s', [value])

        elif module == 'MySQLdb':

            import _mysql

            return _mysql.escape_string(value)

        else: raise NotImplementedError("Can't escape value for "+module

            +' database')

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

        '''Translates known DB errors to typed exceptions:

        See self.DbCursor.execute().'''

        assert query != None

        if not self.caching: cacheable = False

        used_cache = False

        try:

            # Get cursor

            if cacheable:

                query_lookup = _query_lookup(query, params)

                try:

                    cur = self.query_results[query_lookup]

                    used_cache = True

                except KeyError: cur = self.DbCursor(self)

            else: cur = self.db.cursor()

            # Run query

            cur.execute(query, params)

        finally:

            if self.debug: # only compute msg if needed

                if used_cache: cache_status = 'Cache hit'

                elif cacheable: cache_status = 'Cache miss'

                else: cache_status = 'Non-cacheable'

                self.log_debug(cache_status+': '

                    +strings.one_line(str(get_cur_query(cur, query, params))))

        return cur

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

        return _query_lookup(query, params) in self.query_results

    def with_savepoint(self, func):

        savepoint = 'level_'+str(self._savepoint)

        self.run_query('SAVEPOINT '+savepoint)

        self._savepoint += 1

        try: 

            try: return_val = func()

            finally:

                self._savepoint -= 1

                assert self._savepoint >= 0

        except:

            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint)

            raise

        else:

            self.run_query('RELEASE SAVEPOINT '+savepoint)

            self.do_autocommit()

            return return_val

    def do_autocommit(self):

        '''Autocommits if outside savepoint'''

        assert self._savepoint >= 0

        if self.autocommit and self._savepoint == 0:

            self.log_debug('Autocommiting')

            self.db.commit()

connect = DbConn

##### Querying

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

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

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

def mogrify(db, query, params):

    module = util.root_module(db.db)

    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)

    else: raise NotImplementedError("Can't mogrify query for "+module+

        ' database')

##### Recoverable querying

def with_savepoint(db, func): return db.with_savepoint(func)

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

    if recover == None: recover = False

    try:

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

        if recover and not db.is_cached(query, params):

            return with_savepoint(db, run)

        else: return run() # don't need savepoint if cached

    except Exception, e:

        if not recover: raise # need savepoint to run index_cols()

        msg = exc.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(constraint, cols, e)

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

            r'constraint', msg)

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

        match = re.search(r'invalid input (?:syntax|value)\b.*: "(.+)"\n'

            r'(?:(?s).*)\bfunction "(\w+)".*\bat assignment', msg)

        if match:

            raise FunctionValueException(match.group(2),

                strings.to_unicode(match.group(1)), e)

        match = re.search(r'relation "(\w+)" already exists', msg)

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

        match = re.search(r'function "(\w+)" already exists', msg)

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

        raise # no specific exception raised

##### Basic queries

def next_version(name):

    '''Prepends the version # so it won't be removed if the name is truncated'''

    version = 1 # first existing name was version 0

    match = re.match(r'^v(\d+)_(.*)$', name)

    if match:

        version = int(match.group(1))+1

        name = match.group(2)

    return 'v'+str(version)+'_'+name

def run_query_into(db, query, params, into_ref=None, *args, **kw_args):

    '''Outputs a query to a temp table.

    For params, see run_query().

    '''

    if into_ref == None: return run_query(db, query, params, *args, **kw_args)

    else: # place rows in temp table

        check_name(into_ref[0])

        kw_args['recover'] = True

        while True:

            try:

                create_query = 'CREATE'

                if not db.debug: create_query += ' TEMP'

                create_query += ' TABLE '+into_ref[0]+' AS '+query

                return run_query(db, create_query, params, *args, **kw_args)

                    # CREATE TABLE AS sets rowcount to # rows in query

            except DuplicateTableException, e:

                into_ref[0] = next_version(into_ref[0])

                # try again with next version of name

order_by_pkey = object() # tells mk_select() to order by the pkey

distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns

def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,

    start=None, order_by=order_by_pkey, default_table=None):

    '''

    @param tables The single table to select from, or a list of tables to join

        together, with tables after the first being sql_gen.Join objects

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

    @param conds WHERE conditions: dict(compare_left_side=compare_right_side):

        * compare_left_side: Code|str (for col name)

        * compare_right_side: ValueCond|literal value

    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to

        use all columns

    @return tuple(query, params)

    '''

    # Parse tables param

    if not lists.is_seq(tables): tables = [tables]

    tables = list(tables) # don't modify input! (list() copies input)

    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate

    # Parse other params

    if conds == None: conds = {}

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

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

    if order_by is order_by_pkey:

        if distinct_on != []: order_by = None

        else: order_by = pkey(db, table0, recover=True)

    query = 'SELECT'

    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)

    # DISTINCT ON columns

    if distinct_on != []:

        query += ' DISTINCT'

        if distinct_on is not distinct_on_all:

            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'

    # Columns

    query += ' '

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

    else: query += ', '.join(map(parse_col, fields))

    # Main table

    query += ' FROM '+table0.to_str(db)

    # Add joins

    left_table = table0

    for join_ in tables:

        # Parse special values

        table = join_.table

        if join_.type_ is sql_gen.filter_out:

            conds[sql_gen.Col(pkey(db, table), table)] = None # filter no match

        query += ' '+join_.to_str(db, left_table)

        left_table = table

    missing = True

    if conds != {}:

        query += ' WHERE '+(' AND '.join((sql_gen.as_ValueCond(r).to_str(db, l)

            for l, r in conds.iteritems())))

        missing = False

    if order_by != None:

        query += ' ORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)

    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 (query, [])

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

    '''For params, see mk_select() and run_query()'''

    recover = kw_args.pop('recover', None)

    cacheable = kw_args.pop('cacheable', True)

    query, params = mk_select(db, *args, **kw_args)

    return run_query(db, query, params, recover, cacheable)

def mk_insert_select(db, table, cols=None, select_query=None, params=None,

    returning=None, embeddable=False):

    '''

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

    @param embeddable Whether the query should be embeddable as a nested SELECT.

        Warning: If you set this and cacheable=True when the query is run, the

        query will be fully cached, not just if it raises an exception.

    '''

    if select_query == None: select_query = 'DEFAULT VALUES'

    if cols == []: cols = None # no cols (all defaults) = unknown col names

    table = sql_gen.as_Table(table).to_str(db)

    # Build query

    query = 'INSERT INTO '+table

    if cols != None:

        map(check_name, cols)

        query += ' ('+', '.join(cols)+')'

    query += ' '+select_query

    if returning != None:

        check_name(returning)

        query += ' RETURNING '+returning

    if embeddable:

        # Create function

        function_name = '_'.join(map(clean_name, ['insert', table] + cols))

        return_type = 'SETOF '+table+'.'+returning+'%TYPE'

        while True:

            try:

                function = function_name

                if not db.debug: function = 'pg_temp.'+function

                function_query = '''\

CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''

    LANGUAGE sql

    AS $$'''+mogrify(db, query, params)+''';$$;

'''

                run_query(db, function_query, recover=True, cacheable=True)

                break # this version was successful

            except DuplicateFunctionException, e:

                function_name = next_version(function_name)

                # try again with next version of name

        # Return query that uses function

        return mk_select(db, sql_gen.CustomCode(function+'() AS f ('+returning

            +')'), start=0, order_by=None) # AS clause requires function alias

    return (query, params)

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

    '''For params, see mk_insert_select() and run_query_into()

    @param into_ref List with name of temp table to place RETURNING values in

    '''

    into_ref = kw_args.pop('into_ref', None)

    if into_ref != None: kw_args['embeddable'] = True

    recover = kw_args.pop('recover', None)

    cacheable = kw_args.pop('cacheable', True)

    query, params = mk_insert_select(db, *args, **kw_args)

    return run_query_into(db, query, params, into_ref, recover=recover,

        cacheable=cacheable)

default = object() # tells insert() to use the default value for a column

def insert(db, table, row, *args, **kw_args):

    '''For params, see insert_select()'''

    if lists.is_seq(row): cols = None

    else:

        cols = row.keys()

        row = row.values()

    row = list(row) # ensure that "!= []" works

    # Check for special values

    labels = []

    values = []

    for value in row:

        if value is default: labels.append('DEFAULT')

        else:

            labels.append('%s')

            values.append(value)

    # Build query

    if values != []: query = ' VALUES ('+(', '.join(labels))+')'

    else: query = None

    return insert_select(db, table, cols, query, values, *args, **kw_args)

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, schema='public'):

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

##### Database structure queries

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

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

    return col_names(select(db, table, limit=0, order_by=None,

        recover=recover)).next()

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)

    module = util.root_module(db.db)

    if module == 'psycopg2':

        return 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}, cacheable=True)))

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

        ' database')

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

row_num_col = '_row_num'

def add_row_num(db, table):

    '''Adds a row number column to a table. Its name is in row_num_col. It will

    be the primary key.'''

    check_name(table)

    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col

        +' serial NOT NULL PRIMARY KEY')

def tables(db, schema='public', table_like='%'):

    module = util.root_module(db.db)

    params = {'schema': schema, 'table_like': table_like}

    if module == 'psycopg2':

        return values(run_query(db, '''\

SELECT tablename

FROM pg_tables

WHERE

    schemaname = %(schema)s

    AND tablename LIKE %(table_like)s

ORDER BY tablename

''',

            params, cacheable=True))

    elif module == 'MySQLdb':

        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,

            cacheable=True))

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

##### Database management

def empty_db(db, schema='public', **kw_args):

    '''For kw_args, see tables()'''

    for table in tables(db, schema, **kw_args): truncate(db, table, schema)

##### Heuristic queries

def put(db, table, row, pkey_=None, row_ct_ref=None):

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    '''

    if pkey_ == None: pkey_ = pkey(db, table, recover=True)

    try:

        cur = insert(db, table, row, pkey_, recover=True)

        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, limit=1, recover=True))

    except StopIteration:

        if not create: raise

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

def put_table(db, out_table, in_tables, mapping, limit=None, start=0,

    row_ct_ref=None):

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    @param in_tables The main input table to select from, followed by a list of

        tables to join with it using the main input table's pkey

    @param mapping dict(out_table_col=in_table_col, ...)

        * out_table_col: sql_gen.Col|str

        * in_table_col: sql_gen.Col Wrap literal values in a sql_gen.NamedCode

    @return sql_gen.Col Where the output pkeys are made available

    '''

    for in_table_col in mapping.itervalues():

        assert isinstance(in_table_col, sql_gen.Col)

    temp_prefix = clean_name(out_table)

    pkeys_ref = [temp_prefix+'_pkeys']

    # Join together input tables

    in_tables = in_tables[:] # don't modify input!

    in_tables0 = in_tables.pop(0) # first table is separate

    in_pkey = pkey(db, in_tables0, recover=True)

    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)

    insert_joins = [in_tables0]+[sql_gen.Join(v, {in_pkey: sql_gen.join_using})

        for v in in_tables]

    out_pkey = pkey(db, out_table, recover=True)

    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)

    pkeys = [in_pkey, out_pkey]

    pkeys_cols = [in_pkey_col, out_pkey_col]

    pkeys_table_exists_ref = [False]

    def run_query_into_pkeys(query, params):

        if pkeys_table_exists_ref[0]:

            insert_select(db, pkeys_ref[0], pkeys, query, params)

        else:

            run_query_into(db, query, params, into_ref=pkeys_ref)

            pkeys_table_exists_ref[0] = True

    conds = {}

    distinct_on = []

    def mk_main_select(cols):

        return mk_select(db, insert_joins, cols, conds, distinct_on,