/trunk/lib/sql.py - Annotate - BIEN 3 - NCEAS Projects

root/trunk/lib/sql.py @ 13227

-            aaronmk
+# Database access
-            aaronmk
+import copy
-            aaronmk
+import re
-            aaronmk
+import time
-            aaronmk
+import warnings
             aaronmk
-            aaronmk
+import exc
-            aaronmk
+import dicts
-            aaronmk
+import iters
-            aaronmk
+import lists
-            aaronmk
+import profiling
-            aaronmk
+from Proxy import Proxy
-            aaronmk
+import rand
-            aaronmk
+import sql_gen
-            aaronmk
+import strings
-            aaronmk
+import util
             aaronmk
-            aaronmk
+##### Exceptions
-            aaronmk
+def get_cur_query(cur, input_query=None):
-            aaronmk
+    raw_query = None
     if hasattr(cur, 'query'): raw_query = cur.query
     elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
             aaronmk
     if raw_query != None: return raw_query
-            aaronmk
+    else: return '[input] '+strings.ustr(input_query)
             aaronmk
-            aaronmk
+def _add_cursor_info(e, *args, **kw_args):
     '''For params, see get_cur_query()'''
-            aaronmk
+    exc.add_msg(e, 'query: '+strings.ustr(get_cur_query(*args, **kw_args)))
             aaronmk
-            aaronmk
+class DbException(exc.ExceptionWithCause):
-            aaronmk
+    def __init__(self, msg, cause=None, cur=None):
-            aaronmk
+        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
-            aaronmk
+        if cur != None: _add_cursor_info(self, cur)
-            aaronmk
+class ExceptionWithName(DbException):
     def __init__(self, name, cause=None):
-            aaronmk
+        DbException.__init__(self, 'for name: '
             +strings.as_tt(strings.ustr(name)), cause)
-            aaronmk
+        self.name = name
             aaronmk
-            aaronmk
+class ExceptionWithValue(DbException):
     def __init__(self, value, cause=None):
-            aaronmk
+        DbException.__init__(self, 'for value: '
             +strings.as_tt(strings.urepr(value)), cause)
-            aaronmk
+        self.value = value
-            aaronmk
+class ExceptionWithNameType(DbException):
     def __init__(self, type_, name, cause=None):
-            aaronmk
+        DbException.__init__(self, 'for type: '+strings.as_tt(strings.ustr(
             type_))+'; name: '+strings.as_tt(name), cause)
-            aaronmk
+        self.type = type_
         self.name = name
-            aaronmk
+class ConstraintException(DbException):
-            aaronmk
+    def __init__(self, name, cond, cols, cause=None):
         msg = 'Violated '+strings.as_tt(name)+' constraint'
-            aaronmk
+        if cond != None: msg += ' with condition '+strings.as_tt(cond)
-            aaronmk
+        if cols != []: msg += ' on columns: '+strings.as_tt(', '.join(cols))
         DbException.__init__(self, msg, cause)
-            aaronmk
+        self.name = name
-            aaronmk
+        self.cond = cond
-            aaronmk
+        self.cols = cols
             aaronmk
-            aaronmk
+class MissingCastException(DbException):
-            aaronmk
+    def __init__(self, type_, col=None, cause=None):
         msg = 'Missing cast to type '+strings.as_tt(type_)
         if col != None: msg += ' on column: '+strings.as_tt(col)
         DbException.__init__(self, msg, cause)
-            aaronmk
+        self.type = type_
         self.col = col
-            aaronmk
+class EncodingException(ExceptionWithName): pass
-            aaronmk
+class DuplicateKeyException(ConstraintException): pass
             aaronmk
-            aaronmk
+class NullValueException(ConstraintException): pass
             aaronmk
-            aaronmk
+class CheckException(ConstraintException): pass
-            aaronmk
+class InvalidValueException(ExceptionWithValue): pass
             aaronmk
-            aaronmk
+class InvalidTypeException(ExceptionWithNameType): pass
-            aaronmk
+class DuplicateException(ExceptionWithNameType): pass
             aaronmk
-            aaronmk
+class DoesNotExistException(ExceptionWithNameType): pass
-            aaronmk
+class EmptyRowException(DbException): pass
-            aaronmk
+##### Warnings
 class DbWarning(UserWarning): pass
-            aaronmk
+##### 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
-            aaronmk
+##### Escaping
             aaronmk
-            aaronmk
+def esc_name_by_module(module, name):
     if module == 'psycopg2' or module == None: quote = '"'
-            aaronmk
+    elif module == 'MySQLdb': quote = '`'
     else: raise NotImplementedError("Can't escape name for "+module+' database')
-            aaronmk
+    return sql_gen.esc_name(name, quote)
             aaronmk
 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
-            aaronmk
+##### Database connections
             aaronmk
-            aaronmk
+db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
             aaronmk
-            aaronmk
+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']
-            aaronmk
+log_debug_none = lambda msg, level=2: None
             aaronmk
-            aaronmk
+class DbConn:
-            aaronmk
+    def __init__(self, db_config, autocommit=True, caching=True,
-            aaronmk
+        log_debug=log_debug_none, debug_temp=False, src=None):
-            aaronmk
+        '''
         @param debug_temp Whether temporary objects should instead be permanent.
             This assists in debugging the internal objects used by the program.
-            aaronmk
+        @param src In autocommit mode, will be included in a comment in every
             query, to help identify the data source in pg_stat_activity.
-            aaronmk
+        '''
-            aaronmk
+        self.db_config = db_config
-            aaronmk
+        self.autocommit = autocommit
         self.caching = caching
-            aaronmk
+        self.log_debug = log_debug
-            aaronmk
+        self.debug = log_debug != log_debug_none
-            aaronmk
+        self.debug_temp = debug_temp
-            aaronmk
+        self.src = src
-            aaronmk
+        self.autoanalyze = False
-            aaronmk
+        self.autoexplain = False
         self.profile_row_ct = None
             aaronmk
-            aaronmk
+        self._savepoint = 0
-            aaronmk
+        self._reset()
             aaronmk
     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
-            aaronmk
+        state['log_debug'] = None # don't pickle the debug callback
-            aaronmk
+        state['_DbConn__db'] = None # don't pickle the connection
         return state
-            aaronmk
+    def clear_cache(self): self.query_results = {}
-            aaronmk
+    def _reset(self):
-            aaronmk
+        self.clear_cache()
-            aaronmk
+        assert self._savepoint == 0
-            aaronmk
+        self._notices_seen = set()
         self.__db = None
-            aaronmk
+    def connected(self): return self.__db != None
-            aaronmk
+    def close(self):
-            aaronmk
+        if not self.connected(): return
-            aaronmk
+        # Record that the automatic transaction is now closed
-            aaronmk
+        self._savepoint -= 1
             aaronmk
-            aaronmk
+        self.db.close()
-            aaronmk
+        self._reset()
             aaronmk
-            aaronmk
+    def reconnect(self):
         # Do not do this in test mode as it would roll back everything
         if self.autocommit: self.close()
         # Connection will be reopened automatically on first query
-            aaronmk
+    def _db(self):
         if self.__db == None:
             # Process db_config
             db_config = self.db_config.copy() # don't modify input!
-            aaronmk
+            schemas = db_config.pop('schemas', None)
-            aaronmk
+            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)
-            aaronmk
+            # Record that a transaction is already open
             self._savepoint += 1
