Add db checks.

notesserver/test_views.py

@@ -2,6 +2,7 @@ import datetime
 import base64
 from mock import patch, Mock
 from django.core.urlresolvers import reverse
+from django.db import DatabaseError
 from rest_framework.test import APITestCase
 from elasticsearch.exceptions import TransportError
 
@@ -19,7 +20,7 @@ class OperationalEndpointsTest(APITestCase):
         self.assertEquals(response.data, {"OK": True})
 
     @patch('notesserver.views.get_es')
-    def test_heartbeat_failure(self, mocked_get_es):
+    def test_heartbeat_failure_es(self, mocked_get_es):
         """
         Elasticsearch is not reachable.
         """
@@ -28,6 +29,16 @@ class OperationalEndpointsTest(APITestCase):
         self.assertEquals(response.status_code, 500)
         self.assertEquals(response.data, {"OK": False, "check": "es"})
 
+    @patch("django.db.backends.utils.CursorWrapper")
+    def test_heartbeat_failure_db(self, mocked_cursor_wrapper):
+        """
+        Database is not reachable.
+        """
+        mocked_cursor_wrapper.side_effect = Exception
+        response = self.client.get(reverse('heartbeat'))
+        self.assertEquals(response.status_code, 500)
+        self.assertEquals(response.data, {"OK": False, "check": "db"})
+
     def test_root(self):
         """
         Test root endpoint.
@@ -42,15 +53,9 @@ class OperationalEndpointsTest(APITestCase):
             }
         )
 
-
-class OperationalAuthEndpointsTest(APITestCase):
-    """
-    Tests for operational authenticated endpoints.
-    """
-
     def test_selftest_status(self):
         """
-        Test status on authorization success.
+        Test status success.
         """
         response = self.client.get(reverse('selftest'))
         self.assertEquals(response.status_code, 200)
@@ -69,15 +74,27 @@ class OperationalAuthEndpointsTest(APITestCase):
             response.data,
             {
                 "es": {},
+                "db": "OK",
                 "time_elapsed": 0.0
             }
         )
 
     @patch('notesserver.views.get_es')
-    def test_selftest_failure(self, mocked_get_es):
+    def test_selftest_failure_es(self, mocked_get_es):
         """
         Elasticsearch is not reachable on selftest.
         """
         mocked_get_es.return_value.info.side_effect = TransportError()
         response = self.client.get(reverse('selftest'))
         self.assertEquals(response.status_code, 500)
+        self.assertIn('es_error', response.data)
+
+    @patch("django.db.backends.utils.CursorWrapper")
+    def test_selftest_failure_db(self, mocked_cursor_wrapper):
+        """
+        Database is not reachable on selftest.
+        """
+        mocked_cursor_wrapper.side_effect = Exception
+        response = self.client.get(reverse('selftest'))
+        self.assertEquals(response.status_code, 500)
+        self.assertIn('db_error', response.data)

notesserver/views.py

 import traceback
 import datetime
+
+from django.db import connection
 from rest_framework import status
 from rest_framework.permissions import AllowAny
 from rest_framework.response import Response
@@ -25,13 +27,18 @@ def root(request):  # pylint: disable=unused-argument
 @permission_classes([AllowAny])
 def heartbeat(request):  # pylint: disable=unused-argument
     """
-    ElasticSearch is reachable and ready to handle requests.
+    ElasticSearch and database are reachable and ready to handle requests.
     """
-    if get_es().ping():
-        return Response({"OK": True})
-    else:
+    try:
+        db_status()
+    except Exception:
+        return Response({"OK": False, "check": "db"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
+
+    if not get_es().ping():
         return Response({"OK": False, "check": "es"}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
+    return Response({"OK": True})
+
 
 @api_view(['GET'])
 @permission_classes([AllowAny])
@@ -40,6 +47,7 @@ def selftest(request):  # pylint: disable=unused-argument
     Manual test endpoint.
     """
     start = datetime.datetime.now()
+
     try:
         es_status = get_es().info()
     except TransportError:
@@ -48,10 +56,29 @@ def selftest(request):  # pylint: disable=unused-argument
             status=status.HTTP_500_INTERNAL_SERVER_ERROR
         )
 
+    try:
+        db_status()
+        database = "OK"
+    except Exception:
+        return Response(
+            {"db_error": traceback.format_exc()},
+            status=status.HTTP_500_INTERNAL_SERVER_ERROR
+        )
+
     end = datetime.datetime.now()
     delta = end - start
 
     return Response({
         "es": es_status,
+        "db": database,
         "time_elapsed": int(delta.total_seconds() * 1000)  # In milliseconds.
     })
+
+
+def db_status():
+    """
+    Return database status.
+    """
+    with connection.cursor() as cursor:
+        cursor.execute("SELECT 1")
+        cursor.fetchone()