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 DuplicateException(ExceptionWithNameType): pass
86

    
87
class DoesNotExistException(ExceptionWithNameType): pass
88

    
89
class EmptyRowException(DbException): pass
90

    
91
##### Warnings
92

    
93
class DbWarning(UserWarning): pass
94

    
95
##### Result retrieval
96

    
97
def col_names(cur): return (col[0] for col in cur.description)
98

    
99
def rows(cur): return iter(lambda: cur.fetchone(), None)
100

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

    
105
def next_row(cur): return rows(cur).next()
106

    
107
def row(cur):
108
    row_ = next_row(cur)
109
    consume_rows(cur)
110
    return row_
111

    
112
def next_value(cur): return next_row(cur)[0]
113

    
114
def value(cur): return row(cur)[0]
115

    
116
def values(cur): return iters.func_iter(lambda: next_value(cur))
117

    
118
def value_or_none(cur):
119
    try: return value(cur)
120
    except StopIteration: return None
121

    
122
##### Escaping
123

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

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

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

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

    
142
##### Database connections
143

    
144
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
145

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

    
151
DatabaseErrors_set = set([DbException])
152
DatabaseErrors = tuple(DatabaseErrors_set)
153

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

    
159
def db_config_str(db_config):
160
    return db_config['engine']+' database '+db_config['database']
161

    
162
log_debug_none = lambda msg, level=2: None
163

    
164
class DbConn:
165
    def __init__(self, db_config, autocommit=True, caching=True,
166
        log_debug=log_debug_none, debug_temp=False, src=None):
167
        '''
168
        @param debug_temp Whether temporary objects should instead be permanent.
169
            This assists in debugging the internal objects used by the program.
170
        @param src In autocommit mode, will be included in a comment in every
171
            query, to help identify the data source in pg_stat_activity.
172
        '''
173
        self.db_config = db_config
174
        self.autocommit = autocommit
175
        self.caching = caching
176
        self.log_debug = log_debug
177
        self.debug = log_debug != log_debug_none
178
        self.debug_temp = debug_temp
179
        self.src = src
180
        self.autoanalyze = False
181
        self.autoexplain = False
182
        self.profile_row_ct = None
183
        
184
        self._savepoint = 0
185
        self._reset()
186
    
187
    def __getattr__(self, name):
188
        if name == '__dict__': raise Exception('getting __dict__')
189
        if name == 'db': return self._db()
190
        else: raise AttributeError()
191
    
192
    def __getstate__(self):
193
        state = copy.copy(self.__dict__) # shallow copy
194
        state['log_debug'] = None # don't pickle the debug callback
195
        state['_DbConn__db'] = None # don't pickle the connection
196
        return state
197
    
198
    def clear_cache(self): self.query_results = {}
199
    
200
    def _reset(self):
201
        self.clear_cache()
202
        assert self._savepoint == 0
203
        self._notices_seen = set()
204
        self.__db = None
205
    
206
    def connected(self): return self.__db != None
207
    
208
    def close(self):
209
        if not self.connected(): return
210
        
211
        # Record that the automatic transaction is now closed
212
        self._savepoint -= 1
213
        
214
        self.db.close()
215
        self._reset()
216
    
217
    def reconnect(self):
218
        # Do not do this in test mode as it would roll back everything
219
        if self.autocommit: self.close()
220
        # Connection will be reopened automatically on first query
221
    
222
    def _db(self):
223
        if self.__db == None:
224
            # Process db_config
225
            db_config = self.db_config.copy() # don't modify input!
226
            schemas = db_config.pop('schemas', None)
227
            module_name, mappings = db_engines[db_config.pop('engine')]
228
            module = __import__(module_name)
229
            _add_module(module)
230
            for orig, new in mappings.iteritems():
231
                try: util.rename_key(db_config, orig, new)
232
                except KeyError: pass
233
            
234
            # Connect
235
            self.__db = module.connect(**db_config)
236
            
237
            # Record that a transaction is already open
238
            self._savepoint += 1
239
            
240
            # Configure connection
241
            if hasattr(self.db, 'set_isolation_level'):
242
                import psycopg2.extensions
243
                self.db.set_isolation_level(
244
                    psycopg2.extensions.ISOLATION_LEVEL_READ_COMMITTED)
245
            if schemas != None:
246
                search_path = [self.esc_name(s) for s in schemas.split(',')]
247
                search_path.append(value(run_query(self, 'SHOW search_path',
248
                    log_level=4)))
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
        table = sql_gen.Table('columns', 'information_schema')
473
        cols = [sql_gen.Col('data_type'), sql_gen.Col('udt_name'),
474
            'column_default', sql_gen.Cast('boolean',
475
            sql_gen.Col('is_nullable'))]
476
        
477
        conds = [('table_name', col.table.name),
478
            ('column_name', strings.ustr(col.name))]
479
        schema = col.table.schema
480
        if schema != None: conds.append(('table_schema', schema))
481
        
482
        cur = select(self, table, cols, conds, order_by='table_schema', limit=1,
483
            cacheable=cacheable, log_level=4) # TODO: order by search_path order
484
        try: type_, extra_type, default, nullable = row(cur)
485
        except StopIteration: raise sql_gen.NoUnderlyingTableException(col)
