# Database access

import copy

import re

import time

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

    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 '[input] '+strings.ustr(input_query)

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

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

    exc.add_msg(e, 'query: '+strings.ustr(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: '+strings.as_tt(str(name)), cause)

        self.name = name

class ExceptionWithValue(DbException):

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

        DbException.__init__(self, 'for value: '+strings.as_tt(repr(value)),

            cause)

        self.value = value

class ExceptionWithNameType(DbException):

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

        DbException.__init__(self, 'for type: '+strings.as_tt(str(type_))

            +'; name: '+strings.as_tt(name), cause)

        self.type = type_

        self.name = name

class ConstraintException(DbException):

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

        DbException.__init__(self, 'Violated '+strings.as_tt(name)

            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)

        self.name = name

        self.cols = cols

class MissingCastException(DbException):

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

        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)

            +' on column: '+strings.as_tt(col), cause)

        self.type = type_

        self.col = col

class NameException(DbException): pass

class DuplicateKeyException(ConstraintException): pass

class NullValueException(ConstraintException): pass

class InvalidValueException(ExceptionWithValue): pass

class DuplicateException(ExceptionWithNameType): 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

##### Escaping

def esc_name_by_module(module, name):

    if module == 'psycopg2' or module == None: quote = '"'

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

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

    return sql_gen.esc_name(name, 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']

log_debug_none = lambda msg, level=2: None

class DbConn:

    def __init__(self, db_config, autocommit=True, caching=True,

        log_debug=log_debug_none, debug_temp=False, src=None):

        '''

        @param debug_temp Whether temporary objects should instead be permanent.

            This assists in debugging the internal objects used by the program.

        @param src In autocommit mode, will be included in a comment in every

            query, to help identify the data source in pg_stat_activity.

        '''

        self.db_config = db_config

        self.autocommit = autocommit

        self.caching = caching

        self.log_debug = log_debug

        self.debug = log_debug != log_debug_none

        self.debug_temp = debug_temp

        self.src = src

        self.autoanalyze = False

        self._savepoint = 0

        self._reset()

    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 clear_cache(self): self.query_results = {}

    def _reset(self):

        self.clear_cache()

        assert self._savepoint == 0

        self._notices_seen = set()

        self.__db = None

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

    def close(self):

        if not self.connected(): return

        # Record that the automatic transaction is now closed

        self._savepoint -= 1

        self.db.close()

        self._reset()

    def reconnect(self):

        # Do not do this in test mode as it would roll back everything

        if self.autocommit: self.close()

        # Connection will be reopened automatically on first query

    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)

            # Record that a transaction is already open

            self._savepoint += 1

            # Configure connection

            if hasattr(self.db, 'set_isolation_level'):

                import psycopg2.extensions

                self.db.set_isolation_level(

                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)

            if schemas != None:

                search_path = [self.esc_name(s) for s in schemas.split(',')]

                search_path.append(value(run_query(self, 'SHOW search_path',

                    log_level=4)))

                run_query(self, 'SET search_path TO '+(','.join(search_path)),

                    log_level=3)

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

            self._is_insert = query.startswith('INSERT')

            self.query_lookup = query

            try:

                try: cur = self.inner.execute(query)

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

            except Exception, e:

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

                self._cache_result()

                raise

            # Always cache certain queries

            query = sql_gen.lstrip(query)

            if query.startswith('CREATE') or query.startswith('ALTER'):

                # structural changes

                # Rest of query must be unique in the face of name collisions,

                # so don't cache ADD COLUMN unless it has distinguishing comment

                if query.find('ADD COLUMN') < 0 or query.endswith('*/'):

                    self._cache_result()

            elif self.rowcount == 0 and query.startswith('SELECT'): # empty

                consume_rows(self) # fetch all rows so result will be cached

            return cur

        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 that return a result set, don't cache result set since

            # inserts are not idempotent. Other non-SELECT queries don't have

            # their result set read, so only exceptions will be cached (an

            # invalid query 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):

        try: str_ = self.mogrify('%s', [value])

        except NotImplementedError, e:

            module = util.root_module(self.db)

            if module == 'MySQLdb':

                import _mysql

                str_ = _mysql.escape_string(value)

            else: raise e

        return strings.to_unicode(str_)

    def esc_name(self, name): return esc_name(self, name) # calls global func

    def std_code(self, str_):

        '''Standardizes SQL code.

        * Ensures that string literals are prefixed by `E`

        '''

        if str_.startswith("'"): str_ = 'E'+str_

        return str_

    def can_mogrify(self):

        module = util.root_module(self.db)

        return module == 'psycopg2'

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

        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)

        else: raise NotImplementedError("Can't mogrify query")

    def print_notices(self):

        if hasattr(self.db, 'notices'):

            for msg in self.db.notices:

                if msg not in self._notices_seen:

                    self._notices_seen.add(msg)

                    self.log_debug(msg, level=2)

    def run_query(self, query, cacheable=False, log_level=2,

        debug_msg_ref=None):

        '''

        @param log_ignore_excs The log_level will be increased by 2 if the query

            throws one of these exceptions.

        @param debug_msg_ref If specified, the log message will be returned in

            this instead of being output. This allows you to filter log messages

            depending on the result of the query.

        '''

        assert query != None

        if self.autocommit and self.src != None:

            query = sql_gen.esc_comment(self.src)+'\t'+query

        if not self.caching: cacheable = False

        used_cache = False

        start_time = time.time()

        try:

            # Get cursor

            if cacheable:

                try: cur = self.query_results[query]

                except KeyError: cur = self.DbCursor(self)

                else: used_cache = True

            else: cur = self.db.cursor()

            # Run query

            try: cur.execute(query)

            except Exception, e:

                _add_cursor_info(e, self, query)

                raise

            else: self.do_autocommit()

        finally:

            time_ = time.time() - start_time

            self.print_notices()

            if self.debug: # log or return query

                query = str(get_cur_query(cur, query))

                # So that the src comment is hidden when cating the file, and

                # put on a separate line when viewed in a text editor.

                query = query.replace('\t', '\r', 1)

                msg = 'DB query: '

                if used_cache: msg += 'cache hit'

                elif cacheable: msg += 'cache miss'

                else: msg += 'non-cacheable'

                msg += ', took '+str(time_)+'s:\n'+strings.as_code(query, 'SQL')

                if debug_msg_ref != None: debug_msg_ref[0] = msg

                else: self.log_debug(msg, log_level)

        return cur

    def is_cached(self, query): return query in self.query_results

    def with_autocommit(self, func):

        import psycopg2.extensions

        prev_isolation_level = self.db.isolation_level

        self.db.set_isolation_level(

            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)

        try: return func()

        finally: self.db.set_isolation_level(prev_isolation_level)

    def with_savepoint(self, func):

        top = self._savepoint == 0

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

        # Must happen before running queries so they don't get autocommitted

        self._savepoint += 1

        if top: query = 'START TRANSACTION ISOLATION LEVEL READ COMMITTED'

        else: query = 'SAVEPOINT '+savepoint

        self.run_query(query, log_level=4)

        try:

            return func()

            if top: self.run_query('COMMIT', log_level=4)

        except:

            if top: query = 'ROLLBACK'

            else: query = 'ROLLBACK TO SAVEPOINT '+savepoint

            self.run_query(query, log_level=4)

            raise

        finally:

            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,

            # "The savepoint remains valid and can be rolled back to again"

            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).

            if not top:

                self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)

            self._savepoint -= 1

            assert self._savepoint >= 0

            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT

    def do_autocommit(self):

        '''Autocommits if outside savepoint'''

        assert self._savepoint >= 1

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

            self.log_debug('Autocommitting', level=4)

            self.db.commit()

    def col_info(self, col, cacheable=True):

        table = sql_gen.Table('columns', 'information_schema')

        type_ = sql_gen.Coalesce(sql_gen.Nullif(sql_gen.Col('data_type'),

            'USER-DEFINED'), sql_gen.Col('udt_name'))

        cols = [type_, 'column_default',

            sql_gen.Cast('boolean', sql_gen.Col('is_nullable'))]

        conds = [('table_name', col.table.name), ('column_name', col.name)]

        schema = col.table.schema

        if schema != None: conds.append(('table_schema', schema))

        type_, default, nullable = row(select(self, table, cols, conds,

            order_by='table_schema', limit=1, cacheable=cacheable, log_level=4))

            # TODO: order_by search_path schema order

        default = sql_gen.as_Code(default, self)

        return sql_gen.TypedCol(col.name, type_, default, nullable)

    def TempFunction(self, name):

        if self.debug_temp: schema = None

        else: schema = 'pg_temp'

        return sql_gen.Function(name, schema)

