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: '+strings.as_tt(str(name)), cause)
41
        self.name = name
42

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

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

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

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

    
74
class NameException(DbException): pass
75

    
76
class DuplicateKeyException(ConstraintException): pass
77

    
78
class NullValueException(ConstraintException): pass
79

    
80
class CheckException(ConstraintException): pass
81

    
82
class InvalidValueException(ExceptionWithValue): pass
83

    
84
class DuplicateException(ExceptionWithNameType): pass
85

    
86
class DoesNotExistException(ExceptionWithNameType): pass
87

    
88
class EmptyRowException(DbException): pass
89

    
90
##### Warnings
91

    
92
class DbWarning(UserWarning): pass
93

    
94
##### Result retrieval
95

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

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

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

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

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

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

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

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

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

    
121
##### Escaping
122

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

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

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

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

    
141
##### Database connections
142

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

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

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

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

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

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

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

    
491
connect = DbConn
492

    
493
##### Recoverable querying
494

    
495
def with_savepoint(db, func): return db.with_savepoint(func)
496

    
497
def run_query(db, query, recover=None, cacheable=False, log_level=2,
498
    log_ignore_excs=None, **kw_args):
499
    '''For params, see DbConn.run_query()'''
500
    if recover == None: recover = False
501
    if log_ignore_excs == None: log_ignore_excs = ()
502
    log_ignore_excs = tuple(log_ignore_excs)
503
    debug_msg_ref = [None]
504
    
505
    query = with_explain_comment(db, query)
506
    
507
    try:
508
        try:
509
            def run(): return db.run_query(query, cacheable, log_level,
510
                debug_msg_ref, **kw_args)
511
            if recover and not db.is_cached(query):
512
                return with_savepoint(db, run)
513
            else: return run() # don't need savepoint if cached
514
        except Exception, e:
515
            msg = strings.ustr(e.args[0])
516
            msg = re.sub(r'^(?:PL/Python: )?ValueError: ', r'', msg)
517
            
518
            match = re.match(r'^duplicate key value violates unique constraint '
519
                r'"(.+?)"', msg)
520
            if match:
521
                constraint, = match.groups()
522
                cols = []
523
                if recover: # need auto-rollback to run index_cols()
524
                    try: cols = index_cols(db, constraint)
525
                    except NotImplementedError: pass
526
                raise DuplicateKeyException(constraint, None, cols, e)
527
            
528
            match = re.match(r'^null value in column "(.+?)" violates not-null'
529
                r' constraint', msg)
530
            if match:
531
                col, = match.groups()
532
                raise NullValueException('NOT NULL', None, [col], e)
533
            
534
            match = re.match(r'^new row for relation "(.+?)" violates check '
535
                r'constraint "(.+?)"', msg)
536
            if match:
537
                table, constraint = match.groups()
538
                constraint = sql_gen.Col(constraint, table)
539
                cond = None
540
                if recover: # need auto-rollback to run constraint_cond()
541
                    try: cond = constraint_cond(db, constraint)
542
                    except NotImplementedError: pass
543
                raise CheckException(constraint.to_str(db), cond, [], e)
544
            
545
            match = re.match(r'^(?:invalid input (?:syntax|value)\b.*?'
546
                r'|.+? out of range): "(.+?)"', msg)
547
            if match:
548
                value, = match.groups()
549
                raise InvalidValueException(strings.to_unicode(value), e)
550
            
551
            match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
552
                r'is of type', msg)
553
            if match:
554
                col, type_ = match.groups()
555
                raise MissingCastException(type_, col, e)
556
            
557
            match = re.match(r'^could not determine polymorphic type because '
558
                r'input has type "unknown"', msg)
559
            if match: raise MissingCastException('text', None, e)
560
            
561
            typed_name_re = r'^(\S+) "(.+?)"'
562
            
563
            match = re.match(typed_name_re+r'.*? already exists', msg)
564
            if match:
565
                type_, name = match.groups()
