webref/backend/app/core/storage.py

"""MinIO storage client utilities."""

import logging
from io import BytesIO
from typing import BinaryIO

import boto3
from botocore.client import Config
from botocore.exceptions import ClientError

from app.core.config import settings

logger = logging.getLogger(__name__)


class StorageClient:
    """MinIO storage client wrapper."""

    def __init__(self):
        """Initialize MinIO client."""
        self.client = boto3.client(
            "s3",
            endpoint_url=f"{'https' if settings.MINIO_SECURE else 'http'}://{settings.MINIO_ENDPOINT}",
            aws_access_key_id=settings.MINIO_ACCESS_KEY,
            aws_secret_access_key=settings.MINIO_SECRET_KEY,
            config=Config(signature_version="s3v4"),
        )
        self.bucket = settings.MINIO_BUCKET
        self._ensure_bucket_exists()

    def put_object(self, bucket_name: str, object_name: str, data: BinaryIO, length: int, content_type: str):
        """MinIO-compatible put_object method."""
        return self.upload_file(data, object_name, content_type)

    def remove_object(self, bucket_name: str, object_name: str):
        """MinIO-compatible remove_object method."""
        return self.delete_file(object_name)

    def _ensure_bucket_exists(self) -> None:
        """Create bucket if it doesn't exist."""
        try:
            self.client.head_bucket(Bucket=self.bucket)
        except ClientError:
            logger.info(f"Creating bucket: {self.bucket}")
            self.client.create_bucket(Bucket=self.bucket)

    def upload_file(self, file_data: BinaryIO, object_name: str, content_type: str) -> str:
        """Upload file to MinIO.

        Args:
            file_data: File data to upload
            object_name: S3 object name (path)
            content_type: MIME type of the file

        Returns:
            str: Object URL

        Raises:
            Exception: If upload fails
        """
        try:
            self.client.upload_fileobj(
                file_data,
                self.bucket,
                object_name,
                ExtraArgs={"ContentType": content_type},
            )
            return f"{settings.MINIO_ENDPOINT}/{self.bucket}/{object_name}"
        except ClientError as e:
            logger.error(f"Failed to upload file {object_name}: {e}")
            raise

    def download_file(self, object_name: str) -> BytesIO:
        """Download file from MinIO.

        Args:
            object_name: S3 object name (path)

        Returns:
            BytesIO: File data

        Raises:
            Exception: If download fails
        """
        try:
            file_data = BytesIO()
            self.client.download_fileobj(self.bucket, object_name, file_data)
            file_data.seek(0)
            return file_data
        except ClientError as e:
            logger.error(f"Failed to download file {object_name}: {e}")
            raise

    def get_object(self, object_name: str) -> bytes | None:
        """Get object as bytes from MinIO.

        Args:
            object_name: S3 object name (path)

        Returns:
            bytes: File data or None if not found

        Raises:
            Exception: If download fails for reasons other than not found
        """
        try:
            file_data = self.download_file(object_name)
            return file_data.read()
        except ClientError as e:
            if e.response["Error"]["Code"] == "404":
                return None
            logger.error(f"Failed to get object {object_name}: {e}")
            raise

    def delete_file(self, object_name: str) -> None:
        """Delete file from MinIO.

        Args:
            object_name: S3 object name (path)

        Raises:
            Exception: If deletion fails
        """
        try:
            self.client.delete_object(Bucket=self.bucket, Key=object_name)
        except ClientError as e:
            logger.error(f"Failed to delete file {object_name}: {e}")
            raise

    def file_exists(self, object_name: str) -> bool:
        """Check if file exists in MinIO.

        Args:
            object_name: S3 object name (path)

        Returns:
            bool: True if file exists, False otherwise
        """
        try:
            self.client.head_object(Bucket=self.bucket, Key=object_name)
            return True
        except ClientError:
            return False


# Global storage client instance
storage_client = StorageClient()


def get_storage_client() -> StorageClient:
    """Get the global storage client instance."""
    return storage_client


# Compatibility methods for MinIO-style API
def put_object(bucket_name: str, object_name: str, data: BinaryIO, length: int, content_type: str):
    """MinIO-compatible put_object method."""
    storage_client.upload_file(data, object_name, content_type)


def remove_object(bucket_name: str, object_name: str):
    """MinIO-compatible remove_object method."""
    storage_client.delete_file(object_name)