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 sql_gen
15
import strings
16
import util
17

    
18
##### Exceptions
19

    
20
def get_cur_query(cur, input_query=None, input_params=None):
21
    raw_query = None
22
    if hasattr(cur, 'query'): raw_query = cur.query
23
    elif hasattr(cur, '_last_executed'): raw_query = cur._last_executed
24
    
25
    if raw_query != None: return raw_query
26
    else: return repr(input_query)+' % '+repr(input_params)
27

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

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

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

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

    
47
class NameException(DbException): pass
48

    
49
class DuplicateKeyException(ExceptionWithColumns): pass
50

    
51
class NullValueException(ExceptionWithColumns): pass
52

    
53
class DuplicateTableException(ExceptionWithName): pass
54

    
55
class DuplicateFunctionException(ExceptionWithName): pass
56

    
57
class EmptyRowException(DbException): pass
58

    
59
##### Warnings
60

    
61
class DbWarning(UserWarning): pass
62

    
63
##### Result retrieval
64

    
65
def col_names(cur): return (col[0] for col in cur.description)
66

    
67
def rows(cur): return iter(lambda: cur.fetchone(), None)
68

    
69
def consume_rows(cur):
70
    '''Used to fetch all rows so result will be cached'''
71
    iters.consume_iter(rows(cur))
72

    
73
def next_row(cur): return rows(cur).next()
74

    
75
def row(cur):
76
    row_ = next_row(cur)
77
    consume_rows(cur)
78
    return row_
79

    
80
def next_value(cur): return next_row(cur)[0]
81

    
82
def value(cur): return row(cur)[0]
83

    
84
def values(cur): return iters.func_iter(lambda: next_value(cur))
85

    
86
def value_or_none(cur):
87
    try: return value(cur)
88
    except StopIteration: return None
89

    
90
##### Input validation
91

    
92
def clean_name(name): return re.sub(r'\W', r'', name).lower()
93

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

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

    
109
def esc_name_by_engine(engine, name, **kw_args):
110
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
111

    
112
def esc_name(db, name, **kw_args):
113
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
114

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

    
121
##### Database connections
122

    
123
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
124

    
125
db_engines = {
126
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
127
    'PostgreSQL': ('psycopg2', {}),
128
}
129

    
130
DatabaseErrors_set = set([DbException])
131
DatabaseErrors = tuple(DatabaseErrors_set)
132

    
133
def _add_module(module):
134
    DatabaseErrors_set.add(module.DatabaseError)
135
    global DatabaseErrors
136
    DatabaseErrors = tuple(DatabaseErrors_set)
137

    
138
def db_config_str(db_config):
139
    return db_config['engine']+' database '+db_config['database']
140

    
141
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
142

    
143
log_debug_none = lambda msg: None
144

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

    
318
connect = DbConn
319

    
320
##### Querying
321

    
322
def run_raw_query(db, *args, **kw_args):
323
    '''For params, see DbConn.run_query()'''
324
    return db.run_query(*args, **kw_args)
325

    
326
def mogrify(db, query, params):
327
    module = util.root_module(db.db)
328
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
329
    else: raise NotImplementedError("Can't mogrify query for "+module+
330
        ' database')
331

    
332
##### Recoverable querying
333

    
334
def with_savepoint(db, func): return db.with_savepoint(func)
335

    
336
def run_query(db, query, params=None, recover=None, cacheable=False):
337
    if recover == None: recover = False
338
    
339
    try:
340
        def run(): return run_raw_query(db, query, params, cacheable)
341
        if recover and not db.is_cached(query, params):
342
            return with_savepoint(db, run)
343
        else: return run() # don't need savepoint if cached
344
    except Exception, e:
345
        if not recover: raise # need savepoint to run index_cols()
346
        msg = str(e)
347
        match = re.search(r'duplicate key value violates unique constraint '
348
            r'"((_?[^\W_]+)_[^"]+)"', msg)
349
        if match:
350
            constraint, table = match.groups()
351
            try: cols = index_cols(db, table, constraint)
352
            except NotImplementedError: raise e
353
            else: raise DuplicateKeyException(cols, e)
354
        match = re.search(r'null value in column "(\w+)" violates not-null '
355
            'constraint', msg)
356
        if match: raise NullValueException([match.group(1)], e)
357
        match = re.search(r'relation "(\w+)" already exists', msg)
358
        if match: raise DuplicateTableException(match.group(1), e)
359
        match = re.search(r'function "(\w+)" already exists', msg)
360
        if match: raise DuplicateFunctionException(match.group(1), e)
361
        raise # no specific exception raised
