Changes made as per comments

Author: Vansh Kalra

state-manager/app/main.py

@@ -23,6 +23,9 @@
 from .models.db.store import Store
 from .models.db.run import Run
 
+# Define models list
+DOCUMENT_MODELS = [State, GraphTemplate, RegisteredNode, Store, Run]
+
 # injecting routes
 from .routes import router
 
@@ -46,7 +49,7 @@ async def lifespan(app: FastAPI):
     # initializing beanie
     client = AsyncMongoClient(settings.mongo_uri)
     db = client[settings.mongo_database_name]
-    await init_beanie(db, document_models=[State, GraphTemplate, RegisteredNode, Store, Run])
+    await init_beanie(db, document_models=DOCUMENT_MODELS)
     logger.info("beanie dbs initialized")
 
     # initialize secret
@@ -55,7 +58,7 @@ async def lifespan(app: FastAPI):
     logger.info("secret initialized")
 
     # perform database health check
-    await check_database_health(logger)
+    await check_database_health(DOCUMENT_MODELS)
 
     # main logic of the server
     yield

state-manager/app/singletons/logs_manager.py

@@ -37,7 +37,7 @@ def __init__(self):
 
         if is_development:
             # In development mode, set level to WARNING to disable INFO logs
-            logger.setLevel(logging.WARNING)
+            logger.setLevel(logging.INFO)
         else:
             # In production mode, keep INFO level
             logger.setLevel(logging.INFO)

state-manager/app/utils/check_database_health.py

@@ -1,19 +1,16 @@
-# injecting models
-from ..models.db.state import State
-from ..models.db.graph_template_model import GraphTemplate
-from ..models.db.registered_node import RegisteredNode
-from ..models.db.store import Store
-from ..models.db.run import Run
-
-async def check_database_health(logger):
-    models_to_check = [State, GraphTemplate, RegisteredNode, Store, Run]
+from ..singletons.logs_manager import LogsManager
 
+async def check_database_health(models_to_check):
+    logger = LogsManager().get_logger()
+    checks_per_model = 3 
+    
     logger.info("Starting database health check")
 
     for model in models_to_check:
         try:
-            await model.find_one()
-            logger.info(f"Health check passed for {model.__name__}")
+            for i in range(checks_per_model):
+                await model.find_one()
+            logger.info(f"Health check passed for {model.__name__} ({checks_per_model} checks)")
         except Exception as e:
             error_msg = f"Database migrations needed as per the current version of state-manager. Failed to query {model.__name__}: {str(e)}"
             logger.error(error_msg)