Project

General

Profile

1 11 aaronmk
# Database access
2
3 1869 aaronmk
import copy
4 11 aaronmk
import re
5 3238 aaronmk
import time
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 3241 aaronmk
import profiling
13 1889 aaronmk
from Proxy import Proxy
14 1872 aaronmk
import rand
15 2217 aaronmk
import sql_gen
16 862 aaronmk
import strings
17 131 aaronmk
import util
18 11 aaronmk
19 832 aaronmk
##### Exceptions
20
21 2804 aaronmk
def get_cur_query(cur, input_query=None):
22 2168 aaronmk
    raw_query = None
23
    if hasattr(cur, 'query'): raw_query = cur.query
24
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
25 2170 aaronmk
26
    if raw_query != None: return raw_query
27 2804 aaronmk
    else: return '[input] '+strings.ustr(input_query)
28 14 aaronmk
29 2170 aaronmk
def _add_cursor_info(e, *args, **kw_args):
30
    '''For params, see get_cur_query()'''
31 2771 aaronmk
    exc.add_msg(e, 'query: '+strings.ustr(get_cur_query(*args, **kw_args)))
32 135 aaronmk
33 300 aaronmk
class DbException(exc.ExceptionWithCause):
34 14 aaronmk
    def __init__(self, msg, cause=None, cur=None):
35 2145 aaronmk
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
36 14 aaronmk
        if cur != None: _add_cursor_info(self, cur)
37
38 2143 aaronmk
class ExceptionWithName(DbException):
39
    def __init__(self, name, cause=None):
40 4491 aaronmk
        DbException.__init__(self, 'for name: '
41
            +strings.as_tt(strings.ustr(name)), cause)
42 2143 aaronmk
        self.name = name
43 360 aaronmk
44 3109 aaronmk
class ExceptionWithValue(DbException):
45
    def __init__(self, value, cause=None):
46 4492 aaronmk
        DbException.__init__(self, 'for value: '
47
            +strings.as_tt(strings.urepr(value)), cause)
48 2240 aaronmk
        self.value = value
49
50 2945 aaronmk
class ExceptionWithNameType(DbException):
51
    def __init__(self, type_, name, cause=None):
52 4491 aaronmk
        DbException.__init__(self, 'for type: '+strings.as_tt(strings.ustr(
53
            type_))+'; name: '+strings.as_tt(name), cause)
54 2945 aaronmk
        self.type = type_
55
        self.name = name
56
57 2306 aaronmk
class ConstraintException(DbException):
58 3345 aaronmk
    def __init__(self, name, cond, cols, cause=None):
59
        msg = 'Violated '+strings.as_tt(name)+' constraint'
60 5447 aaronmk
        if cond != None: msg += ' with condition '+strings.as_tt(cond)
61 3345 aaronmk
        if cols != []: msg += ' on columns: '+strings.as_tt(', '.join(cols))
62
        DbException.__init__(self, msg, cause)
63 2306 aaronmk
        self.name = name
64 3345 aaronmk
        self.cond = cond
65 468 aaronmk
        self.cols = cols
66 11 aaronmk
67 2523 aaronmk
class MissingCastException(DbException):
68 4139 aaronmk
    def __init__(self, type_, col=None, cause=None):
69
        msg = 'Missing cast to type '+strings.as_tt(type_)
70
        if col != None: msg += ' on column: '+strings.as_tt(col)
71
        DbException.__init__(self, msg, cause)
72 2523 aaronmk
        self.type = type_
73
        self.col = col
74
75 5576 aaronmk
class EncodingException(ExceptionWithName): pass
76
77 2306 aaronmk
class DuplicateKeyException(ConstraintException): pass
78 13 aaronmk
79 2306 aaronmk
class NullValueException(ConstraintException): pass
80 13 aaronmk
81 3346 aaronmk
class CheckException(ConstraintException): pass
82
83 3109 aaronmk
class InvalidValueException(ExceptionWithValue): pass
84 2239 aaronmk
85 2945 aaronmk
class DuplicateException(ExceptionWithNameType): pass
86 2143 aaronmk
87 3419 aaronmk
class DoesNotExistException(ExceptionWithNameType): pass
88
89 89 aaronmk
class EmptyRowException(DbException): pass
90
91 865 aaronmk
##### Warnings
92
93
class DbWarning(UserWarning): pass
94
95 1930 aaronmk
##### Result retrieval
96
97
def col_names(cur): return (col[0] for col in cur.description)
98
99
def rows(cur): return iter(lambda: cur.fetchone(), None)
100
101
def consume_rows(cur):
102
    '''Used to fetch all rows so result will be cached'''
103
    iters.consume_iter(rows(cur))
104
105
def next_row(cur): return rows(cur).next()
106
107
def row(cur):
108
    row_ = next_row(cur)
109
    consume_rows(cur)
110
    return row_
111
112
def next_value(cur): return next_row(cur)[0]
113
114
def value(cur): return row(cur)[0]
115
116
def values(cur): return iters.func_iter(lambda: next_value(cur))
117
118
def value_or_none(cur):
119
    try: return value(cur)
120
    except StopIteration: return None
121
122 2762 aaronmk
##### Escaping
123 2101 aaronmk
124 2573 aaronmk
def esc_name_by_module(module, name):
125
    if module == 'psycopg2' or module == None: quote = '"'
126 2101 aaronmk
    elif module == 'MySQLdb': quote = '`'
127
    else: raise NotImplementedError("Can't escape name for "+module+' database')
128 2500 aaronmk
    return sql_gen.esc_name(name, quote)
129 2101 aaronmk
130
def esc_name_by_engine(engine, name, **kw_args):
131
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
132
133
def esc_name(db, name, **kw_args):
134
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
135
136
def qual_name(db, schema, table):
137
    def esc_name_(name): return esc_name(db, name)
138
    table = esc_name_(table)
139
    if schema != None: return esc_name_(schema)+'.'+table
140
    else: return table
141
142 1869 aaronmk
##### Database connections
143 1849 aaronmk
144 2097 aaronmk
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
145 1926 aaronmk
146 1869 aaronmk
db_engines = {
147
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
148
    'PostgreSQL': ('psycopg2', {}),
149
}
150
151
DatabaseErrors_set = set([DbException])
152
DatabaseErrors = tuple(DatabaseErrors_set)
153
154
def _add_module(module):
155
    DatabaseErrors_set.add(module.DatabaseError)
156
    global DatabaseErrors
157
    DatabaseErrors = tuple(DatabaseErrors_set)
158
159
def db_config_str(db_config):
160
    return db_config['engine']+' database '+db_config['database']
161
162 2448 aaronmk
log_debug_none = lambda msg, level=2: None
163 1901 aaronmk
164 1849 aaronmk
class DbConn:
165 2923 aaronmk
    def __init__(self, db_config, autocommit=True, caching=True,
166 3183 aaronmk
        log_debug=log_debug_none, debug_temp=False, src=None):
167 2915 aaronmk
        '''
168
        @param debug_temp Whether temporary objects should instead be permanent.
169
            This assists in debugging the internal objects used by the program.
170 3183 aaronmk
        @param src In autocommit mode, will be included in a comment in every
171
            query, to help identify the data source in pg_stat_activity.
172 2915 aaronmk
        '''
173 1869 aaronmk
        self.db_config = db_config
174 2190 aaronmk
        self.autocommit = autocommit
175
        self.caching = caching
176 1901 aaronmk
        self.log_debug = log_debug
177 2193 aaronmk
        self.debug = log_debug != log_debug_none
178 2915 aaronmk
        self.debug_temp = debug_temp
179 3183 aaronmk
        self.src = src
180 3074 aaronmk
        self.autoanalyze = False
181 3269 aaronmk
        self.autoexplain = False
182
        self.profile_row_ct = None
183 1869 aaronmk
184 3124 aaronmk
        self._savepoint = 0
185 3120 aaronmk
        self._reset()
186 1869 aaronmk
187
    def __getattr__(self, name):
188
        if name == '__dict__': raise Exception('getting __dict__')
189
        if name == 'db': return self._db()
190
        else: raise AttributeError()
191
192
    def __getstate__(self):
193
        state = copy.copy(self.__dict__) # shallow copy
194 1915 aaronmk
        state['log_debug'] = None # don't pickle the debug callback
195 1869 aaronmk
        state['_DbConn__db'] = None # don't pickle the connection
196
        return state
197
198 3118 aaronmk
    def clear_cache(self): self.query_results = {}
199
200 3120 aaronmk
    def _reset(self):
201 3118 aaronmk
        self.clear_cache()
202 3124 aaronmk
        assert self._savepoint == 0
203 3118 aaronmk
        self._notices_seen = set()
204
        self.__db = None
205
206 2165 aaronmk
    def connected(self): return self.__db != None
207
208 3116 aaronmk
    def close(self):
