mp_message/main.py

from datetime import datetime, timezone

import uvicorn
from fastapi import FastAPI, WebSocket
from starlette.websockets import WebSocketDisconnect

from core.ws.handlers import connection_manager

from loguru import logger
from starlette.middleware.cors import CORSMiddleware

from config import Config

from core.api import router as api_router
from core.errors.errors import CustomExceptionHandler

app = FastAPI()

origins = [
    "*"
]

app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Including routers

app.include_router(api_router)


@app.on_event('startup')
async def on_startup():
    logger.success('Application startup complete at {time}', time=datetime.now(tz=timezone.utc))


@app.exception_handler(CustomExceptionHandler)
async def custom_exception_handler(_, exc):
    return exc.result()


@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await connection_manager.connect(websocket)
    try:
        while True:
            data = await websocket.receive_text()
            await connection_manager.route_message(data, websocket)
    except WebSocketDisconnect:
        connection_manager.disconnect(websocket)

if __name__ == '__main__':
    uvicorn.run('main:app', host=str(Config.host), port=Config.port, reload=True)