Project

General

Profile

1 11 aaronmk
# Database access
2
3 1869 aaronmk
import copy
4 2127 aaronmk
import operator
5 11 aaronmk
import re
6 865 aaronmk
import warnings
7 11 aaronmk
8 300 aaronmk
import exc
9 1909 aaronmk
import dicts
10 1893 aaronmk
import iters
11 1960 aaronmk
import lists
12 1889 aaronmk
from Proxy import Proxy
13 1872 aaronmk
import rand
14 2217 aaronmk
import sql_gen
15 862 aaronmk
import strings
16 131 aaronmk
import util
17 11 aaronmk
18 832 aaronmk
##### Exceptions
19
20 2804 aaronmk
def get_cur_query(cur, input_query=None):
21 2168 aaronmk
    raw_query = None
22
    if hasattr(cur, 'query'): raw_query = cur.query
23
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
24 2170 aaronmk
25
    if raw_query != None: return raw_query
26 2804 aaronmk
    else: return '[input] '+strings.ustr(input_query)
27 14 aaronmk
28 2170 aaronmk
def _add_cursor_info(e, *args, **kw_args):
29
    '''For params, see get_cur_query()'''
30 2771 aaronmk
    exc.add_msg(e, 'query: '+strings.ustr(get_cur_query(*args, **kw_args)))
31 135 aaronmk
32 300 aaronmk
class DbException(exc.ExceptionWithCause):
33 14 aaronmk
    def __init__(self, msg, cause=None, cur=None):
34 2145 aaronmk
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
35 14 aaronmk
        if cur != None: _add_cursor_info(self, cur)
36
37 2143 aaronmk
class ExceptionWithName(DbException):
38
    def __init__(self, name, cause=None):
39 2484 aaronmk
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name)), cause)
40 2143 aaronmk
        self.name = name
41 360 aaronmk
42 2240 aaronmk
class ExceptionWithNameValue(DbException):
43
    def __init__(self, name, value, cause=None):
44 2484 aaronmk
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name))
45
            +'; value: '+strings.as_tt(repr(value)), cause)
46 2240 aaronmk
        self.name = name
47
        self.value = value
48
49 2945 aaronmk
class ExceptionWithNameType(DbException):
50
    def __init__(self, type_, name, cause=None):
51
        DbException.__init__(self, 'for type: '+strings.as_tt(str(type_))
52
            +'; name: '+strings.as_tt(name), cause)
53
        self.type = type_
54
        self.name = name
55
56 2306 aaronmk
class ConstraintException(DbException):
57
    def __init__(self, name, cols, cause=None):
58 2484 aaronmk
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
59
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
60 2306 aaronmk
        self.name = name
61 468 aaronmk
        self.cols = cols
62 11 aaronmk
63 2523 aaronmk
class MissingCastException(DbException):
64
    def __init__(self, type_, col, cause=None):
65
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
66
            +' on column: '+strings.as_tt(col), cause)
67
        self.type = type_
68
        self.col = col
69
70 2143 aaronmk
class NameException(DbException): pass
71
72 2306 aaronmk
class DuplicateKeyException(ConstraintException): pass
73 13 aaronmk
74 2306 aaronmk
class NullValueException(ConstraintException): pass
75 13 aaronmk
76 2240 aaronmk
class FunctionValueException(ExceptionWithNameValue): pass
77 2239 aaronmk
78 2945 aaronmk
class DuplicateException(ExceptionWithNameType): pass
79 2143 aaronmk
80 89 aaronmk
class EmptyRowException(DbException): pass
81
82 865 aaronmk
##### Warnings
83
84
class DbWarning(UserWarning): pass
85
86 1930 aaronmk
##### Result retrieval
87
88
def col_names(cur): return (col[0] for col in cur.description)
89
90
def rows(cur): return iter(lambda: cur.fetchone(), None)
91
92
def consume_rows(cur):
93
    '''Used to fetch all rows so result will be cached'''
94
    iters.consume_iter(rows(cur))
95
96
def next_row(cur): return rows(cur).next()
97
98
def row(cur):
99
    row_ = next_row(cur)
100
    consume_rows(cur)
101
    return row_
102
103
def next_value(cur): return next_row(cur)[0]
104
105
def value(cur): return row(cur)[0]
106
107
def values(cur): return iters.func_iter(lambda: next_value(cur))
108
109
def value_or_none(cur):
110
    try: return value(cur)
111
    except StopIteration: return None
112
113 2762 aaronmk
##### Escaping
114 2101 aaronmk
115 2573 aaronmk
def esc_name_by_module(module, name):
116
    if module == 'psycopg2' or module == None: quote = '"'
117 2101 aaronmk
    elif module == 'MySQLdb': quote = '`'
118
    else: raise NotImplementedError("Can't escape name for "+module+' database')
119 2500 aaronmk
    return sql_gen.esc_name(name, quote)
120 2101 aaronmk
121
def esc_name_by_engine(engine, name, **kw_args):
122
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
123
124
def esc_name(db, name, **kw_args):
125
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
126
127
def qual_name(db, schema, table):
128
    def esc_name_(name): return esc_name(db, name)
129
    table = esc_name_(table)
130
    if schema != None: return esc_name_(schema)+'.'+table
131
    else: return table
132
133 1869 aaronmk
##### Database connections
134 1849 aaronmk
135 2097 aaronmk
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
136 1926 aaronmk
137 1869 aaronmk
db_engines = {
138
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
139
    'PostgreSQL': ('psycopg2', {}),
140
}
141
142
DatabaseErrors_set = set([DbException])
143
DatabaseErrors = tuple(DatabaseErrors_set)
144
145
def _add_module(module):
146
    DatabaseErrors_set.add(module.DatabaseError)
147
    global DatabaseErrors
148
    DatabaseErrors = tuple(DatabaseErrors_set)
149
150
def db_config_str(db_config):
151
    return db_config['engine']+' database '+db_config['database']
152
153 2448 aaronmk
log_debug_none = lambda msg, level=2: None
154 1901 aaronmk
155 1849 aaronmk
class DbConn:
156 2923 aaronmk
    def __init__(self, db_config, autocommit=True, caching=True,
157 2915 aaronmk
        log_debug=log_debug_none, debug_temp=False):
158
        '''
159
        @param debug_temp Whether temporary objects should instead be permanent.
160
            This assists in debugging the internal objects used by the program.
161
        '''
162 1869 aaronmk
        self.db_config = db_config
163 2190 aaronmk
        self.autocommit = autocommit
164
        self.caching = caching
165 1901 aaronmk
        self.log_debug = log_debug
166 2193 aaronmk
        self.debug = log_debug != log_debug_none
167 2915 aaronmk
        self.debug_temp = debug_temp
168 1869 aaronmk
169
        self.__db = None
170 1889 aaronmk
        self.query_results = {}
171 2139 aaronmk
        self._savepoint = 0
172 2671 aaronmk
        self._notices_seen = set()
173 1869 aaronmk
174
    def __getattr__(self, name):
175
        if name == '__dict__': raise Exception('getting __dict__')
176
        if name == 'db': return self._db()
177
        else: raise AttributeError()
178
179
    def __getstate__(self):
180
        state = copy.copy(self.__dict__) # shallow copy
181 1915 aaronmk
        state['log_debug'] = None # don't pickle the debug callback
182 1869 aaronmk
        state['_DbConn__db'] = None # don't pickle the connection
183
        return state
184
185 2165 aaronmk
    def connected(self): return self.__db != None
186
187 1869 aaronmk
    def _db(self):
188
        if self.__db == None:
189
            # Process db_config
190
            db_config = self.db_config.copy() # don't modify input!
191 2097 aaronmk
            schemas = db_config.pop('schemas', None)
192 1869 aaronmk
            module_name, mappings = db_engines[db_config.pop('engine')]
193
            module = __import__(module_name)
194
            _add_module(module)
195
            for orig, new in mappings.iteritems():
196
                try: util.rename_key(db_config, orig, new)
197
                except KeyError: pass
198
199
            # Connect
200
            self.__db = module.connect(**db_config)
201
202
            # Configure connection
203 2906 aaronmk
            if hasattr(self.db, 'set_isolation_level'):
204
                import psycopg2.extensions
