internal.block.management

  1"""Block storage management module."""
  2
  3import asyncio
  4from io import BytesIO
  5
  6from fastapi import HTTPException, UploadFile, status
  7from minio import Minio, S3Error
  8from PIL import Image
  9
 10from internal.settings.env import block_settings
 11
 12PROFILE_IMAGES_BUCKET = "profileimages"
 13BUNDLE_IMAGES_BUCKET = "bundleimages"
 14MAX_SIZE = 5 * 1024 * 1024  # 5mb
 15
 16
 17async def process_image(file: UploadFile) -> bytes:
 18    """Confirm image and compress.
 19
 20    Args:
 21        file: uploaded file
 22
 23    Returns:
 24        compressed image
 25
 26    Raises:
 27        HTTPException: if not jpeg or larger then 5mb
 28    """
 29    if file.content_type != "image/jpeg":
 30        raise HTTPException(
 31            status.HTTP_400_BAD_REQUEST, "only jpeg images are accepted"
 32        )
 33    file_bytes = await file.read()
 34    if len(file_bytes) > MAX_SIZE:
 35        raise HTTPException(status_code=400, detail="image must be under 5mb")
 36    image = Image.open(BytesIO(file_bytes))
 37    output = BytesIO()
 38    image.save(output, format="JPEG", quality=85, optimize=True)
 39    return output.getvalue()
 40
 41
 42class BlockManagement:
 43    """Block storage management class."""
 44
 45    client: Minio
 46
 47    def initialise(self) -> None:
 48        """Initialise block storage connection."""
 49        self.client = Minio(
 50            endpoint=block_settings.url_port,
 51            access_key=block_settings.access_key,
 52            secret_key=block_settings.secret_key,
 53            secure=False,
 54        )
 55        if not self.client.bucket_exists(PROFILE_IMAGES_BUCKET):
 56            self.client.make_bucket(PROFILE_IMAGES_BUCKET)
 57        if not self.client.bucket_exists(BUNDLE_IMAGES_BUCKET):
 58            self.client.make_bucket(BUNDLE_IMAGES_BUCKET)
 59
 60    async def upload_profile_image(self, user_id: int, file: UploadFile) -> None:
 61        """Change user profile image.
 62
 63        Args:
 64            user_id: user id
 65            file: uploaded file
 66        """
 67        image_bytes = await process_image(file)
 68        await asyncio.to_thread(
 69            self.client.put_object,
 70            bucket_name=PROFILE_IMAGES_BUCKET,
 71            object_name=f"{user_id}.jpeg",
 72            data=BytesIO(image_bytes),
 73            length=len(image_bytes),
 74            content_type="image/jpeg",
 75        )
 76
 77    async def upload_bundle_image(self, bundle_id: int, file: UploadFile) -> None:
 78        """Change bundle image.
 79
 80        Args:
 81            bundle_id: bundle id
 82            file: uploaded file
 83        """
 84        image_bytes = await process_image(file)
 85        await asyncio.to_thread(
 86            self.client.put_object,
 87            bucket_name=BUNDLE_IMAGES_BUCKET,
 88            object_name=f"{bundle_id}.jpeg",
 89            data=BytesIO(image_bytes),
 90            length=len(image_bytes),
 91            content_type="image/jpeg",
 92        )
 93
 94    def get_profile_image(self, user_id: int) -> bytes:
 95        """Get user profile image.
 96
 97        Args:
 98            user_id: user id
 99
100        Returns:
101            image bytes
102
103        Raises:
104            HTTPException: if failed to get image
105        """
106        try:
107            image = self.client.get_object(PROFILE_IMAGES_BUCKET, f"{user_id}.jpeg")
108            try:
109                return image.read()
110            finally:
111                image.close()
112                image.release_conn()
113        except S3Error as err:
114            if err.code == "NoSuchKey":
115                raise HTTPException(status.HTTP_404_NOT_FOUND, "image not found")
116            raise HTTPException(
117                status.HTTP_500_INTERNAL_SERVER_ERROR, "failed to get image"
118            )
119
120    def get_bundle_image(self, bundle_id: int) -> bytes:
121        """Get bundle image.
122
123        Args:
124            bundle_id: bundle id
125
126        Returns:
127            image bytes
128
129        Raises:
130            HTTPException: if failed to get image
131        """
132        try:
133            image = self.client.get_object(BUNDLE_IMAGES_BUCKET, f"{bundle_id}.jpeg")
134            try:
135                return image.read()
136            finally:
137                image.close()
138                image.release_conn()
139        except S3Error as err:
140            if err.code == "NoSuchKey":
141                raise HTTPException(status.HTTP_404_NOT_FOUND, "image not found")
142            raise HTTPException(
143                status.HTTP_500_INTERNAL_SERVER_ERROR, "failed to get image"
144            )
145
146
147block_management = BlockManagement()