tweepbot/dbtools/__init__.py
2017-02-08 15:26:21 +01:00

113 lines
2.9 KiB
Python

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