# Database access

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

    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)+' % '+repr(input_params)

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

    '''For params, see get_cur_query()'''

    exc.add_msg(e, 'query: '+str(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: '+str(name), cause)

        self.name = name

class ExceptionWithNameValue(DbException):

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

        DbException.__init__(self,

            'for name: '+str(name)+'; value: '+repr(value), cause)

        self.name = name

        self.value = value

class ConstraintException(DbException):

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

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

            +(', '.join(cols)), cause)

        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

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

##### Input validation

def check_name(name):

    if re.search(r'\W', name) != None: raise NameException('Name "'+name

        +'" may contain only alphanumeric characters and _')

def esc_name_by_module(module, name, ignore_case=False):

    if module == 'psycopg2' or module == None:

        if ignore_case:

            # Don't enclose in quotes because this disables case-insensitivity

            check_name(name)

            return name

        else: quote = '"'

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

    else: raise NotImplementedError("Can't escape name for "+module+' database')

    return quote + name.replace(quote, '') + 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']

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,

        caching=True, log_debug=log_debug_none):

        self.db_config = db_config

        self.serializable = serializable

        self.autocommit = autocommit

        self.caching = caching

        self.log_debug = log_debug

        self.debug = log_debug != log_debug_none

        self.__db = None

        self.query_results = {}

        self._savepoint = 0

    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 connected(self): return self.__db != None

    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)

            # Configure connection

            if self.serializable and not self.autocommit: run_raw_query(self,

                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')

            if schemas != None:

                schemas_ = ''.join((esc_name(self, s)+', '

                    for s in schemas.split(',')))

                run_raw_query(self, "SELECT set_config('search_path', \

%s || current_setting('search_path'), false)", [schemas_])

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

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

            self.query_lookup = _query_lookup(query, params)

            try:

                try:

                    return_value = self.inner.execute(query, params)

                    self.outer.do_autocommit()

                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

                self._cache_result()

                raise

            # Fetch all rows so result will be cached

            if self.rowcount == 0 and not self._is_insert: consume_rows(self)

            return return_value

        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, only cache exceptions since inserts are not

            # idempotent, but an invalid insert 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):

        module = util.root_module(self.db)

        if module == 'psycopg2': str_ = self.db.cursor().mogrify('%s', [value])

        elif module == 'MySQLdb':

            import _mysql

            str_ = _mysql.escape_string(value)

        else: raise NotImplementedError("Can't escape value for "+module

            +' database')

        return strings.to_unicode(str_)

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

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

        '''Translates known DB errors to typed exceptions:

        See self.DbCursor.execute().'''

        assert query != None

        if not self.caching: cacheable = False

        used_cache = False

        try:

            # Get cursor

            if cacheable:

                query_lookup = _query_lookup(query, params)

                try:

                    cur = self.query_results[query_lookup]

                    used_cache = True

                except KeyError: cur = self.DbCursor(self)

            else: cur = self.db.cursor()

            # Run query

            cur.execute(query, params)

        finally:

            if self.debug: # only compute msg if needed

                if used_cache: cache_status = 'Cache hit'

                elif cacheable: cache_status = 'Cache miss'

                else: cache_status = 'Non-cacheable'

                self.log_debug(cache_status+': '

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

        return cur

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

        return _query_lookup(query, params) in self.query_results

    def with_savepoint(self, func):

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

        self.run_query('SAVEPOINT '+savepoint)

        self._savepoint += 1

        try: 

            try: return_val = func()

            finally:

                self._savepoint -= 1

                assert self._savepoint >= 0

        except:

            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint)

            raise

        else:

            self.run_query('RELEASE SAVEPOINT '+savepoint)

            self.do_autocommit()

            return return_val

    def do_autocommit(self):

        '''Autocommits if outside savepoint'''

        assert self._savepoint >= 0

        if self.autocommit and self._savepoint == 0:

            self.log_debug('Autocommiting')

            self.db.commit()

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

    module = util.root_module(db.db)

    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)

    else: raise NotImplementedError("Can't mogrify query for "+module+

        ' database')

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

        def run(): return run_raw_query(db, query, params, cacheable)

        if recover and not db.is_cached(query, params):

            return with_savepoint(db, run)

        else: return run() # don't need savepoint if cached

    except Exception, e:

        if not recover: raise # need savepoint to run index_cols()

        msg = exc.str_(e)

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

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

        if match:

            constraint, table = match.groups()

            try: cols = index_cols(db, table, constraint)

            except NotImplementedError: raise e

            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('NOT NULL', [match.group(1)], e)

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

            r'|date/time field value out of range): "(.+?)"\n'

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

        if match:

            value, name = match.groups()

            raise FunctionValueException(name, strings.to_unicode(value), e)

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

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

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

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

        raise # no specific exception raised

