import os import uvicorn from dotenv import load_dotenv 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 import time load_dotenv() app = FastAPI() SECRET_KEY = os.getenv("SECRET_KEY") ALGORITHM = os.getenv("ALGORITHM", "HS256") class NodeStatusRequest(BaseModel): node_name: str host: str port: int dbname: str user: str password: str class DatabaseRequest(BaseModel): host: str port: int dbname: str user: str password: str def verify_token( authorization: str = Header(None), ): if not authorization: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Authorization header missing", headers={ "WWW-Authenticate": "Bearer" }, ) parts = authorization.split() if len(parts) != 2: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid Authorization header", ) scheme = parts[0] token = parts[1] if scheme.lower() != "bearer": raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Authorization scheme must be Bearer", ) try: payload = jwt.decode( token, SECRET_KEY, algorithms=[ALGORITHM], ) return payload except ExpiredSignatureError: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Token expired", ) except InvalidTokenError: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid token", ) @app.get("/") def root(): return {"message": "PostgreSQL Agent is running"} #{ # "host": "172.16.10.178", # "node_name :"node3", # "port": 5444, # "dbname": "postgres", # "user": "postgres", # "password": "postgres", # "connect_timeout": 3, # } @app.get("/node/status") def node_status(data: NodeStatusRequest,token_data=Depends(verify_token)): try: db_config = { "host": data.host, "port": data.port, "dbname": data.dbname, "user": data.user, "password": data.password, "connect_timeout": 3, } conn = psycopg2.connect(**db_config) conn.close() return { "node": data.node_name, "online": True, "status": "online" } except Exception as e: return { "node": data.node_name, "online": False, "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() # { # "node_name": "node3", # "host": "172.16.10.178", # "port": 5444, # "dbname": "postgres", # "user": "postgres", # "password": "postgres" # } @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") 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") #GET ACTIVE CONNECTIONS def get_active_connections( data: DatabaseRequest ): conn = None cursor = 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 COUNT(*) FROM pg_stat_activity WHERE state = 'active' AND pid <> pg_backend_pid(); """) active_connections = cursor.fetchone()[0] return { "database_name": data.dbname, "active_connections": active_connections } except psycopg2.Error as e: raise HTTPException( status_code=500, detail=f"Failed to fetch active connections: {str(e)}" ) finally: if cursor: cursor.close() if conn: conn.close() # HOW MANY CONNECTIONS ARE BEING USED def get_connection_usage(data: DatabaseRequest): conn = None cursor = 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 COUNT(*) AS current_connections, current_setting('max_connections')::int AS max_connections FROM pg_stat_activity GROUP BY current_setting('max_connections'); """) row = cursor.fetchone() current_connections = row[0] max_connections = row[1] connection_usage_percent = round( (current_connections / max_connections) * 100, 2 ) return { "current_connections": current_connections, "max_connections": max_connections, "connection_usage_percent": connection_usage_percent } except psycopg2.Error as e: return { "current_connections": None, "max_connections": None, "connection_usage_percent": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # WEB SOCKETS FOR ACTIVE CONNECTIONS MONITORING @app.websocket("/ws/database/active-connections") async def database_active_connections_monitor(websocket: WebSocket): await websocket.accept() try: raw_data = await websocket.receive_json() data = DatabaseRequest(**raw_data) while True: try: ( active_connections_result, connection_usage_result ) = await asyncio.wait_for( asyncio.gather( asyncio.to_thread( get_active_connections, data ), asyncio.to_thread( get_connection_usage, data ) ), timeout=5 ) result = { "database_name": data.dbname, "active_connections": active_connections_result[ "active_connections" ], "current_connections": connection_usage_result[ "current_connections" ], "max_connections": connection_usage_result[ "max_connections" ], "connection_usage_percent": connection_usage_result[ "connection_usage_percent" ] } except asyncio.TimeoutError: result = { "database_name": data.dbname, "active_connections": None, "current_connections": None, "max_connections": None, "connection_usage_percent": None } await websocket.send_json(result) await asyncio.sleep(2) except WebSocketDisconnect: print( "Database active connections WebSocket disconnected" ) #TPS def get_database_tps(data: DatabaseRequest): conn = None cursor = None try: conn = psycopg2.connect( host=data.host, port=data.port, dbname=data.dbname, user=data.user, password=data.password, connect_timeout=3 ) conn.autocommit = True cursor = conn.cursor() # First sample cursor.execute(""" SELECT xact_commit + xact_rollback FROM pg_stat_database WHERE datname = %s; """, (data.dbname,)) first_count = cursor.fetchone()[0] start_time = time.monotonic() time.sleep(1) # Second sample cursor.execute(""" SELECT xact_commit + xact_rollback FROM pg_stat_database WHERE datname = %s; """, (data.dbname,)) second_count = cursor.fetchone()[0] end_time = time.monotonic() elapsed_time = end_time - start_time transactions = second_count - first_count tps = round( transactions / elapsed_time, 2 ) return { "database_name": data.dbname, "tps": tps } except psycopg2.Error as e: return { "database_name": data.dbname, "tps": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # CACHE HIT RATIO def get_cache_hit_ratio(data: DatabaseRequest): conn = None cursor = 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 CASE WHEN (blks_hit + blks_read) = 0 THEN 0 ELSE ROUND( (blks_hit::numeric / (blks_hit + blks_read)) * 100, 2 ) END AS cache_hit_ratio FROM pg_stat_database WHERE datname = %s; """, (data.dbname,)) row = cursor.fetchone() cache_hit_ratio = float(row[0]) if row else 0.0 return { "database_name": data.dbname, "cache_hit_ratio": cache_hit_ratio } except psycopg2.Error as e: return { "database_name": data.dbname, "cache_hit_ratio": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() #TPS AND CACHE HIT RATIO WEB SOCKET @app.websocket("/ws/database/performance") async def database_performance_monitor(websocket: WebSocket): await websocket.accept() try: raw_data = await websocket.receive_json() data = DatabaseRequest(**raw_data) while True: try: ( tps_result, cache_hit_result ) = await asyncio.wait_for( asyncio.gather( asyncio.to_thread( get_database_tps, data ), asyncio.to_thread( get_cache_hit_ratio, data ) ), timeout=5 ) result = { "database_name": data.dbname, "tps": tps_result[ "tps" ], "cache_hit_ratio": cache_hit_result[ "cache_hit_ratio" ] } except asyncio.TimeoutError: result = { "database_name": data.dbname, "tps": None, "cache_hit_ratio": None } await websocket.send_json(result) await asyncio.sleep(2) except WebSocketDisconnect: print("Database performance WebSocket 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 )