Project

General

Profile

1
# Database access
2

    
3
import copy
4
import operator
5
import re
6
import warnings
7

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

    
17
##### Exceptions
18

    
19
def get_cur_query(cur, input_query=None, input_params=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 repr(input_query)+' % '+repr(input_params)
26

    
27
def _add_cursor_info(e, *args, **kw_args):
28
    '''For params, see get_cur_query()'''
29
    exc.add_msg(e, 'query: '+str(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: '+str(name), cause)
39
        self.name = name
40

    
41
class ExceptionWithColumns(DbException):
42
    def __init__(self, cols, cause=None):
43
        DbException.__init__(self, 'for columns: '+(', '.join(cols)), cause)
44
        self.cols = cols
45

    
46
class NameException(DbException): pass
47

    
48
class DuplicateKeyException(ExceptionWithColumns): pass
49

    
50
class NullValueException(ExceptionWithColumns): pass
51

    
52
class DuplicateTableException(ExceptionWithName): pass
53

    
54
class EmptyRowException(DbException): pass
55

    
56
##### Warnings
57

    
58
class DbWarning(UserWarning): pass
59

    
60
##### Result retrieval
61

    
62
def col_names(cur): return (col[0] for col in cur.description)
63

    
64
def rows(cur): return iter(lambda: cur.fetchone(), None)
65

    
66
def consume_rows(cur):
67
    '''Used to fetch all rows so result will be cached'''
68
    iters.consume_iter(rows(cur))
69

    
70
def next_row(cur): return rows(cur).next()
71

    
72
def row(cur):
73
    row_ = next_row(cur)
74
    consume_rows(cur)
75
    return row_
76

    
77
def next_value(cur): return next_row(cur)[0]
78

    
79
def value(cur): return row(cur)[0]
80

    
81
def values(cur): return iters.func_iter(lambda: next_value(cur))
82

    
83
def value_or_none(cur):
84
    try: return value(cur)
85
    except StopIteration: return None
86

    
87
##### Input validation
88

    
89
def clean_name(name): return re.sub(r'\W', r'', name)
90

    
91
def check_name(name):
92
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
93
        +'" may contain only alphanumeric characters and _')
94

    
95
def esc_name_by_module(module, name, ignore_case=False):
96
    if module == 'psycopg2':
97
        if ignore_case:
98
            # Don't enclose in quotes because this disables case-insensitivity
99
            check_name(name)
100
            return name
101
        else: quote = '"'
102
    elif module == 'MySQLdb': quote = '`'
103
    else: raise NotImplementedError("Can't escape name for "+module+' database')
104
    return quote + name.replace(quote, '') + quote
105

    
106
def esc_name_by_engine(engine, name, **kw_args):
107
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
108

    
109
def esc_name(db, name, **kw_args):
110
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
111

    
112
def qual_name(db, schema, table):
113
    def esc_name_(name): return esc_name(db, name)
114
    table = esc_name_(table)
115
    if schema != None: return esc_name_(schema)+'.'+table
116
    else: return table
117

    
118
##### Database connections
119

    
120
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
121

    
122
db_engines = {
123
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
124
    'PostgreSQL': ('psycopg2', {}),
125
}
126

    
127
DatabaseErrors_set = set([DbException])
128
DatabaseErrors = tuple(DatabaseErrors_set)
129

    
130
def _add_module(module):
131
    DatabaseErrors_set.add(module.DatabaseError)
132
    global DatabaseErrors
133
    DatabaseErrors = tuple(DatabaseErrors_set)
134

    
135
def db_config_str(db_config):
136
    return db_config['engine']+' database '+db_config['database']
137

    
138
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
139

    
140
log_debug_none = lambda msg: None
141

    
142
class DbConn:
143
    def __init__(self, db_config, serializable=True, log_debug=log_debug_none,
144
        caching=True):
145
        self.db_config = db_config
146
        self.serializable = serializable
147
        self.log_debug = log_debug
148
        self.caching = caching
149
        
150
        self.__db = None
151
        self.query_results = {}
152
        self._savepoint = 0
153
    
154
    def __getattr__(self, name):
155
        if name == '__dict__': raise Exception('getting __dict__')
156
        if name == 'db': return self._db()
157
        else: raise AttributeError()
158
    
159
    def __getstate__(self):
160
        state = copy.copy(self.__dict__) # shallow copy
161
        state['log_debug'] = None # don't pickle the debug callback
162
        state['_DbConn__db'] = None # don't pickle the connection
163
        return state
164
    
165
    def connected(self): return self.__db != None
166
    
167
    def _db(self):
168
        if self.__db == None:
169
            # Process db_config
170
            db_config = self.db_config.copy() # don't modify input!
171
            schemas = db_config.pop('schemas', None)
172
            module_name, mappings = db_engines[db_config.pop('engine')]
173
            module = __import__(module_name)
174
            _add_module(module)
175
            for orig, new in mappings.iteritems():
176
                try: util.rename_key(db_config, orig, new)
177
                except KeyError: pass
178
            
179
            # Connect
180
            self.__db = module.connect(**db_config)
181
            
182
            # Configure connection
183
            if self.serializable: run_raw_query(self,
184
                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
185
            if schemas != None:
186
                schemas_ = ''.join((esc_name(self, s)+', '
187
                    for s in schemas.split(',')))
188
                run_raw_query(self, "SELECT set_config('search_path', \
189
%s || current_setting('search_path'), false)", [schemas_])
190
        
191
        return self.__db
192
    
193
    class DbCursor(Proxy):
194
        def __init__(self, outer):
195
            Proxy.__init__(self, outer.db.cursor())
196
            self.query_results = outer.query_results
197
            self.query_lookup = None
198
            self.result = []
199
        
200
        def execute(self, query, params=None):
201
            self._is_insert = query.upper().find('INSERT') >= 0
202
            self.query_lookup = _query_lookup(query, params)
203
            try:
204
                try: return_value = self.inner.execute(query, params)
205
                finally: self.query = get_cur_query(self.inner)
206
            except Exception, e:
207
                _add_cursor_info(e, self, query, params)
208
                self.result = e # cache the exception as the result
209
                self._cache_result()
210
                raise
211
            # Fetch all rows so result will be cached
212
            if self.rowcount == 0 and not self._is_insert: consume_rows(self)
213
            return return_value
214
        
215
        def fetchone(self):
216
            row = self.inner.fetchone()
217
            if row != None: self.result.append(row)
218
            # otherwise, fetched all rows
219
            else: self._cache_result()
220
            return row
221
        
222
        def _cache_result(self):
223
            # For inserts, only cache exceptions since inserts are not
224
            # idempotent, but an invalid insert will always be invalid
225
            if self.query_results != None and (not self._is_insert
226
                or isinstance(self.result, Exception)):
227
                
228
                assert self.query_lookup != None
229
                self.query_results[self.query_lookup] = self.CacheCursor(
230
                    util.dict_subset(dicts.AttrsDictView(self),
231
                    ['query', 'result', 'rowcount', 'description']))
232
        
233
        class CacheCursor:
234
            def __init__(self, cached_result): self.__dict__ = cached_result
235
            
236
            def execute(self, *args, **kw_args):
237
                if isinstance(self.result, Exception): raise self.result
238
                # otherwise, result is a rows list
239
                self.iter = iter(self.result)
240
            
241
            def fetchone(self):
242
                try: return self.iter.next()
243
                except StopIteration: return None
244
    
245
    def run_query(self, query, params=None, cacheable=False):
246
        '''Translates known DB errors to typed exceptions:
247
        See self.DbCursor.execute().'''
248
        assert query != None
249
        
250
        if not self.caching: cacheable = False
251
        used_cache = False
252
        try:
253
            # Get cursor
254
            if cacheable:
255
                query_lookup = _query_lookup(query, params)
256
                try:
257
                    cur = self.query_results[query_lookup]
258
                    used_cache = True
259
                except KeyError: cur = self.DbCursor(self)
260
            else: cur = self.db.cursor()
261
            
262
            # Run query
263
            cur.execute(query, params)
264
        finally:
265
            if self.log_debug != log_debug_none: # only compute msg if needed
266
                if used_cache: cache_status = 'Cache hit'
267
                elif cacheable: cache_status = 'Cache miss'
268
                else: cache_status = 'Non-cacheable'
269
                self.log_debug(cache_status+': '
270
                    +strings.one_line(str(get_cur_query(cur, query, params))))
271
        
272
        return cur
273
    
274
    def is_cached(self, query, params=None):
275
        return _query_lookup(query, params) in self.query_results
276
    
277
    def with_savepoint(self, func):
278
        savepoint = 'level_'+str(self._savepoint)
279
        self.run_query('SAVEPOINT '+savepoint)
280
        self._savepoint += 1
281
        try: 
282
            try: return_val = func()
283
            finally:
284
                self._savepoint -= 1
285
                assert self._savepoint >= 0
286
        except:
287
            self.run_query('ROLLBACK TO SAVEPOINT '+savepoint)
288
            raise
289
        else:
290
            self.run_query('RELEASE SAVEPOINT '+savepoint)
291
            return return_val
292

    
293
connect = DbConn
294

    
295
##### Querying
296

    
297
def run_raw_query(db, *args, **kw_args):
298
    '''For params, see DbConn.run_query()'''
299
    return db.run_query(*args, **kw_args)
300

    
301
def mogrify(db, query, params):
302
    module = util.root_module(db.db)
303
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
304
    else: raise NotImplementedError("Can't mogrify query for "+module+
305
        ' database')
306

    
307
##### Recoverable querying
308

    
309
def with_savepoint(db, func): return db.with_savepoint(func)
310

    
311
def run_query(db, query, params=None, recover=None, cacheable=False):
312
    if recover == None: recover = False
313
    
314
    try:
315
        def run(): return run_raw_query(db, query, params, cacheable)
316
        if recover and not db.is_cached(query, params):
317
            return with_savepoint(db, run)
318
        else: return run() # don't need savepoint if cached
319
    except Exception, e:
320
        if not recover: raise # need savepoint to run index_cols()
321
        msg = str(e)
322
        match = re.search(r'duplicate key value violates unique constraint '
323
            r'"((_?[^\W_]+)_[^"]+)"', msg)
324
        if match:
325
            constraint, table = match.groups()
326
            try: cols = index_cols(db, table, constraint)
327
            except NotImplementedError: raise e
328
            else: raise DuplicateKeyException(cols, e)
329
        match = re.search(r'null value in column "(\w+)" violates not-null '
330
            'constraint', msg)
331
        if match: raise NullValueException([match.group(1)], e)
332
        match = re.search(r'relation "(\w+)" already exists', msg)
333
        if match: raise DuplicateTableException(match.group(1), e)
334
        raise # no specific exception raised
335

    
336
##### Basic queries
337

    
338
def next_version(name):
339
    '''Prepends the version # so it won't be removed if the name is truncated'''
340
    version = 1 # first existing name was version 0
341
    match = re.match(r'^v(\d+)_(.*)$', name)
342
    if match:
343
        version = int(match.group(1))+1
344
        name = match.group(2)
345
    return 'v'+str(version)+'_'+name
346

    
347
def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
348
    '''Outputs a query to a temp table.
349
    For params, see run_query().
350
    '''
351
    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
352
    else: # place rows in temp table
353
        check_name(into_ref[0])
354
        kw_args['recover'] = True
355
        while True:
356
            try:
357
                return run_query(db, 'CREATE TEMP TABLE '+into_ref[0]+' AS '
358
                    +query, params, *args, **kw_args)
359
                    # CREATE TABLE AS sets rowcount to # rows in query
360
            except DuplicateTableException, e:
361
                into_ref[0] = next_version(into_ref[0])
362
                # try again with next version of name
363

    
364
order_by_pkey = object() # tells mk_select() to order by the pkey
365

    
366
join_using = object() # tells mk_select() to join the column with USING
367

    
368
def mk_select(db, tables, fields=None, conds=None, limit=None, start=None,
369
    order_by=order_by_pkey, table_is_esc=False):
370
    '''
371
    @param tables The single table to select from, or a list of tables to join
372
        together: [table0, (table1, dict(right_col=left_col, ...)), ...]
373
    @param fields Use None to select all fields in the table
374
    @param table_is_esc Whether the table name has already been escaped
375
    @return tuple(query, params)
376
    '''
377
    def esc_name_(name): return esc_name(db, name)
378
    
379
    if not lists.is_seq(tables): tables = [tables]
380
    tables = list(tables) # don't modify input! (list() copies input)
381
    table0 = tables.pop(0) # first table is separate
382
    
383
    if conds == None: conds = {}
384
    assert limit == None or type(limit) == int
385
    assert start == None or type(start) == int
386
    if order_by == order_by_pkey:
387
        order_by = pkey(db, table0, recover=True, table_is_esc=table_is_esc)
388
    if not table_is_esc: table0 = esc_name_(table0)
389
    
390
    params = []
391
    
392
    def parse_col(field, default_table=None):
393
        '''Parses fields'''
394
        if field == None: field = (field,) # for None values, tuple is optional
395
        is_tuple = isinstance(field, tuple)
396
        if is_tuple and len(field) == 1: # field is literal value
397
            value, = field
398
            sql_ = '%s'
399
            params.append(value)
400
        elif is_tuple and len(field) == 2: # field is col with table
401
            table, col = field
402
            if not table_is_esc: table = esc_name_(table)
403
            sql_ = table+'.'+esc_name_(col)
404
        else:
405
            sql_ = esc_name_(field) # field is col name
406
            if default_table != None: sql_ = default_table+'.'+sql_
407
        return sql_
408
    def cond(entry):
409
        '''Parses conditions'''
410
        col, value = entry
411
        cond_ = esc_name_(col)+' '
412
        if value == None: cond_ += 'IS'
413
        else: cond_ += '='
414
        cond_ += ' %s'
415
        return cond_
416
    
417
    query = 'SELECT '
418
    if fields == None: query += '*'
419
    else: query += ', '.join(map(parse_col, fields))
420
    query += ' FROM '+table0
421
    
422
    # Add joins
423
    left_table = table0
424
    for table, joins in tables:
425
        if not table_is_esc: table = esc_name_(table)
426
        query += ' JOIN '+table
427
        
428
        def join(entry):
429
            '''Parses non-USING joins'''
430
            right_col, left_col = entry
431
            right_col = table+'.'+esc_name_(right_col)
432
            sql_ = right_col+' '
433
            
434
            if left_col == None: left_col = (left_col,)
435
                # for None values, tuple is optional
436
            if isinstance(left_col, tuple) and len(left_col) == 1:
437
                # col is literal value
438
                value, = left_col
439
                if value == None: sql_ += 'IS'
440
                else: sql_ += '='
441
                sql_ += ' %s'
442
                params.append(value)
443
            else: # col is name
444
                left_col = parse_col(left_col, left_table)
445
                sql_ += ('= '+left_col+' OR ('+right_col+' IS NULL AND '
446
                    +left_col+' IS NULL)')
447
            
448
            return sql_
449
        
450
        if reduce(operator.and_, (v == join_using for v in joins.itervalues())):
451
            # all cols w/ USING
452
            query += ' USING ('+(', '.join(joins.iterkeys()))+')'
453
        else: query += ' ON '+(' AND '.join(map(join, joins.iteritems())))
454
        
455
        left_table = table
456
    
457
    missing = True
458
    if conds != {}:
459
        query += ' WHERE '+(' AND '.join(map(cond, conds.iteritems())))
460
        params += conds.values()
461
        missing = False
462
    if order_by != None: query += ' ORDER BY '+esc_name_(order_by)
463
    if limit != None: query += ' LIMIT '+str(limit); missing = False
464
    if start != None:
465
        if start != 0: query += ' OFFSET '+str(start)
466
        missing = False
467
    if missing: warnings.warn(DbWarning(
468
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
469
    
470
    return (query, params)
471

    
472
def select(db, *args, **kw_args):
473
    '''For params, see mk_select() and run_query()'''
474
    recover = kw_args.pop('recover', None)
475
    cacheable = kw_args.pop('cacheable', True)
476
    
477
    query, params = mk_select(db, *args, **kw_args)
478
    return run_query(db, query, params, recover, cacheable)
479

    
480
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
481
    returning=None, embeddable=False, table_is_esc=False):
482
    '''
483
    @param returning str|None An inserted column (such as pkey) to return
484
    @param embeddable Whether the query should be embeddable as a nested SELECT.
485
        Warning: If you set this and cacheable=True when the query is run, the
486
        query will be fully cached, not just if it raises an exception.
487
    @param table_is_esc Whether the table name has already been escaped
488
    '''
489
    if select_query == None: select_query = 'DEFAULT VALUES'
490
    if cols == []: cols = None # no cols (all defaults) = unknown col names
491
    if not table_is_esc: check_name(table)
492
    
493
    # Build query
494
    query = 'INSERT INTO '+table
495
    if cols != None:
496
        map(check_name, cols)
497
        query += ' ('+', '.join(cols)+')'
498
    query += ' '+select_query
499
    
500
    if returning != None:
501
        check_name(returning)
502
        query += ' RETURNING '+returning
503
    
504
    if embeddable:
505
        # Create function
506
        function = 'pg_temp.'+('_'.join(map(clean_name,
507
            ['insert', table] + cols)))
508
        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
509
        function_query = '''\
510
CREATE OR REPLACE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
511
    LANGUAGE sql
512
    AS $$'''+mogrify(db, query, params)+''';$$;
513
'''
514
        run_query(db, function_query)
515
            # don't cache because function def could change and then change back
516
        
517
        # Return query that uses function
518
        return mk_select(db, function+'() AS f ('+returning+')', start=0,
519
            order_by=None, table_is_esc=True)# AS clause requires function alias
520
    
521
    return (query, params)
522

    
523
def insert_select(db, *args, **kw_args):
524
    '''For params, see mk_insert_select() and run_query_into()
525
    @param into_ref List with name of temp table to place RETURNING values in
526
    '''
527
    into_ref = kw_args.pop('into_ref', None)
528
    if into_ref != None: kw_args['embeddable'] = True
529
    recover = kw_args.pop('recover', None)
530
    cacheable = kw_args.pop('cacheable', True)
531
    
532
    query, params = mk_insert_select(db, *args, **kw_args)
533
    return run_query_into(db, query, params, into_ref, recover=recover,
534
        cacheable=cacheable)
535

    
536
default = object() # tells insert() to use the default value for a column
537

    
538
def insert(db, table, row, *args, **kw_args):
539
    '''For params, see insert_select()'''
540
    if lists.is_seq(row): cols = None
541
    else:
542
        cols = row.keys()
543
        row = row.values()
544
    row = list(row) # ensure that "!= []" works
545
    
546
    # Check for special values
547
    labels = []
548
    values = []
549
    for value in row:
550
        if value == default: labels.append('DEFAULT')
551
        else:
552
            labels.append('%s')
553
            values.append(value)
554
    
555
    # Build query
556
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
557
    else: query = None
558
    
559
    return insert_select(db, table, cols, query, values, *args, **kw_args)
560

    
561
def last_insert_id(db):
562
    module = util.root_module(db.db)
563
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
564
    elif module == 'MySQLdb': return db.insert_id()
565
    else: return None
566

    
567
def truncate(db, table, schema='public'):
568
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
569

    
570
##### Database structure queries
571

    
572
def pkey(db, table, recover=None, table_is_esc=False):
573
    '''Assumed to be first column in table'''
574
    return col_names(select(db, table, limit=0, order_by=None, recover=recover,
575
        table_is_esc=table_is_esc)).next()
576

    
577
def index_cols(db, table, index):
578
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
579
    automatically created. When you don't know whether something is a UNIQUE
580
    constraint or a UNIQUE index, use this function.'''
581
    check_name(table)
582
    check_name(index)
583
    module = util.root_module(db.db)
584
    if module == 'psycopg2':
585
        return list(values(run_query(db, '''\
586
SELECT attname
587
FROM
588
(
589
        SELECT attnum, attname
590
        FROM pg_index
591
        JOIN pg_class index ON index.oid = indexrelid
592
        JOIN pg_class table_ ON table_.oid = indrelid
593
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
594
        WHERE
595
            table_.relname = %(table)s
596
            AND index.relname = %(index)s
597
    UNION
598
        SELECT attnum, attname
599
        FROM
600
        (
601
            SELECT
602
                indrelid
603
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
604
                    AS indkey
605
            FROM pg_index
606
            JOIN pg_class index ON index.oid = indexrelid
607
            JOIN pg_class table_ ON table_.oid = indrelid
608
            WHERE
609
                table_.relname = %(table)s
610
                AND index.relname = %(index)s
611
        ) s
612
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
613
) s
614
ORDER BY attnum
615
''',
616
            {'table': table, 'index': index}, cacheable=True)))
617
    else: raise NotImplementedError("Can't list index columns for "+module+
618
        ' database')
619

    
620
def constraint_cols(db, table, constraint):
621
    check_name(table)
622
    check_name(constraint)
623
    module = util.root_module(db.db)
624
    if module == 'psycopg2':
625
        return list(values(run_query(db, '''\
626
SELECT attname
627
FROM pg_constraint
628
JOIN pg_class ON pg_class.oid = conrelid
629
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
630
WHERE
631
    relname = %(table)s
632
    AND conname = %(constraint)s
633
ORDER BY attnum
634
''',
635
            {'table': table, 'constraint': constraint})))
636
    else: raise NotImplementedError("Can't list constraint columns for "+module+
637
        ' database')
638

    
639
row_num_col = '_row_num'
640

    
641
def add_row_num(db, table):
642
    '''Adds a row number column to a table. Its name is in row_num_col. It will
643
    be the primary key.'''
644
    check_name(table)
645
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
646
        +' serial NOT NULL PRIMARY KEY')
647

    
648
def tables(db, schema='public', table_like='%'):
649
    module = util.root_module(db.db)
650
    params = {'schema': schema, 'table_like': table_like}
651
    if module == 'psycopg2':
652
        return values(run_query(db, '''\
653
SELECT tablename
654
FROM pg_tables
655
WHERE
656
    schemaname = %(schema)s
657
    AND tablename LIKE %(table_like)s
658
ORDER BY tablename
659
''',
660
            params, cacheable=True))
661
    elif module == 'MySQLdb':
662
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
663
            cacheable=True))
664
    else: raise NotImplementedError("Can't list tables for "+module+' database')
665

    
666
##### Database management
667

    
668
def empty_db(db, schema='public', **kw_args):
669
    '''For kw_args, see tables()'''
670
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
671

    
672
##### Heuristic queries
673

    
674
def put(db, table, row, pkey_=None, row_ct_ref=None):
675
    '''Recovers from errors.
676
    Only works under PostgreSQL (uses INSERT RETURNING).
677
    '''
678
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
679
    
680
    try:
681
        cur = insert(db, table, row, pkey_, recover=True)
682
        if row_ct_ref != None and cur.rowcount >= 0:
683
            row_ct_ref[0] += cur.rowcount
684
        return value(cur)
685
    except DuplicateKeyException, e:
686
        return value(select(db, table, [pkey_],
687
            util.dict_subset_right_join(row, e.cols), recover=True))
688

    
689
def get(db, table, row, pkey, row_ct_ref=None, create=False):
690
    '''Recovers from errors'''
691
    try: return value(select(db, table, [pkey], row, 1, recover=True))
692
    except StopIteration:
693
        if not create: raise
694
        return put(db, table, row, pkey, row_ct_ref) # insert new row
695

    
696
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
697
    row_ct_ref=None, table_is_esc=False):
698
    '''Recovers from errors.
699
    Only works under PostgreSQL (uses INSERT RETURNING).
700
    @param in_tables The main input table to select from, followed by a list of
701
        tables to join with it using the main input table's pkey
702
    @return (table, col) Where the pkeys (from INSERT RETURNING) are made
703
        available
704
    '''
705
    temp_suffix = clean_name(out_table)
706
        # suffix, not prefix, so main name won't be removed if name is truncated
707
    pkeys_ref = ['pkeys_'+temp_suffix]
708
    
709
    # Join together input tables
710
    in_tables = in_tables[:] # don't modify input!
711
    in_tables0 = in_tables.pop(0) # first table is separate
712
    in_pkey = pkey(db, in_tables0, recover=True, table_is_esc=table_is_esc)
713
    insert_joins = [in_tables0]+[(t, {in_pkey: join_using}) for t in in_tables]
714
    
715
    out_pkey = pkey(db, out_table, recover=True, table_is_esc=table_is_esc)
716
    pkeys_cols = [in_pkey, out_pkey]
717
    
718
    def mk_select_(cols):
719
        return mk_select(db, insert_joins, cols, limit=limit, start=start,
720
            table_is_esc=table_is_esc)
721
    
722
    out_pkeys_ref = ['out_pkeys_'+temp_suffix]
723
    def insert_():
724
        '''Inserts and capture output pkeys.'''
725
        cur = insert_select(db, out_table, mapping.keys(),
726
            *mk_select_(mapping.values()), returning=out_pkey,
727
            into_ref=out_pkeys_ref, recover=True, table_is_esc=table_is_esc)
728
        if row_ct_ref != None and cur.rowcount >= 0:
729
            row_ct_ref[0] += cur.rowcount
730
            add_row_num(db, out_pkeys_ref[0]) # for joining it with input pkeys
731
        
732
        # Get input pkeys corresponding to rows in insert
733
        in_pkeys_ref = ['in_pkeys_'+temp_suffix]
734
        run_query_into(db, *mk_select_([in_pkey]), into_ref=in_pkeys_ref)
735
        add_row_num(db, in_pkeys_ref[0]) # for joining it with output pkeys
736
        
737
        # Join together output and input pkeys
738
        run_query_into(db, *mk_select(db,
739
            [in_pkeys_ref[0], (out_pkeys_ref[0], {row_num_col: join_using})],
740
            pkeys_cols, start=0), into_ref=pkeys_ref)
741
    
742
    # Do inserts and selects
743
    try: insert_()
744
    except DuplicateKeyException, e:
745
        join_cols = util.dict_subset_right_join(mapping, e.cols)
746
        select_joins = insert_joins + [(out_table, join_cols)]
747
        
748
        # Get pkeys of already existing rows
749
        run_query_into(db, *mk_select(db, select_joins, pkeys_cols, start=0,
750
            table_is_esc=table_is_esc), into_ref=pkeys_ref, recover=True)
751
    
752
    return (pkeys_ref[0], out_pkey)
753

    
754
##### Data cleanup
755

    
756
def cleanup_table(db, table, cols, table_is_esc=False):
757
    def esc_name_(name): return esc_name(db, name)
758
    
759
    if not table_is_esc: check_name(table)
760
    cols = map(esc_name_, cols)
761
    
762
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
763
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
764
            for col in cols))),
765
        dict(null0='', null1=r'\N'))
(22-22/33)