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 2484 aaronmk
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name)), cause)
41 2143 aaronmk
        self.name = name
42 360 aaronmk
43 3109 aaronmk
class ExceptionWithValue(DbException):
44
    def __init__(self, value, cause=None):
45
        DbException.__init__(self, 'for value: '+strings.as_tt(repr(value)),
46
            cause)
47 2240 aaronmk
        self.value = value
48
49 2945 aaronmk
class ExceptionWithNameType(DbException):
50
    def __init__(self, type_, name, cause=None):
51
        DbException.__init__(self, 'for type: '+strings.as_tt(str(type_))
52
            +'; name: '+strings.as_tt(name), cause)
53
        self.type = type_
54
        self.name = name
55
56 2306 aaronmk
class ConstraintException(DbException):
57
    def __init__(self, name, cols, cause=None):
58 2484 aaronmk
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
59
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
60 2306 aaronmk
        self.name = name
61 468 aaronmk
        self.cols = cols
62 11 aaronmk
63 2523 aaronmk
class MissingCastException(DbException):
64
    def __init__(self, type_, col, cause=None):
65
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
66
            +' on column: '+strings.as_tt(col), cause)
67
        self.type = type_
68
        self.col = col
69
70 2143 aaronmk
class NameException(DbException): pass
71
72 2306 aaronmk
class DuplicateKeyException(ConstraintException): pass
73 13 aaronmk
74 2306 aaronmk
class NullValueException(ConstraintException): pass
75 13 aaronmk
76 3109 aaronmk
class InvalidValueException(ExceptionWithValue): pass
77 2239 aaronmk
78 2945 aaronmk
class DuplicateException(ExceptionWithNameType): pass
79 2143 aaronmk
80 89 aaronmk
class EmptyRowException(DbException): pass
81
82 865 aaronmk
##### Warnings
83
84
class DbWarning(UserWarning): pass
85
86 1930 aaronmk
##### Result retrieval
87
88
def col_names(cur): return (col[0] for col in cur.description)
89
90
def rows(cur): return iter(lambda: cur.fetchone(), None)
91
92
def consume_rows(cur):
93
    '''Used to fetch all rows so result will be cached'''
94
    iters.consume_iter(rows(cur))
95
96
def next_row(cur): return rows(cur).next()
97
98
def row(cur):
99
    row_ = next_row(cur)
100
    consume_rows(cur)
101
    return row_
102
103
def next_value(cur): return next_row(cur)[0]
104
105
def value(cur): return row(cur)[0]
106
107
def values(cur): return iters.func_iter(lambda: next_value(cur))
108
109
def value_or_none(cur):
110
    try: return value(cur)
111
    except StopIteration: return None
112
113 2762 aaronmk
##### Escaping
114 2101 aaronmk
115 2573 aaronmk
def esc_name_by_module(module, name):
116
    if module == 'psycopg2' or module == None: quote = '"'
117 2101 aaronmk
    elif module == 'MySQLdb': quote = '`'
118
    else: raise NotImplementedError("Can't escape name for "+module+' database')
119 2500 aaronmk
    return sql_gen.esc_name(name, quote)
120 2101 aaronmk
121
def esc_name_by_engine(engine, name, **kw_args):
122
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
123
124
def esc_name(db, name, **kw_args):
125
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
126
127
def qual_name(db, schema, table):
128
    def esc_name_(name): return esc_name(db, name)
129
    table = esc_name_(table)
130
    if schema != None: return esc_name_(schema)+'.'+table
131
    else: return table
132
133 1869 aaronmk
##### Database connections
134 1849 aaronmk
135 2097 aaronmk
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
136 1926 aaronmk
137 1869 aaronmk
db_engines = {
138
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
139
    'PostgreSQL': ('psycopg2', {}),
140
}
141
142
DatabaseErrors_set = set([DbException])
143
DatabaseErrors = tuple(DatabaseErrors_set)
144
145
def _add_module(module):
146
    DatabaseErrors_set.add(module.DatabaseError)
147
    global DatabaseErrors
148
    DatabaseErrors = tuple(DatabaseErrors_set)
149
150
def db_config_str(db_config):
151
    return db_config['engine']+' database '+db_config['database']
152
153 2448 aaronmk
log_debug_none = lambda msg, level=2: None
154 1901 aaronmk
155 1849 aaronmk
class DbConn:
156 2923 aaronmk
    def __init__(self, db_config, autocommit=True, caching=True,
157 3183 aaronmk
        log_debug=log_debug_none, debug_temp=False, src=None):
158 2915 aaronmk
        '''
159
        @param debug_temp Whether temporary objects should instead be permanent.
160
            This assists in debugging the internal objects used by the program.
161 3183 aaronmk
        @param src In autocommit mode, will be included in a comment in every
162
            query, to help identify the data source in pg_stat_activity.
163 2915 aaronmk
        '''
164 1869 aaronmk
        self.db_config = db_config
165 2190 aaronmk
        self.autocommit = autocommit
166
        self.caching = caching
167 1901 aaronmk
        self.log_debug = log_debug
168 2193 aaronmk
        self.debug = log_debug != log_debug_none
169 2915 aaronmk
        self.debug_temp = debug_temp
170 3244 aaronmk
        self.profile_row_ct = None
171 3183 aaronmk
        self.src = src
172 3074 aaronmk
        self.autoanalyze = False
173 1869 aaronmk
174 3124 aaronmk
        self._savepoint = 0
175 3120 aaronmk
        self._reset()
176 1869 aaronmk
177
    def __getattr__(self, name):
178
        if name == '__dict__': raise Exception('getting __dict__')
179
        if name == 'db': return self._db()
180
        else: raise AttributeError()
181
182
    def __getstate__(self):
183
        state = copy.copy(self.__dict__) # shallow copy
184 1915 aaronmk
        state['log_debug'] = None # don't pickle the debug callback
185 1869 aaronmk
        state['_DbConn__db'] = None # don't pickle the connection
186
        return state
187
188 3118 aaronmk
    def clear_cache(self): self.query_results = {}
189
190 3120 aaronmk
    def _reset(self):
191 3118 aaronmk
        self.clear_cache()
192 3124 aaronmk
        assert self._savepoint == 0
193 3118 aaronmk
        self._notices_seen = set()
194
        self.__db = None
195
196 2165 aaronmk
    def connected(self): return self.__db != None
197
198 3116 aaronmk
    def close(self):
199 3119 aaronmk
        if not self.connected(): return
200
201 3135 aaronmk
        # Record that the automatic transaction is now closed
202 3136 aaronmk
        self._savepoint -= 1
203 3135 aaronmk
204 3119 aaronmk
        self.db.close()
205 3120 aaronmk
        self._reset()
206 3116 aaronmk
207 3125 aaronmk
    def reconnect(self):
208
        # Do not do this in test mode as it would roll back everything
209
        if self.autocommit: self.close()
210
        # Connection will be reopened automatically on first query
211
212 1869 aaronmk
    def _db(self):
