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
##### Database connections
75

    
76
db_config_names = ['engine', 'host', 'user', 'password', 'database']
77

    
78
db_engines = {
79
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
80
    'PostgreSQL': ('psycopg2', {}),
81
}
82

    
83
DatabaseErrors_set = set([DbException])
84
DatabaseErrors = tuple(DatabaseErrors_set)
85

    
86
def _add_module(module):
87
    DatabaseErrors_set.add(module.DatabaseError)
88
    global DatabaseErrors
89
    DatabaseErrors = tuple(DatabaseErrors_set)
90

    
91
def db_config_str(db_config):
92
    return db_config['engine']+' database '+db_config['database']
93

    
94
def _query_lookup(query, params): return (query, dicts.make_hashable(params))
95

    
96
log_debug_none = lambda msg: None
97

    
98
class DbConn:
99
    def __init__(self, db_config, serializable=True, log_debug=log_debug_none):
100
        self.db_config = db_config
101
        self.serializable = serializable
102
        self.log_debug = log_debug
103
        
104
        self.__db = None
105
        self.query_results = {}
106
    
107
    def __getattr__(self, name):
108
        if name == '__dict__': raise Exception('getting __dict__')
109
        if name == 'db': return self._db()
110
        else: raise AttributeError()
111
    
112
    def __getstate__(self):
113
        state = copy.copy(self.__dict__) # shallow copy
114
        state['log_debug'] = None # don't pickle the debug callback
115
        state['_DbConn__db'] = None # don't pickle the connection
116
        return state
117
    
118
    def _db(self):
119
        if self.__db == None:
120
            # Process db_config
121
            db_config = self.db_config.copy() # don't modify input!
122
            module_name, mappings = db_engines[db_config.pop('engine')]
123
            module = __import__(module_name)
124
            _add_module(module)
125
            for orig, new in mappings.iteritems():
126
                try: util.rename_key(db_config, orig, new)
127
                except KeyError: pass
128
            
129
            # Connect
130
            self.__db = module.connect(**db_config)
131
            
132
            # Configure connection
133
            if self.serializable: run_raw_query(self,
134
                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
135
        
136
        return self.__db
137
    
138
    class DbCursor(Proxy):
139
        def __init__(self, outer):
140
            Proxy.__init__(self, outer.db.cursor())
141
            self.query_results = outer.query_results
142
            self.query_lookup = None
143
            self.result = []
144
        
145
        def execute(self, query, params=None):
146
            self._is_insert = query.upper().find('INSERT') >= 0
147
            self.query_lookup = _query_lookup(query, params)
148
            try: return_value = self.inner.execute(query, params)
149
            except Exception, e:
150
                self.result = e # cache the exception as the result
151
                self._cache_result()
152
                raise
153
            finally: self.query = get_cur_query(self.inner)
154
            # Fetch all rows so result will be cached
155
            if self.rowcount == 0 and not self._is_insert: consume_rows(self)
156
            return return_value
157
        
158
        def fetchone(self):
159
            row = self.inner.fetchone()
160
            if row != None: self.result.append(row)
161
            # otherwise, fetched all rows
162
            else: self._cache_result()
163
            return row
164
        
165
        def _cache_result(self):
166
            # For inserts, only cache exceptions since inserts are not
167
            # idempotent, but an invalid insert will always be invalid
168
            if self.query_results != None and (not self._is_insert
169
                or isinstance(self.result, Exception)):
170
                
171
                assert self.query_lookup != None
172
                self.query_results[self.query_lookup] = self.CacheCursor(
173
                    util.dict_subset(dicts.AttrsDictView(self),
174
                    ['query', 'result', 'rowcount', 'description']))
175
        
176
        class CacheCursor:
177
            def __init__(self, cached_result): self.__dict__ = cached_result
178
            
179
            def execute(self, *args, **kw_args):
180
                if isinstance(self.result, Exception): raise self.result
181
                # otherwise, result is a rows list
182
                self.iter = iter(self.result)
183
            
184
            def fetchone(self):
185
                try: return self.iter.next()
186
                except StopIteration: return None
187
    
188
    def run_query(self, query, params=None, cacheable=False):
189
        used_cache = False
190
        try:
191
            # Get cursor
192
            if cacheable:
193
                query_lookup = _query_lookup(query, params)
194
                try:
195
                    cur = self.query_results[query_lookup]
196
                    used_cache = True
197
                except KeyError: cur = self.DbCursor(self)
198
            else: cur = self.db.cursor()
199
            
200
            # Run query
201
            try: cur.execute(query, params)
202
            except Exception, e:
203
                _add_cursor_info(e, cur)
204
                raise
205
        finally:
206
            if self.log_debug != log_debug_none: # only compute msg if needed
207
                if used_cache: cache_status = 'Cache hit'
208
                elif cacheable: cache_status = 'Cache miss'
209
                else: cache_status = 'Non-cacheable'
210
                self.log_debug(cache_status+': '
211
                    +strings.one_line(get_cur_query(cur)))
212
        
213
        return cur
214
    
215
    def is_cached(self, query, params=None):
216
        return _query_lookup(query, params) in self.query_results
217

    
218
connect = DbConn
219

    
220
##### Input validation
221

    
222
def check_name(name):
223
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
224
        +'" may contain only alphanumeric characters and _')
