refactor: update response object schema

apis/s3.py

@@ -1,11 +1,17 @@
 import os
 import boto3
+import botocore
+import logging
 from dotenv import load_dotenv
+
+from demo.qccamapi.models.response import Response
 load_dotenv()
 
 class S3:
 
-    def __init__(self) -> None:
+    def __init__(self, custom_logging=None) -> None:
+        self.logging = custom_logging if custom_logging is not None else logging
+
         self.aws_access_key_id = os.getenv('AWS_ACCESS_KEY_ID', '')
         self.aws_secret_access_key = os.getenv('AWS_SECRET_ACCESS_KEY', '')
         self.aws_default_region = os.getenv('AWS_DEFAULT_REGION', '')
@@ -28,7 +34,20 @@ class S3:
 
     # Function to upload file to S3
     def upload_file(self, bucket_name, file, name):
-        return self.client.upload_fileobj(file, bucket_name, name)
+        res = Response()
+        try:
+            self.client.upload_fileobj(file, bucket_name, name)
+            res.status = 200
+        except botocore.exceptions.ClientError as error:
+            res.status = 429
+            if error.response['Error']['Code'] == 'LimitExceededException':
+                res.error = ValueError('API call limit exceeded; backing off and retrying...')
+            else:
+                res.error = error
+        except botocore.exceptions.ParamValidationError as error:
+            res.status = 400
+            res.error = ValueError('The parameters you provided are incorrect: {}'.format(error))
+        return res
 
     # Function to delete a file from S3
     def delete_image(self, bucket_name, file_key):

app.py

@@ -4,6 +4,7 @@ import logging
 from fastapi import FastAPI, Header
 from apis.s3 import S3
 from dotenv import load_dotenv
+from demo.qccamapi.models.response import Response
 from models.image_object import ImageObject
 from apis.common import Common
 from fastapi.openapi.docs import (
@@ -18,15 +19,17 @@ AWS_S3_BUCKET_NAME = os.getenv('AWS_S3_BUCKET_NAME', '')
 X_REQUEST_USER = os.environ.get('X_REQUEST_USER')
 X_API_KEY = os.environ.get('X_API_KEY')
 
-s3client = S3()
-common = Common()
-app = FastAPI(docs_url=None, redoc_url=None)
 logging.basicConfig(
     level=logging.DEBUG if IS_DEV else logging.INFO,
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
     datefmt='%Y-%m-%d %H:%M:%S'
 )
 
+s3client = S3(logging)
+common = Common()
+app = FastAPI(docs_url=None, redoc_url=None)
+
+
 @app.get("/docs", include_in_schema=False)
 async def custom_swagger_ui_html():
     return get_swagger_ui_html(
@@ -52,11 +55,11 @@ async def redoc_html():
 
 @app.get("/")
 def root():
-    return {"message": "ok"}
+    return Response(status=200, data='ok')
         
 @app.get("/health")
 def healthcheck():
-    return {"message": "ok"}
+    return Response(status=200, data='ok')
         
 @app.post("/image")
 async def upload(
@@ -64,7 +67,7 @@ async def upload(
     x_request_user: str = Header(...),
     x_api_key: str = Header(...)
 ):
-    res = None
+    res = Response()
     logging.info("--------------------------------")
     logging.info("Received request to upload image")
     
@@ -86,11 +89,11 @@ async def upload(
                 logging.info("File uploaded successfully")
 
             except Exception as e:
-                res = str(e)
-                logging.warning(f"Error during upload: {res}")
+                res.error = str(e)
+                logging.warning(f"Error during upload: {res.error}")
                 return {"error": "Failed to upload content"}
     
-    return {"result": res}
+    return res
 
 def is_valid(u, p):
     return u == X_REQUEST_USER and p == X_API_KEY

models/response.py

@@ -0,0 +1,6 @@
+from pydantic import BaseModel
+
+class Response(BaseModel):
+    status: int = 500
+    data: any = None
+    error: any = None

requirements.txt

@@ -1,4 +1,5 @@
 boto3
+botocore
 gradio
 pillow
 pydantic