205
                self.db.set_isolation_level(
206
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
207 2101 aaronmk
            if schemas != None:
208 2893 aaronmk
                search_path = [self.esc_name(s) for s in schemas.split(',')]
209
                search_path.append(value(run_query(self, 'SHOW search_path',
210
                    log_level=4)))
211
                run_query(self, 'SET search_path TO '+(','.join(search_path)),
212
                    log_level=3)
213 1869 aaronmk
214
        return self.__db
215 1889 aaronmk
216 1891 aaronmk
    class DbCursor(Proxy):
217 1927 aaronmk
        def __init__(self, outer):
218 1891 aaronmk
            Proxy.__init__(self, outer.db.cursor())
219 2191 aaronmk
            self.outer = outer
220 1927 aaronmk
            self.query_results = outer.query_results
221 1894 aaronmk
            self.query_lookup = None
222 1891 aaronmk
            self.result = []
223 1889 aaronmk
224 2802 aaronmk
        def execute(self, query):
225 2764 aaronmk
            self._is_insert = query.startswith('INSERT')
226 2797 aaronmk
            self.query_lookup = query
227 2148 aaronmk
            try:
228 2191 aaronmk
                try:
229 2802 aaronmk
                    cur = self.inner.execute(query)
230 2191 aaronmk
                    self.outer.do_autocommit()
231 2802 aaronmk
                finally: self.query = get_cur_query(self.inner, query)
232 1904 aaronmk
            except Exception, e:
233 2802 aaronmk
                _add_cursor_info(e, self, query)
234 1904 aaronmk
                self.result = e # cache the exception as the result
235
                self._cache_result()
236
                raise
237 2800 aaronmk
            if self.rowcount == 0 and query.startswith('SELECT'): # empty SELECT
238
                consume_rows(self) # fetch all rows so result will be cached
239 2762 aaronmk
            return cur
240 1894 aaronmk
241 1891 aaronmk
        def fetchone(self):
242
            row = self.inner.fetchone()
243 1899 aaronmk
            if row != None: self.result.append(row)
244
            # otherwise, fetched all rows
245 1904 aaronmk
            else: self._cache_result()
246
            return row
247
248
        def _cache_result(self):
249 2948 aaronmk
            # For inserts that return a result set, don't cache result set since
250
            # inserts are not idempotent. Other non-SELECT queries don't have
251
            # their result set read, so only exceptions will be cached (an
252
            # invalid query will always be invalid).
253 1930 aaronmk
            if self.query_results != None and (not self._is_insert
254 1906 aaronmk
                or isinstance(self.result, Exception)):
255
256 1894 aaronmk
                assert self.query_lookup != None
257 1916 aaronmk
                self.query_results[self.query_lookup] = self.CacheCursor(
258
                    util.dict_subset(dicts.AttrsDictView(self),
259
                    ['query', 'result', 'rowcount', 'description']))
260 1906 aaronmk
261 1916 aaronmk
        class CacheCursor:
262
            def __init__(self, cached_result): self.__dict__ = cached_result
263
264 1927 aaronmk
            def execute(self, *args, **kw_args):
265 1916 aaronmk
                if isinstance(self.result, Exception): raise self.result
266
                # otherwise, result is a rows list
267
                self.iter = iter(self.result)
268
269
            def fetchone(self):
270
                try: return self.iter.next()
271
                except StopIteration: return None
272 1891 aaronmk
273 2212 aaronmk
    def esc_value(self, value):
274 2663 aaronmk
        try: str_ = self.mogrify('%s', [value])
275
        except NotImplementedError, e:
276
            module = util.root_module(self.db)
277
            if module == 'MySQLdb':
278
                import _mysql
279
                str_ = _mysql.escape_string(value)
280
            else: raise e
281 2374 aaronmk
        return strings.to_unicode(str_)
282 2212 aaronmk
283 2347 aaronmk
    def esc_name(self, name): return esc_name(self, name) # calls global func
284
285 2814 aaronmk
    def std_code(self, str_):
286
        '''Standardizes SQL code.
287
        * Ensures that string literals are prefixed by `E`
288
        '''
289
        if str_.startswith("'"): str_ = 'E'+str_
290
        return str_
291
292 2665 aaronmk
    def can_mogrify(self):
293 2663 aaronmk
        module = util.root_module(self.db)
294 2665 aaronmk
        return module == 'psycopg2'
295 2663 aaronmk
296 2665 aaronmk
    def mogrify(self, query, params=None):
297
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
298
        else: raise NotImplementedError("Can't mogrify query")
299
300 2671 aaronmk
    def print_notices(self):
301 2725 aaronmk
        if hasattr(self.db, 'notices'):
302
            for msg in self.db.notices:
303
                if msg not in self._notices_seen:
304
                    self._notices_seen.add(msg)
305
                    self.log_debug(msg, level=2)
306 2671 aaronmk
307 2793 aaronmk
    def run_query(self, query, cacheable=False, log_level=2,
308 2464 aaronmk
        debug_msg_ref=None):
309 2445 aaronmk
        '''
310 2464 aaronmk
        @param log_ignore_excs The log_level will be increased by 2 if the query
311
            throws one of these exceptions.
312 2664 aaronmk
        @param debug_msg_ref If specified, the log message will be returned in
313
            this instead of being output. This allows you to filter log messages
314
            depending on the result of the query.
315 2445 aaronmk
        '''
316 2167 aaronmk
        assert query != None
317
318 2047 aaronmk
        if not self.caching: cacheable = False
319 1903 aaronmk
        used_cache = False
320 2664 aaronmk
321
        def log_msg(query):
322
            if used_cache: cache_status = 'cache hit'
323
            elif cacheable: cache_status = 'cache miss'
324
            else: cache_status = 'non-cacheable'
325
            return 'DB query: '+cache_status+':\n'+strings.as_code(query, 'SQL')
326
327 1903 aaronmk
        try:
328 1927 aaronmk
            # Get cursor
329
            if cacheable:
330
                try:
331 2797 aaronmk
                    cur = self.query_results[query]
332 1927 aaronmk
                    used_cache = True
333
                except KeyError: cur = self.DbCursor(self)
334
            else: cur = self.db.cursor()
335
336 2664 aaronmk
            # Log query
337
            if self.debug and debug_msg_ref == None: # log before running
338
                self.log_debug(log_msg(query), log_level)
339
340 1927 aaronmk
            # Run query
341 2793 aaronmk
            cur.execute(query)
342 1903 aaronmk
        finally:
343 2671 aaronmk
            self.print_notices()
344 2664 aaronmk
            if self.debug and debug_msg_ref != None: # return after running
345 2793 aaronmk
                debug_msg_ref[0] = log_msg(str(get_cur_query(cur, query)))
346 1903 aaronmk
347
        return cur
348 1914 aaronmk
349 2797 aaronmk
    def is_cached(self, query): return query in self.query_results
350 2139 aaronmk
351 2907 aaronmk
    def with_autocommit(self, func):
352 2801 aaronmk
        import psycopg2.extensions
353
354
        prev_isolation_level = self.db.isolation_level
355 2907 aaronmk
        self.db.set_isolation_level(
356
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
357 2683 aaronmk
        try: return func()
358 2801 aaronmk
        finally: self.db.set_isolation_level(prev_isolation_level)
359 2683 aaronmk
360 2139 aaronmk
    def with_savepoint(self, func):
361 2171 aaronmk
        savepoint = 'level_'+str(self._savepoint)
362 2443 aaronmk
        self.run_query('SAVEPOINT '+savepoint, log_level=4)
363 2139 aaronmk
        self._savepoint += 1
364 2930 aaronmk
        try: return func()
365 2139 aaronmk
        except:
366 2443 aaronmk
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint, log_level=4)
367 2139 aaronmk
            raise
368 2930 aaronmk
        finally:
369
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
370
            # "The savepoint remains valid and can be rolled back to again"
371
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
372 2443 aaronmk
            self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
373 2930 aaronmk
374
            self._savepoint -= 1
375
            assert self._savepoint >= 0
376
377
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
378 2191 aaronmk
379
    def do_autocommit(self):
380
        '''Autocommits if outside savepoint'''
381
        assert self._savepoint >= 0
382
        if self.autocommit and self._savepoint == 0:
383 2924 aaronmk
            self.log_debug('Autocommitting', level=4)
384 2191 aaronmk
            self.db.commit()
385 2643 aaronmk
386 2819 aaronmk
    def col_info(self, col):
387 2643 aaronmk
        table = sql_gen.Table('columns', 'information_schema')
388 2821 aaronmk
        cols = ['data_type', 'column_default',
389
            cast(self, 'boolean', 'is_nullable')]
390 2643 aaronmk
391
        conds = [('table_name', col.table.name), ('column_name', col.name)]
392
        schema = col.table.schema
393
        if schema != None: conds.append(('table_schema', schema))
394
395 2819 aaronmk
        type_, default, nullable = row(select(self, table, cols, conds,
396 2874 aaronmk
            order_by='table_schema', limit=1, log_level=4))
397 2643 aaronmk
            # TODO: order_by search_path schema order
398 2819 aaronmk
        default = sql_gen.as_Code(default, self)
399
400
        return sql_gen.TypedCol(col.name, type_, default, nullable)
401 2917 aaronmk
402
    def TempFunction(self, name):
403
        if self.debug_temp: schema = None
404
        else: schema = 'pg_temp'
405
        return sql_gen.Function(name, schema)
406 1849 aaronmk
407 1869 aaronmk
connect = DbConn
408
409 832 aaronmk
##### Recoverable querying
410 15 aaronmk
411 2139 aaronmk
def with_savepoint(db, func): return db.with_savepoint(func)
412 11 aaronmk
413 2791 aaronmk
def run_query(db, query, recover=None, cacheable=False, log_level=2,
414
    log_ignore_excs=None, **kw_args):
415 2794 aaronmk
    '''For params, see DbConn.run_query()'''
416 830 aaronmk
    if recover == None: recover = False
417 2464 aaronmk
    if log_ignore_excs == None: log_ignore_excs = ()
418
    log_ignore_excs = tuple(log_ignore_excs)
419 830 aaronmk
420 2666 aaronmk
    debug_msg_ref = None # usually, db.run_query() logs query before running it
421
    # But if filtering with log_ignore_excs, wait until after exception parsing
422
    if log_ignore_excs != () or not db.can_mogrify(): debug_msg_ref = [None]
423
424 2148 aaronmk
    try:
425 2464 aaronmk
        try:
426 2794 aaronmk
            def run(): return db.run_query(query, cacheable, log_level,
427 2793 aaronmk
                debug_msg_ref, **kw_args)
428 2796 aaronmk
            if recover and not db.is_cached(query):
429 2464 aaronmk
                return with_savepoint(db, run)
430
            else: return run() # don't need savepoint if cached
431
        except Exception, e:
432
            if not recover: raise # need savepoint to run index_cols()
433
            msg = exc.str_(e)
434
435
            match = re.search(r'duplicate key value violates unique constraint '
436 2493 aaronmk
                r'"((_?[^\W_]+)_.+?)"', msg)
437 2464 aaronmk
            if match:
438
                constraint, table = match.groups()
439
                try: cols = index_cols(db, table, constraint)
440
                except NotImplementedError: raise e
441
                else: raise DuplicateKeyException(constraint, cols, e)
442
443 2493 aaronmk
            match = re.search(r'null value in column "(.+?)" violates not-null'
444 2464 aaronmk
                r' constraint', msg)
445
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
446
447
            match = re.search(r'\b(?:invalid input (?:syntax|value)\b.*?'
448
                r'|date/time field value out of range): "(.+?)"\n'
449 2535 aaronmk
                r'(?:(?s).*?)\bfunction "(.+?)"', msg)
450 2464 aaronmk
            if match:
451
                value, name = match.groups()
452
                raise FunctionValueException(name, strings.to_unicode(value), e)
453
454 2526 aaronmk
            match = re.search(r'column "(.+?)" is of type (.+?) but expression '
455 2523 aaronmk
                r'is of type', msg)
456
            if match:
457
                col, type_ = match.groups()
458
                raise MissingCastException(type_, col, e)
459
460 2946 aaronmk
            match = re.search(r'\b(\S+) "(.+?)".*? already exists', msg)
461 2945 aaronmk
            if match:
462
                type_, name = match.groups()
463
                raise DuplicateException(type_, name, e)
464 2464 aaronmk
465
            raise # no specific exception raised
466
    except log_ignore_excs:
467
        log_level += 2
468
        raise
469
    finally:
470 2666 aaronmk
        if debug_msg_ref != None and debug_msg_ref[0] != None:
471
            db.log_debug(debug_msg_ref[0], log_level)
472 830 aaronmk
473 832 aaronmk
##### Basic queries
474
475 2153 aaronmk
def next_version(name):
476 2163 aaronmk
    version = 1 # first existing name was version 0
477 2586 aaronmk
    match = re.match(r'^(.*)#(\d+)$', name)
478 2153 aaronmk
    if match:
479 2586 aaronmk
        name, version = match.groups()
480
        version = int(version)+1
481 2932 aaronmk
    return sql_gen.concat(name, '#'+str(version))
482 2153 aaronmk
483 2899 aaronmk
def lock_table(db, table, mode):
484
    table = sql_gen.as_Table(table)
485
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
486
487 2789 aaronmk
def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):
488 2085 aaronmk
    '''Outputs a query to a temp table.
489
    For params, see run_query().
490
    '''
