diff --git a/abots/db/sqlite.py b/abots/db/sqlite.py index e27a858..3a38e48 100644 --- a/abots/db/sqlite.py +++ b/abots/db/sqlite.py @@ -7,8 +7,8 @@ from time import strftime from operator import itemgetter from contextlib import closing, contextmanager from threading import Lock -from sqlite3 import (connect, register_converter, PARSE_DECLTYPES, - Error as SQLiteError) +from sqlite3 import connect, register_converter, PARSE_DECLTYPES +from sqlite3 import ProgrammingError, Error as SQLiteError register_converter("BOOLEAN", lambda value: bool(int(value))) @@ -75,33 +75,46 @@ class SQLite: self.conenction.rollback() self._remove_old_backups(backup_dir, backups) - def execute(self, executor, values=tuple(), fetch=False): - with closing(self.connection.cursor()) as cursor: - if "?" in executor: - result = cursor.execute(executor, values) - else: - result = cursor.execute(executor) - if not fetch: - return True - return cursor.fetchall() + def execute(self, executor, values=tuple(), fetch=False, commit=False): + try: + with closing(self.connection.cursor()) as cursor: + args = (executor, values) if "?" in executor else (executor,) + if commit: + with self.transaction(): + result = cursor.execute(*args) + else: + result = cursor.execute(*args) - def fetch(self, executor, values=tuple()): - return self.execute(executor, values, fetch=True) + if not fetch: + return True + return cursor.fetchall() + except ProgrammingError as e: + eprint(e) + return False - def create_table(self, name, fields): + def fetch(self, executor, values=tuple(), commit=False): + return self.execute(executor, values, fetch=True, commit=commit) + + """ + e.g. create_table("test", ["id INTEGER PRIMARY KEY", "value TEXT"]) + """ + def create_table(self, name, fields, commit=False): fields_string = ",".join(fields) executor = f"CREATE TABLE {name} ({fields_string})" - return self.execute(executor) + return self.execute(executor, commit=commit) - def insert(self, table, insertion): + """ + e.g. insert("test", {"value": "abc"}) + """ + def insert(self, table, insertion, commit=False): assert isinstance(insertion, dict), "Expected dict" keys = ",".join(insertion.keys()) places = ",".join(["?"] * len(insertion)) values = tuple(self._convert(insertion.values())) executor = f"INSERT INTO {table} ({keys}) VALUES({places})" - return self.execute(executor, values) + return self.execute(executor, values, commit=commit) - def update(self, table, modification, where): + def update(self, table, modification, where, commit=False): assert isinstance(modification, dict), "Expected dict" assert isinstance(where, tuple), "Expected tuple" assert len(where) == 2, "Expected length of '2'" @@ -113,7 +126,18 @@ class SQLite: where_values = where[1] values = mod_values + where_values executor = f"UPDATE {table} SET {keys} WHERE {where_query}" - return self.execute(executor, values) + return self.execute(executor, values, commit=commit) + + def delete(self, table, where, commit=False): + assert isinstance(where, tuple), "Expected tuple" + assert len(where) == 2, "Expected length of '2'" + assert isinstance(where[0], str), "Expected str" + assert isinstance(where[1], tuple), "Expected tuple" + where_query = where[0] + where_values = where[1] + values = where_values + executor = f"DELTE FROM {table} WHERE {where_query}" + return self.execute(executor, values, commit=commit) def lookup(self, table, search, where=None): if type(search) != list: @@ -131,6 +155,18 @@ class SQLite: return self.fetch(executor, where_values) return self.fetch(executor) + def lookup_all(self, table, where=None): + return self.lookup(table, ["*"], where) + + def get_all_tables(self): + executor = "SELECT name FROM sqlite_master WHERE type='table'" + tables = self.fetch(executor) + return (table[0] for table in tables) + + def drop_table(self, name, commit=False): + executor = f"DROP TABLE IF EXISTS {name}" + return self.execute(executor, commit=commit) + @contextmanager def transaction(self): with self.lock: @@ -139,4 +175,4 @@ class SQLite: self.connection.commit() except: self.connection.rollback() - raise \ No newline at end of file + raise