566
                raise DuplicateException(type_, name, e)
567
            
568
            match = re.match(r'more than one (\S+) named ""(.+?)""', msg)
569
            if match:
570
                type_, name = match.groups()
571
                raise DuplicateException(type_, name, e)
572
            
573
            match = re.match(typed_name_re+r' does not exist', msg)
574
            if match:
575
                type_, name = match.groups()
576
                raise DoesNotExistException(type_, name, e)
577
            
578
            raise # no specific exception raised
579
    except log_ignore_excs:
580
        log_level += 2
581
        raise
582
    finally:
583
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
584

    
585
##### Basic queries
586

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

    
592
def explain(db, query, **kw_args):
593
    '''
594
    For params, see run_query().
595
    '''
596
    kw_args.setdefault('log_level', 4)
597
    
598
    return strings.ustr(strings.join_lines(values(run_query(db,
599
        'EXPLAIN '+query, recover=True, cacheable=True, **kw_args))))
600
        # not a higher log_level because it's useful to see what query is being
601
        # run before it's executed, which EXPLAIN effectively provides
602

    
603
def has_comment(query): return query.endswith('*/')
604

    
605
def with_explain_comment(db, query, **kw_args):
606
    if db.autoexplain and not has_comment(query) and is_explainable(query):
607
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
608
            +explain(db, query, **kw_args))
609
    return query
610

    
611
def next_version(name):
612
    version = 1 # first existing name was version 0
613
    match = re.match(r'^(.*)#(\d+)$', name)
614
    if match:
615
        name, version = match.groups()
616
        version = int(version)+1
617
    return sql_gen.concat(name, '#'+str(version))
618

    
619
def lock_table(db, table, mode):
620
    table = sql_gen.as_Table(table)
621
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
622

    
623
def run_query_into(db, query, into=None, add_pkey_=False, **kw_args):
624
    '''Outputs a query to a temp table.
625
    For params, see run_query().
626
    '''
627
    if into == None: return run_query(db, query, **kw_args)
628
    
629
    assert isinstance(into, sql_gen.Table)
630
    
631
    into.is_temp = True
632
    # "temporary tables cannot specify a schema name", so remove schema
633
    into.schema = None
634
    
635
    kw_args['recover'] = True
636
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
637
    
638
    temp = not db.debug_temp # tables are permanent in debug_temp mode
639
    
640
    # Create table
641
    while True:
642
        create_query = 'CREATE'
643
        if temp: create_query += ' TEMP'
644
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
645
        
646
        try:
647
            cur = run_query(db, create_query, **kw_args)
648
                # CREATE TABLE AS sets rowcount to # rows in query
649
            break
650
        except DuplicateException, e:
651
            into.name = next_version(into.name)
652
            # try again with next version of name
653
    
654
    if add_pkey_: add_pkey(db, into)
655
    
656
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
657
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
658
    # table is going to be used in complex queries, it is wise to run ANALYZE on
659
    # the temporary table after it is populated."
660
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
661
    # If into is not a temp table, ANALYZE is useful but not required.
662
    analyze(db, into)
663
    
664
    return cur
665

    
666
order_by_pkey = object() # tells mk_select() to order by the pkey
667

    
668
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
669

    
670
def mk_select(db, tables=None, fields=None, conds=None, distinct_on=[],
671
    limit=None, start=None, order_by=order_by_pkey, default_table=None,
672
    explain=True):
673
    '''
674
    @param tables The single table to select from, or a list of tables to join
675
        together, with tables after the first being sql_gen.Join objects
676
    @param fields Use None to select all fields in the table
677
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
678
        * container can be any iterable type
679
        * compare_left_side: sql_gen.Code|str (for col name)
680
        * compare_right_side: sql_gen.ValueCond|literal value
681
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
682
        use all columns
683
    @return query
684
    '''
685
    # Parse tables param
686
    tables = lists.mk_seq(tables)
687
    tables = list(tables) # don't modify input! (list() copies input)
