backend/app/storage.py - develop - Ru5Kasper/Trackly

backend/app/storage.py
171 строка6 KB
DORANDKOR
feat(backend): create migrations, add support for user avatars
31 мар 2026, 07:30
31 мар 2026, 07:30b6b9ba7
import io
from pathlib import Path
from typing import BinaryIO
from uuid import uuid4

from minio import Minio
from minio.error import S3Error
from PIL import Image
from fastapi import HTTPException, UploadFile

from app.config import settings


class MinIOStorage:
    def __init__(self):
        self.client = Minio(
            settings.MINIO_ENDPOINT,
            access_key=settings.MINIO_ACCESS_KEY,
            secret_key=settings.MINIO_SECRET_KEY.get_secret_value(),
            secure=settings.MINIO_SECURE,
        )
        self.bucket_name = settings.MINIO_BUCKET_NAME
        self._ensure_bucket_exists()

    def _ensure_bucket_exists(self):
        """Create bucket if it doesn't exist"""
        try:
            if not self.client.bucket_exists(self.bucket_name):
                self.client.make_bucket(self.bucket_name)
                # Set bucket policy to allow public read access for avatars
                policy = {
                    "Version": "2012-10-17",
                    "Statement": [
                        {
                            "Effect": "Allow",
                            "Principal": {"AWS": "*"},
                            "Action": ["s3:GetObject"],
                            "Resource": [f"arn:aws:s3:::{self.bucket_name}/*"],
                        }
                    ],
                }
                import json

                self.client.set_bucket_policy(self.bucket_name, json.dumps(policy))
        except Exception as e:
            # Avoid crashing app startup if MinIO is not ready yet.
            print(f"Error ensuring bucket exists: {e}")

    async def validate_image(self, file: UploadFile) -> None:
        """Validate that the uploaded file is a valid image"""
        # Check file extension
        allowed_extensions = {".jpg", ".jpeg", ".png", ".gif", ".webp"}
        file_ext = Path(file.filename or "").suffix.lower()

        if file_ext not in allowed_extensions:
            raise HTTPException(
                status_code=400,
                detail=f"Invalid file type. Allowed types: {', '.join(allowed_extensions)}",
            )

        # Check file size (max 5MB)
        max_size = 5 * 1024 * 1024  # 5MB in bytes
        file.file.seek(0, 2)  # Seek to end
        file_size = file.file.tell()
        file.file.seek(0)  # Reset to beginning

        if file_size > max_size:
            raise HTTPException(
                status_code=400,
                detail=f"File too large. Maximum size is 5MB",
            )

        # Validate it's actually an image by trying to open it
        try:
            contents = await file.read()
            image = Image.open(io.BytesIO(contents))
            image.verify()
            # Reset file pointer
            await file.seek(0)
        except Exception as e:
            raise HTTPException(
                status_code=400,
                detail="Invalid image file",
            )

    async def process_and_upload_avatar(self, file: UploadFile, user_id: int) -> str:
        """
        Process and upload an avatar image
        Returns the URL to access the avatar
        """
        await self.validate_image(file)

        # Read and process the image
        contents = await file.read()
        image = Image.open(io.BytesIO(contents))

        # Convert to RGB if necessary (for PNG with transparency)
        if image.mode in ("RGBA", "LA", "P"):
            background = Image.new("RGB", image.size, (255, 255, 255))
            if image.mode == "P":
                image = image.convert("RGBA")
            background.paste(
                image, mask=image.split()[-1] if image.mode == "RGBA" else None
            )
            image = background

        # Resize to a reasonable size (max 500x500, maintain aspect ratio)
        max_size = (500, 500)
        image.thumbnail(max_size, Image.Resampling.LANCZOS)

        # Save to bytes
        output = io.BytesIO()
        image.save(output, format="JPEG", quality=85, optimize=True)
        output.seek(0)

        # Generate unique filename
        file_ext = ".jpg"  # Always save as JPEG after processing
        filename = f"avatars/user_{user_id}_{uuid4().hex}{file_ext}"

        # Upload to MinIO
        try:
            self.client.put_object(
                self.bucket_name,
                filename,
                output,
                length=output.getbuffer().nbytes,
                content_type="image/jpeg",
            )
        except S3Error as e:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to upload avatar: {str(e)}",
            )

        # Return the URL
        return self._get_object_url(filename)

    def _get_object_url(self, object_name: str) -> str:
        """Generate URL for accessing an object"""
        # For public access, construct the URL directly
        protocol = "https" if settings.MINIO_SECURE else "http"
        return (
            f"{protocol}://{settings.MINIO_ENDPOINT}/{self.bucket_name}/{object_name}"
        )

    def delete_avatar(self, avatar_url: str) -> None:
        """Delete an avatar from storage"""
        try:
            # Extract object name from URL
            object_name = avatar_url.split(f"{self.bucket_name}/")[-1]
            self.client.remove_object(self.bucket_name, object_name)
        except S3Error as e:
            print(f"Error deleting avatar: {e}")
            # Don't raise exception, just log it

    def get_avatar_stream(self, avatar_url: str) -> BinaryIO:
        """Get avatar file stream for downloading"""
        try:
            # Extract object name from URL
            object_name = avatar_url.split(f"{self.bucket_name}/")[-1]
            response = self.client.get_object(self.bucket_name, object_name)
            return response
        except S3Error as e:
            raise HTTPException(
                status_code=404,
                detail="Avatar not found",
            )


# Singleton instance
minio_storage = MinIOStorage()