1 files changed, 257 insertions, 0 deletions
diff --git a/velascobot2/database.py b/velascobot2/database.py
new file mode 100644
index 0000000..582495f
--- /dev/null
+++ b/velascobot2/database.py
@@ -0,0 +1,257 @@
+import sqlite3
+import traceback
+from typing import Any
+
+from velascobot2.utils import ChatData
+
+
+class _SQLiteContextMgr(object):
+    def __init__(
+        self,
+        database: str = "sqlitedb",
+        isolation_level: str = "",
+        ignore_exc: bool = False,
+    ):
+        self.database = database
+        self.isolation_level = isolation_level
+        self.ignore_exc = ignore_exc
+        self.connection = None
+        self.cursor = None
+
+    def __enter__(self):
+        try:
+            self.connection = sqlite3.connect(
+                database=self.database, isolation_level=self.isolation_level
+            )
+            self.cursor = self.connection.cursor()
+            return self.cursor
+        except Exception as ex:
+            traceback.print_exc()
+            raise ex
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        try:
+            if exc_type is not None:
+                self.connection.rollback()
+                return self.ignore_exc
+            else:
+                self.connection.commit()
+        except Exception as ex:
+            traceback.print_exc()
+            raise ex
+        finally:
+            self.cursor.close()
+            self.connection.close()
+
+
+class Database:
+    def __init__(self, bot_name: str):
+        self.db_name = f"{bot_name}.sqlite"
+        self._init_database()
+
+    def _init_database(self):
+        """creates schema on first startup"""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute("""
+                CREATE TABLE IF NOT EXISTS chats (
+                    id TEXT PRIMARY KEY,
+                    type TEXT NOT NULL,
+                    name TEXT NOT NULL,
+                    processed_messages INTEGER DEFAULT 0,
+                    message_period INTEGER DEFAULT 10,
+                    answer_chance REAL DEFAULT 0.5,
+                    restricted INTEGER DEFAULT 0,
+                    silenced INTEGER DEFAULT 0
+                )
+            """)
+
+            db.execute("""
+                CREATE TABLE IF NOT EXISTS vocab (
+                    chat_id TEXT NOT NULL,
+                    key TEXT NOT NULL,
+                    word TEXT NOT NULL,
+                    count INTEGER DEFAULT 1,
+                    PRIMARY KEY (chat_id, key, word),
+                    FOREIGN KEY (chat_id) REFERENCES chats(id)
+                )
+            """)
+
+            db.execute("""
+                CREATE INDEX IF NOT EXISTS vocab_idx
+                ON vocab(chat_id, key)
+            """)
+
+            db.execute("""
+                CREATE TABLE IF NOT EXISTS graphics (
+                    chat_id TEXT NOT NULL,
+                    type TEXT NOT NULL,
+                    graph_id TEXT NOT NULL,
+                    count INTEGER DEFAULT 1,
+                    PRIMARY KEY (chat_id, graph_id),
+                    FOREIGN KEY (chat_id) REFERENCES chats(id)
+                )
+            """)
+
+            db.execute("""
+                CREATE INDEX IF NOT EXISTS graphics_idx
+                ON graphics(chat_id)
+            """)
+
+    def ensure_chat_exists(
+        self, chat_id: str, chat_type: str, chat_name: str
+    ) -> bool:
+        """
+        checks if the chat exists, stores it with defaults if not.
+        returns True if newly created, False if already existed
+        """
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            # check existence
+            result = db.execute(
+                "SELECT id FROM chats WHERE id = ?", (chat_id,)
+            ).fetchone()
+
+            if result:
+                return False
+
+            # create new chat config entry
+            db.execute(
+                "INSERT INTO chats (id, type, name) VALUES (?, ?, ?)",
+                (chat_id, chat_type, chat_name),
+            )
+
+        return True
+
+    def add_words(self, chat_id: str, key: str, word: str, count: int):
+        """adds or updates a word ocurrence"""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "INSERT INTO vocab (chat_id, key, word, count) VALUES (?, ?, ?,"
+                " ?) ON CONFLICT(chat_id, key, word) DO UPDATE SET count ="
+                " count + ?",
+                (chat_id, key, word, count, count),
+            )
+
+    def add_graphic(self, chat_id: str, graph_type: str, graph_id: str):
+        """adds or updates a graphic ocurrence"""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "INSERT INTO graphics (chat_id, type, graph_id, count) VALUES"
+                " (?, ?, ?, ?) ON CONFLICT(chat_id, graph_id) DO UPDATE"
+                " SET count = count + ?",
+                (chat_id, graph_type, graph_id, 1, 1),
+            )
+
+    def get_next_words(self, chat_id: str, key: str) -> list[tuple[str, int]]:
+        """returns list of possible words to follow the key n-gram"""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            rows = db.execute(
+                "SELECT word, count FROM vocab WHERE chat_id = ? AND key = ?",
+                (chat_id, key),
+            ).fetchall()
+
+            if not rows:
+                return []
+
+            return rows
+
+    def get_graphics(self, chat_id: str):
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            rows = db.execute(
+                "SELECT type, graph_id, count FROM graphics WHERE chat_id = ?",
+                (chat_id,),
+            ).fetchall()
+
+            if not rows:
+                return []
+
+            return rows
+
+    def _update_chat(self, chat_id: str, field: str, value: Any):
+        """updates a field in the chats table"""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                f"UPDATE chats SET {field} = ? WHERE id = ?", (value, chat_id)
+            )
+
+    def update_chat_name(self, chat_id: str, chat_name: str):
+        """updates a field in the chats table"""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "UPDATE chats SET name = ? WHERE id = ?", (chat_name, chat_id)
+            )
+
+    def update_chat_answer_chance(self, chat_id: str, answer_chance: float):
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "UPDATE chats SET answer_chance = ? WHERE id = ?",
+                (answer_chance, chat_id),
+            )
+
+    def update_chat_message_period(self, chat_id: str, period: int):
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "UPDATE chats SET message_period = ? WHERE id = ?",
+                (period, chat_id),
+            )
+
+    def update_chat_set_restricted(self, chat_id: str, restricted: bool):
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "UPDATE chats SET restricted = ? WHERE id = ?",
+                (int(restricted), chat_id),
+            )
+
+    def update_chat_set_silenced(self, chat_id: str, silenced: bool):
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "UPDATE chats SET silenced = ? WHERE id = ?",
+                (int(silenced), chat_id),
+            )
+
+    def incr_processed_messages(self, chat_id: str):
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            db.execute(
+                "UPDATE chats SET processed_messages = processed_messages + ?"
+                " WHERE id = ?",
+                (1, chat_id),
+            )
+
+    def count_graphics(self, chat_id: str) -> int:
+        """Returns the number of graphics stored for a given chat id."""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            count = (
+                db.execute(
+                    "SELECT SUM(count) FROM graphics WHERE chat_id = ?",
+                    (chat_id,),
+                ).fetchone()[0]
+                or 0
+            )
+        return count
+
+    def count_words(self, chat_id: str, key: str) -> int:
+        """Returns the number of words that follow a given key for a given chat id."""
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            count = (
+                db.execute(
+                    "SELECT SUM(count) FROM vocab WHERE chat_id = ? AND"
+                    " key = ?",
+                    (chat_id, key),
+                ).fetchone()[0]
+                or 0
+            )
+        return count
+
+    def get_all_chats(self) -> list[ChatData]:
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            rows = db.execute("SELECT * FROM chats").fetchall()
+            return [ChatData.from_sql_row(row) for row in rows]
+
+    def get_chat_data(self, chat_id: str) -> ChatData | None:
+        with _SQLiteContextMgr(database=self.db_name) as db:
+            result = db.execute(
+                "SELECT * FROM chats WHERE id = ?", (chat_id,)
+            ).fetchone()
+
+        if result:
+            return ChatData.from_sql_row(result)
+        return None