688
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
689
    
690
    # Parse other params
691
    if conds == None: conds = []
692
    elif dicts.is_dict(conds): conds = conds.items()
693
    conds = list(conds) # don't modify input! (list() copies input)
694
    assert limit == None or isinstance(limit, (int, long))
695
    assert start == None or isinstance(start, (int, long))
696
    if order_by is order_by_pkey:
697
        if table0 == None or distinct_on != []: order_by = None
698
        else: order_by = pkey(db, table0, recover=True)
699
    
700
    query = 'SELECT'
701
    
702
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
703
    
704
    # DISTINCT ON columns
705
    if distinct_on != []:
706
        query += '\nDISTINCT'
707
        if distinct_on is not distinct_on_all:
708
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
709
    
710
    # Columns
711
    if query.find('\n') >= 0: whitespace = '\n'
712
    else: whitespace = ' '
713
    if fields == None: query += whitespace+'*'
714
    else:
715
        assert fields != []
716
        if len(fields) > 1: whitespace = '\n'
717
        query += whitespace+('\n, '.join(map(parse_col, fields)))
718
    
719
    # Main table
720
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
721
    else: whitespace = ' '
722
    if table0 != None: query += whitespace+'FROM '+table0.to_str(db)
723
    
724
    # Add joins
725
    left_table = table0
726
    for join_ in tables:
727
        table = join_.table
728
        
729
        # Parse special values
730
        if join_.type_ is sql_gen.filter_out: # filter no match
731
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
732
                sql_gen.CompareCond(None, '~=')))
733
        
734
        query += '\n'+join_.to_str(db, left_table)
735
        
736
        left_table = table
737
    
738
    missing = True
739
    if conds != []:
740
        if len(conds) == 1: whitespace = ' '
741
        else: whitespace = '\n'
742
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
743
            .to_str(db) for l, r in conds], 'WHERE')
744
    if order_by != None:
745
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
746
    if limit != None: query += '\nLIMIT '+str(limit)
747
    if start != None:
748
        if start != 0: query += '\nOFFSET '+str(start)
749
    
750
    if explain: query = with_explain_comment(db, query)
751
    
752
    return query
753

    
754
def select(db, *args, **kw_args):
755
    '''For params, see mk_select() and run_query()'''
756
    recover = kw_args.pop('recover', None)
757
    cacheable = kw_args.pop('cacheable', True)
758
    log_level = kw_args.pop('log_level', 2)
759
    
760
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
761
        log_level=log_level)
762

    
763
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
764
    embeddable=False, ignore=False, src=None):
765
    '''
766
    @param returning str|None An inserted column (such as pkey) to return
767
    @param embeddable Whether the query should be embeddable as a nested SELECT.
768
        Warning: If you set this and cacheable=True when the query is run, the
769
        query will be fully cached, not just if it raises an exception.
770
    @param ignore Whether to ignore duplicate keys.
771
    @param src Will be included in the name of any created function, to help
772
        identify the data source in pg_stat_activity.
773
    '''
774
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
775
    if cols == []: cols = None # no cols (all defaults) = unknown col names
776
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
777
    if select_query == None: select_query = 'DEFAULT VALUES'
778
    if returning != None: returning = sql_gen.as_Col(returning, table)
779
    
780
    first_line = 'INSERT INTO '+table.to_str(db)
781
    
782
    def mk_insert(select_query):
783
        query = first_line
784
        if cols != None:
785
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
786
        query += '\n'+select_query
787
        
788
        if returning != None:
789
            returning_name_col = sql_gen.to_name_only_col(returning)
790
            query += '\nRETURNING '+returning_name_col.to_str(db)
791
        
792
        return query
793
    
794
    return_type = sql_gen.CustomCode('unknown')
795
    if returning != None: return_type = sql_gen.ColType(returning)
796
    
797
    if ignore:
798
        # Always return something to set the correct rowcount
799
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
800
        