209 3119 aaronmk
        if not self.connected(): return
210
211 3135 aaronmk
        # Record that the automatic transaction is now closed
212 3136 aaronmk
        self._savepoint -= 1
213 3135 aaronmk
214 3119 aaronmk
        self.db.close()
215 3120 aaronmk
        self._reset()
216 3116 aaronmk
217 3125 aaronmk
    def reconnect(self):
218
        # Do not do this in test mode as it would roll back everything
219
        if self.autocommit: self.close()
220
        # Connection will be reopened automatically on first query
221
222 1869 aaronmk
    def _db(self):
223
        if self.__db == None:
224
            # Process db_config
225
            db_config = self.db_config.copy() # don't modify input!
226 2097 aaronmk
            schemas = db_config.pop('schemas', None)
227 1869 aaronmk
            module_name, mappings = db_engines[db_config.pop('engine')]
228
            module = __import__(module_name)
229
            _add_module(module)
230
            for orig, new in mappings.iteritems():
231
                try: util.rename_key(db_config, orig, new)
232
                except KeyError: pass
233
234
            # Connect
235
            self.__db = module.connect(**db_config)
236
237 3161 aaronmk
            # Record that a transaction is already open
238
            self._savepoint += 1
239
240 1869 aaronmk
            # Configure connection
241 2906 aaronmk
            if hasattr(self.db, 'set_isolation_level'):
242
                import psycopg2.extensions
