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