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, cause=None):
68
        DbException.__init__(self, 'Missing cast to type '+strings.as_tt(type_)
69
            +' on column: '+strings.as_tt(col), cause)
70
        self.type = type_
71
        self.col = col
72

    
73
class NameException(DbException): pass
74

    
75
class DuplicateKeyException(ConstraintException): pass
76

    
77
class NullValueException(ConstraintException): pass
78

    
79
class CheckException(ConstraintException): pass
80

    
81
class InvalidValueException(ExceptionWithValue): pass
82

    
83
class DuplicateException(ExceptionWithNameType): pass
84

    
85
class DoesNotExistException(ExceptionWithNameType): pass
86

    
87
class EmptyRowException(DbException): pass
88

    
89
##### Warnings
90

    
91
class DbWarning(UserWarning): pass
92

    
93
##### Result retrieval
94

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

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

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

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

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

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

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

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

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

    
120
##### Escaping
121

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

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

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

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

    
140
##### Database connections
141

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

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

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

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

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

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

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

    
488
connect = DbConn
489

    
490
##### Recoverable querying
491

    
492
def with_savepoint(db, func): return db.with_savepoint(func)
493

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

    
573
##### Basic queries
574

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

    
580
def explain(db, query, **kw_args):
581
    '''
582
    For params, see run_query().
583
    '''
584
    kw_args.setdefault('log_level', 4)
585
    
586
    return strings.join_lines(values(run_query(db, 'EXPLAIN '+query,
587
        recover=True, cacheable=True, **kw_args)))
588
        # not a higher log_level because it's useful to see what query is being
589
        # run before it's executed, which EXPLAIN effectively provides
590

    
591
def has_comment(query): return query.endswith('*/')
592

    
593
def with_explain_comment(db, query, **kw_args):
594
    if db.autoexplain and not has_comment(query) and is_explainable(query):
595
        query += '\n'+sql_gen.esc_comment(' EXPLAIN:\n'
596
            +explain(db, query, **kw_args))
597
    return query
598

    
599
def next_version(name):
600
    version = 1 # first existing name was version 0
601
    match = re.match(r'^(.*)#(\d+)$', name)
602
    if match:
603
        name, version = match.groups()
604
        version = int(version)+1
605
    return sql_gen.concat(name, '#'+str(version))
606

    
607
def lock_table(db, table, mode):
608
    table = sql_gen.as_Table(table)
609
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
610

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

    
654
order_by_pkey = object() # tells mk_select() to order by the pkey
655

    
656
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
657

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

    
741
def select(db, *args, **kw_args):
742
    '''For params, see mk_select() and run_query()'''
743
    recover = kw_args.pop('recover', None)
744
    cacheable = kw_args.pop('cacheable', True)
745
    log_level = kw_args.pop('log_level', 2)
746
    
747
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
748
        log_level=log_level)
749

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

    
822
def insert_select(db, table, *args, **kw_args):
823
    '''For params, see mk_insert_select() and run_query_into()
824
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
825
        values in
826
    '''
827
    returning = kw_args.get('returning', None)
828
    ignore = kw_args.get('ignore', False)
829
    
830
    into = kw_args.pop('into', None)
831
    if into != None: kw_args['embeddable'] = True
832
    recover = kw_args.pop('recover', None)
833
    if ignore: recover = True
834
    cacheable = kw_args.pop('cacheable', True)
835
    log_level = kw_args.pop('log_level', 2)
836
    
837
    rowcount_only = ignore and returning == None # keep NULL rows on server
838
    if rowcount_only: into = sql_gen.Table('rowcount')
839
    
840
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
841
        into, recover=recover, cacheable=cacheable, log_level=log_level)
842
    if rowcount_only: empty_temp(db, into)
843
    autoanalyze(db, table)
844
    return cur
845

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

    
848
def insert(db, table, row, *args, **kw_args):
849
    '''For params, see insert_select()'''
850
    if lists.is_seq(row): cols = None
851
    else:
852
        cols = row.keys()
