- Updated path to use the new layout

- Changed authentication to use a gateway key
- Moved handlers and models to their own files
- Updated logger to syntax to use %s

requirements.txt

@@ -9,6 +9,7 @@ fastapi==0.146.0
 h11==0.16.0
 idna==3.20
 lib-secret-manager==0.3.1
+lib-service-auth==0.1.0
 multidict==7.1.0
 opentelemetry-api==1.45.1
 pamqp==4.0.1

src/auth_config.py

@@ -0,0 +1,11 @@
+from enum import StrEnum
+
+from service_auth import KeySource, ServiceAuth
+
+
+class Caller(StrEnum):
+	GATEWAY = "gateway"
+
+auth = ServiceAuth({
+	Caller.GATEWAY: KeySource("GATEWAY_KEY", "/etc/secrets/gateway_key"),
+})

src/handlers.py

@@ -0,0 +1,53 @@
+
+from fastapi import FastAPI, Request
+from fastapi.responses import JSONResponse
+from starlette.exceptions import HTTPException as StarletteHTTPException
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from metrics_server import REQUEST_COUNTER
+from simple_logger_handler import setup_logger
+
+logger = setup_logger(__name__)
+
+_KNOWN_METHODS = {"GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"}
+
+async def prometheus_middleware(request: Request, call_next):
+	logger.debug(
+		"[FastAPI] Incoming request: %s %s",
+		request.method, request.url.path,
+		)
+	status = 500
+	try:
+		response = await call_next(request)
+		status = response.status_code
+		logger.debug("[FastAPI] Request processed with status: %s", status)
+	except Exception as e:
+		logger.exception("Exception during request processing: %s", e)
+		raise
+	finally:
+		route = request.scope.get("route")
+		path = route.path if route else "unmatched"
+		method = request.method if request.method in _KNOWN_METHODS else "OTHER"
+		REQUEST_COUNTER.labels(method, path, status).inc()
+		logger.debug(
+				"[FastAPI] Prometheus counter incremented: %s %s %s",
+				method, path, status,
+				)
+	return response
+
+async def custom_http_exception_handler(request,exc):
+	logger.debug("[FastAPI] Handling HTTP exception: %s", 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}
+	)
+
+def register_handlers(app: FastAPI) -> None:
+	"""Attach the error handlers and the metrics middleware to the app."""
+	app.add_exception_handler(StarletteHTTPException, custom_http_exception_handler)
+	app.add_middleware(BaseHTTPMiddleware, dispatch=prometheus_middleware)

src/main.py

@@ -5,26 +5,19 @@ from contextlib import asynccontextmanager
 
 import uvicorn
 from aio_pika.exceptions import AMQPException
-from fastapi import FastAPI, HTTPException, Request
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel, Field
-from starlette.exceptions import HTTPException as StarletteHTTPException
+from fastapi import Depends, FastAPI, HTTPException, Request
+from service_auth import UserId
 
-from metrics_server import REQUEST_COUNTER
+from auth_config import Caller, auth
+from handlers import register_handlers
+from metrics_server import metrics_api
+from models import PullResponse
 from rabbitmq_handler import RabbitMQConsumer
 from secret_manager import cleanup_secret_files
 from simple_logger_handler import LOG_LEVEL, setup_logger
 
 logger = setup_logger(__name__)
-MAX_MESSAGES = 100
-
-class PullRequest(BaseModel):
-	user_id: int = Field(gt=0)
-
-class PullResponse(BaseModel):
-	messages: list[dict]
-	count: int
-	has_more: bool
+_MAX_MESSAGES = 10
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -53,77 +46,45 @@ api = FastAPI(
 		description="API to retrieve pull notifications on invocation",
 		version="1.0.0",
 		lifespan=lifespan,
-		docs_url=None,
-		redoc_url=None,
-		openapi_url=None
+		dependencies=[Depends(auth.authenticate)],
 )
 
