app.py

import os
from typing import Annotated
from fastapi import FastAPI, Header
from apis.s3 import S3
from dotenv import load_dotenv
from models.image_object import ImageObject
from models.common_headers import CommonHeaders
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', '')
USERNAME = os.environ.get('USERNAME')
PASSWORD = os.environ.get('PASSWORD')

s3client = S3()
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(
        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")
def upload(
    o: ImageObject,
    headers: Annotated[CommonHeaders, Header()]
):
    res = None
    if headers.x_request_user == USERNAME \
       and headers.x_api_key == PASSWORD:
        key = f'{o.key}/{o.job_no}/{o.name}'
        res = s3client.upload_file(AWS_S3_BUCKET_NAME, o.content, key)
    return res

# # Create a Gradio interface
# with gr.Blocks() as demo:
#     pass
   
# def auth(username, password):
#     u = os.environ.get('USERNAME')
#     p = os.environ.get('PASSWORD')
#     return (username == u and password == p)
 
# # Launch the Gradio app
# demo.launch(auth=auth if not IS_DEV else None)