Project

General

Profile

1
# Database access
2

    
3
import copy
4
import re
5
import warnings
6

    
7
import exc
8
import dicts
9
import iters
10
import lists
11
from Proxy import Proxy
12
import rand
13
import sql_gen
14
import strings
15
import util
16

    
17
##### Exceptions
18

    
19
def get_cur_query(cur, input_query=None):
20
    raw_query = None
21
    if hasattr(cur, 'query'): raw_query = cur.query
22
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
23
    
24
    if raw_query != None: return raw_query
25
    else: return '[input] '+strings.ustr(input_query)
26

    
27
def _add_cursor_info(e, *args, **kw_args):
28
    '''For params, see get_cur_query()'''
29
    exc.add_msg(e, 'query: '+strings.ustr(get_cur_query(*args, **kw_args)))
30

    
31
class DbException(exc.ExceptionWithCause):
32
    def __init__(self, msg, cause=None, cur=None):
33
        exc.ExceptionWithCause.__init__(self, msg, cause, cause_newline=True)
34
        if cur != None: _add_cursor_info(self, cur)
35

    
36
class ExceptionWithName(DbException):
37
    def __init__(self, name, cause=None):
38
        DbException.__init__(self, 'for name: '+strings.as_tt(str(name)), cause)
39
        self.name = name
40

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

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

    
54
class ConstraintException(DbException):
55
    def __init__(self, name, cols, cause=None):
56
        DbException.__init__(self, 'Violated '+strings.as_tt(name)
57
            +' constraint on columns: '+strings.as_tt(', '.join(cols)), cause)
58
        self.name = name
59
        self.cols = cols
60

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

    
68
class NameException(DbException): pass
69

    
70
class DuplicateKeyException(ConstraintException): pass
71

    
72
class NullValueException(ConstraintException): pass
73

    
74
class InvalidValueException(ExceptionWithValue): pass
75

    
76
class DuplicateException(ExceptionWithNameType): pass
77

    
78
class EmptyRowException(DbException): pass
79

    
80
##### Warnings
81

    
82
class DbWarning(UserWarning): pass
83

    
84
##### Result retrieval
85

    
86
def col_names(cur): return (col[0] for col in cur.description)
87

    
88
def rows(cur): return iter(lambda: cur.fetchone(), None)
89

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

    
94
def next_row(cur): return rows(cur).next()
95

    
96
def row(cur):
97
    row_ = next_row(cur)
98
    consume_rows(cur)
99
    return row_
100

    
101
def next_value(cur): return next_row(cur)[0]
102

    
103
def value(cur): return row(cur)[0]
104

    
105
def values(cur): return iters.func_iter(lambda: next_value(cur))
106

    
107
def value_or_none(cur):
108
    try: return value(cur)
109
    except StopIteration: return None
110

    
111
##### Escaping
112

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

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

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

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

    
131
##### Database connections
132

    
133
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
134

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

    
140
DatabaseErrors_set = set([DbException])
141
DatabaseErrors = tuple(DatabaseErrors_set)
142

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

    
148
def db_config_str(db_config):
149
    return db_config['engine']+' database '+db_config['database']
150

    
151
log_debug_none = lambda msg, level=2: None
152

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

    
451
connect = DbConn
452

    
453
##### Recoverable querying
454

    
455
def with_savepoint(db, func): return db.with_savepoint(func)
456

    
457
def run_query(db, query, recover=None, cacheable=False, log_level=2,
458
    log_ignore_excs=None, **kw_args):
459
    '''For params, see DbConn.run_query()'''
460
    if recover == None: recover = False
461
    if log_ignore_excs == None: log_ignore_excs = ()
462
    log_ignore_excs = tuple(log_ignore_excs)
463
    
464
    debug_msg_ref = None # usually, db.run_query() logs query before running it
465
    # But if filtering with log_ignore_excs, wait until after exception parsing
466
    if log_ignore_excs != () or not db.can_mogrify(): debug_msg_ref = [None]
467
    
468
    try:
469
        try:
470
            def run(): return db.run_query(query, cacheable, log_level,
471
                debug_msg_ref, **kw_args)
472
            if recover and not db.is_cached(query):
473
                return with_savepoint(db, run)
474
            else: return run() # don't need savepoint if cached
475
        except Exception, e:
476
            msg = strings.ustr(e.args[0])
477
            
478
            match = re.match(r'^duplicate key value violates unique constraint '
479
                r'"((_?[^\W_]+(?=[._]))?.+?)"', msg)
480
            if match:
481
                constraint, table = match.groups()
482
                cols = []
483
                if recover: # need auto-rollback to run index_cols()
484
                    try: cols = index_cols(db, table, constraint)
485
                    except NotImplementedError: pass
486
                raise DuplicateKeyException(constraint, cols, e)
487
            