213
        if self.__db == None:
214
            # Process db_config
215
            db_config = self.db_config.copy() # don't modify input!
216 2097 aaronmk
            schemas = db_config.pop('schemas', None)
217 1869 aaronmk
            module_name, mappings = db_engines[db_config.pop('engine')]
218
            module = __import__(module_name)
219
            _add_module(module)
220
            for orig, new in mappings.iteritems():
221
                try: util.rename_key(db_config, orig, new)
222
                except KeyError: pass
223
224
            # Connect
225
            self.__db = module.connect(**db_config)
226
227 3161 aaronmk
            # Record that a transaction is already open
228
            self._savepoint += 1
229
230 1869 aaronmk
            # Configure connection
231 2906 aaronmk
            if hasattr(self.db, 'set_isolation_level'):
232
                import psycopg2.extensions
233
                self.db.set_isolation_level(
234
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
235 2101 aaronmk
            if schemas != None:
236 2893 aaronmk
                search_path = [self.esc_name(s) for s in schemas.split(',')]
237
                search_path.append(value(run_query(self, 'SHOW search_path',
238
                    log_level=4)))
239
                run_query(self, 'SET search_path TO '+(','.join(search_path)),
240
                    log_level=3)
241 1869 aaronmk
242
        return self.__db
243 1889 aaronmk
244 1891 aaronmk
    class DbCursor(Proxy):
245 1927 aaronmk
        def __init__(self, outer):
246 1891 aaronmk
            Proxy.__init__(self, outer.db.cursor())
247 2191 aaronmk
            self.outer = outer
248 1927 aaronmk
            self.query_results = outer.query_results
249 1894 aaronmk
            self.query_lookup = None
250 1891 aaronmk
            self.result = []
251 1889 aaronmk
252 2802 aaronmk
        def execute(self, query):
253 2764 aaronmk
            self._is_insert = query.startswith('INSERT')
254 2797 aaronmk
            self.query_lookup = query
255 2148 aaronmk
            try:
256 3162 aaronmk
                try: cur = self.inner.execute(query)
257 2802 aaronmk
                finally: self.query = get_cur_query(self.inner, query)
258 1904 aaronmk
            except Exception, e:
259
                self.result = e # cache the exception as the result
260
                self._cache_result()
261
                raise
262 3004 aaronmk
263
            # Always cache certain queries
264 3183 aaronmk
            query = sql_gen.lstrip(query)
265 3004 aaronmk
            if query.startswith('CREATE') or query.startswith('ALTER'):
266 3007 aaronmk
                # structural changes
267 3040 aaronmk
                # Rest of query must be unique in the face of name collisions,
268
                # so don't cache ADD COLUMN unless it has distinguishing comment
269
                if query.find('ADD COLUMN') < 0 or query.endswith('*/'):
270 3007 aaronmk
                    self._cache_result()
271 3004 aaronmk
            elif self.rowcount == 0 and query.startswith('SELECT'): # empty
272 2800 aaronmk
                consume_rows(self) # fetch all rows so result will be cached
273 3004 aaronmk
274 2762 aaronmk
            return cur
275 1894 aaronmk
276 1891 aaronmk
        def fetchone(self):
277
            row = self.inner.fetchone()
278 1899 aaronmk
            if row != None: self.result.append(row)
279
            # otherwise, fetched all rows
280 1904 aaronmk
            else: self._cache_result()
281
            return row
282
283
        def _cache_result(self):
284 2948 aaronmk
            # For inserts that return a result set, don't cache result set since
285
            # inserts are not idempotent. Other non-SELECT queries don't have
286
            # their result set read, so only exceptions will be cached (an
287
            # invalid query will always be invalid).
288 1930 aaronmk
            if self.query_results != None and (not self._is_insert
289 1906 aaronmk
                or isinstance(self.result, Exception)):
290
291 1894 aaronmk
                assert self.query_lookup != None
292 1916 aaronmk
                self.query_results[self.query_lookup] = self.CacheCursor(
293
                    util.dict_subset(dicts.AttrsDictView(self),
294
                    ['query', 'result', 'rowcount', 'description']))
295 1906 aaronmk
296 1916 aaronmk
        class CacheCursor:
297
            def __init__(self, cached_result): self.__dict__ = cached_result
298
299 1927 aaronmk
            def execute(self, *args, **kw_args):
300 1916 aaronmk
                if isinstance(self.result, Exception): raise self.result
301
                # otherwise, result is a rows list
302
                self.iter = iter(self.result)
303
304
            def fetchone(self):
305
                try: return self.iter.next()
306
                except StopIteration: return None
307 1891 aaronmk
308 2212 aaronmk
    def esc_value(self, value):
309 2663 aaronmk
        try: str_ = self.mogrify('%s', [value])
310
        except NotImplementedError, e:
311
            module = util.root_module(self.db)
312
            if module == 'MySQLdb':
313
                import _mysql
314
                str_ = _mysql.escape_string(value)
315
            else: raise e
316 2374 aaronmk
        return strings.to_unicode(str_)
317 2212 aaronmk
318 2347 aaronmk
    def esc_name(self, name): return esc_name(self, name) # calls global func
319
320 2814 aaronmk
    def std_code(self, str_):
321
        '''Standardizes SQL code.
322
        * Ensures that string literals are prefixed by `E`
323
        '''
324
        if str_.startswith("'"): str_ = 'E'+str_
325
        return str_
326
327 2665 aaronmk
    def can_mogrify(self):
328 2663 aaronmk
        module = util.root_module(self.db)
329 2665 aaronmk
        return module == 'psycopg2'
330 2663 aaronmk
331 2665 aaronmk
    def mogrify(self, query, params=None):
332
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
333
        else: raise NotImplementedError("Can't mogrify query")
334
335 2671 aaronmk
    def print_notices(self):
336 2725 aaronmk
        if hasattr(self.db, 'notices'):
337
            for msg in self.db.notices:
338
                if msg not in self._notices_seen:
339
                    self._notices_seen.add(msg)
340
                    self.log_debug(msg, level=2)
341 2671 aaronmk
342 2793 aaronmk
    def run_query(self, query, cacheable=False, log_level=2,
343 2464 aaronmk
        debug_msg_ref=None):
344 2445 aaronmk
        '''
345 2464 aaronmk
        @param log_ignore_excs The log_level will be increased by 2 if the query
346
            throws one of these exceptions.
347 2664 aaronmk
        @param debug_msg_ref If specified, the log message will be returned in
348
            this instead of being output. This allows you to filter log messages
349
            depending on the result of the query.
350 2445 aaronmk
        '''
351 2167 aaronmk
        assert query != None
352
353 3183 aaronmk
        if self.autocommit and self.src != None:
354 3206 aaronmk
            query = sql_gen.esc_comment(self.src)+'\t'+query
355 3183 aaronmk
356 2047 aaronmk
        if not self.caching: cacheable = False
357 1903 aaronmk
        used_cache = False
358 2664 aaronmk
359 3242 aaronmk
        if self.debug:
360
            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
361 1903 aaronmk
        try:
362 1927 aaronmk
            # Get cursor
363
            if cacheable:
364 3238 aaronmk
                try: cur = self.query_results[query]
365 1927 aaronmk
                except KeyError: cur = self.DbCursor(self)
366 3238 aaronmk
                else: used_cache = True
367 1927 aaronmk
            else: cur = self.db.cursor()
368
369
            # Run query
370 3238 aaronmk
            try: cur.execute(query)
371 3162 aaronmk
            except Exception, e:
372
                _add_cursor_info(e, self, query)
373
                raise
374 3238 aaronmk
            else: self.do_autocommit()
375 1903 aaronmk
        finally:
376 3242 aaronmk
            if self.debug:
377 3244 aaronmk
                profiler.stop(self.profile_row_ct)
378 3242 aaronmk
379
                ## Log or return query
380
381 3239 aaronmk
                query = str(get_cur_query(cur, query))
382
                # So that the src comment is hidden when cating the file, and
383
                # put on a separate line when viewed in a text editor.
384
                query = query.replace('\t', '\r', 1)
385
386 3240 aaronmk
                msg = 'DB query: '
387 3239 aaronmk
388 3240 aaronmk
                if used_cache: msg += 'cache hit'
389
                elif cacheable: msg += 'cache miss'
390
                else: msg += 'non-cacheable'
391 3239 aaronmk
392 3241 aaronmk
                msg += ':\n'+profiler.msg()+'\n'+strings.as_code(query, 'SQL')
393 3240 aaronmk
394 3237 aaronmk
                if debug_msg_ref != None: debug_msg_ref[0] = msg
395
                else: self.log_debug(msg, log_level)
396 3245 aaronmk
397
                self.print_notices()
398 1903 aaronmk
399
        return cur
400 1914 aaronmk
401 2797 aaronmk
    def is_cached(self, query): return query in self.query_results
402 2139 aaronmk
403 2907 aaronmk
    def with_autocommit(self, func):
404 2801 aaronmk
        import psycopg2.extensions
405
406
        prev_isolation_level = self.db.isolation_level
407 2907 aaronmk
        self.db.set_isolation_level(
408
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
409 2683 aaronmk
        try: return func()
410 2801 aaronmk
        finally: self.db.set_isolation_level(prev_isolation_level)
411 2683 aaronmk
412 2139 aaronmk
    def with_savepoint(self, func):
413 3137 aaronmk
        top = self._savepoint == 0
414 2171 aaronmk
        savepoint = 'level_'+str(self._savepoint)
415 3137 aaronmk
416 3160 aaronmk
        # Must happen before running queries so they don't get autocommitted
417
        self._savepoint += 1
418
419 3137 aaronmk
        if top: query = 'START TRANSACTION ISOLATION LEVEL READ COMMITTED'
420
        else: query = 'SAVEPOINT '+savepoint
421
        self.run_query(query, log_level=4)
422
        try:
423
            return func()
424
            if top: self.run_query('COMMIT', log_level=4)
425 2139 aaronmk
        except:
426 3137 aaronmk
            if top: query = 'ROLLBACK'
427
            else: query = 'ROLLBACK TO SAVEPOINT '+savepoint
428
            self.run_query(query, log_level=4)
429
430 2139 aaronmk
            raise
431 2930 aaronmk
        finally:
432
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
433
            # "The savepoint remains valid and can be rolled back to again"
434
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
435 3137 aaronmk
            if not top:
436
                self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
437 2930 aaronmk
438
            self._savepoint -= 1
439
            assert self._savepoint >= 0
440
441
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
442 2191 aaronmk
443
    def do_autocommit(self):
444
        '''Autocommits if outside savepoint'''
445 3135 aaronmk
        assert self._savepoint >= 1
446
        if self.autocommit and self._savepoint == 1:
447 2924 aaronmk
            self.log_debug('Autocommitting', level=4)
448 2191 aaronmk
            self.db.commit()
449 2643 aaronmk
450 3155 aaronmk
    def col_info(self, col, cacheable=True):
451 2643 aaronmk
        table = sql_gen.Table('columns', 'information_schema')
452 3063 aaronmk
        type_ = sql_gen.Coalesce(sql_gen.Nullif(sql_gen.Col('data_type'),
453
            'USER-DEFINED'), sql_gen.Col('udt_name'))
454 3078 aaronmk
        cols = [type_, 'column_default',
455
            sql_gen.Cast('boolean', sql_gen.Col('is_nullable'))]
456 2643 aaronmk
457
        conds = [('table_name', col.table.name), ('column_name', col.name)]
458
        schema = col.table.schema
459
        if schema != None: conds.append(('table_schema', schema))
460
461 2819 aaronmk
        type_, default, nullable = row(select(self, table, cols, conds,
462 3150 aaronmk
            order_by='table_schema', limit=1, cacheable=cacheable, log_level=4))
463 2643 aaronmk
            # TODO: order_by search_path schema order
464 2819 aaronmk
        default = sql_gen.as_Code(default, self)
465
466
        return sql_gen.TypedCol(col.name, type_, default, nullable)
467 2917 aaronmk
468
    def TempFunction(self, name):
469
        if self.debug_temp: schema = None
470
        else: schema = 'pg_temp'
471
        return sql_gen.Function(name, schema)
472 1849 aaronmk
473 1869 aaronmk
connect = DbConn
474
475 832 aaronmk
##### Recoverable querying
476 15 aaronmk
477 2139 aaronmk
def with_savepoint(db, func): return db.with_savepoint(func)
478 11 aaronmk
479 2791 aaronmk
def run_query(db, query, recover=None, cacheable=False, log_level=2,
480
    log_ignore_excs=None, **kw_args):
481 2794 aaronmk
    '''For params, see DbConn.run_query()'''
482 830 aaronmk
    if recover == None: recover = False
483 2464 aaronmk
    if log_ignore_excs == None: log_ignore_excs = ()
484
    log_ignore_excs = tuple(log_ignore_excs)
485 3236 aaronmk
    debug_msg_ref = [None]
486 830 aaronmk
487 3258 aaronmk
    if db.debug and is_explainable(query):
488
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'+explain(db, query))
489
490 2148 aaronmk
    try:
491 2464 aaronmk
        try:
492 2794 aaronmk
            def run(): return db.run_query(query, cacheable, log_level,
493 2793 aaronmk
                debug_msg_ref, **kw_args)
494 2796 aaronmk
            if recover and not db.is_cached(query):
495 2464 aaronmk
                return with_savepoint(db, run)
496
            else: return run() # don't need savepoint if cached
497
        except Exception, e:
498 3095 aaronmk
            msg = strings.ustr(e.args[0])
499 2464 aaronmk
500 3095 aaronmk
            match = re.match(r'^duplicate key value violates unique constraint '
501 3096 aaronmk
                r'"((_?[^\W_]+(?=[._]))?.+?)"', msg)
502 2464 aaronmk
            if match:
503
                constraint, table = match.groups()
504 3025 aaronmk
                cols = []
505
                if recover: # need auto-rollback to run index_cols()
506
                    try: cols = index_cols(db, table, constraint)
507
                    except NotImplementedError: pass
508
                raise DuplicateKeyException(constraint, cols, e)
509 2464 aaronmk
510 3095 aaronmk
            match = re.match(r'^null value in column "(.+?)" violates not-null'
511 2464 aaronmk
                r' constraint', msg)
512
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
513
514 3095 aaronmk
            match = re.match(r'^(?:invalid input (?:syntax|value)\b.*?'
515 3109 aaronmk
                r'|.+? field value out of range): "(.+?)"', msg)
516 2464 aaronmk
            if match:
517 3109 aaronmk
                value, = match.groups()
518
                raise InvalidValueException(strings.to_unicode(value), e)
519 2464 aaronmk
520 3095 aaronmk
            match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
521 2523 aaronmk
                r'is of type', msg)
522
            if match:
523
                col, type_ = match.groups()
524
                raise MissingCastException(type_, col, e)
525
526 3095 aaronmk
            match = re.match(r'^(\S+) "(.+?)".*? already exists', msg)
527 2945 aaronmk
            if match:
528
                type_, name = match.groups()
529
                raise DuplicateException(type_, name, e)
530 2464 aaronmk
531
            raise # no specific exception raised
532
    except log_ignore_excs:
533
        log_level += 2
534
        raise
535
    finally:
536 3236 aaronmk
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
537 830 aaronmk
538 832 aaronmk
##### Basic queries
539
540 3256 aaronmk
def is_explainable(query):
541
    # See <http://www.postgresql.org/docs/8.3/static/sql-explain.html#AEN57749>
542 3257 aaronmk
    return re.match(r'^(?:SELECT|INSERT|UPDATE|DELETE|VALUES|EXECUTE|DECLARE)\b'
543
        , query)
544 3256 aaronmk
545
def explain(db, query):
546
    return strings.join_lines(values(run_query(db, 'EXPLAIN '+query,
547
        recover=True, cacheable=True, log_level=3)))
548
        # not a higher log_level because it's useful to see what query is being
549
        # run before it's executed, which EXPLAIN effectively provides
550
551 2153 aaronmk
def next_version(name):
552 2163 aaronmk
    version = 1 # first existing name was version 0
553 2586 aaronmk
    match = re.match(r'^(.*)#(\d+)$', name)
554 2153 aaronmk
    if match:
555 2586 aaronmk
        name, version = match.groups()
556
        version = int(version)+1
557 2932 aaronmk
    return sql_gen.concat(name, '#'+str(version))
558 2153 aaronmk
559 2899 aaronmk
def lock_table(db, table, mode):
560
    table = sql_gen.as_Table(table)
561
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
562
563 2789 aaronmk
def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):
564 2085 aaronmk
    '''Outputs a query to a temp table.
565
    For params, see run_query().
566
    '''