-            aaronmk
+            # Configure connection
-            aaronmk
+            if hasattr(self.db, 'set_isolation_level'):
                 import psycopg2.extensions
                 self.db.set_isolation_level(
                     psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
-            aaronmk
+            if schemas != None:
-            aaronmk
+                search_path = [self.esc_name(s) for s in schemas.split(',')]
                 run_query(self, 'SET search_path TO '+(','.join(search_path)),
                     log_level=3)
             aaronmk
         return self.__db
             aaronmk
-            aaronmk
+    class DbCursor(Proxy):
-            aaronmk
+        def __init__(self, outer):
-            aaronmk
+            Proxy.__init__(self, outer.db.cursor())
-            aaronmk
+            self.outer = outer
-            aaronmk
+            self.query_results = outer.query_results
-            aaronmk
+            self.query_lookup = None
-            aaronmk
+            self.result = []
             aaronmk
-            aaronmk
+        def execute(self, query):
-            aaronmk
+            self._is_insert = query.startswith('INSERT')
-            aaronmk
+            self.query_lookup = query
-            aaronmk
+            try:
-            aaronmk
+                try: cur = self.inner.execute(query)
-            aaronmk
+                finally: self.query = get_cur_query(self.inner, query)
-            aaronmk
+            except Exception, e:
                 self.result = e # cache the exception as the result
                 self._cache_result()
                 raise
             aaronmk
             # Always cache certain queries
-            aaronmk
+            query = sql_gen.lstrip(query)
-            aaronmk
+            if query.startswith('CREATE') or query.startswith('ALTER'):
-            aaronmk
+                # structural changes
-            aaronmk
+                # Rest of query must be unique in the face of name collisions,
                 # so don't cache ADD COLUMN unless it has distinguishing comment
                 if query.find('ADD COLUMN') < 0 or query.endswith('*/'):
-            aaronmk
+                    self._cache_result()
-            aaronmk
+            elif self.rowcount == 0 and query.startswith('SELECT'): # empty
-            aaronmk
+                consume_rows(self) # fetch all rows so result will be cached
             aaronmk
-            aaronmk
+            return cur
             aaronmk
-            aaronmk
+        def fetchone(self):
             row = self.inner.fetchone()
-            aaronmk
+            if row != None: self.result.append(row)
             # otherwise, fetched all rows
-            aaronmk
+            else: self._cache_result()
             return row
         def _cache_result(self):
-            aaronmk
+            # For inserts that return a result set, don't cache result set since
             # inserts are not idempotent. Other non-SELECT queries don't have
             # their result set read, so only exceptions will be cached (an
             # invalid query will always be invalid).
-            aaronmk
+            if self.query_results != None and (not self._is_insert
-            aaronmk
+                or isinstance(self.result, Exception)):
-            aaronmk
+                assert self.query_lookup != None
-            aaronmk
+                self.query_results[self.query_lookup] = self.CacheCursor(
                     util.dict_subset(dicts.AttrsDictView(self),
                     ['query', 'result', 'rowcount', 'description']))
             aaronmk
-            aaronmk
+        class CacheCursor:
             def __init__(self, cached_result): self.__dict__ = cached_result
-            aaronmk
+            def execute(self, *args, **kw_args):
-            aaronmk
+                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
             aaronmk
-            aaronmk
+    def esc_value(self, value):
-            aaronmk
+        try: str_ = self.mogrify('%s', [value])
         except NotImplementedError, e:
             module = util.root_module(self.db)
             if module == 'MySQLdb':
                 import _mysql
                 str_ = _mysql.escape_string(value)
             else: raise e
-            aaronmk
+        return strings.to_unicode(str_)
             aaronmk
-            aaronmk
+    def esc_name(self, name): return esc_name(self, name) # calls global func
-            aaronmk
+    def std_code(self, str_):
         '''Standardizes SQL code.
         * Ensures that string literals are prefixed by `E`
         '''
         if str_.startswith("'"): str_ = 'E'+str_
         return str_
-            aaronmk
+    def can_mogrify(self):
-            aaronmk
+        module = util.root_module(self.db)
-            aaronmk
+        return module == 'psycopg2'
             aaronmk
-            aaronmk
+    def mogrify(self, query, params=None):
         if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
         else: raise NotImplementedError("Can't mogrify query")
-            aaronmk
+    def set_encoding(self, encoding):
         encoding_str = sql_gen.Literal(encoding)
         run_query(self, 'SET NAMES '+encoding_str.to_str(self))
-            aaronmk
+    def print_notices(self):
-            aaronmk
+        if hasattr(self.db, 'notices'):
             for msg in self.db.notices:
                 if msg not in self._notices_seen:
                     self._notices_seen.add(msg)
                     self.log_debug(msg, level=2)
             aaronmk
-            aaronmk
+    def run_query(self, query, cacheable=False, log_level=2,
-            aaronmk
+        debug_msg_ref=None):
-            aaronmk
+        '''
-            aaronmk
+        @param log_ignore_excs The log_level will be increased by 2 if the query
             throws one of these exceptions.
-            aaronmk
+        @param debug_msg_ref If specified, the log message will be returned in
             this instead of being output. This allows you to filter log messages
             depending on the result of the query.
-            aaronmk
+        '''
-            aaronmk
+        assert query != None
-            aaronmk
+        if self.autocommit and self.src != None:
-            aaronmk
+            query = sql_gen.esc_comment(self.src)+'\t'+query
             aaronmk
-            aaronmk
+        if not self.caching: cacheable = False
-            aaronmk
+        used_cache = False
             aaronmk
-            aaronmk
+        if self.debug:
             profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
-            aaronmk
+        try:
-            aaronmk
+            # Get cursor
             if cacheable:
-            aaronmk
+                try: cur = self.query_results[query]
-            aaronmk
+                except KeyError: cur = self.DbCursor(self)
-            aaronmk
+                else: used_cache = True
-            aaronmk
+            else: cur = self.db.cursor()
             # Run query
-            aaronmk
+            try: cur.execute(query)
-            aaronmk
+            except Exception, e:
                 _add_cursor_info(e, self, query)
                 raise
-            aaronmk
+            else: self.do_autocommit()
-            aaronmk
+        finally:
-            aaronmk
+            if self.debug:
-            aaronmk
+                profiler.stop(self.profile_row_ct)
             aaronmk
                 ## Log or return query
-            aaronmk
+                query = strings.ustr(get_cur_query(cur, query))
-            aaronmk
+                # Put the src comment on a separate line in the log file
                 query = query.replace('\t', '\n', 1)
             aaronmk
-            aaronmk
+                msg = 'DB query: '
             aaronmk
-            aaronmk
+                if used_cache: msg += 'cache hit'
                 elif cacheable: msg += 'cache miss'
                 else: msg += 'non-cacheable'
             aaronmk
-            aaronmk
+                msg += ':\n'+profiler.msg()+'\n'+strings.as_code(query, 'SQL')
             aaronmk
-            aaronmk
+                if debug_msg_ref != None: debug_msg_ref[0] = msg
                 else: self.log_debug(msg, log_level)
             aaronmk
                 self.print_notices()
             aaronmk
         return cur
             aaronmk
-            aaronmk
+    def is_cached(self, query): return query in self.query_results
             aaronmk
-            aaronmk
+    def with_autocommit(self, func):
-            aaronmk
+        import psycopg2.extensions
         prev_isolation_level = self.db.isolation_level
-            aaronmk
+        self.db.set_isolation_level(
             psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
-            aaronmk
+        try: return func()
-            aaronmk
+        finally: self.db.set_isolation_level(prev_isolation_level)
             aaronmk
-            aaronmk
+    def with_savepoint(self, func):
-            aaronmk
+        top = self._savepoint == 0
-            aaronmk
+        savepoint = 'level_'+str(self._savepoint)
             aaronmk
-            aaronmk
+        if self.debug:
-            aaronmk
+            self.log_debug('Begin transaction', level=4)
-            aaronmk
+            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
-            aaronmk
+        # Must happen before running queries so they don't get autocommitted
         self._savepoint += 1
-            aaronmk
+        if top: query = 'START TRANSACTION ISOLATION LEVEL READ COMMITTED'
         else: query = 'SAVEPOINT '+savepoint
         self.run_query(query, log_level=4)
         try:
             return func()
             if top: self.run_query('COMMIT', log_level=4)
-            aaronmk
+        except:
-            aaronmk
+            if top: query = 'ROLLBACK'
             else: query = 'ROLLBACK TO SAVEPOINT '+savepoint
             self.run_query(query, log_level=4)
-            aaronmk
+            raise
-            aaronmk
+        finally:
             # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
             # "The savepoint remains valid and can be rolled back to again"
             # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
-            aaronmk
+            if not top:
                 self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
             aaronmk
             self._savepoint -= 1
             assert self._savepoint >= 0
-            aaronmk
+            if self.debug:
                 profiler.stop(self.profile_row_ct)
-            aaronmk
+                self.log_debug('End transaction\n'+profiler.msg(), level=4)
             aaronmk
-            aaronmk
+            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
             aaronmk
     def do_autocommit(self):
         '''Autocommits if outside savepoint'''
-            aaronmk
+        assert self._savepoint >= 1
         if self.autocommit and self._savepoint == 1:
-            aaronmk
+            self.log_debug('Autocommitting', level=4)
-            aaronmk
+            self.db.commit()
             aaronmk
-            aaronmk
+    def col_info(self, col, cacheable=True):
-            aaronmk
+        module = util.root_module(self.db)
-            aaronmk
+        if module == 'psycopg2':
-            aaronmk
+            qual_table = sql_gen.Literal(col.table.to_str(self))
-            aaronmk
+            col_name_str = sql_gen.Literal(col.name)
             try:
-            aaronmk
+                type_, is_array, default, nullable = row(run_query(self, '''\
-            aaronmk
+SELECT
 format_type(COALESCE(NULLIF(typelem, 0), pg_type.oid), -1) AS type
 , typcategory = 'A' AS type_is_array
-            aaronmk
+, pg_get_expr(pg_attrdef.adbin, attrelid, true) AS default
-            aaronmk
+, NOT pg_attribute.attnotnull AS nullable
 FROM pg_attribute
 LEFT JOIN pg_type ON pg_type.oid = atttypid
 LEFT JOIN pg_attrdef ON adrelid = attrelid AND adnum = attnum
 WHERE
-            aaronmk
+attrelid = '''+qual_table.to_str(self)+'''::regclass
 AND attname = '''+col_name_str.to_str(self)+'''
-            aaronmk
+'''
-            aaronmk
+                    , recover=True, cacheable=cacheable, log_level=4))
-            aaronmk
+            except (DoesNotExistException, StopIteration):
                 raise sql_gen.NoUnderlyingTableException(col)
-            aaronmk
+            if is_array: type_ = sql_gen.ArrayType(type_)
-            aaronmk
+        else:
             table = sql_gen.Table('columns', 'information_schema')
             cols = [sql_gen.Col('data_type'), sql_gen.Col('udt_name'),
                 'column_default', sql_gen.Cast('boolean',
                 sql_gen.Col('is_nullable'))]
             conds = [('table_name', col.table.name),
                 ('column_name', strings.ustr(col.name))]
             schema = col.table.schema
             if schema != None: conds.append(('table_schema', schema))
             cur = select(self, table, cols, conds, order_by='table_schema',
                 limit=1, cacheable=cacheable, log_level=4)
             try: type_, extra_type, default, nullable = row(cur)
             except StopIteration: raise sql_gen.NoUnderlyingTableException(col)
             if type_ == 'USER-DEFINED': type_ = extra_type
             elif type_ == 'ARRAY':
                 type_ = sql_gen.ArrayType(strings.remove_prefix('_', extra_type,
                     require=True))
             aaronmk
-            aaronmk
+        if default != None: default = sql_gen.as_Code(default, self)
-            aaronmk
+        return sql_gen.TypedCol(col.name, type_, default, nullable)
             aaronmk
     def TempFunction(self, name):
         if self.debug_temp: schema = None
         else: schema = 'pg_temp'
         return sql_gen.Function(name, schema)
             aaronmk
-            aaronmk
+connect = DbConn
-            aaronmk
+##### Recoverable querying
             aaronmk
-            aaronmk
+def parse_exception(db, e, recover=False):
     msg = strings.ustr(e.args[0])
     msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)
     match = re.match(r'^invalid byte sequence for encoding "(.+?)":', msg)
     if match:
         encoding, = match.groups()
         raise EncodingException(encoding, e)
-            aaronmk
+    def make_DuplicateKeyException(constraint, e):
-            aaronmk
+        cols = []
         cond = None
         if recover: # need auto-rollback to run index_cols()
             try:
                 cols = index_cols(db, constraint)
                 cond = index_cond(db, constraint)
             except NotImplementedError: pass
-            aaronmk
+            except DatabaseErrors, e2: # e.g. "current transaction is aborted"
                 exc.print_ex(e2)
                 raise e
-            aaronmk
+        return DuplicateKeyException(constraint, cond, cols, e)
             aaronmk
-            aaronmk
+    match = re.match(r'^duplicate key value violates unique constraint "(.+?)"',
         msg)
     if match:
         constraint, = match.groups()
         raise make_DuplicateKeyException(constraint, e)
-            aaronmk
+    match = re.match(r'^could not create unique index "(.+?)"\n'
         r'DETAIL:  Key .+? is duplicated', msg)
     if match:
         constraint, = match.groups()
-            aaronmk
+        raise DuplicateKeyException(constraint, None, [], e)
             aaronmk
-            aaronmk
+    match = re.match(r'^null value in column "(.+?)" violates not-null'
         r' constraint', msg)
     if match:
         col, = match.groups()
         raise NullValueException('NOT NULL', None, [col], e)
     match = re.match(r'^new row for relation "(.+?)" violates check '
         r'constraint "(.+?)"', msg)
     if match:
         table, constraint = match.groups()
         constraint = sql_gen.Col(constraint, table)
         cond = None
         if recover: # need auto-rollback to run constraint_cond()
             try: cond = constraint_cond(db, constraint)
             except NotImplementedError: pass
         raise CheckException(constraint.to_str(db), cond, [], e)
-            aaronmk
+    match = re.match(r'^(?:invalid input (?:syntax|value)\b[^:]*'
-            aaronmk
+        r'|.+? out of range)(?:: "(.+?)")?', msg)
-            aaronmk
+    if match:
         value, = match.groups()
-            aaronmk
+        value = util.do_ignore_none(strings.to_unicode, value)
         raise InvalidValueException(value, e)
             aaronmk
     match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
         r'is of type', msg)
     if match:
         col, type_ = match.groups()
         raise MissingCastException(type_, col, e)
     match = re.match(r'^could not determine polymorphic type because '
         r'input has type "unknown"', msg)
     if match: raise MissingCastException('text', None, e)