486
        default = sql_gen.as_Code(default, self)
487
        if type_ == 'USER-DEFINED': type_ = extra_type
488
        elif type_ == 'ARRAY':
489
            type_ = sql_gen.ArrayType(strings.remove_prefix('_', extra_type,
490
                require=True))
491
        
492
        return sql_gen.TypedCol(col.name, type_, default, nullable)
493
    
494
    def TempFunction(self, name):
495
        if self.debug_temp: schema = None
496
        else: schema = 'pg_temp'
497
        return sql_gen.Function(name, schema)
498

    
499
connect = DbConn
500

    
501
##### Recoverable querying
502

    
503
def parse_exception(db, e, recover=False):
504
    msg = strings.ustr(e.args[0])
505
    msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)
506
    
507
    match = re.match(r'^invalid byte sequence for encoding "(.+?)":', msg)
508
    if match:
509
        encoding, = match.groups()
510
        raise EncodingException(encoding, e)
511
    
512
    match = re.match(r'^duplicate key value violates unique constraint "(.+?)"',
513
        msg)
514
    if match:
515
        constraint, = match.groups()
516
        cols = []
517
        cond = None
518
        if recover: # need auto-rollback to run index_cols()
519
            try:
520
                cols = index_cols(db, constraint)
521
                cond = index_cond(db, constraint)
522
            except NotImplementedError: pass
523
        raise DuplicateKeyException(constraint, cond, cols, e)
524
    
525
    match = re.match(r'^null value in column "(.+?)" violates not-null'
526
        r' constraint', msg)
527
    if match:
528
        col, = match.groups()
529
        raise NullValueException('NOT NULL', None, [col], e)
530
    
531
    match = re.match(r'^new row for relation "(.+?)" violates check '
532
        r'constraint "(.+?)"', msg)
533
    if match:
534
        table, constraint = match.groups()
535
        constraint = sql_gen.Col(constraint, table)
536
        cond = None
537
        if recover: # need auto-rollback to run constraint_cond()
538
            try: cond = constraint_cond(db, constraint)
539
            except NotImplementedError: pass
540
        raise CheckException(constraint.to_str(db), cond, [], e)
541
    
542
    match = re.match(r'^(?:invalid input (?:syntax|value)\b.*?'
543
        r'|.+? out of range)(?:: "(.+?)")?', msg)
544
    if match:
545
        value, = match.groups()
546
        value = util.do_ignore_none(strings.to_unicode, value)
547
        raise InvalidValueException(value, e)
548
    
549
    match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
550
        r'is of type', msg)
551
    if match:
552
        col, type_ = match.groups()
553
        raise MissingCastException(type_, col, e)
554
    
555
    match = re.match(r'^could not determine polymorphic type because '
556
        r'input has type "unknown"', msg)
557
    if match: raise MissingCastException('text', None, e)
558
    
559
    match = re.match(r'^.+? types .+? and .+? cannot be matched', msg)
560
    if match: raise MissingCastException('text', None, e)
561
    
562
    typed_name_re = r'^(\S+) "?(.+?)"?(?: of relation ".+?")?'
563
    
564
    match = re.match(typed_name_re+r'.*? already exists', msg)
565
    if match:
566
        type_, name = match.groups()
567
        raise DuplicateException(type_, name, e)
568
    
569
    match = re.match(r'more than one (\S+) named ""(.+?)""', msg)
570
    if match:
571
        type_, name = match.groups()
572
        raise DuplicateException(type_, name, e)
573
    
574
    match = re.match(typed_name_re+r' does not exist', msg)
575
    if match:
576
        type_, name = match.groups()
577
        if type_ == 'function':
578
            match = re.match(r'^(.+?)\(.*\)$', name)
579
            if match: # includes params, so is call rather than cast to regproc
580
                function_name, = match.groups()
581
                func = sql_gen.Function(function_name)
582
                if function_exists(db, func) and msg.find('CAST') < 0:
583
                    # not found only because of a missing cast
584
                    type_ = function_param0_type(db, func)
585
                    raise MissingCastException(type_, None, e)
586
        raise DoesNotExistException(type_, name, e)
587
    
588
    raise # no specific exception raised
589

    
590
def with_savepoint(db, func): return db.with_savepoint(func)
591

    
592
def run_query(db, query, recover=None, cacheable=False, log_level=2,
593
    log_ignore_excs=None, **kw_args):
594
    '''For params, see DbConn.run_query()'''
595
    if recover == None: recover = False
596
    if log_ignore_excs == None: log_ignore_excs = ()
597
    log_ignore_excs = tuple(log_ignore_excs)
598
    debug_msg_ref = [None]
599
    
600
    query = with_explain_comment(db, query)
601
    
602
    try:
603
        try:
604
            def run(): return db.run_query(query, cacheable, log_level,
605
                debug_msg_ref, **kw_args)
606
            if recover and not db.is_cached(query):
607
                return with_savepoint(db, run)
608
            else: return run() # don't need savepoint if cached
609
        except Exception, e: parse_exception(db, e, recover)
610
    except log_ignore_excs:
611
        log_level += 2
612
        raise
613
    finally:
