From d7b8e0c99a47a4babe923303eabd867a474762b2 Mon Sep 17 00:00:00 2001
From: Tim Loyer <tim.loyer@gmail.com>
Date: Fri, 30 Sep 2022 17:36:03 -0400
Subject: [PATCH] increase error verbosity with logs

---
 django_middleware/healthchecks.py | 100 ++++++++++++++++++++----------
 1 file changed, 68 insertions(+), 32 deletions(-)

diff --git a/django_middleware/healthchecks.py b/django_middleware/healthchecks.py
index 84ed32e..cf1f53e 100644
--- a/django_middleware/healthchecks.py
+++ b/django_middleware/healthchecks.py
@@ -10,7 +10,6 @@
 
 from django.http import HttpResponse, HttpResponseServerError
 
-
 HEALTHZ_ENDPOINT = "/healthz"
 READINESS_ENDPOINT = "/readiness"
 
@@ -28,46 +27,83 @@ def msg_filter(record):
 logger.addFilter(msg_filter)
 
 
-def check_readiness():
-    """Raise as exception if the server is not ready."""
+def databases_ready() -> bool:
+    """Check health of database connections.
 
-    # Connect to each database and do a generic standard SQL query that doesn't
-    # write any data and doesn't depend on any tables being present.
-    try:
-        from django.db import connections
+    Connect to each database and do a generic standard SQL query that doesn't write
+    any data and doesn't depend on any tables being present.
+    """
+    from django.db import connections as dbs
+
+    ready = True
+    databases = {alias: dbs[alias] for alias in dbs}
 
-        for name in connections:
-            cursor = connections[name].cursor()
+    for name, conn in databases.items():
+        try:
+            cursor = conn.cursor()
             cursor.execute("SELECT 1;")
-            row = cursor.fetchone()
-            if row is None:
-                return HttpResponseServerError("db: invalid response")
-    except Exception as e:
-        logger.exception(e)
-        return HttpResponseServerError("db: cannot connect to database.")
-
-    # Call get_stats() to connect to each memcached instance and get its stats.
-    # This can effectively check if each is online.
+
+            if cursor.fetchone() is None:
+                ready = False
+                logger.error(f"Database: Invalid response from '{name}'")
+        except Exception as exc:
+            ready = False
+            logger.error(f"Database: Error while attempting to connect to '{name}'")
+            logger.exception(exc)
+
+    return ready
+
+
+def caches_ready() -> bool:
+    """Check health of cache connections.
+
+    Calls get_stats() to connect to each memcached instance and get its stats.
+    This can effectively check if each is online.
+    """
+    from django.core.cache import caches
+    from django.core.cache.backends.memcached import BaseMemcachedCache
+
+    ready = True
+
+    for cache in caches.all():
+        if not isinstance(cache, BaseMemcachedCache):
+            continue
+
+        try:
+            stats = cache._cache.get_stats()
+            if len(stats) != len(cache._servers):
+                ready = False
+                logger.error(f"Cache: Unable to get stats for '{cache}'")
+        except Exception as exc:
+            ready = False
+            logger.error("Cache: Error while attempting to connect to '{cache}'")
+            logger.exception(exc)
+
+    return ready
+
+
+def check_readiness() -> HttpResponse:
+    """Raise as exception if the server is not ready."""
+    ready = True
+
     try:
-        from django.core.cache import caches
-        from django.core.cache.backends.memcached import BaseMemcachedCache
+        ready = databases_ready() and caches_ready()
+    except ModuleNotFoundError as exc:
+        logger.error("Import: Error while importing Django modules")
+        logger.exception(exc)
+    except Exception as exc:
+        logger.error("Unknown: Critical error")
+        logger.exception(exc)
 
-        for cache in caches.all():
-            if isinstance(cache, BaseMemcachedCache):
-                stats = cache._cache.get_stats()
-                if len(stats) != len(cache._servers):
-                    return HttpResponseServerError("cache: cannot connect to cache.")
-    except Exception as e:
-        logger.exception(e)
-        return HttpResponseServerError("cache: cannot connect to cache.")
+    return HttpResponse("OK") if ready else HttpResponseServerError("Not Ready")
 
 
 class HealthCheckMiddleware:
     """Simple Django health check endpoints.
 
     Endpoints:
-    healthz/    Responds with 200 OK if server can return a simple response.
-    readiness/  Responds with 200 OK if requisite databases and caches are ready.
+    /healthz    Responds with 200 OK if server can return a simple response.
+    /readiness  Responds with 200 OK if requisite databases and caches are ready.
     """
 
     endpoints = (
@@ -78,14 +114,14 @@ class HealthCheckMiddleware:
     def __init__(self, get_response):
         self.get_response = get_response
 
-    def __call__(self, request):
+    def __call__(self, request) -> HttpResponse:
         if request.method != "GET" or request.path not in self.endpoints:
             # Passthrough if we aren't accessing health checks.
             return self.get_response(request)
 
         if request.path == "/readiness":
             # Throw an exception if checks don't pass.
-            check_readiness()
+            return check_readiness()
 
         # Return a simple 200 OK response.
         return HttpResponse("OK")