853
        row = row.values()
854
    row = list(row) # ensure that "== []" works
855
    
856
    if row == []: query = None
857
    else: query = sql_gen.Values(row).to_str(db)
858
    
859
    return insert_select(db, table, cols, query, *args, **kw_args)
860

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

    
897
def update(db, table, *args, **kw_args):
898
    '''For params, see mk_update() and run_query()'''
899
    recover = kw_args.pop('recover', None)
900
    cacheable = kw_args.pop('cacheable', False)
901
    log_level = kw_args.pop('log_level', 2)
902
    
903
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
904
        cacheable, log_level=log_level)
905
    autoanalyze(db, table)
906
    return cur
907

    
908
def mk_delete(db, table, cond=None):
909
    '''
910
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
911
    @return str query
912
    '''
913
    query = 'DELETE FROM '+table.to_str(db)
914
    if cond != None: query += '\nWHERE '+cond.to_str(db)
915
    
916
    query = with_explain_comment(db, query)
917
    
918
    return query
919

    
920
def delete(db, table, *args, **kw_args):
921
    '''For params, see mk_delete() and run_query()'''
922
    recover = kw_args.pop('recover', None)
923
    cacheable = kw_args.pop('cacheable', True)
924
    log_level = kw_args.pop('log_level', 2)
925
    
926
    cur = run_query(db, mk_delete(db, table, *args, **kw_args), recover,
927
        cacheable, log_level=log_level)
928
    autoanalyze(db, table)
929
    return cur
930

    
931
def last_insert_id(db):
932
    module = util.root_module(db.db)
933
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
934
    elif module == 'MySQLdb': return db.insert_id()
935
    else: return None
936

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

    
969
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
970
    '''For params, see mk_flatten_mapping()
971
    @return See return value of mk_flatten_mapping()
972
    '''
973
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
974
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
975
    run_query_into(db, mk_select(db, joins, cols, order_by=None, limit=limit,
976
        start=start), into=into, add_pkey_=True)
977
    return dict(items)
978

    
979
##### Database structure introspection
980

    
981
#### Expressions
982

    
983
bool_re = r'(?:true|false)'
984

    
985
def simplify_expr(expr):
986
    expr = expr.replace('(NULL IS NULL)', 'true')
987
    expr = expr.replace('(NULL IS NOT NULL)', 'false')
988
    expr = re.sub(r' OR '+bool_re, r'', expr)
989
    expr = re.sub(bool_re+r' OR ', r'', expr)
990
    while True:
991
        expr, n = re.subn(r'\((\([^()]*\))\)', r'\1', expr)
992
        if n == 0: break
993
    return expr
994

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

    
997
def parse_expr_col(str_):
998
    match = re.match(r'^\('+name_re+r'\(('+name_re+r').*\)\)$', str_)
999
    if match: str_ = match.group(1)
1000
    return sql_gen.unesc_name(str_)
1001

    
1002
def map_expr(db, expr, mapping, in_cols_found=None):
1003
    '''Replaces output columns with input columns in an expression.
1004
    @param in_cols_found If set, will be filled in with the expr's (input) cols
1005
    '''
1006
    for out, in_ in mapping.iteritems():
1007
        orig_expr = expr
1008
        out = sql_gen.to_name_only_col(out)
1009
        in_str = sql_gen.to_name_only_col(sql_gen.remove_col_rename(in_)
1010
            ).to_str(db)
1011
        
1012
        # Replace out both with and without quotes
1013
        expr = expr.replace(out.to_str(db), in_str)
1014
        expr = re.sub(r'\b'+out.name+r'\b', in_str, expr)
1015
        
1016
        if in_cols_found != None and expr != orig_expr: # replaced something
1017
            in_cols_found.append(in_)
1018
    
1019
    return simplify_expr(expr)
1020

    
1021
#### Tables
1022

    
1023
def tables(db, schema_like='public', table_like='%', exact=False):
1024
    if exact: compare = '='
1025
    else: compare = 'LIKE'
1026
    
