128 lines
4.3 KiB
Python
128 lines
4.3 KiB
Python
import json
|
|
import sqlite3
|
|
import os
|
|
from hashlib import md5
|
|
|
|
class SQLite():
|
|
def __init__(self, db_name: str, schema:dict = None) -> None:
|
|
self.conn: sqlite3.Connection|None = None
|
|
self.cur = None
|
|
self.db_name = db_name
|
|
self.__backup = None
|
|
self.schema = {} if schema is None else schema
|
|
self.DB_VER = md5(json.dumps(self.schema).encode()).hexdigest()
|
|
|
|
def connect(self) -> None:
|
|
try:
|
|
self.conn = sqlite3.connect(self.db_name, check_same_thread = False)
|
|
except Exception:
|
|
path_name = os.path.dirname(self.db_name)
|
|
if not os.path.exists(path_name):
|
|
os.makedirs(path_name)
|
|
with open(self.db_name, 'x'):
|
|
pass
|
|
self.conn = sqlite3.connect(self.db_name, check_same_thread = False)
|
|
def dict_factory(cursor, row):
|
|
d = {}
|
|
for idx, col in enumerate(cursor.description):
|
|
d[col[0]] = row[idx]
|
|
return d
|
|
self.conn.row_factory = dict_factory
|
|
|
|
def backup(self):
|
|
with open(self.db_name, "rb") as f:
|
|
self.__backup = f.read()
|
|
|
|
def restore(self):
|
|
were_active = True if self.conn else False
|
|
self.disconnect()
|
|
if self.__backup:
|
|
with open(self.db_name, "wb") as f:
|
|
f.write(self.__backup)
|
|
self.__backup = None
|
|
if were_active:
|
|
self.connect()
|
|
|
|
def delete_backup(self):
|
|
self.__backup = None
|
|
|
|
def disconnect(self) -> None:
|
|
if self.conn:
|
|
self.conn.close()
|
|
self.conn = None
|
|
|
|
def create_schema(self, tables = {}) -> None:
|
|
if self.conn:
|
|
cur = self.conn.cursor()
|
|
cur.execute("CREATE TABLE IF NOT EXISTS main.keys_values(key VARCHAR(100) PRIMARY KEY, value VARCHAR(100) NOT NULL);")
|
|
for t in tables:
|
|
if t == "QUERY":
|
|
continue
|
|
cur.execute('CREATE TABLE IF NOT EXISTS main.{}({});'.format(t, ''.join([(c + ' ' + tables[t][c] + ', ') for c in tables[t]])[:-2]))
|
|
if "QUERY" in tables:
|
|
[cur.execute(qry) for qry in tables["QUERY"]]
|
|
cur.close()
|
|
|
|
def query(self, query, *values):
|
|
return self.queries([(query, *values)])[0]
|
|
|
|
def queries(self, queries: list[tuple[str, ...]]):
|
|
return list(self.queries_iter(queries))
|
|
|
|
def queries_iter(self, queries: list[tuple[str, ...]]):
|
|
cur = self.conn.cursor()
|
|
try:
|
|
cur.execute("BEGIN")
|
|
for query_data in queries:
|
|
values = []
|
|
str_query = None
|
|
if isinstance(query_data, str):
|
|
str_query = query_data
|
|
elif (isinstance(query_data, tuple) or isinstance(query_data, list)) and len(query_data) > 0 and isinstance(query_data[0], str):
|
|
str_query = query_data[0]
|
|
values = query_data[1:]
|
|
if str_query:
|
|
cur.execute(str_query, values)
|
|
yield cur.fetchall()
|
|
cur.execute("COMMIT")
|
|
except Exception as e:
|
|
cur.execute("ROLLBACK")
|
|
raise e
|
|
finally:
|
|
cur.close()
|
|
try:
|
|
self.conn.commit()
|
|
except Exception:
|
|
pass
|
|
|
|
def delete(self):
|
|
self.disconnect()
|
|
os.remove(self.db_name)
|
|
|
|
def init(self):
|
|
self.connect()
|
|
try:
|
|
if self.get('DB_VERSION') != self.DB_VER:
|
|
raise Exception("DB_VERSION is not correct")
|
|
except Exception:
|
|
self.delete()
|
|
self.connect()
|
|
self.create_schema(self.schema)
|
|
self.put('DB_VERSION', self.DB_VER)
|
|
|
|
def get(self, key, default = None):
|
|
q = self.query('SELECT value FROM keys_values WHERE key = ?', key)
|
|
if len(q) == 0:
|
|
return default
|
|
else:
|
|
return q[0]["value"]
|
|
|
|
def put(self, key, value):
|
|
if self.get(key) is None:
|
|
self.query('INSERT INTO keys_values (key, value) VALUES (?, ?);', key, str(value))
|
|
else:
|
|
self.query('UPDATE keys_values SET value=? WHERE key = ?;', str(value), key)
|
|
|
|
def set(self, key, value):
|
|
return self.put(key, value)
|