801
        embeddable = True # must use function
802
        
803
        if cols == None: row = [sql_gen.Col(sql_gen.all_cols, 'row')]
804
        else: row = [sql_gen.Col(c.name, 'row') for c in cols]
805
        
806
        query = sql_gen.RowExcIgnore(sql_gen.RowType(table), select_query,
807
            sql_gen.ReturnQuery(mk_insert(sql_gen.Values(row).to_str(db))),
808
            cols)
809
    else: query = mk_insert(select_query)
810
    
811
    if embeddable:
812
        # Create function
813
        function_name = sql_gen.clean_name(first_line)
814
        if src != None: function_name = src+': '+function_name
815
        while True:
816
            try:
817
                func = db.TempFunction(function_name)
818
                def_ = sql_gen.FunctionDef(func, sql_gen.SetOf(return_type),
819
                    query)
820
                
821
                run_query(db, def_.to_str(db), recover=True, cacheable=True,
822
                    log_ignore_excs=(DuplicateException,))
823
                break # this version was successful
824
            except DuplicateException, e:
825
                function_name = next_version(function_name)
826
                # try again with next version of name
827
        
828
        # Return query that uses function
829
        cols = None
830
        if returning != None: cols = [returning]
831
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(func), cols)
832
            # AS clause requires function alias
833
        return mk_select(db, func_table, order_by=None)
834
    
835
    return query
836

    
837
def insert_select(db, table, *args, **kw_args):
838
    '''For params, see mk_insert_select() and run_query_into()
839
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
840
        values in
841
    '''
842
    returning = kw_args.get('returning', None)
843
    ignore = kw_args.get('ignore', False)
844
    
845
    into = kw_args.pop('into', None)
846
    if into != None: kw_args['embeddable'] = True
847
    recover = kw_args.pop('recover', None)
848
    if ignore: recover = True
849
    cacheable = kw_args.pop('cacheable', True)
850
    log_level = kw_args.pop('log_level', 2)
851
    
852
    rowcount_only = ignore and returning == None # keep NULL rows on server
853
    if rowcount_only: into = sql_gen.Table('rowcount')
854
    
855
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
856
        into, recover=recover, cacheable=cacheable, log_level=log_level)
857
    if rowcount_only: empty_temp(db, into)
858
    autoanalyze(db, table)
859
    return cur
860

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

    
863
def insert(db, table, row, *args, **kw_args):
864
    '''For params, see insert_select()'''
865
    if lists.is_seq(row): cols = None
866
    else:
867
        cols = row.keys()
868
        row = row.values()
869
    row = list(row) # ensure that "== []" works
870
    
871
    if row == []: query = None
872
    else: query = sql_gen.Values(row).to_str(db)
873
    
874
    return insert_select(db, table, cols, query, *args, **kw_args)
875

    
876
def mk_update(db, table, changes=None, cond=None, in_place=False,
877
    cacheable_=True):
878
    '''
879
    @param changes [(col, new_value),...]
880
        * container can be any iterable type
881
        * col: sql_gen.Code|str (for col name)
882
        * new_value: sql_gen.Code|literal value
883
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
884
    @param in_place If set, locks the table and updates rows in place.
885
        This avoids creating dead rows in PostgreSQL.
886
        * cond must be None
887
    @param cacheable_ Whether column structure information used to generate the
888
        query can be cached
889
    @return str query
890
    '''
891
    table = sql_gen.as_Table(table)
892
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
893
        for c, v in changes]
894
    
895
    if in_place:
896
        assert cond == None
897
        
898
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
899
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '
900
            +db.col_info(sql_gen.with_default_table(c, table), cacheable_).type
901
            +'\nUSING '+v.to_str(db) for c, v in changes))
902
    else:
903
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
904
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
905
            for c, v in changes))
906
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
907
    
908
    query = with_explain_comment(db, query)
909
    
910
    return query
911

    
912
def update(db, table, *args, **kw_args):
913
    '''For params, see mk_update() and run_query()'''