243
                self.db.set_isolation_level(
244
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
245 2101 aaronmk
            if schemas != None:
246 2893 aaronmk
                search_path = [self.esc_name(s) for s in schemas.split(',')]
247
                search_path.append(value(run_query(self, 'SHOW search_path',
248
                    log_level=4)))
249
                run_query(self, 'SET search_path TO '+(','.join(search_path)),
250
                    log_level=3)
251 1869 aaronmk
252
        return self.__db
253 1889 aaronmk
254 1891 aaronmk
    class DbCursor(Proxy):
255 1927 aaronmk
        def __init__(self, outer):
256 1891 aaronmk
            Proxy.__init__(self, outer.db.cursor())
257 2191 aaronmk
            self.outer = outer
258 1927 aaronmk
            self.query_results = outer.query_results
259 1894 aaronmk
            self.query_lookup = None
260 1891 aaronmk
            self.result = []
261 1889 aaronmk
262 2802 aaronmk
        def execute(self, query):
263 2764 aaronmk
            self._is_insert = query.startswith('INSERT')
264 2797 aaronmk
            self.query_lookup = query
265 2148 aaronmk
            try:
266 3162 aaronmk
                try: cur = self.inner.execute(query)
267 2802 aaronmk
                finally: self.query = get_cur_query(self.inner, query)
268 1904 aaronmk
            except Exception, e:
269
                self.result = e # cache the exception as the result
270
                self._cache_result()
271
                raise
272 3004 aaronmk
273
            # Always cache certain queries
274 3183 aaronmk
            query = sql_gen.lstrip(query)
275 3004 aaronmk
            if query.startswith('CREATE') or query.startswith('ALTER'):
276 3007 aaronmk
                # structural changes
277 3040 aaronmk
                # Rest of query must be unique in the face of name collisions,
278
                # so don't cache ADD COLUMN unless it has distinguishing comment
279
                if query.find('ADD COLUMN') < 0 or query.endswith('*/'):
280 3007 aaronmk
                    self._cache_result()
281 3004 aaronmk
            elif self.rowcount == 0 and query.startswith('SELECT'): # empty
282 2800 aaronmk
                consume_rows(self) # fetch all rows so result will be cached
283 3004 aaronmk
284 2762 aaronmk
            return cur
285 1894 aaronmk
286 1891 aaronmk
        def fetchone(self):
287
            row = self.inner.fetchone()
288 1899 aaronmk
            if row != None: self.result.append(row)
289
            # otherwise, fetched all rows
290 1904 aaronmk
            else: self._cache_result()
291
            return row
292
293
        def _cache_result(self):
294 2948 aaronmk
            # For inserts that return a result set, don't cache result set since
295
            # inserts are not idempotent. Other non-SELECT queries don't have
296
            # their result set read, so only exceptions will be cached (an
297
            # invalid query will always be invalid).
298 1930 aaronmk
            if self.query_results != None and (not self._is_insert
299 1906 aaronmk
                or isinstance(self.result, Exception)):
300
301 1894 aaronmk
                assert self.query_lookup != None
302 1916 aaronmk
                self.query_results[self.query_lookup] = self.CacheCursor(
303
                    util.dict_subset(dicts.AttrsDictView(self),
304
                    ['query', 'result', 'rowcount', 'description']))
305 1906 aaronmk
306 1916 aaronmk
        class CacheCursor:
307
            def __init__(self, cached_result): self.__dict__ = cached_result
308
309 1927 aaronmk
            def execute(self, *args, **kw_args):
310 1916 aaronmk
                if isinstance(self.result, Exception): raise self.result
311
                # otherwise, result is a rows list
312
                self.iter = iter(self.result)
313
314
            def fetchone(self):
315
                try: return self.iter.next()
316
                except StopIteration: return None
317 1891 aaronmk
318 2212 aaronmk
    def esc_value(self, value):
319 2663 aaronmk
        try: str_ = self.mogrify('%s', [value])
320
        except NotImplementedError, e:
321
            module = util.root_module(self.db)
322
            if module == 'MySQLdb':
323
                import _mysql
324
                str_ = _mysql.escape_string(value)
325
            else: raise e
326 2374 aaronmk
        return strings.to_unicode(str_)
327 2212 aaronmk
328 2347 aaronmk
    def esc_name(self, name): return esc_name(self, name) # calls global func
329
330 2814 aaronmk
    def std_code(self, str_):
331
        '''Standardizes SQL code.
332
        * Ensures that string literals are prefixed by `E`
333
        '''
334
        if str_.startswith("'"): str_ = 'E'+str_
335
        return str_
336
337 2665 aaronmk
    def can_mogrify(self):
338 2663 aaronmk
        module = util.root_module(self.db)
339 2665 aaronmk
        return module == 'psycopg2'
340 2663 aaronmk
341 2665 aaronmk
    def mogrify(self, query, params=None):
342
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
343
        else: raise NotImplementedError("Can't mogrify query")
344
345 5579 aaronmk
    def set_encoding(self, encoding):
346
        encoding_str = sql_gen.Literal(encoding)
347
        run_query(self, 'SET NAMES '+encoding_str.to_str(self))
348
349 2671 aaronmk
    def print_notices(self):
350 2725 aaronmk
        if hasattr(self.db, 'notices'):
351
            for msg in self.db.notices:
352
                if msg not in self._notices_seen:
353
                    self._notices_seen.add(msg)
354
                    self.log_debug(msg, level=2)
355 2671 aaronmk
356 2793 aaronmk
    def run_query(self, query, cacheable=False, log_level=2,
357 2464 aaronmk
        debug_msg_ref=None):
358 2445 aaronmk
        '''
359 2464 aaronmk
        @param log_ignore_excs The log_level will be increased by 2 if the query
360
            throws one of these exceptions.
361 2664 aaronmk
        @param debug_msg_ref If specified, the log message will be returned in
362
            this instead of being output. This allows you to filter log messages
363
            depending on the result of the query.
364 2445 aaronmk
        '''
365 2167 aaronmk
        assert query != None
366
367 3183 aaronmk
        if self.autocommit and self.src != None:
368 3206 aaronmk
            query = sql_gen.esc_comment(self.src)+'\t'+query
369 3183 aaronmk
370 2047 aaronmk
        if not self.caching: cacheable = False
371 1903 aaronmk
        used_cache = False
372 2664 aaronmk
373 3242 aaronmk
        if self.debug:
374
            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
375 1903 aaronmk
        try:
376 1927 aaronmk
            # Get cursor
377
            if cacheable:
378 3238 aaronmk
                try: cur = self.query_results[query]
379 1927 aaronmk
                except KeyError: cur = self.DbCursor(self)
380 3238 aaronmk
                else: used_cache = True
381 1927 aaronmk
            else: cur = self.db.cursor()
382
383
            # Run query
384 3238 aaronmk
            try: cur.execute(query)
385 3162 aaronmk
            except Exception, e:
386
                _add_cursor_info(e, self, query)
387
                raise
388 3238 aaronmk
            else: self.do_autocommit()
389 1903 aaronmk
        finally:
390 3242 aaronmk
            if self.debug:
391 3244 aaronmk
                profiler.stop(self.profile_row_ct)
392 3242 aaronmk
393
                ## Log or return query
394
395 4491 aaronmk
                query = strings.ustr(get_cur_query(cur, query))
396 3281 aaronmk
                # Put the src comment on a separate line in the log file
397
                query = query.replace('\t', '\n', 1)
398 3239 aaronmk
399 3240 aaronmk
                msg = 'DB query: '
400 3239 aaronmk
401 3240 aaronmk
                if used_cache: msg += 'cache hit'
402
                elif cacheable: msg += 'cache miss'
403
                else: msg += 'non-cacheable'
404 3239 aaronmk
405 3241 aaronmk
                msg += ':\n'+profiler.msg()+'\n'+strings.as_code(query, 'SQL')
406 3240 aaronmk
407 3237 aaronmk
                if debug_msg_ref != None: debug_msg_ref[0] = msg
408
                else: self.log_debug(msg, log_level)
409 3245 aaronmk
410
                self.print_notices()
411 1903 aaronmk
412
        return cur
413 1914 aaronmk
414 2797 aaronmk
    def is_cached(self, query): return query in self.query_results
415 2139 aaronmk
416 2907 aaronmk
    def with_autocommit(self, func):
417 2801 aaronmk
        import psycopg2.extensions
418
419
        prev_isolation_level = self.db.isolation_level
420 2907 aaronmk
        self.db.set_isolation_level(
421
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
422 2683 aaronmk
        try: return func()
423 2801 aaronmk
        finally: self.db.set_isolation_level(prev_isolation_level)
424 2683 aaronmk
425 2139 aaronmk
    def with_savepoint(self, func):
426 3137 aaronmk
        top = self._savepoint == 0
427 2171 aaronmk
        savepoint = 'level_'+str(self._savepoint)
428 3137 aaronmk
429 3272 aaronmk
        if self.debug:
430 3273 aaronmk
            self.log_debug('Begin transaction', level=4)
431 3272 aaronmk
            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
432
433 3160 aaronmk
        # Must happen before running queries so they don't get autocommitted
434
        self._savepoint += 1
435
436 3137 aaronmk
        if top: query = 'START TRANSACTION ISOLATION LEVEL READ COMMITTED'
437
        else: query = 'SAVEPOINT '+savepoint
438
        self.run_query(query, log_level=4)
439
        try:
440
            return func()
441
            if top: self.run_query('COMMIT', log_level=4)
442 2139 aaronmk
        except:
443 3137 aaronmk
            if top: query = 'ROLLBACK'
444
            else: query = 'ROLLBACK TO SAVEPOINT '+savepoint
445
            self.run_query(query, log_level=4)
446
447 2139 aaronmk
            raise
448 2930 aaronmk
        finally:
449
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
450
            # "The savepoint remains valid and can be rolled back to again"
451
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
452 3137 aaronmk
            if not top:
453
                self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
454 2930 aaronmk
455
            self._savepoint -= 1
456
            assert self._savepoint >= 0
457
458 3272 aaronmk
            if self.debug:
459
                profiler.stop(self.profile_row_ct)
460 3273 aaronmk
                self.log_debug('End transaction\n'+profiler.msg(), level=4)
461 3272 aaronmk
462 2930 aaronmk
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
463 2191 aaronmk
464
    def do_autocommit(self):
465
        '''Autocommits if outside savepoint'''
466 3135 aaronmk
        assert self._savepoint >= 1
467
        if self.autocommit and self._savepoint == 1:
468 2924 aaronmk
            self.log_debug('Autocommitting', level=4)
469 2191 aaronmk
            self.db.commit()
470 2643 aaronmk
471 3155 aaronmk
    def col_info(self, col, cacheable=True):
472 6900 aaronmk
        module = util.root_module(self.db)
473 6899 aaronmk
        if module == 'psycopg2':
474 6900 aaronmk
            qual_table = sql_gen.Literal(col.table.to_str(self))
475 6899 aaronmk
            col_name_str = sql_gen.Literal(col.name)
476
            try:
477 6900 aaronmk
                type_, is_array, default, nullable = row(run_query(self, '''\
478 6899 aaronmk
SELECT
479
format_type(COALESCE(NULLIF(typelem, 0), pg_type.oid), -1) AS type
480
, typcategory = 'A' AS type_is_array
481
, pg_attrdef.adsrc AS default
482
, NOT pg_attribute.attnotnull AS nullable
483
FROM pg_attribute
484
LEFT JOIN pg_type ON pg_type.oid = atttypid
485
LEFT JOIN pg_attrdef ON adrelid = attrelid AND adnum = attnum
486
WHERE
487 6901 aaronmk
attrelid = '''+qual_table.to_str(self)+'''::regclass
488
AND attname = '''+col_name_str.to_str(self)+'''
489 6899 aaronmk
'''
490 6902 aaronmk
                    , cacheable=cacheable, log_level=4))
491 6899 aaronmk
            except StopIteration: raise sql_gen.NoUnderlyingTableException(col)
492 6900 aaronmk
            if is_array: type_ = sql_gen.ArrayType(type_)
493 6899 aaronmk
        else:
494
            table = sql_gen.Table('columns', 'information_schema')
495
            cols = [sql_gen.Col('data_type'), sql_gen.Col('udt_name'),
496
                'column_default', sql_gen.Cast('boolean',
497
                sql_gen.Col('is_nullable'))]
498
499
            conds = [('table_name', col.table.name),
500
                ('column_name', strings.ustr(col.name))]
501
            schema = col.table.schema
502
            if schema != None: conds.append(('table_schema', schema))
503
504
            cur = select(self, table, cols, conds, order_by='table_schema',
505
                limit=1, cacheable=cacheable, log_level=4)
506
            try: type_, extra_type, default, nullable = row(cur)
507
            except StopIteration: raise sql_gen.NoUnderlyingTableException(col)
508
            if type_ == 'USER-DEFINED': type_ = extra_type
509
            elif type_ == 'ARRAY':
510
                type_ = sql_gen.ArrayType(strings.remove_prefix('_', extra_type,
511
                    require=True))
512 2643 aaronmk
513 6899 aaronmk
        if default != None: default = sql_gen.as_Code(default, self)
514 2819 aaronmk
        return sql_gen.TypedCol(col.name, type_, default, nullable)
515 2917 aaronmk
516
    def TempFunction(self, name):
517
        if self.debug_temp: schema = None
518
        else: schema = 'pg_temp'
519
        return sql_gen.Function(name, schema)
520 1849 aaronmk
521 1869 aaronmk
connect = DbConn
522
523 832 aaronmk
##### Recoverable querying
524 15 aaronmk
525 5577 aaronmk
def parse_exception(db, e, recover=False):
526
    msg = strings.ustr(e.args[0])
527
    msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)
528
529
    match = re.match(r'^invalid byte sequence for encoding "(.+?)":', msg)
530
    if match:
531
        encoding, = match.groups()
532
        raise EncodingException(encoding, e)
533
534 5763 aaronmk
    def make_DuplicateKeyException(constraint, e):
535 5577 aaronmk
        cols = []
536
        cond = None
537
        if recover: # need auto-rollback to run index_cols()
538
            try:
539
                cols = index_cols(db, constraint)
540
                cond = index_cond(db, constraint)
541
            except NotImplementedError: pass
542 5763 aaronmk
        return DuplicateKeyException(constraint, cond, cols, e)
543 5577 aaronmk
544 5763 aaronmk
    match = re.match(r'^duplicate key value violates unique constraint "(.+?)"',
545
        msg)
546
    if match:
547
        constraint, = match.groups()
548
        raise make_DuplicateKeyException(constraint, e)
549
550 5764 aaronmk
    match = re.match(r'^could not create unique index "(.+?)"\n'
551
        r'DETAIL:  Key .+? is duplicated', msg)
552
    if match:
553
        constraint, = match.groups()
554 5823 aaronmk
        raise DuplicateKeyException(constraint, None, [], e)
555 5764 aaronmk
556 5577 aaronmk
    match = re.match(r'^null value in column "(.+?)" violates not-null'
557
        r' constraint', msg)
558
    if match:
559
        col, = match.groups()
560
        raise NullValueException('NOT NULL', None, [col], e)
561
562
    match = re.match(r'^new row for relation "(.+?)" violates check '
563
        r'constraint "(.+?)"', msg)
564
    if match:
565
        table, constraint = match.groups()
566
        constraint = sql_gen.Col(constraint, table)
567
        cond = None
568
        if recover: # need auto-rollback to run constraint_cond()
569
            try: cond = constraint_cond(db, constraint)
570
            except NotImplementedError: pass
571
        raise CheckException(constraint.to_str(db), cond, [], e)
572
573 5723 aaronmk
    match = re.match(r'^(?:invalid input (?:syntax|value)\b[^:]*'
574 5717 aaronmk
        r'|.+? out of range)(?:: "(.+?)")?', msg)
575 5577 aaronmk
    if match:
576
        value, = match.groups()
577 5717 aaronmk
        value = util.do_ignore_none(strings.to_unicode, value)
578
        raise InvalidValueException(value, e)
579 5577 aaronmk
580
    match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
581
        r'is of type', msg)
