Project

General

Profile

1
# Database access
2

    
3
import copy
4
import re
5
import time
6
import warnings
7

    
8
import exc
9
import dicts
10
import iters
11
import lists
12
import profiling
13
from Proxy import Proxy
14
import rand
15
import sql_gen
16
import strings
17
import util
18

    
19
##### Exceptions
20

    
21
def get_cur_query(cur, input_query=None):
22
    raw_query = None
23
    if hasattr(cur, 'query'): raw_query = cur.query
24
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
25
    
26
    if raw_query != None: return raw_query
27
    else: return '[input] '+strings.ustr(input_query)
28

    
29
def _add_cursor_info(e, *args, **kw_args):
30
    '''For params, see get_cur_query()'''
31
    exc.add_msg(e, 'query: '+strings.ustr(get_cur_query(*args, **kw_args)))
32

    
33
class DbException(exc.ExceptionWithCause):
34
    def __init__(self, msg, cause=None, cur=None):
35
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
36
        if cur != None: _add_cursor_info(self, cur)
37

    
38
class ExceptionWithName(DbException):
39
    def __init__(self, name, cause=None):
40
        DbException.__init__(self, 'for name: '
41
            +strings.as_tt(strings.ustr(name)), cause)
42
        self.name = name
43

    
44
class ExceptionWithValue(DbException):
45
    def __init__(self, value, cause=None):
46
        DbException.__init__(self, 'for value: '
47
            +strings.as_tt(strings.urepr(value)), cause)
48
        self.value = value
49

    
50
class ExceptionWithNameType(DbException):
51
    def __init__(self, type_, name, cause=None):
52
        DbException.__init__(self, 'for type: '+strings.as_tt(strings.ustr(
53
            type_))+'; name: '+strings.as_tt(name), cause)
54
        self.type = type_
55
        self.name = name
56

    
57
class ConstraintException(DbException):
58
    def __init__(self, name, cond, cols, cause=None):
59
        msg = 'Violated '+strings.as_tt(name)+' constraint'
60
        if cond != None: msg += ' with condition '+strings.as_tt(cond)
61
        if cols != []: msg += ' on columns: '+strings.as_tt(', '.join(cols))
62
        DbException.__init__(self, msg, cause)
63
        self.name = name
64
        self.cond = cond
65
        self.cols = cols
66

    
67
class MissingCastException(DbException):
68
    def __init__(self, type_, col=None, cause=None):
69
        msg = 'Missing cast to type '+strings.as_tt(type_)
70
        if col != None: msg += ' on column: '+strings.as_tt(col)
71
        DbException.__init__(self, msg, cause)
72
        self.type = type_
73
        self.col = col
74

    
75
class EncodingException(ExceptionWithName): pass
76

    
77
class DuplicateKeyException(ConstraintException): pass
78

    
79
class NullValueException(ConstraintException): pass
80

    
81
class CheckException(ConstraintException): pass
82

    
83
class InvalidValueException(ExceptionWithValue): pass
84

    
85
class InvalidTypeException(ExceptionWithNameType): pass
86

    
87
class DuplicateException(ExceptionWithNameType): pass
88

    
89
class DoesNotExistException(ExceptionWithNameType): pass
90

    
91
class EmptyRowException(DbException): pass
92

    
93
##### Warnings
94

    
95
class DbWarning(UserWarning): pass
96

    
97
##### Result retrieval
98

    
99
def col_names(cur): return (col[0] for col in cur.description)
100

    
101
def rows(cur): return iter(lambda: cur.fetchone(), None)
102

    
103
def consume_rows(cur):
104
    '''Used to fetch all rows so result will be cached'''
105
    iters.consume_iter(rows(cur))
106

    
107
def next_row(cur): return rows(cur).next()
108

    
109
def row(cur):
110
    row_ = next_row(cur)
111
    consume_rows(cur)
112
    return row_
113

    
114
def next_value(cur): return next_row(cur)[0]
115

    
116
def value(cur): return row(cur)[0]
117

    
118
def values(cur): return iters.func_iter(lambda: next_value(cur))
119

    
120
def value_or_none(cur):
121
    try: return value(cur)
122
    except StopIteration: return None
123

    
124
##### Escaping
125

    
126
def esc_name_by_module(module, name):
127
    if module == 'psycopg2' or module == None: quote = '"'
128
    elif module == 'MySQLdb': quote = '`'
129
    else: raise NotImplementedError("Can't escape name for "+module+' database')
130
    return sql_gen.esc_name(name, quote)
131

    
132
def esc_name_by_engine(engine, name, **kw_args):
133
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
134

    
135
def esc_name(db, name, **kw_args):
136
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
137

    
138
def qual_name(db, schema, table):
139
    def esc_name_(name): return esc_name(db, name)
140
    table = esc_name_(table)
141
    if schema != None: return esc_name_(schema)+'.'+table
142
    else: return table
143

    
144
##### Database connections
145

    
146
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
147

    
148
db_engines = {
149
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
150
    'PostgreSQL': ('psycopg2', {}),
151
}
152

    
153
DatabaseErrors_set = set([DbException])
154
DatabaseErrors = tuple(DatabaseErrors_set)
155

    
156
def _add_module(module):
157
    DatabaseErrors_set.add(module.DatabaseError)
158
    global DatabaseErrors
159
    DatabaseErrors = tuple(DatabaseErrors_set)
160

    
161
def db_config_str(db_config):
162
    return db_config['engine']+' database '+db_config['database']
163

    
164
log_debug_none = lambda msg, level=2: None
165

    
166
class DbConn:
167
    def __init__(self, db_config, autocommit=True, caching=True,
168
        log_debug=log_debug_none, debug_temp=False, src=None):
169
        '''
170
        @param debug_temp Whether temporary objects should instead be permanent.
171
            This assists in debugging the internal objects used by the program.
172
        @param src In autocommit mode, will be included in a comment in every
173
            query, to help identify the data source in pg_stat_activity.
174
        '''
175
        self.db_config = db_config
176
        self.autocommit = autocommit
177
        self.caching = caching
178
        self.log_debug = log_debug
179
        self.debug = log_debug != log_debug_none
180
        self.debug_temp = debug_temp
181
        self.src = src
182
        self.autoanalyze = False
183
        self.autoexplain = False
184
        self.profile_row_ct = None
185
        
186
        self._savepoint = 0
187
        self._reset()
188
    
189
    def __getattr__(self, name):
190
        if name == '__dict__': raise Exception('getting __dict__')
191
        if name == 'db': return self._db()
192
        else: raise AttributeError()
193
    
194
    def __getstate__(self):
195
        state = copy.copy(self.__dict__) # shallow copy
196
        state['log_debug'] = None # don't pickle the debug callback
197
        state['_DbConn__db'] = None # don't pickle the connection
198
        return state
199
    
200
    def clear_cache(self): self.query_results = {}
201
    
202
    def _reset(self):
203
        self.clear_cache()
204
        assert self._savepoint == 0
205
        self._notices_seen = set()
206
        self.__db = None
207
    
208
    def connected(self): return self.__db != None
209
    
210
    def close(self):
211
        if not self.connected(): return
212
        
213
        # Record that the automatic transaction is now closed
214
        self._savepoint -= 1
215
        
216
        self.db.close()
217
        self._reset()
218
    
219
    def reconnect(self):
220
        # Do not do this in test mode as it would roll back everything
221
        if self.autocommit: self.close()
222
        # Connection will be reopened automatically on first query
223
    
224
    def _db(self):
225
        if self.__db == None:
226
            # Process db_config
227
            db_config = self.db_config.copy() # don't modify input!
228
            schemas = db_config.pop('schemas', None)
229
            module_name, mappings = db_engines[db_config.pop('engine')]
230
            module = __import__(module_name)
231
            _add_module(module)
232
            for orig, new in mappings.iteritems():
233
                try: util.rename_key(db_config, orig, new)
234
                except KeyError: pass
235
            
236
            # Connect
237
            self.__db = module.connect(**db_config)
238
            
239
            # Record that a transaction is already open
240
            self._savepoint += 1
241
            
242
            # Configure connection
243
            if hasattr(self.db, 'set_isolation_level'):
244
                import psycopg2.extensions
245
                self.db.set_isolation_level(
246
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
247
            if schemas != None:
248
                search_path = [self.esc_name(s) for s in schemas.split(',')]
249
                run_query(self, 'SET search_path TO '+(','.join(search_path)),
250
                    log_level=3)
251
        
252
        return self.__db
253
    
254
    class DbCursor(Proxy):
255
        def __init__(self, outer):
256
            Proxy.__init__(self, outer.db.cursor())
257
            self.outer = outer
258
            self.query_results = outer.query_results
259
            self.query_lookup = None
260
            self.result = []
261
        
262
        def execute(self, query):
263
            self._is_insert = query.startswith('INSERT')
264
            self.query_lookup = query
265
            try:
266
                try: cur = self.inner.execute(query)
267
                finally: self.query = get_cur_query(self.inner, query)
268
            except Exception, e:
269
                self.result = e # cache the exception as the result
270
                self._cache_result()
271
                raise
272
            
273
            # Always cache certain queries
274
            query = sql_gen.lstrip(query)
275
            if query.startswith('CREATE') or query.startswith('ALTER'):
276
                # structural changes
277
                # Rest of query must be unique in the face of name collisions,
278
                # so don't cache ADD COLUMN unless it has distinguishing comment
279
                if query.find('ADD COLUMN') < 0 or query.endswith('*/'):
280
                    self._cache_result()
281
            elif self.rowcount == 0 and query.startswith('SELECT'): # empty
282
                consume_rows(self) # fetch all rows so result will be cached
283
            
284
            return cur
285
        
286
        def fetchone(self):
287
            row = self.inner.fetchone()
288
            if row != None: self.result.append(row)
289
            # otherwise, fetched all rows
290
            else: self._cache_result()
291
            return row
292
        
293
        def _cache_result(self):
294
            # For inserts that return a result set, don't cache result set since
295
            # inserts are not idempotent. Other non-SELECT queries don't have
296
            # their result set read, so only exceptions will be cached (an
297
            # invalid query will always be invalid).
298
            if self.query_results != None and (not self._is_insert
299
                or isinstance(self.result, Exception)):
300
                
301
                assert self.query_lookup != None
302
                self.query_results[self.query_lookup] = self.CacheCursor(
303
                    util.dict_subset(dicts.AttrsDictView(self),
304
                    ['query', 'result', 'rowcount', 'description']))
305
        
306
        class CacheCursor:
307
            def __init__(self, cached_result): self.__dict__ = cached_result
308
            
309
            def execute(self, *args, **kw_args):
310
                if isinstance(self.result, Exception): raise self.result
311
                # otherwise, result is a rows list
312
                self.iter = iter(self.result)
313
            
314
            def fetchone(self):
315
                try: return self.iter.next()
316
                except StopIteration: return None
317
    
318
    def esc_value(self, value):
319
        try: str_ = self.mogrify('%s', [value])
320
        except NotImplementedError, e:
321
            module = util.root_module(self.db)
322
            if module == 'MySQLdb':
323
                import _mysql
324
                str_ = _mysql.escape_string(value)
325
            else: raise e
326
        return strings.to_unicode(str_)
327
    
328
    def esc_name(self, name): return esc_name(self, name) # calls global func
329
    
330
    def std_code(self, str_):
331
        '''Standardizes SQL code.
332
        * Ensures that string literals are prefixed by `E`
333
        '''
334
        if str_.startswith("'"): str_ = 'E'+str_
335
        return str_
336
    
337
    def can_mogrify(self):
338
        module = util.root_module(self.db)
339
        return module == 'psycopg2'
340
    
341
    def mogrify(self, query, params=None):
342
        if self.can_mogrify(): return self.db.cursor().mogrify(query, params)
343
        else: raise NotImplementedError("Can't mogrify query")
344
    
345
    def set_encoding(self, encoding):
346
        encoding_str = sql_gen.Literal(encoding)
347
        run_query(self, 'SET NAMES '+encoding_str.to_str(self))
348
    
349
    def print_notices(self):
350
        if hasattr(self.db, 'notices'):
351
            for msg in self.db.notices:
352
                if msg not in self._notices_seen:
353
                    self._notices_seen.add(msg)
354
                    self.log_debug(msg, level=2)
355
    
356
    def run_query(self, query, cacheable=False, log_level=2,
357
        debug_msg_ref=None):
358
        '''
359
        @param log_ignore_excs The log_level will be increased by 2 if the query
360
            throws one of these exceptions.
361
        @param debug_msg_ref If specified, the log message will be returned in
362
            this instead of being output. This allows you to filter log messages
363
            depending on the result of the query.
364
        '''
365
        assert query != None
366
        
367
        if self.autocommit and self.src != None:
368
            query = sql_gen.esc_comment(self.src)+'\t'+query
369
        
370
        if not self.caching: cacheable = False
371
        used_cache = False
372
        
373
        if self.debug:
374
            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
375
        try:
376
            # Get cursor
377
            if cacheable:
378
                try: cur = self.query_results[query]
379
                except KeyError: cur = self.DbCursor(self)
380
                else: used_cache = True
381
            else: cur = self.db.cursor()
382
            
383
            # Run query
384
            try: cur.execute(query)
385
            except Exception, e:
386
                _add_cursor_info(e, self, query)
387
                raise
388
            else: self.do_autocommit()
389
        finally:
390
            if self.debug:
391
                profiler.stop(self.profile_row_ct)
392
                
393
                ## Log or return query
394
                
395
                query = strings.ustr(get_cur_query(cur, query))
396
                # Put the src comment on a separate line in the log file
397
                query = query.replace('\t', '\n', 1)
398
                
399
                msg = 'DB query: '
400
                
401
                if used_cache: msg += 'cache hit'
402
                elif cacheable: msg += 'cache miss'
403
                else: msg += 'non-cacheable'
404
                
405
                msg += ':\n'+profiler.msg()+'\n'+strings.as_code(query, 'SQL')
406
                
407
                if debug_msg_ref != None: debug_msg_ref[0] = msg
408
                else: self.log_debug(msg, log_level)
409
                
410
                self.print_notices()
411
        
412
        return cur
413
    
414
    def is_cached(self, query): return query in self.query_results
415
    
416
    def with_autocommit(self, func):
417
        import psycopg2.extensions
418
        
419
        prev_isolation_level = self.db.isolation_level
420
        self.db.set_isolation_level(
421
            psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
422
        try: return func()
423
        finally: self.db.set_isolation_level(prev_isolation_level)
424
    
425
    def with_savepoint(self, func):
426
        top = self._savepoint == 0
427
        savepoint = 'level_'+str(self._savepoint)
428
        
429
        if self.debug:
430
            self.log_debug('Begin transaction', level=4)
431
            profiler = profiling.ItersProfiler(start_now=True, iter_text='row')
432
        
433
        # Must happen before running queries so they don't get autocommitted
434
        self._savepoint += 1
435
        
436
        if top: query = 'START TRANSACTION ISOLATION LEVEL READ COMMITTED'
437
        else: query = 'SAVEPOINT '+savepoint
438
        self.run_query(query, log_level=4)
439
        try:
440
            return func()
441
            if top: self.run_query('COMMIT', log_level=4)
442
        except:
443
            if top: query = 'ROLLBACK'
444
            else: query = 'ROLLBACK TO SAVEPOINT '+savepoint
445
            self.run_query(query, log_level=4)
446
            
447
            raise
448
        finally:
449
            # Always release savepoint, because after ROLLBACK TO SAVEPOINT,
450
            # "The savepoint remains valid and can be rolled back to again"
451
            # (http://www.postgresql.org/docs/8.3/static/sql-rollback-to.html).
452
            if not top:
453
                self.run_query('RELEASE SAVEPOINT '+savepoint, log_level=4)
454
            
455
            self._savepoint -= 1
456
            assert self._savepoint >= 0
457
            
458
            if self.debug:
459
                profiler.stop(self.profile_row_ct)
460
                self.log_debug('End transaction\n'+profiler.msg(), level=4)
461
            
462
            self.do_autocommit() # OK to do this after ROLLBACK TO SAVEPOINT
463
    
464
    def do_autocommit(self):
465
        '''Autocommits if outside savepoint'''
466
        assert self._savepoint >= 1
467
        if self.autocommit and self._savepoint == 1:
468
            self.log_debug('Autocommitting', level=4)
469
            self.db.commit()
470
    
471
    def col_info(self, col, cacheable=True):
472
        module = util.root_module(self.db)
473
        if module == 'psycopg2':
474
            qual_table = sql_gen.Literal(col.table.to_str(self))
475
            col_name_str = sql_gen.Literal(col.name)
476
            try:
477
                type_, is_array, default, nullable = row(run_query(self, '''\
478
SELECT
479
format_type(COALESCE(NULLIF(typelem, 0), pg_type.oid), -1) AS type
480
, typcategory = 'A' AS type_is_array
481
, pg_get_expr(pg_attrdef.adbin, attrelid, true) AS default
482
, NOT pg_attribute.attnotnull AS nullable
483
FROM pg_attribute
484
LEFT JOIN pg_type ON pg_type.oid = atttypid
485
LEFT JOIN pg_attrdef ON adrelid = attrelid AND adnum = attnum
486
WHERE
487
attrelid = '''+qual_table.to_str(self)+'''::regclass
488
AND attname = '''+col_name_str.to_str(self)+'''
489
'''
490
                    , recover=True, cacheable=cacheable, log_level=4))
491
            except (DoesNotExistException, StopIteration):
492
                raise sql_gen.NoUnderlyingTableException(col)
493
            if is_array: type_ = sql_gen.ArrayType(type_)
494
        else:
495
            table = sql_gen.Table('columns', 'information_schema')
496
            cols = [sql_gen.Col('data_type'), sql_gen.Col('udt_name'),
497
                'column_default', sql_gen.Cast('boolean',
498
                sql_gen.Col('is_nullable'))]
499
            
500
            conds = [('table_name', col.table.name),
501
                ('column_name', strings.ustr(col.name))]
502
            schema = col.table.schema
503
            if schema != None: conds.append(('table_schema', schema))
504
            
505
            cur = select(self, table, cols, conds, order_by='table_schema',
506
                limit=1, cacheable=cacheable, log_level=4)
507
            try: type_, extra_type, default, nullable = row(cur)
508
            except StopIteration: raise sql_gen.NoUnderlyingTableException(col)
509
            if type_ == 'USER-DEFINED': type_ = extra_type
510
            elif type_ == 'ARRAY':
511
                type_ = sql_gen.ArrayType(strings.remove_prefix('_', extra_type,
512
                    require=True))
513
        
514
        if default != None: default = sql_gen.as_Code(default, self)
515
        return sql_gen.TypedCol(col.name, type_, default, nullable)
516
    
517
    def TempFunction(self, name):
518
        if self.debug_temp: schema = None
519
        else: schema = 'pg_temp'
520
        return sql_gen.Function(name, schema)
521

    
522
connect = DbConn
523

    
524
##### Recoverable querying
525

    
526
def parse_exception(db, e, recover=False):
527
    msg = strings.ustr(e.args[0])
528
    msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)
529
    
530
    match = re.match(r'^invalid byte sequence for encoding "(.+?)":', msg)
531
    if match:
532
        encoding, = match.groups()
533
        raise EncodingException(encoding, e)
534
    
535
    def make_DuplicateKeyException(constraint, e):
536
        cols = []
537
        cond = None
538
        if recover: # need auto-rollback to run index_cols()
539
            try:
540
                cols = index_cols(db, constraint)
541
                cond = index_cond(db, constraint)
542
            except NotImplementedError: pass
543
            except DatabaseErrors, e2: # e.g. "current transaction is aborted"
544
                exc.print_ex(e2)
545
                raise e
546
        return DuplicateKeyException(constraint, cond, cols, e)
547
    
548
    match = re.match(r'^duplicate key value violates unique constraint "(.+?)"',
549
        msg)
550
    if match:
551
        constraint, = match.groups()
552
        raise make_DuplicateKeyException(constraint, e)
553
    
554
    match = re.match(r'^could not create unique index "(.+?)"\n'
555
        r'DETAIL:  Key .+? is duplicated', msg)
556
    if match:
557
        constraint, = match.groups()
558
        raise DuplicateKeyException(constraint, None, [], e)
559
    
560
    match = re.match(r'^null value in column "(.+?)" violates not-null'
561
        r' constraint', msg)
562
    if match:
563
        col, = match.groups()
564
        raise NullValueException('NOT NULL', None, [col], e)
565
    
566
    match = re.match(r'^new row for relation "(.+?)" violates check '
567
        r'constraint "(.+?)"', msg)
568
    if match:
569
        table, constraint = match.groups()
570
        constraint = sql_gen.Col(constraint, table)
571
        cond = None
572
        if recover: # need auto-rollback to run constraint_cond()
573
            try: cond = constraint_cond(db, constraint)
574
            except NotImplementedError: pass
575
        raise CheckException(constraint.to_str(db), cond, [], e)
576
    
577
    match = re.match(r'^(?:invalid input (?:syntax|value)\b[^:]*'
578
        r'|.+? out of range)(?:: "(.+?)")?', msg)
579
    if match:
580
        value, = match.groups()
581
        value = util.do_ignore_none(strings.to_unicode, value)
582
        raise InvalidValueException(value, e)
583
    
584
    match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
585
        r'is of type', msg)
586
    if match:
587
        col, type_ = match.groups()
588
        raise MissingCastException(type_, col, e)
589
    
590
    match = re.match(r'^could not determine polymorphic type because '
591
        r'input has type "unknown"', msg)
592
    if match: raise MissingCastException('text', None, e)
593
    
594
    match = re.match(r'^.+? types (.+?) and (.+?) cannot be matched', msg)
595
    if match:
596
        type0, type1 = match.groups()
597
        raise MissingCastException(type0, None, e)
598
    
599
    match = re.match(r'^.*?\brelation "(.+?)" is not a (table)', msg)
600
    if match:
601
        name, type_ = match.groups()
602
        raise InvalidTypeException(type_, name, e)
603
    
604
    typed_name_re = r'^(\S+) ("?)(.+?)\2(?: of relation ".+?")?'
605
        # regexp must be followed with text for .*? to match properly
606
    
607
    match = re.match(typed_name_re+r'(?:\(.*?)? already exists', msg)
608
    if match:
609
        type_, quote, name = match.groups()
610
        raise DuplicateException(type_, name, e)
611
    
612
    match = re.match(r'more than one (\S+) named ""(.+?)""', msg)
613
    if match:
614
        type_, name = match.groups()
615
        raise DuplicateException(type_, name, e)
616
    
617
    match = re.match(typed_name_re+r' does not exist', msg)
618
    if match:
619
        type_, quote, name = match.groups()
620
        if type_ == 'function':
621
            match = re.match(r'^(.+?)\(.*\)$', name)
622
            if match: # includes params, so is call rather than cast to regproc
623
                function_name, = match.groups()
624
                func = sql_gen.Function(function_name)
625
                if function_exists(db, func) and msg.find('CAST') < 0:
626
                    # not found only because of a missing cast
627
                    type_ = function_param0_type(db, func)
628
                    col = None
629
                    if type_ == 'anyelement': type_ = 'text'
630
                    elif type_ == 'hstore': # cast just the value param
631
                        type_ = 'text'
632
                        col = 'value'
633
                    raise MissingCastException(type_, col, e)
634
        raise DoesNotExistException(type_, name, e)
635
    
636
    raise # no specific exception raised
637

    
638
def with_savepoint(db, func): return db.with_savepoint(func)
639

    
640
def run_query(db, query, recover=None, cacheable=False, log_level=2,
641
    log_ignore_excs=None, **kw_args):
642
    '''For params, see DbConn.run_query()'''
643
    if recover == None: recover = False
644
    if log_ignore_excs == None: log_ignore_excs = ()
645
    log_ignore_excs = tuple(log_ignore_excs)
646
    debug_msg_ref = [None]
647
    
648
    query = with_explain_comment(db, query)
649
    
650
    try:
651
        try:
652
            def run(): return db.run_query(query, cacheable, log_level,
653
                debug_msg_ref, **kw_args)
654
            if recover and not db.is_cached(query):
655
                return with_savepoint(db, run)
656
            else: return run() # don't need savepoint if cached
657
        except Exception, e:
658
            # Give failed EXPLAIN approximately the log_level of its query
659
            if query.startswith('EXPLAIN'): log_level -= 1
660
            
661
            parse_exception(db, e, recover)
662
    except log_ignore_excs:
663
        log_level += 2
664
        raise
665
    finally:
666
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
667

    
668
##### Basic queries
669

    
670
def is_explainable(query):
671
    # See <http://www.postgresql.org/docs/8.3/static/sql-explain.html#AEN57749>
672
    return re.match(r'^(?:SELECT|INSERT|UPDATE|DELETE|VALUES|EXECUTE|DECLARE)\b'
673
        , query)
674

    
675
def explain(db, query, **kw_args):
676
    '''
677
    For params, see run_query().
678
    '''
679
    kw_args.setdefault('log_level', 4)
680
    
681
    return strings.ustr(strings.join_lines(values(run_query(db,
682
        'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))
683
        # not a higher log_level because it's useful to see what query is being
684
        # run before it's executed, which EXPLAIN effectively provides
685

    
686
def has_comment(query): return query.endswith('*/')
687

    
688
def with_explain_comment(db, query, **kw_args):
689
    if db.autoexplain and not has_comment(query) and is_explainable(query):
690
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
691
            +explain(db, query, **kw_args))
692
    return query
693

    
694
def next_version(name):
695
    version = 1 # first existing name was version 0
696
    match = re.match(r'^(.*)#(\d+)$', name)
697
    if match:
698
        name, version = match.groups()
699
        version = int(version)+1
700
    return sql_gen.concat(name, '#'+str(version))
701

    
702
def lock_table(db, table, mode):
703
    table = sql_gen.as_Table(table)
704
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
705

    
706
def run_query_into(db, query, into=None, add_pkey_=False, **kw_args):
707
    '''Outputs a query to a temp table.
708
    For params, see run_query().
709
    '''
710
    if into == None: return run_query(db, query, **kw_args)
711
    
712
    assert isinstance(into, sql_gen.Table)
713
    
714
    into.is_temp = True
715
    # "temporary tables cannot specify a schema name", so remove schema
716
    into.schema = None
717
    
718
    kw_args['recover'] = True
719
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
720
    
721
    temp = not db.debug_temp # tables are permanent in debug_temp mode
722
    
723
    # Create table
724
    while True:
725
        create_query = 'CREATE'
726
        if temp: create_query += ' TEMP'
727
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
728
        
729
        try:
730
            cur = run_query(db, create_query, **kw_args)
731
                # CREATE TABLE AS sets rowcount to # rows in query
732
            break
733
        except DuplicateException, e:
734
            into.name = next_version(into.name)
735
            # try again with next version of name
736
    
737
    if add_pkey_: add_pkey_or_index(db, into, warn=True)
738
    
739
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
740
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
741
    # table is going to be used in complex queries, it is wise to run ANALYZE on
742
    # the temporary table after it is populated."
743
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
744
    # If into is not a temp table, ANALYZE is useful but not required.
745
    analyze(db, into)
746
    
747
    return cur
748

    
749
order_by_pkey = object() # tells mk_select() to order by the pkey
750

    
751
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
752

    
753
def has_subset_func(db, table):
754
    return sql_gen.is_underlying_table(table) and function_exists(db, table)
755

    
756
def mk_select(db, tables=None, fields=None, conds=None, distinct_on=[],
757
    limit=None, start=None, order_by=order_by_pkey, default_table=None,
758
    explain=True):
759
    '''
760
    @param tables The single table to select from, or a list of tables to join
761
        together, with tables after the first being sql_gen.Join objects
762
    @param fields Use None to select all fields in the table
763
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
764
        * container can be any iterable type
765
        * compare_left_side: sql_gen.Code|str (for col name)
766
        * compare_right_side: sql_gen.ValueCond|literal value
767
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
768
        use all columns
769
    @return query
770
    '''
771
    # Parse tables param
772
    tables = lists.mk_seq(tables)
773
    tables = list(tables) # don't modify input! (list() copies input)
774
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
775
    
776
    # Parse other params
777
    if conds == None: conds = []
778
    elif dicts.is_dict(conds): conds = conds.items()
779
    conds = list(conds) # don't modify input! (list() copies input)
780
    assert limit == None or isinstance(limit, (int, long))
781
    assert start == None or isinstance(start, (int, long))
782
    if limit == 0: order_by = None
783
    if order_by in (None, order_by_pkey) and has_subset_func(db, table0):
784
        # can use subset function for fast querying at large OFFSET values
785
        run_query(db, 'SET LOCAL enable_sort TO off')
786
        table0 = sql_gen.FunctionCall(table0, limit_=limit, offset_=start)
787
        if limit != 0: limit = None # done by function
788
        start = None # done by function
789
        order_by = None # done by function
790
    elif order_by is order_by_pkey:
791
        if lists.is_seq(distinct_on) and distinct_on: order_by = distinct_on[0]
792
        elif table0 != None: order_by = table_order_by(db, table0, recover=True)
793
        else: order_by = None
794
    
795
    query = 'SELECT'
796
    
797
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
798
    
799
    # DISTINCT ON columns
800
    if distinct_on != []:
801
        query += '\nDISTINCT'
802
        if distinct_on is not distinct_on_all:
803
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
804
    
805
    # Columns
806
    if query.find('\n') >= 0: whitespace = '\n'
807
    else: whitespace = ' '
808
    if fields == None: query += whitespace+'*'
809
    else:
810
        assert fields != []
811
        if len(fields) > 1: whitespace = '\n'
812
        query += whitespace+('\n, '.join(map(parse_col, fields)))
813
    
814
    # Main table
815
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
816
    else: whitespace = ' '
817
    if table0 != None: query += whitespace+'FROM '+table0.to_str(db)
818
    
819
    # Add joins
820
    left_table = table0
821
    for join_ in tables:
822
        table = join_.table
823
        
824
        # Parse special values
825
        if join_.type_ is sql_gen.filter_out: # filter no match
826
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
827
                sql_gen.CompareCond(None, '~=')))
828
        
829
        query += '\n'+join_.to_str(db, left_table)
830
        
831
        left_table = table
832
    
833
    missing = True
834
    if conds != []:
835
        if len(conds) == 1: whitespace = ' '
836
        else: whitespace = '\n'
837
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
838
            .to_str(db) for l, r in conds], 'WHERE')
839
    if order_by != None:
840
        query += '\nORDER BY '+sql_gen.as_Col(order_by).to_str(db)
841
    if limit != None: query += '\nLIMIT '+str(limit)
842
    if start != None:
843
        if start != 0: query += '\nOFFSET '+str(start)
844
    
845
    if explain: query = with_explain_comment(db, query)
846
    
847
    return query
848

    
849
def select(db, *args, **kw_args):
850
    '''For params, see mk_select() and run_query()'''
851
    recover = kw_args.pop('recover', None)
852
    cacheable = kw_args.pop('cacheable', True)
853
    log_level = kw_args.pop('log_level', 2)
854
    
855
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
856
        log_level=log_level)
857

    
858
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
859
    embeddable=False, ignore=False, src=None):
860
    '''
861
    @param returning str|None An inserted column (such as pkey) to return
862
    @param embeddable Whether the query should be embeddable as a nested SELECT.
863
        Warning: If you set this and cacheable=True when the query is run, the
864
        query will be fully cached, not just if it raises an exception.
865
    @param ignore Whether to ignore duplicate keys.
866
    @param src Will be included in the name of any created function, to help
867
        identify the data source in pg_stat_activity.
868
    '''
869
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
870
    if cols == []: cols = None # no cols (all defaults) = unknown col names
871
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
872
    if select_query == None: select_query = 'DEFAULT VALUES'
873
    if returning != None: returning = sql_gen.as_Col(returning, table)
874
    
875
    first_line = 'INSERT INTO '+table.to_str(db)
876
    
877
    def mk_insert(select_query):
878
        query = first_line
879
        if cols != None:
880
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
881
        query += '\n'+select_query
882
        
883
        if returning != None:
884
            returning_name_col = sql_gen.to_name_only_col(returning)
885
            query += '\nRETURNING '+returning_name_col.to_str(db)
886
        
887
        return query
888
    
889
    return_type = sql_gen.CustomCode('unknown')
890
    if returning != None: return_type = sql_gen.ColType(returning)
891
    
892
    if ignore:
893
        # Always return something to set the correct rowcount
894
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
895
        
896
        embeddable = True # must use function
897
        
898
        if cols == None: row = [sql_gen.Col(sql_gen.all_cols, 'row')]
899
        else: row = [sql_gen.Col(c.name, 'row') for c in cols]
900
        
901
        query = sql_gen.RowExcIgnore(sql_gen.RowType(table), select_query,
902
            sql_gen.ReturnQuery(mk_insert(sql_gen.Values(row).to_str(db))),
903
            cols)
904
    else: query = mk_insert(select_query)
905
    
906
    if embeddable:
907
        # Create function
908
        function_name = sql_gen.clean_name(first_line)
909
        if src != None: function_name = src+': '+function_name
910
        while True:
911
            try:
912
                func = db.TempFunction(function_name)
913
                def_ = sql_gen.FunctionDef(func, sql_gen.SetOf(return_type),
914
                    query)
915
                
916
                run_query(db, def_.to_str(db), recover=True, cacheable=True,
917
                    log_ignore_excs=(DuplicateException,))
918
                break # this version was successful
919
            except DuplicateException, e:
920
                function_name = next_version(function_name)
921
                # try again with next version of name
922
        
923
        # Return query that uses function
924
        cols = None
925
        if returning != None: cols = [returning]
926
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(func), cols)
927
            # AS clause requires function alias
928
        return mk_select(db, func_table, order_by=None)
929
    
930
    return query
931

    
932
def insert_select(db, table, *args, **kw_args):
933
    '''For params, see mk_insert_select() and run_query_into()
934
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
935
        values in
936
    '''
937
    returning = kw_args.get('returning', None)
938
    ignore = kw_args.get('ignore', False)
939
    
940
    into = kw_args.pop('into', None)
941
    if into != None: kw_args['embeddable'] = True
942
    recover = kw_args.pop('recover', None)
943
    if ignore: recover = True
944
    cacheable = kw_args.pop('cacheable', True)
945
    log_level = kw_args.pop('log_level', 2)
946
    
947
    rowcount_only = ignore and returning == None # keep NULL rows on server
948
    if rowcount_only: into = sql_gen.Table('rowcount')
949
    
950
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
951
        into, recover=recover, cacheable=cacheable, log_level=log_level)
952
    if rowcount_only: empty_temp(db, into)
953
    autoanalyze(db, table)
954
    return cur
955

    
956
default = sql_gen.default # tells insert() to use the default value for a column
957

    
958
def insert(db, table, row, *args, **kw_args):
959
    '''For params, see insert_select()'''
960
    ignore = kw_args.pop('ignore', False)
961
    if ignore: kw_args.setdefault('recover', True)
962
    
963
    if lists.is_seq(row): cols = None
964
    else:
965
        cols = row.keys()
966
        row = row.values()
967
    row = list(row) # ensure that "== []" works
968
    
969
    if row == []: query = None
970
    else: query = sql_gen.Values(row).to_str(db)
971
    
972
    try: return insert_select(db, table, cols, query, *args, **kw_args)
973
    except (DuplicateKeyException, NullValueException):
974
        if not ignore: raise
975
        return None
976

    
977
def mk_update(db, table, changes=None, cond=None, in_place=False,
978
    cacheable_=True):
979
    '''
980
    @param changes [(col, new_value),...]
981
        * container can be any iterable type
982
        * col: sql_gen.Code|str (for col name)
983
        * new_value: sql_gen.Code|literal value
984
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
985
    @param in_place If set, locks the table and updates rows in place.
986
        This avoids creating dead rows in PostgreSQL.
987
        * cond must be None
988
    @param cacheable_ Whether column structure information used to generate the
989
        query can be cached
990
    @return str query
991
    '''
992
    table = sql_gen.as_Table(table)
993
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
994
        for c, v in changes]
995
    
996
    if in_place:
997
        assert cond == None
998
        
999
        def col_type(col):
1000
            return sql_gen.canon_type(db.col_info(
1001
                sql_gen.with_default_table(c, table), cacheable_).type)
1002
        changes = [(c, v, col_type(c)) for c, v in changes]
1003
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
1004
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '+t+'\nUSING '
1005
            +v.to_str(db) for c, v, t in changes))
1006
    else:
1007
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
1008
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
1009
            for c, v in changes))
1010
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
1011
    
1012
    query = with_explain_comment(db, query)
1013
    
1014
    return query
1015

    
1016
def update(db, table, *args, **kw_args):
1017
    '''For params, see mk_update() and run_query()'''
1018
    recover = kw_args.pop('recover', None)
1019
    cacheable = kw_args.pop('cacheable', False)
1020
    log_level = kw_args.pop('log_level', 2)
1021
    
1022
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
1023
        cacheable, log_level=log_level)
1024
    autoanalyze(db, table)
1025
    return cur
1026

    
1027
def mk_delete(db, table, cond=None):
1028
    '''
1029
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
1030
    @return str query
1031
    '''
1032
    query = 'DELETE FROM '+table.to_str(db)
1033
    if cond != None: query += '\nWHERE '+cond.to_str(db)
1034
    
1035
    query = with_explain_comment(db, query)
1036
    
1037
    return query
1038

    
1039
def delete(db, table, *args, **kw_args):
1040
    '''For params, see mk_delete() and run_query()'''
1041
    recover = kw_args.pop('recover', None)
1042
    cacheable = kw_args.pop('cacheable', True)
1043
    log_level = kw_args.pop('log_level', 2)
1044
    
1045
    cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
1046
        cacheable, log_level=log_level)
1047
    autoanalyze(db, table)
1048
    return cur
1049

    
1050
def last_insert_id(db):
1051
    module = util.root_module(db.db)
1052
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
1053
    elif module == 'MySQLdb': return db.insert_id()
1054
    else: return None
1055

    
1056
def define_func(db, def_):
1057
    func = def_.function
1058
    while True:
1059
        try:
1060
            run_query(db, def_.to_str(db), recover=True, cacheable=True,
1061
                log_ignore_excs=(DuplicateException,))
1062
            break # successful
1063
        except DuplicateException:
1064
            func.name = next_version(func.name)
1065
            # try again with next version of name
1066

    
1067
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
1068
    '''Creates a mapping from original column names (which may have collisions)
1069
    to names that will be distinct among the columns' tables.
1070
    This is meant to be used for several tables that are being joined together.
1071
    @param cols The columns to combine. Duplicates will be removed.
1072
    @param into The table for the new columns.
1073
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
1074
        columns will be included in the mapping even if they are not in cols.
1075
        The tables of the provided Col objects will be changed to into, so make
1076
        copies of them if you want to keep the original tables.
1077
    @param as_items Whether to return a list of dict items instead of a dict
1078
    @return dict(orig_col=new_col, ...)
1079
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
1080
        * new_col: sql_gen.Col(orig_col_name, into)
1081
        * All mappings use the into table so its name can easily be
1082
          changed for all columns at once
1083
    '''
1084
    cols = lists.uniqify(cols)
1085
    
1086
    items = []
1087
    for col in preserve:
1088
        orig_col = copy.copy(col)
1089
        col.table = into
1090
        items.append((orig_col, col))
1091
    preserve = set(preserve)
1092
    for col in cols:
1093
        if col not in preserve:
1094
            items.append((col, sql_gen.Col(strings.ustr(col), into, col.srcs)))
1095
    
1096
    if not as_items: items = dict(items)
1097
    return items
1098

    
1099
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
1100
    '''For params, see mk_flatten_mapping()
1101
    @return See return value of mk_flatten_mapping()
1102
    '''
1103
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
1104
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
1105
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
1106
        into=into, add_pkey_=True)
1107
        # don't cache because the temp table will usually be truncated after use
1108
    return dict(items)
1109

    
1110
##### Database structure introspection
1111

    
1112
#### Tables
1113

    
1114
def tables(db, schema_like='public', table_like='%', exact=False,
1115
    cacheable=True):
1116
    if exact: compare = '='
1117
    else: compare = 'LIKE'
1118
    
1119
    module = util.root_module(db.db)
1120
    if module == 'psycopg2':
1121
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1122
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1123
        return values(select(db, 'pg_tables', ['tablename'], conds,
1124
            order_by='tablename', cacheable=cacheable, log_level=4))
1125
    elif module == 'MySQLdb':
1126
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1127
            , cacheable=True, log_level=4))
1128
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1129

    
1130
def table_exists(db, table, cacheable=True):
1131
    table = sql_gen.as_Table(table)
1132
    return list(tables(db, table.schema, table.name, True, cacheable)) != []
1133

    
1134
def table_row_count(db, table, recover=None):
1135
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
1136
        order_by=None), recover=recover, log_level=3))
1137

    
1138
def table_col_names(db, table, recover=None):
1139
    return list(col_names(select(db, table, limit=0, recover=recover,
1140
        log_level=4)))
1141

    
1142
def table_cols(db, table, *args, **kw_args):
1143
    return [sql_gen.as_Col(strings.ustr(c), table)
1144
        for c in table_col_names(db, table, *args, **kw_args)]
1145

    
1146
def table_pkey_index(db, table, recover=None):
1147
    table_str = sql_gen.Literal(table.to_str(db))
1148
    try:
1149
        return sql_gen.Table(value(run_query(db, '''\
1150
SELECT relname
1151
FROM pg_index
1152
JOIN pg_class index ON index.oid = indexrelid
1153
WHERE
1154
indrelid = '''+table_str.to_str(db)+'''::regclass
1155
AND indisprimary
1156
'''
1157
            , recover, cacheable=True, log_level=4)), table.schema)
1158
    except StopIteration: raise DoesNotExistException('primary key', '')
1159

    
1160
def table_pkey_col(db, table, recover=None):
1161
    table = sql_gen.as_Table(table)
1162
    
1163
    module = util.root_module(db.db)
1164
    if module == 'psycopg2':
1165
        return sql_gen.Col(index_cols(db, table_pkey_index(db, table,
1166
            recover))[0], table)
1167
    else:
1168
        join_cols = ['table_schema', 'table_name', 'constraint_schema',
1169
            'constraint_name']
1170
        tables = [sql_gen.Table('key_column_usage', 'information_schema'),
1171
            sql_gen.Join(
1172
                sql_gen.Table('table_constraints', 'information_schema'),
1173
                dict(((c, sql_gen.join_same_not_null) for c in join_cols)))]
1174
        cols = [sql_gen.Col('column_name')]
1175
        
1176
        conds = [('constraint_type', 'PRIMARY KEY'), ('table_name', table.name)]
1177
        schema = table.schema
1178
        if schema != None: conds.append(('table_schema', schema))
1179
        order_by = 'position_in_unique_constraint'
1180
        
1181
        try: return sql_gen.Col(value(select(db, tables, cols, conds,
1182
            order_by=order_by, limit=1, log_level=4)), table)
1183
        except StopIteration: raise DoesNotExistException('primary key', '')
1184

    
1185
def table_has_pkey(db, table, recover=None):
1186
    try: table_pkey_col(db, table, recover)
1187
    except DoesNotExistException: return False
1188
    else: return True
1189

    
1190
def pkey_name(db, table, recover=None):
1191
    '''If no pkey, returns the first column in the table.'''
1192
    return pkey_col(db, table, recover).name
1193

    
1194
def pkey_col(db, table, recover=None):
1195
    '''If no pkey, returns the first column in the table.'''
1196
    try: return table_pkey_col(db, table, recover)
1197
    except DoesNotExistException: return table_cols(db, table, recover)[0]
1198

    
1199
not_null_col = 'not_null_col'
1200

    
1201
def table_not_null_col(db, table, recover=None):
1202
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1203
    if not_null_col in table_col_names(db, table, recover): return not_null_col
1204
    else: return pkey_name(db, table, recover)
1205

    
1206
def constraint_cond(db, constraint):
1207
    module = util.root_module(db.db)
1208
    if module == 'psycopg2':
1209
        table_str = sql_gen.Literal(constraint.table.to_str(db))
1210
        name_str = sql_gen.Literal(constraint.name)
1211
        return value(run_query(db, '''\
1212
SELECT consrc
1213
FROM pg_constraint
1214
WHERE
1215
conrelid = '''+table_str.to_str(db)+'''::regclass
1216
AND conname = '''+name_str.to_str(db)+'''
1217
'''
1218
            , cacheable=True, log_level=4))
1219
    else: raise NotImplementedError("Can't get constraint condition for "
1220
        +module+' database')
1221

    
1222
def index_exprs(db, index):
1223
    index = sql_gen.as_Table(index)
1224
    module = util.root_module(db.db)
1225
    if module == 'psycopg2':
1226
        qual_index = sql_gen.Literal(index.to_str(db))
1227
        return list(values(run_query(db, '''\
1228
SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
1229
FROM pg_index
1230
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1231
'''
1232
            , cacheable=True, log_level=4)))
1233
    else: raise NotImplementedError()
1234

    
1235
def index_cols(db, index):
1236
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1237
    automatically created. When you don't know whether something is a UNIQUE
1238
    constraint or a UNIQUE index, use this function.'''
1239
    return map(sql_gen.parse_expr_col, index_exprs(db, index))
1240

    
1241
def index_cond(db, index):
1242
    index = sql_gen.as_Table(index)
1243
    module = util.root_module(db.db)
1244
    if module == 'psycopg2':
1245
        qual_index = sql_gen.Literal(index.to_str(db))
1246
        return value(run_query(db, '''\
1247
SELECT pg_get_expr(indpred, indrelid, true)
1248
FROM pg_index
1249
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1250
'''
1251
            , cacheable=True, log_level=4))
1252
    else: raise NotImplementedError()
1253

    
1254
def index_order_by(db, index):
1255
    return sql_gen.CustomCode(', '.join(index_exprs(db, index)))
1256

    
1257
def table_cluster_on(db, table, recover=None):
1258
    '''
1259
    @return The table's cluster index, or its pkey if none is set
1260
    '''
1261
    table_str = sql_gen.Literal(table.to_str(db))
1262
    try:
1263
        return sql_gen.Table(value(run_query(db, '''\
1264
SELECT relname
1265
FROM pg_index
1266
JOIN pg_class index ON index.oid = indexrelid
1267
WHERE
1268
indrelid = '''+table_str.to_str(db)+'''::regclass
1269
AND indisclustered
1270
'''
1271
            , recover, cacheable=True, log_level=4)), table.schema)
1272
    except StopIteration: return table_pkey_index(db, table, recover)
1273

    
1274
def table_order_by(db, table, recover=None):
1275
    '''
1276
    @return None if table is view, because table_cluster_on() would return None
1277
    '''
1278
    if table.order_by == None:
1279
        try: table.order_by = index_order_by(db, table_cluster_on(db, table,
1280
            recover))
1281
        except DoesNotExistException: pass
1282
    return table.order_by
1283

    
1284
#### Views
1285

    
1286
def view_exists(db, view):
1287
    view_str = sql_gen.Literal(view.to_str(db))
1288
    try:
1289
        return value(run_query(db, '''\
1290
SELECT relkind = 'v'
1291
FROM pg_class
1292
WHERE oid = '''+view_str.to_str(db)+'''::regclass
1293
'''
1294
            , cacheable=True, log_level=4))
1295
    except DoesNotExistException: return False
1296

    
1297
#### Functions
1298

    
1299
def function_exists(db, function):
1300
    qual_function = sql_gen.Literal(function.to_str(db))
1301
    try:
1302
        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
1303
            recover=True, cacheable=True, log_level=4)
1304
    except DoesNotExistException: return False
1305
    except DuplicateException: return True # overloaded function
1306
    else: return True
1307

    
1308
def function_param0_type(db, function):
1309
    qual_function = sql_gen.Literal(function.to_str(db))
1310
    return value(run_query(db, '''\
1311
SELECT proargtypes[0]::regtype
1312
FROM pg_proc
1313
WHERE oid = '''+qual_function.to_str(db)+'''::regproc
1314
'''
1315
        , cacheable=True, log_level=4))
1316

    
1317
##### Structural changes
1318

    
1319
#### Columns
1320

    
1321
def add_col(db, table, col, comment=None, if_not_exists=False, **kw_args):
1322
    '''
1323
    @param col TypedCol Name may be versioned, so be sure to propagate any
1324
        renaming back to any source column for the TypedCol.
1325
    @param comment None|str SQL comment used to distinguish columns of the same
1326
        name from each other when they contain different data, to allow the
1327
        ADD COLUMN query to be cached. If not set, query will not be cached.
1328
    '''
1329
    assert isinstance(col, sql_gen.TypedCol)
1330
    
1331
    while True:
1332
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1333
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1334
        
1335
        try:
1336
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1337
            break
1338
        except DuplicateException:
1339
            if if_not_exists: raise
1340
            col.name = next_version(col.name)
1341
            # try again with next version of name
1342

    
1343
def add_not_null(db, col):
1344
    table = col.table
1345
    col = sql_gen.to_name_only_col(col)
1346
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1347
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1348

    
1349
def drop_not_null(db, col):
1350
    table = col.table
1351
    col = sql_gen.to_name_only_col(col)
1352
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1353
        +col.to_str(db)+' DROP NOT NULL', cacheable=True, log_level=3)
1354

    
1355
row_num_col = '_row_num'
1356

    
1357
row_num_col_def = sql_gen.TypedCol('', 'serial', nullable=False,
1358
    constraints='PRIMARY KEY')
1359

    
1360
def add_row_num(db, table, name=row_num_col):
1361
    '''Adds a row number column to a table. Its definition is in
1362
    row_num_col_def. It will be the primary key.'''
1363
    col_def = copy.copy(row_num_col_def)
1364
    col_def.name = name
1365
    add_col(db, table, col_def, comment='', if_not_exists=True, log_level=3)
1366

    
1367
#### Indexes
1368

    
1369
def add_pkey(db, table, cols=None, recover=None):
1370
    '''Adds a primary key.
1371
    @param cols [sql_gen.Col,...] The columns in the primary key.
1372
        Defaults to the first column in the table.
1373
    @pre The table must not already have a primary key.
