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 7112 aaronmk
class InvalidTypeException(ExceptionWithNameType): pass
86
87 2945 aaronmk
class DuplicateException(ExceptionWithNameType): pass
88 2143 aaronmk
89 3419 aaronmk
class DoesNotExistException(ExceptionWithNameType): pass
90
91 89 aaronmk
class EmptyRowException(DbException): pass
92
93 865 aaronmk
##### Warnings
94
95
class DbWarning(UserWarning): pass
96
97 1930 aaronmk
##### Result retrieval
98
99
def col_names(cur): return (col[0] for col in cur.description)
100
101
def rows(cur): return iter(lambda: cur.fetchone(), None)
102
103
def consume_rows(cur):
104
    '''Used to fetch all rows so result will be cached'''
105
    iters.consume_iter(rows(cur))
106
107
def next_row(cur): return rows(cur).next()
108
109
def row(cur):
110
    row_ = next_row(cur)
111
    consume_rows(cur)
112
    return row_
113
114
def next_value(cur): return next_row(cur)[0]
115
116
def value(cur): return row(cur)[0]
117
118
def values(cur): return iters.func_iter(lambda: next_value(cur))
119
120
def value_or_none(cur):
121
    try: return value(cur)
122
    except StopIteration: return None
123
124 2762 aaronmk
##### Escaping
125 2101 aaronmk
126 2573 aaronmk
def esc_name_by_module(module, name):
127
    if module == 'psycopg2' or module == None: quote = '"'
128 2101 aaronmk
    elif module == 'MySQLdb': quote = '`'
129
    else: raise NotImplementedError("Can't escape name for "+module+' database')
130 2500 aaronmk
    return sql_gen.esc_name(name, quote)
131 2101 aaronmk
132
def esc_name_by_engine(engine, name, **kw_args):
133
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
134
135
def esc_name(db, name, **kw_args):
136
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
137
138
def qual_name(db, schema, table):
139
    def esc_name_(name): return esc_name(db, name)
140
    table = esc_name_(table)
141
    if schema != None: return esc_name_(schema)+'.'+table
142
    else: return table
143
144 1869 aaronmk
##### Database connections
145 1849 aaronmk
146 2097 aaronmk
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
147 1926 aaronmk
148 1869 aaronmk
db_engines = {
149
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
150
    'PostgreSQL': ('psycopg2', {}),
151
}
152
153
DatabaseErrors_set = set([DbException])
154
DatabaseErrors = tuple(DatabaseErrors_set)
155
156
def _add_module(module):
157
    DatabaseErrors_set.add(module.DatabaseError)
158
    global DatabaseErrors
159
    DatabaseErrors = tuple(DatabaseErrors_set)
160
161
def db_config_str(db_config):
162
    return db_config['engine']+' database '+db_config['database']
163
164 2448 aaronmk
log_debug_none = lambda msg, level=2: None
165 1901 aaronmk
166 1849 aaronmk
class DbConn:
167 2923 aaronmk
    def __init__(self, db_config, autocommit=True, caching=True,
168 3183 aaronmk
        log_debug=log_debug_none, debug_temp=False, src=None):
169 2915 aaronmk
        '''
170
        @param debug_temp Whether temporary objects should instead be permanent.
171
            This assists in debugging the internal objects used by the program.
172 3183 aaronmk
        @param src In autocommit mode, will be included in a comment in every
173
            query, to help identify the data source in pg_stat_activity.
174 2915 aaronmk
        '''
175 1869 aaronmk
        self.db_config = db_config
176 2190 aaronmk
        self.autocommit = autocommit
177
        self.caching = caching
178 1901 aaronmk
        self.log_debug = log_debug
179 2193 aaronmk
        self.debug = log_debug != log_debug_none
180 2915 aaronmk
        self.debug_temp = debug_temp
181 3183 aaronmk
        self.src = src
182 3074 aaronmk
        self.autoanalyze = False
183 3269 aaronmk
        self.autoexplain = False
184
        self.profile_row_ct = None
185 1869 aaronmk
186 3124 aaronmk
        self._savepoint = 0
187 3120 aaronmk
        self._reset()
188 1869 aaronmk
189
    def __getattr__(self, name):
190
        if name == '__dict__': raise Exception('getting __dict__')
191
        if name == 'db': return self._db()
192
        else: raise AttributeError()
193
194
    def __getstate__(self):
195
        state = copy.copy(self.__dict__) # shallow copy
196 1915 aaronmk
        state['log_debug'] = None # don't pickle the debug callback
197 1869 aaronmk
        state['_DbConn__db'] = None # don't pickle the connection
198
        return state
199
200 3118 aaronmk
    def clear_cache(self): self.query_results = {}
201
202 3120 aaronmk
    def _reset(self):
203 3118 aaronmk
        self.clear_cache()
204 3124 aaronmk
        assert self._savepoint == 0
205 3118 aaronmk
        self._notices_seen = set()
206
        self.__db = None
207
208 2165 aaronmk
    def connected(self): return self.__db != None
209
210 3116 aaronmk
    def close(self):
211 3119 aaronmk
        if not self.connected(): return
212
213 3135 aaronmk
        # Record that the automatic transaction is now closed
214 3136 aaronmk
        self._savepoint -= 1
215 3135 aaronmk
216 3119 aaronmk
        self.db.close()
217 3120 aaronmk
        self._reset()
218 3116 aaronmk
219 3125 aaronmk
    def reconnect(self):
220
        # Do not do this in test mode as it would roll back everything
221
        if self.autocommit: self.close()
222
        # Connection will be reopened automatically on first query
223
224 1869 aaronmk
    def _db(self):
225
        if self.__db == None:
226
            # Process db_config
227
            db_config = self.db_config.copy() # don't modify input!
228 2097 aaronmk
            schemas = db_config.pop('schemas', None)
229 1869 aaronmk
            module_name, mappings = db_engines[db_config.pop('engine')]
230
            module = __import__(module_name)
231
            _add_module(module)
232
            for orig, new in mappings.iteritems():
233
                try: util.rename_key(db_config, orig, new)
234
                except KeyError: pass
235
236
            # Connect
237
            self.__db = module.connect(**db_config)
238
239 3161 aaronmk
            # Record that a transaction is already open
240
            self._savepoint += 1
241
242 1869 aaronmk
            # Configure connection
243 2906 aaronmk
            if hasattr(self.db, 'set_isolation_level'):
244
                import psycopg2.extensions
245
                self.db.set_isolation_level(
246
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
247 2101 aaronmk
            if schemas != None:
248 2893 aaronmk
                search_path = [self.esc_name(s) for s in schemas.split(',')]
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 6903 aaronmk
, pg_get_expr(pg_attrdef.adbin, attrelid, true) AS default
482 6899 aaronmk
, 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 6904 aaronmk
                    , recover=True, cacheable=cacheable, log_level=4))
491 6905 aaronmk
            except (DoesNotExistException, StopIteration):
492
                raise sql_gen.NoUnderlyingTableException(col)
493 6900 aaronmk
            if is_array: type_ = sql_gen.ArrayType(type_)
494 6899 aaronmk
        else:
495
            table = sql_gen.Table('columns', 'information_schema')
496
            cols = [sql_gen.Col('data_type'), sql_gen.Col('udt_name'),
497
                'column_default', sql_gen.Cast('boolean',
498
                sql_gen.Col('is_nullable'))]
