PySimpleSocial/main.py

import uvloop
import asyncio
import uvicorn

from fastapi import FastAPI, Request
from fastapi.exceptions import (
    HTTPException,
    RequestValidationError,
    StarletteHTTPException,
)
from slowapi.errors import RateLimitExceeded
from pathlib import Path


from endpoints import users, media
from config import (
    LOGGER,
    LIMITER,
    NotAuthenticated,
    SMTP_HOST,
    SMTP_USER,
    SMTP_PORT,
    SMTP_USE_TLS,
    SMTP_PASSWORD,
    SMTP_TIMEOUT,
    HOST,
    PORT,
    WORKERS,
    PLATFORM_NAME,
)
from orm import create_tables, Media
from util.exception_handlers import (
    http_exception,
    rate_limited,
    request_invalid,
    not_authenticated,
    generic_error,
)
from util.email import test_smtp


with (Path(__file__).parent / "API.md").resolve(strict=True).open() as f:
    description = f.read()


app = FastAPI(
    title=PLATFORM_NAME,
    license_info={"name": "MIT", "url": "https://opensource.org/licenses/MIT"},
    contact={
        "name": "Matt",
        "url": "https://nocturn9x.space",
        "email": "nocturn9x@nocturn9x.space",
    },
    version="0.0.1",
    description=description,
    openapi_tags=[
        {
            "name": "Miscellaneous",
            "description": "Simple endpoints that don't fit anywhere else",
        },
        {
            "name": "Users",
            "description": "Endpoints that handle user-related operations such as"
            " signing in, signing up, settings management and more",
        },
        {
            "name": "Posts",
            "description": "Endpoints that handle post creation, modification and deletion",
        },
        {
            "name": "Media",
            "description": "Endpoints that allow fetching individual media objects by their ID",
        },
    ],
)


@app.get("/", tags=["Miscellaneous"])
@LIMITER.limit("10/second")
async def root(request: Request):
    raise HTTPException(401, detail="Unauthorized")


@app.get("/ping", tags=["Miscellaneous"])
@LIMITER.limit("1/minute")
async def ping(request: Request) -> dict:
    """
    This method simply replies to "ping" requests and
    is used to check whether the API is up and running.
    It also performs a sanity check with the database and
    the SMTP server to ensure that they are functioning correctly.
    For this reason, this endpoint's rate limit is very strict:
    it can only be called once per minute
    """

    LOGGER.info(f"Processing ping request from {request.client.host}")
    try:
        await Media.raw("SELECT 1;")
        await test_smtp(
            SMTP_HOST,
            SMTP_PORT,
            SMTP_USER,
            SMTP_PASSWORD,
            SMTP_TIMEOUT,
            SMTP_USE_TLS,
            True,
        )
        return {"status_code": 200, "msg": "OK"}
    except Exception:
        raise HTTPException(500)


async def startup_checks():
    LOGGER.info("Initializing database")
    try:
        await create_tables()
        await Media.raw("SELECT 1;")
    except Exception as e:
        LOGGER.error(f"An error occurred while trying to initialize the database -> {type(e).__name__}: {e}")
    else:
        LOGGER.info("Database initialized")
    LOGGER.info("Testing SMTP connection")
    try:
        await test_smtp(
            SMTP_HOST,
            SMTP_PORT,
            SMTP_USER,
            SMTP_PASSWORD,
            SMTP_TIMEOUT,
            SMTP_USE_TLS,
            True,
        )
    except Exception as e:
        LOGGER.error(f"An error occurred while trying to connect to the SMTP server -> {type(e).__name__}: {e}")
    else:
        LOGGER.info("SMTP test was successful")


if __name__ == "__main__":
    LOGGER.info("Backend starting up!")
    LOGGER.debug("Including modules")
    app.include_router(users.router)
    app.include_router(media.router)
    app.state.limiter = LIMITER
    LOGGER.debug("Setting exception handlers")
    app.add_exception_handler(RateLimitExceeded, rate_limited)
    app.add_exception_handler(NotAuthenticated, not_authenticated)
    app.add_exception_handler(HTTPException, http_exception)
    app.add_exception_handler(StarletteHTTPException, http_exception)
    app.add_exception_handler(RequestValidationError, request_invalid)
    app.add_exception_handler(Exception, generic_error)
    LOGGER.debug("Installing uvloop")
    uvloop.install()
    log_config = uvicorn.config.LOGGING_CONFIG
    log_config["formatters"]["access"]["datefmt"] = LOGGER.handlers[0].formatter.datefmt
    log_config["formatters"]["default"]["datefmt"] = LOGGER.handlers[0].formatter.datefmt
    log_config["formatters"]["access"]["fmt"] = LOGGER.handlers[0].formatter._fmt
    log_config["formatters"]["default"]["fmt"] = LOGGER.handlers[0].formatter._fmt
    log_config["handlers"]["access"]["stream"] = "ext://sys.stderr"
    asyncio.run(startup_checks())
    uvicorn.run(host=HOST, port=PORT, app=app, log_config=log_config, workers=WORKERS)