##### Basic queries

def next_version(name):

    '''Prepends the version # so it won't be removed if the name is truncated'''

    version = 1 # first existing name was version 0

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

    if match:

        version = int(match.group(1))+1

        name = match.group(2)

    return 'v'+str(version)+'_'+name

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

    '''Outputs a query to a temp table.

    For params, see run_query().

    '''

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

    else: # place rows in temp table

        assert isinstance(into, sql_gen.Table)

        kw_args['recover'] = True

        while True:

            try:

                create_query = 'CREATE'

                if not db.debug: create_query += ' TEMP'

                create_query += ' TABLE '+into.to_str(db)+' AS '+query

                return run_query(db, create_query, params, *args, **kw_args)

                    # CREATE TABLE AS sets rowcount to # rows in query

            except DuplicateTableException, e:

                into.name = next_version(into.name)

                # try again with next version of name

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: Code|str (for col name)

        * compare_right_side: ValueCond|literal value

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

        use all columns

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

    elif isinstance(conds, dict): 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:

        if distinct_on != []: order_by = None

        else: order_by = pkey(db, table0, recover=True)

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

    # Columns

    query += ' '

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

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

    # Main table

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

    # Add joins

    left_table = table0

    for join_ in tables:

        table = join_.table

        # Parse special values

        if join_.type_ is sql_gen.filter_out: # filter no match

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

                None))

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

        left_table = table

    missing = True

    if conds != []:

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

            for l, r in conds)))

        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

    if start != None:

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

        missing = False

    if missing: warnings.warn(DbWarning(

        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))

    return (query, [])

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

    '''For params, see mk_select() and run_query()'''

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

    cacheable = kw_args.pop('cacheable', True)

    query, params = mk_select(db, *args, **kw_args)

    return run_query(db, query, params, recover, cacheable)

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

    returning=None, embeddable=False):

    '''

    @param returning str|None An inserted column (such as pkey) to return

    @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

        query will be fully cached, not just if it raises an exception.

    '''

    table = sql_gen.as_Table(table)

    if cols == []: cols = None # no cols (all defaults) = unknown col names

    if cols != None: cols = [sql_gen.as_Col(v).to_str(db) for v in cols]

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

    if returning != None: returning = sql_gen.as_Col(returning, table)

    # Build query

    query = 'INSERT INTO '+table.to_str(db)

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

    query += ' '+select_query

    if returning != None:

        returning_name = copy.copy(returning)

        returning_name.table = None

        returning_name = returning_name.to_str(db)

        query += ' RETURNING '+returning_name

    if embeddable:

        assert returning != None

        # Create function

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

        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'

        while True:

            try:

                func_schema = None

                if not db.debug: func_schema = 'pg_temp'

                function = sql_gen.Table(function_name, func_schema).to_str(db)

                function_query = '''\

CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''

    LANGUAGE sql

    AS $$'''+mogrify(db, query, params)+''';$$;

'''

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

                break # this version was successful

            except DuplicateFunctionException, e:

                function_name = next_version(function_name)

                # try again with next version of name

        # Return query that uses function

        func_table = sql_gen.NamedTable('f', sql_gen.CustomCode(function+'()'),

            [returning_name]) # AS clause requires function alias

        return mk_select(db, func_table, start=0, order_by=None)

    return (query, params)

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

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

    @param into sql_gen.Table with suggested name of temp table to put RETURNING

        values in

    '''

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

    if into != None: kw_args['embeddable'] = True

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

    cacheable = kw_args.pop('cacheable', True)

    query, params = mk_insert_select(db, *args, **kw_args)

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

        cacheable=cacheable)

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

    else:

        cols = row.keys()

        row = row.values()

    row = list(row) # ensure that "!= []" works

    # Check for special values

    labels = []

    values = []

    for value in row:

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

        else:

            labels.append('%s')

            values.append(value)

    # Build query

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

    else: query = None

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

    elif module == 'MySQLdb': return db.insert_id()

    else: return None

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

    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')

##### Database structure queries

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

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

        recover=recover)))

def mk_flatten_mapping(db, into, cols, preserve=[]):

    '''Creates a mapping from original column names (which may have collisions)

    to names that will be distinct among the given tables.

    This is meant to be used for several tables that are being joined together.

    @param into The table for the new columns

    @param preserve [sql_gen.Col...] Columns not to rename. The tables of the

        provided Col objects will be changed to into, so make copies of

        them if you want to keep the original tables.

    @return dict(orig_col=new_col, ...)

        * orig_col: sql_gen.Col(orig_col_name, orig_table)

        * new_col: sql_gen.Col(orig_col_name, into)

        * All mappings use the into table so its name can easily be

          changed for all columns at once

    '''

    flatten_mapping = {}

    for col in cols: flatten_mapping[col] = sql_gen.Col(str(col), into)

    for col in preserve:

        orig_col = copy.copy(col)

        col.table = into

        flatten_mapping[orig_col] = col

    return flatten_mapping

def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):

    '''For params, see mk_flatten_mapping()

    @return See return value of mk_flatten_mapping()

    '''

    mapping = mk_flatten_mapping(db, into, cols, **kw_args)

    cols = [sql_gen.NamedCol(new.name, old) for old, new in mapping.iteritems()]

    run_query_into(db, *mk_select(db, joins, cols, limit=limit, start=start),

        into=into)

    return mapping

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

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

    return table_cols(db, table, recover)[0]

not_null_col = 'not_null'

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

    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''

    if not_null_col in table_cols(db, table, recover): return not_null_col

    else: return pkey(db, table, recover)

def index_cols(db, table, index):

    '''Can also use this for UNIQUE constraints, because a UNIQUE index is

    automatically created. When you don't know whether something is a UNIQUE

    constraint or a UNIQUE index, use this function.'''

    module = util.root_module(db.db)

    if module == 'psycopg2':

        return list(values(run_query(db, '''\