225

    
226
def esc_name_by_module(module, name, preserve_case=False):
227
    if module == 'psycopg2':
228
        if preserve_case: quote = '"'
229
        # Don't enclose in quotes because this disables case-insensitivity
230
        else: return name
231
    elif module == 'MySQLdb': quote = '`'
232
    else: raise NotImplementedError("Can't escape name for "+module+' database')
233
    return quote + name.replace(quote, '') + quote
234

    
235
def esc_name_by_engine(engine, name, **kw_args):
236
    return esc_name_by_module(db_engines[engine][0], name, **kw_args)
237

    
238
def esc_name(db, name, **kw_args):
239
    return esc_name_by_module(util.root_module(db.db), name, **kw_args)
240

    
241
##### Querying
242

    
243
def run_raw_query(db, *args, **kw_args):
244
    '''For args, see DbConn.run_query()'''
245
    return db.run_query(*args, **kw_args)
246

    
247
##### Recoverable querying
248

    
249
def with_savepoint(db, func):
250
    savepoint = 'savepoint_'+str(rand.rand_int()) # must be unique
251
    run_raw_query(db, 'SAVEPOINT '+savepoint)
252
    try: return_val = func()
253
    except:
254
        run_raw_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)
255
        raise
256
    else:
257
        run_raw_query(db, 'RELEASE SAVEPOINT '+savepoint)
258
        return return_val
259

    
260
def run_query(db, query, params=None, recover=None, cacheable=False):
261
    if recover == None: recover = False
262
    
263
    def run(): return run_raw_query(db, query, params, cacheable)
264
    if recover and not db.is_cached(query, params):
265
        return with_savepoint(db, run)
266
    else: return run() # don't need savepoint if cached
267

    
268
##### Basic queries
269

    
270
def select(db, table, fields=None, conds=None, limit=None, start=None,
271
    recover=None, cacheable=True):
272
    '''@param fields Use None to select all fields in the table'''
273
    if conds == None: conds = {}
274
    assert limit == None or type(limit) == int
275
    assert start == None or type(start) == int
276
    check_name(table)
277
    if fields != None: map(check_name, fields)
278
    map(check_name, conds.keys())
279
    
280
    def cond(entry):
281
        col, value = entry
282
        cond_ = esc_name(db, col)+' '
283
        if value == None: cond_ += 'IS'
284
        else: cond_ += '='
285
        cond_ += ' %s'
286
        return cond_
287
    query = 'SELECT '
288
    if fields == None: query += '*'
289
    else: query += ', '.join([esc_name(db, field) for field in fields])
290
    query += ' FROM '+esc_name(db, table)
291
    
292
    missing = True
293
    if conds != {}:
294
        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))
295
        missing = False
296
    if limit != None: query += ' LIMIT '+str(limit); missing = False
297
    if start != None:
298
        if start != 0: query += ' OFFSET '+str(start)
299
        missing = False