614
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
615

    
616
##### Basic queries
617

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

    
623
def explain(db, query, **kw_args):
624
    '''
625
    For params, see run_query().
626
    '''
627
    kw_args.setdefault('log_level', 4)
628
    
629
    return strings.ustr(strings.join_lines(values(run_query(db,
630
        'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))
631
        # not a higher log_level because it's useful to see what query is being
632
        # run before it's executed, which EXPLAIN effectively provides
633

    
634
def has_comment(query): return query.endswith('*/')
635

    
636
def with_explain_comment(db, query, **kw_args):
637
    if db.autoexplain and not has_comment(query) and is_explainable(query):
638
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
639
            +explain(db, query, **kw_args))
640
    return query
641

    
642
def next_version(name):
643
    version = 1 # first existing name was version 0
644
    match = re.match(r'^(.*)#(\d+)$', name)
645
    if match:
646
        name, version = match.groups()
647
        version = int(version)+1
648
    return sql_gen.concat(name, '#'+str(version))
649

    
650
def lock_table(db, table, mode):
651
    table = sql_gen.as_Table(table)
652
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
653

    
654
def run_query_into(db, query, into=None, add_pkey_=False, **kw_args):
655
    '''Outputs a query to a temp table.
656
    For params, see run_query().
657
    '''
658
    if into == None: return run_query(db, query, **kw_args)
659
    
660
    assert isinstance(into, sql_gen.Table)
661
    
662
    into.is_temp = True
663
    # "temporary tables cannot specify a schema name", so remove schema
664
    into.schema = None
665
    
666
    kw_args['recover'] = True
667
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
668
    
669
    temp = not db.debug_temp # tables are permanent in debug_temp mode
670
    
671
    # Create table
672
    while True:
673
        create_query = 'CREATE'
674
        if temp: create_query += ' TEMP'
675
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
676
        
677
        try:
678
            cur = run_query(db, create_query, **kw_args)
679
                # CREATE TABLE AS sets rowcount to # rows in query
680
            break
681
        except DuplicateException, e:
682
            into.name = next_version(into.name)
683
            # try again with next version of name
684
    
685
    if add_pkey_: add_pkey(db, into)
686
    
687
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
688
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
689
    # table is going to be used in complex queries, it is wise to run ANALYZE on
690
    # the temporary table after it is populated."
691
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
692
    # If into is not a temp table, ANALYZE is useful but not required.
693
    analyze(db, into)
694
    
695
    return cur
696

    
697
order_by_pkey = object() # tells mk_select() to order by the pkey
698

    
699
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
700

    
701
def mk_select(db, tables=None, fields=None, conds=None, distinct_on=[],
702
    limit=None, start=None, order_by=order_by_pkey, default_table=None,
703
    explain=True):
704
    '''
705
    @param tables The single table to select from, or a list of tables to join
706
        together, with tables after the first being sql_gen.Join objects
707
    @param fields Use None to select all fields in the table
708
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
709
        * container can be any iterable type
710
        * compare_left_side: sql_gen.Code|str (for col name)
711
        * compare_right_side: sql_gen.ValueCond|literal value
712
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
713
        use all columns
714
    @return query
715
    '''
716
    # Parse tables param
717
    tables = lists.mk_seq(tables)
718
    tables = list(tables) # don't modify input! (list() copies input)
719
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
720
    
721
    # Parse other params
722
    if conds == None: conds = []
723
    elif dicts.is_dict(conds): conds = conds.items()
724
    conds = list(conds) # don't modify input! (list() copies input)
725
    assert limit == None or isinstance(limit, (int, long))
726
    assert start == None or isinstance(start, (int, long))
727
    if limit == 0: order_by = None
728
    if order_by is order_by_pkey:
729
        if lists.is_seq(distinct_on) and distinct_on: order_by = distinct_on[0]
730
        elif table0 != None: order_by = table_order_by(db, table0, recover=True)
731
        else: order_by = None
732
    
733
    query = 'SELECT'
734
    
735
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
736
    
737
    # DISTINCT ON columns
738
    if distinct_on != []:
739
        query += '\nDISTINCT'
740
        if distinct_on is not distinct_on_all:
741
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
742
    
743
    # Columns
744
    if query.find('\n') >= 0: whitespace = '\n'
745
    else: whitespace = ' '
746
    if fields == None: query += whitespace+'*'
747
    else:
748
        assert fields != []
749
        if len(fields) > 1: whitespace = '\n'
750
        query += whitespace+('\n, '.join(map(parse_col, fields)))
751
    
752
    # Main table
753
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
754
    else: whitespace = ' '
755
    if table0 != None: query += whitespace+'FROM '+table0.to_str(db)
756
    
757
    # Add joins
758
    left_table = table0
759
    for join_ in tables:
760
        table = join_.table
761
        
762
        # Parse special values
763
        if join_.type_ is sql_gen.filter_out: # filter no match
764
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
765
                sql_gen.CompareCond(None, '~=')))
766
        
767
        query += '\n'+join_.to_str(db, left_table)
768
        
769
        left_table = table
770
    
771
    missing = True
772
    if conds != []:
773
        if len(conds) == 1: whitespace = ' '
774
        else: whitespace = '\n'
775
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
776
            .to_str(db) for l, r in conds], 'WHERE')
