diff --git a/main.py b/main.py index 8445dde..2e72bb5 100644 --- a/main.py +++ b/main.py @@ -26,6 +26,14 @@ class NodeStatusRequest(BaseModel): +class DatabaseRequest(BaseModel): + host: str + port: int + dbname: str + user: str + password: str + + def verify_token( authorization: str = Header(None), ): @@ -256,7 +264,137 @@ async def replication_monitor(websocket: WebSocket): +def get_database_status( + data: DatabaseRequest +): + conn = None + + try: + conn = psycopg2.connect( + host=data.host, + port=data.port, + dbname=data.dbname, + user=data.user, + password=data.password, + connect_timeout=3 + ) + + return { + "database_name": data.dbname, + "status": "online" + } + + except psycopg2.Error: + + return { + "database_name": data.dbname, + "status": "offline" + } + + finally: + if conn: + conn.close() + + +def get_database_uptime( + data: DatabaseRequest +): + + conn = None + + try: + conn = psycopg2.connect( + host=data.host, + port=data.port, + dbname=data.dbname, + user=data.user, + password=data.password, + connect_timeout=3 + ) + + cursor = conn.cursor() + + cursor.execute(""" + SELECT EXTRACT( + EPOCH FROM ( + CURRENT_TIMESTAMP - pg_postmaster_start_time() + ) + )::BIGINT; + """) + + uptime_seconds = cursor.fetchone()[0] + + cursor.close() + + return { + "database_name": data.dbname, + "uptime_seconds": uptime_seconds + } + + except psycopg2.Error as e: + + return { + "database_name": data.dbname, + "uptime_seconds": None, + "error": str(e) + } + + finally: + if conn: + conn.close() + + +@app.websocket("/ws/database/status") +async def database_status_monitor(websocket: WebSocket): + + await websocket.accept() + + try: + + raw_data = await websocket.receive_json() + + data = DatabaseRequest(**raw_data) + + while True: + + try: + + status_result, uptime_result = await asyncio.wait_for( + asyncio.gather( + asyncio.to_thread( + get_database_status, + data + ), + asyncio.to_thread( + get_database_uptime, + data + ) + ), + timeout=5 + ) + + result = { + "database_name": data.dbname, + "status": status_result["status"], + "uptime_seconds": uptime_result["uptime_seconds"] + } + + except asyncio.TimeoutError: + + result = { + "database_name": data.dbname, + "status": "offline", + "uptime_seconds": None + } + + await websocket.send_json(result) + + await asyncio.sleep(2) + + except WebSocketDisconnect: + + print("Database status WebSocket disconnected")