Merge branch 'revsys:master' into feature/disable-atomic-requests-for-health-check

Author: TheSuperiorStanislav

README.md

@@ -105,6 +105,20 @@ To use Health Check Subsets, Specify a subset name and associate it with the rel
     }
 ```
 
+To add checks on a specific database, it's possible to parameterize `DatabaseBackend` to use a specific database:
+```python
+    HEALTH_CHECK = {
+        # .....
+        "SUBSETS": {
+            "database-probe": [
+                "DatabaseBackend[default]",  # This is equivalent to "DatabaseBackend"
+                "DatabaseBackend[secondary]",
+            ],
+        },
+        # .....
+    }
+```
+
 To only execute specific subset of health check
 ```shell
 curl -X GET -H "Accept: application/json" http://www.example.com/ht/startup-probe/

health_check/db/apps.py

@@ -1,4 +1,5 @@
 from django.apps import AppConfig
+from django.conf import settings
 
 from health_check.plugins import plugin_dir
 
@@ -11,3 +12,6 @@ def ready(self):
         from .backends import DatabaseBackend
 
         plugin_dir.register(DatabaseBackend)
+
+        for database_name in settings.DATABASES.keys():
+            plugin_dir.register(DatabaseBackend, database_name=database_name)

health_check/db/backends.py

@@ -1,4 +1,4 @@
-from django.db import DatabaseError, IntegrityError
+from django.db import DatabaseError, IntegrityError, transaction
 
 from health_check.backends import BaseHealthCheckBackend
 from health_check.exceptions import ServiceReturnedUnexpectedResult, ServiceUnavailable
@@ -7,13 +7,26 @@
 
 
 class DatabaseBackend(BaseHealthCheckBackend):
+    database_name = ""
+
+    def __init__(self, database_name: str = "") -> None:
+        super().__init__()
+        self.database_name = database_name or None
+
     def check_status(self):
         try:
-            obj = TestModel.objects.create(title="test")
-            obj.title = "newtest"
-            obj.save()
-            obj.delete()
+            with transaction.atomic(using=self.database_name):
+                obj = TestModel.objects.create(title="test")
+                obj.title = "newtest"
+                obj.save()
+                obj.delete()
         except IntegrityError:
             raise ServiceReturnedUnexpectedResult("Integrity Error")
         except DatabaseError:
             raise ServiceUnavailable("Database error")
+
+    def identifier(self) -> str:
+        if not self.database_name:
+            return super().identifier()
+
+        return f"{self.__class__.__name__}[{self.database_name}]"

tests/test_autodiscover.py

@@ -20,7 +20,9 @@ def test_autodiscover(self):
         ]
 
         # The number of installed apps excluding celery should equal to all plugins except celery
-        assert len(non_celery_plugins) == len(health_check_plugins)
+        assert len(non_celery_plugins) == len(health_check_plugins) + len(
+            settings.DATABASES  # Each database creates specific health_check attached to it
+        )
 
     def test_discover_celery_queues(self):
         celery_plugins = [x for x in plugin_dir._registry if issubclass(x[0], CeleryHealthCheck)]

tests/test_db.py

@@ -5,6 +5,7 @@
 from django.test import TestCase
 
 from health_check.db.backends import DatabaseBackend
+from health_check.mixins import CheckMixin
 
 
 class MockDBModel(Model):
@@ -78,3 +79,25 @@ def test_raise_exception(self):
         db_backend = DatabaseBackend()
         with self.assertRaises(Exception):
             db_backend.run_check()
+
+
+class MultipleDatabasesTests(TestCase):
+    """
+    Tests database specific features.
+    Ensures that all databases are registered as their own backend and each backend runs in its own database.
+    """
+
+    def test_all_databases_registered(self):
+        check_mixin = CheckMixin()
+        all_plugins = check_mixin.plugins
+
+        self.assertIn("DatabaseBackend[default]", all_plugins)
+        self.assertIn("DatabaseBackend[other]", all_plugins)
+
+    @patch("django.db.transaction.atomic")
+    def test_checks_on_database(self, mocked_atomic):
+        db1_health_check = DatabaseBackend(database_name="other")
+
+        db1_health_check.run_check()
+
+        mocked_atomic.assert_called_with(using="other")