Merge pull request #32 from mingjie-zhang1998/feature/table_compute_tool

1. 新增表碎片及大表数据统计；2. 当查询系统表时，创建账号需要赋多个数据库的权限

Author: cnwenf

README.md

@@ -139,6 +139,8 @@ Add the following configuration to the MCP client configuration file:
 * `explain_sql`: Execute sql `explain` and return sql result.
 * `show_engine_innodb_status`: Execute sql `show engine innodb status` and return sql result.
 * `show_create_table`: Execute sql `show create table` and return sql result.
+* `show_largest_table`: Query the top few tables with the highest space occupancy.
+* `show_largest_table_fragment`: Query the tables with the largest table fragments.
 * `query_sql`: Execute read-only sql and return sql result.
 
 ### Toolsets

README_CN.md

@@ -136,6 +136,8 @@ git clone https://github.com/aliyun/alibabacloud-rds-openapi-mcp-server.git
 * `explain_sql`：执行 SQL `explain` 命令并返回 SQL 执行计划结果。
 * `show_engine_innodb_status`：执行 SQL `show engine innodb status` 命令并返回 SQL 执行结果。
 * `show_create_table`：执行 SQL `show create table` 命令并返回 SQL 执行结果。
+* `show_largest_table`: 查询空间占用率最高的前几张表。
+* `show_largest_table_fragment`: 查询空间碎片率最高的前几张表。
 * `query_sql`：执行只读 SQL 语句并返回 SQL 执行结果。
 
 

src/alibabacloud_rds_openapi_mcp_server/db_service.py

@@ -49,10 +49,12 @@ class DBService:
     def __init__(self,
                  region_id,
                  instance_id,
-                 database=None, ):
+                 database=None,
+                 privilege_databases=None):
         self.instance_id = instance_id
         self.database = database
         self.region_id = region_id
+        self.privilege_databases = privilege_databases
 
         self.__db_type = None
         self.__account_name, self.__account_password = get_rds_account()
@@ -126,11 +128,22 @@ def _create_temp_account(self):
         self.__client.create_account(request)
 
     def _grant_privilege(self):
+        if self.privilege_databases:
+            dbname = self.privilege_databases
+            cnt = len(self.privilege_databases.split(','))
+            if self.db_type.lower() in ('mysql', 'postgresql'):
+                privilege = ",".join(['ReadOnly' for i in range(cnt)])
+            else:
+                privilege = ",".join(['DBOwner' for i in range(cnt)])
+        else:
+            dbname = self.database
+            privilege = "ReadOnly" if self.db_type.lower() in ('mysql', 'postgresql') else "DBOwner"
+
         req = rds_20140815_models.GrantAccountPrivilegeRequest(
             dbinstance_id=self.instance_id,
             account_name=self.account_name,
-            dbname=self.database,
-            account_privilege="ReadOnly" if self.db_type.lower() in ('mysql', 'postgresql') else "DBOwner"
+            dbname=dbname,
+            account_privilege=privilege
         )
         self.__client.grant_account_privilege(req)
 

src/alibabacloud_rds_openapi_mcp_server/server.py

@@ -16,9 +16,11 @@
 import time
 from datetime import datetime
 from typing import Dict, Any, List, Optional
+from pydantic import Field
 
 from alibabacloud_bssopenapi20171214 import models as bss_open_api_20171214_models
 from alibabacloud_openapi_util.client import Client as OpenApiUtilClient
+from alibabacloud_rds20140815.client import Client as RdsClient
 from alibabacloud_rds20140815 import models as rds_20140815_models
 from alibabacloud_tea_openapi import models as open_api_models
 from alibabacloud_tea_util import models as util_models