488
            match = re.match(r'^null value in column "(.+?)" violates not-null'
489
                r' constraint', msg)
490
            if match: raise NullValueException('NOT NULL', [match.group(1)], e)
491
            
492
            match = re.match(r'^(?:invalid input (?:syntax|value)\b.*?'
493
                r'|.+? field value out of range): "(.+?)"', msg)
494
            if match:
495
                value, = match.groups()
496
                raise InvalidValueException(strings.to_unicode(value), e)
497
            
498
            match = re.match(r'^column "(.+?)" is of type (.+?) but expression '
499
                r'is of type', msg)
500
            if match:
501
                col, type_ = match.groups()
502
                raise MissingCastException(type_, col, e)
503
            
504
            match = re.match(r'^(\S+) "(.+?)".*? already exists', msg)
505
            if match:
506
                type_, name = match.groups()
507
                raise DuplicateException(type_, name, e)
508
            
509
            raise # no specific exception raised
510
    except log_ignore_excs:
511
        log_level += 2
512
        raise
513
    finally:
514
        if debug_msg_ref != None and debug_msg_ref[0] != None:
515
            db.log_debug(debug_msg_ref[0], log_level)
516

    
517
##### Basic queries
518

    
519
def next_version(name):
520
    version = 1 # first existing name was version 0
521
    match = re.match(r'^(.*)#(\d+)$', name)
522
    if match:
523
        name, version = match.groups()
524
        version = int(version)+1
525
    return sql_gen.concat(name, '#'+str(version))
526

    
527
def lock_table(db, table, mode):
528
    table = sql_gen.as_Table(table)
529
    run_query(db, 'LOCK TABLE '+table.to_str(db)+' IN '+mode+' MODE')
530

    
531
def run_query_into(db, query, into=None, add_indexes_=False, **kw_args):
532
    '''Outputs a query to a temp table.
533
    For params, see run_query().
534
    '''
535
    if into == None: return run_query(db, query, **kw_args)
536
    
537
    assert isinstance(into, sql_gen.Table)
538
    
539
    into.is_temp = True
540
    # "temporary tables cannot specify a schema name", so remove schema
541
    into.schema = None
542
    
543
    kw_args['recover'] = True
544
    kw_args.setdefault('log_ignore_excs', (DuplicateException,))
545
    
546
    temp = not db.debug_temp # tables are permanent in debug_temp mode
547
    
548
    # Create table
549
    while True:
550
        create_query = 'CREATE'
551
        if temp: create_query += ' TEMP'
552
        create_query += ' TABLE '+into.to_str(db)+' AS\n'+query
553
        
554
        try:
555
            cur = run_query(db, create_query, **kw_args)
556
                # CREATE TABLE AS sets rowcount to # rows in query
557
            break
558
        except DuplicateException, e:
559
            into.name = next_version(into.name)
560
            # try again with next version of name
561
    
562
    if add_indexes_: add_indexes(db, into)
563
    
564
    # According to the PostgreSQL doc, "The autovacuum daemon cannot access and
565
    # therefore cannot vacuum or analyze temporary tables. [...] if a temporary
566
    # table is going to be used in complex queries, it is wise to run ANALYZE on
567
    # the temporary table after it is populated."
568
    # (http://www.postgresql.org/docs/9.1/static/sql-createtable.html)
569
    # If into is not a temp table, ANALYZE is useful but not required.
570
    analyze(db, into)
571
    
572
    return cur
573

    
574
order_by_pkey = object() # tells mk_select() to order by the pkey
575

    
576
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
577

    
578
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
579
    start=None, order_by=order_by_pkey, default_table=None):
580
    '''
581
    @param tables The single table to select from, or a list of tables to join
582
        together, with tables after the first being sql_gen.Join objects
583
    @param fields Use None to select all fields in the table
584
    @param conds WHERE conditions: [(compare_left_side, compare_right_side),...]
585
        * container can be any iterable type
586
        * compare_left_side: sql_gen.Code|str (for col name)
587
        * compare_right_side: sql_gen.ValueCond|literal value
588
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
589
        use all columns
590
    @return query
591
    '''
592
    # Parse tables param
593
    tables = lists.mk_seq(tables)
594
    tables = list(tables) # don't modify input! (list() copies input)
595
    table0 = sql_gen.as_Table(tables.pop(0)) # first table is separate
596
    
597
    # Parse other params
598
    if conds == None: conds = []
599
    elif dicts.is_dict(conds): conds = conds.items()
600
    conds = list(conds) # don't modify input! (list() copies input)
601
    assert limit == None or isinstance(limit, (int, long))
602
    assert start == None or isinstance(start, (int, long))
603
    if order_by is order_by_pkey:
604
        if distinct_on != []: order_by = None
605
        else: order_by = pkey(db, table0, recover=True)
