diff --git a/database.py b/database.py index ccac189..79282a8 100644 --- a/database.py +++ b/database.py @@ -1,44 +1,43 @@ -import sqlite3 -from typing import List +from pyodbc import connect, Row +from typing import List, Union class BasicDatabase: - def __init__(self, filename: str): - self.conn = sqlite3.connect(filename) - self.conn.row_factory = sqlite3.Row + def __init__(self, connection_string: str): + self.conn = connect(connection_string) 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,)) - return self.cursor.fetchone() + self.cursor.execute("SELECT * FROM chats WHERE id = ?", chat_id) + 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) + - f" WHERE id = {chat_id}", kwargs) + self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in 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)) - return self.cursor.fetchone() + self.cursor.execute("SELECT * FROM users WHERE chat_id = ? AND user_id = ?", chat_id, user_id) + 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,)) - return self.cursor.fetchall() + self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", chat_id) + 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) + - f" WHERE chat_id = {chat_id} AND user_id = {user_id}", kwargs) + sql = "UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) + " WHERE chat_id = ? AND user_id = ?" + 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,)) - return self.cursor.fetchall() + """, chat_id) + 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,)) - return self.cursor.fetchall() + """, chat_id) + 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)) - return self.cursor.fetchall() + """, chat_id, threshold_days) + 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,)) - return self.cursor.fetchall() + """, chat_id) + 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") diff --git a/tg/__main__.py b/tg/__main__.py index bf69272..5c96ff9 100644 --- a/tg/__main__.py +++ b/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() diff --git a/tg/tg_database.py b/tg/tg_database.py index f8924da..eb3c7fa 100644 --- a/tg/tg_database.py +++ b/tg/tg_database.py @@ -2,32 +2,38 @@ import database class TgDatabase(database.BasicDatabase): - def __init__(self): - super().__init__('tg.db') + def __init__(self, connection_string: str): + super().__init__(connection_string) self.cursor.execute(""" CREATE TABLE IF NOT EXISTS chats ( - "id" INTEGER, - "active" INTEGER NOT NULL DEFAULT 0, - "rules" TEXT, - "greeting_join" TEXT, - "ai_prompt" TEXT, - PRIMARY KEY("id")) + id BIGINT, + active TINYINT NOT NULL DEFAULT 0, + rules VARCHAR(2000), + greeting_join VARCHAR(1000), + ai_prompt VARCHAR(2048), + PRIMARY KEY (id)) """) self.cursor.execute(""" CREATE TABLE IF NOT EXISTS users ( - "chat_id" INTEGER, - "user_id" INTEGER, - "last_message" INTEGER NOT NULL DEFAULT 0, - "messages_today" INTEGER NOT NULL DEFAULT 0, - "messages_month" INTEGER NOT NULL DEFAULT 0, - "warnings" INTEGER NOT NULL DEFAULT 0, - "about" TEXT, - PRIMARY KEY("chat_id","user_id")) + chat_id BIGINT, + user_id BIGINT, + last_message BIGINT NOT NULL DEFAULT 0, + messages_today SMALLINT NOT NULL DEFAULT 0, + messages_month SMALLINT NOT NULL DEFAULT 0, + warnings TINYINT NOT NULL DEFAULT 0, + about VARCHAR(1000), + 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) diff --git a/vk/__main__.py b/vk/__main__.py index b028904..dccf44e 100644 --- a/vk/__main__.py +++ b/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)) diff --git a/vk/vk_database.py b/vk/vk_database.py index 3d50ecb..8660c5b 100644 --- a/vk/vk_database.py +++ b/vk/vk_database.py @@ -2,32 +2,33 @@ import database class VkDatabase(database.BasicDatabase): - def __init__(self): - super().__init__('vk.db') + def __init__(self, connection_string: str): + super().__init__(connection_string) self.cursor.execute(""" CREATE TABLE IF NOT EXISTS chats ( - "id" INTEGER, - "active" INTEGER NOT NULL DEFAULT 0, - "rules" TEXT, - "greeting_join" TEXT, - "greeting_rejoin" TEXT, - "birthday_message" TEXT, - "ai_prompt" TEXT, - PRIMARY KEY("id")) + id BIGINT, + active TINYINT NOT NULL DEFAULT 0, + rules VARCHAR(2000), + greeting_join VARCHAR(1000), + greeting_rejoin VARCHAR(1000), + birthday_message VARCHAR(1000), + ai_prompt VARCHAR(2048), + PRIMARY KEY (id)) """) self.cursor.execute(""" CREATE TABLE IF NOT EXISTS users ( - "chat_id" INTEGER, - "user_id" INTEGER, - "last_message" INTEGER NOT NULL DEFAULT 0, - "messages_today" INTEGER NOT NULL DEFAULT 0, - "messages_month" INTEGER NOT NULL DEFAULT 0, - "warnings" INTEGER NOT NULL DEFAULT 0, - "happy_birthday" INTEGER NOT NULL DEFAULT 1, - "about" TEXT, - PRIMARY KEY("chat_id","user_id")) + chat_id BIGINT, + user_id BIGINT, + last_message BIGINT NOT NULL DEFAULT 0, + messages_today SMALLINT NOT NULL DEFAULT 0, + messages_month SMALLINT NOT NULL DEFAULT 0, + warnings TINYINT NOT NULL DEFAULT 0, + happy_birthday TINYINT NOT NULL DEFAULT 1, + about VARCHAR(1000), + 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)