Project

General

Profile

1
# Database access
2

    
3
import copy
4
import re
5
import warnings
6

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

    
16
##### Exceptions
17

    
18
def get_cur_query(cur):
19
    if hasattr(cur, 'query'): return cur.query
20
    elif hasattr(cur, '_last_executed'): return cur._last_executed
21
    else: return None
22

    
23
def _add_cursor_info(e, cur): exc.add_msg(e, 'query: '+get_cur_query(cur))
24

    
25
class DbException(exc.ExceptionWithCause):
26
    def __init__(self, msg, cause=None, cur=None):
27
        exc.ExceptionWithCause.__init__(self, msg, cause)
28
        if cur != None: _add_cursor_info(self, cur)
29

    
30
class NameException(DbException): pass
31

    
32
class ExceptionWithColumns(DbException):
33
    def __init__(self, cols, cause=None):
34
        DbException.__init__(self, 'columns: ' + ', '.join(cols), cause)
35
        self.cols = cols
36

    
37
class DuplicateKeyException(ExceptionWithColumns): pass
38

    
39
class NullValueException(ExceptionWithColumns): pass
40

    
41
class EmptyRowException(DbException): pass
42

    
43
##### Warnings
44

    
45
class DbWarning(UserWarning): pass
46

    
47
##### Result retrieval
48

    
49
def col_names(cur): return (col[0] for col in cur.description)
50

    
51
def rows(cur): return iter(lambda: cur.fetchone(), None)
52

    
53
def consume_rows(cur):
54
    '''Used to fetch all rows so result will be cached'''
55
    iters.consume_iter(rows(cur))
56

    
57
def next_row(cur): return rows(cur).next()
58

    
59
def row(cur):
60
    row_ = next_row(cur)
61
    consume_rows(cur)
62
    return row_
63

    
64
def next_value(cur): return next_row(cur)[0]
65

    
66
def value(cur): return row(cur)[0]
67

    
68
def values(cur): return iters.func_iter(lambda: next_value(cur))
69

    
70
def value_or_none(cur):
71
    try: return value(cur)
72
    except StopIteration: return None
73

    
74
##### Input validation
75

    
76
def clean_name(name): return re.sub(r'\W', r'', name)
77

    
78
def check_name(name):
79
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
80
        +'" may contain only alphanumeric characters and _')
81

    
82
def esc_name_by_module(module, name, ignore_case=False):
83
    if module == 'psycopg2':
84
        if ignore_case:
85
            # Don't enclose in quotes because this disables case-insensitivity
86
            check_name(name)
87
            return name
88
        else: quote = '"'
89
    elif module == 'MySQLdb': quote = '`'
90
    else: raise NotImplementedError("Can't escape name for "+module+' database')
91
    return quote + name.replace(quote, '') + quote
92

    
93
def esc_name_by_engine(engine, name, **kw_args):
94
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
95

    
96
def esc_name(db, name, **kw_args):
97
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
98

    
99
def qual_name(db, schema, table):
100
    def esc_name_(name): return esc_name(db, name)
101
    table = esc_name_(table)
102
    if schema != None: return esc_name_(schema)+'.'+table
103
    else: return table
104

    
105
##### Database connections
106

    
107
db_config_names = ['engine', 'host', 'user', 'password', 'database', 'schemas']
108

    
109
db_engines = {
110
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
111
    'PostgreSQL': ('psycopg2', {}),
112
}
113

    
114
DatabaseErrors_set = set([DbException])
115
DatabaseErrors = tuple(DatabaseErrors_set)
116

    
117
def _add_module(module):
118
    DatabaseErrors_set.add(module.DatabaseError)
119
    global DatabaseErrors
120
    DatabaseErrors = tuple(DatabaseErrors_set)
121

    
122
def db_config_str(db_config):
123
    return db_config['engine']+' database '+db_config['database']
124

    
125
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
126

    
127
log_debug_none = lambda msg: None
128

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

    
258
connect = DbConn
259

    
260
##### Querying
261

    
262
def run_raw_query(db, *args, **kw_args):
263
    '''For params, see DbConn.run_query()'''
264
    return db.run_query(*args, **kw_args)
265

    
266
def mogrify(db, query, params):
267
    module = util.root_module(db.db)
268
    if module == 'psycopg2': return db.db.cursor().mogrify(query, params)
269
    else: raise NotImplementedError("Can't mogrify query for "+module+
270
        ' database')
271

    
272
##### Recoverable querying
273

    
274
def with_savepoint(db, func):
275
    savepoint = 'savepoint_'+str(rand.rand_int()) # must be unique
