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 ConstraintException(DbException):
50
    def __init__(self, name, cols, cause=None):
51
        DbException.__init__(self, 'Violated constraint '+name+ ' on columns: '
52
            +(', '.join(cols)), cause)
53
        self.name = name
54
        self.cols = cols
55

    
56
class NameException(DbException): pass
57

    
58
class DuplicateKeyException(ConstraintException): pass
59

    
60
class NullValueException(ConstraintException): pass
61

    
62
class FunctionValueException(ExceptionWithNameValue): pass
63

    
64
class DuplicateTableException(ExceptionWithName): pass
65

    
66
class DuplicateFunctionException(ExceptionWithName): pass
67

    
68
class EmptyRowException(DbException): pass
69

    
70
##### Warnings
71

    
72
class DbWarning(UserWarning): pass
73

    
74
##### Result retrieval
75

    
76
def col_names(cur): return (col[0] for col in cur.description)
77

    
78
def rows(cur): return iter(lambda: cur.fetchone(), None)
79

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

    
84
def next_row(cur): return rows(cur).next()
85

    
86
def row(cur):
87
    row_ = next_row(cur)
88
    consume_rows(cur)
89
    return row_
90

    
91
def next_value(cur): return next_row(cur)[0]
92

    
93
def value(cur): return row(cur)[0]
94

    
95
def values(cur): return iters.func_iter(lambda: next_value(cur))
96

    
97
def value_or_none(cur):
98
    try: return value(cur)
99
    except StopIteration: return None
100

    
101
##### Input validation
102

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

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

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

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

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

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

    
132
##### Database connections
133

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

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

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

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

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

    
152
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
153

    
154
log_debug_none = lambda msg: None
155

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

    
329
connect = DbConn
330

    
331
##### Querying
332

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

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

    
343
##### Recoverable querying
344

    
345
def with_savepoint(db, func): return db.with_savepoint(func)
346

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

    
385
##### Basic queries
386

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

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

    
416
order_by_pkey = object() # tells mk_select() to order by the pkey
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, default_table=None):
422
    '''
423
    @param tables The single table to select from, or a list of tables to join
424
        together, with tables after the first being sql_gen.Join objects
425
    @param fields Use None to select all fields in the table
426
    @param conds WHERE conditions: dict(compare_left_side=compare_right_side):
427
        * compare_left_side: Code|str (for col name)
428
        * compare_right_side: ValueCond|literal value
429
    @param distinct_on The columns to SELECT DISTINCT ON, or distinct_on_all to
430
        use all columns
431
    @return tuple(query, params)
432
    '''
433
    if not lists.is_seq(tables): tables = [tables]
434
    tables = list(tables) # don't modify input! (list() copies input)
435
    table0 = tables.pop(0) # first table is separate
436
    
437
    if conds == None: conds = {}
438
    assert limit == None or type(limit) == int
439
    assert start == None or type(start) == int
440
    table0 = sql_gen.as_Table(table0)
441
    if order_by is order_by_pkey: order_by = pkey(db, table0, recover=True)
442
    
443
    params = []
444
    
445
    def parse_col(field):
446
        '''Parses fields'''
447
        return sql_gen.as_Col(field, default_table).to_str(db)
448
    
449
    query = 'SELECT'
450
    
451
    # DISTINCT ON columns
452
    if distinct_on != []:
453
        query += ' DISTINCT'
454
        if distinct_on is not distinct_on_all:
455
            query += ' ON ('+(', '.join(map(parse_col, distinct_on)))+')'
456
    
457
    # Columns
458
    query += ' '
459
    if fields == None: query += '*'
460
    else: query += ', '.join(map(parse_col, fields))
461
    
462
    # Main table
463
    query += ' FROM '+table0.to_str(db)
464
    
465
    # Add joins
466
    left_table = table0
467
    for join_ in tables:
468
        # Parse special values
469
        table = join_.table
470
        if join_.type_ is sql_gen.filter_out:
471
            conds[sql_gen.Col(pkey(db, table), table)] = None # filter no match
472
        
473
        query += ' '+join_.to_str(db, left_table)
474
        
475
        left_table = table
476
    
477
    missing = True
478
    if conds != {}:
479
        query += ' WHERE '+(' AND '.join((sql_gen.as_ValueCond(r).to_str(db, l)
480
            for l, r in conds.iteritems())))
481
        params += conds.values()
482
        missing = False
483
    if order_by != None:
484
        query += ' ORDER BY '+sql_gen.as_Col(order_by, table0).to_str(db)
485
    if limit != None: query += ' LIMIT '+str(limit); missing = False
486
    if start != None:
487
        if start != 0: query += ' OFFSET '+str(start)