777
    if order_by != None:
778
        query += '\nORDER BY '+sql_gen.as_Col(order_by).to_str(db)
779
    if limit != None: query += '\nLIMIT '+str(limit)
780
    if start != None:
781
        if start != 0: query += '\nOFFSET '+str(start)
782
    
783
    if explain: query = with_explain_comment(db, query)
784
    
785
    return query
786

    
787
def select(db, *args, **kw_args):
788
    '''For params, see mk_select() and run_query()'''
789
    recover = kw_args.pop('recover', None)
790
    cacheable = kw_args.pop('cacheable', True)
791
    log_level = kw_args.pop('log_level', 2)
792
    
793
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
794
        log_level=log_level)
795

    
796
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
797
    embeddable=False, ignore=False, src=None):
798
    '''
799
    @param returning str|None An inserted column (such as pkey) to return
800
    @param embeddable Whether the query should be embeddable as a nested SELECT.
801
        Warning: If you set this and cacheable=True when the query is run, the
802
        query will be fully cached, not just if it raises an exception.
803
    @param ignore Whether to ignore duplicate keys.
804
    @param src Will be included in the name of any created function, to help
805
        identify the data source in pg_stat_activity.
806
    '''
807
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
808
    if cols == []: cols = None # no cols (all defaults) = unknown col names
809
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
810
    if select_query == None: select_query = 'DEFAULT VALUES'
811
    if returning != None: returning = sql_gen.as_Col(returning, table)
812
    
813
    first_line = 'INSERT INTO '+table.to_str(db)
814
    
815
    def mk_insert(select_query):
816
        query = first_line
817
        if cols != None:
818
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
819
        query += '\n'+select_query
820
        
821
        if returning != None:
822
            returning_name_col = sql_gen.to_name_only_col(returning)
823
            query += '\nRETURNING '+returning_name_col.to_str(db)
824
        
825
        return query
826
    
827
    return_type = sql_gen.CustomCode('unknown')
828
    if returning != None: return_type = sql_gen.ColType(returning)
829
    
830
    if ignore:
831
        # Always return something to set the correct rowcount
832
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
833
        
834
        embeddable = True # must use function
835
        
836
        if cols == None: row = [sql_gen.Col(sql_gen.all_cols, 'row')]
837
        else: row = [sql_gen.Col(c.name, 'row') for c in cols]
838
        
839
        query = sql_gen.RowExcIgnore(sql_gen.RowType(table), select_query,
840
            sql_gen.ReturnQuery(mk_insert(sql_gen.Values(row).to_str(db))),
841
            cols)
842
    else: query = mk_insert(select_query)
843
    
844
    if embeddable:
845
        # Create function
846
        function_name = sql_gen.clean_name(first_line)
847
        if src != None: function_name = src+': '+function_name
848
        while True:
849
            try:
850
                func = db.TempFunction(function_name)
851
                def_ = sql_gen.FunctionDef(func, sql_gen.SetOf(return_type),
852
                    query)
853
                
854
                run_query(db, def_.to_str(db), recover=True, cacheable=True,
855
                    log_ignore_excs=(DuplicateException,))
856
                break # this version was successful
857
            except DuplicateException, e:
858
                function_name = next_version(function_name)
859
                # try again with next version of name
860
        
861
        # Return query that uses function
862
        cols = None
863
        if returning != None: cols = [returning]
864
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(func), cols)
865
            # AS clause requires function alias
866
        return mk_select(db, func_table, order_by=None)
867
    
868
    return query
869

    
870
def insert_select(db, table, *args, **kw_args):
871
    '''For params, see mk_insert_select() and run_query_into()
872
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
873
        values in
874
    '''
875
    returning = kw_args.get('returning', None)
876
    ignore = kw_args.get('ignore', False)
877
    
878
    into = kw_args.pop('into', None)
879
    if into != None: kw_args['embeddable'] = True
880
    recover = kw_args.pop('recover', None)
881
    if ignore: recover = True
882
    cacheable = kw_args.pop('cacheable', True)
883
    log_level = kw_args.pop('log_level', 2)
884
    
885
    rowcount_only = ignore and returning == None # keep NULL rows on server
886
    if rowcount_only: into = sql_gen.Table('rowcount')
887
    
888
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
889
        into, recover=recover, cacheable=cacheable, log_level=log_level)
890
    if rowcount_only: empty_temp(db, into)
891
    autoanalyze(db, table)
892
    return cur
893

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

    
896
def insert(db, table, row, *args, **kw_args):
897
    '''For params, see insert_select()'''
898
    ignore = kw_args.pop('ignore', False)
899
    if ignore: kw_args.setdefault('recover', True)
900
    
901
    if lists.is_seq(row): cols = None
902
    else:
903
        cols = row.keys()
904
        row = row.values()
905
    row = list(row) # ensure that "== []" works
906
    
907
    if row == []: query = None
908
    else: query = sql_gen.Values(row).to_str(db)
909
    
910
    try: return insert_select(db, table, cols, query, *args, **kw_args)
911
    except (DuplicateKeyException, NullValueException):
912
        if not ignore: raise
913
        return None
914

    
915
def mk_update(db, table, changes=None, cond=None, in_place=False,
916
    cacheable_=True):