276
    run_raw_query(db, 'SAVEPOINT '+savepoint)
277
    try: return_val = func()
278
    except:
279
        run_raw_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)
280
        raise
281
    else:
282
        run_raw_query(db, 'RELEASE SAVEPOINT '+savepoint)
283
        return return_val
284

    
285
def run_query(db, query, params=None, recover=None, cacheable=False):
286
    if recover == None: recover = False
287
    
288
    def run(): return run_raw_query(db, query, params, cacheable)
289
    if recover and not db.is_cached(query, params):
290
        return with_savepoint(db, run)
291
    else: return run() # don't need savepoint if cached
292

    
293
##### Basic queries
294

    
295
def run_query_into(db, query, params, into=None, *args, **kw_args):
296
    '''Outputs a query to a temp table.
297
    For params, see run_query().
298
    '''
299
    if into == None: return run_query(db, query, params, *args, **kw_args)
300
    else: # place rows in temp table
301
        check_name(into)
302
        
303
        run_query(db, 'DROP TABLE IF EXISTS '+into+' CASCADE', *args, **kw_args)
304
        return run_query(db, 'CREATE TEMP TABLE '+into+' AS '+query, params,
305
            *args, **kw_args) # CREATE TABLE sets rowcount to # rows in query
306

    
307
def mk_select(db, table, fields=None, conds=None, limit=None, start=None,
308
    table_is_esc=False):
309
    '''
310
    @param fields Use None to select all fields in the table
311
    @param table_is_esc Whether the table name has already been escaped
312
    @return tuple(query, params)
313
    '''
314
    def esc_name_(name): return esc_name(db, name)
315
    
316
    if conds == None: conds = {}
317
    assert limit == None or type(limit) == int
318
    assert start == None or type(start) == int
319
    if not table_is_esc: table = esc_name_(table)
320
    
321
    params = []
322
    
323
    def parse_col(field):
324
        '''Parses fields'''
325
        if isinstance(field, tuple): # field is literal values
326
            value, col = field
327
            sql_ = '%s'
328
            params.append(value)
329
            if col != None: sql_ += ' AS '+esc_name_(col)
330
        else: sql_ = esc_name_(field) # field is col name
331
        return sql_
332
    def cond(entry):
333
        '''Parses conditions'''
334
        col, value = entry
335
        cond_ = esc_name_(col)+' '
336
        if value == None: cond_ += 'IS'
337
        else: cond_ += '='
338
        cond_ += ' %s'
339
        return cond_
340
    
341
    query = 'SELECT '
342
    if fields == None: query += '*'
343
    else: query += ', '.join(map(parse_col, fields))
344
    query += ' FROM '+table
345
    
346
    missing = True
347
    if conds != {}:
348
        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))
349
        params += conds.values()
350
        missing = False
351
    if limit != None: query += ' LIMIT '+str(limit); missing = False
352
    if start != None:
353
        if start != 0: query += ' OFFSET '+str(start)
354
        missing = False