914
    recover = kw_args.pop('recover', None)
915
    cacheable = kw_args.pop('cacheable', False)
916
    log_level = kw_args.pop('log_level', 2)
917
    
918
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
919
        cacheable, log_level=log_level)
920
    autoanalyze(db, table)
921
    return cur
922

    
923
def mk_delete(db, table, cond=None):
924
    '''
925
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
926
    @return str query
927
    '''
928
    query = 'DELETE FROM '+table.to_str(db)
929
    if cond != None: query += '\nWHERE '+cond.to_str(db)
930
    
931
    query = with_explain_comment(db, query)
932
    
933
    return query
934

    
935
def delete(db, table, *args, **kw_args):
936
    '''For params, see mk_delete() and run_query()'''
937
    recover = kw_args.pop('recover', None)
938
    cacheable = kw_args.pop('cacheable', True)
939
    log_level = kw_args.pop('log_level', 2)
940
    
941
    cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
942
        cacheable, log_level=log_level)
943
    autoanalyze(db, table)
944
    return cur
945

    
946
def last_insert_id(db):
947
    module = util.root_module(db.db)
948
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
949
    elif module == 'MySQLdb': return db.insert_id()
950
    else: return None
951

    
952
def define_func(db, def_):
953
    func = def_.function
954
    while True:
955
        try:
956
            run_query(db, def_.to_str(db), recover=True, cacheable=True,
957
                log_ignore_excs=(DuplicateException,))
958
            break # successful
959
        except DuplicateException:
960
            func.name = next_version(func.name)
961
            # try again with next version of name
962

    
963
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
964
    '''Creates a mapping from original column names (which may have collisions)
965
    to names that will be distinct among the columns' tables.
966
    This is meant to be used for several tables that are being joined together.
967
    @param cols The columns to combine. Duplicates will be removed.
968
    @param into The table for the new columns.
969
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
970
        columns will be included in the mapping even if they are not in cols.
971
        The tables of the provided Col objects will be changed to into, so make
972
        copies of them if you want to keep the original tables.
973
    @param as_items Whether to return a list of dict items instead of a dict
974
    @return dict(orig_col=new_col, ...)
975
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
976
        * new_col: sql_gen.Col(orig_col_name, into)
977
        * All mappings use the into table so its name can easily be
978
          changed for all columns at once
979
    '''
980
    cols = lists.uniqify(cols)
981
    
982
    items = []
983
    for col in preserve:
984
        orig_col = copy.copy(col)
985
        col.table = into
986
        items.append((orig_col, col))
987
    preserve = set(preserve)
988
    for col in cols:
989
        if col not in preserve:
990
            items.append((col, sql_gen.Col(strings.ustr(col), into, col.srcs)))
991
    
992
    if not as_items: items = dict(items)
993
    return items
994

    
995
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
996
    '''For params, see mk_flatten_mapping()
997
    @return See return value of mk_flatten_mapping()
998
    '''
999
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
1000
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
1001
    run_query_into(db, mk_select(db, joins, cols, order_by=None, limit=limit,
1002
        start=start), into=into, add_pkey_=True)
1003
        # don't cache because the temp table will usually be truncated after use
1004
    return dict(items)
1005

    
1006
##### Database structure introspection
1007

    
1008
#### Expressions
1009

    
1010
bool_re = r'(?:true|false)'
1011

    
1012
def simplify_expr(expr):
1013
    expr = expr.replace('(NULL IS NULL)', 'true')
1014
    expr = expr.replace('(NULL IS NOT NULL)', 'false')
1015
    expr = re.sub(r' OR '+bool_re, r'', expr)
1016
    expr = re.sub(bool_re+r' OR ', r'', expr)
1017
    while True:
1018
        expr, n = re.subn(r'\((\([^()]*\))\)', r'\1', expr)
1019
        if n == 0: break
1020
    return expr
