Awesome IRC bot
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

114 lines
3.5 KiB

#!/usr/bin/env python
"""
.. module:: MySQL
:synopsis: Module providing a mysql type service
.. moduleauthor:: Dave Pedu <dave@davepedu.com>
"""
from pyircbot.modulebase import ModuleBase
import sys
import pymysql as MySQLdb # python 3.x
class MySQL(ModuleBase):
def __init__(self, bot, moduleName):
ModuleBase.__init__(self, bot, moduleName)
self.services = ["mysql"]
self.connection = self.getConnection()
def getConnection(self):
return Connection(self)
class Connection:
def __init__(self, master):
self.config = master.config
self.log = master.log
self._connect()
# Check if the table requested exists
def tableExists(self, tablename):
c = self.getCursor()
c.execute("SHOW TABLES;")
tables = c.fetchall()
if len(tables) == 0:
return False
key = list(tables[0].keys())[0]
for table in tables:
if table[key] == tablename:
return True
return False
def query(self, queryText, args=()):
"""Execute a MySQL query and return the cursor
:param queryText: the mysql query as a string, using '%s' for token replacement
:type queryText: str
:param args: arguments to be escaped into the query
:type args: tuple
:returns: cursor -- the sql cursor"""
c = self.getCursor()
if len(args) == 0:
c.execute(queryText)
else:
c.execute(queryText, args)
return c
# Returns a cusor object, after checking for connectivity
def getCursor(self):
self.ensureConnected()
if sys.version_info > (3, 0):
c = self.connection.cursor(MySQLdb.cursors.DictCursor)
else:
c = self.connection.cursor(cursorclass=MySQLdb.cursors.DictCursor)
c.execute("USE `%s`;" % self.config["database"])
return c
def escape(self, s):
"""Escape a string using the mysql server
:param s: the string to escape
:type s: str
:returns: str -- the escaped string"""
self.ensureConnected()
return self.connection.escape_string(s)
def ensureConnected(self):
try:
self.connection.ping()
except:
try:
self.connection.close()
except:
pass
del self.connection
self._connect()
def ondisable(self):
self.connection.close()
# Connects to the database server, and selects a database (Or attempts to create it if it doesn't exist yet)
def _connect(self):
self.log.info("MySQL: Connecting to db host at %s" % self.config["host"])
self.connection = MySQLdb.connect(host=self.config["host"], user=self.config["username"],
passwd=self.config["password"])
self.log.info("MySQL: Connected.")
self.connection.autocommit(True)
c = None
if sys.version_info > (3, 0):
c = self.connection.cursor(MySQLdb.cursors.DictCursor)
else:
c = self.connection.cursor(cursorclass=MySQLdb.cursors.DictCursor)
c.execute("SHOW DATABASES")
dblist = c.fetchall()
found = False
for row in dblist:
if row["Database"] == self.config["database"]:
found = True
if not found:
c.execute("CREATE DATABASE `%s`;" % self.config["database"])
c.execute("USE `%s`;" % self.config["database"])
c.close()