Merge pull request aliyun#24 from junliang-xjl/feature/query_sql

Optimize query sql

Author: cnwenf

src/alibabacloud_rds_openapi_mcp_server/db_service.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import random
 import socket
@@ -60,24 +61,24 @@ def __init__(self,
  self.__client = get_rds_client(region_id)
  self.__db_conn = None
 
- def __enter__(self):
- self._get_db_instance_info()
+ async def __aenter__(self):
+ await asyncio.to_thread(self._get_db_instance_info)
  if not self.__account_name or not self.__account_password:
- self._create_temp_account()
+ await asyncio.to_thread(self._create_temp_account)
  if self.database:
- self._grant_privilege()
+ await asyncio.to_thread(self._grant_privilege)
  else:
  self.account_name = self.__account_name
  self.account_password = self.__account_password
  self.__db_conn = DBConn(self)
- self.__db_conn.connect()
+ await asyncio.to_thread(self.__db_conn.connect)
  return self
 
- def __exit__(self, exc_type, exc_val, exc_tb):
+ async def __aexit__(self, exc_type, exc_val, exc_tb):
  if self.__db_conn is not None:
- self.__db_conn.close()
+ await asyncio.to_thread(self.__db_conn.close)
  if not self.__account_name or not self.__account_password:
- self._delete_account()
+ await asyncio.to_thread(self._delete_account)
  self.__client = None
 
  def _get_db_instance_info(self):
@@ -142,8 +143,8 @@ def _delete_account(self):
  )
  self.__client.delete_account(req)
 
- def execute_sql(self, sql):
- return self.__db_conn.execute_sql(sql)
+ async def execute_sql(self, sql):
+ return await asyncio.to_thread(self.__db_conn.execute_sql, sql)
 
  @property
  def user(self):
@@ -214,6 +215,7 @@ def execute_sql(self, sql):
  result = [dict(zip(columns, row)) for row in rows]
  else:
  result = []
+ cursor.close()
  try:
  return json.dumps(result, ensure_ascii=False)
  except Exception as e:

src/alibabacloud_rds_openapi_mcp_server/server.py

@@ -1559,8 +1559,8 @@ async def show_engine_innodb_status(
  the sql result.
  """
  try:
- with DBService(region_id, dbinstance_id) as service:
- return service.execute_sql("show engine innodb status")
+ async with DBService(region_id, dbinstance_id) as service:
+ return await service.execute_sql("show engine innodb status")
  except Exception as e:
  logger.error(f"Error occurred: {str(e)}")
  raise e
@@ -1584,8 +1584,8 @@ async def show_create_table(
  the sql result.
  """
  try:
- with DBService(region_id, dbinstance_id, db_name) as service:
- return service.execute_sql(f"show create table {db_name}.{table_name}")
+ async with DBService(region_id, dbinstance_id, db_name) as service:
+ return await service.execute_sql(f"show create table {db_name}.{table_name}")
  except Exception as e:
  logger.error(f"Error occurred: {str(e)}")
  raise e
@@ -1608,8 +1608,8 @@ async def explain_sql(
  the sql execute plan.
  """
  try:
- with DBService(region_id, dbinstance_id, db_name) as service:
- return service.execute_sql(f"explain {sql}")
+ async with DBService(region_id, dbinstance_id, db_name) as service:
+ return await service.execute_sql(f"explain {sql}")
  except Exception as e:
  logger.error(f"Error occurred: {str(e)}")
  raise e
@@ -1633,8 +1633,8 @@ async def query_sql(
  the sql result.
  """
  try:
- with DBService(region_id, dbinstance_id, db_name) as service:
- return service.execute_sql(sql=sql)
+ async with DBService(region_id, dbinstance_id, db_name) as service:
+ return await service.execute_sql(sql=sql)
  except Exception as e:
  logger.error(f"Error occurred: {str(e)}")
  raise e