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

root/lib/sql.py @ 2294

-            aaronmk
+# Database access
-            aaronmk
+import copy
-            aaronmk
+import operator
-            aaronmk
+import re
-            aaronmk
+import warnings
             aaronmk
-            aaronmk
+import exc
-            aaronmk
+import dicts
-            aaronmk
+import iters
-            aaronmk
+import lists
-            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, input_params=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
     else: return repr(input_query)+' % '+repr(input_params)
             aaronmk
-            aaronmk
+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)))
             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: '+str(name), cause)
-            aaronmk
+        self.name = name
             aaronmk
-            aaronmk
+class ExceptionWithNameValue(DbException):
     def __init__(self, name, value, cause=None):
         DbException.__init__(self,
-            aaronmk
+            'for name: '+str(name)+'; value: '+strings.ustr(value), cause)
-            aaronmk
+        self.name = name
         self.value = value
-            aaronmk
+class ExceptionWithColumns(DbException):
     def __init__(self, cols, cause=None):
-            aaronmk
+        DbException.__init__(self, 'for columns: '+(', '.join(cols)), cause)
-            aaronmk
+        self.cols = cols
             aaronmk
-            aaronmk
+class NameException(DbException): pass
-            aaronmk
+class DuplicateKeyException(ExceptionWithColumns): pass
             aaronmk
-            aaronmk
+class NullValueException(ExceptionWithColumns): pass
             aaronmk
-            aaronmk
+class FunctionValueException(ExceptionWithNameValue): pass
             aaronmk
-            aaronmk
+class DuplicateTableException(ExceptionWithName): pass
-            aaronmk
+class DuplicateFunctionException(ExceptionWithName): 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
+##### Input validation
-            aaronmk
+def clean_name(name): return re.sub(r'\W', r'', name).lower()
             aaronmk
 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':
         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
-            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
+def _query_lookup(query, params): return (query, dicts.make_hashable(params))
             aaronmk
-            aaronmk
+log_debug_none = lambda msg: None
-            aaronmk
+class DbConn:
-            aaronmk
+    def __init__(self, db_config, serializable=True, autocommit=False,
         caching=True, log_debug=log_debug_none):
-            aaronmk
+        self.db_config = db_config
         self.serializable = serializable
-            aaronmk
+        self.autocommit = autocommit
         self.caching = caching
-            aaronmk
+        self.log_debug = log_debug
-            aaronmk
+        self.debug = log_debug != log_debug_none
             aaronmk
         self.__db = None
-            aaronmk
+        self.query_results = {}
-            aaronmk
+        self._savepoint = 0
             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 connected(self): return self.__db != None
-            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)
             # Configure connection
-            aaronmk
+            if self.serializable and not self.autocommit: run_raw_query(self,
                 'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
-            aaronmk
+            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_])
             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, params=None):
-            aaronmk
+            self._is_insert = query.upper().find('INSERT') >= 0
-            aaronmk
+            self.query_lookup = _query_lookup(query, params)
-            aaronmk
+            try:
-            aaronmk
+                try:
                     return_value = self.inner.execute(query, params)
                     self.outer.do_autocommit()
-            aaronmk
+                finally: self.query = get_cur_query(self.inner)
-            aaronmk
+            except Exception, e:
-            aaronmk
+                _add_cursor_info(e, self, query, params)
-            aaronmk
+                self.result = e # cache the exception as the result
                 self._cache_result()
                 raise
-            aaronmk
+            # Fetch all rows so result will be cached
             if self.rowcount == 0 and not self._is_insert: consume_rows(self)
-            aaronmk
+            return return_value
-            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, only cache exceptions since inserts are not
             # idempotent, but an invalid insert 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
+        module = util.root_module(self.db)
         if module == 'psycopg2': return self.db.cursor().mogrify('%s', [value])
-            aaronmk
+        elif module == 'MySQLdb':
             import _mysql
             return _mysql.escape_string(value)
         else: raise NotImplementedError("Can't escape value for "+module
             +' database')