@@ -1642,11 +1644,111 @@ async def query_sql(
         the sql result.
     """
     try:
-        async with DBService(region_id, dbinstance_id, db_name) as service:
+        if db_name == 'information_schema':
+            client = get_rds_client(region_id)
+            databases = _get_db_instance_databases_str(client, region_id, dbinstance_id)
+        else:
+            databases = None
+        async with DBService(region_id, dbinstance_id, db_name, databases) as service:
             return await service.execute_sql(sql=sql)
     except Exception as e:
         logger.error(f"Error occurred: {str(e)}")
         raise e
+    
+@mcp.tool(annotations=READ_ONLY_TOOL,
+          description="Query the top few tables with the highest space occupancy")
+async def show_largest_table(
+        region_id: str = Field(description="The region ID of the RDS instance."),
+        dbinstance_id: str = Field(description="The ID of the RDS instance."),
+        topK: int = Field(default=5, description="To display the number of the top few largest tables")
+) -> str:
+    """
+    Returns:
+        The data table with the largest storage space.
+    """
+    try:
+        client = get_rds_client(region_id)
+        request = rds_20140815_models.DescribeDBInstanceAttributeRequest(dbinstance_id=dbinstance_id)
+        response = client.describe_dbinstance_attribute(request)
+        result = response.body.to_map()
+
+        db_type = result['Items']['DBInstanceAttribute'][0]['Engine']
+        db_version = result['Items']['DBInstanceAttribute'][0]['EngineVersion']
+        if db_type.lower() != 'mysql' and db_version not in ["5.6", "5.7", "8.0"]:
+            return f"Unsupported db version {db_version}."
+
+        db_name = 'information_schema'
+        databases = _get_db_instance_databases_str(client, region_id, dbinstance_id)
+        # 构建SQL
+        # 限制展示最大表数量的上限
+        topK = min(topK, 100)
+        base_query = f"""
+                SELECT 
+                  table_schema AS '数据库',
+                  table_name AS '表名',
+                  ROUND((data_length + index_length) / 1024 / 1024, 2) AS '总大小(MB)',
+                  ROUND(INDEX_LENGTH / 1024 / 1024, 2) AS '索引大小(MB)',
+                  table_rows AS '行数'
+                FROM information_schema.TABLES
+                ORDER BY (data_length + index_length) DESC
+                Limit {topK};
+            """
+
+        async with DBService(region_id, dbinstance_id, db_name, databases) as service:
+            return await service.execute_sql(sql=base_query)
+    except Exception as e:
+        logger.exception("show largest table failed.")
+        logger.error(f"Error occurred: {str(e)}")
+        raise e
+
+
+@mcp.tool(annotations=READ_ONLY_TOOL,
+          description="Query the tables with the largest table fragments")
+async def show_largest_table_fragment(
+        region_id: str = Field(description="The region ID of the RDS instance."),
+        dbinstance_id: str = Field(description="The ID of the RDS instance."),
+        topK: int = Field(default=5, description="To display the number of the top few largest tables")
+) -> str:
+    """
+    Returns:
+        The largest fragmented data table.
+    """
+    try:
+        client = get_rds_client(region_id)
+        request = rds_20140815_models.DescribeDBInstanceAttributeRequest(dbinstance_id=dbinstance_id)
+        response = client.describe_dbinstance_attribute(request)
+        result = response.body.to_map()
+
+        db_type = result['Items']['DBInstanceAttribute'][0]['Engine']
+        db_version = result['Items']['DBInstanceAttribute'][0]['EngineVersion']
+        if db_type.lower() != 'mysql' and db_version not in ["5.6", "5.7", "8.0"]:
+            return f"Unsupported db version {db_version}."
+
+        db_name = 'information_schema'
+        databases = _get_db_instance_databases_str(client, region_id, dbinstance_id)
+        # 构建SQL
+        # 限制展示最大表数量的上限
+        topK = min(topK, 100)
+        base_query = f"""
+                SELECT 
+                  ENGINE,
+                  TABLE_SCHEMA,
+                  TABLE_NAME,
+                  ROUND(DATA_LENGTH / 1024 / 1024,2) AS '总大小(MB)',
+                  ROUND(INDEX_LENGTH / 1024 / 1024,2) AS '索引大小(MB)',
+                  ROUND(DATA_FREE / 1024 / 1024,2) AS '碎片大小(MB)'
+                FROM information_schema.TABLES
+                WHERE DATA_FREE > 0
+                ORDER BY DATA_FREE DESC
+                Limit {topK};
+            """
+
+        async with DBService(region_id, dbinstance_id, db_name, databases) as service:
+            return await service.execute_sql(sql=base_query)
+    except Exception as e:
+        logger.exception("show largest table failed.")
+        logger.error(f"Error occurred: {str(e)}")
+        raise e
 
 async def health_check(request: Request) -> JSONResponse:
     """Health check endpoint for container health monitoring and load balancer probes."""
@@ -1763,6 +1865,24 @@ def _parse_groups_from_source(source: str | None) -> List[str]:
                 expanded_groups.append(g_to_add)
     return expanded_groups or [DEFAULT_TOOL_GROUP]
 
+def _get_db_instance_databases_str(
+        client: RdsClient,
+        region_id: str,
+        db_instance_id: str
+) -> str:
+    try:
+        client = get_rds_client(region_id)
+        request = rds_20140815_models.DescribeDatabasesRequest(
+            dbinstance_id=db_instance_id
+        )
+        response = client.describe_databases(request)
+        databases_info = response.body.databases.database
+
+        result = ",".join([database.dbname for database in databases_info])
+        return result
+    except Exception as e:
+        raise e
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

src/alibabacloud_rds_openapi_mcp_server/utils.py

@@ -56,7 +56,21 @@
         "DiskUsage": ["disk_usage"],
         "IOPSUsage": ["data_iops_usage"],
         "IOBytesPS": ["data_io_bytes_ps"],
-        "MdlLockSession": ["mdl_lock_session"]
+        "MdlLockSession": ["mdl_lock_session"],
+        "RelayLogSize": ["relay_log_size"],
+        "UndoLogSize": ["undolog_size"],
+        "RedoLog_Size": ["redolog_size"],
+        "TempFileSize": ["temp_file_size"],
+        "InsSize": ["ins_size"],
+        "SysDataSize": ["sys_data_size"],
+        "GeneralLogSize": ["general_log_size"],
+        "SlowLogSize": ["slow_log_size"],
+        "BinlogSize": ["binlog_size"],
+        "UserDataSize": ["user_data_size"],
+        "InnodbRowsInsert":["Innodb_rows_inserted_ps"],
+        "MemCpuUsage": ["cpu_usage"],
+        "QPS": ["qps"],
+        "SLowSQL": ["slow_sql"]
     }
 }