491 2789 aaronmk
    if into == None: return run_query(db, query, **kw_args)
492 2790 aaronmk
493
    assert isinstance(into, sql_gen.Table)
494
495
    kw_args['recover'] = True
496 2945 aaronmk
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
497 2790 aaronmk
498 2916 aaronmk
    temp = not db.debug_temp # tables are permanent in debug_temp mode
499 2790 aaronmk
    # "temporary tables cannot specify a schema name", so remove schema
500
    if temp: into.schema = None
501
502
    # Create table
503
    while True:
504
        create_query = 'CREATE'
505
        if temp: create_query += ' TEMP'
506
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
507 2385 aaronmk
508 2790 aaronmk
        try:
509
            cur = run_query(db, create_query, **kw_args)
510
                # CREATE TABLE AS sets rowcount to # rows in query
511
            break
512 2945 aaronmk
        except DuplicateException, e:
513 2790 aaronmk
            into.name = next_version(into.name)
514
            # try again with next version of name
515
516
    if add_indexes_: add_indexes(db, into)
517
518
    return cur
519 2085 aaronmk
520 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
521
522 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
523
524 2233 aaronmk
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
525 2293 aaronmk
    start=None, order_by=order_by_pkey, default_table=None):
526 1981 aaronmk
    '''
527 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
528 2280 aaronmk
        together, with tables after the first being sql_gen.Join objects
529 1981 aaronmk
    @param fields Use None to select all fields in the table
530 2377 aaronmk
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
531 2379 aaronmk
        * container can be any iterable type
532 2399 aaronmk
        * compare_left_side: sql_gen.Code|str (for col name)
533
        * compare_right_side: sql_gen.ValueCond|literal value
534 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
535
        use all columns
536 2786 aaronmk
    @return query
537 1981 aaronmk
    '''
538 2315 aaronmk
    # Parse tables param
539 2121 aaronmk
    if not lists.is_seq(tables): tables = [tables]
540 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
541 2315 aaronmk
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
542 2121 aaronmk
543 2315 aaronmk
    # Parse other params
544 2376 aaronmk
    if conds == None: conds = []
545 2650 aaronmk
    elif dicts.is_dict(conds): conds = conds.items()
546 2379 aaronmk
    conds = list(conds) # don't modify input! (list() copies input)
547 135 aaronmk
    assert limit == None or type(limit) == int
548 865 aaronmk
    assert start == None or type(start) == int
549 2315 aaronmk
    if order_by is order_by_pkey:
550
        if distinct_on != []: order_by = None
551
        else: order_by = pkey(db, table0, recover=True)
552 865 aaronmk
553 2315 aaronmk
    query = 'SELECT'
554 2056 aaronmk
555 2315 aaronmk
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
556 2056 aaronmk
557 2200 aaronmk
    # DISTINCT ON columns
558 2233 aaronmk
    if distinct_on != []:
559 2467 aaronmk
        query += '\nDISTINCT'
560 2254 aaronmk
        if distinct_on is not distinct_on_all:
561 2200 aaronmk
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
562
563
    # Columns
564 2467 aaronmk
    query += '\n'
565 1135 aaronmk
    if fields == None: query += '*'
566 2765 aaronmk
    else:
567
        assert fields != []
568
        query += '\n, '.join(map(parse_col, fields))
569 2200 aaronmk
570
    # Main table
571 2467 aaronmk
    query += '\nFROM '+table0.to_str(db)
572 865 aaronmk
573 2122 aaronmk
    # Add joins
574 2271 aaronmk
    left_table = table0
575 2263 aaronmk
    for join_ in tables:
576
        table = join_.table
577 2238 aaronmk
578 2343 aaronmk
        # Parse special values
579
        if join_.type_ is sql_gen.filter_out: # filter no match
580 2376 aaronmk
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
581 2853 aaronmk
                sql_gen.CompareCond(None, '~=')))
582 2343 aaronmk
583 2467 aaronmk
        query += '\n'+join_.to_str(db, left_table)
584 2122 aaronmk
585
        left_table = table
586
587 865 aaronmk
    missing = True
588 2376 aaronmk
    if conds != []:
589 2576 aaronmk
        if len(conds) == 1: whitespace = ' '
590
        else: whitespace = '\n'
591 2578 aaronmk
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
592
            .to_str(db) for l, r in conds], 'WHERE')
593 865 aaronmk
        missing = False
594 2227 aaronmk
    if order_by != None:
595 2467 aaronmk
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
596
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
597 865 aaronmk
    if start != None:
598 2467 aaronmk
        if start != 0: query += '\nOFFSET '+str(start)
599 865 aaronmk
        missing = False