567 2789 aaronmk
    if into == None: return run_query(db, query, **kw_args)
568 2790 aaronmk
569
    assert isinstance(into, sql_gen.Table)
570
571 2992 aaronmk
    into.is_temp = True
572 3008 aaronmk
    # "temporary tables cannot specify a schema name", so remove schema
573
    into.schema = None
574 2992 aaronmk
575 2790 aaronmk
    kw_args['recover'] = True
576 2945 aaronmk
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
577 2790 aaronmk
578 2916 aaronmk
    temp = not db.debug_temp # tables are permanent in debug_temp mode
579 2790 aaronmk
580
    # Create table
581
    while True:
582
        create_query = 'CREATE'
583
        if temp: create_query += ' TEMP'
584
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
585 2385 aaronmk
586 2790 aaronmk
        try:
587
            cur = run_query(db, create_query, **kw_args)
588
                # CREATE TABLE AS sets rowcount to # rows in query
589
            break
590 2945 aaronmk
        except DuplicateException, e:
591 2790 aaronmk
            into.name = next_version(into.name)
592
            # try again with next version of name
593
594
    if add_indexes_: add_indexes(db, into)
595 3075 aaronmk
596
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
597
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
598
    # table is going to be used in complex queries, it is wise to run ANALYZE on
599
    # the temporary table after it is populated."
