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

    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 ExceptionWithNameValue(DbException):

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

        self.name = name

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 FunctionValueException(ExceptionWithNameValue): pass

class DuplicateTableException(ExceptionWithName): pass

class DuplicateFunctionException(ExceptionWithName): 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, serializable=True, autocommit=False,

        self.db_config = db_config

        self.autocommit = autocommit

        self.log_debug = log_debug

        self.debug = log_debug != log_debug_none

        self.query_results = {}

        self._savepoint = 0

        self._notices_seen = set()

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

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

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

    def _db(self):

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

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

            if self.serializable and not self.autocommit: self.run_query(

                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE', log_level=4)

            if schemas != None:

                self.run_query("SELECT set_config('search_path', "

    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)

                    self.outer.do_autocommit()

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

            except Exception, e:

                _add_cursor_info(e, self, query)

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

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

            return cur

        def fetchone(self):

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

            else: self._cache_result()

            # For inserts, only cache exceptions since inserts are not

            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 not self.caching: cacheable = False

        used_cache = False

        try:

            # Get cursor

                    cur = self.query_results[query]

                    used_cache = True

            # Log query

            # Run query

            cur.execute(query)

        finally:

            self.print_notices()

            if self.debug and debug_msg_ref != None: # return after running

                debug_msg_ref[0] = log_msg(str(get_cur_query(cur, query)))

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

    def with_autocommit(self, func, autocommit=True):

        import psycopg2.extensions

        try: return func()

        finally: self.db.set_isolation_level(prev_isolation_level)

    def with_savepoint(self, func):

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

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

        self._savepoint += 1

        try:

            try: return_val = func()

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

            raise

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

            self.do_autocommit()

            return return_val

            self.log_debug('Autocommitting')

            self.db.commit()

    def col_info(self, col):

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

        cols = ['data_type', 'column_default',

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

            order_by='table_schema', limit=1, 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 # usually, db.run_query() logs query before running it

    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)

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

            if match:

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

                r' constraint', msg)

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

            if match:

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

                r'is of type', msg)

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

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

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

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

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

##### 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.add_suffix(name, '#'+str(version))

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)

        try:

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

    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 = []

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

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

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

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

    if order_by is order_by_pkey:

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

    query += '\n'

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

    else:

    query += '\nFROM '+table0.to_str(db)

    # Add joins

    left_table = table0

    for join_ in tables:

        # Parse special values

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

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

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

    missing = True

    if conds != []:

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

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

        missing = False

    if order_by != None:

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

    if start != None:

        if start != 0: query += '\nOFFSET '+str(start)

        missing = False

    return query

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

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

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

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

    '''

    @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

    '''

    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:

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