diff --git a/README.md b/README.md index ffba6438..b95f2b44 100644 --- 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"], + "": ["= 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 index adaff36b..37c9f9dc 100644 --- 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 index 9d7272c6..e9076bbf 100644 --- 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 index 8687d994..1c7473fd 100644 --- 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 index 1bb7e723..191d50d1 100644 --- a/health_check/management/commands/health_check.py +++ b/health_check/management/commands/health_check.py @@ -1,6 +1,7 @@ import sys from django.core.management.base import BaseCommand +from django.http import Http404 from health_check.mixins import CheckMixin @@ -8,15 +9,24 @@ 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 index 8b3e5313..dd04180c 100644 --- 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 index 5920b133..0d966d67 100644 --- 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 index 92bc7d65..416406be 100644 --- a/health_check/urls.py +++ b/health_check/urls.py @@ -6,4 +6,5 @@ urlpatterns = [ path("", MainView.as_view(), name="health_check_home"), + path("/", MainView.as_view(), name="health_check_subset"), ] diff --git a/health_check/views.py b/health_check/views.py index 8eadf912..3c928e90 100644 --- 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,7 +109,9 @@ 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, @@ -114,10 +119,17 @@ def get(self, request, *args, **kwargs): ) 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 index 59c0ceb2..da2a7aec 100644 --- 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 index ca06b5e1..24bd6243 100644 --- 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 index 879d4b9d..d2ad270c 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -5,12 +5,12 @@ class FakePlugin(BaseHealthCheckBackend): - def check_status(self): + def check_status(self, subset=None): pass class Plugin(BaseHealthCheckBackend): - def check_status(self): + def check_status(self, subset=None): pass diff --git a/tests/test_views.py b/tests/test_views.py index 252e76c9..d9161081 100644 --- a/tests/test_views.py +++ b/tests/test_views.py @@ -92,7 +92,7 @@ def test_success(self, client): def test_error(self, client): class MyBackend(BaseHealthCheckBackend): - def check_status(self): + def check_status(self, subset=None): self.add_error("Super Fail!") plugin_dir.reset() @@ -104,7 +104,7 @@ def check_status(self): def test_warning(self, client): class MyBackend(BaseHealthCheckBackend): - def check_status(self): + def check_status(self, subset=None): raise ServiceWarning("so so") plugin_dir.reset() @@ -124,7 +124,7 @@ def test_non_critical(self, client): class MyBackend(BaseHealthCheckBackend): critical_service = False - def check_status(self): + def check_status(self, subset=None): self.add_error("Super Fail!") plugin_dir.reset() @@ -136,7 +136,7 @@ def check_status(self): def test_success_accept_json(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -147,7 +147,7 @@ def run_check(self): def test_success_prefer_json(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -160,7 +160,7 @@ def run_check(self): def test_success_accept_xhtml(self, client): class SuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -171,7 +171,7 @@ def run_check(self): def test_success_unsupported_accept(self, client): class SuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -186,7 +186,7 @@ def run_check(self): def test_success_unsupported_and_supported_accept(self, client): class SuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -199,7 +199,7 @@ def run_check(self): def test_success_accept_order(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -213,7 +213,7 @@ def run_check(self): def test_success_accept_order__reverse(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -227,7 +227,7 @@ def run_check(self): def test_format_override(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -238,7 +238,7 @@ def run_check(self): def test_format_no_accept_header(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -249,7 +249,7 @@ def run_check(self): def test_error_accept_json(self, client): class JSONErrorBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): self.add_error("JSON Error") plugin_dir.reset() @@ -266,7 +266,7 @@ def run_check(self): def test_success_param_json(self, client): class JSONSuccessBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): pass plugin_dir.reset() @@ -278,9 +278,57 @@ def run_check(self): JSONSuccessBackend().identifier(): JSONSuccessBackend().pretty_status() } + def test_success_subset_define(self, client): + class SuccessOneBackend(BaseHealthCheckBackend): + def run_check(self, subset=None): + pass + + class SuccessTwoBackend(BaseHealthCheckBackend): + def run_check(self, subset=None): + pass + + plugin_dir.reset() + plugin_dir.register(SuccessOneBackend) + plugin_dir.register(SuccessTwoBackend) + + HEALTH_CHECK["SUBSETS"] = { + "startup-probe": ["SuccessOneBackend", "SuccessTwoBackend"], + "liveness-probe": ["SuccessTwoBackend"], + } + + response_startup_probe = client.get( + self.url + "startup-probe/", {"format": "json"} + ) + assert ( + response_startup_probe.status_code == 200 + ), response_startup_probe.content.decode("utf-8") + assert response_startup_probe["content-type"] == "application/json" + assert json.loads(response_startup_probe.content.decode("utf-8")) == { + SuccessOneBackend().identifier(): SuccessOneBackend().pretty_status(), + SuccessTwoBackend().identifier(): SuccessTwoBackend().pretty_status(), + } + + response_liveness_probe = client.get( + self.url + "liveness-probe/", {"format": "json"} + ) + assert ( + response_liveness_probe.status_code == 200 + ), response_liveness_probe.content.decode("utf-8") + assert response_liveness_probe["content-type"] == "application/json" + assert json.loads(response_liveness_probe.content.decode("utf-8")) == { + SuccessTwoBackend().identifier(): SuccessTwoBackend().pretty_status(), + } + + def test_error_subset_not_found(self, client): + plugin_dir.reset() + response = client.get(self.url + "liveness-probe/", {"format": "json"}) + print(f"content: {response.content}") + print(f"code: {response.status_code}") + assert response.status_code == 404, response.content.decode("utf-8") + def test_error_param_json(self, client): class JSONErrorBackend(BaseHealthCheckBackend): - def run_check(self): + def run_check(self, subset=None): self.add_error("JSON Error") plugin_dir.reset()