Source Code for Module gozerbot.db

  1  # gozerbot/db.py 
  2  # 
  3  # 
  4   
  5  """ mysql interface """ 
  6   
  7  __copyright__ = 'this file is in the public domain' 
  8   
  9  from gozerbot.config import config 
 10  from gozerbot.generic import rlog, lockdec, tolatin1, handle_exception 
 11  from gozerbot.datadir import datadir 
 12  import thread, os, time, types 
 13   
 14  dblock = thread.allocate_lock() 
 15  dblocked = lockdec(dblock) 
 16   
 17 -class Db(object): 
 18   
 19      """ this class implements an database connection. it connects to the  
 20          database on initialisation. 
 21      """ 
 22   
 23 -    def __init__(self, doconnect=True, dbtype=None): 
 24          self.dbname = "" 
 25          self.dbhost = "" 
 26          self.dbuser = "" 
 27          self.dbpasswd = "" 
 28          self.connection = None 
 29          self.timeout = 15 
 30          self.oldstyle = "" 
 31          self.dbtype = dbtype or config['dbtype'] or 'mysql' 
 32          if doconnect: 
 33              self.connect() 
 34   
 35      @dblocked 
 36 -    def connect(self, dbname=None, dbhost=None, dbuser=None, dbpasswd=None, \ 
 37  timeout=15, oldstyle=False): 
 38          """ connect to the database """ 
 39          self.dbname = dbname or config['dbname'] 
 40          self.dbhost = dbhost or config['dbhost'] 
 41          self.dbuser = dbuser or config['dbuser'] 
 42          self.dbpasswd = dbpasswd or config['dbpasswd'] 
 43          self.timeout = timeout 
 44          self.oldstyle = oldstyle or config['dboldstyle'] 
 45          if self.dbtype == 'mysql': 
 46              import MySQLdb 
 47              self.connection = MySQLdb.connect(db=self.dbname, \ 
 48  host=self.dbhost, user=self.dbuser, passwd=self.dbpasswd, \ 
 49  connect_timeout=self.timeout, charset='utf8') 
 50          elif self.dbtype == 'sqlite': 
 51              try: 
 52                  import sqlite 
 53                  self.connection = sqlite.connect(datadir + os.sep + \ 
 54  self.dbname) 
 55              except ImportError: 
 56                  import sqlite3 
 57                  self.connection = sqlite3.connect(datadir + os.sep + \ 
 58  self.dbname, check_same_thread=False) 
 59          elif self.dbtype == 'postgres': 
 60              import psycopg2 
 61              rlog(1000, 'db', 'NOTE THAT POSTGRES IS NOT FULLY SUPPORTED') 
 62              self.connection = psycopg2.connect(database=self.dbname, \ 
 63  host=self.dbhost, user=self.dbuser, password=self.dbpasswd) 
 64          else: 
 65              rlog(100, 'db', 'unknown database type %s' % self.dbtype) 
 66              return 0 
 67          rlog(5, 'db', "database ok") 
 68          return 1 
 69   
 70 -    def reconnect(self): 
 71          """ reconnect to the mysql server """ 
 72          if self.dbtype == 'mysql': 
 73              import MySQLdb 
 74              self.connection = MySQLdb.connect(db=self.dbname, \ 
 75  host=self.dbhost, user=self.dbuser, passwd=self.dbpasswd, \ 
 76  connect_timeout=self.timeout, charset='utf8') 
 77          elif self.dbtype == 'sqlite': 
 78              try: 
 79                  import sqlite 
 80                  self.connection = sqlite.connect(datadir + os.sep + \ 
 81  self.dbname) 
 82              except ImportError: 
 83                  import sqlite3 
 84                  self.connection = sqlite3.connect(datadir + os.sep + \ 
 85  self.dbname, check_same_thread=False) 
 86          elif self.dbtype == 'postgres': 
 87              import psycopg2 
 88              rlog(1000, 'db', 'NOTE THAT POSTGRES IS NOT FULLY SUPPORTED') 
 89              self.connection = psycopg2.connect(database=self.dbname, \ 
 90  host=self.dbhost, user=self.dbuser, password=self.dbpasswd) 
 91          else: 
 92              rlog(100, 'db', 'unknown database type %s' % self.dbtype) 
 93              return 0 
 94          rlog(5, 'db', "database ok") 
 95          return 1 
 96   
 97      @dblocked 
 98 -    def execute(self, execstr, args=None): 
 99          """ execute string on database """ 
100          time.sleep(0.001) 
101          result = None 
102          execstr = execstr.strip() 
103          if self.dbtype == 'sqlite': 
104              execstr = execstr.replace('%s', '?') 
105          # first to ping to see if connection is alive .. if not reconnect 
106          if self.dbtype == 'mysql': 
107              try: 
108                  self.ping() 
109              except Exception, ex: 
110                  rlog(10, 'db', "can't ping database: %s" % str(ex)) 
111                  rlog(10, 'db', 'reconnecting') 
112                  try: 
113                      self.reconnect() 
114                  except Exception, ex: 
115                      rlog(10, 'db', 'failed reconnect: %s' % str(ex)) 
116                      return 
117          # get cursor 
118          cursor = self.cursor() 
119          # excecute string on cursor 
120          nr = 0 
121          if args and self.oldstyle: 
122              nargs = [] 
123              for i in args: 
124                  nargs.append(tolatin1(i)) 
125              args = nargs 
126          rlog(-2, 'db', 'exec %s %s' % (execstr, args)) 
127          try: 
128              if args: 
129                  if type(args) == tuple or type(args) == list: 
130                      nr = cursor.execute(execstr, args) 
131                  else: 
132                      nr = cursor.execute(execstr, (args, )) 
133              else: 
134                  nr = cursor.execute(execstr) 
135          except: 
136              if self.dbtype == 'postgres': 
137                  cursor.execute(""" ROLLBACK """) 
138              if self.dbtype == 'sqlite': 
139                  cursor.close() 
140              raise 
141          # see if we need to commit the query 
142          got = False 
143          if execstr.startswith('INSERT'): 
144              nr = cursor.lastrowid or nr 
145              got = True 
146          elif execstr.startswith('UPDATE'): 
147              nr = cursor.rowcount 
148              got = True 
149          elif execstr.startswith('DELETE'): 
150              nr = cursor.rowcount 
151              got = True 
152          if got: 
153              self.commit() 
154          if self.dbtype == 'sqlite' and not got and type(nr) != types.IntType: 
155              nr = cursor.rowcount or cursor.lastrowid 
156              if nr == -1: 
157                  nr = 0 
158          # fetch results 
159          result = None 
160          try: 
161              result = cursor.fetchall() 
162              if not result: 
163                  result = nr 
164          except Exception, ex: 
165              if 'no results to fetch' in str(ex): 
166                  pass 
167              else: 
168                  handle_exception() 
169              result = nr 
170          cursor.close() 
171          return result 
172   
173 -    def cursor(self): 
174          """ return cursor to the database """ 
175          return self.connection.cursor() 
176   
177 -    def commit(self): 
178          """ do a commit on the datase """ 
179          self.connection.commit() 
180   
181 -    def ping(self): 
182          """ do a ping """ 
183          return self.connection.ping() 
184   
185 -    def close(self): 
186          """ close database """ 
187          self.connection.close() 
188   
189  # create default database if enabled 
190  cfg = config['dbenable'] 
191  if cfg: 
192      db = Db() 
193  else: 
194      db = None 
195