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") #LOCKS def get_database_locks(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 total_locks, COUNT(*) FILTER ( WHERE l.granted = TRUE ) AS granted_locks, COUNT(*) FILTER ( WHERE l.granted = FALSE ) AS waiting_locks FROM pg_locks l JOIN pg_stat_activity a ON l.pid = a.pid WHERE a.datname = %s; """, (data.dbname,)) row = cursor.fetchone() return { "database_name": data.dbname, "total_locks": row[0], "granted_locks": row[1], "waiting_locks": row[2] } except psycopg2.Error as e: return { "database_name": data.dbname, "total_locks": None, "granted_locks": None, "waiting_locks": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # DEADLOCK COUNT def get_database_deadlocks(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 deadlocks FROM pg_stat_database WHERE datname = %s; """, (data.dbname,)) row = cursor.fetchone() deadlocks = row[0] if row else 0 return { "database_name": data.dbname, "deadlocks": deadlocks } except psycopg2.Error as e: return { "database_name": data.dbname, "deadlocks": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # LOCKS AND DEADLOCKS WEB SOCKET @app.websocket("/ws/database/locks-deadlocks") async def database_locks_deadlocks_monitor(websocket: WebSocket): await websocket.accept() try: raw_data = await websocket.receive_json() data = DatabaseRequest(**raw_data) while True: try: ( locks_result, deadlocks_result ) = await asyncio.wait_for( asyncio.gather( asyncio.to_thread( get_database_locks, data ), asyncio.to_thread( get_database_deadlocks, data ) ), timeout=5 ) result = { "database_name": data.dbname, "total_locks": locks_result["total_locks"], "granted_locks": locks_result["granted_locks"], "waiting_locks": locks_result["waiting_locks"], "deadlocks": deadlocks_result["deadlocks"] } except asyncio.TimeoutError: result = { "database_name": data.dbname, "total_locks": None, "granted_locks": None, "waiting_locks": None, "deadlocks": None } await websocket.send_json(result) await asyncio.sleep(2) except WebSocketDisconnect: print("Locks and deadlocks WebSocket disconnected") #-========================================================================================================================================================================================= #EXPERIMENTAL FUNTIONS # QUERY LATENCY def get_query_latency(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 SUM(calls) = 0 THEN 0 ELSE ROUND( SUM(total_exec_time)::numeric / SUM(calls), 2 ) END FROM pg_stat_statements WHERE dbid = ( SELECT oid FROM pg_database WHERE datname = %s ); """, (data.dbname,)) row = cursor.fetchone() query_latency_ms = ( float(row[0]) if row and row[0] is not None else 0.0 ) return { "database_name": data.dbname, "query_latency_ms": query_latency_ms } except psycopg2.Error as e: return { "database_name": data.dbname, "query_latency_ms": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() #SLOW RUNNING QUERIES def get_slow_queries( data: DatabaseRequest, threshold_seconds: int = 1 ): 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_factory=RealDictCursor ) cursor.execute(""" SELECT pid, usename, application_name, client_addr, state, query, ROUND( EXTRACT( EPOCH FROM ( CURRENT_TIMESTAMP - query_start ) )::numeric * 1000, 2 ) AS duration_ms, query_start FROM pg_stat_activity WHERE datname = %s AND state = 'active' AND pid <> pg_backend_pid() AND query_start IS NOT NULL AND CURRENT_TIMESTAMP - query_start > (%s * INTERVAL '1 second') ORDER BY query_start ASC; """, ( data.dbname, threshold_seconds )) rows = cursor.fetchall() return { "database_name": data.dbname, "threshold_seconds": threshold_seconds, "slow_query_count": len(rows), "slow_queries": [ { "pid": row["pid"], "user": row["usename"], "application_name": row["application_name"], "client_addr": ( str(row["client_addr"]) if row["client_addr"] else None ), "state": row["state"], "query": row["query"], "duration_ms": float(row["duration_ms"]), "query_start": row["query_start"].isoformat() } for row in rows ] } except psycopg2.Error as e: return { "database_name": data.dbname, "threshold_seconds": threshold_seconds, "slow_query_count": None, "slow_queries": [], "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() #BLOCKED QUERIES def get_blocked_queries(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_factory=RealDictCursor ) cursor.execute(""" SELECT blocked.pid AS blocked_pid, blocked.usename AS blocked_user, blocked.application_name AS blocked_application, blocked.client_addr AS blocked_client_addr, blocked.query AS blocked_query, blocked.query_start AS blocked_query_start, blocked.wait_event_type, blocked.wait_event, ROUND( EXTRACT( EPOCH FROM ( CURRENT_TIMESTAMP - blocked.query_start ) )::numeric * 1000, 2 ) AS blocked_duration_ms, pg_blocking_pids(blocked.pid) AS blocking_pids FROM pg_stat_activity blocked WHERE blocked.datname = %s AND blocked.pid <> pg_backend_pid() AND cardinality( pg_blocking_pids(blocked.pid) ) > 0 ORDER BY blocked.query_start ASC; """, (data.dbname,)) rows = cursor.fetchall() return { "database_name": data.dbname, "blocked_query_count": len(rows), "blocked_queries": [ { "blocked_pid": row["blocked_pid"], "blocked_user": row["blocked_user"], "blocked_application": row["blocked_application"], "blocked_client_addr": str(row["blocked_client_addr"]) if row["blocked_client_addr"] else None, "blocked_query": row["blocked_query"], "blocked_query_start": row["blocked_query_start"].isoformat() if row["blocked_query_start"] else None, "blocked_duration_ms": float(row["blocked_duration_ms"]) if row["blocked_duration_ms"] is not None else None, "wait_event_type": row["wait_event_type"], "wait_event": row["wait_event"], "blocking_pids": row["blocking_pids"] } for row in rows ] } except psycopg2.Error as e: return { "database_name": data.dbname, "blocked_query_count": None, "blocked_queries": [], "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() #======================================================================================================================================================================================================= 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 )