488
        missing = False
489
    if missing: warnings.warn(DbWarning(
490
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
491
    
492
    return (query, params)
493

    
494
def select(db, *args, **kw_args):
495
    '''For params, see mk_select() and run_query()'''
496
    recover = kw_args.pop('recover', None)
497
    cacheable = kw_args.pop('cacheable', True)
498
    
499
    query, params = mk_select(db, *args, **kw_args)
500
    return run_query(db, query, params, recover, cacheable)
501

    
502
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
503
    returning=None, embeddable=False):
504
    '''
505
    @param returning str|None An inserted column (such as pkey) to return
506
    @param embeddable Whether the query should be embeddable as a nested SELECT.
507
        Warning: If you set this and cacheable=True when the query is run, the
508
        query will be fully cached, not just if it raises an exception.
509
    '''
510
    if select_query == None: select_query = 'DEFAULT VALUES'
511
    if cols == []: cols = None # no cols (all defaults) = unknown col names
512
    table = sql_gen.as_Table(table).to_str(db)
513
    
514
    # Build query
515
    query = 'INSERT INTO '+table
516
    if cols != None:
517
        map(check_name, cols)
518
        query += ' ('+', '.join(cols)+')'
519
    query += ' '+select_query
520
    
521
    if returning != None:
522
        check_name(returning)
523
        query += ' RETURNING '+returning
524
    
525
    if embeddable:
526
        # Create function
527
        function_name = '_'.join(map(clean_name, ['insert', table] + cols))
528
        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
529
        while True:
530
            try:
531
                function = function_name
532
                if not db.debug: function = 'pg_temp.'+function
533
                
534
                function_query = '''\
535
CREATE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
536
    LANGUAGE sql
537
    AS $$'''+mogrify(db, query, params)+''';$$;
538
'''
539
                run_query(db, function_query, recover=True, cacheable=True)
540
                break # this version was successful
541
            except DuplicateFunctionException, e:
542
                function_name = next_version(function_name)
543
                # try again with next version of name
544
        
545
        # Return query that uses function
546
        return mk_select(db, sql_gen.CustomCode(function+'() AS f ('+returning
547
            +')'), start=0, order_by=None) # AS clause requires function alias
548
    
549
    return (query, params)
550

    
551
def insert_select(db, *args, **kw_args):
552
    '''For params, see mk_insert_select() and run_query_into()
553
    @param into_ref List with name of temp table to place RETURNING values in
554
    '''
555
    into_ref = kw_args.pop('into_ref', None)
556
    if into_ref != None: kw_args['embeddable'] = True
557
    recover = kw_args.pop('recover', None)
558
    cacheable = kw_args.pop('cacheable', True)
559
    
560
    query, params = mk_insert_select(db, *args, **kw_args)
561
    return run_query_into(db, query, params, into_ref, recover=recover,
562
        cacheable=cacheable)
563

    
564
default = object() # tells insert() to use the default value for a column
565

    
566
def insert(db, table, row, *args, **kw_args):
567
    '''For params, see insert_select()'''
568
    if lists.is_seq(row): cols = None
569
    else:
570
        cols = row.keys()
571
        row = row.values()
572
    row = list(row) # ensure that "!= []" works
573
    
574
    # Check for special values
575
    labels = []
576
    values = []
577
    for value in row:
578
        if value is default: labels.append('DEFAULT')
579
        else:
580
            labels.append('%s')
581
            values.append(value)
582
    
583
    # Build query
584
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
585
    else: query = None
586
    
587
    return insert_select(db, table, cols, query, values, *args, **kw_args)
588

    
589
def last_insert_id(db):
590
    module = util.root_module(db.db)
591
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
592
    elif module == 'MySQLdb': return db.insert_id()
593
    else: return None
594

    
595
def truncate(db, table, schema='public'):
596
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
597

    
598
##### Database structure queries
599

    
600
def pkey(db, table, recover=None):
601
    '''Assumed to be first column in table'''
602
    return col_names(select(db, table, limit=0, order_by=None,
603
        recover=recover)).next()
604

    
605
def index_cols(db, table, index):
606
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
607
    automatically created. When you don't know whether something is a UNIQUE
608
    constraint or a UNIQUE index, use this function.'''
609
    check_name(table)
610
    check_name(index)
611
    module = util.root_module(db.db)
612
    if module == 'psycopg2':
613
        return list(values(run_query(db, '''\
614
SELECT attname
615
FROM
616
(
617
        SELECT attnum, attname
618
        FROM pg_index
619
        JOIN pg_class index ON index.oid = indexrelid
620
        JOIN pg_class table_ ON table_.oid = indrelid
621
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
622
        WHERE
623
            table_.relname = %(table)s
624
            AND index.relname = %(index)s
625
    UNION
626
        SELECT attnum, attname
627
        FROM
628
        (
629
            SELECT
630
                indrelid
631
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
632
                    AS indkey
633
            FROM pg_index
634
            JOIN pg_class index ON index.oid = indexrelid
635
            JOIN pg_class table_ ON table_.oid = indrelid
636
            WHERE
637
                table_.relname = %(table)s
638
                AND index.relname = %(index)s
639
        ) s
640
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
641
) s
642
ORDER BY attnum
643
''',
644
            {'table': table, 'index': index}, cacheable=True)))
645
    else: raise NotImplementedError("Can't list index columns for "+module+
646
        ' database')
647

    
648
def constraint_cols(db, table, constraint):
649
    check_name(table)
650
    check_name(constraint)
651
    module = util.root_module(db.db)
652
    if module == 'psycopg2':
653
        return list(values(run_query(db, '''\
654
SELECT attname
655
FROM pg_constraint
656
JOIN pg_class ON pg_class.oid = conrelid
657
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
658
WHERE
659
    relname = %(table)s
660
    AND conname = %(constraint)s
661
ORDER BY attnum
662
''',
663
            {'table': table, 'constraint': constraint})))
664
    else: raise NotImplementedError("Can't list constraint columns for "+module+
665
        ' database')
666

    
667
row_num_col = '_row_num'
668

    
669
def add_row_num(db, table):
670
    '''Adds a row number column to a table. Its name is in row_num_col. It will
671
    be the primary key.'''
672
    check_name(table)
673
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
674
        +' serial NOT NULL PRIMARY KEY')
675

    
676
def tables(db, schema='public', table_like='%'):
677
    module = util.root_module(db.db)
678
    params = {'schema': schema, 'table_like': table_like}
679
    if module == 'psycopg2':
680
        return values(run_query(db, '''\
681
SELECT tablename
682
FROM pg_tables
683
WHERE
684
    schemaname = %(schema)s
685
    AND tablename LIKE %(table_like)s
686
ORDER BY tablename
687
''',
688
            params, cacheable=True))
689
    elif module == 'MySQLdb':
690
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
691
            cacheable=True))
692
    else: raise NotImplementedError("Can't list tables for "+module+' database')
693

    
694
##### Database management
695

    
696
def empty_db(db, schema='public', **kw_args):
697
    '''For kw_args, see tables()'''
698
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
699

    
700
##### Heuristic queries
701

    
702
def put(db, table, row, pkey_=None, row_ct_ref=None):
703
    '''Recovers from errors.
704
    Only works under PostgreSQL (uses INSERT RETURNING).
705
    '''
706
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
707
    
708
    try:
709
        cur = insert(db, table, row, pkey_, recover=True)
710
        if row_ct_ref != None and cur.rowcount >= 0:
711
            row_ct_ref[0] += cur.rowcount
712
        return value(cur)
713
    except DuplicateKeyException, e:
714
        return value(select(db, table, [pkey_],
715
            util.dict_subset_right_join(row, e.cols), recover=True))
716

    
717
def get(db, table, row, pkey, row_ct_ref=None, create=False):
718
    '''Recovers from errors'''
719
    try: return value(select(db, table, [pkey], row, limit=1, recover=True))
720
    except StopIteration:
721
        if not create: raise
722
        return put(db, table, row, pkey, row_ct_ref) # insert new row
723

    
724
def put_table(db, out_table, in_tables, mapping, limit=None, start=0,
725
    row_ct_ref=None):
726
    '''Recovers from errors.
727
    Only works under PostgreSQL (uses INSERT RETURNING).
728
    @param in_tables The main input table to select from, followed by a list of
729
        tables to join with it using the main input table's pkey
730
    @return sql_gen.Col Where the pkeys (from INSERT RETURNING) are made
731
        available
732
    '''
733
    temp_prefix = clean_name(out_table)
734
    pkeys_ref = [temp_prefix+'_pkeys']
735
    
736
    # Join together input tables
737
    in_tables = in_tables[:] # don't modify input!
738
    in_tables0 = in_tables.pop(0) # first table is separate
739
    in_pkey = pkey(db, in_tables0, recover=True)
740
    in_pkey_col = sql_gen.as_Col(in_pkey, in_tables0)
741
    insert_joins = [in_tables0]+[sql_gen.Join(v, {in_pkey: sql_gen.join_using})
742
        for v in in_tables]
743
    
744
    out_pkey = pkey(db, out_table, recover=True)
745
    out_pkey_col = sql_gen.as_Col(out_pkey, out_table)
746
    
747
    pkeys = [in_pkey, out_pkey]
748
    pkeys_cols = [in_pkey_col, out_pkey_col]
749
    
750
    pkeys_table_exists_ref = [False]
751
    def run_query_into_pkeys(query, params):
752
        if pkeys_table_exists_ref[0]:
753
            insert_select(db, pkeys_ref[0], pkeys, query, params)
754
        else:
755
            run_query_into(db, query, params, into_ref=pkeys_ref)
756
            pkeys_table_exists_ref[0] = True
757
    
758
    conds = {}
759
    distinct_on = []
760
    def mk_main_select(cols):
761
        return mk_select(db, insert_joins, cols, conds, distinct_on,
762
            limit=limit, start=start, default_table=in_tables0)
763
    
764
    def log_exc(e):
765
        db.log_debug('Caught exception: '+exc.str_(e, first_line_only=True))
766
    def log_ignore(in_col, value):
767
        db.log_debug('Ignoring rows with '+in_col+' = '+value)
768
    
769
    # Do inserts and selects
770
    join_cols = {}
771
    filter_join_added = False
772
    out_pkeys_ref = [temp_prefix+'_out_pkeys']
773
    in_pkeys_ref = [temp_prefix+'_in_pkeys']
774
    while True:
775
        has_joins = join_cols != {}
776
        
777
        # Prepare to insert new rows
778
        if has_joins:
779
            distinct_on = filter(util.is_str, join_cols.values())
780
            if not filter_join_added:
781
                insert_joins.append(sql_gen.Join(out_table, join_cols,
782
                    sql_gen.filter_out))
783
                filter_join_added = True
784
            returning = None
785
            into_ref = None
786
        else:
787
            returning = out_pkey
788
            into_ref = out_pkeys_ref
789
        
790
        db.log_debug('Inserting new rows')
791
        try:
792
            cur = insert_select(db, out_table, mapping.keys(),
793
                *mk_main_select(mapping.values()), returning=returning,
794
                into_ref=into_ref, recover=True)
795
        except DuplicateKeyException, e:
796
            log_exc(e)
797
            
798
            old_join_cols = join_cols.copy()
799
            join_cols.update(util.dict_subset_right_join(mapping, e.cols))
800
            db.log_debug('Ignoring existing rows, comparing on '+str(join_cols))
801
            assert join_cols != old_join_cols # avoid infinite loops
802
            
803
            # rerun loop with additional constraints
804
        except NullValueException, e:
805
            log_exc(e)
806
            
807
            out_col, = e.cols
808
            try: in_col = mapping[out_col]
809
            except KeyError: # no mapping for missing col, so every row invalid
810
                db.log_debug('Missing mapping for '+out_col)
811
                run_query_into_pkeys(*mk_select(db, insert_joins,
812
                    [in_pkey_col, sql_gen.NamedCode(out_pkey, None)], start=0))
813
                break
814
            else:
815
                log_ignore(in_col, 'NULL')
816
                conds[in_col] = sql_gen.CompareCond(None, '!=')
817
                # rerun loop with additional constraints
818
        except FunctionValueException, e:
819
            log_exc(e)
820
            
821
            assert e.name == sql_gen.unescape_table(out_table)
822
            out_col = 'value' # assume function param was named "value"
823
            in_col = mapping[out_col]
824
            value = e.value
825
            log_ignore(in_col, value)
826
            conds[in_col] = sql_gen.CompareCond(value, '!=')
827
            # rerun loop with additional constraints
828
        else:
829
            if row_ct_ref != None and cur.rowcount >= 0:
830
                row_ct_ref[0] += cur.rowcount
831
                add_row_num(db, out_pkeys_ref[0]) # for joining with input pkeys
832
            
833
            if has_joins:
834
                select_joins = insert_joins+[sql_gen.Join(out_table, join_cols)]
835
                db.log_debug('Getting output pkeys of existing/inserted rows')
836
                run_query_into_pkeys(*mk_select(db, select_joins, pkeys_cols,
837
                    start=0))
838
            else:
839
                db.log_debug('Getting input pkeys for rows in insert')
840
                run_query_into(db, *mk_main_select([in_pkey]),
841
                    into_ref=in_pkeys_ref)
842
                add_row_num(db, in_pkeys_ref[0]) # for joining with output pkeys
843
                
844
                db.log_debug('Joining together output and input pkeys')
845
                run_query_into_pkeys(*mk_select(db, [in_pkeys_ref[0],
846
                    sql_gen.Join(out_pkeys_ref[0],
847
                    {row_num_col: sql_gen.join_using})], pkeys, start=0))
848
            
849
            break # insert successful
850
    
851
    return sql_gen.Col(out_pkey, pkeys_ref[0])
852

    
853
##### Data cleanup
854

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