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
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()
|
|
|