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 FunctionValueException(ExceptionWithName): pass
54

    
55
class DuplicateTableException(ExceptionWithName): pass
56

    
57
class DuplicateFunctionException(ExceptionWithName): pass
58

    
59
class EmptyRowException(DbException): pass
60

    
61
##### Warnings
62

    
63
class DbWarning(UserWarning): pass
64

    
65
##### Result retrieval
66

    
67
def col_names(cur): return (col[0] for col in cur.description)
68

    
69
def rows(cur): return iter(lambda: cur.fetchone(), None)
70

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

    
75
def next_row(cur): return rows(cur).next()
76

    
77
def row(cur):
78
    row_ = next_row(cur)
79
    consume_rows(cur)
80
    return row_
81

    
82
def next_value(cur): return next_row(cur)[0]
83

    
84
def value(cur): return row(cur)[0]
85

    
86
def values(cur): return iters.func_iter(lambda: next_value(cur))
87

    
88
def value_or_none(cur):
89
    try: return value(cur)
90
    except StopIteration: return None
91

    
92
##### Input validation
93

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

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

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

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

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

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

    
123
##### Database connections
124

    
125
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
126

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

    
132
DatabaseErrors_set = set([DbException])
133
DatabaseErrors = tuple(DatabaseErrors_set)
134

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

    
140
def db_config_str(db_config):
141
    return db_config['engine']+' database '+db_config['database']
142

    
143
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
144

    
145
log_debug_none = lambda msg: None
146

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

    
320
connect = DbConn
321

    
322
##### Querying
323

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

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

    
334
##### Recoverable querying
335

    
336
def with_savepoint(db, func): return db.with_savepoint(func)
337

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

    
368
##### Basic queries
369

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

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

    
399
order_by_pkey = object() # tells mk_select() to order by the pkey
400

    
401
join_using = object() # tells mk_select() to join the column with USING
402

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

    
405
distinct_on_all = object() # tells mk_select() to SELECT DISTINCT ON all columns
406

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

    
526
def select(db, *args, **kw_args):
527
    '''For params, see mk_select() and run_query()'''
528
    recover = kw_args.pop('recover', None)
529
    cacheable = kw_args.pop('cacheable', True)
530
    
531
    query, params = mk_select(db, *args, **kw_args)
532
    return run_query(db, query, params, recover, cacheable)
533

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

    
584
def insert_select(db, *args, **kw_args):
585
    '''For params, see mk_insert_select() and run_query_into()
586
    @param into_ref List with name of temp table to place RETURNING values in
587
    '''
588
    into_ref = kw_args.pop('into_ref', None)
589
    if into_ref != None: kw_args['embeddable'] = True
590
    recover = kw_args.pop('recover', None)
591
    cacheable = kw_args.pop('cacheable', True)
592
    
593
    query, params = mk_insert_select(db, *args, **kw_args)
594
    return run_query_into(db, query, params, into_ref, recover=recover,
595
        cacheable=cacheable)
596

    
597
default = object() # tells insert() to use the default value for a column
598

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

    
622
def last_insert_id(db):
623
    module = util.root_module(db.db)
624
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
625
    elif module == 'MySQLdb': return db.insert_id()
626
    else: return None
627

    
628
def truncate(db, table, schema='public'):
629
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
630

    
631
##### Database structure queries
632

    
633
def pkey(db, table, recover=None, table_is_esc=False):
634
    '''Assumed to be first column in table'''
635
    return col_names(select(db, table, limit=0, order_by=None, recover=recover,
636
        table_is_esc=table_is_esc)).next()
637

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

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

    
700
row_num_col = '_row_num'
701

    
702
def add_row_num(db, table):
703
    '''Adds a row number column to a table. Its name is in row_num_col. It will
704
    be the primary key.'''
705
    check_name(table)
706
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
707
        +' serial NOT NULL PRIMARY KEY')
708

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

    
727
##### Database management
728

    
729
def empty_db(db, schema='public', **kw_args):
730
    '''For kw_args, see tables()'''
731
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
732

    
733
##### Heuristic queries
734

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

    
750
def get(db, table, row, pkey, row_ct_ref=None, create=False):
751
    '''Recovers from errors'''
752
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
753
    except StopIteration:
754
        if not create: raise
755
        return put(db, table, row, pkey, row_ct_ref) # insert new row
756

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

    
851
##### Data cleanup
852

    
853
def cleanup_table(db, table, cols, table_is_esc=False):
854
    def esc_name_(name): return esc_name(db, name)
855
    
856
    if not table_is_esc: check_name(table)
857
    cols = map(esc_name_, cols)
858
    
859
    run_query(db, 'UPDATE '+table+' SET\n'+(',\n'.join(('\n'+col
860
        +' = nullif(nullif(trim(both from '+col+"), %(null0)s), %(null1)s)"
861
            for col in cols))),
862
        dict(null0='', null1=r'\N'))
(22-22/34)