feat:

--Added async websocket API to check connection status

main.py

@@ -1,12 +1,13 @@
 import os
 import uvicorn
 from dotenv import load_dotenv
-from fastapi import FastAPI,Depends, Header, HTTPException, status
+from fastapi import FastAPI,Depends, Header, HTTPException, status, WebSocket, WebSocketDisconnect
 import psycopg2
 from jwt import ExpiredSignatureError, InvalidTokenError
 import jwt
 from pydantic import BaseModel
-
+from psycopg2.extras import RealDictCursor
+import asyncio
 
 load_dotenv()
 app = FastAPI()
@@ -90,19 +91,6 @@ def root():
 
 
 
-if __name__ == "__main__":
-    host = os.getenv("HOST", "127.0.0.1")
-    port = int(os.getenv("PORT", 9000))
-
-    uvicorn.run(
-        "main:app",
-        host=host,
-        port=port,
-        reload=True
-    )
-
-
-
 #{
 #     "host": "172.16.10.178",
 #     "node_name = "node3",
@@ -144,3 +132,146 @@ def node_status(data: NodeStatusRequest,token_data=Depends(verify_token)):
             "status": "offline",
             "error": str(e)
         }
+
+
+
+latest_replication_health = None
+monitor_task = None
+def get_replication_health(data: NodeStatusRequest):
+
+    db_config = {
+        "host": data.host,
+        "port": data.port,
+        "dbname": data.dbname,
+        "user": data.user,
+        "password": data.password,
+        "connect_timeout": 2,
+        "options": "-c statement_timeout=2000"
+    }
+
+    conn = None
+
+    try:
+        conn = psycopg2.connect(**db_config)
+
+        cur = conn.cursor(cursor_factory=RealDictCursor)
+
+        cur.execute("""
+            SELECT
+                subscription_name,
+                status,
+                provider_node
+            FROM spock.sub_show_status();
+        """)
+
+        subscriptions = cur.fetchall()
+
+        healthy = all(
+            sub["status"] == "replicating"
+            for sub in subscriptions
+        )
+
+        return {
+            "node": data.node_name,
+            "connection_healthy": healthy,
+            "subscriptions": subscriptions
+        }
+    except psycopg2.errors.QueryCanceled:
+
+        print("Query cancelled because it exceeded 2 seconds")
+
+        return {
+            "status": "timeout",
+            "message": "PostgreSQL query exceeded timelimit and was cancelled"
+        }
+
+    except Exception as e:
+
+        return {
+            "node": data.node_name,
+            "connection_healthy": False,
+            "error": str(e)
+        }
+
+    finally:
+        if conn:
+            conn.close()
+
+
+@app.websocket("/ws/node/replication-health")
+async def replication_monitor(websocket: WebSocket):
+    await websocket.accept()
+    try:
+        raw_data = await websocket.receive_json()
+        data = NodeStatusRequest(**raw_data)
+        global latest_replication_health
+
+        while True:
+
+            try:
+
+                result = await asyncio.wait_for(
+                    asyncio.to_thread(
+                    get_replication_health,
+                    data
+                    ),
+                    timeout=6
+                    )
+
+                latest_replication_health = result
+
+                print("Replication health:")
+                print(result)
+
+            except asyncio.TimeoutError:
+
+                latest_replication_health = {
+                    "node": data.node_name,
+                    "connection_healthy": False,
+                    "status": "timeout",
+                    "message": "PostgreSQL operation exceeded maximum time"
+                }
+
+            print(latest_replication_health)
+            
+            await websocket.send_json(
+                latest_replication_health
+            )
+
+            # Wait 2 seconds before checking again
+            await asyncio.sleep(2)
+    except WebSocketDisconnect:
+
+        print("WebSocket client disconnected")
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+if __name__ == "__main__":
+    host = os.getenv("HOST", "127.0.0.1")
+    port = int(os.getenv("PORT", 9000))
+
+    uvicorn.run(
+        "main:app",
+        host=host,
+        port=port,
+        reload=False
+    )
+