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, cols, cause=None):
58
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
59
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
60
        self.name = name
61
        self.cols = cols
62

    
63
class MissingCastException(DbException):
64
    def __init__(self, type_, col, cause=None):
65
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
66
            +' on column: '+strings.as_tt(col), cause)
67
        self.type = type_
68
        self.col = col
69

    
70
class NameException(DbException): pass
71

    
72
class DuplicateKeyException(ConstraintException): pass
73

    
74
class NullValueException(ConstraintException): pass
75

    
76
class InvalidValueException(ExceptionWithValue): pass
77

    
78
class DuplicateException(ExceptionWithNameType): pass
79

    
80
class EmptyRowException(DbException): pass
81

    
82
##### Warnings
83

    
84
class DbWarning(UserWarning): pass
85

    
86
##### Result retrieval
87

    
88
def col_names(cur): return (col[0] for col in cur.description)
89

    
90
def rows(cur): return iter(lambda: cur.fetchone(), None)
91

    
92
def consume_rows(cur):
93
    '''Used to fetch all rows so result will be cached'''
94
    iters.consume_iter(rows(cur))
95

    
96
def next_row(cur): return rows(cur).next()
97

    
98
def row(cur):
99
    row_ = next_row(cur)
100
    consume_rows(cur)
101
    return row_
102

    
103
def next_value(cur): return next_row(cur)[0]
104

    
105
def value(cur): return row(cur)[0]
106

    
107
def values(cur): return iters.func_iter(lambda: next_value(cur))
108

    
109
def value_or_none(cur):
110
    try: return value(cur)
111
    except StopIteration: return None
112

    
113
##### Escaping
114

    
115
def esc_name_by_module(module, name):
116
    if module == 'psycopg2' or module == None: quote = '"'
117
    elif module == 'MySQLdb': quote = '`'
118
    else: raise NotImplementedError("Can't escape name for "+module+' database')
119
    return sql_gen.esc_name(name, quote)
120

    
121
def esc_name_by_engine(engine, name, **kw_args):
122
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
123

    
124
def esc_name(db, name, **kw_args):
125
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
126

    
127
def qual_name(db, schema, table):
128
    def esc_name_(name): return esc_name(db, name)
129
    table = esc_name_(table)
130
    if schema != None: return esc_name_(schema)+'.'+table
131
    else: return table
132

    
133
##### Database connections
134

    
135
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
136

    
137
db_engines = {
138
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
139
    'PostgreSQL': ('psycopg2', {}),
140
}
141

    
142
DatabaseErrors_set = set([DbException])
143
DatabaseErrors = tuple(DatabaseErrors_set)
144

    
145
def _add_module(module):
146
    DatabaseErrors_set.add(module.DatabaseError)
147
    global DatabaseErrors
148
    DatabaseErrors = tuple(DatabaseErrors_set)
149

    
150
def db_config_str(db_config):
151
    return db_config['engine']+' database '+db_config['database']
152

    
153
log_debug_none = lambda msg, level=2: None
154

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

    
481
connect = DbConn
482

    
483
##### Recoverable querying
484

    
485
def with_savepoint(db, func): return db.with_savepoint(func)
486

    
487
def run_query(db, query, recover=None, cacheable=False, log_level=2,
488
    log_ignore_excs=None, **kw_args):
489
    '''For params, see DbConn.run_query()'''
490
    if recover == None: recover = False
491
    if log_ignore_excs == None: log_ignore_excs = ()
492
    log_ignore_excs = tuple(log_ignore_excs)
493
    debug_msg_ref = [None]
494
    
495
    query = with_explain_comment(db, query)
496
    
497
    try:
498
        try:
499
            def run(): return db.run_query(query, cacheable, log_level,
500
                debug_msg_ref, **kw_args)
501
            if recover and not db.is_cached(query):
502
                return with_savepoint(db, run)
503
            else: return run() # don't need savepoint if cached
504
        except Exception, e:
505
            msg = strings.ustr(e.args[0])
506
            
507
            match = re.match(r'^duplicate key value violates unique constraint '
508
                r'"((_?[^\W_]+(?=[._]))?.+?)"', msg)
