From b03da9ab37883da02bef3b4269429d9f90b337e5 Mon Sep 17 00:00:00 2001
From: Colin Powell <colin@unbl.ink>
Date: Fri, 21 Feb 2025 10:40:44 -0500
Subject: [PATCH] [middleware] Add health check middleware

---
 vrobbler/health_check.py | 27 +++++++++++++++++++++++++++
 vrobbler/settings.py     |  1 +
 2 files changed, 28 insertions(+)
 create mode 100644 vrobbler/health_check.py

diff --git a/vrobbler/health_check.py b/vrobbler/health_check.py
new file mode 100644
index 0000000..d4e9154
--- /dev/null
+++ b/vrobbler/health_check.py
@@ -0,0 +1,27 @@
+import logging
+from django.http import HttpResponse
+from django.db import connections
+from django.db.utils import OperationalError
+
+logger = logging.getLogger(__name__)
+
+
+class HealthCheckMiddleware:
+    def __init__(self, get_response):
+        self.get_response = get_response
+
+    def __call__(self, request):
+        if request.path == "/health":
+            db_conn = connections["default"]
+            try:
+                c = db_conn.cursor()
+            except OperationalError:
+                is_db_connected = False
+            else:
+                is_db_connected = True
+            logger.info(
+                "[health-check]", extra={"is_db_connected": is_db_connected}
+            )
+            if is_db_connected:
+                return HttpResponse("ok")
+        return self.get_response(request)
diff --git a/vrobbler/settings.py b/vrobbler/settings.py
index bd40d7f..354bab8 100644
--- a/vrobbler/settings.py
+++ b/vrobbler/settings.py
@@ -148,6 +148,7 @@ INSTALLED_APPS = [
 SITE_ID = 1
 
 MIDDLEWARE = [
+    "vrobbler.health_check.HealthCheckMiddleware",
     "django.middleware.security.SecurityMiddleware",
     "django.contrib.sessions.middleware.SessionMiddleware",
     "django.middleware.common.CommonMiddleware",