import logging import sqlite3 from contextlib import closing from config import DB_PATH logging.basicConfig(level=logging.INFO) logger = logging.getLogger("database") def create_connection() -> sqlite3.Connection: conn = sqlite3.connect(DB_PATH) conn.row_factory = sqlite3.Row return conn def create_tables() -> None: with closing(create_connection()) as conn: cursor = conn.cursor() cursor.execute( ''' CREATE TABLE IF NOT EXISTS file_ids ( file_key TEXT PRIMARY KEY, file_id TEXT ) ''' ) cursor.execute( ''' CREATE TABLE IF NOT EXISTS user_data ( user_id INTEGER, chat_id INTEGER, count INTEGER NOT NULL DEFAULT 0, start_time REAL NOT NULL DEFAULT 0, ban_trigger INTEGER NOT NULL DEFAULT 0, ban_time REAL NOT NULL DEFAULT 0, username TEXT, PRIMARY KEY (user_id, chat_id) ) ''' ) cursor.execute( ''' CREATE TABLE IF NOT EXISTS user_data_ban ( user_id INTEGER, chat_id INTEGER, command TEXT, count INTEGER NOT NULL DEFAULT 0, ban_time_inf REAL NOT NULL DEFAULT 0, PRIMARY KEY (user_id, chat_id, command) ) ''' ) conn.commit() def get_file_id(file_key: str) -> str | None: try: with closing(create_connection()) as conn: cursor = conn.cursor() cursor.execute("SELECT file_id FROM file_ids WHERE file_key = ?", (file_key,)) row = cursor.fetchone() return row["file_id"] if row and row["file_id"] else None except Exception as exc: logger.error("Ошибка при получении file_id: %s", exc) return None def save_file_id(file_key: str, file_id: str | None) -> None: try: with closing(create_connection()) as conn: cursor = conn.cursor() if file_id is None: cursor.execute("DELETE FROM file_ids WHERE file_key = ?", (file_key,)) else: cursor.execute( "REPLACE INTO file_ids (file_key, file_id) VALUES (?, ?)", (file_key, file_id), ) conn.commit() except Exception as exc: logger.error("Ошибка при сохранении file_id: %s", exc) def get_user_data(user_id: int, chat_id: int): try: with closing(create_connection()) as conn: cursor = conn.cursor() cursor.execute( "SELECT * FROM user_data WHERE user_id = ? AND chat_id = ?", (user_id, chat_id), ) row = cursor.fetchone() return tuple(row) if row else None except Exception as exc: logger.error("Ошибка при получении данных пользователя: %s", exc) return None def save_user_data( user_id: int, chat_id: int, count: int | None = None, start_time: float | None = None, ban_trigger: bool | None = None, ban_time: float | None = None, username: str | None = None, ) -> None: try: with closing(create_connection()) as conn: cursor = conn.cursor() existing = get_user_data(user_id, chat_id) if existing is None: cursor.execute( ''' INSERT INTO user_data (user_id, chat_id, count, start_time, ban_trigger, ban_time, username) VALUES (?, ?, ?, ?, ?, ?, ?) ''', ( user_id, chat_id, count or 0, start_time or 0, int(bool(ban_trigger)) if ban_trigger is not None else 0, ban_time or 0, username, ), ) else: updates = [] params = [] if count is not None: updates.append("count = ?") params.append(count) if start_time is not None: updates.append("start_time = ?") params.append(start_time) if ban_trigger is not None: updates.append("ban_trigger = ?") params.append(int(bool(ban_trigger))) if ban_time is not None: updates.append("ban_time = ?") params.append(ban_time) if username is not None: updates.append("username = ?") params.append(username) if updates: params.extend([user_id, chat_id]) cursor.execute( f"UPDATE user_data SET {', '.join(updates)} WHERE user_id = ? AND chat_id = ?", params, ) conn.commit() except Exception as exc: logger.error("Ошибка при сохранении данных пользователя: %s", exc) def get_user_data_ban(user_id: int, chat_id: int, command: str): try: with closing(create_connection()) as conn: cursor = conn.cursor() cursor.execute( "SELECT * FROM user_data_ban WHERE user_id = ? AND chat_id = ? AND command = ?", (user_id, chat_id, command), ) row = cursor.fetchone() return tuple(row) if row else None except Exception as exc: logger.error("Ошибка при получении данных о блокировке пользователя: %s", exc) return None def save_user_data_ban( user_id: int, chat_id: int, command: str, count: int | None = None, ban_time_inf: float | None = None, ) -> None: try: with closing(create_connection()) as conn: cursor = conn.cursor() existing = get_user_data_ban(user_id, chat_id, command) if existing is None: cursor.execute( ''' INSERT INTO user_data_ban (user_id, chat_id, command, count, ban_time_inf) VALUES (?, ?, ?, ?, ?) ''', (user_id, chat_id, command, count or 0, ban_time_inf or 0), ) else: updates = [] params = [] if count is not None: updates.append("count = ?") params.append(count) if ban_time_inf is not None: updates.append("ban_time_inf = ?") params.append(ban_time_inf) if updates: params.extend([user_id, chat_id, command]) cursor.execute( f"UPDATE user_data_ban SET {', '.join(updates)} WHERE user_id = ? AND chat_id = ? AND command = ?", params, ) conn.commit() except Exception as exc: logger.error("Ошибка при сохранении данных о блокировке пользователя: %s", exc) create_tables()