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

    else: return '[input] '+strings.ustr(input_query)

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

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

        DbException.__init__(self, 'for name: '

        self.name = name

class ExceptionWithValue(DbException):

        DbException.__init__(self, 'for value: '

        self.value = value

class ExceptionWithNameType(DbException):

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

        self.type = type_

class ConstraintException(DbException):

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

        self.name = name

        self.cond = cond

        self.cols = cols

class MissingCastException(DbException):

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

        self.type = type_

class NameException(DbException): pass

class DuplicateKeyException(ConstraintException): pass

class NullValueException(ConstraintException): pass

class CheckException(ConstraintException): pass

class InvalidValueException(ExceptionWithValue): pass

class DuplicateException(ExceptionWithNameType): pass

class DoesNotExistException(ExceptionWithNameType): pass

class EmptyRowException(DbException): pass

##### Warnings

##### Result retrieval

##### Escaping

def esc_name_by_module(module, name):

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

    return sql_gen.esc_name(name, quote)

##### Database connections

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

db_engines = {

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 src In autocommit mode, will be included in a comment in every

        '''

        self.db_config = db_config

        self.autocommit = autocommit

        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._savepoint = 0

        self._reset()

        state['log_debug'] = None # don't pickle the debug callback

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

    def clear_cache(self): self.query_results = {}

    def _reset(self):

        self.clear_cache()

        assert self._savepoint == 0

        self._notices_seen = set()

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

    def _db(self):

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

            module_name, mappings = db_engines[db_config.pop('engine')]

            # Record that a transaction is already open

            # Configure connection

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

            if schemas != None:

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

    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:

            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,

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

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

            else: self._cache_result()

            # For inserts that return a result set, don't cache result set since

            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(

        class CacheCursor:

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

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

    def esc_value(self, value):

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

        return strings.to_unicode(str_)

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

    def std_code(self, str_):

    def can_mogrify(self):

        module = util.root_module(self.db)

        return module == 'psycopg2'

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

    def print_notices(self):

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

    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

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

        '''

        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:

        try:

            # Get cursor

                try: cur = self.query_results[query]

                except KeyError: cur = self.DbCursor(self)

                else: used_cache = True

            else: cur = self.db.cursor()

            try: cur.execute(query)

            except Exception, e:

            else: self.do_autocommit()

        finally:

            if self.debug:

                profiler.stop(self.profile_row_ct)

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

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

                msg = 'DB query: '

                if used_cache: msg += 'cache hit'

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

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

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

    def with_autocommit(self, func):

        import psycopg2.extensions

        self.db.set_isolation_level(

        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

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

        except:

            if top: query = 'ROLLBACK'

            raise

        finally:

            if not top:

            if self.debug:

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

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

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

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

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

        schema = col.table.schema

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

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

        except StopIteration: raise sql_gen.NoUnderlyingTableException(col)

        default = sql_gen.as_Code(default, self)

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

        elif type_ == 'ARRAY':

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,

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

    if recover == None: recover = False

    if log_ignore_excs == None: 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)

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

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

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

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

            if match:

                constraint, = match.groups()

                cols = []

                    try: cols = index_cols(db, constraint)

                    except NotImplementedError: pass

                raise DuplicateKeyException(constraint, None, cols, e)

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

                r' constraint', msg)

            if match:

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

                constraint = sql_gen.Col(constraint, table)

                cond = None

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

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

            if match:

                value, = match.groups()

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

                r'is of type', msg)

            match = re.match(r'^could not determine polymorphic type because '

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

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

            if match:

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

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

            raise # no specific exception raised

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

##### Basic queries

def is_explainable(query):

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

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

    kw_args.setdefault('log_level', 4)

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

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

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