-            aaronmk
+    def run_query(self, query, params=None, cacheable=False):
-            aaronmk
+        '''Translates known DB errors to typed exceptions:
         See self.DbCursor.execute().'''
-            aaronmk
+        assert query != None
-            aaronmk
+        if not self.caching: cacheable = False
-            aaronmk
+        used_cache = False
         try:
-            aaronmk
+            # 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
-            aaronmk
+            cur.execute(query, params)
-            aaronmk
+        finally:
-            aaronmk
+            if self.debug: # only compute msg if needed
-            aaronmk
+                if used_cache: cache_status = 'Cache hit'
                 elif cacheable: cache_status = 'Cache miss'
                 else: cache_status = 'Non-cacheable'
-            aaronmk
+                self.log_debug(cache_status+': '
-            aaronmk
+                    +strings.one_line(str(get_cur_query(cur, query, params))))
             aaronmk
         return cur
             aaronmk
     def is_cached(self, query, params=None):
         return _query_lookup(query, params) in self.query_results
             aaronmk
     def with_savepoint(self, func):
-            aaronmk
+        savepoint = 'level_'+str(self._savepoint)
-            aaronmk
+        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)
-            aaronmk
+            self.do_autocommit()
-            aaronmk
+            return return_val
             aaronmk
     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()
             aaronmk
-            aaronmk
+connect = DbConn
-            aaronmk
+##### Querying
-            aaronmk
+def run_raw_query(db, *args, **kw_args):
-            aaronmk
+    '''For params, see DbConn.run_query()'''
-            aaronmk
+    return db.run_query(*args, **kw_args)
             aaronmk
-            aaronmk
+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')
-            aaronmk
+##### Recoverable querying
             aaronmk
-            aaronmk
+def with_savepoint(db, func): return db.with_savepoint(func)
             aaronmk
-            aaronmk
+def run_query(db, query, params=None, recover=None, cacheable=False):
-            aaronmk
+    if recover == None: recover = False
-            aaronmk
+    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()
-            aaronmk
+        msg = exc.str_(e)
             aaronmk
-            aaronmk
+        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(cols, e)
             aaronmk
-            aaronmk
+        match = re.search(r'null value in column "(\w+)" violates not-null '
-            aaronmk
+            r'constraint', msg)
-            aaronmk
+        if match: raise NullValueException([match.group(1)], e)
             aaronmk
-            aaronmk
+        match = re.search(r'invalid input (?:syntax|value)\b.*: "(.+)"\n'
-            aaronmk
+            r'(?:(?s).*)\bfunction "(\w+)".*\bat assignment', msg)
-            aaronmk
+        if match:
-            aaronmk
+            raise FunctionValueException(match.group(2),
                 strings.to_unicode(match.group(1)), e)
             aaronmk
-            aaronmk
+        match = re.search(r'relation "(\w+)" already exists', msg)
         if match: raise DuplicateTableException(match.group(1), e)
             aaronmk
-            aaronmk
+        match = re.search(r'function "(\w+)" already exists', msg)
         if match: raise DuplicateFunctionException(match.group(1), e)
             aaronmk
-            aaronmk
+        raise # no specific exception raised
             aaronmk
-            aaronmk
+##### Basic queries
-            aaronmk
+def next_version(name):
     '''Prepends the version # so it won't be removed if the name is truncated'''
-            aaronmk
+    version = 1 # first existing name was version 0
-            aaronmk
+    match = re.match(r'^v(\d+)_(.*)$', name)
     if match:
         version = int(match.group(1))+1
         name = match.group(2)
     return 'v'+str(version)+'_'+name
-            aaronmk
+def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
-            aaronmk
+    '''Outputs a query to a temp table.
     For params, see run_query().
     '''
-            aaronmk
+    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
-            aaronmk
+    else: # place rows in temp table
-            aaronmk
+        check_name(into_ref[0])
-            aaronmk
+        kw_args['recover'] = True
         while True:
             try:
-            aaronmk
+                create_query = 'CREATE'
                 if not db.debug: create_query += ' TEMP'
                 create_query += ' TABLE '+into_ref[0]+' AS '+query
                 return run_query(db, create_query, params, *args, **kw_args)
-            aaronmk
+                    # CREATE TABLE AS sets rowcount to # rows in query
             except DuplicateTableException, e:
                 into_ref[0] = next_version(into_ref[0])
                 # try again with next version of name
             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 mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
-            aaronmk
+    start=None, order_by=order_by_pkey, default_table=None):
-            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: dict(compare_left_side=compare_right_side):
         * compare_left_side: Code|str (for col name)
         * compare_right_side: ValueCond|literal value
-            aaronmk
+    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
         use all columns
-            aaronmk
+    @return tuple(query, params)
-            aaronmk
+    '''
-            aaronmk
+    def esc_name_(name): return esc_name(db, name)
             aaronmk
-            aaronmk
+    if not lists.is_seq(tables): tables = [tables]
-            aaronmk
+    tables = list(tables) # don't modify input! (list() copies input)
-            aaronmk
+    table0 = tables.pop(0) # first table is separate
-            aaronmk
+    if conds == None: conds = {}
-            aaronmk
+    assert limit == None or type(limit) == int
-            aaronmk
+    assert start == None or type(start) == int
-            aaronmk
+    table0 = sql_gen.as_Table(table0)
-            aaronmk
+    if order_by is order_by_pkey: order_by = pkey(db, table0, recover=True)
             aaronmk
-            aaronmk
+    params = []
-            aaronmk
+    def parse_col(field):
-            aaronmk
+        '''Parses fields'''
-            aaronmk
+        return sql_gen.as_Col(field, default_table).to_str(db)
             aaronmk
-            aaronmk
+    query = 'SELECT'
     # DISTINCT ON columns
-            aaronmk
+    if distinct_on != []:
-            aaronmk
+        query += ' DISTINCT'
-            aaronmk
+        if distinct_on is not distinct_on_all:
-            aaronmk
+            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
     # Columns
     query += ' '
-            aaronmk
+    if fields == None: query += '*'
-            aaronmk
+    else: query += ', '.join(map(parse_col, fields))
             aaronmk
     # Main table
-            aaronmk
+    query += ' FROM '+table0.to_str(db)
             aaronmk
-            aaronmk
+    # Add joins
-            aaronmk
+    left_table = table0
-            aaronmk
+    for join_ in tables:
-            aaronmk
+        # Parse special values
-            aaronmk
+        table = join_.table
-            aaronmk
+        if join_.type_ is sql_gen.filter_out:
-            aaronmk
+            conds[sql_gen.Col(pkey(db, table), table)] = None # filter no match
             aaronmk
-            aaronmk
+        query += ' '+join_.to_str(db, left_table)
             aaronmk
         left_table = table
-            aaronmk
+    missing = True
-            aaronmk
+    if conds != {}:
-            aaronmk
+        query += ' WHERE '+(' AND '.join((sql_gen.as_ValueCond(r).to_str(db, l)
             for l, r in conds.iteritems())))
-            aaronmk
+        params += conds.values()
-            aaronmk
+        missing = False
-            aaronmk
+    if order_by != None:
-            aaronmk
+        query += ' ORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
-            aaronmk
+    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))
-            aaronmk
+    return (query, params)
             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)
     query, params = mk_select(db, *args, **kw_args)
     return run_query(db, query, params, recover, cacheable)
-            aaronmk
+def mk_insert_select(db, table, cols=None, select_query=None, params=None,
-            aaronmk
+    returning=None, embeddable=False):
-            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
+    '''
-            aaronmk
+    if select_query == None: select_query = 'DEFAULT VALUES'
     if cols == []: cols = None # no cols (all defaults) = unknown col names
-            aaronmk
+    table = sql_gen.as_Table(table).to_str(db)
             aaronmk
     # Build query
     query = 'INSERT INTO '+table
     if cols != None:
         map(check_name, cols)
         query += ' ('+', '.join(cols)+')'
     query += ' '+select_query
     if returning != None:
         check_name(returning)
         query += ' RETURNING '+returning
