feat(database): add functions to analyze size of tables

agent/database.py

@@ -153,6 +153,21 @@ def modify_user_permissions(self, username: str, mode: str, permissions: dict |
 
         self._run_sql(queries_str, commit=True, allow_all_stmt_types=True)
 
+    def fetch_database_table_sizes(self) -> list[dict]:
+        """
+        SELECT * FROM INFORMATION_SCHEMA.INNODB_SYS_TABLESPACES
+        """
+        data = self._run_sql(
+            "SELECT table_name, data_length, index_length FROM INFORMATION_SCHEMA.TABLES", as_dict=True
+        )
+        if len(data) == 0:
+            return []
+        for d in data[0]["output"]:
+            d["data_length"] = int(d["data_length"])
+            d["index_length"] = int(d["index_length"])
+            d["total_size"] = d["data_length"] + d["index_length"]
+        return data
+
     # Private helper methods
     def _run_sql(  # noqa C901
         self, query: str, commit: bool = False, as_dict: bool = False, allow_all_stmt_types: bool = False
@@ -164,7 +179,7 @@ def _run_sql(  # noqa C901
         Args:
         query: SQL query string
         commit: True if you want to commit the changes. If commit is false, it will rollback the changes and
-                also wouldnt allow to run ddl, dcl or tcl queries
+                also wouldn't allow to run ddl, dcl or tcl queries
         as_dict: True if you want to return the result as a dictionary (like frappe.db.sql).
                 Otherwise it will return a dict of columns and data
         allow_all_stmt_types: True if you want to allow all type of sql statements

agent/site.py

@@ -893,6 +893,9 @@ def _fetch_database_table_schema(self):
             )
         return tables
 
+    def fetch_database_table_sizes(self, root_password: str):
+        return Database(self.host, 3306, "root", root_password, self.database).fetch_database_table_sizes()
+
     def get_database_table_indexes(self):
         command = f"""
         SELECT

agent/web.py

@@ -14,7 +14,7 @@
 from playhouse.shortcuts import model_to_dict
 
 from agent.builder import ImageBuilder, get_image_build_context_directory
-from agent.database import JSONEncoderForSQLQueryResult
+from agent.database import Database, JSONEncoderForSQLQueryResult
 from agent.database_server import DatabaseServer
 from agent.exceptions import BenchNotExistsException, SiteNotExistsException
 from agent.job import JobModel, connection
@@ -557,6 +557,11 @@ def fetch_database_table_schema(bench, site):
     return {"job": job}
 
 
+@application.route("/benches/<string:bench>/sites/<string:site>/database/size", methods=["POST"])
+@validate_bench_and_site
+def fetch_database_table_sizes(bench, site):
+    return Response(json.dumps(Server().benches[bench].sites[site].fetch_database_table_sizes()))
+
 @application.route("/benches/<string:bench>/sites/<string:site>/database/query/execute", methods=["POST"])
 @validate_bench_and_site
 def run_sql(bench, site):