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 ExceptionWithNameValue(DbException):
43
    def __init__(self, name, value, cause=None):
44
        DbException.__init__(self,
45
            'for name: '+str(name)+'; value: '+strings.ustr(value), cause)
46
        self.name = name
47
        self.value = value
48

    
49
class ExceptionWithColumns(DbException):
50
    def __init__(self, cols, cause=None):
51
        DbException.__init__(self, 'for columns: '+(', '.join(cols)), cause)
52
        self.cols = cols
53

    
54
class NameException(DbException): pass
55

    
56
class DuplicateKeyException(ExceptionWithColumns): pass
57

    
58
class NullValueException(ExceptionWithColumns): pass
59

    
60
class FunctionValueException(ExceptionWithNameValue): pass
61

    
62
class DuplicateTableException(ExceptionWithName): pass
63

    
64
class DuplicateFunctionException(ExceptionWithName): pass
65

    
66
class EmptyRowException(DbException): pass
67

    
68
##### Warnings
69

    
70
class DbWarning(UserWarning): pass
71

    
72
##### Result retrieval
73

    
74
def col_names(cur): return (col[0] for col in cur.description)
75

    
76
def rows(cur): return iter(lambda: cur.fetchone(), None)
77

    
78
def consume_rows(cur):
79
    '''Used to fetch all rows so result will be cached'''
80
    iters.consume_iter(rows(cur))
81

    
82
def next_row(cur): return rows(cur).next()
83

    
84
def row(cur):
85
    row_ = next_row(cur)
86
    consume_rows(cur)
87
    return row_
88

    
89
def next_value(cur): return next_row(cur)[0]
90

    
91
def value(cur): return row(cur)[0]
92

    
93
def values(cur): return iters.func_iter(lambda: next_value(cur))
94

    
95
def value_or_none(cur):
96
    try: return value(cur)
97
    except StopIteration: return None
98

    
99
##### Input validation
100

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

    
103
def check_name(name):
104
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
105
        +'" may contain only alphanumeric characters and _')
106

    
107
def esc_name_by_module(module, name, ignore_case=False):
108
    if module == 'psycopg2':
109
        if ignore_case:
110
            # Don't enclose in quotes because this disables case-insensitivity
111
            check_name(name)
112
            return name
113
        else: quote = '"'
114
    elif module == 'MySQLdb': quote = '`'
115
    else: raise NotImplementedError("Can't escape name for "+module+' database')
116
    return quote + name.replace(quote, '') + quote
117

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

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

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

    
130
##### Database connections
131

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

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

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

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

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

    
150
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
151

    
152
log_debug_none = lambda msg: None
153

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

    
327
connect = DbConn
328

    
329
##### Querying
330

    
331
def run_raw_query(db, *args, **kw_args):
332
    '''For params, see DbConn.run_query()'''
333
    return db.run_query(*args, **kw_args)
334

    
335
def mogrify(db, query, params):
336
    module = util.root_module(db.db)
337
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
338
    else: raise NotImplementedError("Can't mogrify query for "+module+
339
        ' database')
340

    
341
##### Recoverable querying
342

    
343
def with_savepoint(db, func): return db.with_savepoint(func)
344

    
345
def run_query(db, query, params=None, recover=None, cacheable=False):
346
    if recover == None: recover = False
347
    
348
    try:
349
        def run(): return run_raw_query(db, query, params, cacheable)
350
        if recover and not db.is_cached(query, params):
351
            return with_savepoint(db, run)
352
        else: return run() # don't need savepoint if cached
353
    except Exception, e:
354
        if not recover: raise # need savepoint to run index_cols()
355
        msg = exc.str_(e)
356
        
357
        match = re.search(r'duplicate key value violates unique constraint '
358
            r'"((_?[^\W_]+)_[^"]+)"', msg)
359
        if match:
360
            constraint, table = match.groups()
361
            try: cols = index_cols(db, table, constraint)
362
            except NotImplementedError: raise e
363
            else: raise DuplicateKeyException(cols, e)
364
        
365
        match = re.search(r'null value in column "(\w+)" violates not-null '
366
            r'constraint', msg)
367
        if match: raise NullValueException([match.group(1)], e)
368
        
369
        match = re.search(r'invalid input (?:syntax|value)\b.*: "(.+)"\n'
370
            r'(?:(?s).*)\bfunction "(\w+)".*\bat assignment', msg)
371
        if match:
372
            raise FunctionValueException(match.group(2),
373
                strings.to_unicode(match.group(1)), e)
374
        
375
        match = re.search(r'relation "(\w+)" already exists', msg)