917
    '''
918
    @param changes [(col, new_value),...]
919
        * container can be any iterable type
920
        * col: sql_gen.Code|str (for col name)
921
        * new_value: sql_gen.Code|literal value
922
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
923
    @param in_place If set, locks the table and updates rows in place.
924
        This avoids creating dead rows in PostgreSQL.
925
        * cond must be None
926
    @param cacheable_ Whether column structure information used to generate the
927
        query can be cached
928
    @return str query
929
    '''
930
    table = sql_gen.as_Table(table)
931
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
932
        for c, v in changes]
933
    
934
    if in_place:
935
        assert cond == None
936
        
937
        def col_type(col):
938
            return sql_gen.canon_type(db.col_info(
939
                sql_gen.with_default_table(c, table), cacheable_).type)
940
        changes = [(c, v, col_type(c)) for c, v in changes]
941
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
942
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '+t+'\nUSING '
943
            +v.to_str(db) for c, v, t in changes))
944
    else:
945
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
946
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
947
            for c, v in changes))
948
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
949
    
950
    query = with_explain_comment(db, query)
951
    
952
    return query
953

    
954
def update(db, table, *args, **kw_args):
955
    '''For params, see mk_update() and run_query()'''
956
    recover = kw_args.pop('recover', None)
957
    cacheable = kw_args.pop('cacheable', False)
958
    log_level = kw_args.pop('log_level', 2)
959
    
960
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
961
        cacheable, log_level=log_level)
962
    autoanalyze(db, table)
963
    return cur
964

    
965
def mk_delete(db, table, cond=None):
966
    '''
967
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
968
    @return str query
969
    '''
970
    query = 'DELETE FROM '+table.to_str(db)
971
    if cond != None: query += '\nWHERE '+cond.to_str(db)
972
    
973
    query = with_explain_comment(db, query)
974
    
975
    return query
976

    
977
def delete(db, table, *args, **kw_args):
978
    '''For params, see mk_delete() and run_query()'''
979
    recover = kw_args.pop('recover', None)
980
    cacheable = kw_args.pop('cacheable', True)
981
    log_level = kw_args.pop('log_level', 2)
982
    
983
    cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
984
        cacheable, log_level=log_level)
985
    autoanalyze(db, table)
986
    return cur
987

    
988
def last_insert_id(db):
989
    module = util.root_module(db.db)
990
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
991
    elif module == 'MySQLdb': return db.insert_id()
992
    else: return None
993

    
994
def define_func(db, def_):
995
    func = def_.function
996
    while True:
997
        try:
998
            run_query(db, def_.to_str(db), recover=True, cacheable=True,
999
                log_ignore_excs=(DuplicateException,))
1000
            break # successful
1001
        except DuplicateException:
1002
            func.name = next_version(func.name)
1003
            # try again with next version of name
1004

    
1005
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
1006
    '''Creates a mapping from original column names (which may have collisions)
1007
    to names that will be distinct among the columns' tables.
1008
    This is meant to be used for several tables that are being joined together.
1009
    @param cols The columns to combine. Duplicates will be removed.
1010
    @param into The table for the new columns.
1011
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
1012
        columns will be included in the mapping even if they are not in cols.
1013
        The tables of the provided Col objects will be changed to into, so make
1014
        copies of them if you want to keep the original tables.
1015
    @param as_items Whether to return a list of dict items instead of a dict
1016
    @return dict(orig_col=new_col, ...)
1017
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
1018
        * new_col: sql_gen.Col(orig_col_name, into)
1019
        * All mappings use the into table so its name can easily be
1020
          changed for all columns at once
1021
    '''
1022
    cols = lists.uniqify(cols)
1023
    
1024
    items = []
1025
    for col in preserve:
1026
        orig_col = copy.copy(col)
1027
        col.table = into
1028
        items.append((orig_col, col))
1029
    preserve = set(preserve)
1030
    for col in cols:
1031
        if col not in preserve:
1032
            items.append((col, sql_gen.Col(strings.ustr(col), into, col.srcs)))
1033
    
1034
    if not as_items: items = dict(items)
1035
    return items
1036

    
1037
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
1038
    '''For params, see mk_flatten_mapping()
1039
    @return See return value of mk_flatten_mapping()
1040
    '''
1041
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
1042
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
1043
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
1044
        into=into, add_pkey_=True)
1045
        # don't cache because the temp table will usually be truncated after use
1046
    return dict(items)
1047

    
1048
##### Database structure introspection
1049

    
1050
#### Tables
1051

    
1052
def tables(db, schema_like='public', table_like='%', exact=False,
1053
    cacheable=True):
1054
    if exact: compare = '='
1055
    else: compare = 'LIKE'
1056
    
1057
    module = util.root_module(db.db)
1058
    if module == 'psycopg2':
1059
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1060
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1061
        return values(select(db, 'pg_tables', ['tablename'], conds,
1062
            order_by='tablename', cacheable=cacheable, log_level=4))
1063
    elif module == 'MySQLdb':
1064
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1065
            , cacheable=True, log_level=4))
1066
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1067

    
1068
def table_exists(db, table, cacheable=True):
1069
    table = sql_gen.as_Table(table)