499
500
            conds = [('table_name', col.table.name),
501
                ('column_name', strings.ustr(col.name))]
502
            schema = col.table.schema
503
            if schema != None: conds.append(('table_schema', schema))
504
505
            cur = select(self, table, cols, conds, order_by='table_schema',
506
                limit=1, cacheable=cacheable, log_level=4)
507
            try: type_, extra_type, default, nullable = row(cur)
508
            except StopIteration: raise sql_gen.NoUnderlyingTableException(col)
509
            if type_ == 'USER-DEFINED': type_ = extra_type
510
            elif type_ == 'ARRAY':
511
                type_ = sql_gen.ArrayType(strings.remove_prefix('_', extra_type,
512
                    require=True))
513 2643 aaronmk
514 6899 aaronmk
        if default != None: default = sql_gen.as_Code(default, self)
515 2819 aaronmk
        return sql_gen.TypedCol(col.name, type_, default, nullable)
516 2917 aaronmk
517
    def TempFunction(self, name):
518
        if self.debug_temp: schema = None
519
        else: schema = 'pg_temp'
520
        return sql_gen.Function(name, schema)
521 1849 aaronmk
522 1869 aaronmk
connect = DbConn
523
524 832 aaronmk
##### Recoverable querying
525 15 aaronmk
526 5577 aaronmk
def parse_exception(db, e, recover=False):
527
    msg = strings.ustr(e.args[0])
528
    msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)
529
530
    match = re.match(r'^invalid byte sequence for encoding "(.+?)":', msg)
531
    if match:
532
        encoding, = match.groups()
533
        raise EncodingException(encoding, e)
534
535 5763 aaronmk
    def make_DuplicateKeyException(constraint, e):
536 5577 aaronmk
        cols = []
537
        cond = None
538
        if recover: # need auto-rollback to run index_cols()
539
            try:
540
                cols = index_cols(db, constraint)
541
                cond = index_cond(db, constraint)
542
            except NotImplementedError: pass
543 9757 aaronmk
            except DatabaseErrors, e2: # e.g. "current transaction is aborted"
544
                exc.print_ex(e2)
545
                raise e
546 5763 aaronmk
        return DuplicateKeyException(constraint, cond, cols, e)
547 5577 aaronmk
548 5763 aaronmk
    match = re.match(r'^duplicate key value violates unique constraint "(.+?)"',
549
        msg)
550
    if match:
551
        constraint, = match.groups()
552
        raise make_DuplicateKeyException(constraint, e)
553
554 5764 aaronmk
    match = re.match(r'^could not create unique index "(.+?)"\n'
555
        r'DETAIL:  Key .+? is duplicated', msg)
556
    if match:
557
        constraint, = match.groups()
558 5823 aaronmk
        raise DuplicateKeyException(constraint, None, [], e)
559 5764 aaronmk
560 5577 aaronmk
    match = re.match(r'^null value in column "(.+?)" violates not-null'
561
        r' constraint', msg)
562
    if match:
563
        col, = match.groups()
564
        raise NullValueException('NOT NULL', None, [col], e)
565
566
    match = re.match(r'^new row for relation "(.+?)" violates check '
567
        r'constraint "(.+?)"', msg)
568
    if match:
569
        table, constraint = match.groups()
570
        constraint = sql_gen.Col(constraint, table)
571
        cond = None
572
        if recover: # need auto-rollback to run constraint_cond()
573
            try: cond = constraint_cond(db, constraint)
574
            except NotImplementedError: pass
575
        raise CheckException(constraint.to_str(db), cond, [], e)
576
577 5723 aaronmk
    match = re.match(r'^(?:invalid input (?:syntax|value)\b[^:]*'
578 5717 aaronmk
        r'|.+? out of range)(?:: "(.+?)")?', msg)
579 5577 aaronmk
    if match:
580
        value, = match.groups()
581 5717 aaronmk
        value = util.do_ignore_none(strings.to_unicode, value)
582
        raise InvalidValueException(value, e)
583 5577 aaronmk
584
    match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
585
        r'is of type', msg)
586
    if match:
587
        col, type_ = match.groups()
588
        raise MissingCastException(type_, col, e)
589
590
    match = re.match(r'^could not determine polymorphic type because '
591
        r'input has type "unknown"', msg)
592
    if match: raise MissingCastException('text', None, e)
593
594 12149 aaronmk
    # type mismatches
595 5724 aaronmk
    match = re.match(r'^.+? types (.+?) and (.+?) cannot be matched', msg)
596
    if match:
597
        type0, type1 = match.groups()
598
        raise MissingCastException(type0, None, e)
599 12149 aaronmk
    match = re.match(r'operator does not exist: (.+?) = (.+)', msg)
600
        # these appear when a staging table column is of the wrong type
601
    if match:
602
        type0, type1 = match.groups()
603
        raise MissingCastException(type0, None, e)
604 5577 aaronmk
605 7112 aaronmk
    match = re.match(r'^.*?\brelation "(.+?)" is not a (table)', msg)
606
    if match:
607
        name, type_ = match.groups()
608
        raise InvalidTypeException(type_, name, e)
609
610 9458 aaronmk
    typed_name_re = r'^(\S+) ("?)(.+?)\2(?: of relation ".+?")?'
611 9457 aaronmk
        # regexp must be followed with text for .*? to match properly
612 5577 aaronmk
613 9458 aaronmk
    match = re.match(typed_name_re+r'(?:\(.*?)? already exists', msg)
614 5577 aaronmk
    if match:
615 9458 aaronmk
        type_, quote, name = match.groups()
616 5577 aaronmk
        raise DuplicateException(type_, name, e)
617
618
    match = re.match(r'more than one (\S+) named ""(.+?)""', msg)
619
    if match:
620
        type_, name = match.groups()
621
        raise DuplicateException(type_, name, e)
622
623 12149 aaronmk
    # type mismatches which appear as "does not exist" errors
624 5577 aaronmk
    match = re.match(typed_name_re+r' does not exist', msg)
625
    if match:
626 9458 aaronmk
        type_, quote, name = match.groups()
627 5683 aaronmk
        if type_ == 'function':
628 5709 aaronmk
            match = re.match(r'^(.+?)\(.*\)$', name)
629 5710 aaronmk
            if match: # includes params, so is call rather than cast to regproc
630 5708 aaronmk
                function_name, = match.groups()
631 5711 aaronmk
                func = sql_gen.Function(function_name)
632 5712 aaronmk
                if function_exists(db, func) and msg.find('CAST') < 0:
633 5683 aaronmk
                    # not found only because of a missing cast
634 10548 aaronmk
                    try: type_ = function_param0_type(db, func)
635
                    except DuplicateException: type_ = 'text'
636 6301 aaronmk
                    col = None
637 5792 aaronmk
                    if type_ == 'anyelement': type_ = 'text'
638 6301 aaronmk
                    elif type_ == 'hstore': # cast just the value param
639
                        type_ = 'text'
640
                        col = 'value'
641
                    raise MissingCastException(type_, col, e)
642 5577 aaronmk
        raise DoesNotExistException(type_, name, e)
643
644
    raise # no specific exception raised
645
646 2139 aaronmk
def with_savepoint(db, func): return db.with_savepoint(func)
647 11 aaronmk
648 2791 aaronmk
def run_query(db, query, recover=None, cacheable=False, log_level=2,
649
    log_ignore_excs=None, **kw_args):