1021

    
1022
name_re = r'(?:\w+|(?:"[^"]*")+)'
1023

    
1024
def parse_expr_col(str_):
1025
    match = re.match(r'^\('+name_re+r'\(('+name_re+r').*\)\)$', str_)
1026
    if match: str_ = match.group(1)
1027
    return sql_gen.unesc_name(str_)
1028

    
1029
def map_expr(db, expr, mapping, in_cols_found=None):
1030
    '''Replaces output columns with input columns in an expression.
1031
    @param in_cols_found If set, will be filled in with the expr's (input) cols
1032
    '''
1033
    for out, in_ in mapping.iteritems():
1034
        orig_expr = expr
1035
        out = sql_gen.to_name_only_col(out)
1036
        in_str = sql_gen.to_name_only_col(sql_gen.remove_col_rename(in_)
1037
            ).to_str(db)
1038
        
1039
        # Replace out both with and without quotes
1040
        expr = expr.replace(out.to_str(db), in_str)
1041
        expr = re.sub(r'\b'+out.name+r'\b', in_str, expr)
1042
        
1043
        if in_cols_found != None and expr != orig_expr: # replaced something
1044
            in_cols_found.append(in_)
1045
    
1046
    return simplify_expr(expr)
1047

    
1048
#### Tables
1049

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

    
1065
def table_exists(db, table):
1066
    table = sql_gen.as_Table(table)
1067
    return list(tables(db, table.schema, table.name, exact=True)) != []
1068

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

    
1073
def table_cols(db, table, recover=None):
1074
    return list(col_names(select(db, table, limit=0, order_by=None,
1075
        recover=recover, log_level=4)))
1076

    
1077
def pkey(db, table, recover=None):
1078
    '''Assumed to be first column in table'''
1079
    return table_cols(db, table, recover)[0]
1080

    
1081
not_null_col = 'not_null_col'
1082

    
1083
def table_not_null_col(db, table, recover=None):
1084
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1085
    if not_null_col in table_cols(db, table, recover): return not_null_col
1086
    else: return pkey(db, table, recover)
1087

    
1088
def constraint_cond(db, constraint):
1089
    module = util.root_module(db.db)
1090
    if module == 'psycopg2':
1091
        table_str = sql_gen.Literal(constraint.table.to_str(db))
1092
        name_str = sql_gen.Literal(constraint.name)
1093
        return value(run_query(db, '''\
1094
SELECT consrc
1095
FROM pg_constraint
1096
WHERE
1097
conrelid = '''+table_str.to_str(db)+'''::regclass
1098
AND conname = '''+name_str.to_str(db)+'''
1099
'''
1100
            , cacheable=True, log_level=4))
1101
    else: raise NotImplementedError("Can't list index columns for "+module+
1102
        ' database')
1103

    
1104
def index_cols(db, index):
1105
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1106
    automatically created. When you don't know whether something is a UNIQUE
1107
    constraint or a UNIQUE index, use this function.'''
1108
    index = sql_gen.as_Table(index)
1109
    module = util.root_module(db.db)
1110
    if module == 'psycopg2':
1111
        qual_index = sql_gen.Literal(index.to_str(db))
1112
        return map(parse_expr_col, values(run_query(db, '''\
1113
SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
1114
FROM pg_index
1115
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1116
'''
1117
            , cacheable=True, log_level=4)))
1118
    else: raise NotImplementedError("Can't list index columns for "+module+
1119
        ' database')
1120

    
1121
#### Functions
1122

    
1123
def function_exists(db, function):
1124
    qual_function = sql_gen.Literal(function.to_str(db))
1125
    try:
1126
        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
1127
            recover=True, cacheable=True, log_level=4)
1128
    except DoesNotExistException: return False
1129
    else: return True
1130

    
1131
##### Structural changes
1132

    
1133
#### Columns
1134

    
1135
def add_col(db, table, col, comment=None, **kw_args):
1136
    '''
1137
    @param col TypedCol Name may be versioned, so be sure to propagate any
