Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add more light weight database check #439

Open
wants to merge 16 commits into
base: master
Choose a base branch
from
3 changes: 3 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -105,3 +105,6 @@ ENV/

.envrc
.direnv

# mac
.DS_Store
2 changes: 2 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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.

Expand Down Expand Up @@ -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',
]
```

Expand Down
4 changes: 4 additions & 0 deletions health_check/contrib/db_heartbeat/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
import django

if django.VERSION < (3, 2):
default_app_config = "health_check.contrib.db_heartbeat.apps.HealthCheckConfig"
12 changes: 12 additions & 0 deletions health_check/contrib/db_heartbeat/apps.py
Original file line number Diff line number Diff line change
@@ -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)
22 changes: 22 additions & 0 deletions health_check/contrib/db_heartbeat/backends.py
Original file line number Diff line number Diff line change
@@ -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}")
27 changes: 27 additions & 0 deletions tests/test_db_heartbeat.py
Original file line number Diff line number Diff line change
@@ -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()
1 change: 1 addition & 0 deletions tests/testapp/settings.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@
"health_check.contrib.migrations",
"health_check.contrib.celery_ping",
"health_check.contrib.s3boto_storage",
"health_check.contrib.db_heartbeat",
"tests",
)

Expand Down
Loading