feat: add an endpoint for health check (#50)

Author: allisson

fastqueue/api.py

@@ -1,5 +1,6 @@
 import uvicorn
 from fastapi import Depends, FastAPI, Request, status
+from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
 from prometheus_fastapi_instrumentator import Instrumentator
 from sqlalchemy.orm import Session
@@ -11,6 +12,7 @@
     CreateMessageSchema,
     CreateQueueSchema,
     CreateTopicSchema,
+    HealthSchema,
     ListMessageSchema,
     ListQueueSchema,
     ListTopicSchema,
@@ -21,7 +23,7 @@
     TopicSchema,
     UpdateQueueSchema,
 )
-from fastqueue.services import MessageService, QueueService, TopicService
+from fastqueue.services import HealthService, MessageService, QueueService, TopicService
 
 tags_metadata = [
     {
@@ -67,7 +69,8 @@ def already_exists_exception_handler(request: Request, exc: AlreadyExistsError):
 
 @app.exception_handler(NotFoundError)
 def not_found_exception_handler(request: Request, exc: NotFoundError):
-    return JSONResponse(status_code=status.HTTP_404_NOT_FOUND, content={"detail": exc.args[0]})
+    response = jsonable_encoder(NotFoundSchema(detail=exc.args[0]))
+    return JSONResponse(status_code=status.HTTP_404_NOT_FOUND, content=response)
 
 
 @app.post("/topics", response_model=TopicSchema, status_code=status.HTTP_201_CREATED, tags=["topics"])
@@ -206,10 +209,22 @@ def nack_message(message_id: str, session: Session = Depends(get_session)):
     return MessageService(session=session).nack(id=message_id)
 
 
+@app.get(
+    "/health",
+    response_model=HealthSchema,
+    status_code=status.HTTP_200_OK,
+    tags=["healthcheck"],
+    responses={500: {"model": HealthSchema}},
+)
+def health(session: Session = Depends(get_session)):
+    response = HealthService(session=session).check()
+    status_code = status.HTTP_200_OK if response.success else status.HTTP_500_INTERNAL_SERVER_ERROR
+    return JSONResponse(status_code=status_code, content=jsonable_encoder(response))
+
+
 def run_server():
     uvicorn.run(
         "fastqueue.api:app",
-        debug=settings.debug,
         host=settings.server_host,
         port=settings.server_port,
         log_level=settings.log_level.lower(),

fastqueue/schemas.py

@@ -137,3 +137,7 @@ class Config:
 
 class ListMessageSchema(Schema):
     data: list[MessageSchema]
+
+
+class HealthSchema(Schema):
+    success: bool

fastqueue/services.py

@@ -2,6 +2,7 @@
 from datetime import datetime, timedelta
 from typing import Any
 
+from sqlalchemy import text
 from sqlalchemy.exc import IntegrityError
 from sqlalchemy.orm import Session
 
@@ -11,6 +12,7 @@
     CreateMessageSchema,
     CreateQueueSchema,
     CreateTopicSchema,
+    HealthSchema,
     ListMessageSchema,
     ListQueueSchema,
     ListTopicSchema,
@@ -332,3 +334,14 @@ def nack(self, id: str) -> None:
         message.scheduled_at = now
         message.updated_at = now
         self.session.commit()
+
+
+class HealthService(Service):
+    def check(self) -> HealthSchema:
+        success = True
+        try:
+            self.session.query(text("1")).from_statement(text("SELECT 1")).all()
+        except Exception:
+            success = False
+
+        return HealthSchema(success=success)

tests/test_api.py

@@ -294,3 +294,11 @@ def test_nack_message(session, message, client):
     response = client.put(f"/messages/{message.id}/nack")
 
     assert response.status_code == status.HTTP_204_NO_CONTENT
+
+
+def test_health(session, client):
+    response = client.get("/health")
+    response_data = response.json()
+
+    assert response.status_code == status.HTTP_200_OK
+    assert response_data == {"success": True}

tests/test_services.py

@@ -13,7 +13,7 @@
     RedriveQueueSchema,
     UpdateQueueSchema,
 )
-from fastqueue.services import MessageService, QueueService, TopicService
+from fastqueue.services import HealthService, MessageService, QueueService, TopicService
 from tests.factories import MessageFactory, QueueFactory, TopicFactory
 
 
@@ -428,3 +428,8 @@ def test_message_service_nack_with_removed_message(session):
     assert session.query(Message).filter_by(id=id).count() == 0
     assert MessageService(session=session).nack(id=id) is None
     assert session.query(Message).filter_by(id=id).count() == 0
+
+
+def test_health_service(session):
+    response = HealthService(session=session).check()
+    assert response.success