import sqlite3 from typing import List class Database: def __init__(self): self.conn = sqlite3.connect('bot.db') self.conn.row_factory = sqlite3.Row self.cursor = self.conn.cursor() self.cursor.execute(""" CREATE TABLE IF NOT EXISTS chats ( "id" INTEGER, "active" INTEGER NOT NULL DEFAULT 0, 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, PRIMARY KEY("chat_id","user_id")) """) self.conn.commit() def get_chat(self, chat_id: int): self.cursor.execute("SELECT * FROM chats WHERE id = ?", (chat_id,)) return self.cursor.fetchone() def add_chat(self, chat_id: int): 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.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() def get_users(self, chat_id: int): self.cursor.execute("SELECT * FROM users WHERE chat_id = ?", (chat_id,)) return 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.conn.commit() def user_set_last_message(self, chat_id: int, user_id: int, last_message: int): self.user_update(chat_id, user_id, last_message=last_message) def user_increment_messages(self, chat_id: int, user_id: int): self.user_increment(chat_id, user_id, ['messages_today', 'messages_month']) self.conn.commit() def user_increment_warnings(self, chat_id: int, user_id: int): self.user_increment(chat_id, user_id, ['warnings']) 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)) 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) self.conn.commit() def get_top_messages_today(self, chat_id: int): self.cursor.execute(""" SELECT user_id, messages_today FROM users WHERE chat_id = ? AND messages_today > 0 ORDER BY messages_today DESC """, (chat_id,)) return self.cursor.fetchall() def get_top_messages_month(self, chat_id: int): self.cursor.execute(""" SELECT user_id, messages_month FROM users WHERE chat_id = ? AND messages_month > 0 ORDER BY messages_month DESC """, (chat_id,)) return self.cursor.fetchall() def get_top_silent(self, chat_id: int, threshold_time: int): self.cursor.execute(""" SELECT user_id, last_message FROM users WHERE chat_id = ? AND last_message < ? ORDER BY last_message ASC """, (chat_id, threshold_time)) return self.cursor.fetchall() def get_top_warnings(self, chat_id: int): self.cursor.execute(""" SELECT user_id, warnings FROM users WHERE chat_id = ? AND warnings > 0 ORDER BY warnings DESC """, (chat_id,)) return self.cursor.fetchall() def reset_messages_today(self): self.cursor.execute("UPDATE users SET messages_today = 0") self.conn.commit() def reset_messages_month(self): self.cursor.execute("UPDATE users SET messages_month = 0") self.conn.commit() DB = Database()