From 5b9f5cd1d6568be4543ed22e19499d882f7db71c Mon Sep 17 00:00:00 2001 From: Kirill Kirilenko Date: Sun, 14 Jun 2026 00:02:29 +0300 Subject: [PATCH] =?UTF-8?q?=D0=9F=D0=B5=D1=80=D0=B5=D1=85=D0=BE=D0=B4=20?= =?UTF-8?q?=D1=81=20pyodbc=20=D0=BD=D0=B0=20mariadb.?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- database.py | 296 ++++++++++++++++++++++++++-------------------- requirements.txt | 2 +- tg/__main__.py | 2 +- tg/tg_database.py | 109 ++++++++--------- vk/__main__.py | 2 +- vk/vk_database.py | 115 +++++++++--------- 6 files changed, 284 insertions(+), 242 deletions(-) diff --git a/database.py b/database.py index 497d48a..cb85f1b 100644 --- a/database.py +++ b/database.py @@ -1,55 +1,76 @@ from datetime import datetime -from typing import Dict, List, Optional +from typing import Dict, List, Optional, Any, Type -from pyodbc import connect, SQL_CHAR, SQL_WCHAR, Row +from mariadb import create_pool class BasicDatabase: - def __init__(self, connection_string: str): - self.conn = connect(connection_string, autocommit=True) - self.conn.setdecoding(SQL_CHAR, encoding='utf-8') - self.conn.setdecoding(SQL_WCHAR, encoding='utf-8') - self.conn.setencoding(encoding='utf-8') - self.cursor = self.conn.cursor() + def __init__(self, hostname: str, user: str, password: str, dbname: str): + self.hostname = hostname + self.user = user + self.password = password + self.dbname = dbname + self.pool = create_pool(host=hostname, user=user, password=password, database=dbname, + autocommit=True, dictionary=True, min_size=1, max_size=5) - def get_bots(self): - self.cursor.execute("SELECT * FROM bots") - return self._to_list(self.cursor.fetchall()) + def get_bots(self) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT * FROM bots") + return cursor.fetchall() - def get_bot(self, bot_id: int): - self.cursor.execute("SELECT * FROM bots WHERE id = ?", bot_id) - return self._to_dict(self.cursor.fetchone()) + def get_bot(self, bot_id: int) -> Optional[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT * FROM bots WHERE id = ?", (bot_id,)) + return cursor.fetchone() - def get_chats(self, bot_id: int): - self.cursor.execute("SELECT * FROM chats WHERE bot_id = ?", bot_id) - return self._to_list(self.cursor.fetchall()) + def get_chats(self, bot_id: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT * FROM chats WHERE bot_id = ?", (bot_id,)) + return cursor.fetchall() - def get_chat(self, bot_id: int, chat_id: int): - self.cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) - return self._to_dict(self.cursor.fetchone()) + def get_chat(self, bot_id: int, chat_id: int) -> Optional[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT * FROM chats WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id)) + return cursor.fetchone() def add_chat(self, bot_id: int, chat_id: int): - self.cursor.execute("INSERT INTO chats (bot_id, chat_id) VALUES (?, ?)", bot_id, chat_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("INSERT INTO chats (bot_id, chat_id) VALUES (?, ?)", (bot_id, chat_id)) def chat_update(self, bot_id: int, chat_id: int, **kwargs): - self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) + - " WHERE bot_id = ? AND chat_id = ?", list(kwargs.values()) + [bot_id, chat_id]) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) + + " WHERE bot_id = ? AND chat_id = ?", list(kwargs.values()) + [bot_id, chat_id]) def chat_delete(self, bot_id: int, chat_id: int): - self.cursor.execute("DELETE FROM chats WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("DELETE FROM chats WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id)) - def get_user(self, bot_id: int, chat_id: int, user_id: int): - self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?", - bot_id, chat_id, user_id) - return self._to_dict(self.cursor.fetchone()) + def get_user(self, bot_id: int, chat_id: int, user_id: int) -> Optional[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?", + (bot_id, chat_id, user_id)) + return cursor.fetchone() - def get_users(self, bot_id: int, chat_id: int): - self.cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) - return self._to_list(self.cursor.fetchall()) + def get_users(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT * FROM users WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id)) + return cursor.fetchall() def add_user(self, bot_id: int, chat_id: int, user_id: int): - self.cursor.execute("INSERT INTO users (bot_id, chat_id, user_id) VALUES (?, ?, ?)", - bot_id, chat_id, user_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("INSERT INTO users (bot_id, chat_id, user_id) VALUES (?, ?, ?)", + (bot_id, chat_id, user_id)) def user_set_last_message(self, bot_id: int, chat_id: int, user_id: int, last_message: int): self.user_update(bot_id, chat_id, user_id, last_message=last_message) @@ -61,83 +82,107 @@ class BasicDatabase: self.user_increment(bot_id, chat_id, user_id, ['warnings']) def user_increment(self, bot_id: int, chat_id: int, user_id: int, fields: List[str]): - self.cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) + - " WHERE bot_id = ? AND chat_id = ? AND user_id = ?", bot_id, chat_id, user_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("UPDATE users SET " + ", ".join(f + " = " + f + " + 1" for f in fields) + + " WHERE bot_id = ? AND chat_id = ? AND user_id = ?", (bot_id, chat_id, user_id)) - def user_update(self, bot_id, chat_id: int, user_id: int, **kwargs): - self.cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) + - " WHERE bot_id = ? AND chat_id = ? AND user_id = ?", - list(kwargs.values()) + [bot_id, chat_id, user_id]) + def user_update(self, bot_id: int, chat_id: int, user_id: int, **kwargs): + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("UPDATE users SET " + ", ".join(f + " = ?" for f in kwargs) + + " WHERE bot_id = ? AND chat_id = ? AND user_id = ?", + list(kwargs.values()) + [bot_id, chat_id, user_id]) def delete_user(self, bot_id: int, chat_id: int, user_id: int): - self.cursor.execute("DELETE FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?", - bot_id, chat_id, user_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("DELETE FROM users WHERE bot_id = ? AND chat_id = ? AND user_id = ?", + (bot_id, chat_id, user_id)) - def get_top_messages_today(self, bot_id: int, chat_id: int): - self.cursor.execute(""" - SELECT user_id, messages_today AS value FROM users - WHERE bot_id = ? AND chat_id = ? AND messages_today > 0 - ORDER BY messages_today DESC - """, bot_id, chat_id) - return self._to_list(self.cursor.fetchall()) + def get_top_messages_today(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + SELECT user_id, messages_today AS value FROM users + WHERE bot_id = ? AND chat_id = ? AND messages_today > 0 + ORDER BY messages_today DESC + """, (bot_id, chat_id)) + return cursor.fetchall() - def get_top_messages_month(self, bot_id: int, chat_id: int): - self.cursor.execute(""" - SELECT user_id, messages_month AS value FROM users - WHERE bot_id = ? AND chat_id = ? AND messages_month > 0 - ORDER BY messages_month DESC - """, bot_id, chat_id) - return self._to_list(self.cursor.fetchall()) + def get_top_messages_month(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + SELECT user_id, messages_month AS value FROM users + WHERE bot_id = ? AND chat_id = ? AND messages_month > 0 + ORDER BY messages_month DESC + """, (bot_id, chat_id)) + return cursor.fetchall() - def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int): - current_time = int(datetime.now().timestamp()) - threshold = current_time - threshold_days * 86400 - self.cursor.execute(""" - SELECT user_id, (? - last_message) DIV 86400 as value - FROM users - WHERE bot_id = ? AND chat_id = ? AND last_message <= ? - ORDER BY last_message ASC - """, current_time, bot_id, chat_id, threshold) - result = self._to_list(self.cursor.fetchall()) - for row in result: - if row['value'] > 3650: - row['value'] = 'никогда' - return result + def get_top_silent(self, bot_id: int, chat_id: int, threshold_days: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + current_time = int(datetime.now().timestamp()) + threshold = current_time - threshold_days * 86400 + cursor.execute(""" + SELECT user_id, (? - last_message) DIV 86400 as value + FROM users + WHERE bot_id = ? AND chat_id = ? AND last_message <= ? + ORDER BY last_message ASC + """, (current_time, bot_id, chat_id, threshold)) + result = cursor.fetchall() + for row in result: + if row['value'] > 3650: + row['value'] = 'никогда' + return result - def get_top_warnings(self, bot_id: int, chat_id: int): - self.cursor.execute(""" - SELECT user_id, warnings AS value FROM users - WHERE bot_id = ? AND chat_id = ? AND warnings > 0 - ORDER BY warnings DESC - """, bot_id, chat_id) - return self._to_list(self.cursor.fetchall()) + def get_top_warnings(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + SELECT user_id, warnings AS value FROM users + WHERE bot_id = ? AND chat_id = ? AND warnings > 0 + ORDER BY warnings DESC + """, (bot_id, chat_id)) + return cursor.fetchall() def reset_messages_today(self, bot_id: int): - self.cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", bot_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("UPDATE users SET messages_today = 0 WHERE bot_id = ?", (bot_id,)) def reset_messages_month(self, bot_id: int): - self.cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", bot_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("UPDATE users SET messages_month = 0 WHERE bot_id = ?", (bot_id,)) - def context_get_messages(self, bot_id: int, chat_id: int) -> List[Dict]: - self.cursor.execute(""" - SELECT role, text, image FROM contexts - WHERE bot_id = ? AND chat_id = ? - ORDER BY id - """, bot_id, chat_id) - result = self._to_list(self.cursor.fetchall()) - return result + def context_get_messages(self, bot_id: int, chat_id: int) -> List[Dict[str, Any]]: + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + SELECT role, text, image FROM contexts + WHERE bot_id = ? AND chat_id = ? + ORDER BY id + """, (bot_id, chat_id)) + return cursor.fetchall() def context_get_count(self, bot_id: int, chat_id: int) -> int: - self.cursor.execute("SELECT COUNT(*) FROM contexts WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) - return self.cursor.fetchval() + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("SELECT COUNT(*) FROM contexts WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id)) + return _to_val(int, cursor.fetchone()) def context_get_last_assistant_message_id(self, bot_id: int, chat_id: int) -> Optional[int]: - return self.cursor.execute(""" - SELECT message_id FROM contexts - WHERE bot_id = ? AND chat_id = ? AND role = 'assistant' - ORDER BY id DESC - LIMIT 1 - """, bot_id, chat_id).fetchval() + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + SELECT message_id FROM contexts + WHERE bot_id = ? AND chat_id = ? AND role = 'assistant' + ORDER BY id DESC + LIMIT 1 + """, (bot_id, chat_id)) + return _to_val(int, cursor.fetchone()) def context_add_message(self, bot_id: int, chat_id: int, role: str, text: Optional[str], image: Optional[bytes], @@ -158,60 +203,55 @@ class BasicDatabase: placeholders = ', '.join(['?' for _ in columns]) values = tuple(data[k] for k in columns) - query = f"INSERT INTO contexts ({', '.join(columns)}) VALUES ({placeholders})" - self.cursor.execute(query, values) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + query = f"INSERT INTO contexts ({', '.join(columns)}) VALUES ({placeholders})" + cursor.execute(query, values) def context_set_last_message_id(self, bot_id: int, chat_id: int, message_id: int): - self.cursor.execute( - "UPDATE contexts SET message_id = ? WHERE bot_id = ? AND chat_id = ? AND message_id IS NULL", - message_id, bot_id, chat_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute( + "UPDATE contexts SET message_id = ? WHERE bot_id = ? AND chat_id = ? AND message_id IS NULL", + (message_id, bot_id, chat_id)) def _context_trim(self, bot_id: int, chat_id: int, max_messages: int): current_count = self.context_get_count(bot_id, chat_id) while current_count >= max_messages: - oldest_message_id = self.cursor.execute( - "SELECT id FROM contexts WHERE bot_id = ? AND chat_id = ? ORDER BY id LIMIT 1", - bot_id, chat_id).fetchval() + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute( + "SELECT id FROM contexts WHERE bot_id = ? AND chat_id = ? ORDER BY id LIMIT 1", + (bot_id, chat_id)) + oldest_message_id = _to_val(int, cursor.fetchone()) - if oldest_message_id: - self.cursor.execute("DELETE FROM contexts WHERE id = ?", oldest_message_id) - current_count -= 1 - else: - break + if oldest_message_id: + cursor.execute("DELETE FROM contexts WHERE id = ?", oldest_message_id) + current_count -= 1 + else: + break def context_clear(self, bot_id: int, chat_id: int): - self.cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ?", bot_id, chat_id) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ?", (bot_id, chat_id)) - def create_chat_if_not_exists(self, bot_id: int, chat_id: int): + def create_chat_if_not_exists(self, bot_id: int, chat_id: int) -> Dict[str, Any]: chat = self.get_chat(bot_id, chat_id) if chat is None: self.add_chat(bot_id, chat_id) chat = self.get_chat(bot_id, chat_id) + assert (chat is not None) return chat - def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int): + def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int) -> Dict[str, Any]: user = self.get_user(bot_id, chat_id, user_id) if user is None: self.add_user(bot_id, chat_id, user_id) user = self.get_user(bot_id, chat_id, user_id) + assert (user is not None) return user - def _to_dict(self, args: Optional[Row]) -> Optional[Dict]: - columns = [column[0] for column in self.cursor.description] - if args is not None: - result = {} - for i, column in enumerate(columns): - result[column] = args[i] - return result - else: - return None - def _to_list(self, args: List[Row]) -> List[Dict]: - columns = [column[0] for column in self.cursor.description] - results: list[dict] = [] - for row in args: - row_dict = {} - for i, column in enumerate(columns): - row_dict[column] = row[i] - results.append(row_dict) - return results +def _to_val[T](_: Type[T], arg: Any) -> T: + return next(iter(arg.values())) if arg is not None else None diff --git a/requirements.txt b/requirements.txt index 6e5f14c..ab7ac62 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,7 +2,7 @@ aiogram~=3.28.2 aiohttp~=3.13.5 vkbottle~=4.8.2 vkbottle-types~=5.199.99.20 -pyodbc~=5.3.0 +mariadb[pool]~=2.0.0rc2 openrouter==0.9.1 replicate~=1.0.7 tavily~=1.1.0 diff --git a/tg/__main__.py b/tg/__main__.py index c5bda8c..9652976 100644 --- a/tg/__main__.py +++ b/tg/__main__.py @@ -22,7 +22,7 @@ async def main() -> None: config = json.load(file) print('Конфигурация загружена.') - database.create_database(config['db_connection_string']) + database.create_database(config['db_hostname'], config['db_user'], config['db_password'], config['db_database']) create_ai_agent(config['openrouter_token'], config['openrouter_model'], config['replicate_token'], config['tavily_token'], diff --git a/tg/tg_database.py b/tg/tg_database.py index 7198a32..45f4f34 100644 --- a/tg/tg_database.py +++ b/tg/tg_database.py @@ -2,67 +2,68 @@ import database class TgDatabase(database.BasicDatabase): - def __init__(self, connection_string: str): - super().__init__(connection_string) + def __init__(self, hostname: str, user: str, password: str, dbname: str): + super().__init__(hostname, user, password, dbname) - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS bots ( - id BIGINT NOT NULL, - owner_id BIGINT NOT NULL, - api_token VARCHAR(64) NOT NULL, - ai_prompt VARCHAR(2000) DEFAULT NULL, - group_chats_allowed TINYINT NOT NULL DEFAULT 1, - private_chats_allowed TINYINT NOT NULL DEFAULT 1, - PRIMARY KEY (id)) - """) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + CREATE TABLE IF NOT EXISTS bots ( + id BIGINT NOT NULL, + owner_id BIGINT NOT NULL, + api_token VARCHAR(64) NOT NULL, + ai_prompt VARCHAR(2000) DEFAULT NULL, + group_chats_allowed TINYINT NOT NULL DEFAULT 1, + private_chats_allowed TINYINT NOT NULL DEFAULT 1, + PRIMARY KEY (id) + )""") - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS chats ( - bot_id BIGINT NOT NULL, - chat_id BIGINT NOT NULL, - active TINYINT NOT NULL DEFAULT 0, - rules VARCHAR(4000), - greeting_join VARCHAR(2000), - ai_prompt VARCHAR(2000), - PRIMARY KEY (bot_id, chat_id), - CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE) - """) + cursor.execute(""" + CREATE TABLE IF NOT EXISTS chats ( + bot_id BIGINT NOT NULL, + chat_id BIGINT NOT NULL, + active TINYINT NOT NULL DEFAULT 0, + rules VARCHAR(4000), + greeting_join VARCHAR(2000), + ai_prompt VARCHAR(2000), + PRIMARY KEY (bot_id, chat_id), + CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) + ON UPDATE CASCADE ON DELETE CASCADE + )""") - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS users ( - bot_id BIGINT NOT NULL, - chat_id BIGINT NOT NULL, - user_id BIGINT NOT NULL, - 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 (bot_id, chat_id, user_id), - CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) - ON UPDATE CASCADE ON DELETE CASCADE - )""") + cursor.execute(""" + CREATE TABLE IF NOT EXISTS users ( + bot_id BIGINT NOT NULL, + chat_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + 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 (bot_id, chat_id, user_id), + CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) + ON UPDATE CASCADE ON DELETE CASCADE + )""") - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS contexts ( - id BIGINT NOT NULL auto_increment, - bot_id BIGINT NOT NULL, - chat_id BIGINT NOT NULL, - message_id BIGINT, - role VARCHAR(16) NOT NULL, - text VARCHAR(4000), - image MEDIUMBLOB, - PRIMARY KEY (id), - CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) - ON UPDATE CASCADE ON DELETE CASCADE - )""") - - self.conn.commit() + cursor.execute(""" + CREATE TABLE IF NOT EXISTS contexts ( + id BIGINT NOT NULL auto_increment, + bot_id BIGINT NOT NULL, + chat_id BIGINT NOT NULL, + message_id BIGINT, + role VARCHAR(16) NOT NULL, + text VARCHAR(4000), + image MEDIUMBLOB, + PRIMARY KEY (id), + CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) + ON UPDATE CASCADE ON DELETE CASCADE + )""") DB: TgDatabase -def create_database(connection_string: str): +def create_database(hostname: str, user: str, password: str, dbname: str): global DB - DB = TgDatabase(connection_string) + DB = TgDatabase(hostname, user, password, dbname) diff --git a/vk/__main__.py b/vk/__main__.py index 49cb904..83720a3 100644 --- a/vk/__main__.py +++ b/vk/__main__.py @@ -22,7 +22,7 @@ if __name__ == '__main__': config = json.load(file) print('Конфигурация загружена.') - database.create_database(config['db_connection_string']) + database.create_database(config['db_hostname'], config['db_user'], config['db_password'], config['db_database']) create_ai_agent(config['openrouter_token'], config['openrouter_model'], config['replicate_token'], config['tavily_token'], diff --git a/vk/vk_database.py b/vk/vk_database.py index 47006b8..88f819c 100644 --- a/vk/vk_database.py +++ b/vk/vk_database.py @@ -2,65 +2,66 @@ import database class VkDatabase(database.BasicDatabase): - def __init__(self, connection_string: str): - super().__init__(connection_string) + def __init__(self, hostname: str, user: str, password: str, dbname: str): + super().__init__(hostname, user, password, dbname) - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS bots ( - id BIGINT NOT NULL, - owner_id BIGINT NOT NULL, - api_token VARCHAR(256) NOT NULL, - ai_prompt VARCHAR(2000) DEFAULT NULL, - group_chats_allowed TINYINT NOT NULL DEFAULT 1, - private_chats_allowed TINYINT NOT NULL DEFAULT 1, - PRIMARY KEY (id)) - """) + with self.pool.acquire() as conn: + with conn.cursor() as cursor: + cursor.execute(""" + CREATE TABLE IF NOT EXISTS bots ( + id BIGINT NOT NULL, + owner_id BIGINT NOT NULL, + api_token VARCHAR(256) NOT NULL, + ai_prompt VARCHAR(2000) DEFAULT NULL, + group_chats_allowed TINYINT NOT NULL DEFAULT 1, + private_chats_allowed TINYINT NOT NULL DEFAULT 1, + PRIMARY KEY (id) + )""") - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS chats ( - bot_id BIGINT NOT NULL, - chat_id BIGINT NOT NULL, - active TINYINT NOT NULL DEFAULT 0, - rules VARCHAR(4000), - greeting_join VARCHAR(2000), - greeting_rejoin VARCHAR(2000), - birthday_message VARCHAR(2000), - ai_prompt VARCHAR(2000), - PRIMARY KEY (bot_id, chat_id), - CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) ON DELETE CASCADE ON UPDATE CASCADE) - """) + cursor.execute(""" + CREATE TABLE IF NOT EXISTS chats ( + bot_id BIGINT NOT NULL, + chat_id BIGINT NOT NULL, + active TINYINT NOT NULL DEFAULT 0, + rules VARCHAR(4000), + greeting_join VARCHAR(2000), + greeting_rejoin VARCHAR(2000), + birthday_message VARCHAR(2000), + ai_prompt VARCHAR(2000), + PRIMARY KEY (bot_id, chat_id), + CONSTRAINT fk_chats_bots FOREIGN KEY (bot_id) REFERENCES bots (id) + ON UPDATE CASCADE ON DELETE CASCADE + )""") - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS users ( - bot_id BIGINT NOT NULL, - chat_id BIGINT NOT NULL, - user_id BIGINT NOT NULL, - 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 (bot_id, chat_id, user_id), - CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) - ON UPDATE CASCADE ON DELETE CASCADE - )""") + cursor.execute(""" + CREATE TABLE IF NOT EXISTS users ( + bot_id BIGINT NOT NULL, + chat_id BIGINT NOT NULL, + user_id BIGINT NOT NULL, + 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 (bot_id, chat_id, user_id), + CONSTRAINT fk_users_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) + ON UPDATE CASCADE ON DELETE CASCADE + )""") - self.cursor.execute(""" - CREATE TABLE IF NOT EXISTS contexts ( - id BIGINT NOT NULL auto_increment, - bot_id BIGINT NOT NULL, - chat_id BIGINT NOT NULL, - message_id BIGINT, - role VARCHAR(16) NOT NULL, - text VARCHAR(4000), - image MEDIUMBLOB, - PRIMARY KEY (id), - CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) - ON UPDATE CASCADE ON DELETE CASCADE - )""") - - self.conn.commit() + cursor.execute(""" + CREATE TABLE IF NOT EXISTS contexts ( + id BIGINT NOT NULL auto_increment, + bot_id BIGINT NOT NULL, + chat_id BIGINT NOT NULL, + message_id BIGINT, + role VARCHAR(16) NOT NULL, + text VARCHAR(4000), + image MEDIUMBLOB, + PRIMARY KEY (id), + CONSTRAINT fk_contexts_chats FOREIGN KEY (bot_id, chat_id) REFERENCES chats (bot_id, chat_id) + ON UPDATE CASCADE ON DELETE CASCADE + )""") def user_toggle_happy_birthday(self, bot_id: int, chat_id: int, user_id: int, happy_birthday: int): self.user_update(bot_id, chat_id, user_id, happy_birthday=happy_birthday) @@ -69,6 +70,6 @@ class VkDatabase(database.BasicDatabase): DB: VkDatabase -def create_database(connection_string: str): +def create_database(hostname: str, user: str, password: str, dbname: str): global DB - DB = VkDatabase(connection_string) + DB = VkDatabase(hostname, user, password, dbname)