376
        if match: raise DuplicateTableException(match.group(1), e)
377
        
378
        match = re.search(r'function "(\w+)" already exists', msg)
379
        if match: raise DuplicateFunctionException(match.group(1), e)
380
        
381
        raise # no specific exception raised
382

    
383
##### Basic queries
384

    
385
def next_version(name):
386
    '''Prepends the version # so it won't be removed if the name is truncated'''
387
    version = 1 # first existing name was version 0
388
    match = re.match(r'^v(\d+)_(.*)$', name)
389
    if match:
390
        version = int(match.group(1))+1
391
        name = match.group(2)
392
    return 'v'+str(version)+'_'+name
393

    
394
def run_query_into(db, query, params, into_ref=None, *args, **kw_args):
395
    '''Outputs a query to a temp table.
396
    For params, see run_query().
397
    '''
398
    if into_ref == None: return run_query(db, query, params, *args, **kw_args)
399
    else: # place rows in temp table
400
        check_name(into_ref[0])
401
        kw_args['recover'] = True
402
        while True:
403
            try:
404
                create_query = 'CREATE'
405
                if not db.debug: create_query += ' TEMP'
406
                create_query += ' TABLE '+into_ref[0]+' AS '+query
407
                
408
                return run_query(db, create_query, params, *args, **kw_args)
409
                    # CREATE TABLE AS sets rowcount to # rows in query
410
            except DuplicateTableException, e:
411
                into_ref[0] = next_version(into_ref[0])
412
                # try again with next version of name
413

    
414
order_by_pkey = object() # tells mk_select() to order by the pkey
415

    
416
join_using = object() # tells mk_select() to join the column with USING
417

    
418
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
419

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

    
533
def select(db, *args, **kw_args):
534
    '''For params, see mk_select() and run_query()'''
535
    recover = kw_args.pop('recover', None)
536
    cacheable = kw_args.pop('cacheable', True)
537
    
538
    query, params = mk_select(db, *args, **kw_args)
539
    return run_query(db, query, params, recover, cacheable)
540

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

    
591
def insert_select(db, *args, **kw_args):
592
    '''For params, see mk_insert_select() and run_query_into()
593
    @param into_ref List with name of temp table to place RETURNING values in
594
    '''
595
    into_ref = kw_args.pop('into_ref', None)
596
    if into_ref != None: kw_args['embeddable'] = True
597
    recover = kw_args.pop('recover', None)
598
    cacheable = kw_args.pop('cacheable', True)
599
    
600
    query, params = mk_insert_select(db, *args, **kw_args)
601
    return run_query_into(db, query, params, into_ref, recover=recover,
602
        cacheable=cacheable)
603

    
604
default = object() # tells insert() to use the default value for a column
605

    
606
def insert(db, table, row, *args, **kw_args):
607
    '''For params, see insert_select()'''
608
    if lists.is_seq(row): cols = None
609
    else:
610
        cols = row.keys()
611
        row = row.values()
612
    row = list(row) # ensure that "!= []" works
613
    
614
    # Check for special values
615
    labels = []
616
    values = []
617
    for value in row:
618
        if value is default: labels.append('DEFAULT')
619
        else:
620
            labels.append('%s')
621
            values.append(value)
622
    
623
    # Build query
624
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
625
    else: query = None
626
    
627
    return insert_select(db, table, cols, query, values, *args, **kw_args)
628

    
629
def last_insert_id(db):
630
    module = util.root_module(db.db)
631
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
632
    elif module == 'MySQLdb': return db.insert_id()
633
    else: return None
634

    
635
def truncate(db, table, schema='public'):
636
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
637

    
638
##### Database structure queries
639

    
640
def pkey(db, table, recover=None, table_is_esc=False):
641
    '''Assumed to be first column in table'''
642
    return col_names(select(db, table, limit=0, order_by=None, recover=recover,
643
        table_is_esc=table_is_esc)).next()
644

    
645
def index_cols(db, table, index):
646
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
647
    automatically created. When you don't know whether something is a UNIQUE
648
    constraint or a UNIQUE index, use this function.'''
649
    check_name(table)
650
    check_name(index)
651
    module = util.root_module(db.db)
652
    if module == 'psycopg2':
653
        return list(values(run_query(db, '''\
