diff --git a/main.py b/main.py index 98c43d6..35fc926 100644 --- a/main.py +++ b/main.py @@ -981,16 +981,171 @@ async def database_locks_deadlocks_monitor(websocket: WebSocket): 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 +@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 +@app.post("/database/largest-tables") +def database_largest_tables( + data: DatabaseRequest, + token_data: dict = Depends(verify_token) +): - - - - + return get_largest_tables( + data, + limit=10 + )