diff --git a/main.py b/main.py index b91c2d8..54a86ac 100644 --- a/main.py +++ b/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", @@ -143,4 +131,147 @@ def node_status(data: NodeStatusRequest,token_data=Depends(verify_token)): "online": False, "status": "offline", "error": str(e) - } \ No newline at end of file + } + + + +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 + ) +