Project

General

Profile

1
# Database access
2

    
3
import random
4
import re
5
import sys
6

    
7
import exc
8
import strings
9
import util
10

    
11
##### Exceptions
12

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

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

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

    
25
class NameException(DbException): pass
26

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

    
32
class DuplicateKeyException(ExceptionWithColumns): pass
33

    
34
class NullValueException(ExceptionWithColumns): pass
35

    
36
class EmptyRowException(DbException): pass
37

    
38
##### Input validation
39

    
40
def check_name(name):
41
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
42
        +'" may contain only alphanumeric characters and _')
43

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

    
52
##### Querying
53

    
54
def run_raw_query(db, query, params=None):
55
    if run_raw_query.debug: sys.stderr.write(strings.one_line(query)+'\n')
56
    cur = db.cursor()
57
    try: cur.execute(query, params)
58
    except Exception, e:
59
        _add_cursor_info(e, cur)
60
        raise
61
    return cur
62

    
63
##### Recoverable querying
64

    
65
def with_savepoint(db, func):
66
    savepoint = 'savepoint_'+str(random.randint(0, sys.maxint)) # must be unique
67
    run_raw_query(db, 'SAVEPOINT '+savepoint)
68
    try: return_val = func()
69
    except:
70
        run_raw_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)
71
        raise
72
    else:
73
        run_raw_query(db, 'RELEASE SAVEPOINT '+savepoint)
74
        return return_val
75

    
76
def run_query(db, query, params=None, recover=None):
77
    if recover == None: recover = False
78
    
79
    def run(): return run_raw_query(db, query, params)
80
    if recover: return with_savepoint(db, run)
81
    else: return run()
82

    
83
##### Result retrieval
84

    
85
def col(cur, idx): return cur.description[idx][0]
86

    
87
def rows(cur): return iter(lambda: cur.fetchone(), None)
88

    
89
def row(cur): return rows(cur).next()
90

    
91
def value(cur): return row(cur)[0]
92

    
93
def values(cur): return iter(lambda: value(cur), None)
94

    
95
def value_or_none(cur):
96
    try: return value(cur)
97
    except StopIteration: return None
98

    
99
##### Basic queries
100

    
101
def select(db, table, fields, conds, limit=None, start=None, recover=None):
102
    assert limit == None or type(limit) == int
103
    check_name(table)
104
    map(check_name, fields)
105
    map(check_name, conds.keys())
106
    def cond(entry):
107
        col, value = entry
108
        cond_ = esc_name(db, col)+' '
109
        if value == None: cond_ += 'IS'
110
        else: cond_ += '='
111
        cond_ += ' %s'
112
        return cond_
113
    query = ('SELECT ' + ', '.join([esc_name(db, field) for field in fields])
114
        + ' FROM '+esc_name(db, table))
115
    if conds != {}:
116
        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))
117
    if limit != None: query += ' LIMIT '+str(limit)
118
    if start != None: query += ' OFFSET '+str(start)
119
    return run_query(db, query, conds.values(), recover)
120

    
121
def insert(db, table, row, recover=None):
122
    check_name(table)
123
    cols = row.keys()
124
    map(check_name, cols)
125
    query = 'INSERT INTO '+table
126
    if row != {}: query += ' ('+', '.join(cols)+') VALUES ('\
127
        +', '.join(['%s']*len(cols))+')'
128
    else: query += ' DEFAULT VALUES'
129
    return run_query(db, query, row.values(), recover)
130

    
131
def last_insert_id(db):
132
    module = util.root_module(db)
133
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
134
    elif module == 'MySQLdb': return db.insert_id()
135
    else: return None
136

    
137
def truncate(db, table):
138
    check_name(table)
139
    return run_query(db, 'TRUNCATE '+table+' CASCADE')
140

    
141
##### Database structure queries
142

    
143
def pkey(db, cache, table, recover=None):
144
    '''Assumed to be first column in table'''
145
    check_name(table)
146
    if table not in cache:
147
        cache[table] = col(run_query(db, 'SELECT * FROM '+table+' LIMIT 0',
148
            recover=recover), 0)
149
    return cache[table]
150

    
151
def index_cols(db, table, index):
152
    '''Can also use this for UNIQUE constraints, because a UNIQUE index is
153
    automatically created. When you don't know whether something is a UNIQUE
154
    constraint or a UNIQUE index, use this function.'''
155
    check_name(table)
156
    check_name(index)
157
    module = util.root_module(db)
158
    if module == 'psycopg2':