582
    if match:
583
        col, type_ = match.groups()
584
        raise MissingCastException(type_, col, e)
585
586
    match = re.match(r'^could not determine polymorphic type because '
587
        r'input has type "unknown"', msg)
588
    if match: raise MissingCastException('text', None, e)
589
590 5724 aaronmk
    match = re.match(r'^.+? types (.+?) and (.+?) cannot be matched', msg)
591
    if match:
592
        type0, type1 = match.groups()
593
        raise MissingCastException(type0, None, e)
594 5577 aaronmk
595 5707 aaronmk
    typed_name_re = r'^(\S+) "?(.+?)"?(?: of relation ".+?")?'
596 5577 aaronmk
597
    match = re.match(typed_name_re+r'.*? already exists', msg)
598
    if match:
599
        type_, name = match.groups()
600
        raise DuplicateException(type_, name, e)
601
602
    match = re.match(r'more than one (\S+) named ""(.+?)""', msg)
603
    if match:
604
        type_, name = match.groups()
605
        raise DuplicateException(type_, name, e)
606
607
    match = re.match(typed_name_re+r' does not exist', msg)
608
    if match:
609
        type_, name = match.groups()
610 5683 aaronmk
        if type_ == 'function':
611 5709 aaronmk
            match = re.match(r'^(.+?)\(.*\)$', name)
612 5710 aaronmk
            if match: # includes params, so is call rather than cast to regproc
613 5708 aaronmk
                function_name, = match.groups()
614 5711 aaronmk
                func = sql_gen.Function(function_name)
615 5712 aaronmk
                if function_exists(db, func) and msg.find('CAST') < 0:
616 5683 aaronmk
                    # not found only because of a missing cast
617 5715 aaronmk
                    type_ = function_param0_type(db, func)
618 6301 aaronmk
                    col = None
619 5792 aaronmk
                    if type_ == 'anyelement': type_ = 'text'
620 6301 aaronmk
                    elif type_ == 'hstore': # cast just the value param
621
                        type_ = 'text'
622
                        col = 'value'
623
                    raise MissingCastException(type_, col, e)
624 5577 aaronmk
        raise DoesNotExistException(type_, name, e)
625
626
    raise # no specific exception raised
627
628 2139 aaronmk
def with_savepoint(db, func): return db.with_savepoint(func)
629 11 aaronmk
630 2791 aaronmk
def run_query(db, query, recover=None, cacheable=False, log_level=2,
631
    log_ignore_excs=None, **kw_args):
632 2794 aaronmk
    '''For params, see DbConn.run_query()'''
633 830 aaronmk
    if recover == None: recover = False
634 2464 aaronmk
    if log_ignore_excs == None: log_ignore_excs = ()
635
    log_ignore_excs = tuple(log_ignore_excs)
636 3236 aaronmk
    debug_msg_ref = [None]
637 830 aaronmk
638 3267 aaronmk
    query = with_explain_comment(db, query)
639 3258 aaronmk
640 2148 aaronmk
    try:
641 2464 aaronmk
        try:
642 2794 aaronmk
            def run(): return db.run_query(query, cacheable, log_level,
643 2793 aaronmk
                debug_msg_ref, **kw_args)
644 2796 aaronmk
            if recover and not db.is_cached(query):
645 2464 aaronmk
                return with_savepoint(db, run)
646
            else: return run() # don't need savepoint if cached
647 5890 aaronmk
        except Exception, e:
648
            # Give failed EXPLAIN approximately the log_level of its query
649
            if query.startswith('EXPLAIN'): log_level -= 1
650
651
            parse_exception(db, e, recover)
652 2464 aaronmk
    except log_ignore_excs:
653
        log_level += 2
654
        raise
655
    finally:
656 3236 aaronmk
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
657 830 aaronmk
658 832 aaronmk
##### Basic queries
659
660 3256 aaronmk
def is_explainable(query):
661
    # See <http://www.postgresql.org/docs/8.3/static/sql-explain.html#AEN57749>
662 3257 aaronmk
    return re.match(r'^(?:SELECT|INSERT|UPDATE|DELETE|VALUES|EXECUTE|DECLARE)\b'
663
        , query)
664 3256 aaronmk
665 3263 aaronmk
def explain(db, query, **kw_args):
666
    '''
667
    For params, see run_query().
668
    '''
669 3267 aaronmk
    kw_args.setdefault('log_level', 4)
670 3263 aaronmk
671 3750 aaronmk
    return strings.ustr(strings.join_lines(values(run_query(db,
672
        'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))
673 3256 aaronmk
        # not a higher log_level because it's useful to see what query is being
674
        # run before it's executed, which EXPLAIN effectively provides
675
676 3265 aaronmk
def has_comment(query): return query.endswith('*/')
677
678
def with_explain_comment(db, query, **kw_args):
679 3269 aaronmk
    if db.autoexplain and not has_comment(query) and is_explainable(query):
680 3265 aaronmk
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
681
            +explain(db, query, **kw_args))
682
    return query
683
684 2153 aaronmk
def next_version(name):
685 2163 aaronmk
    version = 1 # first existing name was version 0
686 2586 aaronmk
    match = re.match(r'^(.*)#(\d+)$', name)
687 2153 aaronmk
    if match:
688 2586 aaronmk
        name, version = match.groups()
689
        version = int(version)+1
690 2932 aaronmk
    return sql_gen.concat(name, '#'+str(version))
691 2153 aaronmk
692 2899 aaronmk
def lock_table(db, table, mode):
693
    table = sql_gen.as_Table(table)
694
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
695
696 3303 aaronmk
def run_query_into(db, query, into=None, add_pkey_=False, **kw_args):
697 2085 aaronmk
    '''Outputs a query to a temp table.
698
    For params, see run_query().
699
    '''
700 2789 aaronmk
    if into == None: return run_query(db, query, **kw_args)
701 2790 aaronmk
702
    assert isinstance(into, sql_gen.Table)
703
704 2992 aaronmk
    into.is_temp = True
705 3008 aaronmk
    # "temporary tables cannot specify a schema name", so remove schema
706
    into.schema = None
707 2992 aaronmk
708 2790 aaronmk
    kw_args['recover'] = True
709 2945 aaronmk
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
710 2790 aaronmk
711 2916 aaronmk
    temp = not db.debug_temp # tables are permanent in debug_temp mode
712 2790 aaronmk
713
    # Create table
714
    while True:
715
        create_query = 'CREATE'
716
        if temp: create_query += ' TEMP'
717
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
718 2385 aaronmk
719 2790 aaronmk
        try:
720
            cur = run_query(db, create_query, **kw_args)
721
                # CREATE TABLE AS sets rowcount to # rows in query
722
            break
723 2945 aaronmk
        except DuplicateException, e:
724 2790 aaronmk
            into.name = next_version(into.name)
725
            # try again with next version of name
726
727 5967 aaronmk
    if add_pkey_: add_pkey_or_index(db, into, warn=True)
728 3075 aaronmk
729
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
730
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
731
    # table is going to be used in complex queries, it is wise to run ANALYZE on
732
    # the temporary table after it is populated."
733
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
734
    # If into is not a temp table, ANALYZE is useful but not required.
735 3073 aaronmk
    analyze(db, into)
736 2790 aaronmk
737
    return cur
738 2085 aaronmk
739 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
740
741 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
742
743 3420 aaronmk
def mk_select(db, tables=None, fields=None, conds=None, distinct_on=[],
744 3494 aaronmk
    limit=None, start=None, order_by=order_by_pkey, default_table=None,
745
    explain=True):
746 1981 aaronmk
    '''
747 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
748 2280 aaronmk
        together, with tables after the first being sql_gen.Join objects
749 1981 aaronmk
    @param fields Use None to select all fields in the table
750 2377 aaronmk
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
751 2379 aaronmk
        * container can be any iterable type
752 2399 aaronmk
        * compare_left_side: sql_gen.Code|str (for col name)
753
        * compare_right_side: sql_gen.ValueCond|literal value
754 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
755
        use all columns
756 2786 aaronmk
    @return query
757 1981 aaronmk
    '''
758 2315 aaronmk
    # Parse tables param
759 2964 aaronmk
    tables = lists.mk_seq(tables)
760 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
761 2315 aaronmk
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
762 2121 aaronmk
763 2315 aaronmk
    # Parse other params
764 2376 aaronmk
    if conds == None: conds = []
765 2650 aaronmk
    elif dicts.is_dict(conds): conds = conds.items()
766 2379 aaronmk
    conds = list(conds) # don't modify input! (list() copies input)
767 3129 aaronmk
    assert limit == None or isinstance(limit, (int, long))
768
    assert start == None or isinstance(start, (int, long))
769 5527 aaronmk
    if limit == 0: order_by = None
770 2315 aaronmk
    if order_by is order_by_pkey:
771 5641 aaronmk
        if lists.is_seq(distinct_on) and distinct_on: order_by = distinct_on[0]
772
        elif table0 != None: order_by = table_order_by(db, table0, recover=True)
773
        else: order_by = None
774 865 aaronmk
775 2315 aaronmk
    query = 'SELECT'
776 2056 aaronmk
777 2315 aaronmk
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
778 2056 aaronmk
779 2200 aaronmk
    # DISTINCT ON columns
780 2233 aaronmk
    if distinct_on != []:
781 2467 aaronmk
        query += '\nDISTINCT'
782 2254 aaronmk
        if distinct_on is not distinct_on_all:
783 2200 aaronmk
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
784
785
    # Columns
786 3185 aaronmk
    if query.find('\n') >= 0: whitespace = '\n'
787
    else: whitespace = ' '
788
    if fields == None: query += whitespace+'*'
789 2765 aaronmk
    else:
790
        assert fields != []
791 3185 aaronmk
        if len(fields) > 1: whitespace = '\n'
792
        query += whitespace+('\n, '.join(map(parse_col, fields)))
793 2200 aaronmk
794
    # Main table
795 3185 aaronmk
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
796
    else: whitespace = ' '
797 3420 aaronmk
    if table0 != None: query += whitespace+'FROM '+table0.to_str(db)
798 865 aaronmk
799 2122 aaronmk
    # Add joins
800 2271 aaronmk
    left_table = table0
801 2263 aaronmk
    for join_ in tables:
802
        table = join_.table
803 2238 aaronmk
804 2343 aaronmk
        # Parse special values
805
        if join_.type_ is sql_gen.filter_out: # filter no match
806 2376 aaronmk
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
807 2853 aaronmk
                sql_gen.CompareCond(None, '~=')))
808 2343 aaronmk
809 2467 aaronmk
        query += '\n'+join_.to_str(db, left_table)
810 2122 aaronmk
811
        left_table = table
812
813 865 aaronmk
    missing = True
814 2376 aaronmk
    if conds != []:
815 2576 aaronmk
        if len(conds) == 1: whitespace = ' '
816
        else: whitespace = '\n'
817 2578 aaronmk
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
818
            .to_str(db) for l, r in conds], 'WHERE')
