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