From 2051936d08b31c024873051fd3b7725635579e47 Mon Sep 17 00:00:00 2001
From: Madison Swain-Bowden <bowdenm@spu.edu>
Date: Mon, 29 Apr 2024 15:38:31 -0700
Subject: [PATCH] Convert OperationalError into unhealthy health check response
 (#4197)

* Convert OperationalError into unhealthy health check response

* Use the appropriate error

* Unify errors and differentiate by code

* Specify service in exception detail
---
 api/api/views/health_views.py            | 12 ++++++++----
 api/test/unit/views/test_health_views.py | 17 +++++++++++++++--
 2 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/api/api/views/health_views.py b/api/api/views/health_views.py
index fb7e1da159a..556eb5cda44 100644
--- a/api/api/views/health_views.py
+++ b/api/api/views/health_views.py
@@ -1,5 +1,6 @@
 from django.conf import settings
 from django.db import connection
+from django.db.utils import OperationalError
 from rest_framework import status
 from rest_framework.exceptions import APIException
 from rest_framework.request import Request
@@ -9,7 +10,7 @@
 from api.utils.throttle import ExemptOAuth2IdRateThrottle, HealthcheckAnonRateThrottle
 
 
-class ElasticsearchHealthcheckException(APIException):
+class HealthCheckException(APIException):
     status_code = status.HTTP_503_SERVICE_UNAVAILABLE
 
 
@@ -32,7 +33,10 @@ def _check_db() -> None:
 
         Returns nothing if everything is OK, throws error otherwise.
         """
-        connection.ensure_connection()
+        try:
+            connection.ensure_connection()
+        except OperationalError as err:
+            raise HealthCheckException(f"postgres: {err}")
 
     @staticmethod
     def _check_es() -> None:
@@ -44,10 +48,10 @@ def _check_es() -> None:
         es_health = settings.ES.cluster.health(timeout="5s")
 
         if es_health["timed_out"]:
-            raise ElasticsearchHealthcheckException("es_timed_out")
+            raise HealthCheckException("elasticsearch: es_timed_out")
 
         if (es_status := es_health["status"]) != "green":
-            raise ElasticsearchHealthcheckException(f"es_status_{es_status}")
+            raise HealthCheckException(f"elasticsearch: es_status_{es_status}")
 
     def get(self, request: Request):
         if "check_es" in request.query_params:
diff --git a/api/test/unit/views/test_health_views.py b/api/test/unit/views/test_health_views.py
index d967e098fd6..12fde2e1b34 100644
--- a/api/test/unit/views/test_health_views.py
+++ b/api/test/unit/views/test_health_views.py
@@ -1,5 +1,7 @@
 from unittest import mock
 
+from django.db.utils import OperationalError
+
 import pook
 import pytest
 
@@ -31,13 +33,24 @@ def test_health_check_calls__check_db(api_client):
         mock_check_db.assert_called_once()
 
 
+def test_health_check_calls__check_db_with_failure(api_client):
+    with mock.patch(
+        "api.views.health_views.connection.ensure_connection"
+    ) as mock_ensure_connection:
+        mock_ensure_connection.side_effect = OperationalError("Database has gone away")
+        res = api_client.get("/healthcheck/")
+        assert res.status_code == 503
+        assert res.json() == {"detail": "postgres: Database has gone away"}
+        mock_ensure_connection.assert_called_once()
+
+
 def test_health_check_es_timed_out(api_client):
     with pook.use():
         mock_health_response(timed_out=True)
         res = api_client.get("/healthcheck/", data={"check_es": True})
 
     assert res.status_code == 503
-    assert res.json()["detail"] == "es_timed_out"
+    assert res.json() == {"detail": "elasticsearch: es_timed_out"}
 
 
 @pytest.mark.parametrize("status", ("yellow", "red"))
@@ -47,7 +60,7 @@ def test_health_check_es_status_bad(status, api_client):
         res = api_client.get("/healthcheck/", data={"check_es": True})
 
     assert res.status_code == 503
-    assert res.json()["detail"] == f"es_status_{status}"
+    assert res.json() == {"detail": f"elasticsearch: es_status_{status}"}
 
 
 @pytest.mark.django_db