1027
    module = util.root_module(db.db)
1028
    if module == 'psycopg2':
1029
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
1030
            ('tablename', sql_gen.CompareCond(table_like, compare))]
1031
        return values(select(db, 'pg_tables', ['tablename'], conds,
1032
            order_by='tablename', log_level=4))
1033
    elif module == 'MySQLdb':
1034
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
1035
            , cacheable=True, log_level=4))
1036
    else: raise NotImplementedError("Can't list tables for "+module+' database')
1037

    
1038
def table_exists(db, table):
1039
    table = sql_gen.as_Table(table)
1040
    return list(tables(db, table.schema, table.name, exact=True)) != []
1041

    
1042
def table_row_count(db, table, recover=None):
1043
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
1044
        order_by=None), recover=recover, log_level=3))
1045

    
1046
def table_cols(db, table, recover=None):
1047
    return list(col_names(select(db, table, limit=0, order_by=None,
1048
        recover=recover, log_level=4)))
1049

    
1050
def pkey(db, table, recover=None):
1051
    '''Assumed to be first column in table'''
1052
    return table_cols(db, table, recover)[0]
1053

    
1054
not_null_col = 'not_null_col'
1055

    
1056
def table_not_null_col(db, table, recover=None):
1057
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
1058
    if not_null_col in table_cols(db, table, recover): return not_null_col
1059
    else: return pkey(db, table, recover)
1060

    
1061
def constraint_cond(db, constraint):
1062
    module = util.root_module(db.db)
1063
    if module == 'psycopg2':
1064
        table_str = sql_gen.Literal(constraint.table.to_str(db))
1065
        name_str = sql_gen.Literal(constraint.name)
1066
        return value(run_query(db, '''\
1067
SELECT consrc
1068
FROM pg_constraint
1069
WHERE
1070
conrelid = '''+table_str.to_str(db)+'''::regclass
1071
AND conname = '''+name_str.to_str(db)+'''
1072
'''
1073
            , cacheable=True, log_level=4))
1074
    else: raise NotImplementedError("Can't list index columns for "+module+
1075
        ' database')
1076

    
1077
def index_cols(db, index):
1078
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
1079
    automatically created. When you don't know whether something is a UNIQUE
1080
    constraint or a UNIQUE index, use this function.'''
1081
    index = sql_gen.as_Table(index)
1082
    module = util.root_module(db.db)
1083
    if module == 'psycopg2':
1084
        qual_index = sql_gen.Literal(index.to_str(db))
1085
        return map(parse_expr_col, values(run_query(db, '''\
1086
SELECT pg_get_indexdef(indexrelid, generate_series(1, indnatts), true)
1087
FROM pg_index
1088
WHERE indexrelid = '''+qual_index.to_str(db)+'''::regclass
1089
'''
1090
            , cacheable=True, log_level=4)))
1091
    else: raise NotImplementedError("Can't list index columns for "+module+
1092
        ' database')
1093

    
1094
#### Functions
1095

    
1096
def function_exists(db, function):
1097
    qual_function = sql_gen.Literal(function.to_str(db))
1098
    try:
1099
        select(db, fields=[sql_gen.Cast('regproc', qual_function)],
1100
            recover=True, cacheable=True, log_level=4)
1101
    except DoesNotExistException: return False
1102
    else: return True
1103

    
1104
##### Structural changes
1105

    
1106
#### Columns
1107

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

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

    
1135
row_num_col = '_row_num'
1136

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

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

    
1145
#### Indexes
1146

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

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

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

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

    
1223
#### Tables
1224

    
1225
### Maintenance
1226

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

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

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

    
1239
### Lifecycle
1240

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

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

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

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

    
1299
### Data
1300

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

    
1306
def empty_temp(db, tables):
1307
    tables = lists.mk_seq(tables)
1308
    for table in tables: truncate(db, table, log_level=3)
1309

    
1310
def empty_db(db, schema='public', **kw_args):
1311
    '''For kw_args, see tables()'''
1312
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1313

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