''' File: db.py Library to facilitate database IO ''' import os import sqlite3 ABS_PATH = os.path.dirname(os.path.abspath(__file__)) class db_base: def __init__(self): ''' Initialize conn/cursor for in-class use ''' self.db_path = os.path.join(ABS_PATH, '../cache/cache.db') self.cache_path = os.path.join(ABS_PATH, '../cache') self.conn = self.db_init() self.cursor = self.conn.cursor() def db_init(self): ''' Initialize database schema if db not found, return conn ''' if not os.path.isdir(self.cache_path): os.mkdir(self.cache_path) if not os.path.isfile(self.db_path): conn = sqlite3.connect(self.db_path) print('database not found, initializing...') conn.execute('''CREATE TABLE cache (hash TEXT, content_type TEXT)''') conn.commit() conn.close() conn = sqlite3.connect(self.db_path) conn.row_factory = sqlite3.Row return conn def cache_add(self, _hash, content_type): self.cursor.execute('''INSERT INTO cache(hash, content_type) VALUES(?,?)''', (_hash, content_type,)) self.save() def cache_del(self, _hash): self.cursor.execute('DELETE FROM cache WHERE hash=?', (_hash,)) self.save() def is_cached(self, _hash): self.cursor.execute('''SELECT COUNT(*) FROM cache WHERE hash=?''', (_hash,)) q_count = self.cursor.fetchone() if q_count[0] > 0: return True return False def get_content_type(self, _hash): self.cursor.execute('SELECT * FROM cache WHERE hash=?', (_hash,)) return self.cursor.fetchall()[0][1] def save(self): self.conn.commit() def close(self): self.conn.close()