819 2227 aaronmk
    if order_by != None:
820 5642 aaronmk
        query += '\nORDER BY '+sql_gen.as_Col(order_by).to_str(db)
821 3297 aaronmk
    if limit != None: query += '\nLIMIT '+str(limit)
822 865 aaronmk
    if start != None:
823 2467 aaronmk
        if start != 0: query += '\nOFFSET '+str(start)
824 865 aaronmk
825 3494 aaronmk
    if explain: query = with_explain_comment(db, query)
826 3266 aaronmk
827 2786 aaronmk
    return query
828 11 aaronmk
829 2054 aaronmk
def select(db, *args, **kw_args):
830
    '''For params, see mk_select() and run_query()'''
831
    recover = kw_args.pop('recover', None)
832
    cacheable = kw_args.pop('cacheable', True)
833 2442 aaronmk
    log_level = kw_args.pop('log_level', 2)
834 2054 aaronmk
835 2791 aaronmk
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
836
        log_level=log_level)
837 2054 aaronmk
838 2788 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
839 3181 aaronmk
    embeddable=False, ignore=False, src=None):
840 1960 aaronmk
    '''
841
    @param returning str|None An inserted column (such as pkey) to return
842 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
843 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
844
        query will be fully cached, not just if it raises an exception.
845 3009 aaronmk
    @param ignore Whether to ignore duplicate keys.
846 3181 aaronmk
    @param src Will be included in the name of any created function, to help
847
        identify the data source in pg_stat_activity.
848 1960 aaronmk
    '''
849 2754 aaronmk
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
850 2318 aaronmk
    if cols == []: cols = None # no cols (all defaults) = unknown col names
851 3010 aaronmk
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
852 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
853 2327 aaronmk
    if returning != None: returning = sql_gen.as_Col(returning, table)
854 2063 aaronmk
855 2497 aaronmk
    first_line = 'INSERT INTO '+table.to_str(db)
856 2063 aaronmk
857 3009 aaronmk
    def mk_insert(select_query):
858
        query = first_line
859 3014 aaronmk
        if cols != None:
860
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
861 3009 aaronmk
        query += '\n'+select_query
862
863
        if returning != None:
864
            returning_name_col = sql_gen.to_name_only_col(returning)
865
            query += '\nRETURNING '+returning_name_col.to_str(db)
866
867
        return query
868 2063 aaronmk
869 3489 aaronmk
    return_type = sql_gen.CustomCode('unknown')
870
    if returning != None: return_type = sql_gen.ColType(returning)
871 3017 aaronmk
872 3009 aaronmk
    if ignore:
873 3017 aaronmk
        # Always return something to set the correct rowcount
874
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
875
876 3009 aaronmk
        embeddable = True # must use function
877 3010 aaronmk
878 3450 aaronmk
        if cols == None: row = [sql_gen.Col(sql_gen.all_cols, 'row')]
879
        else: row = [sql_gen.Col(c.name, 'row') for c in cols]
880 3092 aaronmk
881 3484 aaronmk
        query = sql_gen.RowExcIgnore(sql_gen.RowType(table), select_query,
882 3497 aaronmk
            sql_gen.ReturnQuery(mk_insert(sql_gen.Values(row).to_str(db))),
883
            cols)
884 3009 aaronmk
    else: query = mk_insert(select_query)
885
886 2070 aaronmk
    if embeddable:
887
        # Create function
888 2513 aaronmk
        function_name = sql_gen.clean_name(first_line)
889 3181 aaronmk
        if src != None: function_name = src+': '+function_name
890 2189 aaronmk
        while True:
891
            try:
892 3451 aaronmk
                func = db.TempFunction(function_name)
893 3489 aaronmk
                def_ = sql_gen.FunctionDef(func, sql_gen.SetOf(return_type),
894
                    query)
895 2194 aaronmk
896 3443 aaronmk
                run_query(db, def_.to_str(db), recover=True, cacheable=True,
897 2945 aaronmk
                    log_ignore_excs=(DuplicateException,))
898 2189 aaronmk
                break # this version was successful
899 2945 aaronmk
            except DuplicateException, e:
900 2189 aaronmk
                function_name = next_version(function_name)
901
                # try again with next version of name
902 2070 aaronmk
903 2337 aaronmk
        # Return query that uses function
904 3009 aaronmk
        cols = None
905
        if returning != None: cols = [returning]
906 3451 aaronmk
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(func), cols)
907
            # AS clause requires function alias
908 3298 aaronmk
        return mk_select(db, func_table, order_by=None)
909 2070 aaronmk
910 2787 aaronmk
    return query
911 2066 aaronmk
912 3074 aaronmk
def insert_select(db, table, *args, **kw_args):
913 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
914 2386 aaronmk
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
915
        values in
916 2072 aaronmk
    '''
917 3141 aaronmk
    returning = kw_args.get('returning', None)
918
    ignore = kw_args.get('ignore', False)
919
920 2386 aaronmk
    into = kw_args.pop('into', None)
921
    if into != None: kw_args['embeddable'] = True
922 2066 aaronmk
    recover = kw_args.pop('recover', None)
923 3141 aaronmk
    if ignore: recover = True
924 2066 aaronmk
    cacheable = kw_args.pop('cacheable', True)
925 2673 aaronmk
    log_level = kw_args.pop('log_level', 2)
926 2066 aaronmk
927 3141 aaronmk
    rowcount_only = ignore and returning == None # keep NULL rows on server
928
    if rowcount_only: into = sql_gen.Table('rowcount')
929
930 3074 aaronmk
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
931
        into, recover=recover, cacheable=cacheable, log_level=log_level)
932 3141 aaronmk
    if rowcount_only: empty_temp(db, into)
933 3074 aaronmk
    autoanalyze(db, table)
934
    return cur
935 2063 aaronmk
936 2738 aaronmk
default = sql_gen.default # tells insert() to use the default value for a column
937 2066 aaronmk
938 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
939 2085 aaronmk
    '''For params, see insert_select()'''
940 5050 aaronmk
    ignore = kw_args.pop('ignore', False)
941 5094 aaronmk
    if ignore: kw_args.setdefault('recover', True)
942 5050 aaronmk
943 1960 aaronmk
    if lists.is_seq(row): cols = None
944
    else:
945
        cols = row.keys()
946
        row = row.values()
947 2738 aaronmk
    row = list(row) # ensure that "== []" works
948 1960 aaronmk
949 2738 aaronmk
    if row == []: query = None
950
    else: query = sql_gen.Values(row).to_str(db)
951 1961 aaronmk
952 5050 aaronmk
    try: return insert_select(db, table, cols, query, *args, **kw_args)
953 5057 aaronmk
    except (DuplicateKeyException, NullValueException):
954 5050 aaronmk
        if not ignore: raise
955 5163 aaronmk
        return None
956 11 aaronmk
957 3152 aaronmk
def mk_update(db, table, changes=None, cond=None, in_place=False,
958 3153 aaronmk
    cacheable_=True):
959 2402 aaronmk
    '''