654
SELECT attname
655
FROM
656
(
657
        SELECT attnum, attname
658
        FROM pg_index
659
        JOIN pg_class index ON index.oid = indexrelid
660
        JOIN pg_class table_ ON table_.oid = indrelid
661
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
662
        WHERE
663
            table_.relname = %(table)s
664
            AND index.relname = %(index)s
665
    UNION
666
        SELECT attnum, attname
667
        FROM
668
        (
669
            SELECT
670
                indrelid
671
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
672
                    AS indkey
673
            FROM pg_index
674
            JOIN pg_class index ON index.oid = indexrelid
675
            JOIN pg_class table_ ON table_.oid = indrelid
676
            WHERE
677
                table_.relname = %(table)s
678
                AND index.relname = %(index)s
679
        ) s
680
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
681
) s
682
ORDER BY attnum
683
''',
684
            {'table': table, 'index': index}, cacheable=True)))
685
    else: raise NotImplementedError("Can't list index columns for "+module+
686
        ' database')
687

    
688
def constraint_cols(db, table, constraint):
689
    check_name(table)
690
    check_name(constraint)
691
    module = util.root_module(db.db)
692
    if module == 'psycopg2':
693
        return list(values(run_query(db, '''\
694
SELECT attname
695
FROM pg_constraint
696
JOIN pg_class ON pg_class.oid = conrelid
697
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
698
WHERE
699
    relname = %(table)s
700
    AND conname = %(constraint)s
701
ORDER BY attnum
702
''',
703
            {'table': table, 'constraint': constraint})))
704
    else: raise NotImplementedError("Can't list constraint columns for "+module+
705
        ' database')
706

    
707
row_num_col = '_row_num'
708

    
709
def add_row_num(db, table):
710
    '''Adds a row number column to a table. Its name is in row_num_col. It will
711
    be the primary key.'''
712
    check_name(table)
713
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
714
        +' serial NOT NULL PRIMARY KEY')
715

    
716
def tables(db, schema='public', table_like='%'):
717
    module = util.root_module(db.db)
718
    params = {'schema': schema, 'table_like': table_like}
719
    if module == 'psycopg2':
720
        return values(run_query(db, '''\
721
SELECT tablename
722
FROM pg_tables
723
WHERE
724
    schemaname = %(schema)s
725
    AND tablename LIKE %(table_like)s
726
ORDER BY tablename
727
''',
728
            params, cacheable=True))
729
    elif module == 'MySQLdb':
730
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
731
            cacheable=True))
732
    else: raise NotImplementedError("Can't list tables for "+module+' database')
733

    
734
##### Database management
735

    
736
def empty_db(db, schema='public', **kw_args):
737
    '''For kw_args, see tables()'''
738
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
739

    
740
##### Heuristic queries
741

    
742
def put(db, table, row, pkey_=None, row_ct_ref=None):
743
    '''Recovers from errors.
744
    Only works under PostgreSQL (uses INSERT RETURNING).
745
    '''
746
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
747
    
748
    try:
749
        cur = insert(db, table, row, pkey_, recover=True)
750
        if row_ct_ref != None and cur.rowcount >= 0:
751
            row_ct_ref[0] += cur.rowcount
752
        return value(cur)
753
    except DuplicateKeyException, e:
754
        return value(select(db, table, [pkey_],
755
            util.dict_subset_right_join(row, e.cols), recover=True))
756

    
757
def get(db, table, row, pkey, row_ct_ref=None, create=False):
758
    '''Recovers from errors'''
759
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
760
    except StopIteration:
761
        if not create: raise
762
        return put(db, table, row, pkey, row_ct_ref) # insert new row
763

    
764
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
765
    row_ct_ref=None, table_is_esc=False):
766
    '''Recovers from errors.
767
    Only works under PostgreSQL (uses INSERT RETURNING).
768
    @param in_tables The main input table to select from, followed by a list of
769
        tables to join with it using the main input table's pkey
770
    @return (table, col) Where the pkeys (from INSERT RETURNING) are made
771
        available
