Initial draft
This commit is contained in:
1 parent
3c8eb6c7e9
commit
2bbb9dff3c
2 files changed
+264
No files matched your search
+130
@@ -0,0 +1,130 @@
|
||||
from fastapi import FastAPI, Depends, HTTPException, Request, Response
|
||||
from fastapi.responses import JSONResponse
|
||||
from fastapi.security.api_key import APIKeyHeader
|
||||
from starlette.exceptions import HTTPException as StarletteHTTPException
|
||||
from typing import Dict
|
||||
from pydantic import BaseModel
|
||||
from validator import verify_api_key
|
||||
from db_module import get_db, create_connection_pool, close_connection_pool, start_healthcheck_thread
|
||||
from simple_logger_handler import setup_logger, LOG_LEVEL
|
||||
from rabbitmq_handler import RabbitMQProducer
|
||||
import uvicorn
|
||||
from contextlib import asynccontextmanager
|
||||
from metrics_server import REQUEST_COUNTER
|
||||
import asyncio
|
||||
from uvicorn_logger_config import LOGGING_CONFIG
|
||||
from secret_manager import cleanup_secret_files
|
||||
|
||||
logger = setup_logger(__name__)
|
||||
api_key_header_internal = APIKeyHeader(name="X-API-Key-Internal")
|
||||
|
||||
class Notification(BaseModel):
|
||||
receipent_user_id : int
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
logger.info("Starting application...")
|
||||
|
||||
logger.info("Creating MySQL connection pool...")
|
||||
create_connection_pool("ApiInternalPool")
|
||||
start_healthcheck_thread()
|
||||
logger.info("MySQL healthcheck thread started.")
|
||||
|
||||
producer = RabbitMQProducer()
|
||||
|
||||
try:
|
||||
logger.info("Starting RabbitMQ producer...")
|
||||
await producer.connect()
|
||||
app.state.rmq_producer = producer
|
||||
logger.info("[FastAPI] RabbitMQ producer initialized.")
|
||||
|
||||
SECRET_PATHS = frozenset({
|
||||
"/etc/secrets/db_username",
|
||||
"/etc/secrets/db_password",
|
||||
"/etc/secrets/rmq_username",
|
||||
"/etc/secrets/rmq_password"
|
||||
})
|
||||
cleanup_secret_files(SECRET_PATHS)
|
||||
|
||||
yield
|
||||
finally:
|
||||
logger.info("Closing RabbitMQ producer...")
|
||||
await producer.close()
|
||||
|
||||
logger.info("Closing MySQL connection pool...")
|
||||
close_connection_pool()
|
||||
|
||||
api = FastAPI(
|
||||
title="Internal Notifier API",
|
||||
description="API to forward messages to RabbitMQ",
|
||||
version="1.0.0",
|
||||
lifespan=lifespan
|
||||
)
|
||||
|
||||
@api.middleware("http")
|
||||
async def prometheus_middleware(request: Request, call_next):
|
||||
logger.debug(f"Incoming request: {request.method} {request.url.path}")
|
||||
status = 500
|
||||
try:
|
||||
response = await call_next(request)
|
||||
status = response.status_code
|
||||
logger.debug(f"Request processed with status: {status}")
|
||||
except Exception as e:
|
||||
logger.exception(f"Exception during request processing: {e}")
|
||||
raise
|
||||
finally:
|
||||
REQUEST_COUNTER.labels(request.method, request.url.path, status).inc()
|
||||
logger.debug(f"Prometheus counter incremented: {request.method} {request.url.path} {status}")
|
||||
return response
|
||||
|
||||
def verify_api_key_dependency_internal(db=Depends(get_db), api_key: str = Depends(api_key_header_internal)) -> str:
|
||||
logger.debug(f"Verifying internal API key: {api_key}")
|
||||
cursor = db.cursor()
|
||||
cursor.execute("SELECT program_name, api_key FROM internal_api_keys WHERE status = 'active'")
|
||||
for program_name, hashed_key in cursor.fetchall():
|
||||
if verify_api_key(api_key=api_key, hashed=hashed_key):
|
||||
logger.debug(f"API key verified for program: {program_name}")
|
||||
return program_name
|
||||
logger.warning("Internal API key verification failed")
|
||||
raise HTTPException(status_code=403, detail="Unauthorized")
|
||||
|
||||
@api.exception_handler(StarletteHTTPException)
|
||||
async def custom_http_exception_handler(request,exc):
|
||||
logger.debug(f"Handling HTTP exception: {exc}")
|
||||
if exc.status_code == 404:
|
||||
return JSONResponse(
|
||||
status_code=403,
|
||||
content={"detail": "Unauthorized"}
|
||||
)
|
||||
return JSONResponse(
|
||||
status_code=exc.status_code,
|
||||
content={"detail": exc.detail}
|
||||
)
|
||||
|
||||
@api.get("/notifications/pull")
|
||||
async def receive_notifications(
|
||||
notification_data: Notification,
|
||||
request: Request,
|
||||
db = Depends(get_db),
|
||||
program_name: str = Depends(verify_api_key_dependency_internal)
|
||||
):
|
||||
|
||||
logger.info(f"Received notifcation data from {program_name} for RMQ")
|
||||
|
||||
messages = await request.app.state.rmq_producer.fetch_messages()
|
||||
logger.info("Successfully retrieved messages for ")
|
||||
print (messages)
|
||||
|
||||
async def start_servers():
|
||||
logger.debug("Starting FastAPI and metrics servers")
|
||||
config_main = uvicorn.Config("main:api", host="0.0.0.0", port=8101, log_level=LOG_LEVEL.lower(), log_config=LOGGING_CONFIG)
|
||||
config_metrics = uvicorn.Config("metrics_server:metrics_api", host="0.0.0.0", port=9000, log_level=LOG_LEVEL.lower(), log_config=LOGGING_CONFIG)
|
||||
|
||||
server_main = uvicorn.Server(config_main)
|
||||
server_metrics = uvicorn.Server(config_metrics)
|
||||
|
||||
await asyncio.gather(server_main.serve(), server_metrics.serve())
|
||||
logger.debug("Servers started")
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(start_servers())
|
||||
Reference in new issue
Block a user