650 2794 aaronmk
    '''For params, see DbConn.run_query()'''
651 830 aaronmk
    if recover == None: recover = False
652 2464 aaronmk
    if log_ignore_excs == None: log_ignore_excs = ()
653
    log_ignore_excs = tuple(log_ignore_excs)
654 3236 aaronmk
    debug_msg_ref = [None]
655 830 aaronmk
656 3267 aaronmk
    query = with_explain_comment(db, query)
657 3258 aaronmk
658 2148 aaronmk
    try:
659 2464 aaronmk
        try:
660 2794 aaronmk
            def run(): return db.run_query(query, cacheable, log_level,
661 2793 aaronmk
                debug_msg_ref, **kw_args)
662 2796 aaronmk
            if recover and not db.is_cached(query):
663 2464 aaronmk
                return with_savepoint(db, run)
664
            else: return run() # don't need savepoint if cached
665 5890 aaronmk
        except Exception, e:
666
            # Give failed EXPLAIN approximately the log_level of its query
667
            if query.startswith('EXPLAIN'): log_level -= 1
668
669
            parse_exception(db, e, recover)
670 2464 aaronmk
    except log_ignore_excs:
671
        log_level += 2
672
        raise
673
    finally:
674 3236 aaronmk
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
675 830 aaronmk
676 832 aaronmk
##### Basic queries
677
678 3256 aaronmk
def is_explainable(query):
679
    # See <http://www.postgresql.org/docs/8.3/static/sql-explain.html#AEN57749>
680 3257 aaronmk
    return re.match(r'^(?:SELECT|INSERT|UPDATE|DELETE|VALUES|EXECUTE|DECLARE)\b'
681
        , query)
682 3256 aaronmk
683 3263 aaronmk
def explain(db, query, **kw_args):
684
    '''
685
    For params, see run_query().
686
    '''
687 3267 aaronmk
    kw_args.setdefault('log_level', 4)
688 3263 aaronmk
689 3750 aaronmk
    return strings.ustr(strings.join_lines(values(run_query(db,
690
        'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))
691 3256 aaronmk
        # not a higher log_level because it's useful to see what query is being
692
        # run before it's executed, which EXPLAIN effectively provides
693
694 3265 aaronmk
def has_comment(query): return query.endswith('*/')
695
696
def with_explain_comment(db, query, **kw_args):
697 3269 aaronmk
    if db.autoexplain and not has_comment(query) and is_explainable(query):
698 3265 aaronmk
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
699
            +explain(db, query, **kw_args))
700
    return query
701
702 2153 aaronmk
def next_version(name):
703 2163 aaronmk
    version = 1 # first existing name was version 0
704 2586 aaronmk
    match = re.match(r'^(.*)#(\d+)$', name)
705 2153 aaronmk
    if match:
706 2586 aaronmk
        name, version = match.groups()
707
        version = int(version)+1
708 2932 aaronmk
    return sql_gen.concat(name, '#'+str(version))
709 2153 aaronmk
710 2899 aaronmk
def lock_table(db, table, mode):
711
    table = sql_gen.as_Table(table)
712
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
713
714 3303 aaronmk
def run_query_into(db, query, into=None, add_pkey_=False, **kw_args):
715 2085 aaronmk
    '''Outputs a query to a temp table.
716
    For params, see run_query().
717
    '''
718 2789 aaronmk
    if into == None: return run_query(db, query, **kw_args)
719 2790 aaronmk
720
    assert isinstance(into, sql_gen.Table)
721
722 2992 aaronmk
    into.is_temp = True
723 3008 aaronmk
    # "temporary tables cannot specify a schema name", so remove schema
724
    into.schema = None
725 2992 aaronmk
726 2790 aaronmk
    kw_args['recover'] = True
727 2945 aaronmk
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
728 2790 aaronmk
729 2916 aaronmk
    temp = not db.debug_temp # tables are permanent in debug_temp mode
730 2790 aaronmk
731
    # Create table
732
    while True:
733
        create_query = 'CREATE'
734
        if temp: create_query += ' TEMP'
735
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
736 2385 aaronmk
737 2790 aaronmk
        try:
738
            cur = run_query(db, create_query, **kw_args)
739
                # CREATE TABLE AS sets rowcount to # rows in query
740
            break
741 2945 aaronmk
        except DuplicateException, e:
742 2790 aaronmk
            into.name = next_version(into.name)
743
            # try again with next version of name
744
745 5967 aaronmk
    if add_pkey_: add_pkey_or_index(db, into, warn=True)
746 3075 aaronmk
747
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
748
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
749
    # table is going to be used in complex queries, it is wise to run ANALYZE on
750
    # the temporary table after it is populated."
751
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
752
    # If into is not a temp table, ANALYZE is useful but not required.
753 3073 aaronmk
    analyze(db, into)
754 2790 aaronmk
755
    return cur
756 2085 aaronmk
757 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
758
759 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
760
761 8333 aaronmk
def has_subset_func(db, table):
762
    return sql_gen.is_underlying_table(table) and function_exists(db, table)
763
764 3420 aaronmk
def mk_select(db, tables=None, fields=None, conds=None, distinct_on=[],
765 3494 aaronmk
    limit=None, start=None, order_by=order_by_pkey, default_table=None,
766
    explain=True):
767 1981 aaronmk
    '''
768 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
769 2280 aaronmk
        together, with tables after the first being sql_gen.Join objects
770 1981 aaronmk
    @param fields Use None to select all fields in the table
771 2377 aaronmk
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
772 2379 aaronmk
        * container can be any iterable type
773 2399 aaronmk
        * compare_left_side: sql_gen.Code|str (for col name)
774
        * compare_right_side: sql_gen.ValueCond|literal value
775 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
776
        use all columns
777 2786 aaronmk
    @return query
778 1981 aaronmk
    '''
779 2315 aaronmk
    # Parse tables param
780 2964 aaronmk
    tables = lists.mk_seq(tables)
781 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
782 2315 aaronmk
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
783 2121 aaronmk
784 2315 aaronmk
    # Parse other params
785 2376 aaronmk
    if conds == None: conds = []
786 2650 aaronmk
    elif dicts.is_dict(conds): conds = conds.items()
787 2379 aaronmk
    conds = list(conds) # don't modify input! (list() copies input)
788 3129 aaronmk
    assert limit == None or isinstance(limit, (int, long))
789
    assert start == None or isinstance(start, (int, long))
790 5527 aaronmk
    if limit == 0: order_by = None
791 8334 aaronmk
    if order_by in (None, order_by_pkey) and has_subset_func(db, table0):
792
        # can use subset function for fast querying at large OFFSET values
793 8335 aaronmk
        run_query(db, 'SET LOCAL enable_sort TO off')
794 8334 aaronmk
        table0 = sql_gen.FunctionCall(table0, limit_=limit, offset_=start)
795
        if limit != 0: limit = None # done by function
796
        start = None # done by function
797
        order_by = None # done by function
798
    elif order_by is order_by_pkey:
799 5641 aaronmk
        if lists.is_seq(distinct_on) and distinct_on: order_by = distinct_on[0]
800
        elif table0 != None: order_by = table_order_by(db, table0, recover=True)
801
        else: order_by = None
802 865 aaronmk
803 2315 aaronmk
    query = 'SELECT'
804 2056 aaronmk
805 2315 aaronmk
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
806 2056 aaronmk
807 2200 aaronmk
    # DISTINCT ON columns
808 2233 aaronmk
    if distinct_on != []:
809 2467 aaronmk
        query += '\nDISTINCT'
810 2254 aaronmk
        if distinct_on is not distinct_on_all:
811 2200 aaronmk
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
812
813
    # Columns
814 3185 aaronmk
    if query.find('\n') >= 0: whitespace = '\n'
815
    else: whitespace = ' '
816
    if fields == None: query += whitespace+'*'
817 2765 aaronmk
    else:
818
        assert fields != []
819 3185 aaronmk
        if len(fields) > 1: whitespace = '\n'
820
        query += whitespace+('\n, '.join(map(parse_col, fields)))
821 2200 aaronmk
822
    # Main table
823 3185 aaronmk
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
824
    else: whitespace = ' '
825 3420 aaronmk
    if table0 != None: query += whitespace+'FROM '+table0.to_str(db)
826 865 aaronmk
827 2122 aaronmk
    # Add joins
828 2271 aaronmk
    left_table = table0
829 2263 aaronmk
    for join_ in tables:
830
        table = join_.table
831 2238 aaronmk
832 2343 aaronmk
        # Parse special values
833
        if join_.type_ is sql_gen.filter_out: # filter no match
834 2376 aaronmk
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
835 2853 aaronmk
                sql_gen.CompareCond(None, '~=')))
836 2343 aaronmk
837 2467 aaronmk
        query += '\n'+join_.to_str(db, left_table)
838 2122 aaronmk
839
        left_table = table
840
841 865 aaronmk
    missing = True
842 2376 aaronmk
    if conds != []:
843 2576 aaronmk
        if len(conds) == 1: whitespace = ' '
844
        else: whitespace = '\n'
845 2578 aaronmk
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
846
            .to_str(db) for l, r in conds], 'WHERE')