600
    if missing: warnings.warn(DbWarning(
601
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
602
603 2786 aaronmk
    return query
604 11 aaronmk
605 2054 aaronmk
def select(db, *args, **kw_args):
606
    '''For params, see mk_select() and run_query()'''
607
    recover = kw_args.pop('recover', None)
608
    cacheable = kw_args.pop('cacheable', True)
609 2442 aaronmk
    log_level = kw_args.pop('log_level', 2)
610 2054 aaronmk
611 2791 aaronmk
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
612
        log_level=log_level)
613 2054 aaronmk
614 2788 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
615
    embeddable=False):
616 1960 aaronmk
    '''
617
    @param returning str|None An inserted column (such as pkey) to return
618 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
619 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
620
        query will be fully cached, not just if it raises an exception.
621 1960 aaronmk
    '''
622 2754 aaronmk
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
623 2318 aaronmk
    if cols == []: cols = None # no cols (all defaults) = unknown col names
624 2646 aaronmk
    if cols != None:
625
        cols = [sql_gen.to_name_only_col(v, table).to_str(db) for v in cols]
626 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
627 2327 aaronmk
    if returning != None: returning = sql_gen.as_Col(returning, table)
628 2063 aaronmk
629
    # Build query
630 2497 aaronmk
    first_line = 'INSERT INTO '+table.to_str(db)
631
    query = first_line
632 2467 aaronmk
    if cols != None: query += '\n('+', '.join(cols)+')'
633
    query += '\n'+select_query
634 2063 aaronmk
635
    if returning != None:
636 2743 aaronmk
        query += '\nRETURNING '+sql_gen.to_name_only_col(returning).to_str(db)
637 2063 aaronmk
638 2070 aaronmk
    if embeddable:
639 2327 aaronmk
        assert returning != None
640
641 2070 aaronmk
        # Create function
642 2513 aaronmk
        function_name = sql_gen.clean_name(first_line)
643 2327 aaronmk
        return_type = 'SETOF '+returning.to_str(db)+'%TYPE'
644 2189 aaronmk
        while True:
645
            try:
646 2918 aaronmk
                function = db.TempFunction(function_name)
647 2194 aaronmk
648 2189 aaronmk
                function_query = '''\
649 2698 aaronmk
CREATE FUNCTION '''+function.to_str(db)+'''()
650 2467 aaronmk
RETURNS '''+return_type+'''
651
LANGUAGE sql
652
AS $$
653 2788 aaronmk
'''+query+''';
654 2467 aaronmk
$$;
655 2070 aaronmk
'''
656 2446 aaronmk
                run_query(db, function_query, recover=True, cacheable=True,
657 2945 aaronmk
                    log_ignore_excs=(DuplicateException,))
658 2189 aaronmk
                break # this version was successful
659 2945 aaronmk
            except DuplicateException, e:
660 2189 aaronmk
                function_name = next_version(function_name)
661
                # try again with next version of name
662 2070 aaronmk
663 2337 aaronmk
        # Return query that uses function
664 2698 aaronmk
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
665 2742 aaronmk
            [returning]) # AS clause requires function alias
666 2787 aaronmk
        return mk_select(db, func_table, start=0, order_by=None)
667 2070 aaronmk
668 2787 aaronmk
    return query
669 2066 aaronmk
670
def insert_select(db, *args, **kw_args):
671 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
672 2386 aaronmk
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
673
        values in
674 2072 aaronmk
    '''
675 2386 aaronmk
    into = kw_args.pop('into', None)
676
    if into != None: kw_args['embeddable'] = True
677 2066 aaronmk
    recover = kw_args.pop('recover', None)
678
    cacheable = kw_args.pop('cacheable', True)
679 2673 aaronmk
    log_level = kw_args.pop('log_level', 2)
680 2066 aaronmk
681 2789 aaronmk
    return run_query_into(db, mk_insert_select(db, *args, **kw_args), into,
682 2787 aaronmk
        recover=recover, cacheable=cacheable, log_level=log_level)
683 2063 aaronmk
684 2738 aaronmk
default = sql_gen.default # tells insert() to use the default value for a column
685 2066 aaronmk
686 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
687 2085 aaronmk
    '''For params, see insert_select()'''
688 1960 aaronmk
    if lists.is_seq(row): cols = None
689
    else:
690
        cols = row.keys()
691
        row = row.values()
692 2738 aaronmk
    row = list(row) # ensure that "== []" works
693 1960 aaronmk
694 2738 aaronmk
    if row == []: query = None
695
    else: query = sql_gen.Values(row).to_str(db)
696 1961 aaronmk
697 2788 aaronmk
    return insert_select(db, table, cols, query, *args, **kw_args)
698 11 aaronmk
699 2402 aaronmk
def mk_update(db, table, changes=None, cond=None):
700
    '''
701
    @param changes [(col, new_value),...]
702
        * container can be any iterable type
703
        * col: sql_gen.Code|str (for col name)
704
        * new_value: sql_gen.Code|literal value
705
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
706
    @return str query
707
    '''
708
    query = 'UPDATE '+sql_gen.as_Table(table).to_str(db)+'\nSET\n'
709 2405 aaronmk
    query += ',\n'.join((sql_gen.to_name_only_col(col, table).to_str(db)+' = '
710 2402 aaronmk
        +sql_gen.as_Value(new_value).to_str(db) for col, new_value in changes))
711 2467 aaronmk
    if cond != None: query += '\nWHERE\n'+cond.to_str(db)
712 2402 aaronmk
713
    return query
714
715
def update(db, *args, **kw_args):
716
    '''For params, see mk_update() and run_query()'''
717
    recover = kw_args.pop('recover', None)
718
719 2791 aaronmk
    return run_query(db, mk_update(db, *args, **kw_args), recover)
720 2402 aaronmk
721 135 aaronmk
def last_insert_id(db):
722 1849 aaronmk
    module = util.root_module(db.db)
723 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
724
    elif module == 'MySQLdb': return db.insert_id()
725
    else: return None
726 13 aaronmk
727 2394 aaronmk
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
728 2383 aaronmk
    '''Creates a mapping from original column names (which may have collisions)
729 2415 aaronmk
    to names that will be distinct among the columns' tables.
730 2383 aaronmk
    This is meant to be used for several tables that are being joined together.
731 2415 aaronmk
    @param cols The columns to combine. Duplicates will be removed.
732
    @param into The table for the new columns.
733 2394 aaronmk
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
734
        columns will be included in the mapping even if they are not in cols.
735
        The tables of the provided Col objects will be changed to into, so make
736
        copies of them if you want to keep the original tables.
737
    @param as_items Whether to return a list of dict items instead of a dict
738 2383 aaronmk
    @return dict(orig_col=new_col, ...)
739
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
740 2392 aaronmk
        * new_col: sql_gen.Col(orig_col_name, into)
741
        * All mappings use the into table so its name can easily be
742 2383 aaronmk
          changed for all columns at once
743
    '''
744 2415 aaronmk
    cols = lists.uniqify(cols)
745
746 2394 aaronmk
    items = []
747 2389 aaronmk
    for col in preserve:
748 2390 aaronmk
        orig_col = copy.copy(col)
749 2392 aaronmk
        col.table = into
750 2394 aaronmk
        items.append((orig_col, col))
751
    preserve = set(preserve)
752
    for col in cols:
753 2716 aaronmk
        if col not in preserve:
754
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
755 2394 aaronmk
756
    if not as_items: items = dict(items)
757
    return items
758 2383 aaronmk
759 2393 aaronmk
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
760 2391 aaronmk
    '''For params, see mk_flatten_mapping()
761
    @return See return value of mk_flatten_mapping()
762
    '''
763 2394 aaronmk
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
764
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
765 2786 aaronmk
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
766 2846 aaronmk
        into=into, add_indexes_=True)
767 2394 aaronmk
    return dict(items)
768 2391 aaronmk
769 2751 aaronmk
def mk_track_data_error(db, errors_table, cols, value, error_code, error):
770 2810 aaronmk
    assert cols != ()
771
772 2751 aaronmk
    cols = map(sql_gen.to_name_only_col, cols)
773
774
    columns_cols = ['column']
775
    columns = sql_gen.NamedValues('columns', columns_cols,
776
        [[c.name] for c in cols])
777
    values_cols = ['value', 'error_code', 'error']
778
    values = sql_gen.NamedValues('values', values_cols,
779
        [value, error_code, error])
780
781
    select_cols = columns_cols+values_cols
782
    name_only_cols = map(sql_gen.to_name_only_col, select_cols)
783 2755 aaronmk
    errors_table = sql_gen.NamedTable('errors', errors_table)
784 2751 aaronmk
    joins = [columns, sql_gen.Join(values, type_='CROSS'),
785
        sql_gen.Join(errors_table, dict(zip(name_only_cols, select_cols)),
786
        sql_gen.filter_out)]
787
788
    return mk_insert_select(db, errors_table, name_only_cols,
789 2787 aaronmk
        mk_select(db, joins, select_cols, order_by=None))
790 2751 aaronmk
791 2810 aaronmk
def track_data_error(db, errors_table, cols, *args, **kw_args):
792 2758 aaronmk
    '''