-            aaronmk
+    if embeddable:
         # Create function
-            aaronmk
+        function_name = '_'.join(map(clean_name, ['insert', table] + cols))
-            aaronmk
+        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
-            aaronmk
+        while True:
             try:
-            aaronmk
+                function = function_name
                 if not db.debug: function = 'pg_temp.'+function
-            aaronmk
+                function_query = '''\
 CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
-            aaronmk
+    LANGUAGE sql
     AS $$'''+mogrify(db, query, params)+''';$$;
 '''
-            aaronmk
+                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
             aaronmk
         # Return query that uses function
-            aaronmk
+        return mk_select(db, sql_gen.CustomCode(function+'() AS f ('+returning
             +')'), start=0, order_by=None) # AS clause requires function alias
             aaronmk
-            aaronmk
+    return (query, params)
 def insert_select(db, *args, **kw_args):
-            aaronmk
+    '''For params, see mk_insert_select() and run_query_into()
-            aaronmk
+    @param into_ref List with name of temp table to place RETURNING values in
-            aaronmk
+    '''
-            aaronmk
+    into_ref = kw_args.pop('into_ref', None)
     if into_ref != None: kw_args['embeddable'] = True
-            aaronmk
+    recover = kw_args.pop('recover', None)
     cacheable = kw_args.pop('cacheable', True)
     query, params = mk_insert_select(db, *args, **kw_args)
-            aaronmk
+    return run_query_into(db, query, params, into_ref, recover=recover,
         cacheable=cacheable)
             aaronmk
-            aaronmk
+default = object() # tells insert() to use the default value for a column
-            aaronmk
+def insert(db, table, row, *args, **kw_args):
-            aaronmk
+    '''For params, see insert_select()'''
-            aaronmk
+    if lists.is_seq(row): cols = None
     else:
         cols = row.keys()
         row = row.values()
     row = list(row) # ensure that "!= []" works
-            aaronmk
+    # Check for special values
     labels = []
     values = []
     for value in row:
-            aaronmk
+        if value is default: labels.append('DEFAULT')
-            aaronmk
+        else:
             labels.append('%s')
             values.append(value)
     # Build query
-            aaronmk
+    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
     else: query = None
             aaronmk
-            aaronmk
+    return insert_select(db, table, cols, query, values, *args, **kw_args)
             aaronmk
-            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 truncate(db, table, schema='public'):
     return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
             aaronmk
 ##### Database structure queries
-            aaronmk
+def pkey(db, table, recover=None):
-            aaronmk
+    '''Assumed to be first column in table'''
-            aaronmk
+    return col_names(select(db, table, limit=0, order_by=None,
         recover=recover)).next()
             aaronmk
-            aaronmk
+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.'''
     check_name(table)
     check_name(index)
-            aaronmk
+    module = util.root_module(db.db)
     if module == 'psycopg2':
         return list(values(run_query(db, '''\
