Project

General

Profile

1
# Database access
2

    
3
import random
4
import re
5
import sys
6
import warnings
7

    
8
import exc
9
import strings
10
import util
11

    
12
##### Exceptions
13

    
14
def get_cur_query(cur):
15
    if hasattr(cur, 'query'): return cur.query
16
    elif hasattr(cur, '_last_executed'): return cur._last_executed
17
    else: return None
18

    
19
def _add_cursor_info(e, cur): exc.add_msg(e, 'query: '+get_cur_query(cur))
20

    
21
class DbException(exc.ExceptionWithCause):
22
    def __init__(self, msg, cause=None, cur=None):
23
        exc.ExceptionWithCause.__init__(self, msg, cause)
24
        if cur != None: _add_cursor_info(self, cur)
25

    
26
class NameException(DbException): pass
27

    
28
class ExceptionWithColumns(DbException):
29
    def __init__(self, cols, cause=None):
30
        DbException.__init__(self, 'columns: ' + ', '.join(cols), cause)
31
        self.cols = cols
32

    
33
class DuplicateKeyException(ExceptionWithColumns): pass
34

    
35
class NullValueException(ExceptionWithColumns): pass
36

    
37
class EmptyRowException(DbException): pass
38

    
39
##### Warnings
40

    
41
class DbWarning(UserWarning): pass
42

    
43
##### Input validation
44

    
45
def check_name(name):
46
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
47
        +'" may contain only alphanumeric characters and _')
48

    
49
def esc_name(db, name):
50
    module = util.root_module(db)
51
    if module == 'psycopg2': return name
52
        # Don't enclose in quotes because this disables case-insensitivity
53
    elif module == 'MySQLdb': quote = '`'
54
    else: raise NotImplementedError("Can't escape name for "+module+' database')
55
    return quote + name.replace(quote, '') + quote
56

    
57
##### Querying
58

    
59
def run_raw_query(db, query, params=None):
60
    cur = db.cursor()
61
    try: cur.execute(query, params)
62
    except Exception, e:
63
        _add_cursor_info(e, cur)
64
        raise
65
    if run_raw_query.debug:
66
        sys.stderr.write(strings.one_line(get_cur_query(cur))+'\n')
67
    return cur
68

    
69
##### Recoverable querying
70

    
71
def with_savepoint(db, func):
72
    savepoint = 'savepoint_'+str(random.randint(0, sys.maxint)) # must be unique
73
    run_raw_query(db, 'SAVEPOINT '+savepoint)
74
    try: return_val = func()
75
    except:
76
        run_raw_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)
77
        raise
78
    else:
79
        run_raw_query(db, 'RELEASE SAVEPOINT '+savepoint)
80
        return return_val
81

    
82
def run_query(db, query, params=None, recover=None):
83
    if recover == None: recover = False
84
    
85
    def run(): return run_raw_query(db, query, params)
86
    if recover: return with_savepoint(db, run)
87
    else: return run()
88

    
89
##### Result retrieval
90

    
91
def col(cur, idx): return cur.description[idx][0]
92

    
93
def rows(cur): return iter(lambda: cur.fetchone(), None)
94

    
95
def row(cur): return rows(cur).next()
96

    
97
def value(cur): return row(cur)[0]
98

    
99
def values(cur): return iter(lambda: value(cur), None)
100

    
101
def value_or_none(cur):
102
    try: return value(cur)
103
    except StopIteration: return None
104

    
105
##### Basic queries
106

    
107
def select(db, table, fields, conds, limit=None, start=None, recover=None):
108
    assert limit == None or type(limit) == int
109
    assert start == None or type(start) == int
110
    check_name(table)
111
    map(check_name, fields)
112
    map(check_name, conds.keys())
113
    
114
    def cond(entry):
115
        col, value = entry
116
        cond_ = esc_name(db, col)+' '
117
        if value == None: cond_ += 'IS'
118
        else: cond_ += '='
119
        cond_ += ' %s'
120
        return cond_
121
    query = ('SELECT ' + ', '.join([esc_name(db, field) for field in fields])
122
        + ' FROM '+esc_name(db, table))
123
    
124
    missing = True
125
    if conds != {}:
126
        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))
127
        missing = False
128
    if limit != None: query += ' LIMIT '+str(limit); missing = False
129
    if start != None:
130
        if start != 0: query += ' OFFSET '+str(start)
131
        missing = False
132
    if missing: warnings.warn(DbWarning(
133
        'SELECT statement missing a WHERE, LIMIT, or OFFSET clause: '+query))
134
    
135
    return run_query(db, query, conds.values(), recover)
136

    
137
def insert(db, table, row, recover=None):
138
    check_name(table)
139
    cols = row.keys()
140
    map(check_name, cols)
141
    query = 'INSERT INTO '+table
142
    if row != {}: query += ' ('+', '.join(cols)+') VALUES ('\
143
        +', '.join(['%s']*len(cols))+')'
144
    else: query += ' DEFAULT VALUES'
145
    return run_query(db, query, row.values(), recover)
146

    
147
def last_insert_id(db):
148
    module = util.root_module(db)
149
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
150
    elif module == 'MySQLdb': return db.insert_id()
151
    else: return None
152

    
153
def truncate(db, table):
154
    check_name(table)
155
    return run_raw_query(db, 'TRUNCATE '+table+' CASCADE')
156

    
157
##### Database structure queries
158

    
159
def pkey(db, cache, table, recover=None):
160
    '''Assumed to be first column in table'''
161
    check_name(table)
