libdb.py.in 2.15 KB
Newer Older
1
2
3
4
5
6
7
# -*- python -*-

#
# Database functions
#

import sys
8
sys.path.append("@prefix@/lib")
9
10
11
12
13
import os
import pwd

import MySQLdb

Kirk Webb's avatar
Kirk Webb committed
14
import libtestbed
15
from libtestbed import *
16
17
18
19

__dbName = "@TBDBNAME@"
__dbQueryMaxtries = 1
__dbConnMaxtries = 5
20
__debug = 0
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37

__dbConnection = None

def TBDBConnect():
    global __dbConnection
    
    if __dbConnection:
        return

    # Create a DB username for accounting purposes
    uid = os.getuid()
    try:
        name = pwd.getpwuid(uid)[0]
    except KeyError:
        name = "uid%d" % uid
    dbuser = "%s:%s:%d" % (sys.argv[0], name, os.getpid())

38
    if __debug:
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
        print "Connecting to db %s as %s" % (__dbName, dbuser)

    # Connect, with retries
    for tries in range(__dbConnMaxtries):
        try:
            __dbConnection = MySQLdb.connect(db = __dbName, user = dbuser)
        except:
            time.sleep(1)
        else:
            break
    else:
        raise RuntimeError, "Cannot connect to DB after several attempts!"

def DBQueryFatal(queryPat, querySub = (), asDict = False):
    TBDBConnect()
    
    if asDict:
        cursor = __dbConnection.cursor(MySQLdb.cursors.DictCursor)
    else:
        cursor = __dbConnection.cursor()

60
    if __debug:
61
62
63
64
65
66
67
        print "Executing DB query %s" % queryPat

    tries = __dbQueryMaxtries
    while tries:
        try:
            cursor.execute(queryPat, querySub)
            ret = cursor.fetchall()
68
            if __debug:
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
                rs = `ret`
                if len(rs) > 60:
                    rs = rs[:60] + "..."
                print "Result: %s" % rs
            return ret
        except MySQLdb.MySQLError:
            tries -= 1
            if tries == 0:
                raise
            else:
                time.sleep(1)
                try:
                    __dbConnection.ping()
                except MySQLdb.MySQLError: pass
    tbmsg = "".join(traceback.format_exception(*sys.exc_info()))
    SENDMAIL(TBOPS, "DB query failed", "DB query failed:\n\n%s" % tbmsg, TBOPS)
    raise RuntimeError, "Aah!  Escaped DBQueryFatal loop"

def DBQuery(*args):
    try:
        ret = DBQueryFatal(*args)
90
        return ret
91
92
93
    except MySQLdb.MySQLError:
        return None