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

    
636
def with_savepoint(db, func): return db.with_savepoint(func)
637

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

    
666
##### Basic queries
667

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

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

    
684
def has_comment(query): return query.endswith('*/')
685

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

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

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

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

    
747
order_by_pkey = object() # tells mk_select() to order by the pkey
748

    
749
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
750

    
751
def has_subset_func(db, table):
752
    return sql_gen.is_underlying_table(table) and function_exists(db, table)
753

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    
1108
##### Database structure introspection
1109

    
1110
#### Tables
1111

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

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

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

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

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

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

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

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

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

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

    
1197
not_null_col = 'not_null_col'
1198

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

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

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

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

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

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

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

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

    
1282
#### Views
1283

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

    
1295
#### Functions
1296

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

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

    
1315
##### Structural changes
1316

    
1317
#### Columns
1318

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

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

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

    
1353
row_num_col = '_row_num'
1354

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

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

    
1365
#### Indexes
1366

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

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

    
1429
def add_pkey_index(db, table): add_index(db, pkey_col(db, table), table)
1430

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

    
1437
already_indexed = object() # tells add_indexes() the pkey has already been added
1438

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

    
1451
#### Tables
1452

    
1453
### Maintenance
1454

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

    
1459
def autoanalyze(db, table):
1460
    if db.autoanalyze: analyze(db, table)
1461

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

    
1467
### Lifecycle
1468

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

    
1473
def drop_table(db, table): drop(db, 'TABLE', table)
1474

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

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

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

    
1536
### Data
1537

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

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

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

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