606
    
607
    query = 'SELECT'
608
    
609
    def parse_col(col): return sql_gen.as_Col(col, default_table).to_str(db)
610
    
611
    # DISTINCT ON columns
612
    if distinct_on != []:
613
        query += '\nDISTINCT'
614
        if distinct_on is not distinct_on_all:
615
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
616
    
617
    # Columns
618
    if fields == None:
619
        if query.find('\n') >= 0: whitespace = '\n'
620
        else: whitespace = ' '
621
        query += whitespace+'*'
622
    else:
623
        assert fields != []
624
        query += '\n'+('\n, '.join(map(parse_col, fields)))
625
    
626
    # Main table
627
    query += '\nFROM '+table0.to_str(db)
628
    
629
    # Add joins
630
    left_table = table0
631
    for join_ in tables:
632
        table = join_.table
633
        
634
        # Parse special values
635
        if join_.type_ is sql_gen.filter_out: # filter no match
636
            conds.append((sql_gen.Col(table_not_null_col(db, table), table),
637
                sql_gen.CompareCond(None, '~=')))
638
        
639
        query += '\n'+join_.to_str(db, left_table)
640
        
641
        left_table = table
642
    
643
    missing = True
644
    if conds != []:
645
        if len(conds) == 1: whitespace = ' '
646
        else: whitespace = '\n'
647
        query += '\n'+sql_gen.combine_conds([sql_gen.ColValueCond(l, r)
648
            .to_str(db) for l, r in conds], 'WHERE')
649
        missing = False
650
    if order_by != None:
651
        query += '\nORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
652
    if limit != None: query += '\nLIMIT '+str(limit); missing = False
653
    if start != None:
654
        if start != 0: query += '\nOFFSET '+str(start)
655
        missing = False