509
            if match:
510
                constraint, table = match.groups()
511
                cols = []
512
                if recover: # need auto-rollback to run index_cols()
513
                    try: cols = index_cols(db, table, constraint)
514
                    except NotImplementedError: pass
515
                raise DuplicateKeyException(constraint, cols, e)
516
            
517
            match = re.match(r'^null value in column "(.+?)" violates not-null'
518
                r' constraint', msg)
519
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
520
            
521
            match = re.match(r'^(?:invalid input (?:syntax|value)\b.*?'
522
                r'|.+? field value out of range): "(.+?)"', msg)
523
            if match:
524
                value, = match.groups()
525
                raise InvalidValueException(strings.to_unicode(value), e)
526
            
527
            match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
528
                r'is of type', msg)
529
            if match:
530
                col, type_ = match.groups()
531
                raise MissingCastException(type_, col, e)
532
            
533
            match = re.match(r'^(\S+) "(.+?)".*? already exists', msg)
534
            if match:
535
                type_, name = match.groups()
536
                raise DuplicateException(type_, name, e)
537
            
538
            raise # no specific exception raised
539
    except log_ignore_excs:
540
        log_level += 2
541
        raise
542
    finally:
543
        if debug_msg_ref[0] != None: db.log_debug(debug_msg_ref[0], log_level)
544

    
545
##### Basic queries
546

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

    
552
def explain(db, query, **kw_args):
553
    '''
554
    For params, see run_query().
555
    '''
556
    kw_args.setdefault('log_level', 4)
557
    
558
    return strings.join_lines(values(run_query(db, 'EXPLAIN '+query,
559
        recover=True, cacheable=True, **kw_args)))
560
        # not a higher log_level because it's useful to see what query is being
561
        # run before it's executed, which EXPLAIN effectively provides
562

    
563
def has_comment(query): return query.endswith('*/')
564

    
565
def with_explain_comment(db, query, **kw_args):
566
    if db.autoexplain and not has_comment(query) and is_explainable(query):
567
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
568
            +explain(db, query, **kw_args))
569
    return query
570

    
571
def next_version(name):
572
    version = 1 # first existing name was version 0
573
    match = re.match(r'^(.*)#(\d+)$', name)
574
    if match:
575
        name, version = match.groups()
576
        version = int(version)+1
577
    return sql_gen.concat(name, '#'+str(version))
578

    
579
def lock_table(db, table, mode):
580
    table = sql_gen.as_Table(table)
581
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
582

    
583
def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):
584
    '''Outputs a query to a temp table.
585
    For params, see run_query().
586
    '''
587
    if into == None: return run_query(db, query, **kw_args)
588
    
589
    assert isinstance(into, sql_gen.Table)
590
    
591
    into.is_temp = True
592
    # "temporary tables cannot specify a schema name", so remove schema
593
    into.schema = None
594
    
595
    kw_args['recover'] = True
596
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
597
    
598
    temp = not db.debug_temp # tables are permanent in debug_temp mode
599
    
600
    # Create table
601
    while True:
602
        create_query = 'CREATE'
603
        if temp: create_query += ' TEMP'
604
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
605
        
606
        try:
607
            cur = run_query(db, create_query, **kw_args)
608
                # CREATE TABLE AS sets rowcount to # rows in query
609
            break
610
        except DuplicateException, e:
611
            into.name = next_version(into.name)
612
            # try again with next version of name
613
    
614
    if add_indexes_: add_indexes(db, into)
615
    
616
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
617
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
618
    # table is going to be used in complex queries, it is wise to run ANALYZE on
619
    # the temporary table after it is populated."
620
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
621
    # If into is not a temp table, ANALYZE is useful but not required.
622
    analyze(db, into)
623
    
624
    return cur
625

    
626
order_by_pkey = object() # tells mk_select() to order by the pkey
627

    
628
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
629

    
630
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
631
    start=None, order_by=order_by_pkey, default_table=None):