600
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
601
    # If into is not a temp table, ANALYZE is useful but not required.
602 3073 aaronmk
    analyze(db, into)
603 2790 aaronmk
604
    return cur
605 2085 aaronmk
606 2120 aaronmk
order_by_pkey = object() # tells mk_select() to order by the pkey
607
608 2199 aaronmk
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
609
610 2233 aaronmk
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
611 2293 aaronmk
    start=None, order_by=order_by_pkey, default_table=None):
612 1981 aaronmk
    '''
613 2121 aaronmk
    @param tables The single table to select from, or a list of tables to join
614 2280 aaronmk
        together, with tables after the first being sql_gen.Join objects
615 1981 aaronmk
    @param fields Use None to select all fields in the table
616 2377 aaronmk
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
617 2379 aaronmk
        * container can be any iterable type
618 2399 aaronmk
        * compare_left_side: sql_gen.Code|str (for col name)
619
        * compare_right_side: sql_gen.ValueCond|literal value
620 2199 aaronmk
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
621
        use all columns
622 2786 aaronmk
    @return query
623 1981 aaronmk
    '''
624 2315 aaronmk
    # Parse tables param
625 2964 aaronmk
    tables = lists.mk_seq(tables)
626 2141 aaronmk
    tables = list(tables) # don't modify input! (list() copies input)
627 2315 aaronmk
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
628 2121 aaronmk
629 2315 aaronmk
    # Parse other params
630 2376 aaronmk
    if conds == None: conds = []
631 2650 aaronmk
    elif dicts.is_dict(conds): conds = conds.items()
632 2379 aaronmk
    conds = list(conds) # don't modify input! (list() copies input)
633 3129 aaronmk
    assert limit == None or isinstance(limit, (int, long))
634
    assert start == None or isinstance(start, (int, long))
635 2315 aaronmk
    if order_by is order_by_pkey:
636
        if distinct_on != []: order_by = None
637
        else: order_by = pkey(db, table0, recover=True)
638 865 aaronmk
639 2315 aaronmk
    query = 'SELECT'
640 2056 aaronmk
641 2315 aaronmk
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
642 2056 aaronmk
643 2200 aaronmk
    # DISTINCT ON columns
644 2233 aaronmk
    if distinct_on != []:
645 2467 aaronmk
        query += '\nDISTINCT'
646 2254 aaronmk
        if distinct_on is not distinct_on_all:
647 2200 aaronmk
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
648
649
    # Columns
650 3185 aaronmk
    if query.find('\n') >= 0: whitespace = '\n'
651
    else: whitespace = ' '
652
    if fields == None: query += whitespace+'*'
653 2765 aaronmk
    else:
