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 tweets") return True except: return False def getFLDate(self, val = MIN): if val == MIN: mode = "MIN" else: mode = "MAX" if self.dbtype == SQLITE: return setuptools.getDate(str(list(self.executeQuery("SELECT %s(SUBSTR(timestamp,0,11)) FROM tweets" % mode))[0][0])) else: self.executeQuery("SELECT %s(SUBSTR(timestamp,0,11)) FROM tweets" % mode) return setuptools.getDate(str(self.getNext()["%s(SUBSTR(timestamp,0,11))" % mode])) def getFollowers(db): db.executeQuery("SELECT id FROM followers WHERE `until` IS NULL;") for i in db.getAll(): yield i[0] def getFollowing(db): db.executeQuery("SELECT id FROM following WHERE `until` IS NULL;") for i in db.getAll(): yield i[0] def getLatestMessage(db): db.executeQuery("SELECT max(id) FROM messages") try: return int(db.getNext()[0]) except: return 0 def getLatestTweet(db): db.executeQuery("SELECT max(tweet_id) FROM tweets") try: return int(db.getNext()[0]) except: return 0 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()