793
    @param errors_table If None, does nothing.
794
    '''
795 2810 aaronmk
    if errors_table == None or cols == (): return
796
    run_query(db, mk_track_data_error(db, errors_table, cols, *args, **kw_args),
797 2758 aaronmk
        cacheable=True, log_level=4)
798
799 2704 aaronmk
def cast(db, type_, col, errors_table=None):
800
    '''Casts an (unrenamed) column or value.
801 2769 aaronmk
    If errors_table set and col has srcs, saves errors in errors_table (using
802
    col's srcs attr as the source columns) and converts errors to warnings.
803 2820 aaronmk
    @param col str|sql_gen.Col|sql_gen.Literal
804 2704 aaronmk
    @param errors_table None|sql_gen.Table|str
805
    '''
806 2820 aaronmk
    col = sql_gen.as_Col(col)
807 2769 aaronmk
    save_errors = (errors_table != None and isinstance(col, sql_gen.Col)
808
        and col.srcs != ())
809
    if not save_errors: # can't save errors
810
        return sql_gen.CustomCode(col.to_str(db)+'::'+type_) # just cast
811 2704 aaronmk
812
    assert not isinstance(col, sql_gen.NamedCol)
813
814 2770 aaronmk
    errors_table = sql_gen.as_Table(errors_table)
815
    srcs = map(sql_gen.to_name_only_col, col.srcs)
816
    function_name = str(sql_gen.FunctionCall(type_, *srcs))
817 2918 aaronmk
    function = db.TempFunction(function_name)
818 2704 aaronmk
819 2718 aaronmk
    while True:
820
        # Create function definition
821
        query = '''\
822 2704 aaronmk
CREATE FUNCTION '''+function.to_str(db)+'''(value text)
823
RETURNS '''+type_+'''
824
LANGUAGE plpgsql
825
STRICT
826
AS $$
827
BEGIN
828
    /* The explicit cast to the return type is needed to make the cast happen
829
    inside the try block. (Implicit casts to the return type happen at the end
830
    of the function, outside any block.) */
831
    RETURN value::'''+type_+''';
832
EXCEPTION
833
    WHEN data_exception THEN
834
        -- Save error in errors table.
835 2751 aaronmk
        -- Insert the value and error for *each* source column.
836
'''+mk_track_data_error(db, errors_table, srcs,
837
    *map(sql_gen.CustomCode, ['value', 'SQLSTATE', 'SQLERRM']))+''';
838 2704 aaronmk
839
        RAISE WARNING '%', SQLERRM;
840
        RETURN NULL;
841
END;
842
$$;
843
'''
844 2718 aaronmk
845
        # Create function
846
        try:
847
            run_query(db, query, recover=True, cacheable=True,
848 2945 aaronmk
                log_ignore_excs=(DuplicateException,))
849 2718 aaronmk
            break # successful
850 2945 aaronmk
        except DuplicateException:
851 2770 aaronmk
            function.name = next_version(function.name)
852
            # try again with next version of name
853 2704 aaronmk
854
    return sql_gen.FunctionCall(function, col)
855
856 2414 aaronmk
##### Database structure queries
857
858 2426 aaronmk
def table_row_count(db, table, recover=None):
859 2786 aaronmk
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
860 2443 aaronmk
        order_by=None, start=0), recover=recover, log_level=3))
861 2426 aaronmk
862 2414 aaronmk
def table_cols(db, table, recover=None):
863
    return list(col_names(select(db, table, limit=0, order_by=None,
864 2443 aaronmk
        recover=recover, log_level=4)))
865 2414 aaronmk
866 2291 aaronmk
def pkey(db, table, recover=None):
867 832 aaronmk
    '''Assumed to be first column in table'''
868 2339 aaronmk
    return table_cols(db, table, recover)[0]
869 832 aaronmk
870 2559 aaronmk
not_null_col = 'not_null_col'
871 2340 aaronmk
872
def table_not_null_col(db, table, recover=None):
873
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
874
    if not_null_col in table_cols(db, table, recover): return not_null_col
875
    else: return pkey(db, table, recover)
876
877 853 aaronmk
def index_cols(db, table, index):
878
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
879
    automatically created. When you don't know whether something is a UNIQUE
880
    constraint or a UNIQUE index, use this function.'''
881 1909 aaronmk
    module = util.root_module(db.db)
882
    if module == 'psycopg2':
883
        return list(values(run_query(db, '''\
884 853 aaronmk
SELECT attname
885 866 aaronmk
FROM
886
(
887
        SELECT attnum, attname
888
        FROM pg_index
889
        JOIN pg_class index ON index.oid = indexrelid
890
        JOIN pg_class table_ ON table_.oid = indrelid
891
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
892
        WHERE
893 2782 aaronmk
            table_.relname = '''+db.esc_value(table)+'''
894
            AND index.relname = '''+db.esc_value(index)+'''
895 866 aaronmk
    UNION
896
        SELECT attnum, attname
897
        FROM
898
        (
899
            SELECT
900
                indrelid
901
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
902
                    AS indkey
903
            FROM pg_index
904
            JOIN pg_class index ON index.oid = indexrelid
905
            JOIN pg_class table_ ON table_.oid = indrelid
906
            WHERE
907 2782 aaronmk
                table_.relname = '''+db.esc_value(table)+'''
908
                AND index.relname = '''+db.esc_value(index)+'''
909 866 aaronmk
        ) s
910
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
911
) s
912 853 aaronmk
ORDER BY attnum
913 2782 aaronmk
'''
914
            , cacheable=True, log_level=4)))
915 1909 aaronmk
    else: raise NotImplementedError("Can't list index columns for "+module+
916
        ' database')
917 853 aaronmk
918 464 aaronmk
def constraint_cols(db, table, constraint):
919 1849 aaronmk
    module = util.root_module(db.db)
920 464 aaronmk
    if module == 'psycopg2':
921
        return list(values(run_query(db, '''\
922
SELECT attname
923
FROM pg_constraint
924
JOIN pg_class ON pg_class.oid = conrelid
925
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
926
WHERE
927 2783 aaronmk
    relname = '''+db.esc_value(table)+'''
928
    AND conname = '''+db.esc_value(constraint)+'''
929 464 aaronmk
ORDER BY attnum
930 2783 aaronmk
'''
931
            )))
932 464 aaronmk
    else: raise NotImplementedError("Can't list constraint columns for "+module+
933
        ' database')
934
935 2096 aaronmk
row_num_col = '_row_num'
936
937 2847 aaronmk
def add_index(db, exprs, table=None, unique=False, ensure_not_null=True):
938 2688 aaronmk
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
939 2538 aaronmk
    Currently, only function calls are supported as expressions.
940 2847 aaronmk
    @param ensure_not_null If set, translates NULL values to sentinel values.
941
        This allows indexes to be used for comparisons where NULLs are equal.
942 2538 aaronmk
    '''
943 2688 aaronmk
    if not lists.is_seq(exprs): exprs = [exprs]
944 2538 aaronmk
945 2688 aaronmk
    # Parse exprs
946
    old_exprs = exprs[:]
947
    exprs = []
948
    cols = []
949
    for i, expr in enumerate(old_exprs):
950
        expr = copy.deepcopy(expr) # don't modify input!
951 2823 aaronmk
        expr = sql_gen.as_Col(expr, table)
952 2688 aaronmk
953 2823 aaronmk
        # Handle nullable columns
954 2860 aaronmk
        if ensure_not_null:
955
            try: expr = sql_gen.ensure_not_null(db, expr)
956
            except KeyError: pass # unknown type, so just create plain index
957 2823 aaronmk
958 2688 aaronmk
        # Extract col
959
        if isinstance(expr, sql_gen.FunctionCall):
960
            col = expr.args[0]
961
            expr = sql_gen.Expr(expr)
962
        else: col = expr
963 2823 aaronmk
        assert isinstance(col, sql_gen.Col)
964 2688 aaronmk
965
        # Extract table
966
        if table == None:
967
            assert sql_gen.is_table_col(col)
968
            table = col.table
969
970
        col.table = None
971
972
        exprs.append(expr)
973
        cols.append(col)
974 2408 aaronmk
975 2688 aaronmk
    table = sql_gen.as_Table(table)
976
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
977
978
    str_ = 'CREATE'
979
    if unique: str_ += ' UNIQUE'
980
    str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
981
        ', '.join((v.to_str(db) for v in exprs)))+')'
982
983
    try: run_query(db, str_, recover=True, cacheable=True, log_level=3)
984 2945 aaronmk
    except DuplicateException: pass # index already existed
985 2408 aaronmk
986 2686 aaronmk
def add_pkey(db, table, cols=None, recover=None):
987
    '''Adds a primary key.
988
    @param cols [sql_gen.Col,...] The columns in the primary key.
989
        Defaults to the first column in the table.
990 2406 aaronmk
    @pre The table must not already have a primary key.
991
    '''
992
    table = sql_gen.as_Table(table)
993 2686 aaronmk
    if cols == None: cols = [pkey(db, table, recover)]
994
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
995 2406 aaronmk
996 2781 aaronmk
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
997
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
998 2945 aaronmk
        log_ignore_excs=(DuplicateException,))
999 2406 aaronmk
1000 2772 aaronmk
already_indexed = object() # tells add_indexes() the pkey has already been added
1001
1002
def add_indexes(db, table, has_pkey=True):
1003
    '''Adds an index on all columns in a table.
1004
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1005
        index should be added on the first column.
1006
        * If already_indexed, the pkey is assumed to have already been added
1007
    '''
1008
    cols = table_cols(db, table)
1009
    if has_pkey:
1010
        if has_pkey is not already_indexed: add_pkey(db, table)
1011
        cols = cols[1:]
1012
    for col in cols: add_index(db, col, table)
1013
1014 2935 aaronmk
def add_col(db, table, col, **kw_args):
1015 2934 aaronmk
    assert isinstance(col, sql_gen.TypedCol)
1016 2947 aaronmk
    try: run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '
1017
        +col.to_str(db), recover=True, cacheable=True, **kw_args)
1018
    except DuplicateException: pass # column already existed
1019 2934 aaronmk
1020 2935 aaronmk
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1021
    constraints='PRIMARY KEY')
1022
1023 2086 aaronmk
def add_row_num(db, table):
1024 2117 aaronmk
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1025
    be the primary key.'''
1026 2935 aaronmk
    add_col(db, table, row_num_typed_col, log_level=3)
1027 2086 aaronmk
1028 2936 aaronmk
def cast_temp_col(db, type_, col, errors_table=None):
1029
    '''Like cast(), but creates a new column with the cast values if the input
1030
    is a column.
1031
    @return The new column or cast value
1032
    '''
1033
    def cast_(col): return cast(db, type_, col, errors_table)
1034
1035
    try: col = sql_gen.underlying_col(col)
1036
    except sql_gen.NoUnderlyingTableException: return sql_gen.wrap(cast_, col)
1037
1038
    table = col.table
1039
    new_col = sql_gen.Col(sql_gen.concat(col.name, '::'+type_), table, col.srcs)
1040
    expr = cast_(col)
1041
    add_col(db, table, sql_gen.TypedCol(new_col.name, type_))
1042
    update(db, table, [(new_col, expr)])
1043 2955 aaronmk
    add_index(db, new_col)
1044 2936 aaronmk
1045
    return new_col
1046
1047 2889 aaronmk
def drop_table(db, table):
1048
    table = sql_gen.as_Table(table)
1049
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1050
1051 2760 aaronmk
def create_table(db, table, cols, has_pkey=True, col_indexes=True):
1052 2675 aaronmk
    '''Creates a table.
1053 2681 aaronmk
    @param cols [sql_gen.TypedCol,...] The column names and types
1054
    @param has_pkey If set, the first column becomes the primary key.
1055 2760 aaronmk
    @param col_indexes bool|[ref]
1056
        * If True, indexes will be added on all non-pkey columns.
1057
        * If a list reference, [0] will be set to a function to do this.
1058
          This can be used to delay index creation until the table is populated.
1059 2675 aaronmk
    '''
1060
    table = sql_gen.as_Table(table)
1061
1062 2681 aaronmk
    if has_pkey:
1063
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1064 2872 aaronmk
        pkey.constraints = 'PRIMARY KEY'
1065 2681 aaronmk
1066 2675 aaronmk
    str_ = 'CREATE TABLE '+table.to_str(db)+' (\n'
1067 2681 aaronmk
    str_ += '\n, '.join(v.to_str(db) for v in cols)
1068
    str_ += '\n);\n'
1069 2675 aaronmk
    run_query(db, str_, cacheable=True, log_level=2)
1070 2760 aaronmk
1071
    # Add indexes
1072 2773 aaronmk
    if has_pkey: has_pkey = already_indexed
1073
    def add_indexes_(): add_indexes(db, table, has_pkey)
1074
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1075
    elif col_indexes: add_indexes_() # add now
1076 2675 aaronmk
1077 2684 aaronmk
def vacuum(db, table):
1078
    table = sql_gen.as_Table(table)
1079
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1080
        log_level=3))