-            aaronmk
+SELECT attname
-            aaronmk
+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
-            aaronmk
+ORDER BY attnum
 ''',
-            aaronmk
+            {'table': table, 'index': index}, cacheable=True)))
     else: raise NotImplementedError("Can't list index columns for "+module+
         ' database')
             aaronmk
-            aaronmk
+def constraint_cols(db, table, constraint):
     check_name(table)
     check_name(constraint)
-            aaronmk
+    module = util.root_module(db.db)
-            aaronmk
+    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')
-            aaronmk
+row_num_col = '_row_num'
-            aaronmk
+def add_row_num(db, table):
-            aaronmk
+    '''Adds a row number column to a table. Its name is in row_num_col. It will
     be the primary key.'''
-            aaronmk
+    check_name(table)
-            aaronmk
+    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
-            aaronmk
+        +' serial NOT NULL PRIMARY KEY')
             aaronmk
-            aaronmk
+def tables(db, schema='public', table_like='%'):
-            aaronmk
+    module = util.root_module(db.db)
-            aaronmk
+    params = {'schema': schema, 'table_like': table_like}
-            aaronmk
+    if module == 'psycopg2':
-            aaronmk
+        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))
-            aaronmk
+    else: raise NotImplementedError("Can't list tables for "+module+' database')
             aaronmk
-            aaronmk
+##### Database management
-            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
+##### Heuristic queries
-            aaronmk
+def put(db, table, row, pkey_=None, row_ct_ref=None):
-            aaronmk
+    '''Recovers from errors.
-            aaronmk
+    Only works under PostgreSQL (uses INSERT RETURNING).
     '''
-            aaronmk
+    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
-            aaronmk
+    try:
-            aaronmk
+        cur = insert(db, table, row, pkey_, recover=True)
-            aaronmk
+        if row_ct_ref != None and cur.rowcount >= 0:
             row_ct_ref[0] += cur.rowcount
         return value(cur)
-            aaronmk
+    except DuplicateKeyException, e:
-            aaronmk
+        return value(select(db, table, [pkey_],
-            aaronmk
+            util.dict_subset_right_join(row, e.cols), recover=True))
             aaronmk
-            aaronmk
+def get(db, table, row, pkey, row_ct_ref=None, create=False):
-            aaronmk
+    '''Recovers from errors'''
-            aaronmk
+    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
-            aaronmk
+    except StopIteration:
-            aaronmk
+        if not create: raise
-            aaronmk
+        return put(db, table, row, pkey, row_ct_ref) # insert new row
             aaronmk
-            aaronmk
+def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
-            aaronmk
+    row_ct_ref=None):
-            aaronmk
+    '''Recovers from errors.
     Only works under PostgreSQL (uses INSERT RETURNING).
-            aaronmk
+    @param in_tables The main input table to select from, followed by a list of
         tables to join with it using the main input table's pkey