-            aaronmk
+    # type mismatches
-            aaronmk
+    match = re.match(r'^.+? types (.+?) and (.+?) cannot be matched', msg)
     if match:
         type0, type1 = match.groups()
         raise MissingCastException(type0, None, e)
-            aaronmk
+    match = re.match(r'operator does not exist: (.+?) = (.+)', msg)
         # these appear when a staging table column is of the wrong type
     if match:
         type0, type1 = match.groups()
         raise MissingCastException(type0, None, e)
             aaronmk
-            aaronmk
+    match = re.match(r'^.*?\brelation "(.+?)" is not a (table)', msg)
     if match:
         name, type_ = match.groups()
         raise InvalidTypeException(type_, name, e)
-            aaronmk
+    typed_name_re = r'^(\S+) ("?)(.+?)\2(?: of relation ".+?")?'
-            aaronmk
+        # regexp must be followed with text for .*? to match properly
             aaronmk
-            aaronmk
+    match = re.match(typed_name_re+r'(?:\(.*?)? already exists', msg)
-            aaronmk
+    if match:
-            aaronmk
+        type_, quote, name = match.groups()
-            aaronmk
+        raise DuplicateException(type_, name, e)
     match = re.match(r'more than one (\S+) named ""(.+?)""', msg)
     if match:
         type_, name = match.groups()
         raise DuplicateException(type_, name, e)
-            aaronmk
+    # type mismatches which appear as "does not exist" errors
-            aaronmk
+    match = re.match(typed_name_re+r' does not exist', msg)
     if match:
-            aaronmk
+        type_, quote, name = match.groups()
-            aaronmk
+        if type_ == 'function':
-            aaronmk
+            match = re.match(r'^(.+?)\(.*\)$', name)
-            aaronmk
+            if match: # includes params, so is call rather than cast to regproc
-            aaronmk
+                function_name, = match.groups()
-            aaronmk
+                func = sql_gen.Function(function_name)
-            aaronmk
+                if function_exists(db, func) and msg.find('CAST') < 0:
-            aaronmk
+                    # not found only because of a missing cast
-            aaronmk
+                    try: type_ = function_param0_type(db, func)
                     except DuplicateException: type_ = 'text'