960
    @param changes [(col, new_value),...]
961
        * container can be any iterable type
962
        * col: sql_gen.Code|str (for col name)
963
        * new_value: sql_gen.Code|literal value
964
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
965 3056 aaronmk
    @param in_place If set, locks the table and updates rows in place.
966
        This avoids creating dead rows in PostgreSQL.
967
        * cond must be None
968 3153 aaronmk
    @param cacheable_ Whether column structure information used to generate the
969 3152 aaronmk
        query can be cached
970 2402 aaronmk
    @return str query
971
    '''
972 3057 aaronmk
    table = sql_gen.as_Table(table)
973
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
974
        for c, v in changes]
975
976 3056 aaronmk
    if in_place:
977
        assert cond == None
978 3058 aaronmk
979 5398 aaronmk
        def col_type(col):
980
            return sql_gen.canon_type(db.col_info(
981
                sql_gen.with_default_table(c, table), cacheable_).type)
982
        changes = [(c, v, col_type(c)) for c, v in changes]
983 3065 aaronmk
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
984 5396 aaronmk
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '+t+'\nUSING '
985
            +v.to_str(db) for c, v, t in changes))
986 3058 aaronmk
    else:
987
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
988
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
989
            for c, v in changes))
990
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
991 3056 aaronmk
992 3266 aaronmk
    query = with_explain_comment(db, query)
993
994 2402 aaronmk
    return query
995
996 3074 aaronmk
def update(db, table, *args, **kw_args):
997 2402 aaronmk
    '''For params, see mk_update() and run_query()'''
998
    recover = kw_args.pop('recover', None)
999 3043 aaronmk
    cacheable = kw_args.pop('cacheable', False)
1000 3030 aaronmk
    log_level = kw_args.pop('log_level', 2)
1001 2402 aaronmk
1002 3074 aaronmk
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
1003
        cacheable, log_level=log_level)
1004
    autoanalyze(db, table)
1005
    return cur
1006 2402 aaronmk
1007 3286 aaronmk
def mk_delete(db, table, cond=None):
1008
    '''
1009
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
1010
    @return str query
1011
    '''
1012
    query = 'DELETE FROM '+table.to_str(db)
1013
    if cond != None: query += '\nWHERE '+cond.to_str(db)
1014
1015
    query = with_explain_comment(db, query)
1016
1017
    return query
1018
1019
def delete(db, table, *args, **kw_args):
1020
    '''For params, see mk_delete() and run_query()'''
1021
    recover = kw_args.pop('recover', None)
1022 3295 aaronmk
    cacheable = kw_args.pop('cacheable', True)
1023 3286 aaronmk
    log_level = kw_args.pop('log_level', 2)
1024
1025
    cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
1026
        cacheable, log_level=log_level)
1027
    autoanalyze(db, table)
1028
    return cur
1029
1030 135 aaronmk
def last_insert_id(db):
1031 1849 aaronmk
    module = util.root_module(db.db)
1032 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
1033
    elif module == 'MySQLdb': return db.insert_id()
1034
    else: return None
1035 13 aaronmk
1036 3490 aaronmk
def define_func(db, def_):
1037
    func = def_.function
1038
    while True:
1039
        try:
1040
            run_query(db, def_.to_str(db), recover=True, cacheable=True,
1041
                log_ignore_excs=(DuplicateException,))
1042
            break # successful
1043
        except DuplicateException:
1044 3495 aaronmk
            func.name = next_version(func.name)
1045 3490 aaronmk
            # try again with next version of name
1046
1047 2394 aaronmk
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
1048 2383 aaronmk
    '''Creates a mapping from original column names (which may have collisions)
1049 2415 aaronmk
    to names that will be distinct among the columns' tables.
1050 2383 aaronmk
    This is meant to be used for several tables that are being joined together.
1051 2415 aaronmk
    @param cols The columns to combine. Duplicates will be removed.
1052
    @param into The table for the new columns.
1053 2394 aaronmk
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
1054
        columns will be included in the mapping even if they are not in cols.
1055
        The tables of the provided Col objects will be changed to into, so make
1056
        copies of them if you want to keep the original tables.
1057
    @param as_items Whether to return a list of dict items instead of a dict
1058 2383 aaronmk
    @return dict(orig_col=new_col, ...)
1059
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
1060 2392 aaronmk
        * new_col: sql_gen.Col(orig_col_name, into)
1061
        * All mappings use the into table so its name can easily be
1062 2383 aaronmk
          changed for all columns at once
1063
    '''
1064 2415 aaronmk
    cols = lists.uniqify(cols)
1065
1066 2394 aaronmk
    items = []
1067 2389 aaronmk
    for col in preserve:
1068 2390 aaronmk
        orig_col = copy.copy(col)
1069 2392 aaronmk
        col.table = into
1070 2394 aaronmk
        items.append((orig_col, col))
1071
    preserve = set(preserve)
1072
    for col in cols:
1073 2716 aaronmk
        if col not in preserve:
1074 3750 aaronmk
            items.append((col, sql_gen.Col(strings.ustr(col), into, col.srcs)))
1075 2394 aaronmk
1076
    if not as_items: items = dict(items)
1077
    return items
1078 2383 aaronmk
1079 2393 aaronmk
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
1080 2391 aaronmk
    '''For params, see mk_flatten_mapping()
1081
    @return See return value of mk_flatten_mapping()
1082
    '''
1083 2394 aaronmk
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
1084
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
1085 5523 aaronmk
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
1086
        into=into, add_pkey_=True)
1087 3708 aaronmk
        # don't cache because the temp table will usually be truncated after use
1088 2394 aaronmk
    return dict(items)
1089 2391 aaronmk
1090 3079 aaronmk
##### Database structure introspection
1091 2414 aaronmk
1092 3079 aaronmk
#### Tables
1093
1094 4555 aaronmk
def tables(db, schema_like='public', table_like='%', exact=False,
1095
    cacheable=True):
1096 3079 aaronmk
    if exact: compare = '='
1097
    else: compare = 'LIKE'
1098
1099
    module = util.root_module(db.db)
1100
    if module == 'psycopg2':
1101
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1102
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1103
        return values(select(db, 'pg_tables', ['tablename'], conds,
1104 4555 aaronmk
            order_by='tablename', cacheable=cacheable, log_level=4))
1105 3079 aaronmk
    elif module == 'MySQLdb':
1106
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1107
            , cacheable=True, log_level=4))
1108
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1109
1110 4556 aaronmk
def table_exists(db, table, cacheable=True):
1111 3079 aaronmk
    table = sql_gen.as_Table(table)
1112 4556 aaronmk
    return list(tables(db, table.schema, table.name, True, cacheable)) != []
1113 3079 aaronmk
1114 2426 aaronmk
def table_row_count(db, table, recover=None):
1115 2786 aaronmk
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
1116 3298 aaronmk
        order_by=None), recover=recover, log_level=3))
1117 2426 aaronmk
1118 5337 aaronmk
def table_col_names(db, table, recover=None):
1119 5528 aaronmk
    return list(col_names(select(db, table, limit=0, recover=recover,
1120
        log_level=4)))
1121 2414 aaronmk
1122 5383 aaronmk
def table_cols(db, table, *args, **kw_args):
1123
    return [sql_gen.as_Col(strings.ustr(c), table)
1124
        for c in table_col_names(db, table, *args, **kw_args)]
1125
1126 5521 aaronmk
def table_pkey_index(db, table, recover=None):
1127
    table_str = sql_gen.Literal(table.to_str(db))
1128
    try:
1129
        return sql_gen.Table(value(run_query(db, '''\