654
        assert fields != []
655 3185 aaronmk
        if len(fields) > 1: whitespace = '\n'
656
        query += whitespace+('\n, '.join(map(parse_col, fields)))
657 2200 aaronmk
658
    # Main table
659 3185 aaronmk
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
660
    else: whitespace = ' '
661
    query += whitespace+'FROM '+table0.to_str(db)
662 865 aaronmk
663 2122 aaronmk
    # Add joins
664 2271 aaronmk
    left_table = table0
665 2263 aaronmk
    for join_ in tables:
666
        table = join_.table
667 2238 aaronmk
668 2343 aaronmk
        # Parse special values
669
        if join_.type_ is sql_gen.filter_out: # filter no match
670 2376 aaronmk
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
671 2853 aaronmk
                sql_gen.CompareCond(None, '~=')))
672 2343 aaronmk
673 2467 aaronmk
        query += '\n'+join_.to_str(db, left_table)
674 2122 aaronmk
675
        left_table = table
676
677 865 aaronmk
    missing = True
678 2376 aaronmk
    if conds != []:
679 2576 aaronmk
        if len(conds) == 1: whitespace = ' '
680
        else: whitespace = '\n'
681 2578 aaronmk
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
682
            .to_str(db) for l, r in conds], 'WHERE')
683 865 aaronmk
        missing = False
684 2227 aaronmk
    if order_by != None:
685 2467 aaronmk
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
686
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
687 865 aaronmk
    if start != None:
688 2467 aaronmk
        if start != 0: query += '\nOFFSET '+str(start)
689 865 aaronmk
        missing = False
690
    if missing: warnings.warn(DbWarning(
691
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
692
693 2786 aaronmk
    return query
694 11 aaronmk
695 2054 aaronmk
def select(db, *args, **kw_args):
696
    '''For params, see mk_select() and run_query()'''
697
    recover = kw_args.pop('recover', None)
698
    cacheable = kw_args.pop('cacheable', True)
699 2442 aaronmk
    log_level = kw_args.pop('log_level', 2)
700 2054 aaronmk
701 2791 aaronmk
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
702
        log_level=log_level)
703 2054 aaronmk
704 2788 aaronmk
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
705 3181 aaronmk
    embeddable=False, ignore=False, src=None):
706 1960 aaronmk
    '''
707
    @param returning str|None An inserted column (such as pkey) to return
708 2070 aaronmk
    @param embeddable Whether the query should be embeddable as a nested SELECT.
709 2073 aaronmk
        Warning: If you set this and cacheable=True when the query is run, the
710
        query will be fully cached, not just if it raises an exception.
711 3009 aaronmk
    @param ignore Whether to ignore duplicate keys.
712 3181 aaronmk
    @param src Will be included in the name of any created function, to help
713
        identify the data source in pg_stat_activity.
714 1960 aaronmk
    '''
715 2754 aaronmk
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
716 2318 aaronmk
    if cols == []: cols = None # no cols (all defaults) = unknown col names
717 3010 aaronmk
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
718 2063 aaronmk
    if select_query == None: select_query = 'DEFAULT VALUES'
719 2327 aaronmk
    if returning != None: returning = sql_gen.as_Col(returning, table)
720 2063 aaronmk
721 2497 aaronmk
    first_line = 'INSERT INTO '+table.to_str(db)
722 2063 aaronmk
723 3009 aaronmk
    def mk_insert(select_query):
724
        query = first_line
725 3014 aaronmk
        if cols != None:
726
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
727 3009 aaronmk
        query += '\n'+select_query
728
729
        if returning != None:
730
            returning_name_col = sql_gen.to_name_only_col(returning)
731
            query += '\nRETURNING '+returning_name_col.to_str(db)
732
733
        return query
734 2063 aaronmk
735 3017 aaronmk
    return_type = 'unknown'
736
    if returning != None: return_type = returning.to_str(db)+'%TYPE'
737
738 3009 aaronmk
    lang = 'sql'
739
    if ignore:
740 3017 aaronmk
        # Always return something to set the correct rowcount
741
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
742
743 3009 aaronmk
        embeddable = True # must use function
744
        lang = 'plpgsql'
745 3010 aaronmk
746 3092 aaronmk
        if cols == None:
747
            row = [sql_gen.Col(sql_gen.all_cols, 'row')]
748
            row_vars = [sql_gen.Table('row')]
749
        else:
750
            row_vars = row = [sql_gen.Col(c.name, 'row') for c in cols]
751
752 3009 aaronmk
        query = '''\
753 3010 aaronmk
DECLARE
754 3014 aaronmk
    row '''+table.to_str(db)+'''%ROWTYPE;
755 3009 aaronmk
BEGIN
756 3019 aaronmk
    /* Need an EXCEPTION block for each individual row because "When an error is
757
    caught by an EXCEPTION clause, [...] all changes to persistent database
758
    state within the block are rolled back."
759
    This is unfortunate because "A block containing an EXCEPTION clause is
760
    significantly more expensive to enter and exit than a block without one."
761 3015 aaronmk
    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
762
#PLPGSQL-ERROR-TRAPPING)
763
    */
764 3092 aaronmk
    FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
765 3034 aaronmk
'''+select_query+'''
766
    LOOP
767 3015 aaronmk
        BEGIN
768 3019 aaronmk
            RETURN QUERY
769 3014 aaronmk
'''+mk_insert(sql_gen.Values(row).to_str(db))+'''
770 3010 aaronmk
;
771 3015 aaronmk
        EXCEPTION
772 3019 aaronmk
            WHEN unique_violation THEN NULL; -- continue to next row
773 3015 aaronmk
        END;
774 3010 aaronmk
    END LOOP;
775
END;\
776 3009 aaronmk
'''
777
    else: query = mk_insert(select_query)
778
779 2070 aaronmk
    if embeddable:
780
        # Create function
781 2513 aaronmk
        function_name = sql_gen.clean_name(first_line)
782 3181 aaronmk
        if src != None: function_name = src+': '+function_name
783 2189 aaronmk
        while True:
784
            try:
785 2918 aaronmk
                function = db.TempFunction(function_name)
786 2194 aaronmk
787 2189 aaronmk
                function_query = '''\
788 2698 aaronmk
CREATE FUNCTION '''+function.to_str(db)+'''()
789 3017 aaronmk
RETURNS SETOF '''+return_type+'''
790 3009 aaronmk
LANGUAGE '''+lang+'''
791 2467 aaronmk
AS $$
792 3009 aaronmk
'''+query+'''
793 2467 aaronmk
$$;
794 2070 aaronmk
'''
795 2446 aaronmk
                run_query(db, function_query, recover=True, cacheable=True,
796 2945 aaronmk
                    log_ignore_excs=(DuplicateException,))
797 2189 aaronmk
                break # this version was successful
798 2945 aaronmk
            except DuplicateException, e:
799 2189 aaronmk
                function_name = next_version(function_name)
800
                # try again with next version of name
