From e8c64d6d111a48e6cfa2c15c7dd062eb11223b3a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=9D=8E=E5=BC=BA?= <1206709430@qq.com>
Date: Tue, 19 Dec 2023 00:50:42 +0800
Subject: [PATCH] =?UTF-8?q?feat:=20=E6=96=B0=E5=A2=9E=E5=AD=98=E6=B4=BB?=
 =?UTF-8?q?=E6=A3=80=E6=B5=8B=E4=B8=AD=E9=97=B4=E4=BB=B6?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 backend/application/settings.py     |  1 +
 backend/dvadmin/utils/middleware.py | 57 +++++++++++++++++++++++++++++
 2 files changed, 58 insertions(+)

diff --git a/backend/application/settings.py b/backend/application/settings.py
index 590c144..77837ee 100644
--- a/backend/application/settings.py
+++ b/backend/application/settings.py
@@ -63,6 +63,7 @@ INSTALLED_APPS = [
 ]
 
 MIDDLEWARE = [
+    "dvadmin.utils.middleware.HealthCheckMiddleware",
     "django.middleware.security.SecurityMiddleware",
     "whitenoise.middleware.WhiteNoiseMiddleware",
     "django.contrib.sessions.middleware.SessionMiddleware",
diff --git a/backend/dvadmin/utils/middleware.py b/backend/dvadmin/utils/middleware.py
index 5443d14..f4717a4 100644
--- a/backend/dvadmin/utils/middleware.py
+++ b/backend/dvadmin/utils/middleware.py
@@ -2,9 +2,11 @@
 日志 django中间件
 """
 import json
+import logging
 
 from django.conf import settings
 from django.contrib.auth.models import AnonymousUser
+from django.http import HttpResponse, HttpResponseServerError
 from django.utils.deprecation import MiddlewareMixin
 
 from dvadmin.system.models import OperationLog
@@ -87,3 +89,58 @@ class ApiLoggingMiddleware(MiddlewareMixin):
             if self.methods == 'ALL' or request.method in self.methods:
                 self.__handle_response(request, response)
         return response
+
+logger = logging.getLogger("healthz")
+class HealthCheckMiddleware(object):
+    """
+    存活检查中间件
+    """
+    def __init__(self, get_response):
+        self.get_response = get_response
+        # One-time configuration and initialization.
+
+    def __call__(self, request):
+        if request.method == "GET":
+            if request.path == "/readiness":
+                return self.readiness(request)
+            elif request.path == "/healthz":
+                return self.healthz(request)
+        return self.get_response(request)
+
+    def healthz(self, request):
+        """
+        Returns that the server is alive.
+        """
+        return HttpResponse("OK")
+
+    def readiness(self, request):
+        # Connect to each database and do a generic standard SQL query
+        # that doesn't write any data and doesn't depend on any tables
+        # being present.
+        try:
+            from django.db import connections
+            for name in connections:
+                cursor = connections[name].cursor()
+                cursor.execute("SELECT 1;")
+                row = cursor.fetchone()
+                if row is None:
+                    return HttpResponseServerError("db: invalid response")
+        except Exception as e:
+            logger.exception(e)
+            return HttpResponseServerError("db: cannot connect to database.")
+
+        # Call get_stats() to connect to each memcached instance and get it's stats.
+        # This can effectively check if each is online.
+        try:
+            from django.core.cache import caches
+            from django.core.cache.backends.memcached import BaseMemcachedCache
+            for cache in caches.all():
+                if isinstance(cache, BaseMemcachedCache):
+                    stats = cache._cache.get_stats()
+                    if len(stats) != len(cache._servers):
+                        return HttpResponseServerError("cache: cannot connect to cache.")
+        except Exception as e:
+            logger.exception(e)
+            return HttpResponseServerError("cache: cannot connect to cache.")
+
+        return HttpResponse("OK")