-            aaronmk
+                    col = None
-            aaronmk
+                    if type_ == 'anyelement': type_ = 'text'
-            aaronmk
+                    elif type_ == 'hstore': # cast just the value param
                         type_ = 'text'
                         col = 'value'
                     raise MissingCastException(type_, col, e)
-            aaronmk
+        raise DoesNotExistException(type_, name, e)
     raise # no specific exception raised
-            aaronmk
+def with_savepoint(db, func): return db.with_savepoint(func)
             aaronmk
-            aaronmk
+def run_query(db, query, recover=None, cacheable=False, log_level=2,
     log_ignore_excs=None, **kw_args):
-            aaronmk
+    '''For params, see DbConn.run_query()'''
-            aaronmk
+    if recover == None: recover = False
-            aaronmk
+    if log_ignore_excs == None: log_ignore_excs = ()
     log_ignore_excs = tuple(log_ignore_excs)
-            aaronmk
+    debug_msg_ref = [None]
             aaronmk
-            aaronmk
+    query = with_explain_comment(db, query)
             aaronmk
-            aaronmk
+    try:
-            aaronmk
+        try:
-            aaronmk
+            def run(): return db.run_query(query, cacheable, log_level,
-            aaronmk
+                debug_msg_ref, **kw_args)
-            aaronmk
+            if recover and not db.is_cached(query):
-            aaronmk
+                return with_savepoint(db, run)
             else: return run() # don't need savepoint if cached
-            aaronmk
+        except Exception, e:
             # Give failed EXPLAIN approximately the log_level of its query
             if query.startswith('EXPLAIN'): log_level -= 1
             parse_exception(db, e, recover)
-            aaronmk
+    except log_ignore_excs:
         log_level += 2
         raise
     finally:
-            aaronmk
+        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
             aaronmk
-            aaronmk
+##### Basic queries
-            aaronmk
+def is_explainable(query):
     # See <http://www.postgresql.org/docs/8.3/static/sql-explain.html#AEN57749>
-            aaronmk
+    return re.match(r'^(?:SELECT|INSERT|UPDATE|DELETE|VALUES|EXECUTE|DECLARE)\b'
         , query)
             aaronmk
-            aaronmk
+def explain(db, query, **kw_args):
     '''
     For params, see run_query().
     '''
-            aaronmk
+    kw_args.setdefault('log_level', 4)
             aaronmk
-            aaronmk
+    return strings.ustr(strings.join_lines(values(run_query(db,
         'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))
-            aaronmk
+        # not a higher log_level because it's useful to see what query is being
         # run before it's executed, which EXPLAIN effectively provides
-            aaronmk
+def has_comment(query): return query.endswith('*/')
 def with_explain_comment(db, query, **kw_args):
-            aaronmk
+    if db.autoexplain and not has_comment(query) and is_explainable(query):
-            aaronmk
+        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
             +explain(db, query, **kw_args))
     return query
-            aaronmk
+def next_version(name):
-            aaronmk
+    version = 1 # first existing name was version 0
-            aaronmk
+    match = re.match(r'^(.*)#(\d+)$', name)
-            aaronmk
+    if match:
-            aaronmk
+        name, version = match.groups()
         version = int(version)+1
-            aaronmk
+    return sql_gen.concat(name, '#'+str(version))
             aaronmk
-            aaronmk
+def lock_table(db, table, mode):
     table = sql_gen.as_Table(table)
     run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
-            aaronmk
+def run_query_into(db, query, into=None, add_pkey_=False, add_pkey_warn=True,
     **kw_args):
-            aaronmk
+    '''Outputs a query to a temp table.
     For params, see run_query().
     '''
-            aaronmk
+    if into == None: return run_query(db, query, **kw_args)
             aaronmk
     assert isinstance(into, sql_gen.Table)
-            aaronmk
+    into.is_temp = True
-            aaronmk
+    # "temporary tables cannot specify a schema name", so remove schema
     into.schema = None
             aaronmk
-            aaronmk
+    kw_args['recover'] = True
-            aaronmk
+    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
             aaronmk
-            aaronmk
+    temp = not db.debug_temp # tables are permanent in debug_temp mode
             aaronmk
     # Create table
     while True:
         create_query = 'CREATE'
         if temp: create_query += ' TEMP'
         create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
             aaronmk
-            aaronmk
+        try:
             cur = run_query(db, create_query, **kw_args)
                 # CREATE TABLE AS sets rowcount to # rows in query
             break
-            aaronmk
+        except DuplicateException, e:
-            aaronmk
+            into.name = next_version(into.name)
             # try again with next version of name
-            aaronmk
+    if add_pkey_: add_pkey_or_index(db, into, warn=add_pkey_warn)
             aaronmk
     # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
     # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
     # table is going to be used in complex queries, it is wise to run ANALYZE on
     # the temporary table after it is populated."
     # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
     # If into is not a temp table, ANALYZE is useful but not required.
-            aaronmk
+    analyze(db, into)
             aaronmk
     return cur
             aaronmk
-            aaronmk
+order_by_pkey = object() # tells mk_select() to order by the pkey
-            aaronmk
+distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
-            aaronmk
+def has_subset_func(db, table):
     return sql_gen.is_underlying_table(table) and function_exists(db, table)
-            aaronmk
+def mk_select(db, tables=None, fields=None, conds=None, distinct_on=[],
-            aaronmk
+    limit=None, start=None, order_by=order_by_pkey, default_table=None,
     explain=True):
-            aaronmk
+    '''
-            aaronmk
+    @param tables The single table to select from, or a list of tables to join
-            aaronmk
+        together, with tables after the first being sql_gen.Join objects
-            aaronmk
+    @param fields Use None to select all fields in the table
-            aaronmk
+    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
-            aaronmk
+        * container can be any iterable type
-            aaronmk
+        * compare_left_side: sql_gen.Code|str (for col name)
         * compare_right_side: sql_gen.ValueCond|literal value
-            aaronmk
+    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
         use all columns
-            aaronmk
+    @return query
-            aaronmk
+    '''
-            aaronmk
+    # Parse tables param
-            aaronmk
+    tables = lists.mk_seq(tables)
-            aaronmk
+    tables = list(tables) # don't modify input! (list() copies input)
-            aaronmk
+    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
             aaronmk
-            aaronmk
+    # Parse other params
-            aaronmk
+    if conds == None: conds = []
-            aaronmk
+    elif dicts.is_dict(conds): conds = conds.items()
-            aaronmk
+    conds = list(conds) # don't modify input! (list() copies input)
-            aaronmk
+    assert limit == None or isinstance(limit, (int, long))
     assert start == None or isinstance(start, (int, long))
-            aaronmk
+    if limit == 0: order_by = None
-            aaronmk
+    if order_by in (None, order_by_pkey) and has_subset_func(db, table0):
         # can use subset function for fast querying at large OFFSET values
-            aaronmk
+        run_query(db, 'SET LOCAL enable_sort TO off')
-            aaronmk
+        table0 = sql_gen.FunctionCall(table0, limit_=limit, offset_=start)
         if limit != 0: limit = None # done by function
         start = None # done by function
         order_by = None # done by function
     elif order_by is order_by_pkey:
-            aaronmk
+        if lists.is_seq(distinct_on) and distinct_on: order_by = distinct_on[0]
         elif table0 != None: order_by = table_order_by(db, table0, recover=True)
         else: order_by = None
             aaronmk
-            aaronmk
+    query = 'SELECT'
             aaronmk
-            aaronmk
+    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
             aaronmk
-            aaronmk
+    # DISTINCT ON columns
-            aaronmk
+    if distinct_on != []:
-            aaronmk
+        query += '\nDISTINCT'
-            aaronmk
+        if distinct_on is not distinct_on_all:
-            aaronmk
+            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
     # Columns
-            aaronmk
+    if query.find('\n') >= 0: whitespace = '\n'
     else: whitespace = ' '
     if fields == None: query += whitespace+'*'
-            aaronmk
+    else:
         assert fields != []
-            aaronmk
+        if len(fields) > 1: whitespace = '\n'
         query += whitespace+('\n, '.join(map(parse_col, fields)))
             aaronmk
     # Main table
-            aaronmk
+    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
     else: whitespace = ' '
-            aaronmk
+    if table0 != None: query += whitespace+'FROM '+table0.to_str(db)
             aaronmk
-            aaronmk
+    # Add joins
-            aaronmk
+    left_table = table0
-            aaronmk
+    for join_ in tables:
         table = join_.table
             aaronmk
-            aaronmk
+        # Parse special values
         if join_.type_ is sql_gen.filter_out: # filter no match
-            aaronmk
+            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
-            aaronmk
+                sql_gen.CompareCond(None, '~=')))
             aaronmk
-            aaronmk
+        query += '\n'+join_.to_str(db, left_table)
             aaronmk
         left_table = table
-            aaronmk
+    missing = True
-            aaronmk
+    if conds != []:
-            aaronmk
+        if len(conds) == 1: whitespace = ' '
         else: whitespace = '\n'
-            aaronmk
+        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
             .to_str(db) for l, r in conds], 'WHERE')