847 2227 aaronmk
    if order_by != None:
848 5642 aaronmk
        query += '\nORDER BY '+sql_gen.as_Col(order_by).to_str(db)
849 3297 aaronmk
    if limit != None: query += '\nLIMIT '+str(limit)
850 865 aaronmk
    if start != None:
851 2467 aaronmk
        if start != 0: query += '\nOFFSET '+str(start)
852 865 aaronmk
853 3494 aaronmk
    if explain: query = with_explain_comment(db, query)
854 3266 aaronmk
855 2786 aaronmk
    return query
856 11 aaronmk
857 2054 aaronmk
def select(db, *args, **kw_args):
858
    '''For params, see mk_select() and run_query()'''
859
    recover = kw_args.pop('recover', None)
860
    cacheable = kw_args.pop('cacheable', True)
861 2442 aaronmk
    log_level = kw_args.pop('log_level', 2)
862 2054 aaronmk
863 2791 aaronmk
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
864
        log_level=log_level)
865 2054 aaronmk
866 2788 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
867 3181 aaronmk
    embeddable=False, ignore=False, src=None):
868 1960 aaronmk
    '''
869
    @param returning str|None An inserted column (such as pkey) to return
870 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
871 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
872
        query will be fully cached, not just if it raises an exception.
873 3009 aaronmk
    @param ignore Whether to ignore duplicate keys.
874 3181 aaronmk
    @param src Will be included in the name of any created function, to help
875
        identify the data source in pg_stat_activity.
876 1960 aaronmk
    '''
877 2754 aaronmk
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
878 2318 aaronmk
    if cols == []: cols = None # no cols (all defaults) = unknown col names
879 3010 aaronmk
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
880 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
881 2327 aaronmk
    if returning != None: returning = sql_gen.as_Col(returning, table)
882 2063 aaronmk
883 2497 aaronmk
    first_line = 'INSERT INTO '+table.to_str(db)
884 2063 aaronmk
885 3009 aaronmk
    def mk_insert(select_query):
886
        query = first_line
887 3014 aaronmk
        if cols != None:
888
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
889 3009 aaronmk
        query += '\n'+select_query
890
891
        if returning != None:
892
            returning_name_col = sql_gen.to_name_only_col(returning)
893
            query += '\nRETURNING '+returning_name_col.to_str(db)
894
895
        return query
896 2063 aaronmk
897 3489 aaronmk
    return_type = sql_gen.CustomCode('unknown')
898
    if returning != None: return_type = sql_gen.ColType(returning)
899 3017 aaronmk
900 3009 aaronmk
    if ignore:
901 3017 aaronmk
        # Always return something to set the correct rowcount
902
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
903
904 3009 aaronmk
        embeddable = True # must use function
905 3010 aaronmk
906 3450 aaronmk
        if cols == None: row = [sql_gen.Col(sql_gen.all_cols, 'row')]
907
        else: row = [sql_gen.Col(c.name, 'row') for c in cols]
908 3092 aaronmk
909 3484 aaronmk
        query = sql_gen.RowExcIgnore(sql_gen.RowType(table), select_query,
910 3497 aaronmk
            sql_gen.ReturnQuery(mk_insert(sql_gen.Values(row).to_str(db))),
911
            cols)
912 3009 aaronmk
    else: query = mk_insert(select_query)
913
914 2070 aaronmk
    if embeddable:
915
        # Create function
916 2513 aaronmk
        function_name = sql_gen.clean_name(first_line)
917 3181 aaronmk
        if src != None: function_name = src+': '+function_name
918 2189 aaronmk
        while True:
919
            try:
920 3451 aaronmk
                func = db.TempFunction(function_name)
921 3489 aaronmk
                def_ = sql_gen.FunctionDef(func, sql_gen.SetOf(return_type),
922
                    query)
923 2194 aaronmk
924 3443 aaronmk
                run_query(db, def_.to_str(db), recover=True, cacheable=True,
925 2945 aaronmk
                    log_ignore_excs=(DuplicateException,))
926 2189 aaronmk
                break # this version was successful
927 2945 aaronmk
            except DuplicateException, e:
928 2189 aaronmk
                function_name = next_version(function_name)
929
                # try again with next version of name
930 2070 aaronmk
931 2337 aaronmk
        # Return query that uses function
932 3009 aaronmk
        cols = None
933
        if returning != None: cols = [returning]
934 3451 aaronmk
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(func), cols)
935
            # AS clause requires function alias
936 3298 aaronmk
        return mk_select(db, func_table, order_by=None)
937 2070 aaronmk
938 2787 aaronmk
    return query
939 2066 aaronmk
940 3074 aaronmk
def insert_select(db, table, *args, **kw_args):
941 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
942 2386 aaronmk
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
943
        values in
