Support multiple databases in heartbeat check

health_check/contrib/db_heartbeat/backends.py

@@ -1,4 +1,4 @@
-from django.db import connection
+from django.db import connections
 
 from health_check.backends import BaseHealthCheckBackend
 from health_check.exceptions import ServiceUnavailable
@@ -7,7 +7,16 @@
 class DatabaseHeartBeatCheck(BaseHealthCheckBackend):
     """Health check that runs a simple SELECT 1; query to test if the database connection is alive."""
 
+    def __init__(self, backend="default"):
+        super().__init__()
+        self.backend = backend
+
+    def identifier(self):
+        return f"Database heartbeat: {self.backend}"
+
     def check_status(self):
+        connection = connections[self.backend]
+
         try:
             result = None
             with connection.cursor() as cursor:

tests/test_db_heartbeat.py

@@ -6,21 +6,23 @@
 
 
 class TestDatabaseHeartBeatCheck(unittest.TestCase):
-    @patch("health_check.contrib.db_heartbeat.backends.connection")
-    def test_check_status_success(self, mock_connection):
+    @patch("health_check.contrib.db_heartbeat.backends.connections")
+    def test_check_status_success(self, mock_connections):
         mock_cursor = MagicMock()
         mock_cursor.fetchone.return_value = (1,)
-        mock_connection.cursor.return_value.__enter__.return_value = mock_cursor
+        mock_connections[
+            "default"
+        ].cursor.return_value.__enter__.return_value = mock_cursor
 
         health_check = DatabaseHeartBeatCheck()
         try:
             health_check.check_status()
         except Exception as e:
             self.fail(f"check_status() raised an exception unexpectedly: {e}")
 
-    @patch("health_check.contrib.db_heartbeat.backends.connection")
-    def test_check_status_service_unavailable(self, mock_connection):
-        mock_connection.cursor.side_effect = Exception("Database error")
+    @patch("health_check.contrib.db_heartbeat.backends.connections")
+    def test_check_status_service_unavailable(self, mock_connections):
+        mock_connections["default"].cursor.side_effect = Exception("Database error")
 
         health_check = DatabaseHeartBeatCheck()
         with self.assertRaises(ServiceUnavailable):