-            aaronmk
+    if order_by != None:
-            aaronmk
+        query += '\nORDER BY '+sql_gen.as_Col(order_by).to_str(db)
-            aaronmk
+    if limit != None: query += '\nLIMIT '+str(limit)
-            aaronmk
+    if start != None:
-            aaronmk
+        if start != 0: query += '\nOFFSET '+str(start)
             aaronmk
-            aaronmk
+    if explain: query = with_explain_comment(db, query)
             aaronmk
-            aaronmk
+    return query
             aaronmk
-            aaronmk
+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)
-            aaronmk
+    log_level = kw_args.pop('log_level', 2)
             aaronmk
-            aaronmk
+    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
         log_level=log_level)
             aaronmk
-            aaronmk
+def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
-            aaronmk
+    embeddable=False, ignore=False, src=None):
-            aaronmk
+    '''
     @param returning str|None An inserted column (such as pkey) to return
-            aaronmk
+    @param embeddable Whether the query should be embeddable as a nested SELECT.
-            aaronmk
+        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.
-            aaronmk
+    @param ignore Whether to ignore duplicate keys.
-            aaronmk
+    @param src Will be included in the name of any created function, to help
         identify the data source in pg_stat_activity.
-            aaronmk
+    '''
-            aaronmk
+    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
-            aaronmk
+    if cols == []: cols = None # no cols (all defaults) = unknown col names
-            aaronmk
+    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
-            aaronmk
+    if select_query == None: select_query = 'DEFAULT VALUES'
-            aaronmk
+    if returning != None: returning = sql_gen.as_Col(returning, table)
             aaronmk
-            aaronmk
+    first_line = 'INSERT INTO '+table.to_str(db)
             aaronmk
-            aaronmk
+    def mk_insert(select_query):
         query = first_line
-            aaronmk
+        if cols != None:
             query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
-            aaronmk
+        query += '\n'+select_query
         if returning != None:
             returning_name_col = sql_gen.to_name_only_col(returning)
             query += '\nRETURNING '+returning_name_col.to_str(db)
         return query
             aaronmk
-            aaronmk
+    return_type = sql_gen.CustomCode('unknown')
     if returning != None: return_type = sql_gen.ColType(returning)
             aaronmk
-            aaronmk
+    if ignore:
-            aaronmk
+        # Always return something to set the correct rowcount
         if returning == None: returning = sql_gen.NamedCol('NULL', None)
-            aaronmk
+        embeddable = True # must use function
             aaronmk
-            aaronmk
+        if cols == None: row = [sql_gen.Col(sql_gen.all_cols, 'row')]
         else: row = [sql_gen.Col(c.name, 'row') for c in cols]
             aaronmk
-            aaronmk
+        query = sql_gen.RowExcIgnore(sql_gen.RowType(table), select_query,
-            aaronmk
+            sql_gen.ReturnQuery(mk_insert(sql_gen.Values(row).to_str(db))),
             cols)
-            aaronmk
+    else: query = mk_insert(select_query)
-            aaronmk
+    if embeddable:
         # Create function
-            aaronmk
+        function_name = sql_gen.clean_name(first_line)
-            aaronmk
+        if src != None: function_name = src+': '+function_name
-            aaronmk
+        while True:
             try:
-            aaronmk
+                func = db.TempFunction(function_name)
-            aaronmk
+                def_ = sql_gen.FunctionDef(func, sql_gen.SetOf(return_type),
                     query)
             aaronmk
-            aaronmk
+                run_query(db, def_.to_str(db), recover=True, cacheable=True,
-            aaronmk
+                    log_ignore_excs=(DuplicateException,))
-            aaronmk
+                break # this version was successful
-            aaronmk
+            except DuplicateException, e:
-            aaronmk
+                function_name = next_version(function_name)
                 # try again with next version of name
             aaronmk
-            aaronmk
+        # Return query that uses function
-            aaronmk
+        cols = None
         if returning != None: cols = [returning]
-            aaronmk
+        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(func), cols)
             # AS clause requires function alias
-            aaronmk
+        return mk_select(db, func_table, order_by=None)
             aaronmk
-            aaronmk
+    return query
             aaronmk
-            aaronmk
+def insert_select(db, table, *args, **kw_args):
-            aaronmk
+    '''For params, see mk_insert_select() and run_query_into()
-            aaronmk
+    @param into sql_gen.Table with suggested name of temp table to put RETURNING
         values in
-            aaronmk
+    '''
-            aaronmk
+    returning = kw_args.get('returning', None)
     ignore = kw_args.get('ignore', False)
-            aaronmk
+    into = kw_args.pop('into', None)
     if into != None: kw_args['embeddable'] = True
-            aaronmk
+    recover = kw_args.pop('recover', None)
-            aaronmk
+    if ignore: recover = True
-            aaronmk
+    cacheable = kw_args.pop('cacheable', True)
-            aaronmk
+    log_level = kw_args.pop('log_level', 2)
             aaronmk
-            aaronmk
+    rowcount_only = ignore and returning == None # keep NULL rows on server
     if rowcount_only: into = sql_gen.Table('rowcount')
-            aaronmk
+    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
         into, recover=recover, cacheable=cacheable, log_level=log_level)
-            aaronmk
+    if rowcount_only: empty_temp(db, into)
-            aaronmk
+    autoanalyze(db, table)
     return cur
             aaronmk
-            aaronmk
+default = sql_gen.default # tells insert() to use the default value for a column
             aaronmk
-            aaronmk
+def insert(db, table, row, *args, **kw_args):
-            aaronmk
+    '''For params, see insert_select()'''
-            aaronmk
+    ignore = kw_args.pop('ignore', False)
-            aaronmk
+    if ignore: kw_args.setdefault('recover', True)
             aaronmk
-            aaronmk
+    if lists.is_seq(row): cols = None
     else:
         cols = row.keys()
         row = row.values()
-            aaronmk
+    row = list(row) # ensure that "== []" works
             aaronmk
-            aaronmk
+    if row == []: query = None
     else: query = sql_gen.Values(row).to_str(db)
             aaronmk
-            aaronmk
+    try: return insert_select(db, table, cols, query, *args, **kw_args)
-            aaronmk
+    except (DuplicateKeyException, NullValueException):
-            aaronmk
+        if not ignore: raise
-            aaronmk
+        return None
             aaronmk
-            aaronmk
+def mk_update(db, table, changes=None, cond=None, in_place=False,
-            aaronmk
+    cacheable_=True):
-            aaronmk
+    '''
     @param changes [(col, new_value),...]
         * container can be any iterable type
         * col: sql_gen.Code|str (for col name)
         * new_value: sql_gen.Code|literal value
     @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
-            aaronmk
+    @param in_place If set, locks the table and updates rows in place.
         This avoids creating dead rows in PostgreSQL.
         * cond must be None
-            aaronmk
+    @param cacheable_ Whether column structure information used to generate the
-            aaronmk
+        query can be cached
-            aaronmk
+    @return str query
     '''