944 2072 aaronmk
    '''
945 3141 aaronmk
    returning = kw_args.get('returning', None)
946
    ignore = kw_args.get('ignore', False)
947
948 2386 aaronmk
    into = kw_args.pop('into', None)
949
    if into != None: kw_args['embeddable'] = True
950 2066 aaronmk
    recover = kw_args.pop('recover', None)
951 3141 aaronmk
    if ignore: recover = True
952 2066 aaronmk
    cacheable = kw_args.pop('cacheable', True)
953 2673 aaronmk
    log_level = kw_args.pop('log_level', 2)
954 2066 aaronmk
955 3141 aaronmk
    rowcount_only = ignore and returning == None # keep NULL rows on server
956
    if rowcount_only: into = sql_gen.Table('rowcount')
957
958 3074 aaronmk
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
959
        into, recover=recover, cacheable=cacheable, log_level=log_level)
960 3141 aaronmk
    if rowcount_only: empty_temp(db, into)
961 3074 aaronmk
    autoanalyze(db, table)
962
    return cur
963 2063 aaronmk
964 2738 aaronmk
default = sql_gen.default # tells insert() to use the default value for a column
965 2066 aaronmk
966 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
967 2085 aaronmk
    '''For params, see insert_select()'''
968 5050 aaronmk
    ignore = kw_args.pop('ignore', False)
969 5094 aaronmk
    if ignore: kw_args.setdefault('recover', True)
970 5050 aaronmk
971 1960 aaronmk
    if lists.is_seq(row): cols = None
972
    else:
973
        cols = row.keys()
974
        row = row.values()
975 2738 aaronmk
    row = list(row) # ensure that "== []" works
976 1960 aaronmk
977 2738 aaronmk
    if row == []: query = None
978
    else: query = sql_gen.Values(row).to_str(db)
979 1961 aaronmk
980 5050 aaronmk
    try: return insert_select(db, table, cols, query, *args, **kw_args)
981 5057 aaronmk
    except (DuplicateKeyException, NullValueException):
982 5050 aaronmk
        if not ignore: raise
983 5163 aaronmk
        return None
984 11 aaronmk
985 3152 aaronmk
def mk_update(db, table, changes=None, cond=None, in_place=False,
986 3153 aaronmk
    cacheable_=True):
987 2402 aaronmk
    '''
988
    @param changes [(col, new_value),...]
989
        * container can be any iterable type
990
        * col: sql_gen.Code|str (for col name)
991
        * new_value: sql_gen.Code|literal value
992
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
993 3056 aaronmk
    @param in_place If set, locks the table and updates rows in place.
994
        This avoids creating dead rows in PostgreSQL.
995
        * cond must be None
996 3153 aaronmk
    @param cacheable_ Whether column structure information used to generate the
997 3152 aaronmk
        query can be cached
998 2402 aaronmk
    @return str query
999
    '''
1000 3057 aaronmk
    table = sql_gen.as_Table(table)
1001
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
1002
        for c, v in changes]
1003
1004 3056 aaronmk
    if in_place:
1005
        assert cond == None
1006 3058 aaronmk
1007 5398 aaronmk
        def col_type(col):
1008
            return sql_gen.canon_type(db.col_info(
1009
                sql_gen.with_default_table(c, table), cacheable_).type)
1010
        changes = [(c, v, col_type(c)) for c, v in changes]
1011 3065 aaronmk
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
1012 5396 aaronmk
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '+t+'\nUSING '
1013
            +v.to_str(db) for c, v, t in changes))
1014 3058 aaronmk
    else:
1015
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
1016
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
1017
            for c, v in changes))
1018
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
1019 3056 aaronmk
1020 3266 aaronmk
    query = with_explain_comment(db, query)
1021
1022 2402 aaronmk
    return query
1023
1024 3074 aaronmk
def update(db, table, *args, **kw_args):
1025 2402 aaronmk
    '''For params, see mk_update() and run_query()'''
1026
    recover = kw_args.pop('recover', None)
1027 3043 aaronmk
    cacheable = kw_args.pop('cacheable', False)
1028 3030 aaronmk
    log_level = kw_args.pop('log_level', 2)
1029 2402 aaronmk
1030 3074 aaronmk
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
1031
        cacheable, log_level=log_level)
1032
    autoanalyze(db, table)
1033
    return cur
1034 2402 aaronmk
1035 3286 aaronmk
def mk_delete(db, table, cond=None):
1036
    '''
1037
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
1038
    @return str query
1039
    '''
1040
    query = 'DELETE FROM '+table.to_str(db)
1041
    if cond != None: query += '\nWHERE '+cond.to_str(db)
1042
1043
    query = with_explain_comment(db, query)
1044
1045
    return query
1046
1047
def delete(db, table, *args, **kw_args):
1048
    '''For params, see mk_delete() and run_query()'''
1049
    recover = kw_args.pop('recover', None)
1050 3295 aaronmk
    cacheable = kw_args.pop('cacheable', True)
1051 3286 aaronmk
    log_level = kw_args.pop('log_level', 2)
1052
1053
    cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
1054
        cacheable, log_level=log_level)
1055
    autoanalyze(db, table)
1056
    return cur
1057
1058 135 aaronmk
def last_insert_id(db):
1059 1849 aaronmk
    module = util.root_module(db.db)
1060 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
1061
    elif module == 'MySQLdb': return db.insert_id()
1062
    else: return None
1063 13 aaronmk
1064 3490 aaronmk
def define_func(db, def_):
1065
    func = def_.function
1066
    while True:
1067
        try:
1068
            run_query(db, def_.to_str(db), recover=True, cacheable=True,
1069
                log_ignore_excs=(DuplicateException,))
1070
            break # successful
1071
        except DuplicateException:
1072 3495 aaronmk
            func.name = next_version(func.name)
1073 3490 aaronmk
            # try again with next version of name
1074
1075 2394 aaronmk
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
1076 2383 aaronmk
    '''Creates a mapping from original column names (which may have collisions)
1077 2415 aaronmk
    to names that will be distinct among the columns' tables.
1078 2383 aaronmk
    This is meant to be used for several tables that are being joined together.
1079 2415 aaronmk
    @param cols The columns to combine. Duplicates will be removed.
1080
    @param into The table for the new columns.
1081 2394 aaronmk
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
1082
        columns will be included in the mapping even if they are not in cols.
1083
        The tables of the provided Col objects will be changed to into, so make
1084
        copies of them if you want to keep the original tables.
1085
    @param as_items Whether to return a list of dict items instead of a dict
1086 2383 aaronmk
    @return dict(orig_col=new_col, ...)
1087
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
1088 2392 aaronmk
        * new_col: sql_gen.Col(orig_col_name, into)
1089
        * All mappings use the into table so its name can easily be
1090 2383 aaronmk
          changed for all columns at once
1091
    '''
1092 2415 aaronmk
    cols = lists.uniqify(cols)
1093
1094 2394 aaronmk
    items = []
1095 2389 aaronmk
    for col in preserve:
1096 2390 aaronmk
        orig_col = copy.copy(col)
1097 2392 aaronmk
        col.table = into
1098 2394 aaronmk
        items.append((orig_col, col))
1099
    preserve = set(preserve)
1100
    for col in cols:
1101 2716 aaronmk
        if col not in preserve:
1102 3750 aaronmk
            items.append((col, sql_gen.Col(strings.ustr(col), into, col.srcs)))
1103 2394 aaronmk
1104
    if not as_items: items = dict(items)
1105
    return items
1106 2383 aaronmk
1107 2393 aaronmk
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
1108 2391 aaronmk
    '''For params, see mk_flatten_mapping()
1109
    @return See return value of mk_flatten_mapping()
