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

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

        self.name = name

class ExceptionWithValue(DbException):

        self.value = value

class ExceptionWithNameType(DbException):

class ConstraintException(DbException):

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

        self.name = name

        self.cols = cols

class MissingCastException(DbException):

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

##### 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._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

        start_time = time.time()

        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:

            time_ = time.time() - start_time

            self.print_notices()

            if self.debug: # log or return query

                query = str(get_cur_query(cur, query))

                msg = 'DB query: '

                if used_cache: msg += 'cache hit'

                msg += ', took '+str(time_)+'s:\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)

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

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

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

        cols = [type_, 'column_default',

        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)

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]

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

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

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

            if match:

                cols = []

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

                r' constraint', msg)

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

                r'|.+? field value 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'^(\S+) "(.+?)".*? already exists', msg)

            if match:

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

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

def lock_table(db, table, mode):

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

    '''Outputs a query to a temp table.

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

    into.is_temp = True

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

    kw_args['recover'] = True

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

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

        try:

        except DuplicateException, e:

            into.name = next_version(into.name)

    analyze(db, into)

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)

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

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

    if order_by is order_by_pkey: