import os import uvicorn from dotenv import load_dotenv from fastapi import FastAPI,Depends, Header, HTTPException, status, WebSocket, WebSocketDisconnect, WebSocketException 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", ) # JWT TOKEN VERIFICATION FOR WEBSOCKET CONNECTIONS async def verify_websocket_token(websocket: WebSocket): token = websocket.query_params.get("token") if not token: raise WebSocketException( code=status.WS_1008_POLICY_VIOLATION, reason="Authentication token missing" ) try: payload = jwt.decode( token, SECRET_KEY, algorithms=[ALGORITHM], ) return payload except ExpiredSignatureError: raise WebSocketException( code=status.WS_1008_POLICY_VIOLATION, reason="Token expired" ) except InvalidTokenError: raise WebSocketException( code=status.WS_1008_POLICY_VIOLATION, reason="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" # } #DONE @app.websocket("/ws/node/replication-health") async def replication_monitor(websocket: WebSocket): token_data = await verify_websocket_token(websocket) # if token_data is None: # return 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() # DONE @app.websocket("/ws/database/status") async def database_status_monitor(websocket: WebSocket): token_data = await verify_websocket_token(websocket) # if token_data is None: # return 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 # DONE @app.websocket("/ws/database/active-connections") async def database_active_connections_monitor(websocket: WebSocket): token_data = await verify_websocket_token(websocket) # if token_data is None: # return 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 # DONE @app.websocket("/ws/database/performance") async def database_performance_monitor(websocket: WebSocket): token_data = await verify_websocket_token(websocket) # if token_data is None: # return 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): token_data = await verify_websocket_token(websocket) # if token_data is None: # return 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") #GET DATABASE SIZE def get_database_size(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 pg_database_size(%s) AS size_bytes, ROUND( pg_database_size(%s)::numeric / 1024 / 1024 / 1024, 2 ) AS size_gb; """, ( data.dbname, data.dbname )) row = cursor.fetchone() return { "database_name": data.dbname, "size_bytes": row[0], "size_gb": float(row[1]) } except psycopg2.Error as e: return { "database_name": data.dbname, "size_bytes": None, "size_gb": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # API ENDPOINTS FOR DATABASE SIZE #DONE @app.post("/database/size") def database_size( data: DatabaseRequest, token_data: dict = Depends(verify_token) ): return get_database_size(data) # API ENDPOINTS FOR LARGEST TABLES def get_largest_tables( data: DatabaseRequest, limit: int = 10 ): 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 schemaname, relname AS table_name, pg_total_relation_size( quote_ident(schemaname) || '.' || quote_ident(relname) ) AS size_bytes, ROUND( pg_total_relation_size( quote_ident(schemaname) || '.' || quote_ident(relname) )::numeric / 1024 / 1024, 2 ) AS size_mb FROM pg_stat_user_tables ORDER BY size_bytes DESC LIMIT %s; """, (limit,)) rows = cursor.fetchall() return { "database_name": data.dbname, "table_count": len(rows), "tables": [ { "schema": row["schemaname"], "table_name": row["table_name"], "size_bytes": row["size_bytes"], "size_mb": float(row["size_mb"]) } for row in rows ] } except psycopg2.Error as e: return { "database_name": data.dbname, "table_count": None, "tables": [], "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # SHOWS THE LIST OF TABLES IN THE DATABASE SORTED BY SIZE # DONE @app.post("/database/largest-tables") def database_largest_tables( data: DatabaseRequest, token_data: dict = Depends(verify_token) ): return get_largest_tables( data, limit=10 ) # FUNCTION TO GET INDEX SIZE def get_index_size(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 COALESCE( SUM( pg_indexes_size( quote_ident(schemaname) || '.' || quote_ident(relname) ) ), 0 ) AS index_size_bytes FROM pg_stat_user_tables; """) row = cursor.fetchone() index_size_bytes = row[0] index_size_gb = round( index_size_bytes / 1024 / 1024 / 1024, 2 ) return { "database_name": data.dbname, "index_size_bytes": index_size_bytes, "index_size_gb": index_size_gb } except psycopg2.Error as e: return { "database_name": data.dbname, "index_size_bytes": None, "index_size_gb": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # API ENDPOINT FOR INDEX SIZE REST API # DONE @app.post("/database/index-size") def database_index_size( data: DatabaseRequest, token_data: dict = Depends(verify_token) ): return get_index_size(data) # GET TUPLE READS AND WRITES PER SECOND def get_tuple_reads_writes(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 tup_returned, tup_inserted, tup_updated, tup_deleted FROM pg_stat_database WHERE datname = %s; """, (data.dbname,)) first = cursor.fetchone() start_time = time.monotonic() time.sleep(1) # Second sample cursor.execute(""" SELECT tup_returned, tup_inserted, tup_updated, tup_deleted FROM pg_stat_database WHERE datname = %s; """, (data.dbname,)) second = cursor.fetchone() end_time = time.monotonic() elapsed_time = end_time - start_time # Reads reads = second[0] - first[0] # Writes inserts = second[1] - first[1] updates = second[2] - first[2] deletes = second[3] - first[3] writes = inserts + updates + deletes reads_per_second = round( reads / elapsed_time, 2 ) writes_per_second = round( writes / elapsed_time, 2 ) return { "database_name": data.dbname, "tuple_reads_per_second": reads_per_second, "tuple_writes_per_second": writes_per_second, "inserts_per_second": round(inserts / elapsed_time, 2), "updates_per_second": round(updates / elapsed_time, 2), "deletes_per_second": round(deletes / elapsed_time, 2) } except psycopg2.Error as e: return { "database_name": data.dbname, "tuple_reads_per_second": None, "tuple_writes_per_second": None, "inserts_per_second": None, "updates_per_second": None, "deletes_per_second": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # TUPLE READS AND WRITES PER SECOND WEB SOCKET # DONE @app.websocket("/ws/database/tuple-activity") async def database_tuple_activity_monitor(websocket: WebSocket): token_data = await verify_websocket_token(websocket) if token_data is None: return await websocket.accept() try: raw_data = await websocket.receive_json() data = DatabaseRequest(**raw_data) while True: try: tuple_result = await asyncio.wait_for( asyncio.to_thread( get_tuple_reads_writes, data ), timeout=5 ) result = { "database_name": data.dbname, "tuple_reads_per_second": tuple_result["tuple_reads_per_second"], "tuple_writes_per_second": tuple_result["tuple_writes_per_second"], "inserts_per_second": tuple_result["inserts_per_second"], "updates_per_second": tuple_result["updates_per_second"], "deletes_per_second": tuple_result["deletes_per_second"] } except asyncio.TimeoutError: result = { "database_name": data.dbname, "tuple_reads_per_second": None, "tuple_writes_per_second": None, "inserts_per_second": None, "updates_per_second": None, "deletes_per_second": None } await websocket.send_json(result) await asyncio.sleep(2) except WebSocketDisconnect: print("Tuple activity WebSocket disconnected") # GET DEAD TUPLES def get_dead_tuples(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 schemaname, relname AS table_name, n_live_tup, n_dead_tup, last_vacuum, last_autovacuum, pg_total_relation_size( quote_ident(schemaname) || '.' || quote_ident(relname) ) AS size_bytes, ROUND( pg_total_relation_size( quote_ident(schemaname) || '.' || quote_ident(relname) )::numeric / 1024 / 1024, 2 ) AS size_mb FROM pg_stat_user_tables ORDER BY size_bytes DESC; """) rows = cursor.fetchall() total_dead_tuples = sum( row["n_dead_tup"] or 0 for row in rows ) total_size_bytes = sum( row["size_bytes"] or 0 for row in rows ) total_size_mb = round( total_size_bytes / 1024 / 1024, 2 ) return { "database_name": data.dbname, "table_count": len(rows), "total_dead_tuples": total_dead_tuples, "total_table_size_bytes": total_size_bytes, "total_table_size_mb": total_size_mb, "tables": [ { "schema": row["schemaname"], "table_name": row["table_name"], "live_tuples": row["n_live_tup"], "dead_tuples": row["n_dead_tup"], "size_bytes": row["size_bytes"], "size_mb": float(row["size_mb"]), "last_vacuum": row["last_vacuum"].isoformat() if row["last_vacuum"] else None, "last_autovacuum": row["last_autovacuum"].isoformat() if row["last_autovacuum"] else None } for row in rows ] } except psycopg2.Error as e: return { "database_name": data.dbname, "table_count": None, "total_dead_tuples": None, "total_table_size_bytes": None, "total_table_size_mb": None, "tables": [], "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() #API ENDPOINT FOR DEAD TUPLES FOR EACH TABLE IN THE DATABASE @app.post("/database/dead-tuples") def database_dead_tuples( data: DatabaseRequest, token_data: dict = Depends(verify_token) ): return get_dead_tuples(data) # AUTO VACUUM STATUS def get_autovacuum_activity(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 ) # Check whether autovacuum is globally enabled cursor.execute(""" SHOW autovacuum; """) autovacuum_enabled = cursor.fetchone()["autovacuum"] == "on" # Check currently running autovacuum workers cursor.execute(""" SELECT COUNT(*) AS active_workers FROM pg_stat_activity WHERE datname = %s AND backend_type = 'autovacuum worker'; """, (data.dbname,)) active_workers = cursor.fetchone()["active_workers"] # Get autovacuum statistics for tables cursor.execute(""" SELECT COUNT(*) AS total_tables, COUNT(*) FILTER ( WHERE autovacuum_count > 0 ) AS tables_autovacuumed, COALESCE( SUM(autovacuum_count), 0 ) AS total_autovacuum_runs, MAX(last_autovacuum) AS last_autovacuum FROM pg_stat_user_tables; """) stats = cursor.fetchone() if not autovacuum_enabled: status = "disabled" elif active_workers > 0: status = "running" else: status = "idle" return { "database_name": data.dbname, "status": status, "autovacuum_enabled": autovacuum_enabled, "active_workers": active_workers, "total_tables": stats["total_tables"], "tables_autovacuumed": stats["tables_autovacuumed"], "total_autovacuum_runs": stats["total_autovacuum_runs"], "last_autovacuum": stats["last_autovacuum"].isoformat() if stats["last_autovacuum"] else None } except psycopg2.Error as e: return { "database_name": data.dbname, "status": "unknown", "autovacuum_enabled": None, "active_workers": None, "total_tables": None, "tables_autovacuumed": None, "total_autovacuum_runs": None, "last_autovacuum": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # API ENDPOINT FOR AUTO VACUUM STATUS @app.post("/database/autovacuum-activity") def database_autovacuum_activity( data: DatabaseRequest, token_data: dict = Depends(verify_token) ): return get_autovacuum_activity(data) # WAL GENERATED PER SECOND def get_wal_generated_per_second(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 wal_bytes FROM pg_stat_wal; """) first_wal_bytes = float(cursor.fetchone()[0]) start_time = time.monotonic() time.sleep(1) # Second sample cursor.execute(""" SELECT wal_bytes FROM pg_stat_wal; """) second_wal_bytes = float(cursor.fetchone()[0]) end_time = time.monotonic() elapsed_time = end_time - start_time generated_bytes = ( second_wal_bytes - first_wal_bytes ) wal_bytes_per_second = ( generated_bytes / elapsed_time ) wal_mb_per_second = round( wal_bytes_per_second / 1024 / 1024, 2 ) return { "database_name": data.dbname, "wal_generated_bytes_per_second": round(wal_bytes_per_second, 2), "wal_generated_mb_per_second": wal_mb_per_second } except psycopg2.Error as e: return { "database_name": data.dbname, "wal_generated_bytes_per_second": None, "wal_generated_mb_per_second": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() @app.websocket("/ws/database/wal-generated") async def database_wal_generated_monitor(websocket: WebSocket): await websocket.accept() try: raw_data = await websocket.receive_json() data = DatabaseRequest(**raw_data) while True: try: wal_result = await asyncio.wait_for( asyncio.to_thread( get_wal_generated_per_second, data ), timeout=5 ) result = { "database_name": data.dbname, "wal_generated_bytes_per_second": wal_result["wal_generated_bytes_per_second"], "wal_generated_mb_per_second": wal_result["wal_generated_mb_per_second"] } except asyncio.TimeoutError: result = { "database_name": data.dbname, "wal_generated_bytes_per_second": None, "wal_generated_mb_per_second": None } await websocket.send_json(result) await asyncio.sleep(2) except WebSocketDisconnect: print("WAL generated WebSocket disconnected") # GET WAL SIZE def get_wal_size(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 COALESCE(SUM(size), 0) FROM pg_ls_waldir(); """) wal_size_bytes = cursor.fetchone()[0] wal_size_mb = round( wal_size_bytes / 1024 / 1024, 2 ) return { "database_name": data.dbname, "wal_size_bytes": wal_size_bytes, "wal_size_mb": wal_size_mb } except psycopg2.Error as e: return { "database_name": data.dbname, "wal_size_bytes": None, "wal_size_mb": None, "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() #REST API ENDPOINT FOR WAL SIZE @app.post("/database/wal-size") def database_wal_size( data: DatabaseRequest, token_data: dict = Depends(verify_token) ): return get_wal_size(data) #REPLICATION LAG def get_replication_lag(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 pid, application_name, client_addr, state, sent_lsn, write_lsn, flush_lsn, replay_lsn, CASE WHEN replay_lag IS NOT NULL THEN ROUND( EXTRACT(EPOCH FROM replay_lag)::numeric * 1000, 2 ) ELSE NULL END AS lag_ms, CASE WHEN COALESCE( replay_lsn, flush_lsn, write_lsn, sent_lsn ) IS NOT NULL THEN ROUND( ( pg_wal_lsn_diff( pg_current_wal_lsn(), COALESCE( replay_lsn, flush_lsn, write_lsn, sent_lsn ) ) / 1024 / 1024 )::numeric, 2 ) ELSE NULL END AS lag_mb FROM pg_stat_replication ORDER BY application_name; """) rows = cursor.fetchall() replicas = [] for row in rows: replicas.append({ "pid": row["pid"], "application_name": row["application_name"], "client_addr": str(row["client_addr"]) if row["client_addr"] else None, "state": row["state"], "lag_ms": float(row["lag_ms"]) if row["lag_ms"] is not None else None, "lag_mb": float(row["lag_mb"]) if row["lag_mb"] is not None else None, "sent_lsn": str(row["sent_lsn"]) if row["sent_lsn"] else None, "write_lsn": str(row["write_lsn"]) if row["write_lsn"] else None, "flush_lsn": str(row["flush_lsn"]) if row["flush_lsn"] else None, "replay_lsn": str(row["replay_lsn"]) if row["replay_lsn"] else None }) return { "database_name": data.dbname, "replica_count": len(replicas), "replicas": replicas } except psycopg2.Error as e: return { "database_name": data.dbname, "replica_count": None, "replicas": [], "error": str(e) } finally: if cursor: cursor.close() if conn: conn.close() # WEBSOCKET FOR REPLICATION LAG MONITORING #OUTPUT #{"database_name":"postgres","replica_count":2,"replicas":[{"pid":51339,"application_name":"spk_postgres_node3_node1_from_node3","client_addr":"172.16.10.123","state":"streaming","lag_ms":null,"lag_mb":0.0,"sent_lsn":"0/56681A0","write_lsn":"0/56681A0","flush_lsn":"0/56681A0","replay_lsn":"0/56681A0"},{"pid":51328,"application_name":"spk_postgres_node3_node2_from_node3","client_addr":"172.16.10.110","state":"streaming","lag_ms":null,"lag_mb":0.0,"sent_lsn":"0/56681A0","write_lsn":"0/56681A0","flush_lsn":"0/56681A0","replay_lsn":"0/56681A0"}]} @app.websocket("/ws/database/replication-lag") async def database_replication_lag_monitor(websocket: WebSocket): await websocket.accept() try: raw_data = await websocket.receive_json() data = DatabaseRequest(**raw_data) while True: try: replication_result = await asyncio.wait_for( asyncio.to_thread( get_replication_lag, data ), timeout=5 ) result = replication_result except asyncio.TimeoutError: result = { "database_name": data.dbname, "replica_count": None, "replicas": [], "error": "Replication lag check timed out" } await websocket.send_json(result) await asyncio.sleep(2) except WebSocketDisconnect: print("Replication lag 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 )