159
        return list(values(run_query(db, '''\
160
SELECT attname
161
FROM pg_index
162
JOIN pg_class index ON index.oid = indexrelid
163
JOIN pg_class table_ ON table_.oid = indrelid
164
JOIN pg_attribute ON attrelid = indrelid AND attnum = ANY (indkey)
165
WHERE
166
    table_.relname = %(table)s
167
    AND index.relname = %(index)s
168
ORDER BY attnum
169
''',
170
            {'table': table, 'index': index})))
171
    else: raise NotImplementedError("Can't list index columns for "+module+
172
        ' database')
173

    
174
def constraint_cols(db, table, constraint):
175
    check_name(table)
176
    check_name(constraint)
177
    module = util.root_module(db)
178
    if module == 'psycopg2':
179
        return list(values(run_query(db, '''\
180
SELECT attname
181
FROM pg_constraint
182
JOIN pg_class ON pg_class.oid = conrelid
183
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
184
WHERE
185
    relname = %(table)s
186
    AND conname = %(constraint)s
187
ORDER BY attnum
188
''',
189
            {'table': table, 'constraint': constraint})))
190
    else: raise NotImplementedError("Can't list constraint columns for "+module+
191
        ' database')
192

    
193
def tables(db):
194
    module = util.root_module(db)
195
    if module == 'psycopg2':
196
        return values(run_query(db, "SELECT tablename from pg_tables "
197
            "WHERE schemaname = 'public' ORDER BY tablename"))
198
    elif module == 'MySQLdb': return values(run_query(db, 'SHOW TABLES'))
199
    else: raise NotImplementedError("Can't list tables for "+module+' database')
200

    
201
##### Database management
202

    
203
def empty_db(db):
204
    for table in tables(db): truncate(db, table)
205

    
206
##### Heuristic queries
207

    
208
def try_insert(db, table, row):
209
    '''Recovers from errors'''
210
    try: return insert(db, table, row, recover=True)
211
    except Exception, e:
212
        msg = str(e)
213
        match = re.search(r'duplicate key value violates unique constraint '
214
            r'"(([^\W_]+)_[^"]+)"', msg)
215
        if match:
216
            constraint, table = match.groups()
217
            try: cols = index_cols(db, table, constraint)
218
            except NotImplementedError: raise e
219
            else: raise DuplicateKeyException(cols, e)
220
        match = re.search(r'null value in column "(\w+)" violates not-null '
221
            'constraint', msg)
222
        if match: raise NullValueException([match.group(1)], e)
223
        raise # no specific exception raised
224

    
225
def put(db, table, row, pkey, row_ct_ref=None):
226
    '''Recovers from errors'''
227
    try:
228
        row_ct = try_insert(db, table, row).rowcount
229
        if row_ct_ref != None and row_ct >= 0: row_ct_ref[0] += row_ct
230
        return last_insert_id(db)
231
    except DuplicateKeyException, e:
232
        return value(select(db, table, [pkey], util.dict_subset(row, e.cols),
233
            recover=True))
234

    
235
def get(db, table, row, pkey, row_ct_ref=None, create=False):
236
    '''Recovers from errors'''
237
    try: return value(select(db, table, [pkey], row, 1, recover=True))
238
    except StopIteration:
239
        if not create: raise
240
        return put(db, table, row, pkey, row_ct_ref) # insert new row
241

    
242
##### Database connections
243

    
244
db_engines = {
245
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
246
    'PostgreSQL': ('psycopg2', {}),
247
}
248

    
249
DatabaseErrors_set = set([DbException])
250
DatabaseErrors = tuple(DatabaseErrors_set)
251

    
252
def _add_module(module):
253
    DatabaseErrors_set.add(module.DatabaseError)
254
    global DatabaseErrors
255
    DatabaseErrors = tuple(DatabaseErrors_set)
256

    
257
def connect(db_config, serializable=True):
258
    db_config = db_config.copy() # don't modify input!
259
    module_name, mappings = db_engines[db_config.pop('engine')]
260
    module = __import__(module_name)
261
    _add_module(module)
262
    for orig, new in mappings.iteritems():
263
        try: util.rename_key(db_config, orig, new)
264
        except KeyError: pass
265
    db = module.connect(**db_config)
266
    if serializable:
267
        run_raw_query(db, 'SET TRANSACTION ISOLATION LEVEL SERIALIZABLE')
268
    return db
269

    
270
def db_config_str(db_config):
271
    return db_config['engine']+' database '+db_config['database']
(8-8/14)