1081
1082 2732 aaronmk
def truncate(db, table, schema='public'):
1083 2777 aaronmk
    table = sql_gen.as_Table(table, schema)
1084 2733 aaronmk
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE')
1085 2732 aaronmk
1086 2729 aaronmk
def tables(db, schema_like='public', table_like='%', exact=False):
1087
    if exact: compare = '='
1088
    else: compare = 'LIKE'
1089
1090 1849 aaronmk
    module = util.root_module(db.db)
1091 832 aaronmk
    if module == 'psycopg2':
1092 2779 aaronmk
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1093
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1094
        return values(select(db, 'pg_tables', ['tablename'], conds,
1095
            order_by='tablename', log_level=4))
1096 1968 aaronmk
    elif module == 'MySQLdb':
1097 2779 aaronmk
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1098
            , cacheable=True, log_level=4))
1099 832 aaronmk
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1100 830 aaronmk
1101 2726 aaronmk
def table_exists(db, table):
1102
    table = sql_gen.as_Table(table)
1103 2730 aaronmk
    return list(tables(db, table.schema, table.name, exact=True)) != []
1104 2726 aaronmk
1105 2913 aaronmk
def function_exists(db, function):
1106
    function = sql_gen.as_Function(function)
1107
1108
    info_table = sql_gen.Table('routines', 'information_schema')
1109
    conds = [('routine_name', function.name)]
1110
    schema = function.schema
1111
    if schema != None: conds.append(('routine_schema', schema))
1112 2914 aaronmk
    # Exclude trigger functions, since they cannot be called directly
1113
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1114 2913 aaronmk
1115
    return list(values(select(db, info_table, ['routine_name'], conds,
1116
        order_by='routine_schema', limit=1, log_level=4))) != []
1117
        # TODO: order_by search_path schema order
1118
1119 2726 aaronmk
def errors_table(db, table, if_exists=True):
1120
    '''
1121
    @param if_exists If set, returns None if the errors table doesn't exist
1122
    @return None|sql_gen.Table
1123
    '''
1124
    table = sql_gen.as_Table(table)
1125
    if table.srcs != (): table = table.srcs[0]
1126
1127
    errors_table = sql_gen.suffixed_table(table, '.errors')
1128
    if if_exists and not table_exists(db, errors_table): return None
1129
    return errors_table
1130
1131 833 aaronmk
##### Database management
1132
1133 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
1134
    '''For kw_args, see tables()'''
1135
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1136 833 aaronmk
1137 832 aaronmk
##### Heuristic queries
1138
1139 2104 aaronmk
def put(db, table, row, pkey_=None, row_ct_ref=None):
1140 1554 aaronmk
    '''Recovers from errors.
1141 2077 aaronmk
    Only works under PostgreSQL (uses INSERT RETURNING).
1142
    '''
1143 2642 aaronmk
    row = sql_gen.ColDict(db, table, row)
1144 2104 aaronmk
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
1145
1146 471 aaronmk
    try:
1147 2149 aaronmk
        cur = insert(db, table, row, pkey_, recover=True)
1148 1554 aaronmk
        if row_ct_ref != None and cur.rowcount >= 0:
1149
            row_ct_ref[0] += cur.rowcount
1150
        return value(cur)
1151 471 aaronmk
    except DuplicateKeyException, e:
1152 2653 aaronmk
        row = sql_gen.ColDict(db, table,
1153
            util.dict_subset_right_join(row, e.cols))
1154
        return value(select(db, table, [pkey_], row, recover=True))
1155 471 aaronmk
1156 473 aaronmk
def get(db, table, row, pkey, row_ct_ref=None, create=False):
1157 830 aaronmk
    '''Recovers from errors'''
1158 2209 aaronmk
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
1159 14 aaronmk
    except StopIteration:
1160 40 aaronmk
        if not create: raise
1161 471 aaronmk
        return put(db, table, row, pkey, row_ct_ref) # insert new row
1162 2078 aaronmk
1163 2593 aaronmk
def is_func_result(col):
1164
    return col.table.name.find('(') >= 0 and col.name == 'result'
1165
1166 2592 aaronmk
def into_table_name(out_table, in_tables0, mapping, is_func):
1167 2615 aaronmk
    def in_col_str(in_col):
1168
        in_col = sql_gen.remove_col_rename(in_col)
1169
        if isinstance(in_col, sql_gen.Col):
1170
            table = in_col.table
1171
            if table == in_tables0:
1172
                in_col = sql_gen.to_name_only_col(in_col)
1173
            elif is_func_result(in_col): in_col = table # omit col name
1174
        return str(in_col)
1175
1176 2580 aaronmk
    str_ = str(out_table)
1177
    if is_func:
1178 2615 aaronmk
        str_ += '('
1179 2580 aaronmk
1180 2615 aaronmk
        try: value_in_col = mapping['value']
1181
        except KeyError:
1182 2654 aaronmk
            str_ += ', '.join((str(k)+'='+in_col_str(v)
1183 2615 aaronmk
                for k, v in mapping.iteritems()))
1184
        else: str_ += in_col_str(value_in_col)
1185
1186
        str_ += ')'
1187 2812 aaronmk
    else:
1188
        out_col = 'rank'
1189
        try: in_col = mapping[out_col]
1190
        except KeyError: str_ += '_pkeys'
1191
        else: # has a rank column, so hierarchical
1192
            str_ += '['+str(out_col)+'='+in_col_str(in_col)+']'
1193 2580 aaronmk
    return str_
1194
1195 2508 aaronmk
def put_table(db, out_table, in_tables, mapping, row_ct_ref=None, into=None,
1196 2805 aaronmk
    default=None, is_func=False, on_error=exc.raise_):
1197 2078 aaronmk
    '''Recovers from errors.
1198
    Only works under PostgreSQL (uses INSERT RETURNING).
1199 2909 aaronmk
    IMPORTANT: Must be run at the *beginning* of a transaction.
1200 2131 aaronmk
    @param in_tables The main input table to select from, followed by a list of
1201
        tables to join with it using the main input table's pkey
1202 2312 aaronmk
    @param mapping dict(out_table_col=in_table_col, ...)
1203 2612 aaronmk
        * out_table_col: str (*not* sql_gen.Col)
1204
        * in_table_col: sql_gen.Col|literal-value
1205 2489 aaronmk
    @param into The table to contain the output and input pkeys.
1206 2574 aaronmk
        Defaults to `out_table.name+'_pkeys'`.
1207 2509 aaronmk
    @param default The *output* column to use as the pkey for missing rows.
1208
        If this output column does not exist in the mapping, uses None.
1209 2552 aaronmk
    @param is_func Whether out_table is the name of a SQL function, not a table
1210 2312 aaronmk
    @return sql_gen.Col Where the output pkeys are made available
1211 2078 aaronmk
    '''
1212 2329 aaronmk
    out_table = sql_gen.as_Table(out_table)
1213 2312 aaronmk
1214 2450 aaronmk
    def log_debug(msg): db.log_debug(msg, level=1.5)
1215 2505 aaronmk
    def col_ustr(str_):
1216 2567 aaronmk
        return strings.repr_no_u(sql_gen.remove_col_rename(str_))
1217 2450 aaronmk
1218 2910 aaronmk
    log_debug('********** New iteration **********')
1219
    log_debug('Inserting these input columns into '+strings.as_tt(
1220
        out_table.to_str(db))+':\n'+strings.as_table(mapping, ustr=col_ustr))
1221
1222 2942 aaronmk
    is_function = function_exists(db, out_table)
1223
1224
    if is_function: out_pkey = 'result'
1225
    else: out_pkey = pkey(db, out_table, recover=True)
1226 2768 aaronmk
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
1227
1228
    if mapping == {}: # need at least one column for INSERT SELECT
1229
        mapping = {out_pkey: None} # ColDict will replace with default value
1230
1231 2382 aaronmk
    # Create input joins from list of input tables
1232
    in_tables_ = in_tables[:] # don't modify input!
1233
    in_tables0 = in_tables_.pop(0) # first table is separate
1234 2735 aaronmk
    errors_table_ = errors_table(db, in_tables0)
1235 2279 aaronmk
    in_pkey = pkey(db, in_tables0, recover=True)
1236 2285 aaronmk
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
1237 2460 aaronmk
    input_joins = [in_tables0]+[sql_gen.Join(v,
1238
        {in_pkey: sql_gen.join_same_not_null}) for v in in_tables_]
1239 2131 aaronmk
1240 2592 aaronmk
    if into == None:
1241
        into = into_table_name(out_table, in_tables0, mapping, is_func)
1242
    into = sql_gen.as_Table(into)
1243
1244 2702 aaronmk
    # Set column sources
1245
    in_cols = filter(sql_gen.is_table_col, mapping.values())
1246
    for col in in_cols:
1247
        if col.table == in_tables0: col.set_srcs(sql_gen.src_self)
1248
1249 2486 aaronmk
    log_debug('Joining together input tables into temp table')
1250 2395 aaronmk
    # Place in new table for speed and so don't modify input if values edited
1251 2584 aaronmk
    in_table = sql_gen.Table('in')
1252 2702 aaronmk
    mapping = dicts.join(mapping, flatten(db, in_table, input_joins, in_cols,
1253
        preserve=[in_pkey_col], start=0))
1254 2395 aaronmk
    input_joins = [in_table]
1255 2486 aaronmk
    db.log_debug('Temp table: '+strings.as_tt(in_table.to_str(db)), level=2)
1256 2395 aaronmk
1257 2656 aaronmk
    mapping = sql_gen.ColDict(db, out_table, mapping)
1258
        # after applying dicts.join() because that returns a plain dict
1259
1260 2509 aaronmk
    # Resolve default value column
1261 2841 aaronmk
    if default != None:
1262
        try: default = mapping[default]
1263
        except KeyError:
1264 2509 aaronmk
            db.log_debug('Default value column '
1265
                +strings.as_tt(strings.repr_no_u(default))
1266 2511 aaronmk
                +' does not exist in mapping, falling back to None', level=2.1)
1267 2509 aaronmk
            default = None
1268
1269 2387 aaronmk
    pkeys_names = [in_pkey, out_pkey]
1270 2236 aaronmk
    pkeys_cols = [in_pkey_col, out_pkey_col]
1271
1272 2201 aaronmk
    pkeys_table_exists_ref = [False]
1273 2960 aaronmk
    def insert_into_pkeys(joins, cols, distinct=False):
1274
        kw_args = {}
1275
        if distinct: kw_args.update(dict(distinct_on=[in_pkey_col]))
1276
        query = mk_select(db, joins, cols, order_by=None, start=0, **kw_args)
1277
1278 2201 aaronmk
        if pkeys_table_exists_ref[0]:
1279 2786 aaronmk
            insert_select(db, into, pkeys_names, query)
1280 2201 aaronmk
        else:
1281 2786 aaronmk
            run_query_into(db, query, into=into)
1282 2201 aaronmk
            pkeys_table_exists_ref[0] = True
1283
1284 2429 aaronmk
    limit_ref = [None]
1285 2380 aaronmk
    conds = set()
1286 2811 aaronmk
    distinct_on = sql_gen.ColDict(db, out_table)
1287 2325 aaronmk
    def mk_main_select(joins, cols):
1288 2811 aaronmk
        distinct_on_cols = [c.to_Col() for c in distinct_on.values()]
1289
        return mk_select(db, joins, cols, conds, distinct_on_cols,
1290 2429 aaronmk
            limit=limit_ref[0], start=0)
1291 2132 aaronmk
1292 2519 aaronmk
    exc_strs = set()
1293 2309 aaronmk
    def log_exc(e):
1294 2519 aaronmk
        e_str = exc.str_(e, first_line_only=True)
1295
        log_debug('Caught exception: '+e_str)
1296
        assert e_str not in exc_strs # avoid infinite loops
1297
        exc_strs.add(e_str)
1298 2735 aaronmk
1299 2451 aaronmk
    def remove_all_rows():
1300 2896 aaronmk
        log_debug('Ignoring all rows')
1301 2429 aaronmk
        limit_ref[0] = 0 # just create an empty pkeys table
1302 2735 aaronmk
1303 2758 aaronmk
    def ignore(in_col, value, e):
1304
        track_data_error(db, errors_table_, in_col.srcs, value, e.cause.pgcode,
1305
            e.cause.pgerror)
1306 2895 aaronmk
        log_debug('Ignoring rows with '+strings.as_tt(repr(in_col))+' = '
1307 2545 aaronmk
            +strings.as_tt(repr(value)))
1308 2895 aaronmk
1309 2758 aaronmk
    def remove_rows(in_col, value, e):
1310
        ignore(in_col, value, e)
1311 2378 aaronmk
        cond = (in_col, sql_gen.CompareCond(value, '!='))
1312
        assert cond not in conds # avoid infinite loops
1313 2380 aaronmk
        conds.add(cond)
1314 2895 aaronmk
1315 2758 aaronmk
    def invalid2null(in_col, value, e):
1316
        ignore(in_col, value, e)
1317 2403 aaronmk
        update(db, in_table, [(in_col, None)],
1318
            sql_gen.ColValueCond(in_col, value))