656
    if missing: warnings.warn(DbWarning(
657
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
658
    
659
    return query
660

    
661
def select(db, *args, **kw_args):
662
    '''For params, see mk_select() and run_query()'''
663
    recover = kw_args.pop('recover', None)
664
    cacheable = kw_args.pop('cacheable', True)
665
    log_level = kw_args.pop('log_level', 2)
666
    
667
    return run_query(db, mk_select(db, *args, **kw_args), recover, cacheable,
668
        log_level=log_level)
669

    
670
def mk_insert_select(db, table, cols=None, select_query=None, returning=None,
671
    embeddable=False, ignore=False):
672
    '''
673
    @param returning str|None An inserted column (such as pkey) to return
674
    @param embeddable Whether the query should be embeddable as a nested SELECT.
675
        Warning: If you set this and cacheable=True when the query is run, the
676
        query will be fully cached, not just if it raises an exception.
677
    @param ignore Whether to ignore duplicate keys.
678
    '''
679
    table = sql_gen.remove_table_rename(sql_gen.as_Table(table))
680
    if cols == []: cols = None # no cols (all defaults) = unknown col names
681
    if cols != None: cols = [sql_gen.to_name_only_col(c, table) for c in cols]
682
    if select_query == None: select_query = 'DEFAULT VALUES'
683
    if returning != None: returning = sql_gen.as_Col(returning, table)
684
    
685
    first_line = 'INSERT INTO '+table.to_str(db)
686
    
687
    def mk_insert(select_query):
688
        query = first_line
689
        if cols != None:
690
            query += '\n('+(', '.join((c.to_str(db) for c in cols)))+')'
691
        query += '\n'+select_query
692
        
693
        if returning != None:
694
            returning_name_col = sql_gen.to_name_only_col(returning)
695
            query += '\nRETURNING '+returning_name_col.to_str(db)
696
        
697
        return query
698
    
699
    return_type = 'unknown'
700
    if returning != None: return_type = returning.to_str(db)+'%TYPE'
701
    
702
    lang = 'sql'
703
    if ignore:
704
        # Always return something to set the correct rowcount
705
        if returning == None: returning = sql_gen.NamedCol('NULL', None)
706
        
707
        embeddable = True # must use function
708
        lang = 'plpgsql'
709
        
710
        if cols == None:
711
            row = [sql_gen.Col(sql_gen.all_cols, 'row')]
712
            row_vars = [sql_gen.Table('row')]
713
        else:
714
            row_vars = row = [sql_gen.Col(c.name, 'row') for c in cols]
715
        
716
        query = '''\
717
DECLARE
718
    row '''+table.to_str(db)+'''%ROWTYPE;
719
BEGIN
720
    /* Need an EXCEPTION block for each individual row because "When an error is
721
    caught by an EXCEPTION clause, [...] all changes to persistent database
722
    state within the block are rolled back."
723
    This is unfortunate because "A block containing an EXCEPTION clause is
724
    significantly more expensive to enter and exit than a block without one."
725
    (http://www.postgresql.org/docs/8.3/static/plpgsql-control-structures.html\
726
#PLPGSQL-ERROR-TRAPPING)
727
    */
728
    FOR '''+(', '.join((v.to_str(db) for v in row_vars)))+''' IN
729
'''+select_query+'''
730
    LOOP
731
        BEGIN
732
            RETURN QUERY
733
'''+mk_insert(sql_gen.Values(row).to_str(db))+'''
734
;
735
        EXCEPTION
736
            WHEN unique_violation THEN NULL; -- continue to next row
737
        END;
738
    END LOOP;
739
END;\
740
'''
741
    else: query = mk_insert(select_query)
742
    
743
    if embeddable:
744
        # Create function
745
        function_name = sql_gen.clean_name(first_line)
746
        while True:
747
            try:
748
                function = db.TempFunction(function_name)
749
                
750
                function_query = '''\
751
CREATE FUNCTION '''+function.to_str(db)+'''()
752
RETURNS SETOF '''+return_type+'''
753
LANGUAGE '''+lang+'''
754
AS $$
755
'''+query+'''
756
$$;
757
'''
758
                run_query(db, function_query, recover=True, cacheable=True,
759
                    log_ignore_excs=(DuplicateException,))
760
                break # this version was successful
761
            except DuplicateException, e:
762
                function_name = next_version(function_name)
763
                # try again with next version of name
764
        
765
        # Return query that uses function
766
        cols = None
767
        if returning != None: cols = [returning]
768
        func_table = sql_gen.NamedTable('f', sql_gen.FunctionCall(function),
769
            cols) # AS clause requires function alias
770
        return mk_select(db, func_table, start=0, order_by=None)
771
    
772
    return query
773

    
774
def insert_select(db, table, *args, **kw_args):
775
    '''For params, see mk_insert_select() and run_query_into()
776
    @param into sql_gen.Table with suggested name of temp table to put RETURNING
777
        values in
778
    '''
779
    returning = kw_args.get('returning', None)
780
    ignore = kw_args.get('ignore', False)
781
    
782
    into = kw_args.pop('into', None)
783
    if into != None: kw_args['embeddable'] = True
784
    recover = kw_args.pop('recover', None)
785
    if ignore: recover = True
786
    cacheable = kw_args.pop('cacheable', True)
787
    log_level = kw_args.pop('log_level', 2)
788
    
789
    rowcount_only = ignore and returning == None # keep NULL rows on server
790
    if rowcount_only: into = sql_gen.Table('rowcount')
791
    
792
    cur = run_query_into(db, mk_insert_select(db, table, *args, **kw_args),
793
        into, recover=recover, cacheable=cacheable, log_level=log_level)
794
    if rowcount_only: empty_temp(db, into)
795
    autoanalyze(db, table)
796
    return cur
797

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

    
800
def insert(db, table, row, *args, **kw_args):
801
    '''For params, see insert_select()'''
802
    if lists.is_seq(row): cols = None
803
    else:
804
        cols = row.keys()
805
        row = row.values()
806
    row = list(row) # ensure that "== []" works
807
    
808
    if row == []: query = None
809
    else: query = sql_gen.Values(row).to_str(db)
810
    
811
    return insert_select(db, table, cols, query, *args, **kw_args)
812

    
813
def mk_update(db, table, changes=None, cond=None, in_place=False):
814
    '''
815
    @param changes [(col, new_value),...]
816
        * container can be any iterable type
817
        * col: sql_gen.Code|str (for col name)
818
        * new_value: sql_gen.Code|literal value
819
    @param cond sql_gen.Code WHERE condition. e.g. use sql_gen.*Cond objects.
820
    @param in_place If set, locks the table and updates rows in place.
821
        This avoids creating dead rows in PostgreSQL.
822
        * cond must be None
823
    @return str query
824
    '''
825
    table = sql_gen.as_Table(table)
826
    changes = [(sql_gen.to_name_only_col(c, table), sql_gen.as_Value(v))
827
        for c, v in changes]
828
    
829
    if in_place:
830
        assert cond == None
831
        
832
        query = 'ALTER TABLE '+table.to_str(db)+'\n'
833
        query += ',\n'.join(('ALTER COLUMN '+c.to_str(db)+' TYPE '
834
            +db.col_info(sql_gen.with_default_table(c, table)).type
835
            +'\nUSING '+v.to_str(db) for c, v in changes))
836
    else:
837
        query = 'UPDATE '+table.to_str(db)+'\nSET\n'
838
        query += ',\n'.join((c.to_str(db)+' = '+v.to_str(db)
839
            for c, v in changes))
840
        if cond != None: query += '\nWHERE\n'+cond.to_str(db)
841
    
842
    return query
843

    
844
def update(db, table, *args, **kw_args):
845
    '''For params, see mk_update() and run_query()'''
846
    recover = kw_args.pop('recover', None)
847
    cacheable = kw_args.pop('cacheable', False)
848
    log_level = kw_args.pop('log_level', 2)
849
    
850
    cur = run_query(db, mk_update(db, table, *args, **kw_args), recover,
851
        cacheable, log_level=log_level)
852
    autoanalyze(db, table)
853
    return cur
854

    
855
def last_insert_id(db):
856
    module = util.root_module(db.db)
857
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
858
    elif module == 'MySQLdb': return db.insert_id()
859
    else: return None
860

    
861
def mk_flatten_mapping(db, into, cols, preserve=[], as_items=False):
862
    '''Creates a mapping from original column names (which may have collisions)
863
    to names that will be distinct among the columns' tables.
864
    This is meant to be used for several tables that are being joined together.
865
    @param cols The columns to combine. Duplicates will be removed.
866
    @param into The table for the new columns.
867
    @param preserve [sql_gen.Col...] Columns not to rename. Note that these
868
        columns will be included in the mapping even if they are not in cols.
869
        The tables of the provided Col objects will be changed to into, so make
870
        copies of them if you want to keep the original tables.
871
    @param as_items Whether to return a list of dict items instead of a dict
872
    @return dict(orig_col=new_col, ...)
873
        * orig_col: sql_gen.Col(orig_col_name, orig_table)
874
        * new_col: sql_gen.Col(orig_col_name, into)
875
        * All mappings use the into table so its name can easily be
876
          changed for all columns at once
877
    '''
878
    cols = lists.uniqify(cols)
879
    
880
    items = []
881
    for col in preserve:
882
        orig_col = copy.copy(col)
883
        col.table = into
884
        items.append((orig_col, col))
885
    preserve = set(preserve)
886
    for col in cols:
887
        if col not in preserve:
888
            items.append((col, sql_gen.Col(str(col), into, col.srcs)))
889
    
890
    if not as_items: items = dict(items)
891
    return items
892

    
893
def flatten(db, into, joins, cols, limit=None, start=None, **kw_args):
894
    '''For params, see mk_flatten_mapping()
895
    @return See return value of mk_flatten_mapping()
896
    '''
897
    items = mk_flatten_mapping(db, into, cols, as_items=True, **kw_args)
898
    cols = [sql_gen.NamedCol(new.name, old) for old, new in items]
899
    run_query_into(db, mk_select(db, joins, cols, limit=limit, start=start),
900
        into=into, add_indexes_=True)
901
    return dict(items)
902

    
903
##### Database structure introspection
904

    
905
#### Tables
906

    
907
def tables(db, schema_like='public', table_like='%', exact=False):
908
    if exact: compare = '='
909
    else: compare = 'LIKE'
910
    
911
    module = util.root_module(db.db)
912
    if module == 'psycopg2':
913
        conds = [('schemaname', sql_gen.CompareCond(schema_like, compare)),
914
            ('tablename', sql_gen.CompareCond(table_like, compare))]
915
        return values(select(db, 'pg_tables', ['tablename'], conds,
916
            order_by='tablename', log_level=4))
917
    elif module == 'MySQLdb':
918
        return values(run_query(db, 'SHOW TABLES LIKE '+db.esc_value(table_like)
919
            , cacheable=True, log_level=4))
920
    else: raise NotImplementedError("Can't list tables for "+module+' database')
921

    
922
def table_exists(db, table):
923
    table = sql_gen.as_Table(table)
924
    return list(tables(db, table.schema, table.name, exact=True)) != []
925

    
926
def table_row_count(db, table, recover=None):
927
    return value(run_query(db, mk_select(db, table, [sql_gen.row_count],
928
        order_by=None, start=0), recover=recover, log_level=3))
929

    
930
def table_cols(db, table, recover=None):
931
    return list(col_names(select(db, table, limit=0, order_by=None,
932
        recover=recover, log_level=4)))
933

    
934
def pkey(db, table, recover=None):
935
    '''Assumed to be first column in table'''
936
    return table_cols(db, table, recover)[0]
937

    
938
not_null_col = 'not_null_col'
939

    
940
def table_not_null_col(db, table, recover=None):
941
    '''Name assumed to be the value of not_null_col. If not found, uses pkey.'''
942
    if not_null_col in table_cols(db, table, recover): return not_null_col
943
    else: return pkey(db, table, recover)
944

    
945
def index_cols(db, table, index):
946
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
947
    automatically created. When you don't know whether something is a UNIQUE
948
    constraint or a UNIQUE index, use this function.'''
949
    module = util.root_module(db.db)
950
    if module == 'psycopg2':
951
        return list(values(run_query(db, '''\
952
SELECT attname
953
FROM
954
(
955
        SELECT attnum, attname
956
        FROM pg_index
957
        JOIN pg_class index ON index.oid = indexrelid
958
        JOIN pg_class table_ ON table_.oid = indrelid
959
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
960
        WHERE
961
            table_.relname = '''+db.esc_value(table)+'''
962
            AND index.relname = '''+db.esc_value(index)+'''
963
    UNION
964
        SELECT attnum, attname
965
        FROM
966
        (
967
            SELECT
968
                indrelid
969
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
970
                    AS indkey
971
            FROM pg_index
972
            JOIN pg_class index ON index.oid = indexrelid
973
            JOIN pg_class table_ ON table_.oid = indrelid
974
            WHERE
975
                table_.relname = '''+db.esc_value(table)+'''
976
                AND index.relname = '''+db.esc_value(index)+'''
977
        ) s
978
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
979
) s
980
ORDER BY attnum
981
'''
982
            , cacheable=True, log_level=4)))
983
    else: raise NotImplementedError("Can't list index columns for "+module+
984
        ' database')
985

    
986
def constraint_cols(db, table, constraint):
987
    module = util.root_module(db.db)
988
    if module == 'psycopg2':
989
        return list(values(run_query(db, '''\
990
SELECT attname
991
FROM pg_constraint
992
JOIN pg_class ON pg_class.oid = conrelid
993
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
994
WHERE
995
    relname = '''+db.esc_value(table)+'''
996
    AND conname = '''+db.esc_value(constraint)+'''
997
ORDER BY attnum
998
'''
999
            )))
1000
    else: raise NotImplementedError("Can't list constraint columns for "+module+
1001
        ' database')
1002

    
1003
#### Functions
1004

    
1005
def function_exists(db, function):
1006
    function = sql_gen.as_Function(function)
1007
    
1008
    info_table = sql_gen.Table('routines', 'information_schema')
1009
    conds = [('routine_name', function.name)]
1010
    schema = function.schema
1011
    if schema != None: conds.append(('routine_schema', schema))
1012
    # Exclude trigger functions, since they cannot be called directly
1013
    conds.append(('data_type', sql_gen.CompareCond('trigger', '!=')))
1014
    
1015
    return list(values(select(db, info_table, ['routine_name'], conds,
1016
        order_by='routine_schema', limit=1, log_level=4))) != []
1017
        # TODO: order_by search_path schema order
1018

    
1019
##### Structural changes
1020

    
1021
#### Columns
1022

    
1023
def add_col(db, table, col, comment=None, **kw_args):
1024
    '''
1025
    @param col TypedCol Name may be versioned, so be sure to propagate any
1026
        renaming back to any source column for the TypedCol.
1027
    @param comment None|str SQL comment used to distinguish columns of the same
1028
        name from each other when they contain different data, to allow the
1029
        ADD COLUMN query to be cached. If not set, query will not be cached.
1030
    '''
1031
    assert isinstance(col, sql_gen.TypedCol)
1032
    
1033
    while True:
1034
        str_ = 'ALTER TABLE '+table.to_str(db)+' ADD COLUMN '+col.to_str(db)
1035
        if comment != None: str_ += ' '+sql_gen.esc_comment(comment)
1036
        
1037
        try:
1038
            run_query(db, str_, recover=True, cacheable=True, **kw_args)
1039
            break
1040
        except DuplicateException:
1041
            col.name = next_version(col.name)
1042
            # try again with next version of name
1043

    
1044
def add_not_null(db, col):
1045
    table = col.table
1046
    col = sql_gen.to_name_only_col(col)
1047
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ALTER COLUMN '
1048
        +col.to_str(db)+' SET NOT NULL', cacheable=True, log_level=3)
1049

    
1050
row_num_col = '_row_num'
1051

    
1052
row_num_typed_col = sql_gen.TypedCol(row_num_col, 'serial', nullable=False,
1053
    constraints='PRIMARY KEY')
1054

    
1055
def add_row_num(db, table):
1056
    '''Adds a row number column to a table. Its name is in row_num_col. It will
1057
    be the primary key.'''
1058
    add_col(db, table, row_num_typed_col, log_level=3)
1059

    
1060
#### Indexes
1061

    
1062
def add_pkey(db, table, cols=None, recover=None):
1063
    '''Adds a primary key.
1064
    @param cols [sql_gen.Col,...] The columns in the primary key.
1065
        Defaults to the first column in the table.
1066
    @pre The table must not already have a primary key.
1067
    '''
1068
    table = sql_gen.as_Table(table)
1069
    if cols == None: cols = [pkey(db, table, recover)]
1070
    col_strs = [sql_gen.to_name_only_col(v).to_str(db) for v in cols]
1071
    
1072
    run_query(db, 'ALTER TABLE '+table.to_str(db)+' ADD PRIMARY KEY ('
1073
        +(', '.join(col_strs))+')', recover=True, cacheable=True, log_level=3,
1074
        log_ignore_excs=(DuplicateException,))
1075

    
1076
def add_index(db, exprs, table=None, unique=False, ensure_not_null_=True):
1077
    '''Adds an index on column(s) or expression(s) if it doesn't already exist.
1078
    Currently, only function calls are supported as expressions.
1079
    @param ensure_not_null_ If set, translates NULL values to sentinel values.
1080
        This allows indexes to be used for comparisons where NULLs are equal.
1081
    '''
1082
    exprs = lists.mk_seq(exprs)
1083
    
1084
    # Parse exprs
1085
    old_exprs = exprs[:]
1086
    exprs = []
1087
    cols = []
1088
    for i, expr in enumerate(old_exprs):
1089
        expr = sql_gen.as_Col(expr, table)
1090
        
1091
        # Handle nullable columns
1092
        if ensure_not_null_:
1093
            try: expr = ensure_not_null(db, expr)
1094
            except KeyError: pass # unknown type, so just create plain index
1095
        
1096
        # Extract col
1097
        expr = copy.deepcopy(expr) # don't modify input!
1098
        if isinstance(expr, sql_gen.FunctionCall):
1099
            col = expr.args[0]
1100
            expr = sql_gen.Expr(expr)
1101
        else: col = expr
1102
        assert isinstance(col, sql_gen.Col)
1103
        
1104
        # Extract table
1105
        if table == None:
1106
            assert sql_gen.is_table_col(col)
1107
            table = col.table
1108
        
1109
        col.table = None
1110
        
1111
        exprs.append(expr)
1112
        cols.append(col)
1113
    
1114
    table = sql_gen.as_Table(table)
1115
    index = sql_gen.Table(str(sql_gen.Col(','.join(map(str, cols)), table)))
1116
    
1117
    # Add index
1118
    while True:
1119
        str_ = 'CREATE'
1120
        if unique: str_ += ' UNIQUE'
1121
        str_ += ' INDEX '+index.to_str(db)+' ON '+table.to_str(db)+' ('+(
1122
            ', '.join((v.to_str(db) for v in exprs)))+')'
1123
        
1124
        try:
1125
            run_query(db, str_, recover=True, cacheable=True, log_level=3,
1126
                log_ignore_excs=(DuplicateException,))
1127
            break
1128
        except DuplicateException:
1129
            index.name = next_version(index.name)
1130
            # try again with next version of name
1131

    
1132
def add_index_col(db, col, suffix, expr, nullable=True):
1133
    if sql_gen.index_col(col) != None: return # already has index col
1134
    
1135
    new_col = sql_gen.suffixed_col(col, suffix)
1136
    
1137
    # Add column
1138
    new_typed_col = sql_gen.TypedCol(new_col.name, db.col_info(col).type)
1139
    add_col(db, col.table, new_typed_col, comment='src: '+repr(col),
1140
        log_level=3)
1141
    new_col.name = new_typed_col.name # propagate any renaming
1142
    
1143
    update(db, col.table, [(new_col, expr)], in_place=True, cacheable=True,
1144
        log_level=3)
1145
    if not nullable: add_not_null(db, new_col)
1146
    add_index(db, new_col)
1147
    
1148
    col.table.index_cols[col.name] = new_col.name
1149

    
1150
# Controls when ensure_not_null() will use index columns
1151
not_null_index_cols_min_rows = 0 # rows; initially always use index columns
1152

    
1153
def ensure_not_null(db, col):
1154
    '''For params, see sql_gen.ensure_not_null()'''
1155
    expr = sql_gen.ensure_not_null(db, col)
1156
    
1157
    # If a nullable column in a temp table, add separate index column instead.
1158
    # Note that for small datasources, this adds 6-25% to the total import time.
1159
    if (sql_gen.is_temp_col(col) and isinstance(expr, sql_gen.EnsureNotNull)
1160
        and table_row_count(db, col.table) >= not_null_index_cols_min_rows):
1161
        add_index_col(db, col, '::NOT NULL', expr, nullable=False)
1162
        expr = sql_gen.index_col(col)
1163
    
1164
    return expr
1165

    
1166
already_indexed = object() # tells add_indexes() the pkey has already been added
1167

    
1168
def add_indexes(db, table, has_pkey=True):
1169
    '''Adds an index on all columns in a table.
1170
    @param has_pkey bool|already_indexed Whether a pkey instead of a regular
1171
        index should be added on the first column.
1172
        * If already_indexed, the pkey is assumed to have already been added
1173
    '''
1174
    cols = table_cols(db, table)
1175
    if has_pkey:
1176
        if has_pkey is not already_indexed: add_pkey(db, table)
1177
        cols = cols[1:]
1178
    for col in cols: add_index(db, col, table)
1179

    
1180
#### Tables
1181

    
1182
### Maintenance
1183

    
1184
def analyze(db, table):
1185
    table = sql_gen.as_Table(table)
1186
    run_query(db, 'ANALYZE '+table.to_str(db), log_level=3)
1187

    
1188
def autoanalyze(db, table):
1189
    if db.autoanalyze: analyze(db, table)
1190

    
1191
def vacuum(db, table):
1192
    table = sql_gen.as_Table(table)
1193
    db.with_autocommit(lambda: run_query(db, 'VACUUM ANALYZE '+table.to_str(db),
1194
        log_level=3))
1195

    
1196
### Lifecycle
1197

    
1198
def drop_table(db, table):
1199
    table = sql_gen.as_Table(table)
1200
    return run_query(db, 'DROP TABLE IF EXISTS '+table.to_str(db)+' CASCADE')
1201

    
1202
def create_table(db, table, cols=[], has_pkey=True, col_indexes=True,
1203
    like=None):
1204
    '''Creates a table.
1205
    @param cols [sql_gen.TypedCol,...] The column names and types
1206
    @param has_pkey If set, the first column becomes the primary key.
1207
    @param col_indexes bool|[ref]
1208
        * If True, indexes will be added on all non-pkey columns.
1209
        * If a list reference, [0] will be set to a function to do this.
1210
          This can be used to delay index creation until the table is populated.
1211
    '''
1212
    table = sql_gen.as_Table(table)
1213
    
1214
    if like != None:
1215
        cols = [sql_gen.CustomCode('LIKE '+like.to_str(db)+' INCLUDING ALL')
1216
            ]+cols
1217
    if has_pkey:
1218
        cols[0] = pkey = copy.copy(cols[0]) # don't modify input!
1219
        pkey.constraints = 'PRIMARY KEY'
1220
    
1221
    temp = table.is_temp and not db.debug_temp
1222
        # temp tables permanent in debug_temp mode
1223
    
1224
    # Create table
1225
    while True:
1226
        str_ = 'CREATE'
1227
        if temp: str_ += ' TEMP'
1228
        str_ += ' TABLE '+table.to_str(db)+' (\n'
1229
        str_ += '\n, '.join(c.to_str(db) for c in cols)
1230
        str_ += '\n);'
1231
        
1232
        try:
1233
            run_query(db, str_, recover=True, cacheable=True, log_level=2,
1234
                log_ignore_excs=(DuplicateException,))
1235
            break
1236
        except DuplicateException:
1237
            table.name = next_version(table.name)
1238
            # try again with next version of name
1239
    
1240
    # Add indexes
1241
    if has_pkey: has_pkey = already_indexed
1242
    def add_indexes_(): add_indexes(db, table, has_pkey)
1243
    if isinstance(col_indexes, list): col_indexes[0] = add_indexes_ # defer
1244
    elif col_indexes: add_indexes_() # add now
1245

    
1246
def copy_table_struct(db, src, dest):
1247
    '''Creates a structure-only copy of a table. (Does not copy data.)'''
1248
    create_table(db, dest, has_pkey=False, col_indexes=False, like=src)
1249

    
1250
### Data
1251

    
1252
def truncate(db, table, schema='public', **kw_args):
1253
    '''For params, see run_query()'''
1254
    table = sql_gen.as_Table(table, schema)
1255
    return run_query(db, 'TRUNCATE '+table.to_str(db)+' CASCADE', **kw_args)
1256

    
1257
def empty_temp(db, tables):
1258
    if db.debug_temp: return # leave temp tables there for debugging
1259
    tables = lists.mk_seq(tables)
1260
    for table in tables: truncate(db, table, log_level=3)
1261

    
1262
def empty_db(db, schema='public', **kw_args):
1263
    '''For kw_args, see tables()'''
1264
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
1265

    
1266
def distinct_table(db, table, distinct_on):
1267
    '''Creates a copy of a temp table which is distinct on the given columns.
1268
    The old and new tables will both get an index on these columns, to
1269
    facilitate merge joins.
1270
    @param distinct_on If empty, creates a table with one row. This is useful if
1271
        your distinct_on columns are all literal values.
1272
    @return The new table.
1273
    '''
1274
    new_table = sql_gen.suffixed_table(table, '_distinct')
1275
    
1276
    copy_table_struct(db, table, new_table)
1277
    
1278
    limit = None
1279
    if distinct_on == []: limit = 1 # one sample row
1280
    else:
1281
        add_index(db, distinct_on, new_table, unique=True)
1282
        add_index(db, distinct_on, table) # for join optimization
1283
    
1284
    insert_select(db, new_table, None, mk_select(db, table, start=0,
1285
        limit=limit), ignore=True)
1286
    analyze(db, new_table)
1287
    
1288
    return new_table
(24-24/37)