Add ability to define and run a subset health checks (#390)

panteparak · web-flow · commit 4da1fd9d7cef · 2023-12-13T10:59:45.000-06:00
* feat: add ability to define health check subsets

* feat: adjust default health check to conform with new structure

* chore: add usage instruction in readme

* feat: allow subset argument to be specify using Django Command

* feat: fix broken tests with missing subset argument

* feat: move getattr from django settings to conf

* feat: remove unused variable

* feat: add tests

* feat: add subset not found should return HTTP 404

* feat: ensure plugins are sorted by name

* feat: ensure plugins are sorted by name

* feat: adjust health check config structure

* feat: add more tests

* feat: add commmand tests

* feat: fix merge conflicts

* feat: fix merge conflicts
diff --git a/README.md b/README.md
@@ -90,6 +90,24 @@ one of these checks, set its value to `None`.
     }
 ```
 
+To use Health Check Subsets, Specify a subset name and associate it with the relevant health check services to utilize Health Check Subsets.
+```python
+    HEALTH_CHECK = {
+        # .....
+        "SUBSETS": {
+            "startup-probe": ["MigrationsHealthCheck", "DatabaseBackend"],
+            "liveness-probe": ["DatabaseBackend"],
+            "<SUBSET_NAME>": ["<Health_Check_Service_Name"]
+        },
+        # .....
+    }
+```
+
+To only execute specific subset of health check
+```shell
+curl -X GET -H "Accept: application/json" http://www.example.com/ht/startup-probe/
+```
+
 If using the DB check, run migrations:
 
 ```shell
diff --git a/health_check/backends.py b/health_check/backends.py
@@ -20,14 +20,14 @@ class BaseHealthCheckBackend:
     def __init__(self):
         self.errors = []
 
-    def check_status(self):
+    def check_status(self, subset=None):
         raise NotImplementedError
 
-    def run_check(self):
+    def run_check(self, subset=None):
         start = timer()
         self.errors = []
         try:
-            self.check_status()
+            self.check_status(subset=subset)
         except HealthCheckException as e:
             self.add_error(e, e)
         except BaseException:
diff --git a/health_check/cache/backends.py b/health_check/cache/backends.py
@@ -28,7 +28,7 @@ def __init__(self, backend="default"):
     def identifier(self):
         return f"Cache backend: {self.backend}"
 
-    def check_status(self):
+    def check_status(self, subset=None):
         cache = caches[self.backend]
 
         try:
diff --git a/health_check/conf.py b/health_check/conf.py
@@ -4,4 +4,5 @@
 HEALTH_CHECK.setdefault("DISK_USAGE_MAX", 90)
 HEALTH_CHECK.setdefault("MEMORY_MIN", 100)
 HEALTH_CHECK.setdefault("WARNINGS_AS_ERRORS", True)
+HEALTH_CHECK.setdefault("SUBSETS", {})
 HEALTH_CHECK.setdefault("DISABLE_THREADING", False)