300
    if missing: warnings.warn(DbWarning(
301
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
302
    
303
    return run_query(db, query, conds.values(), recover, cacheable)
304

    
305
default = object() # tells insert() to use the default value for a column
306

    
307
def insert(db, table, row, returning=None, recover=None, cacheable=True,
308
    table_is_esc=False):
309
    '''
310
    @param returning str|None An inserted column (such as pkey) to return
311
    @param table_is_esc Whether the table name has already been escaped
312
    '''
313
    if not table_is_esc: check_name(table)
314
    if lists.is_seq(row): cols = None
315
    else:
316
        cols = row.keys()
317
        row = row.values()
318
        map(check_name, cols)
319
    row = list(row) # ensure that "!= []" works
320
    
321
    # Check for special values
322
    labels = []
323
    values = []
324
    for value in row:
325
        if value == default: labels.append('DEFAULT')
326
        else:
327
            labels.append('%s')
328
            values.append(value)
329
    
330
    # Build query
331
    query = 'INSERT INTO '+table
332
    if values != []:
333
        if cols != None: query += ' ('+', '.join(cols)+')'
334
        query += ' VALUES ('+(', '.join(labels))+')'
335
    else: query += ' DEFAULT VALUES'
336
    
337
    if returning != None:
338
        check_name(returning)
339
        query += ' RETURNING '+returning
340
    
341
    return run_query(db, query, values, recover, cacheable)
342

    
343
def last_insert_id(db):
344
    module = util.root_module(db.db)
345
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
346
    elif module == 'MySQLdb': return db.insert_id()
347
    else: return None
348

    
349
def truncate(db, table):
350
    check_name(table)
351
    return run_raw_query(db, 'TRUNCATE '+table+' CASCADE')
352

    
353
##### Database structure queries
354

    
355
def pkey(db, table, recover=None):
356
    '''Assumed to be first column in table'''
357
    check_name(table)
358
    return col_names(select(db, table, limit=0, recover=recover)).next()
359

    
360
def index_cols(db, table, index):
361
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
362
    automatically created. When you don't know whether something is a UNIQUE
363
    constraint or a UNIQUE index, use this function.'''
364
    check_name(table)
365
    check_name(index)
366
    module = util.root_module(db.db)
367
    if module == 'psycopg2':
368
        return list(values(run_query(db, '''\
369
SELECT attname
370
FROM
371
(
372
        SELECT attnum, attname
373
        FROM pg_index
374
        JOIN pg_class index ON index.oid = indexrelid
375
        JOIN pg_class table_ ON table_.oid = indrelid
376
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
377
        WHERE
378
            table_.relname = %(table)s
379
            AND index.relname = %(index)s
380
    UNION
381
        SELECT attnum, attname
382
        FROM
383
        (
384
            SELECT
385
                indrelid
386
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
387
                    AS indkey
388
            FROM pg_index
389
            JOIN pg_class index ON index.oid = indexrelid
390
            JOIN pg_class table_ ON table_.oid = indrelid
391
            WHERE
392
                table_.relname = %(table)s
393
                AND index.relname = %(index)s
394
        ) s
395
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
396
) s
397
ORDER BY attnum
398
''',
399
            {'table': table, 'index': index}, cacheable=True)))
400
    else: raise NotImplementedError("Can't list index columns for "+module+
401
        ' database')
402

    
403
def constraint_cols(db, table, constraint):
404
    check_name(table)
405
    check_name(constraint)
406
    module = util.root_module(db.db)
407
    if module == 'psycopg2':
408
        return list(values(run_query(db, '''\
409
SELECT attname
410
FROM pg_constraint
411
JOIN pg_class ON pg_class.oid = conrelid
412
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
413
WHERE
414
    relname = %(table)s
415
    AND conname = %(constraint)s
416
ORDER BY attnum
417
''',
418
            {'table': table, 'constraint': constraint})))
419
    else: raise NotImplementedError("Can't list constraint columns for "+module+
420
        ' database')
421

    
422
def tables(db):
423
    module = util.root_module(db.db)
424
    if module == 'psycopg2':
425
        return values(run_query(db, "SELECT tablename from pg_tables "
426
            "WHERE schemaname = 'public' ORDER BY tablename"))
427
    elif module == 'MySQLdb': return values(run_query(db, 'SHOW TABLES'))
428
    else: raise NotImplementedError("Can't list tables for "+module+' database')
429

    
430
##### Database management
431

    
432
def empty_db(db):
433
    for table in tables(db): truncate(db, table)
434

    
435
##### Heuristic queries
436

    
437
def try_insert(db, table, row, returning=None):
438
    '''Recovers from errors'''
439
    try: return insert(db, table, row, returning, recover=True)
440
    except Exception, e:
441
        msg = str(e)
442
        match = re.search(r'duplicate key value violates unique constraint '
443
            r'"(([^\W_]+)_[^"]+)"', msg)
444
        if match:
445
            constraint, table = match.groups()
446
            try: cols = index_cols(db, table, constraint)
447
            except NotImplementedError: raise e
448
            else: raise DuplicateKeyException(cols, e)
449
        match = re.search(r'null value in column "(\w+)" violates not-null '
450
            'constraint', msg)
451
        if match: raise NullValueException([match.group(1)], e)
452
        raise # no specific exception raised
453

    
454
def put(db, table, row, pkey, row_ct_ref=None):
455
    '''Recovers from errors.
456
    Only works under PostgreSQL (uses `INSERT ... RETURNING`)'''
457
    try:
458
        cur = try_insert(db, table, row, pkey)
459
        if row_ct_ref != None and cur.rowcount >= 0:
460
            row_ct_ref[0] += cur.rowcount
461
        return value(cur)
462
    except DuplicateKeyException, e:
463
        return value(select(db, table, [pkey],
464
            util.dict_subset_right_join(row, e.cols), recover=True))
465

    
466
def get(db, table, row, pkey, row_ct_ref=None, create=False):
467
    '''Recovers from errors'''
468
    try: return value(select(db, table, [pkey], row, 1, recover=True))
469
    except StopIteration:
470
        if not create: raise
471
        return put(db, table, row, pkey, row_ct_ref) # insert new row
(22-22/33)