Project

General

Profile

1
# Database access
2

    
3
import random
4
import re
5
import sys
6

    
7
import exc
8
import util
9

    
10
def get_cur_query(cur):
11
    if hasattr(cur, 'query'): return cur.query
12
    elif hasattr(cur, '_last_executed'): return cur._last_executed
13
    else: return None
14

    
15
def _add_cursor_info(e, cur): exc.add_msg(e, 'query: '+get_cur_query(cur))
16

    
17
class DbException(exc.ExceptionWithCause):
18
    def __init__(self, msg, cause=None, cur=None):
19
        exc.ExceptionWithCause.__init__(self, msg, cause)
20
        if cur != None: _add_cursor_info(self, cur)
21

    
22
class NameException(DbException): pass
23

    
24
class ExceptionWithColumns(DbException):
25
    def __init__(self, cols, cause=None):
26
        DbException.__init__(self, 'columns: ' + ', '.join(cols), cause)
27
        self.cols = cols
28

    
29
class DuplicateKeyException(ExceptionWithColumns): pass
30

    
31
class NullValueException(ExceptionWithColumns): pass
32

    
33
class EmptyRowException(DbException): pass
34

    
35
def check_name(name):
36
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
37
        +'" may contain only alphanumeric characters and _')
38

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

    
47
def run_query(db, query, params=None):
48
    cur = db.cursor()
49
    try: cur.execute(query, params)
50
    except Exception, e:
51
        _add_cursor_info(e, cur)
52
        raise
53
    return cur
54

    
55
def col(cur, idx): return cur.description[idx][0]
56

    
57
def rows(cur): return iter(lambda: cur.fetchone(), None)
58

    
59
def row(cur): return rows(cur).next()
60

    
61
def value(cur): return row(cur)[0]
62

    
63
def values(cur): return iter(lambda: value(cur), None)
64

    
65
def value_or_none(cur):
66
    try: return value(cur)
67
    except StopIteration: return None
68

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

    
80
def select(db, table, fields, conds, limit=None):
81
    assert limit == None or type(limit) == int
82
    check_name(table)
83
    map(check_name, fields)
84
    map(check_name, conds.keys())
85
    def cond(entry):
86
        col, value = entry
87
        cond_ = esc_name(db, col)+' '
88
        if value == None: cond_ += 'IS'
89
        else: cond_ += '='
90
        cond_ += ' %s'
91
        return cond_
92
    query = ('SELECT ' + ', '.join([esc_name(db, field) for field in fields])
93
        + ' FROM '+esc_name(db, table))
94
    if conds != {}:
95
        query += ' WHERE '+' AND '.join(map(cond, conds.iteritems()))
96
    if limit != None: query += ' LIMIT '+str(limit)
97
    return run_query(db, query, conds.values())
98

    
99
def insert(db, table, row):
100
    check_name(table)
101
    cols = row.keys()
102
    map(check_name, cols)
103
    query = 'INSERT INTO '+table
104
    if row != {}: query += ' ('+', '.join(cols)+') VALUES ('\
105
        +', '.join(['%s']*len(cols))+')'
106
    else: query += ' DEFAULT VALUES'
107
    return run_query(db, query, row.values())
108

    
109
def last_insert_id(db):
110
    module = util.root_module(db)
111
    if module == 'psycopg2': return value(run_query(db, 'SELECT lastval()'))
112
    elif module == 'MySQLdb': return db.insert_id()
113
    else: return None
114

    
115
def constraint_cols(db, table, constraint):
116
    check_name(table)
117
    check_name(constraint)
118
    module = util.root_module(db)
119
    if module == 'psycopg2':
120
        return list(values(run_query(db, '''\
121
SELECT attname
122
FROM pg_constraint
123
JOIN pg_class ON pg_class.oid = conrelid
124
JOIN pg_attribute ON attrelid = conrelid AND attnum = ANY (conkey)
125
WHERE
126
    relname = %(table)s
127
    AND conname = %(constraint)s
128
ORDER BY attnum
129
''',
130
            {'table': table, 'constraint': constraint})))
131
    else: raise NotImplementedError("Can't list constraint columns for "+module+
132
        ' database')
133

    
134
def try_insert(db, table, row):
135
    try: return with_savepoint(db, lambda: insert(db, table, row))
136
    except Exception, e:
137
        msg = str(e)
138
        match = re.search(r'duplicate key value violates unique constraint '
139
            r'"(([^\W_]+)_[^"]+)"', msg)
140
        if match:
141
            constraint, table = match.groups()
142
            try: cols = constraint_cols(db, table, constraint)
143
            except NotImplementedError: raise e
144
            else: raise DuplicateKeyException(cols[0], e)
145
        match = re.search(r'null value in column "(\w+)" violates not-null '
146
            'constraint', msg)
147
        if match: raise NullValueException([match.group(1)], e)
148
        raise # no specific exception raised
149

    
150
def pkey(db, cache, table): # Assumed to be first column in table
151
    check_name(table)
152
    if table not in cache:
153
        cache[table] = col(run_query(db, 'SELECT * FROM '+table+' LIMIT 0'), 0)
154
    return cache[table]
155

    
156
def put(db, table, row, pkey, row_ct_ref=None):
157
    try:
158
        row_ct = try_insert(db, table, row).rowcount
159
        if row_ct_ref != None and row_ct >= 0: row_ct_ref[0] += row_ct
160
        return last_insert_id(db)
161
    except DuplicateKeyException, e:
162
        return value(select(db, table, [pkey], util.dict_subset(row, e.cols)))
163

    
164
def get(db, table, row, pkey, row_ct_ref=None, create=False):
165
    try: return value(select(db, table, [pkey], row, 1))
166
    except StopIteration:
167
        if not create: raise
168
        return put(db, table, row, pkey, row_ct_ref) # insert new row
169

    
170

    
171
def truncate(db, table):
172
    check_name(table)
173
    return run_query(db, 'TRUNCATE '+table+' CASCADE')
174

    
175
def tables(db):
176
    module = util.root_module(db)
177
    if module == 'psycopg2':
178
        return values(run_query(db, "SELECT tablename from pg_tables "
179
            "WHERE schemaname = 'public' ORDER BY tablename"))
180
    elif module == 'MySQLdb': return values(run_query(db, 'SHOW TABLES'))
181
    else: raise NotImplementedError("Can't list tables for "+module+' database')
182

    
183
def empty_db(db):
184
    for table in tables(db): truncate(db, table)
185

    
186
db_engines = {
187
    'MySQL': ('MySQLdb', {'password': 'passwd', 'database': 'db'}),
188
    'PostgreSQL': ('psycopg2', {}),
189
}
190

    
191
DatabaseErrors_set = set([DbException])
192
DatabaseErrors = tuple(DatabaseErrors_set)
193

    
194
def _add_module(module):
195
    DatabaseErrors_set.add(module.DatabaseError)
196
    global DatabaseErrors
197
    DatabaseErrors = tuple(DatabaseErrors_set)
198

    
199
def connect(db_config):
200
    db_config = db_config.copy() # don't modify input!
201
    module_name, mappings = db_engines[db_config.pop('engine')]
202
    module = __import__(module_name)
203
    _add_module(module)
204
    for orig, new in mappings.iteritems():
205
        try: util.rename_key(db_config, orig, new)
206
        except KeyError: pass
207
    return module.connect(**db_config)
(7-7/13)