1 files changed, 64 insertions, 0 deletions
diff --git a/lib/db.py b/lib/db.py
new file mode 100644
index 0000000..4c99795
--- /dev/null
+++ b/lib/db.py
@@ -0,0 +1,64 @@
+'''
+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 cash_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()