aboutsummaryrefslogtreecommitdiff
path: root/lib/db.py
blob: 4c99795ab3de062b6fa1ab95045e949d61fda8ca (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
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()