1130
SELECT relname
1131
FROM pg_index
1132
JOIN pg_class index ON index.oid = indexrelid
1133
WHERE
1134
indrelid = '''+table_str.to_str(db)+'''::regclass
1135
AND indisprimary
1136
'''
1137
            , recover, cacheable=True, log_level=4)), table.schema)
1138
    except StopIteration: raise DoesNotExistException('primary key', '')
1139
1140 5389 aaronmk
def table_pkey_col(db, table, recover=None):
1141 5061 aaronmk
    table = sql_gen.as_Table(table)
1142
1143 5989 aaronmk
    module = util.root_module(db.db)
1144
    if module == 'psycopg2':
1145
        return sql_gen.Col(index_cols(db, table_pkey_index(db, table,
1146
            recover))[0], table)
1147
    else:
1148
        join_cols = ['table_schema', 'table_name', 'constraint_schema',
1149
            'constraint_name']
1150
        tables = [sql_gen.Table('key_column_usage', 'information_schema'),
1151
            sql_gen.Join(
1152
                sql_gen.Table('table_constraints', 'information_schema'),
1153
                dict(((c, sql_gen.join_same_not_null) for c in join_cols)))]
1154
        cols = [sql_gen.Col('column_name')]
1155
1156
        conds = [('constraint_type', 'PRIMARY KEY'), ('table_name', table.name)]
1157
        schema = table.schema
1158
        if schema != None: conds.append(('table_schema', schema))
1159
        order_by = 'position_in_unique_constraint'
1160
1161
        try: return sql_gen.Col(value(select(db, tables, cols, conds,
1162
            order_by=order_by, limit=1, log_level=4)), table)
1163
        except StopIteration: raise DoesNotExistException('primary key', '')
1164 5389 aaronmk
1165 5990 aaronmk
def table_has_pkey(db, table, recover=None):
1166
    try: table_pkey_col(db, table, recover)
1167
    except DoesNotExistException: return False
1168
    else: return True
1169
1170 5389 aaronmk
def pkey_name(db, table, recover=None):
1171
    '''If no pkey, returns the first column in the table.'''
1172 5392 aaronmk
    return pkey_col(db, table, recover).name
1173 832 aaronmk
1174 5390 aaronmk
def pkey_col(db, table, recover=None):
1175 5391 aaronmk
    '''If no pkey, returns the first column in the table.'''
1176 5392 aaronmk
    try: return table_pkey_col(db, table, recover)
1177 5393 aaronmk
    except DoesNotExistException: return table_cols(db, table, recover)[0]
1178 5128 aaronmk
1179 2559 aaronmk
not_null_col = 'not_null_col'
1180 2340 aaronmk
1181
def table_not_null_col(db, table, recover=None):
1182
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1183 5337 aaronmk
    if not_null_col in table_col_names(db, table, recover): return not_null_col
1184 5388 aaronmk
    else: return pkey_name(db, table, recover)
1185 2340 aaronmk
1186 3348 aaronmk
def constraint_cond(db, constraint):
1187
    module = util.root_module(db.db)
1188
    if module == 'psycopg2':
1189
        table_str = sql_gen.Literal(constraint.table.to_str(db))
1190
        name_str = sql_gen.Literal(constraint.name)
1191
        return value(run_query(db, '''\
1192
SELECT consrc
1193
FROM pg_constraint
1194
WHERE
1195
conrelid = '''+table_str.to_str(db)+'''::regclass
1196
AND conname = '''+name_str.to_str(db)+'''
1197
'''
1198
            , cacheable=True, log_level=4))
1199 5443 aaronmk
    else: raise NotImplementedError("Can't get constraint condition for "
1200
        +module+' database')
1201 3348 aaronmk
1202 5520 aaronmk
def index_exprs(db, index):
1203 3322 aaronmk
    index = sql_gen.as_Table(index)
1204 1909 aaronmk
    module = util.root_module(db.db)
1205
    if module == 'psycopg2':
1206 3322 aaronmk
        qual_index = sql_gen.Literal(index.to_str(db))
1207 5520 aaronmk
        return list(values(run_query(db, '''\
1208 3322 aaronmk
SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
1209
FROM pg_index
1210
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1211 2782 aaronmk
'''
1212
            , cacheable=True, log_level=4)))
1213 5520 aaronmk
    else: raise NotImplementedError()
1214 853 aaronmk
1215 5520 aaronmk
def index_cols(db, index):
1216
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1217
    automatically created. When you don't know whether something is a UNIQUE
1218
    constraint or a UNIQUE index, use this function.'''
1219
    return map(sql_gen.parse_expr_col, index_exprs(db, index))
1220
1221 5445 aaronmk
def index_cond(db, index):
1222
    index = sql_gen.as_Table(index)
1223
    module = util.root_module(db.db)
1224
    if module == 'psycopg2':
1225
        qual_index = sql_gen.Literal(index.to_str(db))
1226
        return value(run_query(db, '''\
1227
SELECT pg_get_expr(indpred, indrelid, true)
1228
FROM pg_index
1229
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1230
'''
1231
            , cacheable=True, log_level=4))
1232
    else: raise NotImplementedError()
1233
1234 5521 aaronmk
def index_order_by(db, index):
1235
    return sql_gen.CustomCode(', '.join(index_exprs(db, index)))
1236
1237
def table_cluster_on(db, table, recover=None):
1238
    '''
1239
    @return The table's cluster index, or its pkey if none is set
1240
    '''
1241
    table_str = sql_gen.Literal(table.to_str(db))
1242
    try:
1243
        return sql_gen.Table(value(run_query(db, '''\
1244
SELECT relname
1245
FROM pg_index
1246
JOIN pg_class index ON index.oid = indexrelid
1247
WHERE
1248
indrelid = '''+table_str.to_str(db)+'''::regclass
1249
AND indisclustered
1250
'''
1251
            , recover, cacheable=True, log_level=4)), table.schema)
1252
    except StopIteration: return table_pkey_index(db, table, recover)
1253
1254
def table_order_by(db, table, recover=None):
1255 5525 aaronmk
    if table.order_by == None:
1256
        try: table.order_by = index_order_by(db, table_cluster_on(db, table,
1257
            recover))
1258
        except DoesNotExistException: pass
1259
    return table.order_by
1260 5521 aaronmk
1261 3079 aaronmk
#### Functions
1262
1263
def function_exists(db, function):
1264 3423 aaronmk
    qual_function = sql_gen.Literal(function.to_str(db))
1265
    try:
1266 3425 aaronmk
        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
1267
            recover=True, cacheable=True, log_level=4)
1268 3423 aaronmk
    except DoesNotExistException: return False
1269 4146 aaronmk
    except DuplicateException: return True # overloaded function
1270 3423 aaronmk
    else: return True
1271 3079 aaronmk
1272 5713 aaronmk
def function_param0_type(db, function):
1273
    qual_function = sql_gen.Literal(function.to_str(db))
1274
    return value(run_query(db, '''\
1275
SELECT proargtypes[0]::regtype
1276
FROM pg_proc
1277
WHERE oid = '''+qual_function.to_str(db)+'''::regproc
1278
'''
1279
        , cacheable=True, log_level=4))
1280
1281 3079 aaronmk
##### Structural changes
1282
1283
#### Columns
1284
1285 5020 aaronmk
def add_col(db, table, col, comment=None, if_not_exists=False, **kw_args):
1286 3079 aaronmk
    '''
1287
    @param col TypedCol Name may be versioned, so be sure to propagate any
1288
        renaming back to any source column for the TypedCol.
1289
    @param comment None|str SQL comment used to distinguish columns of the same
1290
        name from each other when they contain different data, to allow the
1291
        ADD COLUMN query to be cached. If not set, query will not be cached.
1292
    '''
1293
    assert isinstance(col, sql_gen.TypedCol)
1294
1295
    while True:
1296
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1297
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1298
1299
        try:
1300
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1301
            break
1302
        except DuplicateException:
1303 5020 aaronmk
            if if_not_exists: raise
1304 3079 aaronmk
            col.name = next_version(col.name)
1305
            # try again with next version of name
1306
1307
def add_not_null(db, col):
1308
    table = col.table
1309
    col = sql_gen.to_name_only_col(col)
1310
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1311
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1312
1313 4443 aaronmk
def drop_not_null(db, col):
1314
    table = col.table
1315
    col = sql_gen.to_name_only_col(col)
1316
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1317
        +col.to_str(db)+' DROP NOT NULL', cacheable=True, log_level=3)
1318
1319 2096 aaronmk
row_num_col = '_row_num'
1320
1321 4997 aaronmk
row_num_col_def = sql_gen.TypedCol('', 'serial', nullable=False,
1322 3079 aaronmk
    constraints='PRIMARY KEY')
1323
1324 4997 aaronmk
def add_row_num(db, table, name=row_num_col):
1325
    '''Adds a row number column to a table. Its definition is in
1326
    row_num_col_def. It will be the primary key.'''
1327
    col_def = copy.copy(row_num_col_def)
1328
    col_def.name = name
1329 5021 aaronmk
    add_col(db, table, col_def, comment='', if_not_exists=True, log_level=3)
1330 3079 aaronmk
1331
#### Indexes
1332
1333
def add_pkey(db, table, cols=None, recover=None):
1334
    '''Adds a primary key.
1335
    @param cols [sql_gen.Col,...] The columns in the primary key.
1336
        Defaults to the first column in the table.
1337
    @pre The table must not already have a primary key.
1338
    '''
1339
    table = sql_gen.as_Table(table)
1340 5388 aaronmk
    if cols == None: cols = [pkey_name(db, table, recover)]
1341 3079 aaronmk
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1342
1343
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1344
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1345
        log_ignore_excs=(DuplicateException,))