362

    
363
##### Basic queries
364

    
365
def next_version(name):
366
    '''Prepends the version # so it won't be removed if the name is truncated'''
367
    version = 1 # first existing name was version 0
368
    match = re.match(r'^v(\d+)_(.*)$', name)
369
    if match:
370
        version = int(match.group(1))+1
371
        name = match.group(2)
372
    return 'v'+str(version)+'_'+name
373

    
374
def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
375
    '''Outputs a query to a temp table.
376
    For params, see run_query().
377
    '''
378
    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
379
    else: # place rows in temp table
380
        check_name(into_ref[0])
381
        kw_args['recover'] = True
382
        while True:
383
            try:
384
                create_query = 'CREATE'
385
                if not db.debug: create_query += ' TEMP'
386
                create_query += ' TABLE '+into_ref[0]+' AS '+query
387
                
388
                return run_query(db, create_query, params, *args, **kw_args)
389
                    # CREATE TABLE AS sets rowcount to # rows in query
390
            except DuplicateTableException, e:
391
                into_ref[0] = next_version(into_ref[0])
392
                # try again with next version of name
393

    
394
order_by_pkey = object() # tells mk_select() to order by the pkey
395

    
396
join_using = object() # tells mk_select() to join the column with USING
397

    
398
filter_out = object() # tells mk_select() to filter out rows that match the join
399

    
400
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
401

    
402
def mk_select(db, tables, fields=None, conds=None, distinct_on=[], limit=None,
403
    start=None, order_by=order_by_pkey, table_is_esc=False, default_table=None):
404
    '''
405
    @param tables The single table to select from, or a list of tables to join
406
        together: [table0, (table1, joins), ...]
407
        
408
        joins has the format: dict(right_col=left_col, ...)
409
        * if left_col is join_using, left_col is set to right_col
410
        * if left_col is filter_out, the tables are LEFT JOINed together and the
411
          query is filtered by `right_col IS NULL` (indicating no match)
412
    @param fields Use None to select all fields in the table
413
    @param conds WHERE conditions: dict(compare_left_side=compare_right_side):
414
        * compare_left_side: Code|str (for col name)
415
        * compare_right_side: ValueCond|literal value
416
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
417
        use all columns
418
    @param table_is_esc Whether the table name has already been escaped
419
    @return tuple(query, params)
420
    '''
421
    def esc_name_(name): return esc_name(db, name)
422
    
423
    if not lists.is_seq(tables): tables = [tables]
424
    tables = list(tables) # don't modify input! (list() copies input)
425
    table0 = tables.pop(0) # first table is separate
426
    
427
    if conds == None: conds = {}
428
    assert limit == None or type(limit) == int
429
    assert start == None or type(start) == int
430
    if order_by == order_by_pkey:
431
        order_by = pkey(db, table0, recover=True, table_is_esc=table_is_esc)
432
    if not table_is_esc: table0 = esc_name_(table0)
433
    
434
    params = []
435
    
436
    def parse_col(field):
437
        '''Parses fields'''
438
        return sql_gen.value2sql_gen(field, default_table, table_is_esc,
439
            assume_col=True).to_str(db)
440
    def cond(entry):
441
        '''Parses conditions'''
442
        left, right = entry
443
        return sql_gen.cond2sql_gen(right).to_str(db, sql_gen.col2sql_gen(left))
444
    
445
    query = 'SELECT'
446
    
447
    # DISTINCT ON columns
448
    if distinct_on != []:
449
        query += ' DISTINCT'
450
        if distinct_on != distinct_on_all:
451
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
452
    
453
    # Columns
454
    query += ' '
455
    if fields == None: query += '*'
456
    else: query += ', '.join(map(parse_col, fields))
457
    
458
    # Main table
459
    query += ' FROM '+table0
460
    
461
    # Add joins
462
    left_table = table0
463
    for table, joins in tables:
464
        if not table_is_esc: table = esc_name_(table)
465
        
466
        left_join_ref = [False]
467
        
468
        def join(entry):
469
            '''Parses non-USING joins'''
470
            right_col, left_col = entry
471
            
472
            # Parse special values
473
            if left_col == None: left_col = (left_col,)
474
                # for None values, tuple is optional
475
            elif left_col == join_using: left_col = right_col
476
            elif left_col == filter_out:
477
                left_col = right_col
478
                left_join_ref[0] = True
479
                conds[(table, right_col)] = None # filter query by no match
480
            
481
            return cond(((table, right_col),
482
                sql_gen.value2sql_gen(left_col, left_table,
483
                    table_is_esc=table_is_esc, assume_col=True)))
484
        
485
        # Create join condition and determine join type
