database/db.py - master - huyarim/TgBot

database/db.py
223 строки9 KB
anastasia_lesnikova
удаление файлов, переименование и счетчик анализов
20 янв 2026, 17:05
20 янв 2026, 17:0553e9c43
"""
Работа с базой данных MySQL.
Используем пул подключений, методы покрывают базовые операции с пользователями.
"""
import logging
from contextlib import contextmanager

import mysql.connector
from mysql.connector import Error as MySQLError
from mysql.connector.pooling import MySQLConnectionPool

logger = logging.getLogger(__name__)


class Database:
    """Класс для работы с БД"""

    def __init__(self, host: str, dbname: str, user: str, password: str, port: int = 3306):
        self.config = {
            "host": host,
            "database": dbname,  # MySQL использует 'database' вместо 'dbname'
            "user": user,
            "password": password,
            "port": port,
        }
        try:
            self.pool = MySQLConnectionPool(
                pool_name="tgbot_pool",
                pool_size=5,
                pool_reset_session=True,
                **self.config,
            )
            logger.info("Подключение к БД установлено (%s)", host)
        except MySQLError as exc:
            logger.error("Не удалось подключиться к БД: %s", exc)
            raise

    @contextmanager
    def _get_conn(self):
        conn = self.pool.get_connection()
        try:
            yield conn
        finally:
            conn.close()

    def init_db(self):
        """Создает необходимые таблицы, если их нет"""
        create_users = """
        CREATE TABLE IF NOT EXISTS users (
            telegram_id BIGINT PRIMARY KEY,
            username TEXT,
            first_name TEXT,
            last_name TEXT,
            full_name TEXT,
            language_code VARCHAR(10),
            is_bot BOOLEAN DEFAULT FALSE,
            is_premium BOOLEAN DEFAULT FALSE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
            updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
        );
        """
        create_user_data = """
        CREATE TABLE IF NOT EXISTS user_data (
            telegram_id BIGINT NOT NULL,
            data_key VARCHAR(255) NOT NULL,
            data_value TEXT,
            PRIMARY KEY (telegram_id, data_key)
        );
        """
        create_files = """
        CREATE TABLE IF NOT EXISTS files (
            id INT AUTO_INCREMENT PRIMARY KEY,
            telegram_id BIGINT,
            file_type TEXT,
            file_name TEXT,
            storage_path TEXT,
            size_bytes BIGINT,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
        """
        self._execute(create_users)
        self._execute(create_user_data)
        self._execute(create_files)
        logger.info("Таблицы users, user_data и files готовы")

    def _execute(self, query: str, params=None, fetchone=False, fetchall=False, return_lastrowid=False):
        params = params or ()
        with self._get_conn() as conn:
            try:
                cursor = conn.cursor(dictionary=True)  # MySQL возвращает результаты как словари
                cursor.execute(query, params)
                result = None
                if fetchone:
                    result = cursor.fetchone()
                elif fetchall:
                    result = cursor.fetchall()
                elif return_lastrowid:
                    result = cursor.lastrowid
                conn.commit()
                cursor.close()
                return result
            except Exception as exc:
                conn.rollback()
                logger.error("Ошибка при выполнении запроса: %s", exc, exc_info=True)
                raise

    def save_user(self, user_id, username=None, first_name=None, last_name=None, 
                  full_name=None, language_code=None, is_bot=False, is_premium=False):
        """Создает или обновляет запись пользователя"""
        query = """
        INSERT INTO users (telegram_id, username, first_name, last_name, full_name, 
                          language_code, is_bot, is_premium)
        VALUES (%s, %s, %s, %s, %s, %s, %s, %s)
        ON DUPLICATE KEY UPDATE
            username = VALUES(username),
            first_name = VALUES(first_name),
            last_name = VALUES(last_name),
            full_name = VALUES(full_name),
            language_code = VALUES(language_code),
            is_bot = VALUES(is_bot),
            is_premium = VALUES(is_premium);
        """
        self._execute(query, (user_id, username, first_name, last_name, full_name, 
                             language_code, is_bot, is_premium))
        logger.debug("Пользователь %s сохранен/обновлен", user_id)

    def get_user(self, user_id):
        """Возвращает словарь с данными пользователя или None"""
        query = "SELECT * FROM users WHERE telegram_id = %s;"
        return self._execute(query, (user_id,), fetchone=True)

    def save_user_data(self, user_id, key, value):
        """Сохраняет произвольную пару ключ/значение для пользователя"""
        query = """
        INSERT INTO user_data (telegram_id, data_key, data_value)
        VALUES (%s, %s, %s)
        ON DUPLICATE KEY UPDATE
            data_value = VALUES(data_value);
        """
        self._execute(query, (user_id, key, value))
        logger.debug("user_data %s:%s сохранено", user_id, key)

    def get_user_data(self, user_id, key):
        """Возвращает значение по ключу или None"""
        query = "SELECT data_value FROM user_data WHERE telegram_id = %s AND data_key = %s;"
        row = self._execute(query, (user_id, key), fetchone=True)
        return row["data_value"] if row else None

    def increment_user_data(self, user_id, key, step: int = 1):
        """Инкрементирует числовое значение user_data"""
        query = """
        INSERT INTO user_data (telegram_id, data_key, data_value)
        VALUES (%s, %s, %s)
        ON DUPLICATE KEY UPDATE
            data_value = CAST(data_value AS UNSIGNED) + VALUES(data_value);
        """
        self._execute(query, (user_id, key, str(step)))
        logger.debug("user_data %s:%s увеличено на %s", user_id, key, step)

    def close(self):
        """Закрывает пул подключений"""
        if self.pool:
            # MySQL connector pool не имеет метода closeall, но соединения закроются автоматически
            logger.info("Пул подключений к БД закрыт")

    def save_file_record(self, telegram_id, file_type, file_name, storage_path, size_bytes):
        """Сохраняет информацию о загруженном файле"""
        query = """
        INSERT INTO files (telegram_id, file_type, file_name, storage_path, size_bytes)
        VALUES (%s, %s, %s, %s, %s);
        """
        return self._execute(
            query,
            (telegram_id, file_type, file_name, storage_path, size_bytes),
            return_lastrowid=True,
        )
        logger.debug("Информация о файле %s сохранена", file_name)

    def update_file_record(self, file_id, telegram_id, file_name, storage_path):
        """Обновляет имя и путь файла"""
        query = """
        UPDATE files
        SET file_name = %s, storage_path = %s
        WHERE id = %s AND telegram_id = %s;
        """
        self._execute(query, (file_name, storage_path, file_id, telegram_id))
        logger.debug("Информация о файле %s обновлена", file_name)

    def get_user_files(self, telegram_id, file_type=None, limit=20):
        """Получает список файлов пользователя"""
        if file_type:
            query = """
            SELECT id, file_type, file_name, storage_path, size_bytes, created_at
            FROM files
            WHERE telegram_id = %s AND file_type = %s
            ORDER BY created_at DESC
            LIMIT %s;
            """
            return self._execute(query, (telegram_id, file_type, limit), fetchall=True)
        else:
            query = """
            SELECT id, file_type, file_name, storage_path, size_bytes, created_at
            FROM files
            WHERE telegram_id = %s
            ORDER BY created_at DESC
            LIMIT %s;
            """
            return self._execute(query, (telegram_id, limit), fetchall=True)
    
    def get_file_by_id(self, file_id, telegram_id):
        """Получает информацию о конкретном файле по ID"""
        query = """
        SELECT id, file_type, file_name, storage_path, size_bytes, created_at
        FROM files
        WHERE id = %s AND telegram_id = %s;
        """
        return self._execute(query, (file_id, telegram_id), fetchone=True)

    def delete_file_record(self, file_id, telegram_id):
        """Удаляет запись о файле"""
        query = "DELETE FROM files WHERE id = %s AND telegram_id = %s;"
        self._execute(query, (file_id, telegram_id))
        logger.debug("Файл %s удален из БД", file_id)