-            aaronmk
+    table = sql_gen.as_Table(table)
     changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
         for c, v in changes]
-            aaronmk
+    if in_place:
         assert cond == None
             aaronmk
-            aaronmk
+        def col_type(col):
             return sql_gen.canon_type(db.col_info(
                 sql_gen.with_default_table(c, table), cacheable_).type)
         changes = [(c, v, col_type(c)) for c, v in changes]
-            aaronmk
+        query = 'ALTER TABLE '+table.to_str(db)+'\n'
-            aaronmk
+        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '+t+'\nUSING '
             +v.to_str(db) for c, v, t in changes))
-            aaronmk
+    else:
         query = 'UPDATE '+table.to_str(db)+'\nSET\n'
         query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
             for c, v in changes))
         if cond != None: query += '\nWHERE\n'+cond.to_str(db)
             aaronmk
-            aaronmk
+    query = with_explain_comment(db, query)
-            aaronmk
+    return query
-            aaronmk
+def update(db, table, *args, **kw_args):
-            aaronmk
+    '''For params, see mk_update() and run_query()'''
     recover = kw_args.pop('recover', None)
-            aaronmk
+    cacheable = kw_args.pop('cacheable', False)
-            aaronmk
+    log_level = kw_args.pop('log_level', 2)
             aaronmk
-            aaronmk
+    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
         cacheable, log_level=log_level)
     autoanalyze(db, table)
     return cur
             aaronmk
-            aaronmk
+def mk_delete(db, table, cond=None):
     '''
     @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
     @return str query
     '''
     query = 'DELETE FROM '+table.to_str(db)
     if cond != None: query += '\nWHERE '+cond.to_str(db)
     query = with_explain_comment(db, query)
     return query
 def delete(db, table, *args, **kw_args):
     '''For params, see mk_delete() and run_query()'''
     recover = kw_args.pop('recover', None)
-            aaronmk
+    cacheable = kw_args.pop('cacheable', True)
-            aaronmk
+    log_level = kw_args.pop('log_level', 2)
     cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
         cacheable, log_level=log_level)
     autoanalyze(db, table)
     return cur
-            aaronmk
+def last_insert_id(db):
-            aaronmk
+    module = util.root_module(db.db)
-            aaronmk
+    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
     elif module == 'MySQLdb': return db.insert_id()
     else: return None
             aaronmk
-            aaronmk
+def define_func(db, def_):
     func = def_.function
     while True:
         try:
             run_query(db, def_.to_str(db), recover=True, cacheable=True,
                 log_ignore_excs=(DuplicateException,))
             break # successful
         except DuplicateException:
-            aaronmk
+            func.name = next_version(func.name)
-            aaronmk
+            # try again with next version of name
-            aaronmk
+def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
-            aaronmk
+    '''Creates a mapping from original column names (which may have collisions)
-            aaronmk
+    to names that will be distinct among the columns' tables.
-            aaronmk
+    This is meant to be used for several tables that are being joined together.
-            aaronmk
+    @param cols The columns to combine. Duplicates will be removed.
     @param into The table for the new columns.
-            aaronmk
+    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
         columns will be included in the mapping even if they are not in cols.
         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.
     @param as_items Whether to return a list of dict items instead of a dict
-            aaronmk
+    @return dict(orig_col=new_col, ...)
         * orig_col: sql_gen.Col(orig_col_name, orig_table)
-            aaronmk
+        * new_col: sql_gen.Col(orig_col_name, into)
         * All mappings use the into table so its name can easily be
-            aaronmk
+          changed for all columns at once
     '''
-            aaronmk
+    cols = lists.uniqify(cols)
-            aaronmk
+    items = []
-            aaronmk
+    for col in preserve:
-            aaronmk
+        orig_col = copy.copy(col)
-            aaronmk
+        col.table = into
-            aaronmk
+        items.append((orig_col, col))
     preserve = set(preserve)
     for col in cols:
-            aaronmk
+        if col not in preserve:
-            aaronmk
+            items.append((col, sql_gen.Col(strings.ustr(col), into, col.srcs)))
             aaronmk
     if not as_items: items = dict(items)
     return items
             aaronmk
-            aaronmk
+def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
-            aaronmk
+    '''For params, see mk_flatten_mapping()
     @return See return value of mk_flatten_mapping()
     '''
-            aaronmk
+    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
     cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
-            aaronmk
+    # don't warn if can't create pkey, because this just indicates that a
     # set-returning function was used
-            aaronmk
+    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
-            aaronmk
+        into=into, add_pkey_=True, add_pkey_warn=False)
-            aaronmk
+        # don't cache because the temp table will usually be truncated after use
-            aaronmk
+    return dict(items)
             aaronmk
-            aaronmk
+##### Database structure introspection
             aaronmk
-            aaronmk
+#### Tables
-            aaronmk
+def tables(db, schema_like='public', table_like='%', exact=False,
     cacheable=True):
-            aaronmk
+    if exact: compare = '='
     else: compare = 'LIKE'
     module = util.root_module(db.db)
     if module == 'psycopg2':
         conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
             ('tablename', sql_gen.CompareCond(table_like, compare))]
         return values(select(db, 'pg_tables', ['tablename'], conds,
-            aaronmk
+            order_by='tablename', cacheable=cacheable, log_level=4))
-            aaronmk
+    elif module == 'MySQLdb':
         return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
             , cacheable=True, log_level=4))
     else: raise NotImplementedError("Can't list tables for "+module+' database')
-            aaronmk
+def table_exists(db, table, cacheable=True):
-            aaronmk
+    table = sql_gen.as_Table(table)
-            aaronmk
+    return list(tables(db, table.schema, table.name, True, cacheable)) != []
             aaronmk
-            aaronmk
+def table_row_count(db, table, recover=None):
-            aaronmk
+    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
-            aaronmk
+        order_by=None), recover=recover, log_level=3))
             aaronmk
-            aaronmk
+def table_col_names(db, table, recover=None):
-            aaronmk
+    return list(col_names(select(db, table, limit=0, recover=recover,
         log_level=4)))
             aaronmk
-            aaronmk
+def table_cols(db, table, *args, **kw_args):
     return [sql_gen.as_Col(strings.ustr(c), table)
         for c in table_col_names(db, table, *args, **kw_args)]
-            aaronmk
+def table_pkey_index(db, table, recover=None):
     table_str = sql_gen.Literal(table.to_str(db))
     try:
         return sql_gen.Table(value(run_query(db, '''\
 SELECT relname
 FROM pg_index
 JOIN pg_class index ON index.oid = indexrelid
 WHERE
 indrelid = '''+table_str.to_str(db)+'''::regclass
 AND indisprimary
 '''
             , recover, cacheable=True, log_level=4)), table.schema)
     except StopIteration: raise DoesNotExistException('primary key', '')
-            aaronmk
+def table_pkey_col(db, table, recover=None):
-            aaronmk
+    table = sql_gen.as_Table(table)
-            aaronmk
+    module = util.root_module(db.db)
     if module == 'psycopg2':
         return sql_gen.Col(index_cols(db, table_pkey_index(db, table,
             recover))[0], table)
     else:
         join_cols = ['table_schema', 'table_name', 'constraint_schema',
             'constraint_name']
         tables = [sql_gen.Table('key_column_usage', 'information_schema'),
             sql_gen.Join(
                 sql_gen.Table('table_constraints', 'information_schema'),
                 dict(((c, sql_gen.join_same_not_null) for c in join_cols)))]
         cols = [sql_gen.Col('column_name')]
         conds = [('constraint_type', 'PRIMARY KEY'), ('table_name', table.name)]
         schema = table.schema
         if schema != None: conds.append(('table_schema', schema))
         order_by = 'position_in_unique_constraint'
         try: return sql_gen.Col(value(select(db, tables, cols, conds,
             order_by=order_by, limit=1, log_level=4)), table)
         except StopIteration: raise DoesNotExistException('primary key', '')
             aaronmk
-            aaronmk
+def table_has_pkey(db, table, recover=None):
     try: table_pkey_col(db, table, recover)
     except DoesNotExistException: return False
     else: return True