connect = DbConn

##### Recoverable querying

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

def run_query(db, query, recover=None, cacheable=False, log_level=2,

    log_ignore_excs=None, **kw_args):

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

    if recover == None: recover = False

    if log_ignore_excs == None: log_ignore_excs = ()

    log_ignore_excs = tuple(log_ignore_excs)

    debug_msg_ref = [None]

    try:

        try:

            def run(): return db.run_query(query, cacheable, log_level,

                debug_msg_ref, **kw_args)

            if recover and not db.is_cached(query):

                return with_savepoint(db, run)

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

        except Exception, e:

            msg = strings.ustr(e.args[0])

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

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

            if match:

                constraint, table = match.groups()

                cols = []

                if recover: # need auto-rollback to run index_cols()

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

                    except NotImplementedError: pass

                raise DuplicateKeyException(constraint, cols, e)

            match = re.match(r'^null value in column "(.+?)" violates not-null'

                r' constraint', msg)

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

            match = re.match(r'^(?:invalid input (?:syntax|value)\b.*?'

                r'|.+? field value out of range): "(.+?)"', msg)

            if match:

                value, = match.groups()

                raise InvalidValueException(strings.to_unicode(value), e)

            match = re.match(r'^column "(.+?)" is of type (.+?) but expression '

                r'is of type', msg)

            if match:

                col, type_ = match.groups()

                raise MissingCastException(type_, col, e)

            match = re.match(r'^(\S+) "(.+?)".*? already exists', msg)

            if match:

                type_, name = match.groups()

                raise DuplicateException(type_, name, e)

            raise # no specific exception raised

    except log_ignore_excs:

        log_level += 2

        raise

    finally:

        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)