1346
1347 2998 aaronmk
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1348 2688 aaronmk
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1349 3356 aaronmk
    Currently, only function calls and literal values are supported expressions.
1350 2998 aaronmk
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1351 2847 aaronmk
        This allows indexes to be used for comparisons where NULLs are equal.
1352 2538 aaronmk
    '''
1353 2964 aaronmk
    exprs = lists.mk_seq(exprs)
1354 2538 aaronmk
1355 2688 aaronmk
    # Parse exprs
1356
    old_exprs = exprs[:]
1357
    exprs = []
1358
    cols = []
1359
    for i, expr in enumerate(old_exprs):
1360 2823 aaronmk
        expr = sql_gen.as_Col(expr, table)
1361 2688 aaronmk
1362 2823 aaronmk
        # Handle nullable columns
1363 2998 aaronmk
        if ensure_not_null_:
1364 3164 aaronmk
            try: expr = sql_gen.ensure_not_null(db, expr)
1365 2860 aaronmk
            except KeyError: pass # unknown type, so just create plain index
1366 2823 aaronmk
1367 2688 aaronmk
        # Extract col
1368 3002 aaronmk
        expr = copy.deepcopy(expr) # don't modify input!
1369 3356 aaronmk
        col = expr
1370
        if isinstance(expr, sql_gen.FunctionCall): col = expr.args[0]
1371
        expr = sql_gen.cast_literal(expr)
1372
        if not isinstance(expr, (sql_gen.Expr, sql_gen.Col)):
1373 2688 aaronmk
            expr = sql_gen.Expr(expr)
1374 3356 aaronmk
1375 2688 aaronmk
1376
        # Extract table
1377
        if table == None:
1378
            assert sql_gen.is_table_col(col)
1379
            table = col.table
1380
1381 3356 aaronmk
        if isinstance(col, sql_gen.Col): col.table = None
1382 2688 aaronmk
1383
        exprs.append(expr)
1384
        cols.append(col)
1385 2408 aaronmk
1386 2688 aaronmk
    table = sql_gen.as_Table(table)
1387
1388 3005 aaronmk
    # Add index
1389 3148 aaronmk
    str_ = 'CREATE'
1390
    if unique: str_ += ' UNIQUE'
1391
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1392
        ', '.join((v.to_str(db) for v in exprs)))+')'
1393
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1394 2408 aaronmk
1395 5987 aaronmk
def add_pkey_index(db, table): add_index(db, pkey_col(db, table), table)
1396
1397 5765 aaronmk
def add_pkey_or_index(db, table, cols=None, recover=None, warn=False):
1398
    try: add_pkey(db, table, cols, recover)
1399
    except DuplicateKeyException, e:
1400
        if warn: warnings.warn(UserWarning(exc.str_(e)))
1401 5988 aaronmk
        add_pkey_index(db, table)
1402 5765 aaronmk
1403 3083 aaronmk
already_indexed = object() # tells add_indexes() the pkey has already been added
1404
1405
def add_indexes(db, table, has_pkey=True):
1406
    '''Adds an index on all columns in a table.
1407
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1408
        index should be added on the first column.
1409
        * If already_indexed, the pkey is assumed to have already been added
1410
    '''
1411 5337 aaronmk
    cols = table_col_names(db, table)
1412 3083 aaronmk
    if has_pkey:
1413
        if has_pkey is not already_indexed: add_pkey(db, table)
1414
        cols = cols[1:]
1415
    for col in cols: add_index(db, col, table)
1416
1417 3079 aaronmk
#### Tables
1418 2772 aaronmk
1419 3079 aaronmk
### Maintenance
1420 2772 aaronmk
1421 3079 aaronmk
def analyze(db, table):
1422
    table = sql_gen.as_Table(table)
1423
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1424 2934 aaronmk
1425 3079 aaronmk
def autoanalyze(db, table):
1426
    if db.autoanalyze: analyze(db, table)
1427 2935 aaronmk
1428 3079 aaronmk
def vacuum(db, table):
1429
    table = sql_gen.as_Table(table)
1430
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1431
        log_level=3))
1432 2086 aaronmk
1433 3079 aaronmk
### Lifecycle
1434
1435 3247 aaronmk
def drop(db, type_, name):
1436
    name = sql_gen.as_Name(name)
1437
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1438 2889 aaronmk
1439 3247 aaronmk
def drop_table(db, table): drop(db, 'TABLE', table)
1440
1441 3082 aaronmk
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1442
    like=None):
1443 2675 aaronmk
    '''Creates a table.
1444 2681 aaronmk
    @param cols [sql_gen.TypedCol,...] The column names and types
1445
    @param has_pkey If set, the first column becomes the primary key.
1446 2760 aaronmk
    @param col_indexes bool|[ref]
1447
        * If True, indexes will be added on all non-pkey columns.
1448
        * If a list reference, [0] will be set to a function to do this.
1449
          This can be used to delay index creation until the table is populated.
1450 2675 aaronmk
    '''
1451
    table = sql_gen.as_Table(table)
1452
1453 3082 aaronmk
    if like != None:
1454
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1455
            ]+cols
1456 5525 aaronmk
        table.order_by = like.order_by
1457 2681 aaronmk
    if has_pkey:
1458
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1459 2872 aaronmk
        pkey.constraints = 'PRIMARY KEY'
1460 2681 aaronmk
1461 3085 aaronmk
    temp = table.is_temp and not db.debug_temp
1462
        # temp tables permanent in debug_temp mode
1463 2760 aaronmk
1464 3085 aaronmk
    # Create table
1465 3383 aaronmk
    def create():
1466 3085 aaronmk
        str_ = 'CREATE'
1467
        if temp: str_ += ' TEMP'
1468
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1469
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1470 3126 aaronmk
        str_ += '\n);'
1471 3085 aaronmk
1472 3383 aaronmk
        run_query(db, str_, recover=True, cacheable=True, log_level=2,
1473
            log_ignore_excs=(DuplicateException,))
1474
    if table.is_temp:
1475
        while True:
1476
            try:
1477
                create()
1478
                break
1479
            except DuplicateException:
1480
                table.name = next_version(table.name)
1481
                # try again with next version of name
1482
    else: create()
1483 3085 aaronmk
1484 2760 aaronmk
    # Add indexes
1485 2773 aaronmk
    if has_pkey: has_pkey = already_indexed
1486
    def add_indexes_(): add_indexes(db, table, has_pkey)
1487
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1488
    elif col_indexes: add_indexes_() # add now
1489 2675 aaronmk
1490 3084 aaronmk
def copy_table_struct(db, src, dest):
1491
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1492 3085 aaronmk
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1493 3084 aaronmk
1494 5529 aaronmk
def copy_table(db, src, dest):
1495
    '''Creates a copy of a table, including data'''
1496
    copy_table_struct(db, src, dest)
1497
    insert_select(db, dest, None, mk_select(db, src))
1498
1499 3079 aaronmk
### Data
1500 2684 aaronmk
1501 2970 aaronmk
def truncate(db, table, schema='public', **kw_args):
1502
    '''For params, see run_query()'''
1503 2777 aaronmk
    table = sql_gen.as_Table(table, schema)
1504 2970 aaronmk
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1505 2732 aaronmk
1506 2965 aaronmk
def empty_temp(db, tables):
1507
    tables = lists.mk_seq(tables)
1508 2971 aaronmk
    for table in tables: truncate(db, table, log_level=3)
1509 2965 aaronmk
1510 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
1511
    '''For kw_args, see tables()'''
1512
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1513 3094 aaronmk
1514
def distinct_table(db, table, distinct_on):
1515
    '''Creates a copy of a temp table which is distinct on the given columns.
1516 5902 aaronmk
    Adds an index on table's distinct_on columns, to facilitate merge joins.
1517 3097 aaronmk
    @param distinct_on If empty, creates a table with one row. This is useful if
1518
        your distinct_on columns are all literal values.
1519 3099 aaronmk
    @return The new table.
1520 3094 aaronmk
    '''
1521 3099 aaronmk
    new_table = sql_gen.suffixed_table(table, '_distinct')
1522 3411 aaronmk
    distinct_on = filter(sql_gen.is_table_col, distinct_on)
1523 3094 aaronmk
1524 3099 aaronmk
    copy_table_struct(db, table, new_table)
1525 3097 aaronmk
1526
    limit = None
1527
    if distinct_on == []: limit = 1 # one sample row
1528 5903 aaronmk
    else: add_index(db, distinct_on, table) # for join optimization
1529 3097 aaronmk
1530 5903 aaronmk
    insert_select(db, new_table, None, mk_select(db, table,
1531
        distinct_on=distinct_on, order_by=None, limit=limit))
1532 3099 aaronmk
    analyze(db, new_table)
1533 3094 aaronmk
1534 3099 aaronmk
    return new_table