SELECT attname

FROM

(

        SELECT attnum, attname

        FROM pg_index

        JOIN pg_class index ON index.oid = indexrelid

        JOIN pg_class table_ ON table_.oid = indrelid

        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)

        WHERE

            table_.relname = %(table)s

            AND index.relname = %(index)s

    UNION

        SELECT attnum, attname

        FROM

        (

            SELECT

                indrelid

                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int

                    AS indkey

            FROM pg_index

            JOIN pg_class index ON index.oid = indexrelid

            JOIN pg_class table_ ON table_.oid = indrelid

            WHERE

                table_.relname = %(table)s

                AND index.relname = %(index)s

        ) s

        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey

) s

ORDER BY attnum

''',

            {'table': table, 'index': index}, cacheable=True)))

    else: raise NotImplementedError("Can't list index columns for "+module+

        ' database')

def constraint_cols(db, table, constraint):

    module = util.root_module(db.db)

    if module == 'psycopg2':

        return list(values(run_query(db, '''\

SELECT attname

FROM pg_constraint

JOIN pg_class ON pg_class.oid = conrelid

JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)

WHERE

    relname = %(table)s

    AND conname = %(constraint)s

ORDER BY attnum

''',

            {'table': table, 'constraint': constraint})))

    else: raise NotImplementedError("Can't list constraint columns for "+module+

        ' database')

row_num_col = '_row_num'

def add_row_num(db, table):

    '''Adds a row number column to a table. Its name is in row_num_col. It will

    be the primary key.'''

    table = sql_gen.as_Table(table).to_str(db)

    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col

        +' serial NOT NULL PRIMARY KEY')

def tables(db, schema='public', table_like='%'):

    module = util.root_module(db.db)

    params = {'schema': schema, 'table_like': table_like}

    if module == 'psycopg2':

        return values(run_query(db, '''\

SELECT tablename

FROM pg_tables

WHERE

    schemaname = %(schema)s

    AND tablename LIKE %(table_like)s

ORDER BY tablename

''',

            params, cacheable=True))

    elif module == 'MySQLdb':

        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,

            cacheable=True))

    else: raise NotImplementedError("Can't list tables for "+module+' database')

##### Database management

def empty_db(db, schema='public', **kw_args):

    '''For kw_args, see tables()'''

    for table in tables(db, schema, **kw_args): truncate(db, table, schema)

##### Heuristic queries

def put(db, table, row, pkey_=None, row_ct_ref=None):

    '''Recovers from errors.

    Only works under PostgreSQL (uses INSERT RETURNING).

    '''

    if pkey_ == None: pkey_ = pkey(db, table, recover=True)

    try:

        cur = insert(db, table, row, pkey_, recover=True)

        if row_ct_ref != None and cur.rowcount >= 0:

            row_ct_ref[0] += cur.rowcount

        return value(cur)

    except DuplicateKeyException, e: