Переход с SQLite3 на ODBC.

database.py

@@ -1,44 +1,43 @@
-import sqlite3
+from pyodbc import connect, Row
-from typing import List
+from typing import List, Union
 
 
 class BasicDatabase:
-    def __init__(self, filename: str):
+    def __init__(self, connection_string: str):
-        self.conn = sqlite3.connect(filename)
+        self.conn = connect(connection_string)
-        self.conn.row_factory = sqlite3.Row
         self.cursor = self.conn.cursor()
 
     def get_chats(self):
         self.cursor.execute("SELECT * FROM chats")
-        return self.cursor.fetchall()
+        return self._to_dict(self.cursor.fetchall())
 
     def get_chat(self, chat_id: int):
-        self.cursor.execute("SELECT * FROM chats WHERE id = ?", (chat_id,))
+        self.cursor.execute("SELECT * FROM chats WHERE id = ?", chat_id)
-        return self.cursor.fetchone()
+        return self._to_dict(self.cursor.fetchone())
 
     def add_chat(self, chat_id: int):
-        self.cursor.execute("INSERT INTO chats (id) VALUES (?)", (chat_id,))
+        self.cursor.execute("INSERT INTO chats (id) VALUES (?)", chat_id)
         self.conn.commit()
 
     def chat_update(self, chat_id: int, **kwargs):
-        self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = :" + f for f in kwargs) +
+        self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) +
-                            f" WHERE id = {chat_id}", kwargs)
+                            " WHERE id = ?", list(kwargs.values()) + [chat_id])
         self.conn.commit()
 
     def chat_delete(self, chat_id: int):
-        self.cursor.execute("DELETE FROM chats WHERE id = ?", (chat_id,))
+        self.cursor.execute("DELETE FROM chats WHERE id = ?", chat_id)
         self.conn.commit()
 
     def get_user(self, chat_id: int, user_id: int):
-        self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
+        self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
-        return self.cursor.fetchone()
+        return self._to_dict(self.cursor.fetchone())
 
     def get_users(self, chat_id: int):
-        self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,))
+        self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", chat_id)
-        return self.cursor.fetchall()
+        return self._to_dict(self.cursor.fetchall())
 
     def add_user(self, chat_id: int, user_id: int):
-        self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", (chat_id, user_id))
+        self.cursor.execute("INSERT INTO users (chat_id, user_id) VALUES (?, ?)", chat_id, user_id)
         self.conn.commit()
 
     def user_set_last_message(self, chat_id: int, user_id: int, last_message: int):
@@ -53,16 +52,18 @@ class BasicDatabase:
 
     def user_increment(self, chat_id: int, user_id: int, fields: List[str]):
         self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) +
-                            " WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
+                            " WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
         self.conn.commit()
 
     def user_update(self, chat_id: int, user_id: int, **kwargs):
-        self.cursor.execute("UPDATE users SET " + ", ".join(f + " = :" + f for f in kwargs) +
+        sql = "UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) + " WHERE chat_id = ? AND user_id = ?"
-                            f" WHERE chat_id = {chat_id} AND user_id = {user_id}", kwargs)
+        args = [chat_id, user_id] + list(kwargs.values())
+        self.cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) +
+                            " WHERE chat_id = ? AND user_id = ?", list(kwargs.values()) + [chat_id, user_id])
         self.conn.commit()
 
     def delete_user(self, chat_id: int, user_id: int):
-        self.cursor.execute("DELETE FROM users WHERE chat_id = ? AND user_id = ?", (chat_id, user_id))
+        self.cursor.execute("DELETE FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id)
         self.conn.commit()
 
     def get_top_messages_today(self, chat_id: int):
@@ -70,16 +71,16 @@ class BasicDatabase:
             SELECT user_id, messages_today AS value FROM users
             WHERE chat_id = ? AND messages_today > 0
             ORDER BY messages_today DESC
-            """, (chat_id,))
+            """, chat_id)
-        return self.cursor.fetchall()
+        return self._to_dict(self.cursor.fetchall())
 
     def get_top_messages_month(self, chat_id: int):
         self.cursor.execute("""
             SELECT user_id, messages_month AS value FROM users
             WHERE chat_id = ? AND messages_month > 0
             ORDER BY messages_month DESC
-            """, (chat_id,))
+            """, chat_id)
-        return self.cursor.fetchall()
+        return self._to_dict(self.cursor.fetchall())
 
     def get_top_silent(self, chat_id: int, threshold_days: int):
         # noinspection SpellCheckingInspection
@@ -92,19 +93,19 @@ class BasicDatabase:
             FROM users
             WHERE chat_id = ? AND value >= ?
             ORDER BY value DESC
-            """, (chat_id, threshold_days))
+            """, chat_id, threshold_days)
-        return self.cursor.fetchall()
+        return self._to_dict(self.cursor.fetchall())
 
     def get_top_warnings(self, chat_id: int):
         self.cursor.execute("""
             SELECT user_id, warnings AS value FROM users
             WHERE chat_id = ? AND warnings > 0
             ORDER BY warnings DESC
-            """, (chat_id,))
+            """, chat_id)
-        return self.cursor.fetchall()
+        return self._to_dict(self.cursor.fetchall())
 
     def delete_users(self, chat_id: int):
-        self.cursor.execute("DELETE FROM users WHERE chat_id = ?", (chat_id,))
+        self.cursor.execute("DELETE FROM users WHERE chat_id = ?", chat_id)
         self.conn.commit()
 
     def reset_messages_today(self):
@@ -128,3 +129,23 @@ class BasicDatabase:
             self.add_user(chat_id, user_id)
             user = self.get_user(chat_id, user_id)
         return user
+
+    def _to_dict(self, args: Union[Row, List[Row], None]):
+        columns = [column[0] for column in self.cursor.description]
+        if args is None:
+            return None
+        elif isinstance(args, Row):
+            result = {}
+            for i, column in enumerate(columns):
+                result[column] = args[i]
+            return result
+        elif isinstance(args, list) and all(isinstance(item, Row) for item in args):
+            results = []
+            for row in args:
+                row_dict = {}
+                for i, column in enumerate(columns):
+                    row_dict[column] = row[i]
+                results.append(row_dict)
+            return results
+        else:
+            raise TypeError("unexpected type")

tg/__main__.py

@@ -4,6 +4,7 @@ import json
 from aiogram import Bot, Dispatcher
 
 from ai_agent import create_ai_agent
+from tg.tg_database import create_database
 
 from . import handlers
 from . import tasks
@@ -15,6 +16,7 @@ async def main() -> None:
     print('Конфигурация загружена.')
 
     bot = Bot(token=config['api_token'])
+    create_database(config['db_connection_string'])
     create_ai_agent(config['openrouter_token'])
 
     dp = Dispatcher()

tg/tg_database.py

@@ -2,32 +2,38 @@ import database
 
 
 class TgDatabase(database.BasicDatabase):
-    def __init__(self):
+    def __init__(self, connection_string: str):
-        super().__init__('tg.db')
+        super().__init__(connection_string)
 
         self.cursor.execute("""
         CREATE TABLE IF NOT EXISTS chats (
-            "id"	    INTEGER,
+            id BIGINT,
-            "active"	INTEGER NOT NULL DEFAULT 0,
+            active TINYINT NOT NULL DEFAULT 0,
-            "rules"             TEXT,
+            rules VARCHAR(2000),
-            "greeting_join"	    TEXT,
+            greeting_join VARCHAR(1000),
-            "ai_prompt"         TEXT,
+            ai_prompt VARCHAR(2048),
-            PRIMARY KEY("id"))
+            PRIMARY KEY (id))
         """)
 
         self.cursor.execute("""
         CREATE TABLE IF NOT EXISTS users (
-            "chat_id"	INTEGER,
+            chat_id	BIGINT,
-            "user_id"	INTEGER,
+            user_id	BIGINT,
-            "last_message"	INTEGER NOT NULL DEFAULT 0,
+            last_message BIGINT NOT NULL DEFAULT 0,
-            "messages_today"	INTEGER NOT NULL DEFAULT 0,
+            messages_today SMALLINT NOT NULL DEFAULT 0,
-            "messages_month"	INTEGER NOT NULL DEFAULT 0,
+            messages_month SMALLINT NOT NULL DEFAULT 0,
-            "warnings"	INTEGER NOT NULL DEFAULT 0,
+            warnings TINYINT NOT NULL DEFAULT 0,
-            "about"   TEXT,
+            about VARCHAR(1000),
-            PRIMARY KEY("chat_id","user_id"))
+            PRIMARY KEY (chat_id, user_id),
+            CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE)
         """)
 
         self.conn.commit()
 
 
-DB = TgDatabase()
+DB: TgDatabase
+
+
+def create_database(connection_string: str):
+    global DB
+    DB = TgDatabase(connection_string)

vk/__main__.py

@@ -3,6 +3,7 @@ import json
 from vkbottle.bot import Bot as VkBot
 
 from ai_agent import create_ai_agent
+from vk.vk_database import create_database
 
 from . import handlers
 from . import tasks
@@ -14,6 +15,7 @@ if __name__ == '__main__':
     print('Конфигурация загружена.')
 
     bot = VkBot(config['api_token'], labeler=handlers.labeler)
+    create_database(config['db_connection_string'])
     create_ai_agent(config["openrouter_token"])
 
     bot.loop_wrapper.on_startup.append(tasks.startup_task(bot.api))

vk/vk_database.py

@@ -2,32 +2,33 @@ import database
 
 
 class VkDatabase(database.BasicDatabase):
-    def __init__(self):
+    def __init__(self, connection_string: str):
-        super().__init__('vk.db')
+        super().__init__(connection_string)
 
         self.cursor.execute("""
         CREATE TABLE IF NOT EXISTS chats (
-            "id"	    INTEGER,
+            id BIGINT,
-            "active"	INTEGER NOT NULL DEFAULT 0,
+            active TINYINT NOT NULL DEFAULT 0,
-            "rules"             TEXT,
+            rules VARCHAR(2000),
-            "greeting_join"	    TEXT,
+            greeting_join VARCHAR(1000),
-            "greeting_rejoin"	TEXT,
+            greeting_rejoin VARCHAR(1000),
-            "birthday_message"  TEXT,
+            birthday_message VARCHAR(1000),
-            "ai_prompt"         TEXT,
+            ai_prompt VARCHAR(2048),
-            PRIMARY KEY("id"))
+            PRIMARY KEY (id))
         """)
 
         self.cursor.execute("""
         CREATE TABLE IF NOT EXISTS users (
-            "chat_id"	INTEGER,
+            chat_id	BIGINT,
-            "user_id"	INTEGER,
+            user_id	BIGINT,
-            "last_message"	INTEGER NOT NULL DEFAULT 0,
+            last_message BIGINT NOT NULL DEFAULT 0,
-            "messages_today"	INTEGER NOT NULL DEFAULT 0,
+            messages_today SMALLINT NOT NULL DEFAULT 0,
-            "messages_month"	INTEGER NOT NULL DEFAULT 0,
+            messages_month SMALLINT NOT NULL DEFAULT 0,
-            "warnings"	INTEGER NOT NULL DEFAULT 0,
+            warnings TINYINT NOT NULL DEFAULT 0,
-            "happy_birthday" INTEGER NOT NULL DEFAULT 1,
+            happy_birthday TINYINT NOT NULL DEFAULT 1,
-            "about"   TEXT,
+            about VARCHAR(1000),
-            PRIMARY KEY("chat_id","user_id"))
+            PRIMARY KEY (chat_id, user_id),
+            CONSTRAINT fk_users_chats FOREIGN KEY (chat_id) REFERENCES chats (id) ON DELETE CASCADE)
         """)
 
         self.conn.commit()
@@ -36,4 +37,9 @@ class VkDatabase(database.BasicDatabase):
         self.user_update(chat_id, user_id, happy_birthday=happy_birthday)
 
 
-DB = VkDatabase()
+DB: VkDatabase
+
+
+def create_database(connection_string: str):
+    global DB
+    DB = VkDatabase(connection_string)