-            aaronmk
+def pkey_name(db, table, recover=None):
     '''If no pkey, returns the first column in the table.'''
-            aaronmk
+    return pkey_col(db, table, recover).name
             aaronmk
-            aaronmk
+def pkey_col(db, table, recover=None):
-            aaronmk
+    '''If no pkey, returns the first column in the table.'''
-            aaronmk
+    try: return table_pkey_col(db, table, recover)
-            aaronmk
+    except DoesNotExistException: return table_cols(db, table, recover)[0]
             aaronmk
-            aaronmk
+not_null_col = 'not_null_col'
             aaronmk
 def table_not_null_col(db, table, recover=None):
     '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
-            aaronmk
+    if not_null_col in table_col_names(db, table, recover): return not_null_col
-            aaronmk
+    else: return pkey_name(db, table, recover)
             aaronmk
-            aaronmk
+def constraint_cond(db, constraint):
     module = util.root_module(db.db)
     if module == 'psycopg2':
         table_str = sql_gen.Literal(constraint.table.to_str(db))
         name_str = sql_gen.Literal(constraint.name)
         return value(run_query(db, '''\
 SELECT consrc
 FROM pg_constraint
 WHERE
 conrelid = '''+table_str.to_str(db)+'''::regclass
 AND conname = '''+name_str.to_str(db)+'''
 '''
             , cacheable=True, log_level=4))
-            aaronmk
+    else: raise NotImplementedError("Can't get constraint condition for "
         +module+' database')
             aaronmk
-            aaronmk
+def index_exprs(db, index):
-            aaronmk
+    index = sql_gen.as_Table(index)
-            aaronmk
+    module = util.root_module(db.db)
     if module == 'psycopg2':
-            aaronmk
+        qual_index = sql_gen.Literal(index.to_str(db))
-            aaronmk
+        return list(values(run_query(db, '''\
-            aaronmk
+SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
 FROM pg_index
 WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
-            aaronmk
+'''
             , cacheable=True, log_level=4)))
-            aaronmk
+    else: raise NotImplementedError()
             aaronmk
-            aaronmk
+def index_cols(db, 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.'''
     return map(sql_gen.parse_expr_col, index_exprs(db, index))
-            aaronmk
+def index_cond(db, index):
     index = sql_gen.as_Table(index)
     module = util.root_module(db.db)
     if module == 'psycopg2':
         qual_index = sql_gen.Literal(index.to_str(db))
         return value(run_query(db, '''\
 SELECT pg_get_expr(indpred, indrelid, true)
 FROM pg_index
 WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
 '''
             , cacheable=True, log_level=4))
     else: raise NotImplementedError()
-            aaronmk
+def index_order_by(db, index):
     return sql_gen.CustomCode(', '.join(index_exprs(db, index)))
 def table_cluster_on(db, table, recover=None):
     '''
     @return The table's cluster index, or its pkey if none is set
     '''
     table_str = sql_gen.Literal(table.to_str(db))
     try:
         return sql_gen.Table(value(run_query(db, '''\
 SELECT relname
 FROM pg_index
 JOIN pg_class index ON index.oid = indexrelid
 WHERE
 indrelid = '''+table_str.to_str(db)+'''::regclass
 AND indisclustered
 '''
             , recover, cacheable=True, log_level=4)), table.schema)
     except StopIteration: return table_pkey_index(db, table, recover)
 def table_order_by(db, table, recover=None):
-            aaronmk
+    '''
     @return None if table is view, because table_cluster_on() would return None
     '''
-            aaronmk
+    if table.order_by == None:
         try: table.order_by = index_order_by(db, table_cluster_on(db, table,
             recover))
         except DoesNotExistException: pass
     return table.order_by
             aaronmk
-            aaronmk
+#### Views
 def view_exists(db, view):
     view_str = sql_gen.Literal(view.to_str(db))
     try:
         return value(run_query(db, '''\
 SELECT relkind = 'v'
 FROM pg_class
 WHERE oid = '''+view_str.to_str(db)+'''::regclass
 '''
             , cacheable=True, log_level=4))
     except DoesNotExistException: return False
-            aaronmk
+#### Columns
 def col_default_value(db, col):
     assert isinstance(col, sql_gen.Col)
     return value(run_query(db,
         'SELECT util.col_default_value('+sql_gen.col2col_ref(db, col)+')'))
 def col_is_constant(db, col):
     assert isinstance(col, sql_gen.Col)
     return value(run_query(db,
         'SELECT util.is_constant('+sql_gen.col2col_ref(db, col)+')'))
-            aaronmk
+#### Functions
 def function_exists(db, function):
-            aaronmk
+    qual_function = sql_gen.Literal(function.to_str(db))
     try:
-            aaronmk
+        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
             recover=True, cacheable=True, log_level=4)
-            aaronmk
+    except DoesNotExistException: return False
-            aaronmk
+    except DuplicateException: return True # overloaded function
-            aaronmk
+    else: return True
             aaronmk
-            aaronmk
+def function_param0_type(db, function):
     qual_function = sql_gen.Literal(function.to_str(db))
     return value(run_query(db, '''\
 SELECT proargtypes[0]::regtype
 FROM pg_proc
 WHERE oid = '''+qual_function.to_str(db)+'''::regproc
 '''
         , cacheable=True, log_level=4))
-            aaronmk
+##### Structural changes
 #### Columns
-            aaronmk
+def add_col(db, table, col, comment=None, if_not_exists=False, **kw_args):
-            aaronmk
+    '''
     @param col TypedCol Name may be versioned, so be sure to propagate any
         renaming back to any source column for the TypedCol.
     @param comment None|str SQL comment used to distinguish columns of the same
         name from each other when they contain different data, to allow the
         ADD COLUMN query to be cached. If not set, query will not be cached.
     '''
     assert isinstance(col, sql_gen.TypedCol)
     while True:
         str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
         if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
         try:
             run_query(db, str_, recover=True, cacheable=True, **kw_args)
             break
         except DuplicateException:
-            aaronmk
+            if if_not_exists: raise
-            aaronmk
+            col.name = next_version(col.name)
             # try again with next version of name
 def add_not_null(db, col):
     table = col.table
     col = sql_gen.to_name_only_col(col)
     run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
         +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
-            aaronmk
+def drop_not_null(db, col):
     table = col.table
     col = sql_gen.to_name_only_col(col)
     run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
         +col.to_str(db)+' DROP NOT NULL', cacheable=True, log_level=3)
-            aaronmk
+row_num_col = '_row_num'
-            aaronmk
+row_num_col_def = sql_gen.TypedCol('', 'serial', nullable=False,
-            aaronmk
+    constraints='PRIMARY KEY')
-            aaronmk
+def add_row_num(db, table, name=row_num_col):
     '''Adds a row number column to a table. Its definition is in
     row_num_col_def. It will be the primary key.'''
     col_def = copy.copy(row_num_col_def)
     col_def.name = name
-            aaronmk
+    add_col(db, table, col_def, comment='', if_not_exists=True, log_level=3)
             aaronmk
 #### Indexes
 def add_pkey(db, table, cols=None, recover=None):
     '''Adds a primary key.
     @param cols [sql_gen.Col,...] The columns in the primary key.
         Defaults to the first column in the table.
     @pre The table must not already have a primary key.
     '''
     table = sql_gen.as_Table(table)
-            aaronmk
+    if cols == None: cols = [pkey_name(db, table, recover)]
-            aaronmk
+    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
     run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
         +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
         log_ignore_excs=(DuplicateException,))
-            aaronmk
+def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
-            aaronmk
+    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
-            aaronmk
+    Currently, only function calls and literal values are supported expressions.
-            aaronmk
+    @param ensure_not_null_ If set, translates NULL values to sentinel values.
-            aaronmk
+        This allows indexes to be used for comparisons where NULLs are equal.
-            aaronmk
+    '''
-            aaronmk
+    exprs = lists.mk_seq(exprs)
             aaronmk
-            aaronmk
+    # Parse exprs
     old_exprs = exprs[:]
     exprs = []
     cols = []
     for i, expr in enumerate(old_exprs):
-            aaronmk
+        expr = sql_gen.as_Col(expr, table)
             aaronmk
-            aaronmk
+        # Handle nullable columns
-            aaronmk
+        if ensure_not_null_:
-            aaronmk
+            try: expr = sql_gen.ensure_not_null(db, expr)
-            aaronmk
+            except KeyError: pass # unknown type, so just create plain index
             aaronmk
-            aaronmk
+        # Extract col
-            aaronmk
+        expr = copy.deepcopy(expr) # don't modify input!
-            aaronmk
+        col = expr
         if isinstance(expr, sql_gen.FunctionCall): col = expr.args[0]
         expr = sql_gen.cast_literal(expr)
         if not isinstance(expr, (sql_gen.Expr, sql_gen.Col)):
-            aaronmk
+            expr = sql_gen.Expr(expr)
             aaronmk
             aaronmk
         # Extract table
         if table == None:
             assert sql_gen.is_table_col(col)
             table = col.table
-            aaronmk
+        if isinstance(col, sql_gen.Col): col.table = None
             aaronmk
         exprs.append(expr)
         cols.append(col)
             aaronmk
-            aaronmk
+    table = sql_gen.as_Table(table)
-            aaronmk
+    # Add index
-            aaronmk
+    str_ = 'CREATE'
     if unique: str_ += ' UNIQUE'
     str_ += ' INDEX ON '+table.to_str(db)+' ('+(
         ', '.join((v.to_str(db) for v in exprs)))+')'
     run_query(db, str_, recover=True, cacheable=True, log_level=3)
             aaronmk
-            aaronmk
+def add_pkey_index(db, table): add_index(db, pkey_col(db, table), table)
-            aaronmk
+def add_pkey_or_index(db, table, cols=None, recover=None, warn=False):
     try: add_pkey(db, table, cols, recover)
     except DuplicateKeyException, e:
         if warn: warnings.warn(UserWarning(exc.str_(e)))
-            aaronmk
+        add_pkey_index(db, table)
             aaronmk
-            aaronmk
+already_indexed = object() # tells add_indexes() the pkey has already been added
 def add_indexes(db, table, has_pkey=True):
     '''Adds an index on all columns in a table.
     @param has_pkey bool|already_indexed Whether a pkey instead of a regular
         index should be added on the first column.
         * If already_indexed, the pkey is assumed to have already been added
     '''
-            aaronmk
+    cols = table_col_names(db, table)
-            aaronmk
+    if has_pkey:
         if has_pkey is not already_indexed: add_pkey(db, table)
         cols = cols[1:]
     for col in cols: add_index(db, col, table)
-            aaronmk
+#### Tables
             aaronmk
-            aaronmk
+### Maintenance
             aaronmk
-            aaronmk
+def analyze(db, table):
     table = sql_gen.as_Table(table)
     run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
             aaronmk
-            aaronmk
+def autoanalyze(db, table):
     if db.autoanalyze: analyze(db, table)
             aaronmk
-            aaronmk
+def vacuum(db, table):
     table = sql_gen.as_Table(table)
     db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
         log_level=3))
             aaronmk