1070
    return list(tables(db, table.schema, table.name, True, cacheable)) != []
1071

    
1072
def table_row_count(db, table, recover=None):
1073
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
1074
        order_by=None), recover=recover, log_level=3))
1075

    
1076
def table_col_names(db, table, recover=None):
1077
    return list(col_names(select(db, table, limit=0, recover=recover,
1078
        log_level=4)))
1079

    
1080
def table_cols(db, table, *args, **kw_args):
1081
    return [sql_gen.as_Col(strings.ustr(c), table)
1082
        for c in table_col_names(db, table, *args, **kw_args)]
1083

    
1084
def table_pkey_index(db, table, recover=None):
1085
    table_str = sql_gen.Literal(table.to_str(db))
1086
    try:
1087
        return sql_gen.Table(value(run_query(db, '''\
1088
SELECT relname
1089
FROM pg_index
1090
JOIN pg_class index ON index.oid = indexrelid
1091
WHERE
1092
indrelid = '''+table_str.to_str(db)+'''::regclass
1093
AND indisprimary
1094
'''
1095
            , recover, cacheable=True, log_level=4)), table.schema)
1096
    except StopIteration: raise DoesNotExistException('primary key', '')
1097

    
1098
def table_pkey_col(db, table, recover=None):
1099
    table = sql_gen.as_Table(table)
1100
    
1101
    join_cols = ['table_schema', 'table_name', 'constraint_schema',
1102
        'constraint_name']
1103
    tables = [sql_gen.Table('key_column_usage', 'information_schema'),
1104
        sql_gen.Join(sql_gen.Table('table_constraints', 'information_schema'),
1105
            dict(((c, sql_gen.join_same_not_null) for c in join_cols)))]
1106
    cols = [sql_gen.Col('column_name')]
1107
    
1108
    conds = [('constraint_type', 'PRIMARY KEY'), ('table_name', table.name)]
1109
    schema = table.schema
1110
    if schema != None: conds.append(('table_schema', schema))
1111
    order_by = 'position_in_unique_constraint'
1112
    
1113
    try: return sql_gen.Col(value(select(db, tables, cols, conds,
1114
        order_by=order_by, limit=1, log_level=4)), table)
1115
    except StopIteration: raise DoesNotExistException('primary key', '')
1116

    
1117
def pkey_name(db, table, recover=None):
1118
    '''If no pkey, returns the first column in the table.'''
1119
    return pkey_col(db, table, recover).name
1120

    
1121
def pkey_col(db, table, recover=None):
1122
    '''If no pkey, returns the first column in the table.'''
1123
    try: return table_pkey_col(db, table, recover)
1124
    except DoesNotExistException: return table_cols(db, table, recover)[0]
1125

    
1126
not_null_col = 'not_null_col'
1127

    
1128
def table_not_null_col(db, table, recover=None):
1129
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1130
    if not_null_col in table_col_names(db, table, recover): return not_null_col
1131
    else: return pkey_name(db, table, recover)
1132

    
1133
def constraint_cond(db, constraint):
1134
    module = util.root_module(db.db)
1135
    if module == 'psycopg2':
1136
        table_str = sql_gen.Literal(constraint.table.to_str(db))
1137
        name_str = sql_gen.Literal(constraint.name)
1138
        return value(run_query(db, '''\
1139
SELECT consrc
1140
FROM pg_constraint
1141
WHERE
1142
conrelid = '''+table_str.to_str(db)+'''::regclass
1143
AND conname = '''+name_str.to_str(db)+'''
1144
'''
1145
            , cacheable=True, log_level=4))
1146
    else: raise NotImplementedError("Can't get constraint condition for "
1147
        +module+' database')
1148

    
1149
def index_exprs(db, index):
1150
    index = sql_gen.as_Table(index)
1151
    module = util.root_module(db.db)
1152
    if module == 'psycopg2':
1153
        qual_index = sql_gen.Literal(index.to_str(db))
1154
        return list(values(run_query(db, '''\
1155
SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
1156
FROM pg_index
1157
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1158
'''
1159
            , cacheable=True, log_level=4)))
1160
    else: raise NotImplementedError()
1161

    
1162
def index_cols(db, index):
1163
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1164
    automatically created. When you don't know whether something is a UNIQUE
1165
    constraint or a UNIQUE index, use this function.'''
1166
    return map(sql_gen.parse_expr_col, index_exprs(db, index))
1167

    
1168
def index_cond(db, index):
1169
    index = sql_gen.as_Table(index)
1170
    module = util.root_module(db.db)
1171
    if module == 'psycopg2':
1172
        qual_index = sql_gen.Literal(index.to_str(db))
1173
        return value(run_query(db, '''\
1174
SELECT pg_get_expr(indpred, indrelid, true)
1175
FROM pg_index
1176
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1177
'''
1178
            , cacheable=True, log_level=4))
1179
    else: raise NotImplementedError()
1180

    
1181
def index_order_by(db, index):
1182
    return sql_gen.CustomCode(', '.join(index_exprs(db, index)))
1183

    
1184
def table_cluster_on(db, table, recover=None):
1185
    '''
1186
    @return The table's cluster index, or its pkey if none is set