801 2070 aaronmk
802 2337 aaronmk
        # Return query that uses function
803 3009 aaronmk
        cols = None
804
        if returning != None: cols = [returning]
805 2698 aaronmk
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
806 3009 aaronmk
            cols) # AS clause requires function alias
807 2787 aaronmk
        return mk_select(db, func_table, start=0, order_by=None)
808 2070 aaronmk
809 2787 aaronmk
    return query
810 2066 aaronmk
811 3074 aaronmk
def insert_select(db, table, *args, **kw_args):
812 2085 aaronmk
    '''For params, see mk_insert_select() and run_query_into()
813 2386 aaronmk
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
814
        values in
815 2072 aaronmk
    '''
816 3141 aaronmk
    returning = kw_args.get('returning', None)
817
    ignore = kw_args.get('ignore', False)
818
819 2386 aaronmk
    into = kw_args.pop('into', None)
820
    if into != None: kw_args['embeddable'] = True
821 2066 aaronmk
    recover = kw_args.pop('recover', None)
822 3141 aaronmk
    if ignore: recover = True
823 2066 aaronmk
    cacheable = kw_args.pop('cacheable', True)
824 2673 aaronmk
    log_level = kw_args.pop('log_level', 2)
825 2066 aaronmk
826 3141 aaronmk
    rowcount_only = ignore and returning == None # keep NULL rows on server
827
    if rowcount_only: into = sql_gen.Table('rowcount')
828
829 3074 aaronmk
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
830
        into, recover=recover, cacheable=cacheable, log_level=log_level)
831 3141 aaronmk
    if rowcount_only: empty_temp(db, into)
832 3074 aaronmk
    autoanalyze(db, table)
833
    return cur
834 2063 aaronmk
835 2738 aaronmk
default = sql_gen.default # tells insert() to use the default value for a column
836 2066 aaronmk
837 2063 aaronmk
def insert(db, table, row, *args, **kw_args):
838 2085 aaronmk
    '''For params, see insert_select()'''
839 1960 aaronmk
    if lists.is_seq(row): cols = None
840
    else:
841
        cols = row.keys()
842
        row = row.values()
843 2738 aaronmk
    row = list(row) # ensure that "== []" works
844 1960 aaronmk
845 2738 aaronmk
    if row == []: query = None
846
    else: query = sql_gen.Values(row).to_str(db)
847 1961 aaronmk
848 2788 aaronmk
    return insert_select(db, table, cols, query, *args, **kw_args)
849 11 aaronmk
850 3152 aaronmk
def mk_update(db, table, changes=None, cond=None, in_place=False,
851 3153 aaronmk
    cacheable_=True):
852 2402 aaronmk
    '''
853
    @param changes [(col, new_value),...]
854
        * container can be any iterable type
855
        * col: sql_gen.Code|str (for col name)
856
        * new_value: sql_gen.Code|literal value
857
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
858 3056 aaronmk
    @param in_place If set, locks the table and updates rows in place.
859
        This avoids creating dead rows in PostgreSQL.
860
        * cond must be None
861 3153 aaronmk
    @param cacheable_ Whether column structure information used to generate the
862 3152 aaronmk
        query can be cached
863 2402 aaronmk
    @return str query
864
    '''
865 3057 aaronmk
    table = sql_gen.as_Table(table)
866
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
867
        for c, v in changes]
868
869 3056 aaronmk
    if in_place:
870
        assert cond == None
871 3058 aaronmk
872 3065 aaronmk
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
873
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '
874 3153 aaronmk
            +db.col_info(sql_gen.with_default_table(c, table), cacheable_).type
875 3065 aaronmk
            +'\nUSING '+v.to_str(db) for c, v in changes))
876 3058 aaronmk
    else:
877
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
878
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
879
            for c, v in changes))
880
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
881 3056 aaronmk
882 2402 aaronmk
    return query
883
884 3074 aaronmk
def update(db, table, *args, **kw_args):
885 2402 aaronmk
    '''For params, see mk_update() and run_query()'''
886
    recover = kw_args.pop('recover', None)
887 3043 aaronmk
    cacheable = kw_args.pop('cacheable', False)
888 3030 aaronmk
    log_level = kw_args.pop('log_level', 2)
889 2402 aaronmk
890 3074 aaronmk
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
891
        cacheable, log_level=log_level)
892
    autoanalyze(db, table)
893
    return cur
894 2402 aaronmk
895 135 aaronmk
def last_insert_id(db):
896 1849 aaronmk
    module = util.root_module(db.db)
897 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
898
    elif module == 'MySQLdb': return db.insert_id()
899
    else: return None
900 13 aaronmk
901 2394 aaronmk
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
902 2383 aaronmk
    '''Creates a mapping from original column names (which may have collisions)
903 2415 aaronmk
    to names that will be distinct among the columns' tables.
904 2383 aaronmk
    This is meant to be used for several tables that are being joined together.
905 2415 aaronmk
    @param cols The columns to combine. Duplicates will be removed.
906
    @param into The table for the new columns.
907 2394 aaronmk
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
908
        columns will be included in the mapping even if they are not in cols.
909
        The tables of the provided Col objects will be changed to into, so make
910
        copies of them if you want to keep the original tables.
911
    @param as_items Whether to return a list of dict items instead of a dict
912 2383 aaronmk
    @return dict(orig_col=new_col, ...)
913
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
914 2392 aaronmk
        * new_col: sql_gen.Col(orig_col_name, into)
915
        * All mappings use the into table so its name can easily be
916 2383 aaronmk
          changed for all columns at once
917
    '''
918 2415 aaronmk
    cols = lists.uniqify(cols)
919
920 2394 aaronmk
    items = []
921 2389 aaronmk
    for col in preserve:
922 2390 aaronmk
        orig_col = copy.copy(col)
923 2392 aaronmk
        col.table = into
924 2394 aaronmk
        items.append((orig_col, col))
925
    preserve = set(preserve)
926
    for col in cols:
927 2716 aaronmk
        if col not in preserve:
928
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
929 2394 aaronmk
930
    if not as_items: items = dict(items)
931
    return items
932 2383 aaronmk
933 2393 aaronmk
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
934 2391 aaronmk
    '''For params, see mk_flatten_mapping()
935
    @return See return value of mk_flatten_mapping()
936
    '''
937 2394 aaronmk
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
938
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
939 2786 aaronmk
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
940 3166 aaronmk
        into=into)
941 2394 aaronmk
    return dict(items)
942 2391 aaronmk
943 3079 aaronmk
##### Database structure introspection
944 2414 aaronmk
945 3079 aaronmk
#### Tables
946
947
def tables(db, schema_like='public', table_like='%', exact=False):
948
    if exact: compare = '='
949
    else: compare = 'LIKE'
950
951
    module = util.root_module(db.db)
952
    if module == 'psycopg2':
953
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
954
            ('tablename', sql_gen.CompareCond(table_like, compare))]
955
        return values(select(db, 'pg_tables', ['tablename'], conds,
956
            order_by='tablename', log_level=4))