632
    '''
633
    @param tables The single table to select from, or a list of tables to join
634
        together, with tables after the first being sql_gen.Join objects
635
    @param fields Use None to select all fields in the table
636
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
637
        * container can be any iterable type
638
        * compare_left_side: sql_gen.Code|str (for col name)
639
        * compare_right_side: sql_gen.ValueCond|literal value
640
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
641
        use all columns
642
    @return query
643
    '''
644
    # Parse tables param
645
    tables = lists.mk_seq(tables)
646
    tables = list(tables) # don't modify input! (list() copies input)
647
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
648
    
649
    # Parse other params
650
    if conds == None: conds = []
651
    elif dicts.is_dict(conds): conds = conds.items()
652
    conds = list(conds) # don't modify input! (list() copies input)
653
    assert limit == None or isinstance(limit, (int, long))
654
    assert start == None or isinstance(start, (int, long))
655
    if order_by is order_by_pkey:
656
        if distinct_on != []: order_by = None
657
        else: order_by = pkey(db, table0, recover=True)
658
    
659
    query = 'SELECT'
660
    
661
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
662
    
663
    # DISTINCT ON columns
664
    if distinct_on != []:
665
        query += '\nDISTINCT'
666
        if distinct_on is not distinct_on_all:
667
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
668
    
669
    # Columns
670
    if query.find('\n') >= 0: whitespace = '\n'
671
    else: whitespace = ' '
672
    if fields == None: query += whitespace+'*'
673
    else:
674
        assert fields != []
675
        if len(fields) > 1: whitespace = '\n'
676
        query += whitespace+('\n, '.join(map(parse_col, fields)))
677
    
678
    # Main table
679
    if query.find('\n') >= 0 or len(tables) > 0: whitespace = '\n'
680
    else: whitespace = ' '
681
    query += whitespace+'FROM '+table0.to_str(db)
682
    
683
    # Add joins
684
    left_table = table0
685
    for join_ in tables:
686
        table = join_.table
687
        
688
        # Parse special values
689
        if join_.type_ is sql_gen.filter_out: # filter no match
690
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
691
                sql_gen.CompareCond(None, '~=')))
692
        
693
        query += '\n'+join_.to_str(db, left_table)
694
        
695
        left_table = table
696
    
697
    missing = True
698
    if conds != []:
699
        if len(conds) == 1: whitespace = ' '
700
        else: whitespace = '\n'
701
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
702
            .to_str(db) for l, r in conds], 'WHERE')
703
        missing = False
704
    if order_by != None:
705
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
706
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
707
    if start != None:
708
        if start != 0: query += '\nOFFSET '+str(start)
709
        missing = False
710
    if missing: warnings.warn(DbWarning(
711
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
712
    
713
    query = with_explain_comment(db, query)
714
    
715
    return query
716

    
717
def select(db, *args, **kw_args):
718
    '''For params, see mk_select() and run_query()'''
719
    recover = kw_args.pop('recover', None)
720
    cacheable = kw_args.pop('cacheable', True)
721
    log_level = kw_args.pop('log_level', 2)
722
    
723
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
724
        log_level=log_level)
725

    
726
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
727
    embeddable=False, ignore=False, src=None):
728
    '''
729
    @param returning str|None An inserted column (such as pkey) to return
730
    @param embeddable Whether the query should be embeddable as a nested SELECT.
731
        Warning: If you set this and cacheable=True when the query is run, the
732
        query will be fully cached, not just if it raises an exception.
733
    @param ignore Whether to ignore duplicate keys.
734
    @param src Will be included in the name of any created function, to help
735
        identify the data source in pg_stat_activity.
736
    '''
737
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
738
    if cols == []: cols = None # no cols (all defaults) = unknown col names
739
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
740
    if select_query == None: select_query = 'DEFAULT VALUES'
741
    if returning != None: returning = sql_gen.as_Col(returning, table)
742
    
743
    first_line = 'INSERT INTO '+table.to_str(db)
744
    
745
    def mk_insert(select_query):
746
        query = first_line
747
        if cols != None:
748
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
749
        query += '\n'+select_query
750
        
751
        if returning != None:
752
            returning_name_col = sql_gen.to_name_only_col(returning)
753
            query += '\nRETURNING '+returning_name_col.to_str(db)
