Project

General

Profile

1 11 aaronmk
# Database access
2
3 1869 aaronmk
import copy
4 11 aaronmk
import re
5
import sys
6 865 aaronmk
import warnings
7 11 aaronmk
8 300 aaronmk
import exc
9 1889 aaronmk
from Proxy import Proxy
10 1872 aaronmk
import rand
11 862 aaronmk
import strings
12 131 aaronmk
import util
13 11 aaronmk
14 832 aaronmk
##### Exceptions
15
16 135 aaronmk
def get_cur_query(cur):
17
    if hasattr(cur, 'query'): return cur.query
18
    elif hasattr(cur, '_last_executed'): return cur._last_executed
19
    else: return None
20 14 aaronmk
21 300 aaronmk
def _add_cursor_info(e, cur): exc.add_msg(e, 'query: '+get_cur_query(cur))
22 135 aaronmk
23 300 aaronmk
class DbException(exc.ExceptionWithCause):
24 14 aaronmk
    def __init__(self, msg, cause=None, cur=None):
25 300 aaronmk
        exc.ExceptionWithCause.__init__(self, msg, cause)
26 14 aaronmk
        if cur != None: _add_cursor_info(self, cur)
27
28 360 aaronmk
class NameException(DbException): pass
29
30 468 aaronmk
class ExceptionWithColumns(DbException):
31
    def __init__(self, cols, cause=None):
32
        DbException.__init__(self, 'columns: ' + ', '.join(cols), cause)
33
        self.cols = cols
34 11 aaronmk
35 468 aaronmk
class DuplicateKeyException(ExceptionWithColumns): pass
36 13 aaronmk
37 468 aaronmk
class NullValueException(ExceptionWithColumns): pass
38 13 aaronmk
39 89 aaronmk
class EmptyRowException(DbException): pass
40
41 865 aaronmk
##### Warnings
42
43
class DbWarning(UserWarning): pass
44
45 832 aaronmk
##### Input validation
46
47 11 aaronmk
def check_name(name):
48
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
49
        +'" may contain only alphanumeric characters and _')
50
51 643 aaronmk
def esc_name(db, name):
52 1849 aaronmk
    module = util.root_module(db.db)
53 645 aaronmk
    if module == 'psycopg2': return name
54
        # Don't enclose in quotes because this disables case-insensitivity
55 643 aaronmk
    elif module == 'MySQLdb': quote = '`'
56 645 aaronmk
    else: raise NotImplementedError("Can't escape name for "+module+' database')
57 643 aaronmk
    return quote + name.replace(quote, '') + quote
58
59 1869 aaronmk
##### Database connections
60 1849 aaronmk
61 1869 aaronmk
db_engines = {
62
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
63
    'PostgreSQL': ('psycopg2', {}),
64
}
65
66
DatabaseErrors_set = set([DbException])
67
DatabaseErrors = tuple(DatabaseErrors_set)
68
69
def _add_module(module):
70
    DatabaseErrors_set.add(module.DatabaseError)
71
    global DatabaseErrors
72
    DatabaseErrors = tuple(DatabaseErrors_set)
73
74
def db_config_str(db_config):
75
    return db_config['engine']+' database '+db_config['database']
76
77 1849 aaronmk
class DbConn:
78 1889 aaronmk
    def __init__(self, db_config, serializable=True, debug=False):
79 1869 aaronmk
        self.db_config = db_config
80
        self.serializable = serializable
81 1889 aaronmk
        self.debug = debug
82 1869 aaronmk
83
        self.__db = None
84 1849 aaronmk
        self.pkeys = {}
85
        self.index_cols = {}
86 1889 aaronmk
        self.query_results = {}
87 1869 aaronmk
88
    def __getattr__(self, name):
89
        if name == '__dict__': raise Exception('getting __dict__')
90
        if name == 'db': return self._db()
91
        else: raise AttributeError()
92
93
    def __getstate__(self):
94
        state = copy.copy(self.__dict__) # shallow copy
95
        state['_DbConn__db'] = None # don't pickle the connection