-            aaronmk
+    @return sql_gen.Col Where the pkeys (from INSERT RETURNING) are made
-            aaronmk
+        available
     '''
-            aaronmk
+    temp_prefix = clean_name(out_table)
     pkeys_ref = [temp_prefix+'_pkeys']
             aaronmk
-            aaronmk
+    # Join together input tables
-            aaronmk
+    in_tables = in_tables[:] # don't modify input!
     in_tables0 = in_tables.pop(0) # first table is separate
-            aaronmk
+    in_pkey = pkey(db, in_tables0, recover=True)
-            aaronmk
+    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
-            aaronmk
+    insert_joins = [in_tables0]+[sql_gen.Join(v, {in_pkey: sql_gen.join_using})
         for v in in_tables]
             aaronmk
-            aaronmk
+    out_pkey = pkey(db, out_table, recover=True)
-            aaronmk
+    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
             aaronmk
-            aaronmk
+    pkeys = [in_pkey, out_pkey]
     pkeys_cols = [in_pkey_col, out_pkey_col]
-            aaronmk
+    pkeys_table_exists_ref = [False]
     def run_query_into_pkeys(query, params):
         if pkeys_table_exists_ref[0]:
-            aaronmk
+            insert_select(db, pkeys_ref[0], pkeys, query, params)
-            aaronmk
+        else:
             run_query_into(db, query, params, into_ref=pkeys_ref)
             pkeys_table_exists_ref[0] = True
-            aaronmk
+    conds = {}
-            aaronmk
+    distinct_on = []
-            aaronmk
+    def mk_main_select(cols):
-            aaronmk
+        return mk_select(db, insert_joins, cols, conds, distinct_on,
-            aaronmk
+            order_by=None, limit=limit, start=start, default_table=in_tables0)
             aaronmk
-            aaronmk
+    def log_ignore(in_col, value):
         db.log_debug('Ignoring rows with '+in_col+' = '+value)
-            aaronmk
+    # Do inserts and selects
-            aaronmk
+    join_cols = {}
     filter_join_added = False
-            aaronmk
+    out_pkeys_ref = [temp_prefix+'_out_pkeys']
-            aaronmk
+    in_pkeys_ref = [temp_prefix+'_in_pkeys']
-            aaronmk
+    while True:
         try:
-            aaronmk
+            if join_cols != {}:
-            aaronmk
+                select_joins = insert_joins+[sql_gen.Join(out_table, join_cols)]
-            aaronmk
+                db.log_debug('Getting pkeys of already existing rows')
                 run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
-            aaronmk
+                    order_by=None, start=0))
             aaronmk
                 # Prepare to insert new rows
                 input_join_cols = join_cols.values()
                 distinct_on = filter(util.is_str, input_join_cols)
                 if not filter_join_added:
-            aaronmk
+                    insert_joins.append(sql_gen.Join(pkeys_ref[0],
                         {in_pkey: in_pkey}, sql_gen.filter_out))
-            aaronmk
+                    filter_join_added = True
-            aaronmk
+            db.log_debug('Inserting new rows')
-            aaronmk
+            cur = insert_select(db, out_table, mapping.keys(),
-            aaronmk
+                *mk_main_select(mapping.values()), returning=out_pkey,
-            aaronmk
+                into_ref=out_pkeys_ref, recover=True)
-            aaronmk
+            if row_ct_ref != None and cur.rowcount >= 0:
                 row_ct_ref[0] += cur.rowcount
                 add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
-            aaronmk
+            db.log_debug('Getting input pkeys corresponding to rows in insert')
-            aaronmk
+            run_query_into(db, *mk_main_select([in_pkey]),
                 into_ref=in_pkeys_ref)
-            aaronmk
+            add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
-            aaronmk
+            db.log_debug('Joining together output and input pkeys')
-            aaronmk
+            run_query_into_pkeys(*mk_select(db, [in_pkeys_ref[0],
-            aaronmk
+                sql_gen.Join(out_pkeys_ref[0],
                 {row_num_col: sql_gen.join_using})], pkeys, start=0))
             aaronmk
             break # insert successful
         except DuplicateKeyException, e:
-            aaronmk
+            old_join_cols = join_cols.copy()
-            aaronmk
+            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
-            aaronmk
+            db.log_debug('Ignoring existing rows, comparing on '+str(join_cols))
-            aaronmk
+            assert join_cols != old_join_cols # avoid infinite loops
-            aaronmk
+            # rerun loop with additional constraints
-            aaronmk
+        except NullValueException, e:
             out_col, = e.cols
             try: in_col = mapping[out_col]
             except KeyError: # no mapping for missing col, so every row invalid
-            aaronmk
+                db.log_debug('Missing mapping for '+out_col)
-            aaronmk
+                run_query_into_pkeys(*mk_select(db, insert_joins,
-            aaronmk
+                    [in_pkey_col, sql_gen.NamedCode(out_pkey, None)],
-            aaronmk
+                    order_by=None, start=0))
-            aaronmk
+                break
             else:
-            aaronmk
+                log_ignore(in_col, 'NULL')
-            aaronmk
+                conds[in_col] = sql_gen.CompareCond(None, '!=')
                 # rerun loop with additional constraints
-            aaronmk
+        except FunctionValueException, e:
             assert e.name == sql_gen.unescape_table(out_table)
             out_col = 'value' # assume function param was named "value"
             in_col = mapping[out_col]
-            aaronmk
+            value = e.value
             log_ignore(in_col, value)
             conds[in_col] = sql_gen.CompareCond(value, '!=')
-            aaronmk
+            # rerun loop with additional constraints
             aaronmk
-            aaronmk
+    return sql_gen.Col(out_pkey, pkeys_ref[0])
             aaronmk
 ##### Data cleanup
-            aaronmk
+def cleanup_table(db, table, cols):
-            aaronmk
+    def esc_name_(name): return esc_name(db, name)
-            aaronmk
+    table = sql_gen.as_Table(table).to_str(db)
-            aaronmk
+    cols = map(esc_name_, cols)
     run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
         +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
             for col in cols))),
         dict(null0='', null1=r'\N'))

Project

General

Profile