754
        
755
        return query
756
    
757
    return_type = 'unknown'
758
    if returning != None: return_type = returning.to_str(db)+'%TYPE'
759
    
760
    lang = 'sql'
761
    if ignore:
762
        # Always return something to set the correct rowcount
763
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
764
        
765
        embeddable = True # must use function
766
        lang = 'plpgsql'
767
        
768
        if cols == None:
769
            row = [sql_gen.Col(sql_gen.all_cols, 'row')]
770
            row_vars = [sql_gen.Table('row')]
771
        else:
772
            row_vars = row = [sql_gen.Col(c.name, 'row') for c in cols]
773
        
774
        query = '''\
775
DECLARE
776
    row '''+table.to_str(db)+'''%ROWTYPE;
777
BEGIN
778
    /* Need an EXCEPTION block for each individual row because "When an error is
779
    caught by an EXCEPTION clause, [...] all changes to persistent database
780
    state within the block are rolled back."
781
    This is unfortunate because "A block containing an EXCEPTION clause is
782
    significantly more expensive to enter and exit than a block without one."
783
    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
784
#PLPGSQL-ERROR-TRAPPING)
785
    */
786
    FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
787
'''+select_query+'''
788
    LOOP
789
        BEGIN
790
            RETURN QUERY
791
'''+mk_insert(sql_gen.Values(row).to_str(db))+'''
792
;
793
        EXCEPTION
794
            WHEN unique_violation THEN NULL; -- continue to next row
795
        END;
796
    END LOOP;
797
END;\
798
'''
799
    else: query = mk_insert(select_query)
800
    
801
    if embeddable:
802
        # Create function
803
        function_name = sql_gen.clean_name(first_line)
804
        if src != None: function_name = src+': '+function_name
805
        while True:
806
            try:
807
                function = db.TempFunction(function_name)
808
                
809
                function_query = '''\
810
CREATE FUNCTION '''+function.to_str(db)+'''()
811
RETURNS SETOF '''+return_type+'''
812
LANGUAGE '''+lang+'''
813
AS $$
814
'''+query+'''
815
$$;
816
'''
817
                run_query(db, function_query, recover=True, cacheable=True,
818
                    log_ignore_excs=(DuplicateException,))
819
                break # this version was successful
820
            except DuplicateException, e:
821
                function_name = next_version(function_name)
822
                # try again with next version of name
823
        
824
        # Return query that uses function
825
        cols = None
826
        if returning != None: cols = [returning]
827
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
828
            cols) # AS clause requires function alias
829
        return mk_select(db, func_table, start=0, order_by=None)
830
    
831
    return query
832

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

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

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

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

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

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

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

    
942
def last_insert_id(db):
943
    module = util.root_module(db.db)
944
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
945
    elif module == 'MySQLdb': return db.insert_id()
946
    else: return None
947

    
948
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
949
    '''Creates a mapping from original column names (which may have collisions)
950
    to names that will be distinct among the columns' tables.
951
    This is meant to be used for several tables that are being joined together.
952
    @param cols The columns to combine. Duplicates will be removed.
953
    @param into The table for the new columns.
954
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
955
        columns will be included in the mapping even if they are not in cols.
956
        The tables of the provided Col objects will be changed to into, so make
957
        copies of them if you want to keep the original tables.
958
    @param as_items Whether to return a list of dict items instead of a dict
959
    @return dict(orig_col=new_col, ...)
960
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
961
        * new_col: sql_gen.Col(orig_col_name, into)
962
        * All mappings use the into table so its name can easily be
963
          changed for all columns at once
964
    '''
965
    cols = lists.uniqify(cols)
966
    
967
    items = []
968
    for col in preserve:
969
        orig_col = copy.copy(col)
970
        col.table = into
971
        items.append((orig_col, col))
972
    preserve = set(preserve)
973
    for col in cols:
974
        if col not in preserve:
975
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
976
    
977
    if not as_items: items = dict(items)
978
    return items
979

    
980
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
981
    '''For params, see mk_flatten_mapping()
982
    @return See return value of mk_flatten_mapping()
983
    '''