96
        return state
97
98
    def _db(self):
99
        if self.__db == None:
100
            # Process db_config
101
            db_config = self.db_config.copy() # don't modify input!
102
            module_name, mappings = db_engines[db_config.pop('engine')]
103
            module = __import__(module_name)
104
            _add_module(module)
105
            for orig, new in mappings.iteritems():
106
                try: util.rename_key(db_config, orig, new)
107
                except KeyError: pass
108
109
            # Connect
110
            self.__db = module.connect(**db_config)
111
112
            # Configure connection
113
            if self.serializable: run_raw_query(self,
114
                'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
115
116
        return self.__db
117 1889 aaronmk
118 1891 aaronmk
    class DbCursor(Proxy):
119
        def __init__(self, outer):
120
            Proxy.__init__(self, outer.db.cursor())
121
            self.outer = outer
122
            self.result = []
123 1889 aaronmk
124 1891 aaronmk
        def fetchone(self):
125
            row = self.inner.fetchone()
126
            if row == None: # fetched all rows
127
                self.outer.query_results[self.inner.query] = self.result
128
            else: self.result.append(row)
129
            return row
130 1889 aaronmk
131 1891 aaronmk
    class CacheCursor:
132
        def __init__(self, result): self.iter = iter(result)
133
134
        def fetchone(self):
135
            try: return self.iter.next()
136
            except StopIteration: return None
137
138 1889 aaronmk
    def run_query(self, query, params=None):
139 1891 aaronmk
        cur = self.DbCursor(self)
140
        query = cur.mogrify(query, params)
141
        try: return self.CacheCursor(self.query_results[query])
142
        except KeyError:
143
            try: cur.execute(query)
144
            except Exception, e:
145
                _add_cursor_info(e, cur)
146
                raise
147
            if self.debug:
148
                sys.stderr.write(strings.one_line(get_cur_query(cur))+'\n')
149
            return cur
150 1849 aaronmk
151 1869 aaronmk
connect = DbConn
152
153 832 aaronmk
##### Querying
154
155 1889 aaronmk
def run_raw_query(db, query, params=None): return db.run_query(query, params)
156 11 aaronmk
157 832 aaronmk
##### Recoverable querying
158 15 aaronmk
159 11 aaronmk
def with_savepoint(db, func):
160 1872 aaronmk
    savepoint = 'savepoint_'+str(rand.rand_int()) # must be unique
161 830 aaronmk
    run_raw_query(db, 'SAVEPOINT '+savepoint)
162 11 aaronmk
    try: return_val = func()
163
    except:
164 830 aaronmk
        run_raw_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)
165 11 aaronmk
        raise
166
    else:
167 830 aaronmk
        run_raw_query(db, 'RELEASE SAVEPOINT '+savepoint)
168 11 aaronmk
        return return_val
169
170 830 aaronmk
def run_query(db, query, params=None, recover=None):
171
    if recover == None: recover = False
172
173
    def run(): return run_raw_query(db, query, params)
174
    if recover: return with_savepoint(db, run)
175
    else: return run()
176
177 832 aaronmk
##### Result retrieval
178
179 1135 aaronmk
def col_names(cur): return (col[0] for col in cur.description)
180 832 aaronmk
181
def rows(cur): return iter(lambda: cur.fetchone(), None)
182
183
def row(cur): return rows(cur).next()
184
185
def value(cur): return row(cur)[0]
186
187
def values(cur): return iter(lambda: value(cur), None)
188
189
def value_or_none(cur):
190
    try: return value(cur)
191
    except StopIteration: return None
192
193
##### Basic queries
194
195 1135 aaronmk
def select(db, table, fields=None, conds=None, limit=None, start=None,
196
    recover=None):
197
    '''@param fields Use None to select all fields in the table'''
198
    if conds == None: conds = {}
199 135 aaronmk
    assert limit == None or type(limit) == int
200 865 aaronmk
    assert start == None or type(start) == int
201 15 aaronmk
    check_name(table)
202 1135 aaronmk
    if fields != None: map(check_name, fields)