1110
    '''
1111 2394 aaronmk
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
1112
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
1113 5523 aaronmk
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
1114
        into=into, add_pkey_=True)
1115 3708 aaronmk
        # don't cache because the temp table will usually be truncated after use
1116 2394 aaronmk
    return dict(items)
1117 2391 aaronmk
1118 3079 aaronmk
##### Database structure introspection
1119 2414 aaronmk
1120 3079 aaronmk
#### Tables
1121
1122 4555 aaronmk
def tables(db, schema_like='public', table_like='%', exact=False,
1123
    cacheable=True):
1124 3079 aaronmk
    if exact: compare = '='
1125
    else: compare = 'LIKE'
1126
1127
    module = util.root_module(db.db)
1128
    if module == 'psycopg2':
1129
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1130
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1131
        return values(select(db, 'pg_tables', ['tablename'], conds,
1132 4555 aaronmk
            order_by='tablename', cacheable=cacheable, log_level=4))
1133 3079 aaronmk
    elif module == 'MySQLdb':
1134
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1135
            , cacheable=True, log_level=4))
1136
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1137
1138 4556 aaronmk
def table_exists(db, table, cacheable=True):
1139 3079 aaronmk
    table = sql_gen.as_Table(table)
1140 4556 aaronmk
    return list(tables(db, table.schema, table.name, True, cacheable)) != []
1141 3079 aaronmk
1142 2426 aaronmk
def table_row_count(db, table, recover=None):
1143 2786 aaronmk
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
1144 3298 aaronmk
        order_by=None), recover=recover, log_level=3))
1145 2426 aaronmk
1146 5337 aaronmk
def table_col_names(db, table, recover=None):
1147 5528 aaronmk
    return list(col_names(select(db, table, limit=0, recover=recover,
1148
        log_level=4)))
1149 2414 aaronmk
1150 5383 aaronmk
def table_cols(db, table, *args, **kw_args):
1151
    return [sql_gen.as_Col(strings.ustr(c), table)
1152
        for c in table_col_names(db, table, *args, **kw_args)]
1153
1154 5521 aaronmk
def table_pkey_index(db, table, recover=None):
1155
    table_str = sql_gen.Literal(table.to_str(db))
1156
    try:
1157
        return sql_gen.Table(value(run_query(db, '''\
1158
SELECT relname
1159
FROM pg_index
1160
JOIN pg_class index ON index.oid = indexrelid
1161
WHERE
1162
indrelid = '''+table_str.to_str(db)+'''::regclass
1163
AND indisprimary
1164
'''
1165
            , recover, cacheable=True, log_level=4)), table.schema)
1166
    except StopIteration: raise DoesNotExistException('primary key', '')
1167
1168 5389 aaronmk
def table_pkey_col(db, table, recover=None):
1169 5061 aaronmk
    table = sql_gen.as_Table(table)
1170
1171 5989 aaronmk
    module = util.root_module(db.db)
1172
    if module == 'psycopg2':
1173
        return sql_gen.Col(index_cols(db, table_pkey_index(db, table,
1174
            recover))[0], table)
1175
    else:
1176
        join_cols = ['table_schema', 'table_name', 'constraint_schema',
1177
            'constraint_name']
1178
        tables = [sql_gen.Table('key_column_usage', 'information_schema'),
1179
            sql_gen.Join(
1180
                sql_gen.Table('table_constraints', 'information_schema'),
1181
                dict(((c, sql_gen.join_same_not_null) for c in join_cols)))]
1182
        cols = [sql_gen.Col('column_name')]
1183
1184
        conds = [('constraint_type', 'PRIMARY KEY'), ('table_name', table.name)]
1185
        schema = table.schema
1186
        if schema != None: conds.append(('table_schema', schema))
1187
        order_by = 'position_in_unique_constraint'
1188
1189
        try: return sql_gen.Col(value(select(db, tables, cols, conds,
1190
            order_by=order_by, limit=1, log_level=4)), table)
1191
        except StopIteration: raise DoesNotExistException('primary key', '')
1192 5389 aaronmk
1193 5990 aaronmk
def table_has_pkey(db, table, recover=None):
1194
    try: table_pkey_col(db, table, recover)
1195
    except DoesNotExistException: return False
1196
    else: return True
1197
1198 5389 aaronmk
def pkey_name(db, table, recover=None):
1199
    '''If no pkey, returns the first column in the table.'''
1200 5392 aaronmk
    return pkey_col(db, table, recover).name
1201 832 aaronmk
1202 5390 aaronmk
def pkey_col(db, table, recover=None):
1203 5391 aaronmk
    '''If no pkey, returns the first column in the table.'''
1204 5392 aaronmk
    try: return table_pkey_col(db, table, recover)
1205 5393 aaronmk
    except DoesNotExistException: return table_cols(db, table, recover)[0]
1206 5128 aaronmk
1207 2559 aaronmk
not_null_col = 'not_null_col'
1208 2340 aaronmk
1209
def table_not_null_col(db, table, recover=None):
1210
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1211 5337 aaronmk
    if not_null_col in table_col_names(db, table, recover): return not_null_col
1212 5388 aaronmk
    else: return pkey_name(db, table, recover)
1213 2340 aaronmk
1214 3348 aaronmk
def constraint_cond(db, constraint):
1215
    module = util.root_module(db.db)
1216
    if module == 'psycopg2':
1217
        table_str = sql_gen.Literal(constraint.table.to_str(db))
1218
        name_str = sql_gen.Literal(constraint.name)
1219
        return value(run_query(db, '''\
1220
SELECT consrc
1221
FROM pg_constraint
1222
WHERE
1223
conrelid = '''+table_str.to_str(db)+'''::regclass
1224
AND conname = '''+name_str.to_str(db)+'''
1225
'''
1226
            , cacheable=True, log_level=4))
1227 5443 aaronmk
    else: raise NotImplementedError("Can't get constraint condition for "
1228
        +module+' database')
1229 3348 aaronmk
1230 5520 aaronmk
def index_exprs(db, index):
1231 3322 aaronmk
    index = sql_gen.as_Table(index)
1232 1909 aaronmk
    module = util.root_module(db.db)
1233
    if module == 'psycopg2':
1234 3322 aaronmk
        qual_index = sql_gen.Literal(index.to_str(db))
1235 5520 aaronmk
        return list(values(run_query(db, '''\
1236 3322 aaronmk
SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
1237
FROM pg_index
1238
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1239 2782 aaronmk
'''
1240
            , cacheable=True, log_level=4)))
1241 5520 aaronmk
    else: raise NotImplementedError()
1242 853 aaronmk
1243 5520 aaronmk
def index_cols(db, index):
1244
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1245
    automatically created. When you don't know whether something is a UNIQUE
1246
    constraint or a UNIQUE index, use this function.'''
1247
    return map(sql_gen.parse_expr_col, index_exprs(db, index))
1248
1249 5445 aaronmk
def index_cond(db, index):
1250
    index = sql_gen.as_Table(index)
1251
    module = util.root_module(db.db)
1252
    if module == 'psycopg2':
1253
        qual_index = sql_gen.Literal(index.to_str(db))
1254
        return value(run_query(db, '''\
1255
SELECT pg_get_expr(indpred, indrelid, true)
1256
FROM pg_index
1257
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1258
'''
1259
            , cacheable=True, log_level=4))
1260
    else: raise NotImplementedError()
1261
1262 5521 aaronmk
def index_order_by(db, index):
1263
    return sql_gen.CustomCode(', '.join(index_exprs(db, index)))
1264
1265
def table_cluster_on(db, table, recover=None):
1266
    '''
1267
    @return The table's cluster index, or its pkey if none is set
1268
    '''
1269
    table_str = sql_gen.Literal(table.to_str(db))
1270
    try:
1271
        return sql_gen.Table(value(run_query(db, '''\
1272
SELECT relname
1273
FROM pg_index
1274
JOIN pg_class index ON index.oid = indexrelid
1275
WHERE
1276
indrelid = '''+table_str.to_str(db)+'''::regclass
1277
AND indisclustered
1278
'''
1279
            , recover, cacheable=True, log_level=4)), table.schema)