957
    elif module == 'MySQLdb':
958
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
959
            , cacheable=True, log_level=4))
960
    else: raise NotImplementedError("Can't list tables for "+module+' database')
961
962
def table_exists(db, table):
963
    table = sql_gen.as_Table(table)
964
    return list(tables(db, table.schema, table.name, exact=True)) != []
965
966 2426 aaronmk
def table_row_count(db, table, recover=None):
967 2786 aaronmk
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
968 2443 aaronmk
        order_by=None, start=0), recover=recover, log_level=3))
969 2426 aaronmk
970 2414 aaronmk
def table_cols(db, table, recover=None):
971
    return list(col_names(select(db, table, limit=0, order_by=None,
972 2443 aaronmk
        recover=recover, log_level=4)))
973 2414 aaronmk
974 2291 aaronmk
def pkey(db, table, recover=None):
975 832 aaronmk
    '''Assumed to be first column in table'''
976 2339 aaronmk
    return table_cols(db, table, recover)[0]
977 832 aaronmk
978 2559 aaronmk
not_null_col = 'not_null_col'
979 2340 aaronmk
980
def table_not_null_col(db, table, recover=None):
981
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
982
    if not_null_col in table_cols(db, table, recover): return not_null_col
983
    else: return pkey(db, table, recover)
984
985 853 aaronmk
def index_cols(db, table, index):
986
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
987
    automatically created. When you don't know whether something is a UNIQUE
988
    constraint or a UNIQUE index, use this function.'''
989 1909 aaronmk
    module = util.root_module(db.db)
990
    if module == 'psycopg2':
991
        return list(values(run_query(db, '''\
992 853 aaronmk
SELECT attname
993 866 aaronmk
FROM
994
(
995
        SELECT attnum, attname
996
        FROM pg_index
997
        JOIN pg_class index ON index.oid = indexrelid
998
        JOIN pg_class table_ ON table_.oid = indrelid
999
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
1000
        WHERE
1001 2782 aaronmk
            table_.relname = '''+db.esc_value(table)+'''
1002
            AND index.relname = '''+db.esc_value(index)+'''
1003 866 aaronmk
    UNION
1004
        SELECT attnum, attname
1005
        FROM
1006
        (
1007
            SELECT
1008
                indrelid
1009
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
1010
                    AS indkey
1011
            FROM pg_index
1012
            JOIN pg_class index ON index.oid = indexrelid
1013
            JOIN pg_class table_ ON table_.oid = indrelid
1014
            WHERE
1015 2782 aaronmk
                table_.relname = '''+db.esc_value(table)+'''
1016
                AND index.relname = '''+db.esc_value(index)+'''
1017 866 aaronmk
        ) s
1018
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
1019
) s
1020 853 aaronmk
ORDER BY attnum
1021 2782 aaronmk
'''
1022
            , cacheable=True, log_level=4)))
1023 1909 aaronmk
    else: raise NotImplementedError("Can't list index columns for "+module+
1024
        ' database')
1025 853 aaronmk
1026 464 aaronmk
def constraint_cols(db, table, constraint):
1027 1849 aaronmk
    module = util.root_module(db.db)
1028 464 aaronmk
    if module == 'psycopg2':
1029
        return list(values(run_query(db, '''\
1030
SELECT attname
1031
FROM pg_constraint
1032
JOIN pg_class ON pg_class.oid = conrelid
1033
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
1034
WHERE
1035 2783 aaronmk
    relname = '''+db.esc_value(table)+'''
1036
    AND conname = '''+db.esc_value(constraint)+'''
1037 464 aaronmk
ORDER BY attnum
1038 2783 aaronmk
'''
1039
            )))
1040 464 aaronmk
    else: raise NotImplementedError("Can't list constraint columns for "+module+
1041
        ' database')
1042
1043 3079 aaronmk
#### Functions
1044
1045
def function_exists(db, function):
1046
    function = sql_gen.as_Function(function)
1047
1048
    info_table = sql_gen.Table('routines', 'information_schema')
1049
    conds = [('routine_name', function.name)]
1050
    schema = function.schema
1051
    if schema != None: conds.append(('routine_schema', schema))
1052
    # Exclude trigger functions, since they cannot be called directly
1053
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1054
1055
    return list(values(select(db, info_table, ['routine_name'], conds,
1056
        order_by='routine_schema', limit=1, log_level=4))) != []
1057
        # TODO: order_by search_path schema order
1058
1059
##### Structural changes
1060
1061
#### Columns
1062
1063
def add_col(db, table, col, comment=None, **kw_args):
1064
    '''
1065
    @param col TypedCol Name may be versioned, so be sure to propagate any
1066
        renaming back to any source column for the TypedCol.
1067
    @param comment None|str SQL comment used to distinguish columns of the same
1068
        name from each other when they contain different data, to allow the
1069
        ADD COLUMN query to be cached. If not set, query will not be cached.
1070
    '''
1071
    assert isinstance(col, sql_gen.TypedCol)
1072
1073
    while True:
1074
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1075
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1076
1077
        try:
1078
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1079
            break
1080
        except DuplicateException:
1081
            col.name = next_version(col.name)
1082
            # try again with next version of name
1083
1084
def add_not_null(db, col):
1085
    table = col.table
1086
    col = sql_gen.to_name_only_col(col)
1087
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1088
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1089
1090 2096 aaronmk
row_num_col = '_row_num'
1091
1092 3079 aaronmk
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1093
    constraints='PRIMARY KEY')
1094
1095
def add_row_num(db, table):
1096
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1097
    be the primary key.'''
1098
    add_col(db, table, row_num_typed_col, log_level=3)
1099
1100
#### Indexes
1101
1102
def add_pkey(db, table, cols=None, recover=None):
1103
    '''Adds a primary key.
1104
    @param cols [sql_gen.Col,...] The columns in the primary key.
1105
        Defaults to the first column in the table.
1106
    @pre The table must not already have a primary key.
1107
    '''
1108
    table = sql_gen.as_Table(table)
1109
    if cols == None: cols = [pkey(db, table, recover)]
1110
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1111
1112
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1113
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1114
        log_ignore_excs=(DuplicateException,))
1115
1116 2998 aaronmk
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1117 2688 aaronmk
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1118 2538 aaronmk
    Currently, only function calls are supported as expressions.
1119 2998 aaronmk
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1120 2847 aaronmk
        This allows indexes to be used for comparisons where NULLs are equal.