203 15 aaronmk
    map(check_name, conds.keys())
204 865 aaronmk
205 11 aaronmk
    def cond(entry):
206 13 aaronmk
        col, value = entry
207 644 aaronmk
        cond_ = esc_name(db, col)+' '
208 11 aaronmk
        if value == None: cond_ += 'IS'
209
        else: cond_ += '='
210
        cond_ += ' %s'
211
        return cond_
212 1135 aaronmk
    query = 'SELECT '
213
    if fields == None: query += '*'
214
    else: query += ', '.join([esc_name(db, field) for field in fields])
215
    query += ' FROM '+esc_name(db, table)
216 865 aaronmk
217
    missing = True
218 89 aaronmk
    if conds != {}:
219
        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))
220 865 aaronmk
        missing = False
221
    if limit != None: query += ' LIMIT '+str(limit); missing = False
222
    if start != None:
223
        if start != 0: query += ' OFFSET '+str(start)
224
        missing = False
225
    if missing: warnings.warn(DbWarning(
226
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
227
228 830 aaronmk
    return run_query(db, query, conds.values(), recover)
229 11 aaronmk
230 1554 aaronmk
def insert(db, table, row, returning=None, recover=None):
231
    '''@param returning str|None An inserted column (such as pkey) to return'''
232 11 aaronmk
    check_name(table)
233 13 aaronmk
    cols = row.keys()
234 15 aaronmk
    map(check_name, cols)
235 89 aaronmk
    query = 'INSERT INTO '+table
236 1554 aaronmk
237 89 aaronmk
    if row != {}: query += ' ('+', '.join(cols)+') VALUES ('\
238
        +', '.join(['%s']*len(cols))+')'
239
    else: query += ' DEFAULT VALUES'
240 1554 aaronmk
241
    if returning != None:
242
        check_name(returning)
243
        query += ' RETURNING '+returning
244
245 830 aaronmk
    return run_query(db, query, row.values(), recover)
246 11 aaronmk
247 135 aaronmk
def last_insert_id(db):
248 1849 aaronmk
    module = util.root_module(db.db)
249 135 aaronmk
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
250
    elif module == 'MySQLdb': return db.insert_id()
251
    else: return None
252 13 aaronmk
253 832 aaronmk
def truncate(db, table):
254
    check_name(table)
255 869 aaronmk
    return run_raw_query(db, 'TRUNCATE '+table+' CASCADE')
256 832 aaronmk
257
##### Database structure queries
258
259 1850 aaronmk
def pkey(db, table, recover=None):
260 832 aaronmk
    '''Assumed to be first column in table'''
261
    check_name(table)
262 1850 aaronmk
    if table not in db.pkeys:
263
        db.pkeys[table] = col_names(run_query(db,
264 1135 aaronmk
            'SELECT * FROM '+table+' LIMIT 0', recover=recover)).next()
265 1850 aaronmk
    return db.pkeys[table]
266 832 aaronmk
267 853 aaronmk
def index_cols(db, table, index):
268
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
269
    automatically created. When you don't know whether something is a UNIQUE
270
    constraint or a UNIQUE index, use this function.'''
271
    check_name(table)
272
    check_name(index)
273 1852 aaronmk
    lookup = (table, index)
274
    if lookup not in db.index_cols:
275
        module = util.root_module(db.db)
276
        if module == 'psycopg2':
277
            db.index_cols[lookup] = list(values(run_query(db, '''\
278 853 aaronmk
SELECT attname
279 866 aaronmk
FROM
280
(
281
        SELECT attnum, attname
282
        FROM pg_index
283
        JOIN pg_class index ON index.oid = indexrelid
284
        JOIN pg_class table_ ON table_.oid = indrelid
285
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
286
        WHERE
287
            table_.relname = %(table)s
288
            AND index.relname = %(index)s
289
    UNION
290
        SELECT attnum, attname
291
        FROM
292
        (
293
            SELECT
294
                indrelid
295
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
296
                    AS indkey
297
            FROM pg_index
298
            JOIN pg_class index ON index.oid = indexrelid
299
            JOIN pg_class table_ ON table_.oid = indrelid
300
            WHERE
301
                table_.relname = %(table)s
302
                AND index.relname = %(index)s
303
        ) s
304
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
305
) s
306 853 aaronmk
ORDER BY attnum
307
''',
308 1852 aaronmk
                {'table': table, 'index': index})))
309
        else: raise NotImplementedError("Can't list index columns for "+module+
310
            ' database')
311
    return db.index_cols[lookup]
312 853 aaronmk
313 464 aaronmk
def constraint_cols(db, table, constraint):
314
    check_name(table)
315
    check_name(constraint)
316 1849 aaronmk
    module = util.root_module(db.db)
317 464 aaronmk
    if module == 'psycopg2':
318
        return list(values(run_query(db, '''\
319
SELECT attname
320
FROM pg_constraint
321
JOIN pg_class ON pg_class.oid = conrelid
322
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
323
WHERE
324
    relname = %(table)s
325
    AND conname = %(constraint)s
326
ORDER BY attnum
327
''',
328
            {'table': table, 'constraint': constraint})))
329
    else: raise NotImplementedError("Can't list constraint columns for "+module+
330
        ' database')
331
332 832 aaronmk
def tables(db):
333 1849 aaronmk
    module = util.root_module(db.db)
334 832 aaronmk
    if module == 'psycopg2':
335
        return values(run_query(db, "SELECT tablename from pg_tables "
336
            "WHERE schemaname = 'public' ORDER BY tablename"))
337
    elif module == 'MySQLdb': return values(run_query(db, 'SHOW TABLES'))
338
    else: raise NotImplementedError("Can't list tables for "+module+' database')
339 830 aaronmk
340 833 aaronmk
##### Database management
341
342
def empty_db(db):
343
    for table in tables(db): truncate(db, table)
344
345 832 aaronmk
##### Heuristic queries
346
347 1554 aaronmk
def try_insert(db, table, row, returning=None):
348 830 aaronmk
    '''Recovers from errors'''
349 1554 aaronmk
    try: return insert(db, table, row, returning, recover=True)
350 46 aaronmk
    except Exception, e:
351
        msg = str(e)
352 465 aaronmk
        match = re.search(r'duplicate key value violates unique constraint '
353
            r'"(([^\W_]+)_[^"]+)"', msg)
354
        if match:
355
            constraint, table = match.groups()
356 854 aaronmk
            try: cols = index_cols(db, table, constraint)
357 465 aaronmk
            except NotImplementedError: raise e
358 851 aaronmk
            else: raise DuplicateKeyException(cols, e)
359 13 aaronmk
        match = re.search(r'null value in column "(\w+)" violates not-null '
360
            'constraint', msg)
361 470 aaronmk
        if match: raise NullValueException([match.group(1)], e)
362 13 aaronmk
        raise # no specific exception raised
363 11 aaronmk
364 471 aaronmk
def put(db, table, row, pkey, row_ct_ref=None):
365 1554 aaronmk
    '''Recovers from errors.
366
    Only works under PostgreSQL (uses `INSERT ... RETURNING`)'''
367 471 aaronmk
    try:
368 1554 aaronmk
        cur = try_insert(db, table, row, pkey)
369
        if row_ct_ref != None and cur.rowcount >= 0:
370
            row_ct_ref[0] += cur.rowcount
371
        return value(cur)
372 471 aaronmk
    except DuplicateKeyException, e:
373 1069 aaronmk
        return value(select(db, table, [pkey],
374
            util.dict_subset_right_join(row, e.cols), recover=True))
375 471 aaronmk
376 473 aaronmk
def get(db, table, row, pkey, row_ct_ref=None, create=False):
377 830 aaronmk
    '''Recovers from errors'''
378
    try: return value(select(db, table, [pkey], row, 1, recover=True))
379 14 aaronmk
    except StopIteration:
380 40 aaronmk
        if not create: raise
381 471 aaronmk
        return put(db, table, row, pkey, row_ct_ref) # insert new row