##### Basic queries

def next_version(name):

    version = 1 # first existing name was version 0

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

    if match:

        name, version = match.groups()

        version = int(version)+1

    return sql_gen.concat(name, '#'+str(version))

def lock_table(db, table, mode):

    table = sql_gen.as_Table(table)

    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')

def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):

    '''Outputs a query to a temp table.

    For params, see run_query().

    '''

    if into == None: return run_query(db, query, **kw_args)

    assert isinstance(into, sql_gen.Table)

    into.is_temp = True

    # "temporary tables cannot specify a schema name", so remove schema

    into.schema = None

    kw_args['recover'] = True

    kw_args.setdefault('log_ignore_excs', (DuplicateException,))

    temp = not db.debug_temp # tables are permanent in debug_temp mode

    # Create table

    while True:

        create_query = 'CREATE'

        if temp: create_query += ' TEMP'

        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query

        try:

            cur = run_query(db, create_query, **kw_args)

                # CREATE TABLE AS sets rowcount to # rows in query

            break

        except DuplicateException, e:

            into.name = next_version(into.name)

            # try again with next version of name

    if add_indexes_: add_indexes(db, into)

    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and

    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary

    # table is going to be used in complex queries, it is wise to run ANALYZE on

    # the temporary table after it is populated."

    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)

    # If into is not a temp table, ANALYZE is useful but not required.

    analyze(db, into)

    return cur

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: [(compare_left_side, compare_right_side),...]

        * container can be any iterable type

        * compare_left_side: sql_gen.Code|str (for col name)

        * compare_right_side: sql_gen.ValueCond|literal value

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

        use all columns

    @return query

    '''

    # Parse tables param

    tables = lists.mk_seq(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 = []

    elif dicts.is_dict(conds): conds = conds.items()

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

    assert limit == None or isinstance(limit, (int, long))

    assert start == None or isinstance(start, (int, long))

    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 += '\nDISTINCT'

        if distinct_on is not distinct_on_all:

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

    # Columns

    if query.find('\n') >= 0: whitespace = '\n'

    else: whitespace = ' '

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

    else:

        assert fields != []

        if len(fields) > 1: whitespace = '\n'

        query += whitespace+('\n, '.join(map(parse_col, fields)))

    # Main table

    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'

    else: whitespace = ' '

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

    # Add joins

    left_table = table0

    for join_ in tables:

        table = join_.table

        # Parse special values

        if join_.type_ is sql_gen.filter_out: # filter no match

            conds.append((sql_gen.Col(table_not_null_col(db, table), table),

                sql_gen.CompareCond(None, '~=')))

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

        left_table = table

    missing = True

    if conds != []:

        if len(conds) == 1: whitespace = ' '

        else: whitespace = '\n'

        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)

            .to_str(db) for l, r in conds], 'WHERE')

        missing = False

    if order_by != None:

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

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

    if start != None:

        if start != 0: query += '\nOFFSET '+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)

    log_level = kw_args.pop('log_level', 2)

    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,

        log_level=log_level)

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

    embeddable=False, ignore=False, src=None):

    '''

    @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.

    @param ignore Whether to ignore duplicate keys.

    @param src Will be included in the name of any created function, to help

        identify the data source in pg_stat_activity.

    '''

    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))

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

    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]

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

    if returning != None: returning = sql_gen.as_Col(returning, table)

    first_line = 'INSERT INTO '+table.to_str(db)

    def mk_insert(select_query):

        query = first_line

        if cols != None:

            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'

        query += '\n'+select_query

        if returning != None:

            returning_name_col = sql_gen.to_name_only_col(returning)

            query += '\nRETURNING '+returning_name_col.to_str(db)

        return query

    return_type = 'unknown'

    if returning != None: return_type = returning.to_str(db)+'%TYPE'

    lang = 'sql'

    if ignore:

        # Always return something to set the correct rowcount

        if returning == None: returning = sql_gen.NamedCol('NULL', None)

        embeddable = True # must use function

        lang = 'plpgsql'

        if cols == None:

            row = [sql_gen.Col(sql_gen.all_cols, 'row')]

            row_vars = [sql_gen.Table('row')]

        else:

            row_vars = row = [sql_gen.Col(c.name, 'row') for c in cols]

        query = '''\

DECLARE

    row '''+table.to_str(db)+'''%ROWTYPE;

BEGIN

    /* Need an EXCEPTION block for each individual row because "When an error is

    caught by an EXCEPTION clause, [...] all changes to persistent database

    state within the block are rolled back."

    This is unfortunate because "A block containing an EXCEPTION clause is

    significantly more expensive to enter and exit than a block without one."

    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\

#PLPGSQL-ERROR-TRAPPING)

    */

    FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN

'''+select_query+'''

    LOOP

        BEGIN

            RETURN QUERY

'''+mk_insert(sql_gen.Values(row).to_str(db))+'''

;

        EXCEPTION

            WHEN unique_violation THEN NULL; -- continue to next row

        END;

    END LOOP;

END;\

'''

    else: query = mk_insert(select_query)

    if embeddable:

        # Create function

        function_name = sql_gen.clean_name(first_line)

        if src != None: function_name = src+': '+function_name

        while True:

            try:

                function = db.TempFunction(function_name)