diff --git a/.gitignore b/.gitignore index 427607d3..84debe81 100644 --- a/.gitignore +++ b/.gitignore @@ -105,3 +105,6 @@ ENV/ .envrc .direnv + +# mac +.DS_Store \ No newline at end of file diff --git a/README.md b/README.md index 833fd270..b408d952 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,7 @@ The following health checks are bundled with this project: - Celery ping - RabbitMQ - Migrations +- Database Heartbeat (Lightweight version of `health_check.db`) Writing your own custom health checks is also very quick and easy. @@ -74,6 +75,7 @@ Add the `health_check` applications to your `INSTALLED_APPS`: 'health_check.contrib.s3boto3_storage', # requires boto3 and S3BotoStorage backend 'health_check.contrib.rabbitmq', # requires RabbitMQ broker 'health_check.contrib.redis', # requires Redis broker + 'health_check.contrib.db_heartbeat', ] ``` diff --git a/health_check/contrib/db_heartbeat/__init__.py b/health_check/contrib/db_heartbeat/__init__.py new file mode 100644 index 00000000..3a3efc65 --- /dev/null +++ b/health_check/contrib/db_heartbeat/__init__.py @@ -0,0 +1,4 @@ +import django + +if django.VERSION < (3, 2): + default_app_config = "health_check.contrib.db_heartbeat.apps.HealthCheckConfig" diff --git a/health_check/contrib/db_heartbeat/apps.py b/health_check/contrib/db_heartbeat/apps.py new file mode 100644 index 00000000..d05fc806 --- /dev/null +++ b/health_check/contrib/db_heartbeat/apps.py @@ -0,0 +1,12 @@ +from django.apps import AppConfig + +from health_check.plugins import plugin_dir + + +class HealthCheckConfig(AppConfig): + name = "health_check.contrib.db_heartbeat" + + def ready(self): + from .backends import DatabaseHeartBeatCheck + + plugin_dir.register(DatabaseHeartBeatCheck) diff --git a/health_check/contrib/db_heartbeat/backends.py b/health_check/contrib/db_heartbeat/backends.py new file mode 100644 index 00000000..8c6ba30a --- /dev/null +++ b/health_check/contrib/db_heartbeat/backends.py @@ -0,0 +1,22 @@ +from django.db import connection + +from health_check.backends import BaseHealthCheckBackend +from health_check.exceptions import ServiceUnavailable + + +class DatabaseHeartBeatCheck(BaseHealthCheckBackend): + """Health check that runs a simple SELECT 1; query to test if the database connection is alive.""" + + def check_status(self): + try: + result = None + with connection.cursor() as cursor: + cursor.execute("SELECT 1;") + result = cursor.fetchone() + + if result != (1,): + raise ServiceUnavailable( + "Health Check query did not return the expected result." + ) + except Exception as e: + raise ServiceUnavailable(f"Database health check failed: {e}") diff --git a/tests/test_db_heartbeat.py b/tests/test_db_heartbeat.py new file mode 100644 index 00000000..abf5e7ca --- /dev/null +++ b/tests/test_db_heartbeat.py @@ -0,0 +1,27 @@ +import unittest +from unittest.mock import MagicMock, patch + +from health_check.contrib.db_heartbeat.backends import DatabaseHeartBeatCheck +from health_check.exceptions import ServiceUnavailable + + +class TestDatabaseHeartBeatCheck(unittest.TestCase): + @patch("health_check.contrib.db_heartbeat.backends.connection") + def test_check_status_success(self, mock_connection): + mock_cursor = MagicMock() + mock_cursor.fetchone.return_value = (1,) + mock_connection.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") + + health_check = DatabaseHeartBeatCheck() + with self.assertRaises(ServiceUnavailable): + health_check.check_status() diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index 620e7190..46c16aad 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -30,6 +30,7 @@ "health_check.contrib.migrations", "health_check.contrib.celery_ping", "health_check.contrib.s3boto_storage", + "health_check.contrib.db_heartbeat", "tests", )