1121 2538 aaronmk
    '''
1122 2964 aaronmk
    exprs = lists.mk_seq(exprs)
1123 2538 aaronmk
1124 2688 aaronmk
    # Parse exprs
1125
    old_exprs = exprs[:]
1126
    exprs = []
1127
    cols = []
1128
    for i, expr in enumerate(old_exprs):
1129 2823 aaronmk
        expr = sql_gen.as_Col(expr, table)
1130 2688 aaronmk
1131 2823 aaronmk
        # Handle nullable columns
1132 2998 aaronmk
        if ensure_not_null_:
1133 3164 aaronmk
            try: expr = sql_gen.ensure_not_null(db, expr)
1134 2860 aaronmk
            except KeyError: pass # unknown type, so just create plain index
1135 2823 aaronmk
1136 2688 aaronmk
        # Extract col
1137 3002 aaronmk
        expr = copy.deepcopy(expr) # don't modify input!
1138 2688 aaronmk
        if isinstance(expr, sql_gen.FunctionCall):
1139
            col = expr.args[0]
1140
            expr = sql_gen.Expr(expr)
1141
        else: col = expr
1142 2823 aaronmk
        assert isinstance(col, sql_gen.Col)
1143 2688 aaronmk
1144
        # Extract table
1145
        if table == None:
1146
            assert sql_gen.is_table_col(col)
1147
            table = col.table
1148
1149
        col.table = None
1150
1151
        exprs.append(expr)
1152
        cols.append(col)
1153 2408 aaronmk
1154 2688 aaronmk
    table = sql_gen.as_Table(table)
1155
1156 3005 aaronmk
    # Add index
1157 3148 aaronmk
    str_ = 'CREATE'
1158
    if unique: str_ += ' UNIQUE'
1159
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1160
        ', '.join((v.to_str(db) for v in exprs)))+')'
1161
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1162 2408 aaronmk
1163 3083 aaronmk
already_indexed = object() # tells add_indexes() the pkey has already been added
1164
1165
def add_indexes(db, table, has_pkey=True):
1166
    '''Adds an index on all columns in a table.
1167
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1168
        index should be added on the first column.
1169
        * If already_indexed, the pkey is assumed to have already been added
1170
    '''
1171
    cols = table_cols(db, table)
1172
    if has_pkey:
1173
        if has_pkey is not already_indexed: add_pkey(db, table)
1174
        cols = cols[1:]
1175
    for col in cols: add_index(db, col, table)
1176
1177 3079 aaronmk
#### Tables
1178 2772 aaronmk
1179 3079 aaronmk
### Maintenance
1180 2772 aaronmk
1181 3079 aaronmk
def analyze(db, table):
1182
    table = sql_gen.as_Table(table)
1183
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1184 2934 aaronmk
1185 3079 aaronmk
def autoanalyze(db, table):
1186
    if db.autoanalyze: analyze(db, table)
1187 2935 aaronmk
1188 3079 aaronmk
def vacuum(db, table):
1189
    table = sql_gen.as_Table(table)
1190
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1191
        log_level=3))
1192 2086 aaronmk
1193 3079 aaronmk
### Lifecycle
1194
1195 3247 aaronmk
def drop(db, type_, name):
1196
    name = sql_gen.as_Name(name)
1197
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1198 2889 aaronmk
1199 3247 aaronmk
def drop_table(db, table): drop(db, 'TABLE', table)
1200
1201 3082 aaronmk
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1202
    like=None):
1203 2675 aaronmk
    '''Creates a table.
1204 2681 aaronmk
    @param cols [sql_gen.TypedCol,...] The column names and types
1205
    @param has_pkey If set, the first column becomes the primary key.
1206 2760 aaronmk
    @param col_indexes bool|[ref]
1207
        * If True, indexes will be added on all non-pkey columns.
1208
        * If a list reference, [0] will be set to a function to do this.
1209
          This can be used to delay index creation until the table is populated.
1210 2675 aaronmk
    '''
1211
    table = sql_gen.as_Table(table)
1212
1213 3082 aaronmk
    if like != None:
1214
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1215
            ]+cols
1216 2681 aaronmk
    if has_pkey:
1217
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1218 2872 aaronmk
        pkey.constraints = 'PRIMARY KEY'
1219 2681 aaronmk
1220 3085 aaronmk
    temp = table.is_temp and not db.debug_temp
1221
        # temp tables permanent in debug_temp mode
1222 2760 aaronmk
1223 3085 aaronmk
    # Create table
1224
    while True:
1225
        str_ = 'CREATE'
1226
        if temp: str_ += ' TEMP'
1227
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1228
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1229 3126 aaronmk
        str_ += '\n);'
1230 3085 aaronmk
1231
        try:
1232 3127 aaronmk
            run_query(db, str_, recover=True, cacheable=True, log_level=2,
1233 3085 aaronmk
                log_ignore_excs=(DuplicateException,))
1234
            break
1235
        except DuplicateException:
1236
            table.name = next_version(table.name)
1237
            # try again with next version of name
1238
1239 2760 aaronmk
    # Add indexes
1240 2773 aaronmk
    if has_pkey: has_pkey = already_indexed
1241
    def add_indexes_(): add_indexes(db, table, has_pkey)
1242
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1243
    elif col_indexes: add_indexes_() # add now
1244 2675 aaronmk
1245 3084 aaronmk
def copy_table_struct(db, src, dest):
1246
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1247 3085 aaronmk
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1248 3084 aaronmk
1249 3079 aaronmk
### Data
1250 2684 aaronmk
1251 2970 aaronmk
def truncate(db, table, schema='public', **kw_args):
1252
    '''For params, see run_query()'''
1253 2777 aaronmk
    table = sql_gen.as_Table(table, schema)
1254 2970 aaronmk
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1255 2732 aaronmk
1256 2965 aaronmk
def empty_temp(db, tables):
1257
    tables = lists.mk_seq(tables)
1258 2971 aaronmk
    for table in tables: truncate(db, table, log_level=3)
1259 2965 aaronmk
1260 1968 aaronmk
def empty_db(db, schema='public', **kw_args):
1261
    '''For kw_args, see tables()'''
1262
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1263 3094 aaronmk
1264
def distinct_table(db, table, distinct_on):
1265
    '''Creates a copy of a temp table which is distinct on the given columns.
1266 3099 aaronmk
    The old and new tables will both get an index on these columns, to
1267
    facilitate merge joins.
1268 3097 aaronmk
    @param distinct_on If empty, creates a table with one row. This is useful if
1269
        your distinct_on columns are all literal values.
1270 3099 aaronmk
    @return The new table.
1271 3094 aaronmk
    '''
1272 3099 aaronmk
    new_table = sql_gen.suffixed_table(table, '_distinct')
1273 3094 aaronmk
1274 3099 aaronmk
    copy_table_struct(db, table, new_table)
1275 3097 aaronmk
1276
    limit = None
1277
    if distinct_on == []: limit = 1 # one sample row
1278 3099 aaronmk
    else:
1279
        add_index(db, distinct_on, new_table, unique=True)
1280
        add_index(db, distinct_on, table) # for join optimization
1281 3097 aaronmk
1282 3099 aaronmk
    insert_select(db, new_table, None, mk_select(db, table, start=0,
1283 3097 aaronmk
        limit=limit), ignore=True)
1284 3099 aaronmk
    analyze(db, new_table)
1285 3094 aaronmk
1286 3099 aaronmk
    return new_table