app.py

import os, io, base64
import uvicorn
import logging
from fastapi import FastAPI, Header
from apis.s3 import S3
from dotenv import load_dotenv
from models.image_object import ImageObject
from apis.common import Common
from fastapi.openapi.docs import (
    get_redoc_html,
    get_swagger_ui_html,
    get_swagger_ui_oauth2_redirect_html,
)

load_dotenv()
IS_DEV = os.environ.get('ENV', 'DEV') != 'PROD'
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'
)

@app.get("/docs", include_in_schema=False)
async def custom_swagger_ui_html():
    return get_swagger_ui_html(
        openapi_url=app.openapi_url,
        title=app.title + " - Swagger UI",
        oauth2_redirect_url=app.swagger_ui_oauth2_redirect_url,
        swagger_js_url="https://unpkg.com/swagger-ui-dist@5/swagger-ui-bundle.js",
        swagger_css_url="https://unpkg.com/swagger-ui-dist@5/swagger-ui.css",
    )


@app.get(app.swagger_ui_oauth2_redirect_url, include_in_schema=False)
async def swagger_ui_redirect():
    return get_swagger_ui_oauth2_redirect_html()

@app.get("/redoc", include_in_schema=False)
async def redoc_html():
    return get_redoc_html(
        openapi_url=app.openapi_url,
        title=app.title + " - ReDoc",
        redoc_js_url="https://unpkg.com/redoc@next/bundles/redoc.standalone.js",
    )

@app.get("/")
def root():
    return {"message": "ok"}
        
@app.get("/health")
def healthcheck():
    return {"message": "ok"}
        
@app.post("/image")
async def upload(
    o: ImageObject,
    x_request_user: str = Header(...),
    x_api_key: str = Header(...)
):
    res = None
    logging.info("--------------------------------")
    logging.info("Received request to upload image")
    if is_valid(x_request_user, x_api_key):
        key = f'{o.key}/{o.job_no}/{o.name}'
        logging.info(f'Key for S3 upload: {key}')
        if o.content is not None:
            try:
                # Decode base64 content
                decoded_content = base64.b64decode(o.content)
                logging.info(f'Decoded content length: {len(decoded_content)} bytes')

                # Wrap bytes in BytesIO to create a file-like object
                file_obj = io.BytesIO(decoded_content)

                # Upload file object to S3
                # logging.info(f"Uploading file to S3 bucket '{AWS_S3_BUCKET_NAME}' with key '{key}'")
                res = s3client.upload_file(AWS_S3_BUCKET_NAME, file_obj, key)
                logging.info("File uploaded successfully")

            except Exception as e:
                res = str(e)
                logging.warning(f"Error during upload: {res}")
                return {"error": "Failed to upload content"}
    
    return {"result": res}

def is_valid(u, p):
    return u == X_REQUEST_USER and p == X_API_KEY

if __name__=='__main__':
    uvicorn.run('app:app', host='0.0.0.0', port=7860, reload=True)