1280
    except StopIteration: return table_pkey_index(db, table, recover)
1281
1282
def table_order_by(db, table, recover=None):
1283 8223 aaronmk
    '''
1284
    @return None if table is view, because table_cluster_on() would return None
1285
    '''
1286 5525 aaronmk
    if table.order_by == None:
1287
        try: table.order_by = index_order_by(db, table_cluster_on(db, table,
1288
            recover))
1289
        except DoesNotExistException: pass
1290
    return table.order_by
1291 5521 aaronmk
1292 7242 aaronmk
#### Views
1293
1294
def view_exists(db, view):
1295
    view_str = sql_gen.Literal(view.to_str(db))
1296
    try:
1297
        return value(run_query(db, '''\
1298
SELECT relkind = 'v'
1299
FROM pg_class
1300
WHERE oid = '''+view_str.to_str(db)+'''::regclass
1301
'''
1302
            , cacheable=True, log_level=4))
1303
    except DoesNotExistException: return False
1304
1305 10139 aaronmk
#### Columns
1306
1307
def col_default_value(db, col):
1308
    assert isinstance(col, sql_gen.Col)
1309
    return value(run_query(db,
1310
        'SELECT util.col_default_value('+sql_gen.col2col_ref(db, col)+')'))
1311
1312
def col_is_constant(db, col):
1313
    assert isinstance(col, sql_gen.Col)
1314
    return value(run_query(db,
1315
        'SELECT util.is_constant('+sql_gen.col2col_ref(db, col)+')'))
1316
1317 3079 aaronmk
#### Functions
1318
1319
def function_exists(db, function):
1320 3423 aaronmk
    qual_function = sql_gen.Literal(function.to_str(db))
1321
    try:
1322 3425 aaronmk
        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
1323
            recover=True, cacheable=True, log_level=4)
1324 3423 aaronmk
    except DoesNotExistException: return False
1325 4146 aaronmk
    except DuplicateException: return True # overloaded function
1326 3423 aaronmk
    else: return True
1327 3079 aaronmk
1328 5713 aaronmk
def function_param0_type(db, function):
1329
    qual_function = sql_gen.Literal(function.to_str(db))
1330
    return value(run_query(db, '''\
1331
SELECT proargtypes[0]::regtype
1332
FROM pg_proc
1333
WHERE oid = '''+qual_function.to_str(db)+'''::regproc
1334
'''
1335
        , cacheable=True, log_level=4))
1336
1337 3079 aaronmk
##### Structural changes
1338
1339
#### Columns
1340
1341 5020 aaronmk
def add_col(db, table, col, comment=None, if_not_exists=False, **kw_args):
1342 3079 aaronmk
    '''
1343
    @param col TypedCol Name may be versioned, so be sure to propagate any
1344
        renaming back to any source column for the TypedCol.
1345
    @param comment None|str SQL comment used to distinguish columns of the same
1346
        name from each other when they contain different data, to allow the
1347
        ADD COLUMN query to be cached. If not set, query will not be cached.
1348
    '''
1349
    assert isinstance(col, sql_gen.TypedCol)
1350
1351
    while True:
1352
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1353
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1354
1355
        try:
1356
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1357
            break
1358
        except DuplicateException:
1359 5020 aaronmk
            if if_not_exists: raise
1360 3079 aaronmk
            col.name = next_version(col.name)
1361
            # try again with next version of name
1362
1363
def add_not_null(db, col):
1364
    table = col.table
1365
    col = sql_gen.to_name_only_col(col)
1366
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1367
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1368
1369 4443 aaronmk
def drop_not_null(db, col):
1370
    table = col.table
1371
    col = sql_gen.to_name_only_col(col)
1372
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1373
        +col.to_str(db)+' DROP NOT NULL', cacheable=True, log_level=3)
1374
1375 2096 aaronmk
row_num_col = '_row_num'
1376
1377 4997 aaronmk
row_num_col_def = sql_gen.TypedCol('', 'serial', nullable=False,
1378 3079 aaronmk
    constraints='PRIMARY KEY')
1379
1380 4997 aaronmk
def add_row_num(db, table, name=row_num_col):
1381
    '''Adds a row number column to a table. Its definition is in
1382
    row_num_col_def. It will be the primary key.'''
1383
    col_def = copy.copy(row_num_col_def)
1384
    col_def.name = name
1385 5021 aaronmk
    add_col(db, table, col_def, comment='', if_not_exists=True, log_level=3)
1386 3079 aaronmk
1387
#### Indexes
1388
1389
def add_pkey(db, table, cols=None, recover=None):
1390
    '''Adds a primary key.
1391
    @param cols [sql_gen.Col,...] The columns in the primary key.
1392
        Defaults to the first column in the table.
1393
    @pre The table must not already have a primary key.
1394
    '''
1395
    table = sql_gen.as_Table(table)
1396 5388 aaronmk
    if cols == None: cols = [pkey_name(db, table, recover)]
1397 3079 aaronmk
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1398
1399
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1400
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1401
        log_ignore_excs=(DuplicateException,))
1402
1403 2998 aaronmk
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1404 2688 aaronmk
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1405 3356 aaronmk
    Currently, only function calls and literal values are supported expressions.
1406 2998 aaronmk
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1407 2847 aaronmk
        This allows indexes to be used for comparisons where NULLs are equal.
1408 2538 aaronmk
    '''
1409 2964 aaronmk
    exprs = lists.mk_seq(exprs)
1410 2538 aaronmk
1411 2688 aaronmk
    # Parse exprs
1412
    old_exprs = exprs[:]
1413
    exprs = []
1414
    cols = []
1415
    for i, expr in enumerate(old_exprs):
1416 2823 aaronmk
        expr = sql_gen.as_Col(expr, table)
1417 2688 aaronmk
1418 2823 aaronmk
        # Handle nullable columns
1419 2998 aaronmk
        if ensure_not_null_:
1420 3164 aaronmk
            try: expr = sql_gen.ensure_not_null(db, expr)
1421 2860 aaronmk
            except KeyError: pass # unknown type, so just create plain index
1422 2823 aaronmk
1423 2688 aaronmk
        # Extract col
1424 3002 aaronmk
        expr = copy.deepcopy(expr) # don't modify input!
1425 3356 aaronmk
        col = expr
1426
        if isinstance(expr, sql_gen.FunctionCall): col = expr.args[0]
1427
        expr = sql_gen.cast_literal(expr)
1428
        if not isinstance(expr, (sql_gen.Expr, sql_gen.Col)):
1429 2688 aaronmk
            expr = sql_gen.Expr(expr)
1430 3356 aaronmk
1431 2688 aaronmk
1432
        # Extract table
1433
        if table == None:
1434
            assert sql_gen.is_table_col(col)
1435
            table = col.table
1436
1437 3356 aaronmk
        if isinstance(col, sql_gen.Col): col.table = None
1438 2688 aaronmk
1439
        exprs.append(expr)
1440
        cols.append(col)
1441 2408 aaronmk
1442 2688 aaronmk
    table = sql_gen.as_Table(table)
1443
1444 3005 aaronmk
    # Add index
1445 3148 aaronmk
    str_ = 'CREATE'
1446
    if unique: str_ += ' UNIQUE'
1447
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1448
        ', '.join((v.to_str(db) for v in exprs)))+')'
1449
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1450 2408 aaronmk
1451 5987 aaronmk
def add_pkey_index(db, table): add_index(db, pkey_col(db, table), table)
1452
1453 5765 aaronmk
def add_pkey_or_index(db, table, cols=None, recover=None, warn=False):
1454
    try: add_pkey(db, table, cols, recover)
1455
    except DuplicateKeyException, e:
1456
        if warn: warnings.warn(UserWarning(exc.str_(e)))
1457 5988 aaronmk
        add_pkey_index(db, table)
1458 5765 aaronmk
1459 3083 aaronmk
already_indexed = object() # tells add_indexes() the pkey has already been added
1460
1461
def add_indexes(db, table, has_pkey=True):
1462
    '''Adds an index on all columns in a table.
