Added lookup_all, get_all_tables, delete, commit argument, and some docstrings

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
+                raise