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, input_params=None):

    raw_query = None

def _add_cursor_info(e, *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: '+str(name), cause)

        self.name = name

class ExceptionWithNameValue(DbException):

            'for name: '+str(name)+'; value: '+strings.ustr(value), cause)

        self.name = name

class ConstraintException(DbException):

        DbException.__init__(self, 'Violated constraint '+name+ ' on columns: '

        self.name = name

        self.cols = cols

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

##### Input validation

def clean_name(name): return re.sub(r'\W', r'', name).lower()

##### Database connections

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

db_engines = {

def _query_lookup(query, params): return (query, dicts.make_hashable(params))

log_debug_none = lambda msg: 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

        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: run_raw_query(self,

            if schemas != None:

    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, params=None):

            self._is_insert = query.upper().find('INSERT') >= 0

            self.query_lookup = _query_lookup(query, params)

            try:

                try:

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

            except Exception, e:

                _add_cursor_info(e, self, query, params)

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

            # Fetch all rows so result will be cached

            return return_value

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

        module = util.root_module(self.db)

        elif module == 'MySQLdb':

    def run_query(self, query, params=None, cacheable=False):

        '''Translates known DB errors to typed exceptions:

        assert query != None

        if not self.caching: cacheable = False

        used_cache = False

            # Get cursor

            cur.execute(query, params)

        finally:

            if self.debug: # only compute msg if needed

                if used_cache: cache_status = 'Cache hit'

                self.log_debug(cache_status+': '

                    +strings.one_line(str(get_cur_query(cur, query, params))))

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

        self.run_query('SAVEPOINT '+savepoint)

            self.do_autocommit()

            return return_val

connect = DbConn

##### Querying

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

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

    return db.run_query(*args, **kw_args)

def mogrify(db, query, params):

##### Recoverable querying

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

def run_query(db, query, params=None, recover=None, cacheable=False):

    if recover == None: recover = False

    try:

        msg = exc.str_(e)

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

            else: raise DuplicateKeyException(constraint, cols, e)

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

            r'constraint', msg)

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

        match = re.search(r'invalid input (?:syntax|value)\b.*: "(.+)"\n'

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

        if match:

            raise FunctionValueException(match.group(2),

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

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

        raise # no specific exception raised

##### Basic queries

def next_version(name):

    version = 1 # first existing name was version 0

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

def run_query_into(db, query, params, into_ref=None, *args, **kw_args):

    '''Outputs a query to a temp table.

    if into_ref == None: return run_query(db, query, params, *args, **kw_args)

    else: # place rows in temp table

        kw_args['recover'] = True

                create_query = 'CREATE'

                into = sql_gen.as_Table(into_ref[0]).to_str(db)

                    # CREATE TABLE AS sets rowcount to # rows in query

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: dict(compare_left_side=compare_right_side):

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

    @return tuple(query, params)

    '''

    # 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 = {}

    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 += ' DISTINCT'

        if distinct_on is not distinct_on_all:

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

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

    else: query += ', '.join(map(parse_col, fields))

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

    # Add joins

    left_table = table0

    for join_ in tables:

        # Parse special values

        table = join_.table

        if join_.type_ is sql_gen.filter_out:

            conds[sql_gen.Col(pkey(db, table), table)] = None # filter no match

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

    missing = True

    if conds != {}:

        query += ' WHERE '+(' AND '.join((sql_gen.as_ValueCond(r).to_str(db, l)

        missing = False

    if order_by != None:

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

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

    return (query, [])

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

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

    returning=None, embeddable=False):

    '''

    @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.as_Table(table).to_str(db)

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

    if cols != None: query += ' ('+', '.join(cols)+')'

    query += ' '+select_query

        query += ' RETURNING '+sql_gen.as_Col(returning).to_str(db)

    if embeddable:

        function_name = '_'.join(map(clean_name, ['insert', table] + cols))

        return_type = 'SETOF '+table+'.'+returning+'%TYPE'

        while True:

                function = function_name

                function_query = '''\

    LANGUAGE sql

                run_query(db, function_query, recover=True, cacheable=True)

        return mk_select(db, sql_gen.CustomCode(function+'() AS f ('+returning

    return (query, params)

    '''For params, see mk_insert_select() and run_query_into()

    @param into_ref List with name of temp table to place RETURNING values in

    '''

    into_ref = kw_args.pop('into_ref', None)

    recover = kw_args.pop('recover', None)

    return run_query_into(db, query, params, into_ref, recover=recover,

default = object() # tells insert() to use the default value for a column

def insert(db, table, row, *args, **kw_args):

    '''For params, see insert_select()'''

    if lists.is_seq(row): cols = None

    # Check for special values

        if value is default: labels.append('DEFAULT')

        else:

    if values != []: query = ' VALUES ('+(', '.join(labels))+')'

    return insert_select(db, table, cols, query, values, *args, **kw_args)

def last_insert_id(db):

    module = util.root_module(db.db)

    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))

def truncate(db, table, schema='public'):

def pkey(db, table, recover=None):

    '''Assumed to be first column in table'''

    return col_names(select(db, table, limit=0, order_by=None,

def index_cols(db, table, index):