From 3722e71219d91ca7ae20795794bc91babe280dcd Mon Sep 17 00:00:00 2001 From: George Date: Wed, 23 Sep 2026 07:28:36 +0530 Subject: [PATCH] feat: --ADDED API to get index size (REST) --ADDED API for tules read and write per second (WS) --ADDED API for get Dead tuples (REST API) --ADDED API for auto vaccume (REST API) --- main.py | 478 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 478 insertions(+) diff --git a/main.py b/main.py index 35fc926..1226ca1 100644 --- a/main.py +++ b/main.py @@ -1149,6 +1149,484 @@ def database_largest_tables( +# 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 +@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 +@app.websocket("/ws/database/tuple-activity") +async def database_tuple_activity_monitor(websocket: WebSocket): + + 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 + + FROM pg_stat_user_tables + WHERE n_dead_tup > 0 + ORDER BY n_dead_tup DESC; + """) + + rows = cursor.fetchall() + + total_dead_tuples = sum( + row["n_dead_tup"] or 0 + for row in rows + ) + + return { + "database_name": data.dbname, + "total_dead_tuples": total_dead_tuples, + + "tables": [ + { + "schema": row["schemaname"], + "table_name": row["table_name"], + "live_tuples": row["n_live_tup"], + "dead_tuples": row["n_dead_tup"], + + "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, + "total_dead_tuples": 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) + +