Project

General

Profile

1
# Database access
2

    
3
import random
4
import re
5
import sys
6

    
7
import ex_util
8

    
9
def _add_cursor_info(ex, cur): ex_util.add_msg(ex, 'query: '+cur.query)
10

    
11
class NameException(Exception): pass
12

    
13
class DbException(ex_util.ExceptionWithCause):
14
    def __init__(self, msg, cause=None, cur=None):
15
        ex_util.ExceptionWithCause.__init__(self, msg, cause)
16
        if cur != None: _add_cursor_info(self, cur)
17

    
18
class ExceptionWithColumn(DbException):
19
    def __init__(self, col, cause=None):
20
        DbException.__init__(self, 'column: '+col, cause)
21
        self.col = col
22

    
23
class DuplicateKeyException(ExceptionWithColumn): pass
24

    
25
class NullValueException(ExceptionWithColumn): pass
26

    
27
def check_name(name):
28
    if re.search(r'\W', name) != None: raise NameException('Name "'+name
29
        +'" may contain only alphanumeric characters and _')
30

    
31
def run_query(db, query, params=None):
32
    cur = db.cursor()
33
    try: cur.execute(query, params)
34
    except Exception, ex:
35
        _add_cursor_info(ex, cur)
36
        raise
37
    return cur
38

    
39
def row(cur): return iter(lambda: cur.fetchone(), None).next()
40

    
41
def value(cur): return row(cur)[0]
42

    
43
def with_savepoint(db, func):
44
    savepoint = 'savepoint_'+str(random.randint(0, sys.maxint)) # must be unique
45
    run_query(db, 'SAVEPOINT '+savepoint)
46
    try: return_val = func()
47
    except:
48
        run_query(db, 'ROLLBACK TO SAVEPOINT '+savepoint)
49
        raise
50
    else:
51
        run_query(db, 'RELEASE SAVEPOINT '+savepoint)
52
        return return_val
53

    
54
def select(db, table, fields, conds):
55
    for field in fields: check_name(field)
56
    for col in conds.keys(): check_name(col)
57
    def cond(entry):
58
        col, value = entry
59
        cond_ = col+' '
60
        if value == None: cond_ += 'IS'
61
        else: cond_ += '='
62
        cond_ += ' %s'
63
        return cond_
64
    return run_query(db, 'SELECT '+', '.join(fields)+' FROM '+table+' WHERE '
65
        +' AND '.join(map(cond, conds.iteritems())), conds.values())
66

    
67
def insert(db, table, row):
68
    check_name(table)
69
    cols = row.keys()
70
    for col in cols: check_name(col)
71
    return run_query(db, 'INSERT INTO '+table+' ('+', '.join(cols)
72
        +') VALUES ('+', '.join(['%s']*len(cols))+')', row.values())
73

    
74
def last_insert_id(db): return value(run_query(db, 'SELECT lastval()'))
75

    
76
def try_insert(db, table, row):
77
    try: return with_savepoint(db, lambda: insert(db, table, row))
78
    except Exception, ex:
79
        msg = str(ex)
80
        match = re.search(r'duplicate key value violates unique constraint "'
81
            +table+'_(\w+)_index"', msg)
82
        if match: raise DuplicateKeyException(match.group(1), ex)
83
        match = re.search(r'null value in column "(\w+)" violates not-null '
84
            'constraint', msg)
85
        if match: raise NullValueException(match.group(1), ex)
86
        raise # no specific exception raised
87

    
88
def insert_or_get(db, table, row, pkey, row_ct_ref=None):
89
    try: return value(select(db, table, [pkey], row))
90
    except StopIteration:
91
        try:
92
            row_ct = try_insert(db, table, row).rowcount
93
            if row_ct_ref != None and row_ct >= 0: row_ct_ref[0] += row_ct
94
            return last_insert_id(db)
95
        except DuplicateKeyException, ex:
96
            return value(select(db, table, [pkey], {ex.col: row[ex.col]}))
(1-1/8)