diff --git a/main.py b/main.py index 1226ca1..804082e 100644 --- a/main.py +++ b/main.py @@ -1626,19 +1626,389 @@ def database_autovacuum_activity( 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")