1463
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1464
        index should be added on the first column.
1465
        * If already_indexed, the pkey is assumed to have already been added
1466
    '''
1467 5337 aaronmk
    cols = table_col_names(db, table)
1468 3083 aaronmk
    if has_pkey:
1469
        if has_pkey is not already_indexed: add_pkey(db, table)
1470
        cols = cols[1:]
1471
    for col in cols: add_index(db, col, table)
1472
1473 3079 aaronmk
#### Tables
1474 2772 aaronmk
1475 3079 aaronmk
### Maintenance
1476 2772 aaronmk
1477 3079 aaronmk
def analyze(db, table):
1478
    table = sql_gen.as_Table(table)
1479
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1480 2934 aaronmk
1481 3079 aaronmk
def autoanalyze(db, table):
1482
    if db.autoanalyze: analyze(db, table)
1483 2935 aaronmk
1484 3079 aaronmk
def vacuum(db, table):
1485
    table = sql_gen.as_Table(table)
1486
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1487
        log_level=3))
1488 2086 aaronmk
1489 3079 aaronmk
### Lifecycle
1490
1491 3247 aaronmk
def drop(db, type_, name):
1492
    name = sql_gen.as_Name(name)
1493
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1494 2889 aaronmk
1495 3247 aaronmk
def drop_table(db, table): drop(db, 'TABLE', table)
1496
1497 3082 aaronmk
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1498 12151 aaronmk
    like=None, **query_opts):
1499 2675 aaronmk
    '''Creates a table.
1500 2681 aaronmk
    @param cols [sql_gen.TypedCol,...] The column names and types
1501
    @param has_pkey If set, the first column becomes the primary key.
1502 2760 aaronmk
    @param col_indexes bool|[ref]
1503
        * If True, indexes will be added on all non-pkey columns.
1504
        * If a list reference, [0] will be set to a function to do this.
1505
          This can be used to delay index creation until the table is populated.
1506 2675 aaronmk
    '''
1507 12151 aaronmk
    query_opts.setdefault('cacheable', True)
1508
1509 2675 aaronmk
    table = sql_gen.as_Table(table)
1510
1511 3082 aaronmk
    if like != None:
1512
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1513
            ]+cols
1514 5525 aaronmk
        table.order_by = like.order_by
1515 2681 aaronmk
    if has_pkey:
1516
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1517 2872 aaronmk
        pkey.constraints = 'PRIMARY KEY'
1518 2681 aaronmk
1519 3085 aaronmk
    temp = table.is_temp and not db.debug_temp
1520
        # temp tables permanent in debug_temp mode
1521 2760 aaronmk
1522 3085 aaronmk
    # Create table
1523 3383 aaronmk
    def create():
1524 3085 aaronmk
        str_ = 'CREATE'
1525
        if temp: str_ += ' TEMP'
1526
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1527
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1528 3126 aaronmk
        str_ += '\n);'
1529 3085 aaronmk
1530 12151 aaronmk
        opts = dict(recover=True, log_level=2,
1531
            log_ignore_excs=(DuplicateException,), **query_opts)
1532 7113 aaronmk
        try: run_query(db, str_, **opts)
1533
        except InvalidTypeException: # try again as view
1534
            run_query_into(db, mk_select(db, like, limit=0), into=table, **opts)
1535 3383 aaronmk
    if table.is_temp:
1536
        while True:
1537
            try:
1538
                create()
1539
                break
1540
            except DuplicateException:
1541
                table.name = next_version(table.name)
1542
                # try again with next version of name
1543
    else: create()
1544 3085 aaronmk
1545 2760 aaronmk
    # Add indexes
1546 2773 aaronmk
    if has_pkey: has_pkey = already_indexed
1547
    def add_indexes_(): add_indexes(db, table, has_pkey)
1548
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1549
    elif col_indexes: add_indexes_() # add now
1550 2675 aaronmk
1551 12151 aaronmk
def copy_table_struct(db, src, dest, **opts):
1552 3084 aaronmk
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1553 12151 aaronmk
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src, **opts)
1554 3084 aaronmk
1555 5529 aaronmk
def copy_table(db, src, dest):
1556
    '''Creates a copy of a table, including data'''
1557
    copy_table_struct(db, src, dest)
1558
    insert_select(db, dest, None, mk_select(db, src))
1559
1560 3079 aaronmk
### Data
1561 2684 aaronmk
1562 2970 aaronmk
def truncate(db, table, schema='public', **kw_args):
1563
    '''For params, see run_query()'''
1564 2777 aaronmk
    table = sql_gen.as_Table(table, schema)
1565 2970 aaronmk
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1566 2732 aaronmk
1567 2965 aaronmk
def empty_temp(db, tables):
1568
    tables = lists.mk_seq(tables)
1569 2971 aaronmk
    for table in tables: truncate(db, table, log_level=3)
1570 2965 aaronmk
1571 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
1572
    '''For kw_args, see tables()'''
1573
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1574 3094 aaronmk
1575 10841 aaronmk
def distinct_table(db, table, distinct_on, joins=None, conds=None):
1576 3094 aaronmk
    '''Creates a copy of a temp table which is distinct on the given columns.
1577 5902 aaronmk
    Adds an index on table's distinct_on columns, to facilitate merge joins.
1578 3097 aaronmk
    @param distinct_on If empty, creates a table with one row. This is useful if
1579
        your distinct_on columns are all literal values.
1580 7179 aaronmk
    @param joins The joins to use when creating the new table
1581 10841 aaronmk
    @param conds Any additional filters to use when creating the new table
1582 3099 aaronmk
    @return The new table.
1583 3094 aaronmk
    '''
1584 7179 aaronmk
    if joins == None: joins = [table]
1585
1586 3099 aaronmk
    new_table = sql_gen.suffixed_table(table, '_distinct')
1587 3411 aaronmk
    distinct_on = filter(sql_gen.is_table_col, distinct_on)
1588 3094 aaronmk
1589 12152 aaronmk
    copy_table_struct(db, table, new_table, cacheable=False)
1590 3097 aaronmk
1591
    limit = None
1592
    if distinct_on == []: limit = 1 # one sample row
1593 5903 aaronmk
    else: add_index(db, distinct_on, table) # for join optimization
1594 3097 aaronmk
1595 7179 aaronmk
    insert_select(db, new_table, None, mk_select(db, joins,
1596 10841 aaronmk
        [sql_gen.Col(sql_gen.all_cols, table)], conds, distinct_on,
1597 7179 aaronmk
        order_by=None, limit=limit))
1598 3099 aaronmk
    analyze(db, new_table)
1599 3094 aaronmk
1600 3099 aaronmk
    return new_table