import setuptools import sqlite3, pymysql, pymysql.cursors SQLITE = 0 MYSQL = 1 MARIADB = MYSQL MIN = 0 MAX = 1 class dbObject: # --------------------------------------------- Initialization ------------------------------------------------- def initMySQL(self, host, port, user, pwd, db): self.conn = pymysql.connect(host = host, port = port, user = user, password = pwd, db = db, charset = "utf8mb4", cursorclass = pymysql.cursors.DictCursor) self.cur = self.conn.cursor() self.dbtype = MYSQL self.host = host self.port = port self.user = user self.pwd = pwd self.db = db def initSQLite(self, path): self.conn = sqlite3.connect(path) self.cur = self.conn.cursor() self.dbtype = SQLITE self.path = path def __init__(self, dbtype = SQLITE, path = None, host = None, port = None, user = None, pwd = None, db = None): if dbtype == SQLITE: self.initSQLite(path or 'Database.db') elif dbtype == MYSQL: self.initMySQL(host or 'localhost', port or 3306, user, pwd, db) else: raise ValueError("Unknown database type %s." % str(dbtype)) # ---------------------------------------------- No more initialization ---------------------------------------- def closeConnection(self): return self.conn.close() def commit(self): return self.conn.commit() def executeQuery(self, query): return self.cur.execute(query) def getAll(self): return self.cur.fetchall() def getNext(self): return self.cur.fetchone() def isInitialized(self): try: self.executeQuery("SELECT * FROM tokens") return True except: return False def deleteUser(self, cid): self.executeQuery("DELETE FROM tokens WHERE cid = %i;" % int(cid)) self.commit() def storeUser(self, cid, ato, ase): self.executeQuery("DELETE FROM tokens WHERE cid = %i;" % int(cid)) self.executeQuery("INSERT INTO tokens(cid, ato, ase) VALUES(%i, '%s', '%s');" % (int(cid), ato, ase)) self.commit() def ato(self, cid): try: self.executeQuery("SELECT ato FROM tokens WHERE cid = %i;" % int(cid)) return self.cur.fetchone()[0] except: return False def ase(self, cid): try: self.executeQuery("SELECT ase FROM tokens WHERE cid = %i;" % int(cid)) return self.cur.fetchone()[0] except: return False def getTStatus(self, cid): try: self.executeQuery("SELECT tweet FROM tokens WHERE cid = %i;" % int(cid)) return True if self.cur.fetchone()[0] == 1 else False except: raise ValueError("No such user: %i" % int(cid)) def toggleTweet(self, cid): self.executeQuery("UPDATE tokens SET tweet = NOT tweet WHERE cid = %i;" % int(cid)) self.commit() return getTStatus def addFish(self, cid): self.executeQuery("UPDATE tokens SET fish = fish + 1 WHERE cid = %i;" % int(cid)) self.commit() def storeToken(self, cid, ato): self.executeQuery('INSERT INTO tokens(cid, ato) VALUES(%i, "%s");' % (int(cid), ato)) self.commit() def dbHelper(): if setuptools.dbtype() == SQLITE: return dbObject(dbtype=SQLITE, path=setuptools.dbpath()) elif setuptools.dbtype() == MYSQL: return dbObject(dbtype=MYSQL, host=setuptools.dbhost(), user=setuptools.dbuser(), pwd=setuptools.dbpass(), db=setuptools.dbname()) else: raise setuptools.SetupException()