telegram-helper-bot/helper_bot/utils/s3_storage.py

"""
Сервис для работы с S3 хранилищем.
"""

import os
import tempfile
from pathlib import Path
from typing import Optional

import aioboto3

from logs.custom_logger import logger


class S3StorageService:
    """Сервис для работы с S3 хранилищем."""

    def __init__(
        self,
        endpoint_url: str,
        access_key: str,
        secret_key: str,
        bucket_name: str,
        region: str = "us-east-1",
    ):
        self.endpoint_url = endpoint_url
        self.access_key = access_key
        self.secret_key = secret_key
        self.bucket_name = bucket_name
        self.region = region
        self.session = aioboto3.Session()

    async def upload_file(
        self, file_path: str, s3_key: str, content_type: Optional[str] = None
    ) -> bool:
        """Загружает файл в S3."""
        try:
            async with self.session.client(
                "s3",
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
            ) as s3:
                extra_args = {}
                if content_type:
                    extra_args["ContentType"] = content_type

                await s3.upload_file(
                    file_path, self.bucket_name, s3_key, ExtraArgs=extra_args
                )
                logger.info(f"Файл загружен в S3: {s3_key}")
                return True
        except Exception as e:
            logger.error(f"Ошибка загрузки файла в S3 {s3_key}: {e}")
            return False

    async def upload_fileobj(
        self, file_obj, s3_key: str, content_type: Optional[str] = None
    ) -> bool:
        """Загружает файл из объекта в S3."""
        try:
            async with self.session.client(
                "s3",
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
            ) as s3:
                extra_args = {}
                if content_type:
                    extra_args["ContentType"] = content_type

                await s3.upload_fileobj(
                    file_obj, self.bucket_name, s3_key, ExtraArgs=extra_args
                )
                logger.info(f"Файл загружен в S3 из объекта: {s3_key}")
                return True
        except Exception as e:
            logger.error(f"Ошибка загрузки файла в S3 из объекта {s3_key}: {e}")
            return False

    async def download_file(self, s3_key: str, local_path: str) -> bool:
        """Скачивает файл из S3 на локальный диск."""
        try:
            async with self.session.client(
                "s3",
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
            ) as s3:
                # Создаем директорию если её нет
                os.makedirs(os.path.dirname(local_path), exist_ok=True)

                await s3.download_file(self.bucket_name, s3_key, local_path)
                logger.info(f"Файл скачан из S3: {s3_key} -> {local_path}")
                return True
        except Exception as e:
            logger.error(f"Ошибка скачивания файла из S3 {s3_key}: {e}")
            return False

    async def download_to_temp(self, s3_key: str) -> Optional[str]:
        """Скачивает файл из S3 во временный файл. Возвращает путь к временному файлу."""
        try:
            # Определяем расширение из ключа
            ext = Path(s3_key).suffix or ".bin"
            temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=ext)
            temp_path = temp_file.name
            temp_file.close()

            success = await self.download_file(s3_key, temp_path)
            if success:
                return temp_path
            else:
                # Удаляем временный файл при ошибке
                try:
                    os.remove(temp_path)
                except:
                    pass
                return None
        except Exception as e:
            logger.error(
                f"Ошибка скачивания файла из S3 во временный файл {s3_key}: {e}"
            )
            return None

    async def file_exists(self, s3_key: str) -> bool:
        """Проверяет существование файла в S3."""
        try:
            async with self.session.client(
                "s3",
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
            ) as s3:
                await s3.head_object(Bucket=self.bucket_name, Key=s3_key)
                return True
        except:
            return False

    async def delete_file(self, s3_key: str) -> bool:
        """Удаляет файл из S3."""
        try:
            async with self.session.client(
                "s3",
                endpoint_url=self.endpoint_url,
                aws_access_key_id=self.access_key,
                aws_secret_access_key=self.secret_key,
                region_name=self.region,
            ) as s3:
                await s3.delete_object(Bucket=self.bucket_name, Key=s3_key)
                logger.info(f"Файл удален из S3: {s3_key}")
                return True
        except Exception as e:
            logger.error(f"Ошибка удаления файла из S3 {s3_key}: {e}")
            return False

    def generate_s3_key(self, content_type: str, file_id: str) -> str:
        """Генерирует S3 ключ для файла. Один и тот же для всех постов с этим file_id."""
        type_folders = {
            "photo": "photos",
            "video": "videos",
            "audio": "music",
            "voice": "voice",
            "video_note": "video_notes",
        }

        folder = type_folders.get(content_type, "other")
        # Определяем расширение из file_id или используем дефолтное
        ext = (
            ".jpg"
            if content_type == "photo"
            else (
                ".mp4"
                if content_type == "video"
                else (
                    ".mp3"
                    if content_type == "audio"
                    else (
                        ".ogg"
                        if content_type == "voice"
                        else ".mp4" if content_type == "video_note" else ".bin"
                    )
                )
            )
        )

        return f"{folder}/{file_id}{ext}"