from pyodbc import connect, Row from typing import List, Union class BasicDatabase: def __init__(self, connection_string: str): self.conn = connect(connection_string, autocommit=True) self.cursor = self.conn.cursor() def get_chats(self): self.cursor.execute("SELECT * FROM chats") 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._to_dict(self.cursor.fetchone()) def add_chat(self, chat_id: int): self.cursor.execute("INSERT INTO chats (id) VALUES (?)", chat_id) def chat_update(self, chat_id: int, **kwargs): self.cursor.execute("UPDATE chats SET " + ", ".join(f + " = ?" for f in kwargs) + " WHERE id = ?", list(kwargs.values()) + [chat_id]) def chat_delete(self, chat_id: int): self.cursor.execute("DELETE FROM chats WHERE id = ?", chat_id) 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._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._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) 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']) 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) def user_update(self, chat_id: int, user_id: int, **kwargs): 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]) 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) def get_top_messages_today(self, chat_id: int): self.cursor.execute(""" 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._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._to_dict(self.cursor.fetchall()) def get_top_silent(self, chat_id: int, threshold_days: int): # noinspection SpellCheckingInspection self.cursor.execute(""" SELECT user_id, CASE WHEN last_message = 0 THEN 'никогда' ELSE (unixepoch() - last_message) / 86400 END AS value FROM users WHERE chat_id = ? AND value >= ? ORDER BY value DESC """, 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._to_dict(self.cursor.fetchall()) def reset_messages_today(self): self.cursor.execute("UPDATE users SET messages_today = 0") def reset_messages_month(self): self.cursor.execute("UPDATE users SET messages_month = 0") def create_chat_if_not_exists(self, chat_id: int): chat = self.get_chat(chat_id) if chat is None: self.add_chat(chat_id) chat = self.get_chat(chat_id) return chat def create_user_if_not_exists(self, chat_id: int, user_id: int): user = self.get_user(chat_id, user_id) if user is None: 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")