1187
    '''
1188
    table_str = sql_gen.Literal(table.to_str(db))
1189
    try:
1190
        return sql_gen.Table(value(run_query(db, '''\
1191
SELECT relname
1192
FROM pg_index
1193
JOIN pg_class index ON index.oid = indexrelid
1194
WHERE
1195
indrelid = '''+table_str.to_str(db)+'''::regclass
1196
AND indisclustered
1197
'''
1198
            , recover, cacheable=True, log_level=4)), table.schema)
1199
    except StopIteration: return table_pkey_index(db, table, recover)
1200

    
1201
def table_order_by(db, table, recover=None):
1202
    if table.order_by == None:
1203
        try: table.order_by = index_order_by(db, table_cluster_on(db, table,
1204
            recover))
1205
        except DoesNotExistException: pass
1206
    return table.order_by
1207

    
1208
#### Functions
1209

    
1210
def function_exists(db, function):
1211
    qual_function = sql_gen.Literal(function.to_str(db))
1212
    try:
1213
        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
1214
            recover=True, cacheable=True, log_level=4)
1215
    except DoesNotExistException: return False
1216
    except DuplicateException: return True # overloaded function
1217
    else: return True
1218

    
1219
def function_param0_type(db, function):
1220
    qual_function = sql_gen.Literal(function.to_str(db))
1221
    return value(run_query(db, '''\
1222
SELECT proargtypes[0]::regtype
1223
FROM pg_proc
1224
WHERE oid = '''+qual_function.to_str(db)+'''::regproc
1225
'''
1226
        , cacheable=True, log_level=4))
1227

    
1228
##### Structural changes
1229

    
1230
#### Columns
1231

    
1232
def add_col(db, table, col, comment=None, if_not_exists=False, **kw_args):
1233
    '''
1234
    @param col TypedCol Name may be versioned, so be sure to propagate any
1235
        renaming back to any source column for the TypedCol.
1236
    @param comment None|str SQL comment used to distinguish columns of the same
1237
        name from each other when they contain different data, to allow the
1238
        ADD COLUMN query to be cached. If not set, query will not be cached.
1239
    '''
1240
    assert isinstance(col, sql_gen.TypedCol)
1241
    
1242
    while True:
1243
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1244
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1245
        
1246
        try:
1247
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1248
            break
1249
        except DuplicateException:
1250
            if if_not_exists: raise
1251
            col.name = next_version(col.name)
1252
            # try again with next version of name
1253

    
1254
def add_not_null(db, col):
1255
    table = col.table
1256
    col = sql_gen.to_name_only_col(col)
1257
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1258
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1259

    
1260
def drop_not_null(db, col):
1261
    table = col.table
1262
    col = sql_gen.to_name_only_col(col)
1263
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1264
        +col.to_str(db)+' DROP NOT NULL', cacheable=True, log_level=3)
1265

    
1266
row_num_col = '_row_num'
1267

    
1268
row_num_col_def = sql_gen.TypedCol('', 'serial', nullable=False,
1269
    constraints='PRIMARY KEY')
1270

    
1271
def add_row_num(db, table, name=row_num_col):
1272
    '''Adds a row number column to a table. Its definition is in
1273
    row_num_col_def. It will be the primary key.'''
1274
    col_def = copy.copy(row_num_col_def)
1275
    col_def.name = name
1276
    add_col(db, table, col_def, comment='', if_not_exists=True, log_level=3)
1277

    
1278
#### Indexes
1279

    
1280
def add_pkey(db, table, cols=None, recover=None):
1281
    '''Adds a primary key.
1282
    @param cols [sql_gen.Col,...] The columns in the primary key.
1283
        Defaults to the first column in the table.
1284
    @pre The table must not already have a primary key.
1285
    '''
1286
    table = sql_gen.as_Table(table)
1287
    if cols == None: cols = [pkey_name(db, table, recover)]
1288
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1289
    
1290
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1291
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1292
        log_ignore_excs=(DuplicateException,))
1293

    
1294
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1295
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1296
    Currently, only function calls and literal values are supported expressions.
1297
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1298
        This allows indexes to be used for comparisons where NULLs are equal.
1299
    '''
1300
    exprs = lists.mk_seq(exprs)
1301
    
1302
    # Parse exprs
1303
    old_exprs = exprs[:]
1304
    exprs = []
1305
    cols = []
1306
    for i, expr in enumerate(old_exprs):
1307
        expr = sql_gen.as_Col(expr, table)
1308
        
1309
        # Handle nullable columns
1310
        if ensure_not_null_:
1311
            try: expr = sql_gen.ensure_not_null(db, expr)
1312
            except KeyError: pass # unknown type, so just create plain index
1313
        
1314
        # Extract col
1315
        expr = copy.deepcopy(expr) # don't modify input!
1316
        col = expr
1317
        if isinstance(expr, sql_gen.FunctionCall): col = expr.args[0]
1318
        expr = sql_gen.cast_literal(expr)
1319
        if not isinstance(expr, (sql_gen.Expr, sql_gen.Col)):
1320
            expr = sql_gen.Expr(expr)
1321
            
1322
        
1323
        # Extract table
1324
        if table == None:
1325
            assert sql_gen.is_table_col(col)
1326
            table = col.table
1327
        