1374
    '''
1375
    table = sql_gen.as_Table(table)
1376
    if cols == None: cols = [pkey_name(db, table, recover)]
1377
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1378
    
1379
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1380
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1381
        log_ignore_excs=(DuplicateException,))
1382

    
1383
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1384
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1385
    Currently, only function calls and literal values are supported expressions.
1386
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1387
        This allows indexes to be used for comparisons where NULLs are equal.
1388
    '''
1389
    exprs = lists.mk_seq(exprs)
1390
    
1391
    # Parse exprs
1392
    old_exprs = exprs[:]
1393
    exprs = []
1394
    cols = []
1395
    for i, expr in enumerate(old_exprs):
1396
        expr = sql_gen.as_Col(expr, table)
1397
        
1398
        # Handle nullable columns
1399
        if ensure_not_null_:
1400
            try: expr = sql_gen.ensure_not_null(db, expr)
1401
            except KeyError: pass # unknown type, so just create plain index
1402
        
1403
        # Extract col
1404
        expr = copy.deepcopy(expr) # don't modify input!
1405
        col = expr
1406
        if isinstance(expr, sql_gen.FunctionCall): col = expr.args[0]
1407
        expr = sql_gen.cast_literal(expr)
1408
        if not isinstance(expr, (sql_gen.Expr, sql_gen.Col)):
1409
            expr = sql_gen.Expr(expr)
1410
            
1411
        
1412
        # Extract table
1413
        if table == None:
1414
            assert sql_gen.is_table_col(col)
1415
            table = col.table
1416
        
1417
        if isinstance(col, sql_gen.Col): col.table = None
1418
        
1419
        exprs.append(expr)
1420
        cols.append(col)
1421
    
1422
    table = sql_gen.as_Table(table)
1423
    
1424
    # Add index
1425
    str_ = 'CREATE'
1426
    if unique: str_ += ' UNIQUE'
1427
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1428
        ', '.join((v.to_str(db) for v in exprs)))+')'
1429
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1430

    
1431
def add_pkey_index(db, table): add_index(db, pkey_col(db, table), table)
1432

    
1433
def add_pkey_or_index(db, table, cols=None, recover=None, warn=False):
1434
    try: add_pkey(db, table, cols, recover)
1435
    except DuplicateKeyException, e:
1436
        if warn: warnings.warn(UserWarning(exc.str_(e)))
1437
        add_pkey_index(db, table)
1438

    
1439
already_indexed = object() # tells add_indexes() the pkey has already been added
1440

    
1441
def add_indexes(db, table, has_pkey=True):
1442
    '''Adds an index on all columns in a table.
1443
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1444
        index should be added on the first column.
1445
        * If already_indexed, the pkey is assumed to have already been added
1446
    '''
1447
    cols = table_col_names(db, table)
1448
    if has_pkey:
1449
        if has_pkey is not already_indexed: add_pkey(db, table)
1450
        cols = cols[1:]
1451
    for col in cols: add_index(db, col, table)
1452

    
1453
#### Tables
1454

    
1455
### Maintenance
1456

    
1457
def analyze(db, table):
1458
    table = sql_gen.as_Table(table)
1459
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1460

    
1461
def autoanalyze(db, table):
1462
    if db.autoanalyze: analyze(db, table)
1463

    
1464
def vacuum(db, table):
1465
    table = sql_gen.as_Table(table)
1466
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1467
        log_level=3))
1468

    
1469
### Lifecycle
1470

    
1471
def drop(db, type_, name):
1472
    name = sql_gen.as_Name(name)
1473
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1474

    
1475
def drop_table(db, table): drop(db, 'TABLE', table)
1476

    
1477
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1478
    like=None):
1479
    '''Creates a table.
1480
    @param cols [sql_gen.TypedCol,...] The column names and types
1481
    @param has_pkey If set, the first column becomes the primary key.
1482
    @param col_indexes bool|[ref]
1483
        * If True, indexes will be added on all non-pkey columns.
1484
        * If a list reference, [0] will be set to a function to do this.
1485
          This can be used to delay index creation until the table is populated.
1486
    '''
1487
    table = sql_gen.as_Table(table)
1488
    
1489
    if like != None:
1490
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1491
            ]+cols
1492
        table.order_by = like.order_by
1493
    if has_pkey:
1494
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1495
        pkey.constraints = 'PRIMARY KEY'
1496
    
1497
    temp = table.is_temp and not db.debug_temp
1498
        # temp tables permanent in debug_temp mode
1499
    
1500
    # Create table
1501
    def create():
1502
        str_ = 'CREATE'
1503
        if temp: str_ += ' TEMP'
1504
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1505
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1506
        str_ += '\n);'
1507
        
1508
        opts = dict(recover=True, cacheable=True, log_level=2,
1509
            log_ignore_excs=(DuplicateException,))
1510
        try: run_query(db, str_, **opts)
1511
        except InvalidTypeException: # try again as view
1512
            run_query_into(db, mk_select(db, like, limit=0), into=table, **opts)
1513
    if table.is_temp:
1514
        while True:
1515
            try:
1516
                create()
1517
                break
1518
            except DuplicateException:
1519
                table.name = next_version(table.name)
1520
                # try again with next version of name
1521
    else: create()
1522
    
1523
    # Add indexes
1524
    if has_pkey: has_pkey = already_indexed
1525
    def add_indexes_(): add_indexes(db, table, has_pkey)
1526
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1527
    elif col_indexes: add_indexes_() # add now
1528

    
1529
def copy_table_struct(db, src, dest):
1530
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1531
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1532

    
1533
def copy_table(db, src, dest):
1534
    '''Creates a copy of a table, including data'''
1535
    copy_table_struct(db, src, dest)
1536
    insert_select(db, dest, None, mk_select(db, src))
1537

    
1538
### Data
1539

    
1540
def truncate(db, table, schema='public', **kw_args):
1541
    '''For params, see run_query()'''
1542
    table = sql_gen.as_Table(table, schema)
1543
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1544

    
1545
def empty_temp(db, tables):
1546
    tables = lists.mk_seq(tables)
1547
    for table in tables: truncate(db, table, log_level=3)
1548

    
1549
def empty_db(db, schema='public', **kw_args):
1550
    '''For kw_args, see tables()'''
1551
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1552

    
1553
def distinct_table(db, table, distinct_on, joins=None):
1554
    '''Creates a copy of a temp table which is distinct on the given columns.
1555
    Adds an index on table's distinct_on columns, to facilitate merge joins.
1556
    @param distinct_on If empty, creates a table with one row. This is useful if
1557
        your distinct_on columns are all literal values.
1558
    @param joins The joins to use when creating the new table
1559
    @return The new table.
1560
    '''
1561
    if joins == None: joins = [table]
1562
    
1563
    new_table = sql_gen.suffixed_table(table, '_distinct')
1564
    distinct_on = filter(sql_gen.is_table_col, distinct_on)
1565
    
1566
    copy_table_struct(db, table, new_table)
1567
    
1568
    limit = None
1569
    if distinct_on == []: limit = 1 # one sample row
1570
    else: add_index(db, distinct_on, table) # for join optimization
1571
    
1572
    insert_select(db, new_table, None, mk_select(db, joins,
1573
        [sql_gen.Col(sql_gen.all_cols, table)], distinct_on=distinct_on,
1574
        order_by=None, limit=limit))
1575
    analyze(db, new_table)
1576
    
1577
    return new_table
(36-36/50)