1138
        renaming back to any source column for the TypedCol.
1139
    @param comment None|str SQL comment used to distinguish columns of the same
1140
        name from each other when they contain different data, to allow the
1141
        ADD COLUMN query to be cached. If not set, query will not be cached.
1142
    '''
1143
    assert isinstance(col, sql_gen.TypedCol)
1144
    
1145
    while True:
1146
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1147
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1148
        
1149
        try:
1150
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1151
            break
1152
        except DuplicateException:
1153
            col.name = next_version(col.name)
1154
            # try again with next version of name
1155

    
1156
def add_not_null(db, col):
1157
    table = col.table
1158
    col = sql_gen.to_name_only_col(col)
1159
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1160
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1161

    
1162
row_num_col = '_row_num'
1163

    
1164
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1165
    constraints='PRIMARY KEY')
1166

    
1167
def add_row_num(db, table):
1168
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1169
    be the primary key.'''
1170
    add_col(db, table, row_num_typed_col, log_level=3)
1171

    
1172
#### Indexes
1173

    
1174
def add_pkey(db, table, cols=None, recover=None):
1175
    '''Adds a primary key.
1176
    @param cols [sql_gen.Col,...] The columns in the primary key.
1177
        Defaults to the first column in the table.
1178
    @pre The table must not already have a primary key.
1179
    '''
1180
    table = sql_gen.as_Table(table)
1181
    if cols == None: cols = [pkey(db, table, recover)]
1182
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1183
    
1184
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1185
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1186
        log_ignore_excs=(DuplicateException,))
1187

    
1188
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1189
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1190
    Currently, only function calls and literal values are supported expressions.
1191
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1192
        This allows indexes to be used for comparisons where NULLs are equal.
1193
    '''
1194
    exprs = lists.mk_seq(exprs)
1195
    
1196
    # Parse exprs
1197
    old_exprs = exprs[:]
1198
    exprs = []
1199
    cols = []
1200
    for i, expr in enumerate(old_exprs):
1201
        expr = sql_gen.as_Col(expr, table)
1202
        
1203
        # Handle nullable columns
1204
        if ensure_not_null_:
1205
            try: expr = sql_gen.ensure_not_null(db, expr)
1206
            except KeyError: pass # unknown type, so just create plain index
1207
        
1208
        # Extract col
1209
        expr = copy.deepcopy(expr) # don't modify input!
1210
        col = expr
1211
        if isinstance(expr, sql_gen.FunctionCall): col = expr.args[0]
1212
        expr = sql_gen.cast_literal(expr)
1213
        if not isinstance(expr, (sql_gen.Expr, sql_gen.Col)):
1214
            expr = sql_gen.Expr(expr)
1215
            
1216
        
1217
        # Extract table
1218
        if table == None:
1219
            assert sql_gen.is_table_col(col)
1220
            table = col.table
1221
        
1222
        if isinstance(col, sql_gen.Col): col.table = None
1223
        
1224
        exprs.append(expr)
1225
        cols.append(col)
1226
    
1227
    table = sql_gen.as_Table(table)
1228
    
1229
    # Add index
1230
    str_ = 'CREATE'
1231
    if unique: str_ += ' UNIQUE'
1232
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1233
        ', '.join((v.to_str(db) for v in exprs)))+')'
1234
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1235

    
1236
already_indexed = object() # tells add_indexes() the pkey has already been added
1237

    
1238
def add_indexes(db, table, has_pkey=True):
1239
    '''Adds an index on all columns in a table.
1240
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1241
        index should be added on the first column.
1242
        * If already_indexed, the pkey is assumed to have already been added