984
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
985
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
986
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
987
        into=into)
988
    return dict(items)
989

    
990
##### Database structure introspection
991

    
992
#### Tables
993

    
994
def tables(db, schema_like='public', table_like='%', exact=False):
995
    if exact: compare = '='
996
    else: compare = 'LIKE'
997
    
998
    module = util.root_module(db.db)
999
    if module == 'psycopg2':
1000
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1001
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1002
        return values(select(db, 'pg_tables', ['tablename'], conds,
1003
            order_by='tablename', log_level=4))
1004
    elif module == 'MySQLdb':
1005
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1006
            , cacheable=True, log_level=4))
1007
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1008

    
1009
def table_exists(db, table):
1010
    table = sql_gen.as_Table(table)
1011
    return list(tables(db, table.schema, table.name, exact=True)) != []
1012

    
1013
def table_row_count(db, table, recover=None):
1014
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
1015
        order_by=None, start=0), recover=recover, log_level=3))
1016

    
1017
def table_cols(db, table, recover=None):
1018
    return list(col_names(select(db, table, limit=0, order_by=None,
1019
        recover=recover, log_level=4)))
1020

    
1021
def pkey(db, table, recover=None):
1022
    '''Assumed to be first column in table'''
1023
    return table_cols(db, table, recover)[0]
1024

    
1025
not_null_col = 'not_null_col'
1026

    
1027
def table_not_null_col(db, table, recover=None):
1028
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1029
    if not_null_col in table_cols(db, table, recover): return not_null_col
1030
    else: return pkey(db, table, recover)
1031

    
1032
def index_cols(db, table, index):
1033
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1034
    automatically created. When you don't know whether something is a UNIQUE
1035
    constraint or a UNIQUE index, use this function.'''
1036
    module = util.root_module(db.db)
1037
    if module == 'psycopg2':
1038
        return list(values(run_query(db, '''\
1039
SELECT attname
1040
FROM
1041
(
1042
        SELECT attnum, attname
1043
        FROM pg_index
1044
        JOIN pg_class index ON index.oid = indexrelid
1045
        JOIN pg_class table_ ON table_.oid = indrelid
1046
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
1047
        WHERE
1048
            table_.relname = '''+db.esc_value(table)+'''
1049
            AND index.relname = '''+db.esc_value(index)+'''
1050
    UNION
1051
        SELECT attnum, attname
1052
        FROM
1053
        (
1054
            SELECT
1055
                indrelid
1056
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
1057
                    AS indkey
1058
            FROM pg_index
1059
            JOIN pg_class index ON index.oid = indexrelid
1060
            JOIN pg_class table_ ON table_.oid = indrelid
1061
            WHERE
1062
                table_.relname = '''+db.esc_value(table)+'''
1063
                AND index.relname = '''+db.esc_value(index)+'''
1064
        ) s
1065
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
1066
) s
1067
ORDER BY attnum
1068
'''
1069
            , cacheable=True, log_level=4)))
1070
    else: raise NotImplementedError("Can't list index columns for "+module+
1071
        ' database')
1072

    
1073
def constraint_cols(db, table, constraint):
1074
    module = util.root_module(db.db)
1075
    if module == 'psycopg2':
1076
        return list(values(run_query(db, '''\
1077
SELECT attname
1078
FROM pg_constraint
1079
JOIN pg_class ON pg_class.oid = conrelid
1080
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
1081
WHERE
1082
    relname = '''+db.esc_value(table)+'''
1083
    AND conname = '''+db.esc_value(constraint)+'''
1084
ORDER BY attnum
1085
'''
1086
            )))
1087
    else: raise NotImplementedError("Can't list constraint columns for "+module+
1088
        ' database')
1089

    
1090
#### Functions
1091

    
1092
def function_exists(db, function):
1093
    function = sql_gen.as_Function(function)
1094
    
1095
    info_table = sql_gen.Table('routines', 'information_schema')
1096
    conds = [('routine_name', function.name)]
1097
    schema = function.schema
1098
    if schema != None: conds.append(('routine_schema', schema))
1099
    # Exclude trigger functions, since they cannot be called directly
1100
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1101
    
1102
    return list(values(select(db, info_table, ['routine_name'], conds,
1103
        order_by='routine_schema', limit=1, log_level=4))) != []
1104
        # TODO: order_by search_path schema order
1105

    
1106
##### Structural changes
1107

    
1108
#### Columns
1109

    
1110
def add_col(db, table, col, comment=None, **kw_args):
1111
    '''
