From 2b1a92b311a1f99df1670b73d14c6763c6207ead Mon Sep 17 00:00:00 2001
From: Pan Teparak <panteparak@me.com>
Date: Wed, 16 Oct 2024 16:47:05 +0700
Subject: [PATCH] feat: add db_heartbeat

---
 health_check/contrib/db_heartbeat/__init__.py |  4 ++++
 health_check/contrib/db_heartbeat/apps.py     | 10 +++++++++
 health_check/contrib/db_heartbeat/backends.py | 21 +++++++++++++++++++
 3 files changed, 35 insertions(+)
 create mode 100644 health_check/contrib/db_heartbeat/__init__.py
 create mode 100644 health_check/contrib/db_heartbeat/apps.py
 create mode 100644 health_check/contrib/db_heartbeat/backends.py

diff --git a/health_check/contrib/db_heartbeat/__init__.py b/health_check/contrib/db_heartbeat/__init__.py
new file mode 100644
index 00000000..7aad707f
--- /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"
\ No newline at end of file
diff --git a/health_check/contrib/db_heartbeat/apps.py b/health_check/contrib/db_heartbeat/apps.py
new file mode 100644
index 00000000..896c2472
--- /dev/null
+++ b/health_check/contrib/db_heartbeat/apps.py
@@ -0,0 +1,10 @@
+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..ca9bb318
--- /dev/null
+++ b/health_check/contrib/db_heartbeat/backends.py
@@ -0,0 +1,21 @@
+from health_check.backends import BaseHealthCheckBackend
+from health_check.exceptions import ServiceUnavailable, ServiceReturnedUnexpectedResult
+from django.db import connection
+
+
+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 ServiceReturnedUnexpectedResult("Health Check query did not return the expected result.")
+        except Exception as e:
+            raise ServiceUnavailable(f"Database health check failed: {e}")