1328
        if isinstance(col, sql_gen.Col): col.table = None
1329
        
1330
        exprs.append(expr)
1331
        cols.append(col)
1332
    
1333
    table = sql_gen.as_Table(table)
1334
    
1335
    # Add index
1336
    str_ = 'CREATE'
1337
    if unique: str_ += ' UNIQUE'
1338
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1339
        ', '.join((v.to_str(db) for v in exprs)))+')'
1340
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1341

    
1342
already_indexed = object() # tells add_indexes() the pkey has already been added
1343

    
1344
def add_indexes(db, table, has_pkey=True):
1345
    '''Adds an index on all columns in a table.
1346
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1347
        index should be added on the first column.
1348
        * If already_indexed, the pkey is assumed to have already been added
1349
    '''
1350
    cols = table_col_names(db, table)
1351
    if has_pkey:
1352
        if has_pkey is not already_indexed: add_pkey(db, table)
1353
        cols = cols[1:]
1354
    for col in cols: add_index(db, col, table)
1355

    
1356
#### Tables
1357

    
1358
### Maintenance
1359

    
1360
def analyze(db, table):
1361
    table = sql_gen.as_Table(table)
1362
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1363

    
1364
def autoanalyze(db, table):
1365
    if db.autoanalyze: analyze(db, table)
1366

    
1367
def vacuum(db, table):
1368
    table = sql_gen.as_Table(table)
1369
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1370
        log_level=3))
1371

    
1372
### Lifecycle
1373

    
1374
def drop(db, type_, name):
1375
    name = sql_gen.as_Name(name)
1376
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1377

    
1378
def drop_table(db, table): drop(db, 'TABLE', table)
1379

    
1380
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1381
    like=None):
1382
    '''Creates a table.
1383
    @param cols [sql_gen.TypedCol,...] The column names and types
1384
    @param has_pkey If set, the first column becomes the primary key.
1385
    @param col_indexes bool|[ref]
1386
        * If True, indexes will be added on all non-pkey columns.
1387
        * If a list reference, [0] will be set to a function to do this.
1388
          This can be used to delay index creation until the table is populated.
1389
    '''
1390
    table = sql_gen.as_Table(table)
1391
    
1392
    if like != None:
1393
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1394
            ]+cols
1395
        table.order_by = like.order_by
1396
    if has_pkey:
1397
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1398
        pkey.constraints = 'PRIMARY KEY'
1399
    
1400
    temp = table.is_temp and not db.debug_temp
1401
        # temp tables permanent in debug_temp mode
1402
    
1403
    # Create table
1404
    def create():
1405
        str_ = 'CREATE'
1406
        if temp: str_ += ' TEMP'
1407
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1408
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1409
        str_ += '\n);'
1410
        
1411
        run_query(db, str_, recover=True, cacheable=True, log_level=2,
1412
            log_ignore_excs=(DuplicateException,))
1413
    if table.is_temp:
1414
        while True:
1415
            try:
1416
                create()
1417
                break
1418
            except DuplicateException:
1419
                table.name = next_version(table.name)
1420
                # try again with next version of name
1421
    else: create()
1422
    
1423
    # Add indexes
1424
    if has_pkey: has_pkey = already_indexed
1425
    def add_indexes_(): add_indexes(db, table, has_pkey)
1426
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1427
    elif col_indexes: add_indexes_() # add now
1428

    
1429
def copy_table_struct(db, src, dest):
1430
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1431
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1432

    
1433
def copy_table(db, src, dest):
1434
    '''Creates a copy of a table, including data'''
1435
    copy_table_struct(db, src, dest)
1436
    insert_select(db, dest, None, mk_select(db, src))
1437

    
1438
### Data
1439

    
1440
def truncate(db, table, schema='public', **kw_args):
1441
    '''For params, see run_query()'''
1442
    table = sql_gen.as_Table(table, schema)
1443
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1444

    
1445
def empty_temp(db, tables):
1446
    tables = lists.mk_seq(tables)
1447
    for table in tables: truncate(db, table, log_level=3)
1448

    
1449
def empty_db(db, schema='public', **kw_args):
1450
    '''For kw_args, see tables()'''
1451
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1452

    
1453
def distinct_table(db, table, distinct_on):
1454
    '''Creates a copy of a temp table which is distinct on the given columns.
1455
    The old and new tables will both get an index on these columns, to
1456
    facilitate merge joins.
1457
    @param distinct_on If empty, creates a table with one row. This is useful if
1458
        your distinct_on columns are all literal values.
1459
    @return The new table.
1460
    '''
1461
    new_table = sql_gen.suffixed_table(table, '_distinct')
1462
    distinct_on = filter(sql_gen.is_table_col, distinct_on)
1463
    
1464
    copy_table_struct(db, table, new_table)
1465
    
1466
    limit = None
1467
    if distinct_on == []: limit = 1 # one sample row
1468
    else:
1469
        add_index(db, distinct_on, new_table, unique=True)
1470
        add_index(db, distinct_on, table) # for join optimization
1471
    
1472
    insert_select(db, new_table, None, mk_select(db, table, order_by=None,
1473
        limit=limit), ignore=True)
1474
    analyze(db, new_table)
1475
    
1476
    return new_table
(28-28/42)