1112
    @param col TypedCol Name may be versioned, so be sure to propagate any
1113
        renaming back to any source column for the TypedCol.
1114
    @param comment None|str SQL comment used to distinguish columns of the same
1115
        name from each other when they contain different data, to allow the
1116
        ADD COLUMN query to be cached. If not set, query will not be cached.
1117
    '''
1118
    assert isinstance(col, sql_gen.TypedCol)
1119
    
1120
    while True:
1121
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1122
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1123
        
1124
        try:
1125
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1126
            break
1127
        except DuplicateException:
1128
            col.name = next_version(col.name)
1129
            # try again with next version of name
1130

    
1131
def add_not_null(db, col):
1132
    table = col.table
1133
    col = sql_gen.to_name_only_col(col)
1134
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1135
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1136

    
1137
row_num_col = '_row_num'
1138

    
1139
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1140
    constraints='PRIMARY KEY')
1141

    
1142
def add_row_num(db, table):
1143
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1144
    be the primary key.'''
1145
    add_col(db, table, row_num_typed_col, log_level=3)
1146

    
1147
#### Indexes
1148

    
1149
def add_pkey(db, table, cols=None, recover=None):
1150
    '''Adds a primary key.
1151
    @param cols [sql_gen.Col,...] The columns in the primary key.
1152
        Defaults to the first column in the table.
1153
    @pre The table must not already have a primary key.
1154
    '''
1155
    table = sql_gen.as_Table(table)
1156
    if cols == None: cols = [pkey(db, table, recover)]
1157
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1158
    
1159
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1160
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1161
        log_ignore_excs=(DuplicateException,))
1162

    
1163
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1164
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1165
    Currently, only function calls are supported as expressions.
1166
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1167
        This allows indexes to be used for comparisons where NULLs are equal.
1168
    '''
1169
    exprs = lists.mk_seq(exprs)
1170
    
1171
    # Parse exprs
1172
    old_exprs = exprs[:]
1173
    exprs = []
1174
    cols = []
1175
    for i, expr in enumerate(old_exprs):
1176
        expr = sql_gen.as_Col(expr, table)
1177
        
1178
        # Handle nullable columns
1179
        if ensure_not_null_:
1180
            try: expr = sql_gen.ensure_not_null(db, expr)
1181
            except KeyError: pass # unknown type, so just create plain index
1182
        
1183
        # Extract col
1184
        expr = copy.deepcopy(expr) # don't modify input!
1185
        if isinstance(expr, sql_gen.FunctionCall):
1186
            col = expr.args[0]
1187
            expr = sql_gen.Expr(expr)
1188
        else: col = expr
1189
        assert isinstance(col, sql_gen.Col)
1190
        
1191
        # Extract table
1192
        if table == None:
1193
            assert sql_gen.is_table_col(col)
1194
            table = col.table
1195
        
1196
        col.table = None
1197
        
1198
        exprs.append(expr)
1199
        cols.append(col)
1200
    
1201
    table = sql_gen.as_Table(table)
1202
    
1203
    # Add index
1204
    str_ = 'CREATE'
1205
    if unique: str_ += ' UNIQUE'
1206
    str_ += ' INDEX ON '+table.to_str(db)+' ('+(
1207
        ', '.join((v.to_str(db) for v in exprs)))+')'
1208
    run_query(db, str_, recover=True, cacheable=True, log_level=3)
1209

    
1210
already_indexed = object() # tells add_indexes() the pkey has already been added
1211

    
1212
def add_indexes(db, table, has_pkey=True):
1213
    '''Adds an index on all columns in a table.
1214
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1215
        index should be added on the first column.
1216
        * If already_indexed, the pkey is assumed to have already been added