772
    '''
773
    temp_prefix = clean_name(out_table)
774
    pkeys_ref = [temp_prefix+'_pkeys']
775
    
776
    # Join together input tables
777
    in_tables = in_tables[:] # don't modify input!
778
    in_tables0 = in_tables.pop(0) # first table is separate
779
    in_pkey = pkey(db, in_tables0, recover=True, table_is_esc=table_is_esc)
780
    in_pkey_col = sql_gen.col2sql_gen(in_pkey, in_tables0, table_is_esc)
781
    insert_joins = [in_tables0]+[(t, {in_pkey: join_using}) for t in in_tables]
782
    
783
    out_pkey = pkey(db, out_table, recover=True, table_is_esc=table_is_esc)
784
    out_pkey_col = sql_gen.col2sql_gen(out_pkey, out_table, table_is_esc)
785
    
786
    pkeys = [in_pkey, out_pkey]
787
    pkeys_cols = [in_pkey_col, out_pkey_col]
788
    
789
    pkeys_table_exists_ref = [False]
790
    def run_query_into_pkeys(query, params):
791
        if pkeys_table_exists_ref[0]:
792
            insert_select(db, pkeys_ref[0], pkeys, query, params)
793
        else:
794
            run_query_into(db, query, params, into_ref=pkeys_ref)
795
            pkeys_table_exists_ref[0] = True
796
    
797
    conds = {}
798
    distinct_on = []
799
    def mk_main_select(cols):
800
        return mk_select(db, insert_joins, cols, conds, distinct_on,
801
            order_by=None, limit=limit, start=start, table_is_esc=table_is_esc,
802
            default_table=in_tables0)
803
    
804
    def log_ignore(in_col, value):
805
        db.log_debug('Ignoring rows with '+in_col+' = '+value)
806
    
807
    # Do inserts and selects
808
    join_cols = {}
809
    filter_join_added = False
810
    out_pkeys_ref = [temp_prefix+'_out_pkeys']
811
    in_pkeys_ref = [temp_prefix+'_in_pkeys']
812
    while True:
813
        try:
814
            if join_cols != {}:
815
                select_joins = insert_joins + [(out_table, join_cols)]
816
                db.log_debug('Getting pkeys of already existing rows')
817
                run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
818
                    order_by=None, start=0, table_is_esc=table_is_esc))
819
                
820
                # Prepare to insert new rows
821
                input_join_cols = join_cols.values()
822
                distinct_on = filter(util.is_str, input_join_cols)
823
                if not filter_join_added:
824
                    insert_joins.append(sql_gen.Join(pkeys_ref[0],
825
                        {in_pkey: in_pkey}, sql_gen.filter_out))
826
                    filter_join_added = True
827
            
828
            db.log_debug('Inserting new rows')
829
            cur = insert_select(db, out_table, mapping.keys(),
830
                *mk_main_select(mapping.values()), returning=out_pkey,
831
                into_ref=out_pkeys_ref, recover=True, table_is_esc=table_is_esc)
832
            if row_ct_ref != None and cur.rowcount >= 0:
833
                row_ct_ref[0] += cur.rowcount
834
                add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
835
            
836
            db.log_debug('Getting input pkeys corresponding to rows in insert')
837
            run_query_into(db, *mk_main_select([in_pkey]),
838
                into_ref=in_pkeys_ref)
839
            add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
840
            
841
            db.log_debug('Joining together output and input pkeys')
842
            run_query_into_pkeys(*mk_select(db, [in_pkeys_ref[0],
843
                (out_pkeys_ref[0], {row_num_col: join_using})], pkeys, start=0))
844
            
845
            break # insert successful
846
        except DuplicateKeyException, e:
847
            old_join_cols = join_cols.copy()
848
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
849
            db.log_debug('Ignoring existing rows, comparing on '+str(join_cols))
850
            assert join_cols != old_join_cols # avoid infinite loops
851
            # rerun loop with additional constraints
852
        except NullValueException, e:
853
            out_col, = e.cols
854
            try: in_col = mapping[out_col]
855
            except KeyError: # no mapping for missing col, so every row invalid
856
                db.log_debug('Missing mapping for '+out_col)
857
                run_query_into_pkeys(*mk_select(db, insert_joins,
858
                    [in_pkey_col, sql_gen.NamedCode(out_pkey, None)],
859
                    order_by=None, start=0, table_is_esc=table_is_esc))
860
                break
861
            else:
862
                log_ignore(in_col, 'NULL')
863
                conds[in_col] = sql_gen.CompareCond(None, '!=')
864
                # rerun loop with additional constraints
865
        except FunctionValueException, e:
866
            assert e.name == sql_gen.unescape_table(out_table)
867
            out_col = 'value' # assume function param was named "value"
868
            in_col = mapping[out_col]
869
            value = e.value
870
            log_ignore(in_col, value)
871
            conds[in_col] = sql_gen.CompareCond(value, '!=')
872
            # rerun loop with additional constraints
873
    
874
    return sql_gen.Col(out_pkey, pkeys_ref[0])
875

    
876
##### Data cleanup
877

    
878
def cleanup_table(db, table, cols, table_is_esc=False):
879
    def esc_name_(name): return esc_name(db, name)
880
    
881
    if not table_is_esc: check_name(table)
882
    cols = map(esc_name_, cols)
883
    
884
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
885
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
886
            for col in cols))),
887
        dict(null0='', null1=r'\N'))
(22-22/34)