diff --git a/health_check/contrib/celery/backends.py b/health_check/contrib/celery/backends.py
@@ -8,7 +8,7 @@
 
 
 class CeleryHealthCheck(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         timeout = getattr(settings, "HEALTHCHECK_CELERY_TIMEOUT", 3)
         result_timeout = getattr(settings, "HEALTHCHECK_CELERY_RESULT_TIMEOUT", timeout)
         queue_timeout = getattr(settings, "HEALTHCHECK_CELERY_QUEUE_TIMEOUT", timeout)
diff --git a/health_check/contrib/celery_ping/backends.py b/health_check/contrib/celery_ping/backends.py
@@ -8,7 +8,7 @@
 class CeleryPingHealthCheck(BaseHealthCheckBackend):
     CORRECT_PING_RESPONSE = {"ok": "pong"}
 
-    def check_status(self):
+    def check_status(self, subset=None):
         timeout = getattr(settings, "HEALTHCHECK_CELERY_PING_TIMEOUT", 1)
 
         try:
diff --git a/health_check/contrib/migrations/backends.py b/health_check/contrib/migrations/backends.py
@@ -14,7 +14,7 @@ class MigrationsHealthCheck(BaseHealthCheckBackend):
     def get_migration_plan(self, executor):
         return executor.migration_plan(executor.loader.graph.leaf_nodes())
 
-    def check_status(self):
+    def check_status(self, subset=None):
         db_alias = getattr(settings, "HEALTHCHECK_MIGRATIONS_DB", DEFAULT_DB_ALIAS)
         try:
             executor = MigrationExecutor(connections[db_alias])
diff --git a/health_check/contrib/psutil/backends.py b/health_check/contrib/psutil/backends.py
@@ -14,7 +14,7 @@
 
 
 class DiskUsage(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         try:
             du = psutil.disk_usage("/")
             if DISK_USAGE_MAX and du.percent >= DISK_USAGE_MAX:
@@ -28,7 +28,7 @@ def check_status(self):
 
 
 class MemoryUsage(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         try:
             memory = psutil.virtual_memory()
             if MEMORY_MIN and memory.available < (MEMORY_MIN * 1024 * 1024):
diff --git a/health_check/contrib/rabbitmq/backends.py b/health_check/contrib/rabbitmq/backends.py
@@ -15,7 +15,7 @@ class RabbitMQHealthCheck(BaseHealthCheckBackend):
 
     namespace = None
 
-    def check_status(self):
+    def check_status(self, subset=None):
         """Check RabbitMQ service by opening and closing a broker channel."""
         logger.debug("Checking for a broker_url on django settings...")
 
diff --git a/health_check/contrib/redis/backends.py b/health_check/contrib/redis/backends.py
@@ -14,7 +14,7 @@ class RedisHealthCheck(BaseHealthCheckBackend):
 
     redis_url = getattr(settings, "REDIS_URL", "redis://localhost/1")
 
-    def check_status(self):
+    def check_status(self, subset=None):
         """Check Redis service by pinging the redis instance with a redis connection."""
         logger.debug("Got %s as the redis_url. Connecting to redis...", self.redis_url)
 
diff --git a/health_check/db/backends.py b/health_check/db/backends.py
@@ -7,7 +7,7 @@
 
 
 class DatabaseBackend(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         try:
             obj = TestModel.objects.create(title="test")
             obj.title = "newtest"
diff --git a/health_check/management/commands/health_check.py b/health_check/management/commands/health_check.py
@@ -1,22 +1,32 @@
 import sys
 
 from django.core.management.base import BaseCommand
+from django.http import Http404
 
 from health_check.mixins import CheckMixin
 
 
 class Command(CheckMixin, BaseCommand):
     help = "Run health checks and exit 0 if everything went well."
 
+    def add_arguments(self, parser):
+        parser.add_argument("-s", "--subset", type=str, nargs=1)
+
     def handle(self, *args, **options):
         # perform all checks
-        errors = self.errors
+        subset = options.get("subset", [])
+        subset = subset[0] if subset else None
+        try:
+            errors = self.check(subset=subset)
+        except Http404 as e:
+            self.stdout.write(str(e))
+            sys.exit(1)
 
-        for plugin in self.plugins:
+        for plugin_identifier, plugin in self.filter_plugins(subset=subset).items():
             style_func = self.style.SUCCESS if not plugin.errors else self.style.ERROR
             self.stdout.write(
                 "{:<24} ... {}\n".format(
-                    plugin.identifier(), style_func(plugin.pretty_status())
+                    plugin_identifier, style_func(plugin.pretty_status())
                 )
             )
 
diff --git a/health_check/mixins.py b/health_check/mixins.py
@@ -1,6 +1,9 @@
 import copy
+from collections import OrderedDict
 from concurrent.futures import ThreadPoolExecutor
 
+from django.http import Http404
+
 from health_check.conf import HEALTH_CHECK
 from health_check.exceptions import ServiceWarning
 from health_check.plugins import plugin_dir
@@ -16,19 +19,43 @@ def errors(self):
             self._errors = self.run_check()
         return self._errors
 
+    def check(self, subset=None):
+        return self.run_check(subset=subset)
+
     @property
     def plugins(self):
+        if not plugin_dir._registry:
+            return OrderedDict({})
+
         if not self._plugins:
-            self._plugins = sorted(
-                (
-                    plugin_class(**copy.deepcopy(options))
-                    for plugin_class, options in plugin_dir._registry
-                ),
-                key=lambda plugin: plugin.identifier(),
+            registering_plugins = (
+                plugin_class(**copy.deepcopy(options))
+                for plugin_class, options in plugin_dir._registry
+            )
+            registering_plugins = sorted(
+                registering_plugins, key=lambda plugin: plugin.identifier()
+            )
+            self._plugins = OrderedDict(
+                {plugin.identifier(): plugin for plugin in registering_plugins}
             )
         return self._plugins
 
-    def run_check(self):
+    def filter_plugins(self, subset=None):
+        if subset is None:
+            return self.plugins
+
+        health_check_subsets = HEALTH_CHECK["SUBSETS"]
+        if subset not in health_check_subsets or not self.plugins:
+            raise Http404(f"Specify subset: '{subset}' does not exists.")
+
+        selected_subset = set(health_check_subsets[subset])
+        return {
+            plugin_identifier: v
+            for plugin_identifier, v in self.plugins.items()
+            if plugin_identifier in selected_subset
+        }
+
+    def run_check(self, subset=None):
         errors = []
 
         def _run(plugin):
@@ -49,13 +76,15 @@ def _collect_errors(plugin):
                 else:
                     errors.extend(plugin.errors)
 
+        plugins = self.filter_plugins(subset=subset)
+        plugin_instances = plugins.values()
+
         if HEALTH_CHECK["DISABLE_THREADING"]:
-            for plugin in self.plugins:
+            for plugin in plugin_instances:
                 _run(plugin)
                 _collect_errors(plugin)
         else:
-            with ThreadPoolExecutor(max_workers=len(self.plugins) or 1) as executor:
-                for plugin in executor.map(_run, self.plugins):
+            with ThreadPoolExecutor(max_workers=len(plugin_instances) or 1) as executor:
+                for plugin in executor.map(_run, plugin_instances):
                     _collect_errors(plugin)
-
         return errors
diff --git a/health_check/storage/backends.py b/health_check/storage/backends.py
@@ -67,7 +67,7 @@ def check_delete(self, file_name):
         if storage.exists(file_name):
             raise ServiceUnavailable("File was not deleted")
 
-    def check_status(self):
+    def check_status(self, subset=None):
         try:
             # write the file to the storage backend
             file_name = self.get_file_name()
diff --git a/health_check/urls.py b/health_check/urls.py
@@ -6,4 +6,5 @@
 
 urlpatterns = [
     path("", MainView.as_view(), name="health_check_home"),
+    path("<str:subset>/", MainView.as_view(), name="health_check_subset"),
 ]
diff --git a/health_check/views.py b/health_check/views.py
@@ -88,12 +88,15 @@ class MainView(CheckMixin, TemplateView):
 
     @method_decorator(never_cache)
     def get(self, request, *args, **kwargs):
-        status_code = 500 if self.errors else 200
-
+        subset = kwargs.get("subset", None)
+        health_check_has_error = self.check(subset)
+        status_code = 500 if health_check_has_error else 200
         format_override = request.GET.get("format")
 
         if format_override == "json":
-            return self.render_to_response_json(self.plugins, status_code)
+            return self.render_to_response_json(
+                self.filter_plugins(subset=subset), status_code
+            )
 
         accept_header = request.META.get("HTTP_ACCEPT", "*/*")
         for media in MediaType.parse_header(accept_header):
@@ -106,18 +109,27 @@ def get(self, request, *args, **kwargs):
                 context = self.get_context_data(**kwargs)
                 return self.render_to_response(context, status=status_code)
             elif media.mime_type in ("application/json", "application/*"):
-                return self.render_to_response_json(self.plugins, status_code)
+                return self.render_to_response_json(
+                    self.filter_plugins(subset=subset), status_code
+                )
         return HttpResponse(
             "Not Acceptable: Supported content types: text/html, application/json",
             status=406,
             content_type="text/plain",
         )
 
     def get_context_data(self, **kwargs):
-        return {**super().get_context_data(**kwargs), "plugins": self.plugins}
+        subset = kwargs.get("subset", None)
+        return {
+            **super().get_context_data(**kwargs),
+            "plugins": self.filter_plugins(subset=subset).values(),
+        }
 
     def render_to_response_json(self, plugins, status):
         return JsonResponse(
-            {str(p.identifier()): str(p.pretty_status()) for p in plugins},
+            {
+                str(plugin_identifier): str(p.pretty_status())
+                for plugin_identifier, p in plugins.items()
+            },
             status=status,
         )
diff --git a/tests/test_commands.py b/tests/test_commands.py
@@ -4,16 +4,17 @@
 from django.core.management import call_command
 
 from health_check.backends import BaseHealthCheckBackend
+from health_check.conf import HEALTH_CHECK
 from health_check.plugins import plugin_dir
 
 
 class FailPlugin(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         self.add_error("Oops")
 
 
 class OkPlugin(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         pass
 
 
@@ -35,3 +36,44 @@ def test_command(self):
             "FailPlugin               ... unknown error: Oops\n"
             "OkPlugin                 ... working\n"
         )
+
+    def test_command_with_subset(self):
+        SUBSET_NAME_1 = "subset-1"
+        SUBSET_NAME_2 = "subset-2"
+        HEALTH_CHECK["SUBSETS"] = {
+            SUBSET_NAME_1: ["OkPlugin"],
+            SUBSET_NAME_2: ["OkPlugin", "FailPlugin"],
+        }
+
+        stdout = StringIO()
+        call_command("health_check", f"--subset={SUBSET_NAME_1}", stdout=stdout)
+        stdout.seek(0)
+        assert stdout.read() == ("OkPlugin                 ... working\n")
+
+    def test_command_with_failed_check_subset(self):
+        SUBSET_NAME = "subset-2"
+        HEALTH_CHECK["SUBSETS"] = {SUBSET_NAME: ["OkPlugin", "FailPlugin"]}
+
+        stdout = StringIO()
+        with pytest.raises(SystemExit):
+            call_command("health_check", f"--subset={SUBSET_NAME}", stdout=stdout)
+        stdout.seek(0)
+        assert stdout.read() == (
+            "FailPlugin               ... unknown error: Oops\n"
+            "OkPlugin                 ... working\n"
+        )
+
+    def test_command_with_non_existence_subset(self):
+        SUBSET_NAME = "subset-2"
+        NON_EXISTENCE_SUBSET_NAME = "abcdef12"
+        HEALTH_CHECK["SUBSETS"] = {SUBSET_NAME: ["OkPlugin"]}
+
+        stdout = StringIO()
+        with pytest.raises(SystemExit):
+            call_command(
+                "health_check", f"--subset={NON_EXISTENCE_SUBSET_NAME}", stdout=stdout
+            )
+        stdout.seek(0)
+        assert stdout.read() == (
+            f"Specify subset: '{NON_EXISTENCE_SUBSET_NAME}' does not exists.\n"
+        )
diff --git a/tests/test_mixins.py b/tests/test_mixins.py
@@ -9,12 +9,12 @@
 
 
 class FailPlugin(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         self.add_error("Oops")
 
 
 class OkPlugin(BaseHealthCheckBackend):
-    def check_status(self):
+    def check_status(self, subset=None):
         pass
 
 
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
diff --git a/tests/test_views.py b/tests/test_views.py