1217
    '''
1218
    cols = table_cols(db, table)
1219
    if has_pkey:
1220
        if has_pkey is not already_indexed: add_pkey(db, table)
1221
        cols = cols[1:]
1222
    for col in cols: add_index(db, col, table)
1223

    
1224
#### Tables
1225

    
1226
### Maintenance
1227

    
1228
def analyze(db, table):
1229
    table = sql_gen.as_Table(table)
1230
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1231

    
1232
def autoanalyze(db, table):
1233
    if db.autoanalyze: analyze(db, table)
1234

    
1235
def vacuum(db, table):
1236
    table = sql_gen.as_Table(table)
1237
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1238
        log_level=3))
1239

    
1240
### Lifecycle
1241

    
1242
def drop(db, type_, name):
1243
    name = sql_gen.as_Name(name)
1244
    run_query(db, 'DROP '+type_+' IF EXISTS '+name.to_str(db)+' CASCADE')
1245

    
1246
def drop_table(db, table): drop(db, 'TABLE', table)
1247

    
1248
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1249
    like=None):
1250
    '''Creates a table.
1251
    @param cols [sql_gen.TypedCol,...] The column names and types
1252
    @param has_pkey If set, the first column becomes the primary key.
1253
    @param col_indexes bool|[ref]
1254
        * If True, indexes will be added on all non-pkey columns.
1255
        * If a list reference, [0] will be set to a function to do this.
1256
          This can be used to delay index creation until the table is populated.
1257
    '''
1258
    table = sql_gen.as_Table(table)
1259
    
1260
    if like != None:
1261
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1262
            ]+cols
1263
    if has_pkey:
1264
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1265
        pkey.constraints = 'PRIMARY KEY'
1266
    
1267
    temp = table.is_temp and not db.debug_temp
1268
        # temp tables permanent in debug_temp mode
1269
    
1270
    # Create table
1271
    while True:
1272
        str_ = 'CREATE'
1273
        if temp: str_ += ' TEMP'
1274
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1275
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1276
        str_ += '\n);'
1277
        
1278
        try:
1279
            run_query(db, str_, recover=True, cacheable=True, log_level=2,
1280
                log_ignore_excs=(DuplicateException,))
1281
            break
1282
        except DuplicateException:
1283
            table.name = next_version(table.name)
1284
            # try again with next version of name
1285
    
1286
    # Add indexes
1287
    if has_pkey: has_pkey = already_indexed
1288
    def add_indexes_(): add_indexes(db, table, has_pkey)
1289
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1290
    elif col_indexes: add_indexes_() # add now
1291

    
1292
def copy_table_struct(db, src, dest):
1293
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1294
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1295

    
1296
### Data
1297

    
1298
def truncate(db, table, schema='public', **kw_args):
1299
    '''For params, see run_query()'''
1300
    table = sql_gen.as_Table(table, schema)
1301
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1302

    
1303
def empty_temp(db, tables):
1304
    tables = lists.mk_seq(tables)
1305
    for table in tables: truncate(db, table, log_level=3)
1306

    
1307
def empty_db(db, schema='public', **kw_args):
1308
    '''For kw_args, see tables()'''
1309
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1310

    
1311
def distinct_table(db, table, distinct_on):
1312
    '''Creates a copy of a temp table which is distinct on the given columns.
1313
    The old and new tables will both get an index on these columns, to
1314
    facilitate merge joins.
1315
    @param distinct_on If empty, creates a table with one row. This is useful if
1316
        your distinct_on columns are all literal values.
1317
    @return The new table.
1318
    '''
1319
    new_table = sql_gen.suffixed_table(table, '_distinct')
1320
    
1321
    copy_table_struct(db, table, new_table)
1322
    
1323
    limit = None
1324
    if distinct_on == []: limit = 1 # one sample row
1325
    else:
1326
        add_index(db, distinct_on, new_table, unique=True)
1327
        add_index(db, distinct_on, table) # for join optimization
1328
    
1329
    insert_select(db, new_table, None, mk_select(db, table, start=0,
1330
        limit=limit), ignore=True)
1331
    analyze(db, new_table)
1332
    
1333
    return new_table
(24-24/37)