bot-telegram/database/database.py

"""
Модуль для подключения к базе данных PostgreSQL через SQLAlchemy
"""

from sqlalchemy import create_engine, MetaData
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from contextlib import contextmanager
import os
from typing import Generator
import logging

# Настройка логирования
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Базовый класс для всех моделей
Base = declarative_base()

class DatabaseConfig:
    """Конфигурация базы данных"""
    
    def __init__(self):
        # Параметры подключения к PostgreSQL
        self.DB_USER = os.getenv('DB_USER', 'postgres')
        self.DB_PASSWORD = os.getenv('DB_PASSWORD', 'password')
        self.DB_HOST = os.getenv('DB_HOST', 'localhost')
        self.DB_PORT = os.getenv('DB_PORT', '5432')
        self.DB_NAME = os.getenv('DB_NAME', 'telegram_bot')
        
        # Строка подключения
        self.DATABASE_URL = f"postgresql://{self.DB_USER}:{self.DB_PASSWORD}@{self.DB_HOST}:{self.DB_PORT}/{self.DB_NAME}"
        
        # Настройки SQLAlchemy
        self.SQLALCHEMY_ENGINE_OPTIONS = {
            'pool_pre_ping': True,  # Проверка соединения перед использованием
            'pool_recycle': 300,    # Переподключение каждые 5 минут
            'pool_size': 10,        # Размер пула соединений
            'max_overflow': 20,     # Максимальное количество дополнительных соединений
            'echo': False,          # Логирование SQL запросов (для отладки)
        }

class Database:
    """Класс для управления подключением к базе данных"""
    
    def __init__(self, config: DatabaseConfig = None):
        self.config = config or DatabaseConfig()
        self.engine = None
        self.SessionLocal = None
        self._initialize()
    
    def _initialize(self):
        """Инициализация подключения к базе данных"""
        try:
            # Создание движка SQLAlchemy
            self.engine = create_engine(
                self.config.DATABASE_URL,
                **self.config.SQLALCHEMY_ENGINE_OPTIONS
            )
            
            # Создание фабрики сессий
            self.SessionLocal = sessionmaker(
                autocommit=False,
                autoflush=False,
                bind=self.engine
            )
            
            logger.info("Подключение к базе данных успешно инициализировано")
            
        except Exception as e:
            logger.error(f"Ошибка при инициализации базы данных: {e}")
            raise
    
    def create_tables(self):
        """Создание всех таблиц в базе данных"""
        try:
            Base.metadata.create_all(bind=self.engine)
            logger.info("Таблицы успешно созданы")
        except Exception as e:
            logger.error(f"Ошибка при создании таблиц: {e}")
            raise
    
    def drop_tables(self):
        """Удаление всех таблиц из базы данных"""
        try:
            Base.metadata.drop_all(bind=self.engine)
            logger.info("Таблицы успешно удалены")
        except Exception as e:
            logger.error(f"Ошибка при удалении таблиц: {e}")
            raise
    
    @contextmanager
    def get_session(self) -> Generator:
        """Контекстный менеджер для получения сессии базы данных"""
        session = self.SessionLocal()
        try:
            yield session
            session.commit()
        except Exception as e:
            session.rollback()
            logger.error(f"Ошибка в сессии базы данных: {e}")
            raise
        finally:
            session.close()
    
    def get_session_sync(self):
        """Получение синхронной сессии (не забудьте закрыть!)"""
        return self.SessionLocal()
    
    def test_connection(self) -> bool:
        """Проверка подключения к базе данных"""
        try:
            with self.get_session() as session:
                session.execute("SELECT 1")
            logger.info("Подключение к базе данных работает корректно")
            return True
        except Exception as e:
            logger.error(f"Ошибка подключения к базе данных: {e}")
            return False
    
    def close(self):
        """Закрытие подключения к базе данных"""
        if self.engine:
            self.engine.dispose()
            logger.info("Подключение к базе данных закрыто")

# Глобальный экземпляр базы данных
db = Database()

# Удобные функции для использования в приложении
def get_db_session():
    """Получение сессии базы данных для использования в приложении"""
    return db.get_session()

def init_database():
    """Инициализация базы данных (создание таблиц)"""
    db.create_tables()

def test_db_connection():
    """Тестирование подключения к базе данных"""
    return db.test_connection()

def load_image_from_file(file_path: str, alt_text: str = None):
    """
    Быстрая загрузка изображения из файла в базу данных
    
    Args:
        file_path: Путь к файлу изображения
        alt_text: Альтернативный текст
        
    Returns:
        Image: Объект изображения или None при ошибке
    """
    import os
    import mimetypes
    from database.repositories import ImageRepository
    
    try:
        if not os.path.isfile(file_path):
            return None
        
        filename = os.path.basename(file_path)
        mime_type, _ = mimetypes.guess_type(filename)
        if not mime_type:
            mime_type = 'image/jpeg'
        
        with open(file_path, 'rb') as f:
            image_data = f.read()
        
        with get_db_session() as session:
            image = ImageRepository.create_image(
                session, filename, mime_type, image_data, alt_text
            )
            session.commit()
            return image
            
    except Exception as e:
        logger.error(f"Ошибка загрузки изображения {file_path}: {e}")
        return None