-            aaronmk
+### Lifecycle
-            aaronmk
+def drop(db, type_, name):
     name = sql_gen.as_Name(name)
     run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
             aaronmk
-            aaronmk
+def drop_table(db, table): drop(db, 'TABLE', table)
-            aaronmk
+def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
-            aaronmk
+    like=None, **query_opts):
-            aaronmk
+    '''Creates a table.
-            aaronmk
+    @param cols [sql_gen.TypedCol,...] The column names and types
     @param has_pkey If set, the first column becomes the primary key.
-            aaronmk
+    @param col_indexes bool|[ref]
         * If True, indexes will be added on all non-pkey columns.
         * If a list reference, [0] will be set to a function to do this.
           This can be used to delay index creation until the table is populated.
-            aaronmk
+    '''
-            aaronmk
+    query_opts.setdefault('cacheable', True)
-            aaronmk
+    table = sql_gen.as_Table(table)
-            aaronmk
+    if like != None:
         cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
             ]+cols
-            aaronmk
+        table.order_by = like.order_by
-            aaronmk
+    if has_pkey:
         cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
-            aaronmk
+        pkey.constraints = 'PRIMARY KEY'
             aaronmk
-            aaronmk
+    temp = table.is_temp and not db.debug_temp
         # temp tables permanent in debug_temp mode
             aaronmk
-            aaronmk
+    # Create table
-            aaronmk
+    def create():
-            aaronmk
+        str_ = 'CREATE'
         if temp: str_ += ' TEMP'
         str_ += ' TABLE '+table.to_str(db)+' (\n'
         str_ += '\n, '.join(c.to_str(db) for c in cols)
-            aaronmk
+        str_ += '\n);'
             aaronmk
-            aaronmk
+        opts = dict(recover=True, log_level=2,
             log_ignore_excs=(DuplicateException,), **query_opts)
-            aaronmk
+        try: run_query(db, str_, **opts)
         except InvalidTypeException: # try again as view
             run_query_into(db, mk_select(db, like, limit=0), into=table, **opts)
-            aaronmk
+    if table.is_temp:
         while True:
             try:
                 create()
                 break
             except DuplicateException:
                 table.name = next_version(table.name)
                 # try again with next version of name
     else: create()
             aaronmk
-            aaronmk
+    # Add indexes
-            aaronmk
+    if has_pkey: has_pkey = already_indexed
     def add_indexes_(): add_indexes(db, table, has_pkey)
     if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
     elif col_indexes: add_indexes_() # add now
             aaronmk
-            aaronmk
+def copy_table_struct(db, src, dest, **opts):
-            aaronmk
+    '''Creates a structure-only copy of a table. (Does not copy data.)'''
-            aaronmk
+    create_table(db, dest, has_pkey=False, col_indexes=False, like=src, **opts)
             aaronmk
-            aaronmk
+def copy_table(db, src, dest):
     '''Creates a copy of a table, including data'''
     copy_table_struct(db, src, dest)
     insert_select(db, dest, None, mk_select(db, src))
-            aaronmk
+### Data
             aaronmk
-            aaronmk
+def truncate(db, table, schema='public', **kw_args):
     '''For params, see run_query()'''
-            aaronmk
+    table = sql_gen.as_Table(table, schema)
-            aaronmk
+    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
             aaronmk
-            aaronmk
+def empty_temp(db, tables):
     tables = lists.mk_seq(tables)
-            aaronmk
+    for table in tables: truncate(db, table, log_level=3)
             aaronmk
-            aaronmk
+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)
             aaronmk
-            aaronmk
+def distinct_table(db, table, distinct_on, joins=None, conds=None):
-            aaronmk
+    '''Creates a copy of a temp table which is distinct on the given columns.
-            aaronmk
+    Adds an index on table's distinct_on columns, to facilitate merge joins.
-            aaronmk
+    @param distinct_on If empty, creates a table with one row. This is useful if
         your distinct_on columns are all literal values.
-            aaronmk
+    @param joins The joins to use when creating the new table
-            aaronmk
+    @param conds Any additional filters to use when creating the new table
-            aaronmk
+    @return The new table.
-            aaronmk
+    '''
-            aaronmk
+    if joins == None: joins = [table]
-            aaronmk
+    new_table = sql_gen.suffixed_table(table, '_distinct')
-            aaronmk
+    distinct_on = filter(sql_gen.is_table_col, distinct_on)
             aaronmk
-            aaronmk
+    copy_table_struct(db, table, new_table, cacheable=False)
             aaronmk
     limit = None
     if distinct_on == []: limit = 1 # one sample row
-            aaronmk
+    else: add_index(db, distinct_on, table) # for join optimization
             aaronmk
-            aaronmk
+    insert_select(db, new_table, None, mk_select(db, joins,
-            aaronmk
+        [sql_gen.Col(sql_gen.all_cols, table)], conds, distinct_on,
-            aaronmk
+        order_by=None, limit=limit))
-            aaronmk
+    analyze(db, new_table)
             aaronmk
-            aaronmk
+    return new_table

Project

General

Profile

root/trunk/lib/sql.py @ 13227