355
    if missing: warnings.warn(DbWarning(
356
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
357
    
358
    return (query, params)
359

    
360
def select(db, *args, **kw_args):
361
    '''For params, see mk_select() and run_query()'''
362
    recover = kw_args.pop('recover', None)
363
    cacheable = kw_args.pop('cacheable', True)
364
    
365
    query, params = mk_select(db, *args, **kw_args)
366
    return run_query(db, query, params, recover, cacheable)
367

    
368
def mk_insert_select(db, table, cols=None, select_query=None, params=None,
369
    returning=None, embeddable=False, table_is_esc=False):
370
    '''
371
    @param returning str|None An inserted column (such as pkey) to return
372
    @param embeddable Whether the query should be embeddable as a nested SELECT.
373
        Warning: If you set this and cacheable=True when the query is run, the
374
        query will be fully cached, not just if it raises an exception.
375
    @param table_is_esc Whether the table name has already been escaped
376
    '''
377
    if select_query == None: select_query = 'DEFAULT VALUES'
378
    if cols == []: cols = None # no cols (all defaults) = unknown col names
379
    if not table_is_esc: check_name(table)
380
    
381
    # Build query
382
    query = 'INSERT INTO '+table
383
    if cols != None:
384
        map(check_name, cols)
385
        query += ' ('+', '.join(cols)+')'
386
    query += ' '+select_query
387
    
388
    if returning != None:
389
        check_name(returning)
390
        query += ' RETURNING '+returning
391
    
392
    if embeddable:
393
        # Create function
394
        function = 'pg_temp.'+('_'.join(map(clean_name,
395
            ['insert', table] + cols)))
396
        return_type = 'SETOF '+table+'.'+returning+'%TYPE'
397
        function_query = '''\
398
CREATE OR REPLACE FUNCTION '''+function+'''() RETURNS '''+return_type+'''
399
    LANGUAGE sql
400
    AS $$'''+mogrify(db, query, params)+''';$$;
401
'''
402
        run_query(db, function_query, cacheable=True)
403
        
404
        # Return query that uses function
405
        return mk_select(db, function+'() AS f ('+returning+')',
406
            table_is_esc=True) # function alias is required in AS clause
407
    
408
    return (query, params)
409

    
410
def insert_select(db, *args, **kw_args):
411
    '''For params, see mk_insert_select() and run_query_into()
412
    @param into Name of temp table to place RETURNING values in
413
    '''
414
    into = kw_args.pop('into', None)
415
    if into != None: kw_args['embeddable'] = True
416
    recover = kw_args.pop('recover', None)
417
    cacheable = kw_args.pop('cacheable', True)
418
    
419
    query, params = mk_insert_select(db, *args, **kw_args)
420
    return run_query_into(db, query, params, into, recover, cacheable)
421

    
422
default = object() # tells insert() to use the default value for a column
423

    
424
def insert(db, table, row, *args, **kw_args):
425
    '''For params, see insert_select()'''
426
    if lists.is_seq(row): cols = None
427
    else:
428
        cols = row.keys()
429
        row = row.values()
430
    row = list(row) # ensure that "!= []" works
431
    
432
    # Check for special values
433
    labels = []
434
    values = []
435
    for value in row:
436
        if value == default: labels.append('DEFAULT')
437
        else:
438
            labels.append('%s')
439
            values.append(value)
440
    
441
    # Build query
442
    if values != []: query = ' VALUES ('+(', '.join(labels))+')'
443
    else: query = None
444
    
445
    return insert_select(db, table, cols, query, values, *args, **kw_args)
446

    
447
def last_insert_id(db):
448
    module = util.root_module(db.db)
449
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
450
    elif module == 'MySQLdb': return db.insert_id()
451
    else: return None
452

    
453
def truncate(db, table, schema='public'):
454
    return run_query(db, 'TRUNCATE '+qual_name(db, schema, table)+' CASCADE')
455

    
456
##### Database structure queries
457

    
458
def pkey(db, table, recover=None, table_is_esc=False):
459
    '''Assumed to be first column in table'''
460
    return col_names(select(db, table, limit=0, recover=recover,
461
        table_is_esc=table_is_esc)).next()
462

    
463
def index_cols(db, table, index):
464
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
465
    automatically created. When you don't know whether something is a UNIQUE
466
    constraint or a UNIQUE index, use this function.'''
467
    check_name(table)
468
    check_name(index)
469
    module = util.root_module(db.db)
470
    if module == 'psycopg2':
471
        return list(values(run_query(db, '''\
472
SELECT attname
473
FROM
474
(
475
        SELECT attnum, attname
476
        FROM pg_index
477
        JOIN pg_class index ON index.oid = indexrelid
478
        JOIN pg_class table_ ON table_.oid = indrelid
479
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
480
        WHERE
481
            table_.relname = %(table)s
482
            AND index.relname = %(index)s
483
    UNION
484
        SELECT attnum, attname
485
        FROM
486
        (
487
            SELECT
488
                indrelid
489
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
490
                    AS indkey
491
            FROM pg_index
492
            JOIN pg_class index ON index.oid = indexrelid
493
            JOIN pg_class table_ ON table_.oid = indrelid
494
            WHERE
495
                table_.relname = %(table)s
496
                AND index.relname = %(index)s
497
        ) s
498
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
499
) s
500
ORDER BY attnum
501
''',
502
            {'table': table, 'index': index}, cacheable=True)))
503
    else: raise NotImplementedError("Can't list index columns for "+module+
504
        ' database')
505

    
506
def constraint_cols(db, table, constraint):
507
    check_name(table)
508
    check_name(constraint)
509
    module = util.root_module(db.db)
510
    if module == 'psycopg2':
511
        return list(values(run_query(db, '''\
512
SELECT attname
513
FROM pg_constraint
514
JOIN pg_class ON pg_class.oid = conrelid
515
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
516
WHERE
517
    relname = %(table)s
518
    AND conname = %(constraint)s
519
ORDER BY attnum
520
''',
521
            {'table': table, 'constraint': constraint})))
522
    else: raise NotImplementedError("Can't list constraint columns for "+module+
523
        ' database')
524

    
525
row_num_col = '_row_num'
526

    
527
def add_row_num(db, table):
528
    '''Adds a row number column to a table. Its name is in row_num_col.'''
529
    check_name(table)
530
    run_query(db, 'ALTER TABLE '+table+' ADD COLUMN '+row_num_col
531
        +' serial NOT NULL')
532

    
533
def tables(db, schema='public', table_like='%'):
534
    module = util.root_module(db.db)
535
    params = {'schema': schema, 'table_like': table_like}
536
    if module == 'psycopg2':
537
        return values(run_query(db, '''\
538
SELECT tablename
539
FROM pg_tables
540
WHERE
541
    schemaname = %(schema)s
542
    AND tablename LIKE %(table_like)s
543
ORDER BY tablename
544
''',
545
            params, cacheable=True))
546
    elif module == 'MySQLdb':
547
        return values(run_query(db, 'SHOW TABLES LIKE %(table_like)s', params,
548
            cacheable=True))
549
    else: raise NotImplementedError("Can't list tables for "+module+' database')
550

    
551
##### Database management
552

    
553
def empty_db(db, schema='public', **kw_args):
554
    '''For kw_args, see tables()'''
555
    for table in tables(db, schema, **kw_args): truncate(db, table, schema)
556

    
557
##### Heuristic queries
558

    
559
def with_parsed_errors(db, func):
560
    '''Translates known DB errors to typed exceptions'''
561
    try: return func()
562
    except Exception, e:
563
        msg = str(e)
564
        match = re.search(r'duplicate key value violates unique constraint '
565
            r'"(([^\W_]+)_[^"]+)"', msg)
566
        if match:
567
            constraint, table = match.groups()
568
            try: cols = index_cols(db, table, constraint)
569
            except NotImplementedError: raise e
570
            else: raise DuplicateKeyException(cols, e)
571
        match = re.search(r'null value in column "(\w+)" violates not-null '
572
            'constraint', msg)
573
        if match: raise NullValueException([match.group(1)], e)
574
        raise # no specific exception raised
575

    
576
def try_insert(db, table, row, returning=None):
577
    '''Recovers from errors'''
578
    return with_parsed_errors(db, lambda: insert(db, table, row, returning,
579
        recover=True))
580

    
581
def put(db, table, row, pkey_=None, row_ct_ref=None):
582
    '''Recovers from errors.
583
    Only works under PostgreSQL (uses INSERT RETURNING).
584
    '''
585
    if pkey_ == None: pkey_ = pkey(db, table, recover=True)
586
    
587
    try:
588
        cur = try_insert(db, table, row, pkey_)
589
        if row_ct_ref != None and cur.rowcount >= 0:
590
            row_ct_ref[0] += cur.rowcount
591
        return value(cur)
592
    except DuplicateKeyException, e:
593
        return value(select(db, table, [pkey_],
594
            util.dict_subset_right_join(row, e.cols), recover=True))
595

    
596
def get(db, table, row, pkey, row_ct_ref=None, create=False):
597
    '''Recovers from errors'''
598
    try: return value(select(db, table, [pkey], row, 1, recover=True))
599
    except StopIteration:
600
        if not create: raise
601
        return put(db, table, row, pkey, row_ct_ref) # insert new row
602

    
603
def put_table(db, out_table, out_cols, in_tables, in_cols, pkey,
604
    row_ct_ref=None, table_is_esc=False):
605
    '''Recovers from errors.
606
    Only works under PostgreSQL (uses INSERT RETURNING).
607
    @return Name of the table where the pkeys (from INSERT RETURNING) are made
608
        available
609
    '''
610
    pkeys_table = clean_name(out_table)+'_pkeys'
611
    def insert_():
612
        return insert_select(db, out_table, out_cols,
613
            *mk_select(db, in_tables[0], in_cols, table_is_esc=table_is_esc),
614
            returning=pkey, into=pkeys_table, recover=True,
615
            table_is_esc=table_is_esc)
616
    try:
617
        cur = with_parsed_errors(db, insert_)
618
        if row_ct_ref != None and cur.rowcount >= 0:
619
            row_ct_ref[0] += cur.rowcount
620
        
621
        # Add row_num to pkeys_table, so it can be joined with in_table's pkeys
622
        add_row_num(db, pkeys_table)
623
        
624
        return pkeys_table
625
    except DuplicateKeyException, e: raise
(22-22/33)