208 lines
9.6 KiB
Python
208 lines
9.6 KiB
Python
from datetime import datetime
|
|
from typing import List, Optional, Union
|
|
|
|
from pyodbc import connect, SQL_CHAR, SQL_WCHAR, Row
|
|
|
|
|
|
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 get_bots(self):
|
|
self.cursor.execute("SELECT * FROM bots")
|
|
return self._to_dict(self.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_chats(self, bot_id: int):
|
|
self.cursor.execute("SELECT * FROM chats WHERE bot_id = ?", bot_id)
|
|
return self._to_dict(self.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 add_chat(self, bot_id: int, chat_id: int):
|
|
self.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])
|
|
|
|
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)
|
|
|
|
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_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_dict(self.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)
|
|
|
|
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)
|
|
|
|
def user_increment_messages(self, bot_id: int, chat_id: int, user_id: int):
|
|
self.user_increment(bot_id, chat_id, user_id, ['messages_today', 'messages_month'])
|
|
|
|
def user_increment_warnings(self, bot_id: int, chat_id: int, user_id: int):
|
|
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)
|
|
|
|
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 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)
|
|
|
|
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_dict(self.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_dict(self.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_dict(self.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_dict(self.cursor.fetchall())
|
|
|
|
def reset_messages_today(self, bot_id: int):
|
|
self.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)
|
|
|
|
def context_get_messages(self, bot_id: int, chat_id: int) -> list[dict]:
|
|
self.cursor.execute("""
|
|
SELECT role, content FROM contexts
|
|
WHERE bot_id = ? AND chat_id = ? AND message_id IS NOT NULL
|
|
ORDER BY message_id
|
|
""", bot_id, chat_id)
|
|
return self._to_dict(self.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()
|
|
|
|
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' AND message_id IS NOT NULL
|
|
ORDER BY message_id DESC
|
|
LIMIT 1
|
|
""", bot_id, chat_id).fetchval()
|
|
|
|
def context_add_message(self, bot_id: int, chat_id: int, role: str, content: str, message_id: Optional[int], max_messages: int):
|
|
self._context_trim(bot_id, chat_id, max_messages)
|
|
|
|
if message_id is not None:
|
|
self.cursor.execute(
|
|
"INSERT INTO contexts (bot_id, chat_id, message_id, role, content) VALUES (?, ?, ?, ?, ?)",
|
|
bot_id, chat_id, message_id, role, content)
|
|
else:
|
|
self.cursor.execute("INSERT INTO contexts (bot_id, chat_id, role, content) VALUES (?, ?, ?, ?)",
|
|
bot_id, chat_id, role, content)
|
|
|
|
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)
|
|
|
|
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 message_id FROM contexts
|
|
WHERE bot_id = ? AND chat_id = ? AND message_id IS NOT NULL
|
|
ORDER BY message_id ASC
|
|
LIMIT 1
|
|
""", bot_id, chat_id).fetchval()
|
|
|
|
if oldest_message_id:
|
|
self.cursor.execute("DELETE FROM contexts WHERE bot_id = ? AND chat_id = ? AND message_id = ?",
|
|
bot_id, chat_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)
|
|
|
|
def create_chat_if_not_exists(self, bot_id: int, chat_id: int):
|
|
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)
|
|
return chat
|
|
|
|
def create_user_if_not_exists(self, bot_id: int, chat_id: int, user_id: int):
|
|
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)
|
|
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: 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
|
|
else:
|
|
raise TypeError("unexpected type")
|