-KNOWN_METHODS = {"GET", "POST", "PUT", "PATCH", "DELETE", "HEAD", "OPTIONS"}
-
-@api.middleware("http")
-async def prometheus_middleware(request: Request, call_next):
-	logger.debug(f"[FastAPI] Incoming request: {request.method} {request.url.path}")
-	status = 500
-	try:
-		response = await call_next(request)
-		status = response.status_code
-		logger.debug(f"[FastAPI] Request processed with status: {status}")
-	except Exception as e:
-		logger.exception(f"Exception during request processing: {e}")
-		raise
-	finally:
-		route = request.scope.get("route")
-		path = route.path if route else "unmatched"
-		method = request.method if request.method in KNOWN_METHODS else "OTHER"
-		REQUEST_COUNTER.labels(method, path, status).inc()
-		logger.debug(f"[FastAPI] Prometheus counter incremented: {method} {path} {status}")
-	return response
-
-@api.exception_handler(StarletteHTTPException)
-async def custom_http_exception_handler(request,exc):
-	logger.debug(f"[FastAPI] 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.post("/internal/pull-notifications", response_model=PullResponse)
+@api.post("/notifications/pull", response_model=PullResponse)
 async def pull_notifications(
-	pull_request: PullRequest,
+	user_id: UserId,
 	request: Request
 ):
-	user_id = pull_request.user_id
-	logger.debug(f"[FastAPI] User id {user_id} is requesting messages")
+	logger.debug("[FastAPI] User id %s is requesting messages", user_id)
 
 	try:
 		messages = await request.app.state.rmq_consumer.fetch_messages(
-			user_id=user_id, max_messages=MAX_MESSAGES
+			user_id=user_id, max_messages=_MAX_MESSAGES
 		)
 	except (AMQPException, OSError) as e:
-		logger.error(f"[FastAPI] RabbitMQ unavailable for user_id={user_id}: {e}")
+		logger.error("[FastAPI] RabbitMQ unavailable for user_id=%s: %s", user_id, e)
 		raise HTTPException(status_code=503, detail="Message broker unavailable") from e
 	except RuntimeError as e:
-		logger.error(f"[FastAPI] Consumer not ready: {e}")
+		logger.error("[FastAPI] Consumer not ready: %s", e)
 		raise HTTPException(status_code=503, detail="Service not ready") from e
 
-	logger.info(f"[FastAPI] Retrieved {len(messages)} message(s) for user_id={user_id}")
-	return PullResponse(messages=messages, count=len(messages), has_more=len(messages) >= MAX_MESSAGES)
+	logger.info("[FastAPI] Retrieved %s message(s) for user_id=%s",len(messages), user_id)
+	return PullResponse(messages=messages, count=len(messages), has_more=len(messages) >= _MAX_MESSAGES)
 
+register_handlers(api)
 
 class QuietServer(uvicorn.Server):
 	def install_signal_handlers(self):
 		pass
 
 async def start_servers():
-	logger.info("[Server] Starting main API (port 8101) and metrics server (port 9000)...")
-	logger.info(f"[Server] Build commit id={os.getenv("COMMIT_ID","unknown")}")
-	config_main = uvicorn.Config(api, host="0.0.0.0", port=8101, log_level=LOG_LEVEL.lower(), log_config=None)
-	config_metrics = uvicorn.Config("metrics_server:metrics_api", host="0.0.0.0", port=9000, log_level=LOG_LEVEL.lower(), log_config=None)
+	logger.info("[Server] Starting main API (port 5000) and metrics server (port 9000)...")
+	logger.info("[Server] Build commit id=%s", os.getenv("COMMIT_ID","unknown"))
+
+	auth.load_keys(api, Caller.GATEWAY)
+	cleanup_secret_files(auth.key_paths(Caller.GATEWAY))
+
+	config_main = uvicorn.Config(api, host="0.0.0.0", port=5000, log_level=LOG_LEVEL.lower(), log_config=None)
+	config_metrics = uvicorn.Config(metrics_api, host="0.0.0.0", port=9000, log_level=LOG_LEVEL.lower(), log_config=None)
 	servers = [QuietServer(config_main), QuietServer(config_metrics)]
 
 	def stop():

src/models.py

@@ -0,0 +1,11 @@
+
+from pydantic import BaseModel, Field
+
+
+class PullRequest(BaseModel):
+	user_id: int = Field(gt=0)
+
+class PullResponse(BaseModel):
+	messages: list[dict]
+	count: int
+	has_more: bool

src/rabbitmq_handler.py

@@ -26,7 +26,7 @@ class RabbitMQConsumer:
 		self.channel: aio_pika.RobustChannel | None = None
 		self.exchange: aio_pika.Exchange | None = None
 		self._closing = False
-		logger.debug(f"[RabbitMQ] Initialized consumer for exchange {self.rmq_exchange}")
+		logger.debug("[RabbitMQ] Initialized consumer for exchange %s", self.rmq_exchange)
 
 	def _lock_for(self, user_id: int) -> asyncio.Lock:
 		"""
@@ -47,7 +47,7 @@ class RabbitMQConsumer:
 			f"@{self.rmq_host}/{quote(self.rmq_vhost, safe='')}"
 		)
 
-		logger.info(f"[RabbitMQ] Connecting to {self.rmq_host}...")
+		logger.info("[RabbitMQ] Connecting to %s...", self.rmq_host)
 		self.connection = await aio_pika.connect_robust(url)
 		self.channel = await self.connection.channel()
 		self.exchange = await self.channel.declare_exchange(
@@ -105,7 +105,7 @@ class RabbitMQConsumer:
 					try:
 						data = json.loads(message.body.decode())
 					except (json.JSONDecodeError, UnicodeDecodeError) as e:
-						logger.error(f"[RabbitMQ] Bad message, discarding: {e}", exc_info=True)
+						logger.error("[RabbitMQ] Bad message, discarding: %s",e, exc_info=True)
 						await message.reject(requeue=False)
 						continue
 
@@ -115,11 +115,11 @@ class RabbitMQConsumer:
 				for message in to_ack:
 					await message.ack()
 
-				logger.debug(f"[RabbitMQ] Fetched {len(results)} message(s)")
+				logger.debug("[RabbitMQ] Fetched %s message(s)", len(results))
 				return results
 
 			except (AMQPException, OSError) as e:
-				logger.error(f"[RabbitMQ] Error while fetching, requeueing batch: {e}", exc_info=True)
+				logger.error("[RabbitMQ] Error while fetching, requeueing batch: %s",e, exc_info=True)
 				for message in to_ack:
 					try:
 						await message.nack(requeue=True)