486
        if reduce(operator.and_, (v == join_using for v in joins.itervalues())):
487
            # all cols w/ USING, so can use simpler USING syntax
488
            join_cond = 'USING ('+(', '.join(joins.iterkeys()))+')'
489
        else: join_cond = 'ON '+(' AND '.join(map(join, joins.iteritems())))
490
        
491
        # Create join
492
        if left_join_ref[0]: query += ' LEFT'
493
        query += ' JOIN '+table+' '+join_cond
494
        
495
        left_table = table
496
    
497
    missing = True
498
    if conds != {}:
499
        query += ' WHERE '+(' AND '.join(map(cond, conds.iteritems())))
500
        params += conds.values()
501
        missing = False
502
    if order_by != None:
503
        query += ' ORDER BY '+sql_gen.col2sql_gen(order_by, table0,
504
            table_is_esc).to_str(db)
505
    if limit != None: query += ' LIMIT '+str(limit); missing = False
506
    if start != None:
507
        if start != 0: query += ' OFFSET '+str(start)
508
        missing = False
509
    if missing: warnings.warn(DbWarning(
510
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
511
    
512
    return (query, params)
513

    
514
def select(db, *args, **kw_args):
515
    '''For params, see mk_select() and run_query()'''
516
    recover = kw_args.pop('recover', None)
517
    cacheable = kw_args.pop('cacheable', True)
518
    
519
    query, params = mk_select(db, *args, **kw_args)
520
    return run_query(db, query, params, recover, cacheable)
521

    
522
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
523
    returning=None, embeddable=False, table_is_esc=False):
524
    '''
525
    @param returning str|None An inserted column (such as pkey) to return
526
    @param embeddable Whether the query should be embeddable as a nested SELECT.
527
        Warning: If you set this and cacheable=True when the query is run, the
528
        query will be fully cached, not just if it raises an exception.
529
    @param table_is_esc Whether the table name has already been escaped
530
    '''
531
    if select_query == None: select_query = 'DEFAULT VALUES'
532
    if cols == []: cols = None # no cols (all defaults) = unknown col names
533
    if not table_is_esc: check_name(table)
534
    
535
    # Build query
536
    query = 'INSERT INTO '+table
537
    if cols != None:
538
        map(check_name, cols)
539
        query += ' ('+', '.join(cols)+')'
540
    query += ' '+select_query
541
    
542
    if returning != None:
543
        check_name(returning)
544
        query += ' RETURNING '+returning
545
    
546
    if embeddable:
547
        # Create function
548
        function_name = '_'.join(map(clean_name, ['insert', table] + cols))
549
        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
550
        while True:
551
            try:
552
                function = function_name
553
                if not db.debug: function = 'pg_temp.'+function
554
                
555
                function_query = '''\
556
CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
557
    LANGUAGE sql
558
    AS $$'''+mogrify(db, query, params)+''';$$;
559
'''
560
                run_query(db, function_query, recover=True, cacheable=True)
561
                break # this version was successful
562
            except DuplicateFunctionException, e:
563
                function_name = next_version(function_name)
564
                # try again with next version of name
565
        
566
        # Return query that uses function
567
        return mk_select(db, function+'() AS f ('+returning+')', start=0,
568
            order_by=None, table_is_esc=True)# AS clause requires function alias
569
    
570
    return (query, params)
571

    
572
def insert_select(db, *args, **kw_args):
573
    '''For params, see mk_insert_select() and run_query_into()
574
    @param into_ref List with name of temp table to place RETURNING values in
575
    '''
576
    into_ref = kw_args.pop('into_ref', None)
577
    if into_ref != None: kw_args['embeddable'] = True
578
    recover = kw_args.pop('recover', None)
579
    cacheable = kw_args.pop('cacheable', True)
580
    
581
    query, params = mk_insert_select(db, *args, **kw_args)
582
    return run_query_into(db, query, params, into_ref, recover=recover,
583
        cacheable=cacheable)
584

    
585
default = object() # tells insert() to use the default value for a column
586

    
587
def insert(db, table, row, *args, **kw_args):
588
    '''For params, see insert_select()'''
589
    if lists.is_seq(row): cols = None
590
    else:
591
        cols = row.keys()
592
        row = row.values()
593
    row = list(row) # ensure that "!= []" works
594
    
595
    # Check for special values
596
    labels = []
597
    values = []
598
    for value in row:
599
        if value == default: labels.append('DEFAULT')
600
        else:
601
            labels.append('%s')
602
            values.append(value)
603
    
604
    # Build query
605
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
606
    else: query = None
607
    
608
    return insert_select(db, table, cols, query, values, *args, **kw_args)
609

    
610
def last_insert_id(db):
611
    module = util.root_module(db.db)
612
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
613
    elif module == 'MySQLdb': return db.insert_id()
614
    else: return None
615

    
616
def truncate(db, table, schema='public'):
617
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
618

    
619
##### Database structure queries
620

    
621
def pkey(db, table, recover=None, table_is_esc=False):
622
    '''Assumed to be first column in table'''
623
    return col_names(select(db, table, limit=0, order_by=None, recover=recover,
624
        table_is_esc=table_is_esc)).next()
625

    
626
def index_cols(db, table, index):
627
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
628
    automatically created. When you don't know whether something is a UNIQUE
629
    constraint or a UNIQUE index, use this function.'''
630
    check_name(table)
631
    check_name(index)
632
    module = util.root_module(db.db)
633
    if module == 'psycopg2':
634
        return list(values(run_query(db, '''\
635
SELECT attname
636
FROM
637
(
638
        SELECT attnum, attname
639
        FROM pg_index
640
        JOIN pg_class index ON index.oid = indexrelid
641
        JOIN pg_class table_ ON table_.oid = indrelid
642
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
643
        WHERE
644
            table_.relname = %(table)s
645
            AND index.relname = %(index)s
646
    UNION
647
        SELECT attnum, attname
648
        FROM
649
        (
650
            SELECT
651
                indrelid
652
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
653
                    AS indkey
654
            FROM pg_index
655
            JOIN pg_class index ON index.oid = indexrelid
656
            JOIN pg_class table_ ON table_.oid = indrelid
657
            WHERE
658
                table_.relname = %(table)s
659
                AND index.relname = %(index)s
660
        ) s
661
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
662
) s
663
ORDER BY attnum
664
''',
665
            {'table': table, 'index': index}, cacheable=True)))
666
    else: raise NotImplementedError("Can't list index columns for "+module+
667
        ' database')
668

    
669
def constraint_cols(db, table, constraint):
670
    check_name(table)
671
    check_name(constraint)
672
    module = util.root_module(db.db)
673
    if module == 'psycopg2':
674
        return list(values(run_query(db, '''\
675
SELECT attname
676
FROM pg_constraint
677
JOIN pg_class ON pg_class.oid = conrelid
678
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
679
WHERE
680
    relname = %(table)s
681
    AND conname = %(constraint)s
682
ORDER BY attnum
683
''',
684
            {'table': table, 'constraint': constraint})))
685
    else: raise NotImplementedError("Can't list constraint columns for "+module+
686
        ' database')
687

    
688
row_num_col = '_row_num'
689

    
690
def add_row_num(db, table):
691
    '''Adds a row number column to a table. Its name is in row_num_col. It will
692
    be the primary key.'''
693
    check_name(table)
694
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
695
        +' serial NOT NULL PRIMARY KEY')
696

    
697
def tables(db, schema='public', table_like='%'):
698
    module = util.root_module(db.db)
699
    params = {'schema': schema, 'table_like': table_like}
700
    if module == 'psycopg2':
701
        return values(run_query(db, '''\
702
SELECT tablename
703
FROM pg_tables
704
WHERE
705
    schemaname = %(schema)s
706
    AND tablename LIKE %(table_like)s
707
ORDER BY tablename
708
''',
709
            params, cacheable=True))
710
    elif module == 'MySQLdb':
711
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
712
            cacheable=True))
713
    else: raise NotImplementedError("Can't list tables for "+module+' database')
714

    
715
##### Database management
716

    
717
def empty_db(db, schema='public', **kw_args):
718
    '''For kw_args, see tables()'''
719
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
720

    
721
##### Heuristic queries
722

    
723
def put(db, table, row, pkey_=None, row_ct_ref=None):
724
    '''Recovers from errors.
725
    Only works under PostgreSQL (uses INSERT RETURNING).
726
    '''
727
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
728
    
729
    try:
730
        cur = insert(db, table, row, pkey_, recover=True)
731
        if row_ct_ref != None and cur.rowcount >= 0:
732
            row_ct_ref[0] += cur.rowcount
733
        return value(cur)
734
    except DuplicateKeyException, e:
735
        return value(select(db, table, [pkey_],
736
            util.dict_subset_right_join(row, e.cols), recover=True))
737

    
738
def get(db, table, row, pkey, row_ct_ref=None, create=False):
739
    '''Recovers from errors'''
740
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
741
    except StopIteration:
742
        if not create: raise
743
        return put(db, table, row, pkey, row_ct_ref) # insert new row
744

    
745
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
746
    row_ct_ref=None, table_is_esc=False):
747
    '''Recovers from errors.
748
    Only works under PostgreSQL (uses INSERT RETURNING).
749
    @param in_tables The main input table to select from, followed by a list of
750
        tables to join with it using the main input table's pkey
751
    @return (table, col) Where the pkeys (from INSERT RETURNING) are made
752
        available
753
    '''
754
    temp_suffix = clean_name(out_table)
755
        # suffix, not prefix, so main name won't be removed if name is truncated
756
    pkeys_ref = ['pkeys_'+temp_suffix]
757
    
758
    # Join together input tables
759
    in_tables = in_tables[:] # don't modify input!
760
    in_tables0 = in_tables.pop(0) # first table is separate
761
    in_pkey = pkey(db, in_tables0, recover=True, table_is_esc=table_is_esc)
762
    insert_joins = [in_tables0]+[(t, {in_pkey: join_using}) for t in in_tables]
763
    
764
    out_pkey = pkey(db, out_table, recover=True, table_is_esc=table_is_esc)
765
    pkeys_cols = [in_pkey, out_pkey]
766
    
767
    pkeys_table_exists_ref = [False]
768
    def run_query_into_pkeys(query, params):
769
        if pkeys_table_exists_ref[0]:
770
            insert_select(db, pkeys_ref[0], pkeys_cols, query, params)
771
        else:
772
            run_query_into(db, query, params, into_ref=pkeys_ref)
773
            pkeys_table_exists_ref[0] = True
774
    
775
    conds = {}
776
    distinct_on = []
777
    def mk_main_select(cols):
778
        return mk_select(db, insert_joins, cols, conds, distinct_on,
779
            order_by=None, limit=limit, start=start, table_is_esc=table_is_esc,
780
            default_table=in_tables0)
781
    
782
    # Do inserts and selects
783
    out_pkeys_ref = ['out_pkeys_'+temp_suffix]
784
    while True:
785
        try:
786
            cur = insert_select(db, out_table, mapping.keys(),
787
                *mk_main_select(mapping.values()), returning=out_pkey,
788
                into_ref=out_pkeys_ref, recover=True, table_is_esc=table_is_esc)
789
            if row_ct_ref != None and cur.rowcount >= 0:
790
                row_ct_ref[0] += cur.rowcount
791
                add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
792
            
793
            # Get input pkeys corresponding to rows in insert
794
            in_pkeys_ref = ['in_pkeys_'+temp_suffix]
795
            run_query_into(db, *mk_main_select([in_pkey]),
796
                into_ref=in_pkeys_ref)
797
            add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
798
            
799
            # Join together output and input pkeys
800
            run_query_into_pkeys(*mk_select(db, [in_pkeys_ref[0],
801
                (out_pkeys_ref[0], {row_num_col: join_using})], pkeys_cols,
802
                start=0))
803
            
804
            break # insert successful
805
        except DuplicateKeyException, e:
806
            join_cols = util.dict_subset_right_join(mapping, e.cols)
807
            select_joins = insert_joins + [(out_table, join_cols)]
808
            
809
            # Get pkeys of already existing rows
810
            run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
811
                order_by=None, start=0, table_is_esc=table_is_esc))
812
            
813
            # Save existing pkeys in temp table for joining on
814
            existing_pkeys_ref = ['existing_pkeys_'+temp_suffix]
815
            run_query_into(db, *mk_select(db, pkeys_ref[0], [in_pkey],
816
                order_by=None, start=0, table_is_esc=True),
817
                into_ref=existing_pkeys_ref)
818
                # need table_is_esc=True to make table name case-insensitive
819
            
820
            # rerun loop with additional constraints
821
            break # but until NullValueExceptions are handled, end loop here
822
        except NullValueException, e:
823
            out_col, = e.cols
824
            try: in_col = mapping[out_col]
825
            except KeyError: # no mapping for missing col, so every row invalid
826
                run_query_into_pkeys(*mk_select(db, insert_joins,
827
                    [in_pkey, sql_gen.NamedCode(out_pkey, None)],
828
                    order_by=None, start=0, table_is_esc=table_is_esc))
829
                break
830
            else:
831
                conds[in_col] = sql_gen.CompareCond(None, '!=')
832
                # rerun loop with additional constraints
833
    
834
    return (pkeys_ref[0], out_pkey)
835

    
836
##### Data cleanup
837

    
838
def cleanup_table(db, table, cols, table_is_esc=False):
839
    def esc_name_(name): return esc_name(db, name)
840
    
841
    if not table_is_esc: check_name(table)
842
    cols = map(esc_name_, cols)
843
    
844
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
845
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
846
            for col in cols))),
847
        dict(null0='', null1=r'\N'))
(22-22/34)