162
    if table not in cache:
163
        cache[table] = col(run_query(db, 'SELECT * FROM '+table+' LIMIT 0',
164
            recover=recover), 0)
165
    return cache[table]
166

    
167
def index_cols(db, table, index):
168
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
169
    automatically created. When you don't know whether something is a UNIQUE
170
    constraint or a UNIQUE index, use this function.'''
171
    check_name(table)
172
    check_name(index)
173
    module = util.root_module(db)
174
    if module == 'psycopg2':
175
        return list(values(run_query(db, '''\
176
SELECT attname
177
FROM
178
(
179
        SELECT attnum, attname
180
        FROM pg_index
181
        JOIN pg_class index ON index.oid = indexrelid
182
        JOIN pg_class table_ ON table_.oid = indrelid
183
        JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
184
        WHERE
185
            table_.relname = %(table)s
186
            AND index.relname = %(index)s
187
    UNION
188
        SELECT attnum, attname
189
        FROM
190
        (
191
            SELECT
192
                indrelid
193
                , (regexp_matches(indexprs, E':varattno (\\\\d+)', 'g'))[1]::int
194
                    AS indkey
195
            FROM pg_index
196
            JOIN pg_class index ON index.oid = indexrelid
197
            JOIN pg_class table_ ON table_.oid = indrelid
198
            WHERE
199
                table_.relname = %(table)s
200
                AND index.relname = %(index)s
201
        ) s
202
        JOIN pg_attribute ON attrelid = indrelid AND attnum = indkey
203
) s
204
ORDER BY attnum
205
''',
206
            {'table': table, 'index': index})))
207
    else: raise NotImplementedError("Can't list index columns for "+module+
208
        ' database')
209

    
210
def constraint_cols(db, table, constraint):
211
    check_name(table)
212
    check_name(constraint)
213
    module = util.root_module(db)
214
    if module == 'psycopg2':
215
        return list(values(run_query(db, '''\
216
SELECT attname
217
FROM pg_constraint
218
JOIN pg_class ON pg_class.oid = conrelid
219
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
220
WHERE
221
    relname = %(table)s
222
    AND conname = %(constraint)s
223
ORDER BY attnum
224
''',
225
            {'table': table, 'constraint': constraint})))
226
    else: raise NotImplementedError("Can't list constraint columns for "+module+
227
        ' database')
228

    
229
def tables(db):
230
    module = util.root_module(db)
231
    if module == 'psycopg2':
232
        return values(run_query(db, "SELECT tablename from pg_tables "
233
            "WHERE schemaname = 'public' ORDER BY tablename"))
234
    elif module == 'MySQLdb': return values(run_query(db, 'SHOW TABLES'))
235
    else: raise NotImplementedError("Can't list tables for "+module+' database')
236

    
237
##### Database management
238

    
239
def empty_db(db):
240
    for table in tables(db): truncate(db, table)
241

    
242
##### Heuristic queries
243

    
244
def try_insert(db, table, row):
245
    '''Recovers from errors'''
246
    try: return insert(db, table, row, recover=True)
247
    except Exception, e:
248
        msg = str(e)
249
        match = re.search(r'duplicate key value violates unique constraint '
250
            r'"(([^\W_]+)_[^"]+)"', msg)
251
        if match:
252
            constraint, table = match.groups()
253
            try: cols = index_cols(db, table, constraint)
254
            except NotImplementedError: raise e
255
            else: raise DuplicateKeyException(cols, e)
256
        match = re.search(r'null value in column "(\w+)" violates not-null '
257
            'constraint', msg)
258
        if match: raise NullValueException([match.group(1)], e)
259
        raise # no specific exception raised
260

    
261
def put(db, table, row, pkey, row_ct_ref=None):
262
    '''Recovers from errors'''
263
    try:
264
        row_ct = try_insert(db, table, row).rowcount
265
        if row_ct_ref != None and row_ct >= 0: row_ct_ref[0] += row_ct
266
        return last_insert_id(db)
267
    except DuplicateKeyException, e:
268
        return value(select(db, table, [pkey], util.dict_subset(row, e.cols),
269
            recover=True))
270

    
271
def get(db, table, row, pkey, row_ct_ref=None, create=False):
272
    '''Recovers from errors'''
273
    try: return value(select(db, table, [pkey], row, 1, recover=True))
274
    except StopIteration:
275
        if not create: raise
276
        return put(db, table, row, pkey, row_ct_ref) # insert new row
277

    
278
##### Database connections
279

    
280
db_engines = {
281
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
282
    'PostgreSQL': ('psycopg2', {}),
283
}
284

    
285
DatabaseErrors_set = set([DbException])
286
DatabaseErrors = tuple(DatabaseErrors_set)
287

    
288
def _add_module(module):
289
    DatabaseErrors_set.add(module.DatabaseError)
290
    global DatabaseErrors
291
    DatabaseErrors = tuple(DatabaseErrors_set)
292

    
293
def connect(db_config, serializable=True):
294
    db_config = db_config.copy() # don't modify input!
295
    module_name, mappings = db_engines[db_config.pop('engine')]
296
    module = __import__(module_name)
297
    _add_module(module)
298
    for orig, new in mappings.iteritems():
299
        try: util.rename_key(db_config, orig, new)
300
        except KeyError: pass
301
    db = module.connect(**db_config)
302
    if serializable:
303
        run_raw_query(db, 'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
304
    return db
305

    
306
def db_config_str(db_config):
307
    return db_config['engine']+' database '+db_config['database']
(9-9/15)