From 7eb77fb911980c1c442bed48d4e9c889f5a76b6e Mon Sep 17 00:00:00 2001
From: ivan <user025@host025.local>
Date: Tue, 18 Jun 2024 16:25:47 +0500
Subject: [PATCH] Added health check with check connection to db adapter

---
 core/db_adapter/aioredis_adapter.py            | 10 ++++++++--
 core/db_adapter/db_adapter.py                  |  4 ++++
 core/db_adapter/ignite_adapter.py              |  7 +++++++
 core/db_adapter/memory_adapter.py              |  4 ++++
 smart_kit/start_points/main_loop_async_http.py | 16 ++++++++++++++++
 5 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/core/db_adapter/aioredis_adapter.py b/core/db_adapter/aioredis_adapter.py
index de9b9a34..20b4a2fb 100644
--- a/core/db_adapter/aioredis_adapter.py
+++ b/core/db_adapter/aioredis_adapter.py
@@ -1,6 +1,6 @@
 import copy
 import importlib
-import typing
+from typing import Optional
 
 from core.db_adapter import error
 from core.db_adapter.db_adapter import AsyncDBAdapter
@@ -13,7 +13,7 @@ def __init__(self, config=None):
         super().__init__(config)
         self.aioredis = importlib.import_module("redis", "asyncio")
         redis_type = self.aioredis.Redis
-        self._redis: typing.Optional[redis_type] = None
+        self._redis: Optional[redis_type] = None
 
         try:
             del self.config["type"]
@@ -70,3 +70,9 @@ async def _path_exists(self, path):
 
     async def _on_prepare(self):
         pass
+
+    async def is_alive(self) -> Optional[bool]:
+        try:
+            return self._redis.ping()
+        except Exception:
+            log("AIORedisAdapter is not alive", level="INFO")
diff --git a/core/db_adapter/db_adapter.py b/core/db_adapter/db_adapter.py
index 28c14a90..f10e9863 100644
--- a/core/db_adapter/db_adapter.py
+++ b/core/db_adapter/db_adapter.py
@@ -1,5 +1,6 @@
 # coding: utf-8
 import asyncio
+from typing import Optional
 
 import core.logging.logger_constants as log_const
 from core.logging.logger_utils import log
@@ -91,6 +92,9 @@ def _on_all_tries_fail(self):
 class AsyncDBAdapter(DBAdapter):
     IS_ASYNC = True
 
+    async def is_alive(self) -> Optional[bool]:
+        raise NotImplementedError
+
     async def _on_all_tries_fail(self):
         raise
 
diff --git a/core/db_adapter/ignite_adapter.py b/core/db_adapter/ignite_adapter.py
index 09c4c4fb..19fd30f1 100644
--- a/core/db_adapter/ignite_adapter.py
+++ b/core/db_adapter/ignite_adapter.py
@@ -88,3 +88,10 @@ def _on_prepare(self):
 
     async def _get_counter_name(self):
         return "ignite_async_adapter"
+
+    async def is_alive(self) -> Optional[bool]:
+        try:
+            await self.connect()
+            return True
+        except Exception:
+            log("IgniteAdapter is not alive", level="INFO")
diff --git a/core/db_adapter/memory_adapter.py b/core/db_adapter/memory_adapter.py
index 2df3a824..840eecee 100644
--- a/core/db_adapter/memory_adapter.py
+++ b/core/db_adapter/memory_adapter.py
@@ -2,6 +2,7 @@
 from copy import copy
 from cachetools import TTLCache
 from core.db_adapter import error
+from typing import Optional
 
 
 class MemoryAdapter(AsyncDBAdapter):
@@ -49,3 +50,6 @@ async def _get(self, id):
 
     async def _list_dir(self, path):
         pass
+
+    async def is_alive(self) -> Optional[bool]:
+        return True
diff --git a/smart_kit/start_points/main_loop_async_http.py b/smart_kit/start_points/main_loop_async_http.py
index a0349ea4..0d92ed50 100644
--- a/smart_kit/start_points/main_loop_async_http.py
+++ b/smart_kit/start_points/main_loop_async_http.py
@@ -21,6 +21,7 @@ def __init__(self, *args, **kwargs):
         self.app = aiohttp.web.Application()
         super().__init__(*args, **kwargs)
         self.app.add_routes([aiohttp.web.route('*', '/health', self.get_health_check)])
+        self.app.add_routes([aiohttp.web.route('*', '/health_db_adapter', self.get_health_db_adapter_check)])
         self.app.add_routes([aiohttp.web.route('*', '/{tail:.*}', self.iterate)])
 
     async def async_init(self):
@@ -174,6 +175,21 @@ async def get_health_check(self, request: aiohttp.web.Request):
             status=status, reason=reason, data=answer,
         )
 
+    async def get_health_db_adapter_check(self, request: aiohttp.web.Request):
+        status, reason, answer = 200, "OK", "ok"
+        try:
+            is_alive = await self.db_adapter.is_alive()
+            if not is_alive:
+                status, reason, answer = 500, "ERROR", "db adapter connection error"
+        except NotImplementedError:
+            status, reason, answer = 500, "ERROR", f"Method is not implemented in {type(self.db_adapter)}"
+        except Exception as e:
+            status, reason, answer = 500, "ERROR", str(e)
+
+        return aiohttp.web.json_response(
+            status=status, reason=reason, data=answer,
+        )
+
     async def iterate(self, request: aiohttp.web.Request):
         headers = self._get_headers(request.headers)
         body = await request.text()