1243
    '''
1244
    cols = table_cols(db, table)
1245
    if has_pkey:
1246
        if has_pkey is not already_indexed: add_pkey(db, table)
1247
        cols = cols[1:]
1248
    for col in cols: add_index(db, col, table)
1249

    
1250
#### Tables
1251

    
1252
### Maintenance
1253

    
1254
def analyze(db, table):
1255
    table = sql_gen.as_Table(table)
1256
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1257

    
1258
def autoanalyze(db, table):
1259
    if db.autoanalyze: analyze(db, table)
1260

    
1261
def vacuum(db, table):
1262
    table = sql_gen.as_Table(table)
1263
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1264
        log_level=3))
1265

    
1266
### Lifecycle
1267

    
1268
def drop(db, type_, name):
1269
    name = sql_gen.as_Name(name)
1270
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1271

    
1272
def drop_table(db, table): drop(db, 'TABLE', table)
1273

    
1274
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1275
    like=None):
1276
    '''Creates a table.
1277
    @param cols [sql_gen.TypedCol,...] The column names and types
1278
    @param has_pkey If set, the first column becomes the primary key.
1279
    @param col_indexes bool|[ref]
1280
        * If True, indexes will be added on all non-pkey columns.
1281
        * If a list reference, [0] will be set to a function to do this.
1282
          This can be used to delay index creation until the table is populated.
1283
    '''
1284
    table = sql_gen.as_Table(table)
1285
    
1286
    if like != None:
1287
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1288
            ]+cols
1289
    if has_pkey:
1290
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1291
        pkey.constraints = 'PRIMARY KEY'
1292
    
1293
    temp = table.is_temp and not db.debug_temp
1294
        # temp tables permanent in debug_temp mode
1295
    
1296
    # Create table
1297
    def create():
1298
        str_ = 'CREATE'
1299
        if temp: str_ += ' TEMP'
1300
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1301
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1302
        str_ += '\n);'
1303
        
1304
        run_query(db, str_, recover=True, cacheable=True, log_level=2,
1305
            log_ignore_excs=(DuplicateException,))
1306
    if table.is_temp:
1307
        while True:
1308
            try:
1309
                create()
1310
                break
1311
            except DuplicateException:
1312
                table.name = next_version(table.name)
1313
                # try again with next version of name
1314
    else: create()
1315
    
1316
    # Add indexes
1317
    if has_pkey: has_pkey = already_indexed
1318
    def add_indexes_(): add_indexes(db, table, has_pkey)
1319
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1320
    elif col_indexes: add_indexes_() # add now
1321

    
1322
def copy_table_struct(db, src, dest):
1323
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1324
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1325

    
1326
### Data
1327

    
1328
def truncate(db, table, schema='public', **kw_args):
1329
    '''For params, see run_query()'''
1330
    table = sql_gen.as_Table(table, schema)
1331
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1332

    
1333
def empty_temp(db, tables):
1334
    tables = lists.mk_seq(tables)
1335
    for table in tables: truncate(db, table, log_level=3)
1336

    
1337
def empty_db(db, schema='public', **kw_args):
1338
    '''For kw_args, see tables()'''
1339
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1340

    
1341
def distinct_table(db, table, distinct_on):
1342
    '''Creates a copy of a temp table which is distinct on the given columns.
1343
    The old and new tables will both get an index on these columns, to
1344
    facilitate merge joins.
1345
    @param distinct_on If empty, creates a table with one row. This is useful if
1346
        your distinct_on columns are all literal values.
1347
    @return The new table.
1348
    '''
1349
    new_table = sql_gen.suffixed_table(table, '_distinct')
1350
    distinct_on = filter(sql_gen.is_table_col, distinct_on)
1351
    
1352
    copy_table_struct(db, table, new_table)
1353
    
1354
    limit = None
1355
    if distinct_on == []: limit = 1 # one sample row
1356
    else:
1357
        add_index(db, distinct_on, new_table, unique=True)
1358
        add_index(db, distinct_on, table) # for join optimization
1359
    
1360
    insert_select(db, new_table, None, mk_select(db, table, order_by=None,
1361
        limit=limit), ignore=True)
1362
    analyze(db, new_table)
1363
    
1364
    return new_table
(27-27/40)