1319 2245 aaronmk
1320 2589 aaronmk
    def insert_pkeys_table(which):
1321 2932 aaronmk
        return sql_gen.Table(sql_gen.concat(in_table.name,
1322 2589 aaronmk
            '_insert_'+which+'_pkeys'))
1323
    insert_out_pkeys = insert_pkeys_table('out')
1324
    insert_in_pkeys = insert_pkeys_table('in')
1325
1326 2206 aaronmk
    # Do inserts and selects
1327 2642 aaronmk
    join_cols = sql_gen.ColDict(db, out_table)
1328 2206 aaronmk
    while True:
1329 2521 aaronmk
        if limit_ref[0] == 0: # special case
1330
            log_debug('Creating an empty pkeys table')
1331 2786 aaronmk
            cur = run_query_into(db, mk_select(db, out_table, [out_pkey],
1332 2521 aaronmk
                limit=limit_ref[0]), into=insert_out_pkeys)
1333
            break # don't do main case
1334
1335 2303 aaronmk
        has_joins = join_cols != {}
1336
1337 2903 aaronmk
        log_debug('Trying to insert new rows')
1338
1339 2305 aaronmk
        # Prepare to insert new rows
1340 2325 aaronmk
        insert_joins = input_joins[:] # don't modify original!
1341 2403 aaronmk
        insert_args = dict(recover=True, cacheable=False)
1342 2303 aaronmk
        if has_joins:
1343 2325 aaronmk
            insert_joins.append(sql_gen.Join(out_table, join_cols,
1344
                sql_gen.filter_out))
1345
        else:
1346 2404 aaronmk
            insert_args.update(dict(returning=out_pkey, into=insert_out_pkeys))
1347 2786 aaronmk
        main_select = mk_main_select(insert_joins, mapping.values())
1348 2303 aaronmk
1349 2929 aaronmk
        def main_insert():
1350 2942 aaronmk
            if is_function:
1351
                log_debug('Calling function on input rows')
1352
                args = dict(((k.name, v) for k, v in mapping.iteritems()))
1353
                func_call = sql_gen.NamedCol(out_pkey,
1354
                    sql_gen.FunctionCall(out_table, **args))
1355
                insert_into_pkeys(input_joins, [in_pkey_col, func_call])
1356
                return None
1357
            else:
1358
                lock_table(db, out_table, 'EXCLUSIVE')
1359
                return insert_select(db, out_table, mapping.keys(), main_select,
1360
                    **insert_args)
1361 2929 aaronmk
1362 2206 aaronmk
        try:
1363 2929 aaronmk
            cur = with_savepoint(db, main_insert)
1364 2357 aaronmk
            break # insert successful
1365 2931 aaronmk
        except MissingCastException, e:
1366
            log_exc(e)
1367
1368
            out_col = e.col
1369
            type_ = e.type
1370
1371
            log_debug('Casting '+strings.as_tt(out_col)+' input to '
1372
                +strings.as_tt(type_))
1373 2937 aaronmk
            mapping[out_col] = cast_temp_col(db, type_, mapping[out_col],
1374
                errors_table_)
1375 2206 aaronmk
        except DuplicateKeyException, e:
1376 2309 aaronmk
            log_exc(e)
1377
1378 2258 aaronmk
            old_join_cols = join_cols.copy()
1379 2811 aaronmk
            distinct_on.update(util.dict_subset(mapping, e.cols))
1380 2565 aaronmk
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
1381 2486 aaronmk
            log_debug('Ignoring existing rows, comparing on these columns:\n'
1382 2505 aaronmk
                +strings.as_inline_table(join_cols, ustr=col_ustr))
1383 2258 aaronmk
            assert join_cols != old_join_cols # avoid infinite loops
1384 2230 aaronmk
        except NullValueException, e:
1385 2309 aaronmk
            log_exc(e)
1386
1387 2230 aaronmk
            out_col, = e.cols
1388
            try: in_col = mapping[out_col]
1389 2356 aaronmk
            except KeyError:
1390 2486 aaronmk
                log_debug('Missing mapping for NOT NULL column '+out_col)
1391 2451 aaronmk
                remove_all_rows()
1392 2758 aaronmk
            else: remove_rows(in_col, None, e)
1393 2542 aaronmk
        except FunctionValueException, e:
1394
            log_exc(e)
1395
1396
            func_name = e.name
1397
            value = e.value
1398
            for out_col, in_col in mapping.iteritems():
1399 2758 aaronmk
                in_col = sql_gen.unwrap_func_call(in_col, func_name)
1400
                invalid2null(in_col, value, e)
1401 2429 aaronmk
        except DatabaseErrors, e:
1402
            log_exc(e)
1403
1404 2808 aaronmk
            log_debug('No handler for exception')
1405
            on_error(e)
1406 2451 aaronmk
            remove_all_rows()
1407 2358 aaronmk
        # after exception handled, rerun loop with additional constraints
1408 2132 aaronmk
1409 2942 aaronmk
    if cur != None and row_ct_ref != None and cur.rowcount >= 0:
1410 2357 aaronmk
        row_ct_ref[0] += cur.rowcount
1411
1412 2942 aaronmk
    if is_function: pass # pkeys table already created
1413
    elif has_joins:
1414 2357 aaronmk
        select_joins = input_joins+[sql_gen.Join(out_table, join_cols)]
1415 2486 aaronmk
        log_debug('Getting output table pkeys of existing/inserted rows')
1416 2960 aaronmk
        insert_into_pkeys(select_joins, pkeys_cols, distinct=True)
1417 2357 aaronmk
    else:
1418 2404 aaronmk
        add_row_num(db, insert_out_pkeys) # for joining with input pkeys
1419 2357 aaronmk
1420 2486 aaronmk
        log_debug('Getting input table pkeys of inserted rows')
1421 2786 aaronmk
        run_query_into(db, mk_main_select(input_joins, [in_pkey]),
1422 2404 aaronmk
            into=insert_in_pkeys)
1423
        add_row_num(db, insert_in_pkeys) # for joining with output pkeys
1424 2357 aaronmk
1425 2428 aaronmk
        assert table_row_count(db, insert_out_pkeys) == table_row_count(db,
1426
            insert_in_pkeys)
1427
1428 2486 aaronmk
        log_debug('Combining output and input pkeys in inserted order')
1429 2404 aaronmk
        pkey_joins = [insert_in_pkeys, sql_gen.Join(insert_out_pkeys,
1430 2357 aaronmk
            {row_num_col: sql_gen.join_same_not_null})]
1431 2420 aaronmk
        insert_into_pkeys(pkey_joins, pkeys_names)
1432 2959 aaronmk
1433
        if not db.debug_temp:
1434
            log_debug('Emptying unneeded temp tables to free up memory')
1435
            truncate(db, insert_out_pkeys)
1436
            truncate(db, insert_in_pkeys)
1437 2357 aaronmk
1438 2486 aaronmk
    db.log_debug('Adding pkey on pkeys table to enable fast joins', level=2.5)
1439 2594 aaronmk
    add_pkey(db, into)
1440 2407 aaronmk
1441 2508 aaronmk
    log_debug('Setting pkeys of missing rows to '+strings.as_tt(repr(default)))
1442 2489 aaronmk
    missing_rows_joins = input_joins+[sql_gen.Join(into,
1443 2357 aaronmk
        {in_pkey: sql_gen.join_same_not_null}, sql_gen.filter_out)]
1444
        # must use join_same_not_null or query will take forever
1445 2420 aaronmk
    insert_into_pkeys(missing_rows_joins,
1446 2508 aaronmk
        [in_pkey_col, sql_gen.NamedCol(out_pkey, default)])
1447 2357 aaronmk
1448 2489 aaronmk
    assert table_row_count(db, into) == table_row_count(db, in_table)
1449 2428 aaronmk
1450 2959 aaronmk
    if not db.debug_temp:
1451
        log_debug('Emptying unneeded temp tables to free up memory')
1452
        truncate(db, in_table)
1453
1454 2702 aaronmk
    srcs = []
1455
    if is_func: srcs = sql_gen.cols_srcs(in_cols)
1456
    return sql_gen.Col(out_pkey, into, srcs)
1457 2115 aaronmk
1458
##### Data cleanup
1459
1460 2290 aaronmk
def cleanup_table(db, table, cols):
1461 2677 aaronmk
    table = sql_gen.as_Table(table)
1462
    cols = map(sql_gen.as_Col, cols)
1463 2115 aaronmk
1464 2677 aaronmk
    expr = ('nullif(nullif(trim(both from %s), '+db.esc_value('')+'), '
1465
        +db.esc_value(r'\N')+')')
1466
    changes = [(v, sql_gen.CustomCode(expr % v.to_str(db)))
1467
        for v in cols]
1468 2115 aaronmk
1469 2677 aaronmk
    update(db, table, changes)