# Database access

import copy

import re

import time

import warnings

import exc

import dicts

import iters

import lists

import profiling

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(strings.ustr(name)), cause)

        self.name = name

class ExceptionWithValue(DbException):

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

        DbException.__init__(self, 'for value: '

            +strings.as_tt(strings.urepr(value)), cause)

        self.value = value

class ExceptionWithNameType(DbException):

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

        DbException.__init__(self, 'for type: '+strings.as_tt(strings.ustr(

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

        self.type = type_

        self.name = name

class ConstraintException(DbException):

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

        msg = 'Violated '+strings.as_tt(name)+' constraint'

        if cond != None: msg += ' with condition '+strings.as_tt(cond)

        if cols != []: msg += ' on columns: '+strings.as_tt(', '.join(cols))

        DbException.__init__(self, msg, cause)

        self.name = name

        self.cond = cond

        self.cols = cols

class MissingCastException(DbException):

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

        msg = 'Missing cast to type '+strings.as_tt(type_)

        if col != None: msg += ' on column: '+strings.as_tt(col)

        DbException.__init__(self, msg, cause)

        self.type = type_

        self.col = col

class EncodingException(ExceptionWithName): pass

class DuplicateKeyException(ConstraintException): pass

class NullValueException(ConstraintException): pass

class CheckException(ConstraintException): pass

class InvalidValueException(ExceptionWithValue): pass

class InvalidTypeException(ExceptionWithNameType): pass

class DuplicateException(ExceptionWithNameType): pass

class DoesNotExistException(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.autoexplain = False

        self.profile_row_ct = None

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

                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 set_encoding(self, encoding):

        encoding_str = sql_gen.Literal(encoding)

        run_query(self, 'SET NAMES '+encoding_str.to_str(self))

    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

        if self.debug:

            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')

        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:

            if self.debug:

                profiler.stop(self.profile_row_ct)

                ## Log or return query

                query = strings.ustr(get_cur_query(cur, query))

                # Put the src comment on a separate line in the log file

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

                msg = 'DB query: '

                if used_cache: msg += 'cache hit'

                elif cacheable: msg += 'cache miss'

                else: msg += 'non-cacheable'

                msg += ':\n'+profiler.msg()+'\n'+strings.as_code(query, 'SQL')

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

                else: self.log_debug(msg, log_level)

                self.print_notices()

        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)

        if self.debug:

            self.log_debug('Begin transaction', level=4)

            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')

        # 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

            if self.debug:

                profiler.stop(self.profile_row_ct)

                self.log_debug('End transaction\n'+profiler.msg(), level=4)

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

        module = util.root_module(self.db)

        if module == 'psycopg2':

            qual_table = sql_gen.Literal(col.table.to_str(self))

            col_name_str = sql_gen.Literal(col.name)

            try:

                type_, is_array, default, nullable = row(run_query(self, '''\

SELECT

format_type(COALESCE(NULLIF(typelem, 0), pg_type.oid), -1) AS type

, typcategory = 'A' AS type_is_array

, pg_get_expr(pg_attrdef.adbin, attrelid, true) AS default

, NOT pg_attribute.attnotnull AS nullable

FROM pg_attribute

LEFT JOIN pg_type ON pg_type.oid = atttypid

LEFT JOIN pg_attrdef ON adrelid = attrelid AND adnum = attnum

WHERE

attrelid = '''+qual_table.to_str(self)+'''::regclass

AND attname = '''+col_name_str.to_str(self)+'''

'''

                    , recover=True, cacheable=cacheable, log_level=4))

            except (DoesNotExistException, StopIteration):

                raise sql_gen.NoUnderlyingTableException(col)

            if is_array: type_ = sql_gen.ArrayType(type_)

        else:

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

            cols = [sql_gen.Col('data_type'), sql_gen.Col('udt_name'),

                'column_default', sql_gen.Cast('boolean',

                sql_gen.Col('is_nullable'))]

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

                ('column_name', strings.ustr(col.name))]

            schema = col.table.schema

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

            cur = select(self, table, cols, conds, order_by='table_schema',

                limit=1, cacheable=cacheable, log_level=4)

            try: type_, extra_type, default, nullable = row(cur)

            except StopIteration: raise sql_gen.NoUnderlyingTableException(col)

            if type_ == 'USER-DEFINED': type_ = extra_type

            elif type_ == 'ARRAY':

                type_ = sql_gen.ArrayType(strings.remove_prefix('_', extra_type,

                    require=True))

        if default != None: 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 parse_exception(db, e, recover=False):

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

    msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)

    match = re.match(r'^invalid byte sequence for encoding "(.+?)":', msg)

    if match:

        encoding, = match.groups()

        raise EncodingException(encoding, e)

    def make_DuplicateKeyException(constraint, e):

        cols = []

        cond = None

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

            try:

                cols = index_cols(db, constraint)

                cond = index_cond(db, constraint)

            except NotImplementedError: pass

            except DatabaseErrors, e2: # e.g. "current transaction is aborted"

                exc.print_ex(e2)

                raise e

        return DuplicateKeyException(constraint, cond, cols, e)

    match = re.match(r'^duplicate key value violates unique constraint "(.+?)"',

        msg)

    if match:

        constraint, = match.groups()

        raise make_DuplicateKeyException(constraint, e)

    match = re.match(r'^could not create unique index "(.+?)"\n'

        r'DETAIL:  Key .+? is duplicated', msg)

    if match:

        constraint, = match.groups()

        raise DuplicateKeyException(constraint, None, [], e)

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

        r' constraint', msg)

    if match:

        col, = match.groups()

        raise NullValueException('NOT NULL', None, [col], e)

    match = re.match(r'^new row for relation "(.+?)" violates check '

        r'constraint "(.+?)"', msg)

    if match:

        table, constraint = match.groups()

        constraint = sql_gen.Col(constraint, table)

        cond = None

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

            try: cond = constraint_cond(db, constraint)

            except NotImplementedError: pass

        raise CheckException(constraint.to_str(db), cond, [], e)

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

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

    if match:

        value, = match.groups()

        value = util.do_ignore_none(strings.to_unicode, value)

        raise InvalidValueException(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'^could not determine polymorphic type because '

        r'input has type "unknown"', msg)

    if match: raise MissingCastException('text', None, e)

    # type mismatches

    match = re.match(r'^.+? types (.+?) and (.+?) cannot be matched', msg)

    if match:

        type0, type1 = match.groups()

        raise MissingCastException(type0, None, e)

    match = re.match(r'operator does not exist: (.+?) = (.+)', msg)

        # these appear when a staging table column is of the wrong type

    if match:

        type0, type1 = match.groups()

        raise MissingCastException(type0, None, e)

    match = re.match(r'^.*?\brelation "(.+?)" is not a (table)', msg)

    if match:

        name, type_ = match.groups()

        raise InvalidTypeException(type_, name, e)

    typed_name_re = r'^(\S+) ("?)(.+?)\2(?: of relation ".+?")?'

        # regexp must be followed with text for .*? to match properly

    match = re.match(typed_name_re+r'(?:\(.*?)? already exists', msg)

    if match:

        type_, quote, name = match.groups()

        raise DuplicateException(type_, name, e)

    match = re.match(r'more than one (\S+) named ""(.+?)""', msg)

    if match:

        type_, name = match.groups()

        raise DuplicateException(type_, name, e)

    # type mismatches which appear as "does not exist" errors

    match = re.match(typed_name_re+r' does not exist', msg)

    if match:

        type_, quote, name = match.groups()

        if type_ == 'function':

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

            if match: # includes params, so is call rather than cast to regproc

                function_name, = match.groups()

                func = sql_gen.Function(function_name)

                if function_exists(db, func) and msg.find('CAST') < 0:

                    # not found only because of a missing cast

                    try: type_ = function_param0_type(db, func)

                    except DuplicateException: type_ = 'text'

                    col = None

                    if type_ == 'anyelement': type_ = 'text'

                    elif type_ == 'hstore': # cast just the value param

                        type_ = 'text'

                        col = 'value'

                    raise MissingCastException(type_, col, e)

        raise DoesNotExistException(type_, name, e)

    raise # no specific exception raised

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]

    query = with_explain_comment(db, query)

    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:

            # Give failed EXPLAIN approximately the log_level of its query

            if query.startswith('EXPLAIN'): log_level -= 1

            parse_exception(db, e, recover)

    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 is_explainable(query):

    # See <http://www.postgresql.org/docs/8.3/static/sql-explain.html#AEN57749>

    return re.match(r'^(?:SELECT|INSERT|UPDATE|DELETE|VALUES|EXECUTE|DECLARE)\b'

        , query)

def explain(db, query, **kw_args):

    '''

    For params, see run_query().

    '''

    kw_args.setdefault('log_level', 4)

    return strings.ustr(strings.join_lines(values(run_query(db,

        'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))

        # not a higher log_level because it's useful to see what query is being

        # run before it's executed, which EXPLAIN effectively provides

def has_comment(query): return query.endswith('*/')

def with_explain_comment(db, query, **kw_args):

    if db.autoexplain and not has_comment(query) and is_explainable(query):

        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'

            +explain(db, query, **kw_args))

    return query

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_pkey_=False, add_pkey_warn=True,

    **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_pkey_: add_pkey_or_index(db, into, warn=add_pkey_warn)

    # 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