mcp-dbutils 0.21.0__py3-none-any.whl → 0.23.0__py3-none-any.whl

mcp_dbutils/base.py

@@ -3,7 +3,7 @@
 import json
 from abc import ABC, abstractmethod
 from contextlib import asynccontextmanager
-from datetime import datetime
+from datetime import datetime, timedelta
 from importlib.metadata import metadata
 from typing import Any, AsyncContextManager, Dict
 
@@ -474,6 +474,29 @@ class ConnectionServer:
                     },
                     "required": ["connection", "sql"],
                 },
+                annotations={
+                    "examples": [
+                        {
+                            "input": {
+                                "connection": "example_db",
+                                "sql": "SELECT id, name, email FROM users LIMIT 10"
+                            },
+                            "output": "Results showing first 10 users with their IDs, names, and email addresses"
+                        },
+                        {
+                            "input": {
+                                "connection": "example_db",
+                                "sql": "SELECT department, COUNT(*) as employee_count FROM employees GROUP BY department ORDER BY employee_count DESC"
+                            },
+                            "output": "Results showing departments and their employee counts in descending order"
+                        }
+                    ],
+                    "usage_tips": [
+                        "Always use SELECT statements only - other SQL operations are not permitted",
+                        "Use LIMIT to restrict large result sets",
+                        "For complex queries, consider using dbutils-explain-query first to understand query execution plan"
+                    ]
+                }
             ),
             types.Tool(
                 name="dbutils-list-tables",
@@ -488,6 +511,19 @@ class ConnectionServer:
                     },
                     "required": ["connection"],
                 },
+                annotations={
+                    "examples": [
+                        {
+                            "input": {"connection": "example_db"},
+                            "output": "List of tables in the example_db database with their URIs and descriptions"
+                        }
+                    ],
+                    "usage_tips": [
+                        "Use this tool first when exploring a new database to understand its structure",
+                        "After listing tables, use dbutils-describe-table to get detailed information about specific tables",
+                        "Table URIs can be used with other database tools for further operations"
+                    ]
+                }
             ),
             types.Tool(
                 name="dbutils-describe-table",
@@ -999,5 +1035,8 @@ class ConnectionServer:
         """Run server"""
         async with mcp.server.stdio.stdio_server() as streams:
             await self.server.run(
-                streams[0], streams[1], self.server.create_initialization_options()
+                streams[0],
+                streams[1],
+                self.server.create_initialization_options(),
+                read_timeout_seconds=timedelta(seconds=30)  # 设置30秒超时
             )

mcp_dbutils/mysql/handler.py

@@ -537,3 +537,15 @@ class MySQLHandler(ConnectionHandler):
         """Cleanup resources"""
         # Log final stats before cleanup
         self.log("info", f"Final MySQL handler stats: {self.stats.to_dict()}")
+
+        # 主动关闭连接
+        if hasattr(self, '_connection') and self._connection:
+            try:
+                self.log("debug", "Closing MySQL connection")
+                self._connection.close()
+                self._connection = None
+            except Exception as e:
+                self.log("warning", f"Error closing MySQL connection: {str(e)}")
+
+        # 清理其他资源
+        self.log("debug", "MySQL handler cleanup complete")

mcp_dbutils/mysql/server.py

@@ -27,12 +27,12 @@ class MySQLServer(ConnectionServer):
             conn_params = config.get_connection_params()
             masked_params = config.get_masked_connection_info()
             self.log("debug", f"正在连接数据库，参数: {masked_params}")
-            
+
             # 测试连接
             test_conn = mysql.connector.connect(**conn_params)
             test_conn.close()
             self.log("info", "测试连接成功")
-            
+
             # 创建连接池配置
             pool_config = {
                 'pool_name': 'mypool',
@@ -51,7 +51,7 @@ class MySQLServer(ConnectionServer):
             conn = self.pool.get_connection()
             with conn.cursor(dictionary=True) as cur:  # NOSONAR - dictionary参数是正确的，用于返回字典格式的结果
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         table_name,
                         table_comment as description
                     FROM information_schema.tables
@@ -81,7 +81,7 @@ class MySQLServer(ConnectionServer):
             with conn.cursor(dictionary=True) as cur:  # NOSONAR - dictionary参数是正确的，用于返回字典格式的结果
                 # 获取列信息
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         column_name,
                         data_type,
                         is_nullable,
@@ -157,6 +157,9 @@ class MySQLServer(ConnectionServer):
 
         connection = arguments.get("connection")
         conn = None
+        results = []
+        columns = []
+
         try:
             if connection and self.config_path:
                 # 使用指定的数据库连接
@@ -177,19 +180,22 @@ class MySQLServer(ConnectionServer):
                     cur.execute(sql)
                     results = cur.fetchall()
                     columns = [desc[0] for desc in cur.description]
-                    result_text = str({
-                        'type': 'mysql',
-                        'config_name': connection or 'default',
-                        'query_result': {
-                            'columns': columns,
-                            'rows': results,
-                            'row_count': len(results)
-                        }
-                    })
-                    self.log("info", f"查询完成，返回{len(results)}行结果")
-                    return [types.TextContent(type="text", text=result_text)]
                 finally:
                     cur.execute("ROLLBACK")
+
+            # 处理结果（在连接操作完成后）
+            result_text = str({
+                'type': 'mysql',
+                'config_name': connection or 'default',
+                'query_result': {
+                    'columns': columns,
+                    'rows': results,
+                    'row_count': len(results)
+                }
+            })
+            self.log("info", f"查询完成，返回{len(results)}行结果")
+            return [types.TextContent(type="text", text=result_text)]
+
         except Exception as e:
             error = f"查询执行失败: {str(e)}"
             error_msg = str({
@@ -200,12 +206,28 @@ class MySQLServer(ConnectionServer):
             self.log("error", error_msg)
             return [types.TextContent(type="text", text=error_msg)]
         finally:
+            # 确保连接始终被关闭
             if conn:
-                conn.close()  # 关闭连接（连接池会自动处理）
+                try:
+                    conn.close()  # 关闭连接（连接池会自动处理）
+                except Exception as e:
+                    self.log("warning", f"关闭连接时出错: {str(e)}")
 
     async def cleanup(self):
         """清理资源"""
         if hasattr(self, 'pool'):
-            self.log("info", "关闭连接池")
-            # MySQL连接池没有直接的closeall方法
-            # 当对象被销毁时，连接池会自动关闭
+            self.log("info", "关闭MySQL连接池")
+            # MySQL连接池没有直接的closeall方法，但我们可以主动关闭连接
+            try:
+                # 尝试获取并关闭所有活动连接
+                for _ in range(5):  # 假设最多有5个连接在池中
+                    try:
+                        conn = self.pool.get_connection()
+                        conn.close()
+                    except Exception as e:
+                        # 如果没有更多连接或出现其他错误，跳出循环
+                        self.log("debug", f"连接池清理: {str(e)}")
+                        break
+                self.log("info", "MySQL连接池清理完成")
+            except Exception as e:
+                self.log("warning", f"清理MySQL连接池时出错: {str(e)}")

mcp_dbutils/postgres/handler.py

@@ -610,3 +610,15 @@ class PostgreSQLHandler(ConnectionHandler):
         """Cleanup resources"""
         # Log final stats before cleanup
         self.log("info", f"Final PostgreSQL handler stats: {self.stats.to_dict()}")
+
+        # 主动关闭连接
+        if hasattr(self, '_connection') and self._connection:
+            try:
+                self.log("debug", "Closing PostgreSQL connection")
+                self._connection.close()
+                self._connection = None
+            except Exception as e:
+                self.log("warning", f"Error closing PostgreSQL connection: {str(e)}")
+
+        # 清理其他资源
+        self.log("debug", "PostgreSQL handler cleanup complete")

mcp_dbutils/postgres/server.py

@@ -148,9 +148,14 @@ class PostgreSQLServer(ConnectionServer):
         # 仅允许SELECT语句
         if not sql.lower().startswith("select"):
             raise ValueError("仅支持SELECT查询")
+
         connection = arguments.get("connection")
         use_pool = True
         conn = None
+        results = []
+        columns = []
+        formatted_results = []
+
         try:
             if connection and self.config_path:
                 # 使用指定的数据库连接
@@ -163,6 +168,7 @@ class PostgreSQLServer(ConnectionServer):
             else:
                 # 使用现有连接池
                 conn = self.pool.getconn()
+
             self.log("info", f"执行查询: {sql}")
             with conn.cursor() as cur:
                 # 启动只读事务
@@ -172,19 +178,22 @@ class PostgreSQLServer(ConnectionServer):
                     results = cur.fetchall()
                     columns = [desc[0] for desc in cur.description]
                     formatted_results = [dict(zip(columns, row)) for row in results]
-                    result_text = str({
-                        'type': 'postgres',
-                        'config_name': connection or 'default',
-                        'query_result': {
-                            'columns': columns,
-                            'rows': formatted_results,
-                            'row_count': len(results)
-                        }
-                    })
-                    self.log("info", f"查询完成，返回{len(results)}行结果")
-                    return [types.TextContent(type="text", text=result_text)]
                 finally:
                     cur.execute("ROLLBACK")
+
+            # 处理结果（在连接操作完成后）
+            result_text = str({
+                'type': 'postgres',
+                'config_name': connection or 'default',
+                'query_result': {
+                    'columns': columns,
+                    'rows': formatted_results,
+                    'row_count': len(results)
+                }
+            })
+            self.log("info", f"查询完成，返回{len(results)}行结果")
+            return [types.TextContent(type="text", text=result_text)]
+
         except Exception as e:
             if isinstance(e, psycopg2.Error):
                 error = f"查询执行失败: [Code: {e.pgcode}] {e.pgerror or str(e)}"
@@ -198,11 +207,15 @@ class PostgreSQLServer(ConnectionServer):
             self.log("error", error_msg)
             return [types.TextContent(type="text", text=error_msg)]
         finally:
+            # 确保连接始终被正确处理
             if conn:
-                if use_pool:
-                    self.pool.putconn(conn)
-                else:
-                    conn.close()
+                try:
+                    if use_pool:
+                        self.pool.putconn(conn)
+                    else:
+                        conn.close()
+                except Exception as e:
+                    self.log("warning", f"关闭连接时出错: {str(e)}")
     async def cleanup(self):
         """清理资源"""
         if hasattr(self, 'pool'):

mcp_dbutils/sqlite/handler.py

@@ -475,3 +475,15 @@ class SQLiteHandler(ConnectionHandler):
         """Cleanup resources"""
         # Log final stats
         self.log("info", f"Final SQLite handler stats: {self.stats.to_dict()}")
+
+        # 主动关闭连接
+        if hasattr(self, '_connection') and self._connection:
+            try:
+                self.log("debug", "Closing SQLite connection")
+                self._connection.close()
+                self._connection = None
+            except Exception as e:
+                self.log("warning", f"Error closing SQLite connection: {str(e)}")
+
+        # 清理其他资源
+        self.log("debug", "SQLite handler cleanup complete")

mcp_dbutils/sqlite/server.py

@@ -141,8 +141,13 @@ class SQLiteServer(ConnectionServer):
             raise ValueError("仅支持SELECT查询")
 
         conn = None
+        connection = arguments.get("connection")
+        config_name = connection if isinstance(connection, str) else 'default'
+        results = []
+        columns = []
+        formatted_results = []
+
         try:
-            connection = arguments.get("connection")
             if connection and self.config_path:
                 # 使用指定的数据库连接
                 config = SQLiteConfig.from_yaml(self.config_path, connection)
@@ -159,28 +164,24 @@ class SQLiteServer(ConnectionServer):
                 self.log("info", f"执行查询: {sql}")
                 cursor = conn.execute(sql)
                 results = cursor.fetchall()
-
                 columns = [desc[0] for desc in cursor.description]
                 formatted_results = [dict(zip(columns, row)) for row in results]
 
-                # 使用更通用的方法确定配置名称
-                config_name = connection if isinstance(connection, str) else 'default'
-                result_text = json.dumps({
-                    'type': 'sqlite',
-                    'config_name': config_name,
-                    'query_result': {
-                        'columns': columns,
-                        'rows': formatted_results,
-                        'row_count': len(results)
-                    }
-                })
-
-                self.log("info", f"查询完成，返回{len(results)}行结果")
-                return [types.TextContent(type="text", text=result_text)]
+            # 处理结果（在连接操作完成后）
+            result_text = json.dumps({
+                'type': 'sqlite',
+                'config_name': config_name,
+                'query_result': {
+                    'columns': columns,
+                    'rows': formatted_results,
+                    'row_count': len(results)
+                }
+            })
+
+            self.log("info", f"查询完成，返回{len(results)}行结果")
+            return [types.TextContent(type="text", text=result_text)]
 
         except sqlite3.Error as e:
-            # 使用更通用的方法确定配置名称
-            config_name = connection if isinstance(connection, str) else 'default'
             error_msg = json.dumps({
                 'type': 'sqlite',
                 'config_name': config_name,
@@ -188,6 +189,13 @@ class SQLiteServer(ConnectionServer):
             })
             self.log("error", error_msg)
             return [types.TextContent(type="text", text=error_msg)]
+        finally:
+            # 确保连接被正确关闭（如果不是使用with语句）
+            if conn and not isinstance(connection, str):
+                try:
+                    conn.close()
+                except Exception as e:
+                    self.log("warning", f"关闭连接时出错: {str(e)}")
 
     async def cleanup(self):
         """清理资源"""

{mcp_dbutils-0.21.0.dist-info → mcp_dbutils-0.23.0.dist-info}/METADATA

@@ -1,16 +1,17 @@
 Metadata-Version: 2.4
 Name: mcp-dbutils
-Version: 0.21.0
+Version: 0.23.0
 Summary: MCP Database Utilities Service
 Author: Dong Hao
 License-Expression: MIT
 License-File: LICENSE
 Requires-Python: >=3.10
-Requires-Dist: mcp>=1.2.1
+Requires-Dist: mcp>=1.7.1
 Requires-Dist: mysql-connector-python>=8.2.0
 Requires-Dist: psycopg2-binary>=2.9.10
 Requires-Dist: python-dotenv>=1.0.1
 Requires-Dist: pyyaml>=6.0.2
+Requires-Dist: typer>=0.9.0
 Provides-Extra: test
 Requires-Dist: aiosqlite>=0.19.0; extra == 'test'
 Requires-Dist: docker>=7.0.0; extra == 'test'

{mcp_dbutils-0.21.0.dist-info → mcp_dbutils-0.23.0.dist-info}/RECORD

@@ -1,22 +1,22 @@
 mcp_dbutils/__init__.py,sha256=6LLccQv7je2L4IpY_I3OzSJZcK32VUDJv2IY31y6eYg,1900
-mcp_dbutils/base.py,sha256=7UCw9gBZdzS_LnW3_QflWPfEV00eUFpUMDoaQJCN2aA,40012
+mcp_dbutils/base.py,sha256=A9p79jqKxzCmOErFTwabQtUzIftkl1_qMFQ2rlnFB3g,42147
 mcp_dbutils/config.py,sha256=bmXpOd1fyYfoyUS75I035ChT6t3wP5AyEnJ06e2ZS2o,1848
 mcp_dbutils/log.py,sha256=mqxi6I_IL-MF1F_pxBtnYZQKOHbGBJ74gsvZHVelr1w,823
 mcp_dbutils/stats.py,sha256=wMqWPfGnEOg9v5YBtTsARV-1YsFUMM_pKdzitzSU9x4,7137
 mcp_dbutils/mysql/__init__.py,sha256=gNhoHaxK1qhvMAH5AVl1vfV1rUpcbV9KZWUQb41aaQk,129
 mcp_dbutils/mysql/config.py,sha256=BTPPFqlhoTp7EBFIeLJZh8x6bCn3q9NivHYz9yZHziw,9820
-mcp_dbutils/mysql/handler.py,sha256=Nu2YySImnYKOSLYe04ipHYWW1kNndFKiP-gdX5L5wpU,22325
-mcp_dbutils/mysql/server.py,sha256=1bWAu7qHYXVeTZu4wdEpS6gSVB0RoXKI3Smy_ix-y8A,8586
+mcp_dbutils/mysql/handler.py,sha256=CJl6Mcs8kloojtS7i-2yD_-InRSmtU09LfYnHDZDhn0,22783
+mcp_dbutils/mysql/server.py,sha256=4kWkrZby8fZq-Em65Fobs5KDDPJnrOxbsK1zLK1jzp0,9412
 mcp_dbutils/postgres/__init__.py,sha256=-2zYuEJEQ2AMvmGhH5Z_umerSvt7S4xOa_XV4wgvGfI,154
 mcp_dbutils/postgres/config.py,sha256=NyQOVhkXJ1S-JD0w-ePNjTKI1Ja-aZQkDUdHi6U7Vl4,7752
-mcp_dbutils/postgres/handler.py,sha256=rHSvgz93UXIfnVzHFOTGhgnMMZy-s2tB6nR40p2G26Q,24710
-mcp_dbutils/postgres/server.py,sha256=_CiJC9PitpI1NB99Q1Bcs5TYADNgDpYMwv88fRHQunE,8640
+mcp_dbutils/postgres/handler.py,sha256=SnzgXjcO3iuvpQ-dzggNOF0bewAKB5BwL3QS11a5qWw,25183
+mcp_dbutils/postgres/server.py,sha256=qden51sQu5Ht2ibl2jSDyxg7r_N9qDVhg83Pc-q98yc,8887
 mcp_dbutils/sqlite/__init__.py,sha256=fK_3-WylCBYpBAzwuopi8hlwoIGJm2TPAlwcPWG46I0,134
 mcp_dbutils/sqlite/config.py,sha256=j67TJ8mQJ2D886MthSa-zYMtvUUYyyxYLMlNxkYoqZE,4509
-mcp_dbutils/sqlite/handler.py,sha256=nCgBeBp3zjpE2HNohMbh3Jpz5tNeMczt5K87JOhVWzY,19320
-mcp_dbutils/sqlite/server.py,sha256=jqpE8d9vJETMs5xYGB7P0tvNDPes6Yn5ZM_iCCF7Tv4,7181
-mcp_dbutils-0.21.0.dist-info/METADATA,sha256=0ZPTcDD34ecEh38aATtkvhGRj4Z1GNwX1Dp2cE9dSMY,8413
-mcp_dbutils-0.21.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-mcp_dbutils-0.21.0.dist-info/entry_points.txt,sha256=XTjt0QmYRgKOJQT6skR9bp1EMUfIrgpHeZJPZ3CJffs,49
-mcp_dbutils-0.21.0.dist-info/licenses/LICENSE,sha256=1A_CwpWVlbjrKdVEYO77vYfnXlW7oxcilZ8FpA_BzCI,1065
-mcp_dbutils-0.21.0.dist-info/RECORD,,
+mcp_dbutils/sqlite/handler.py,sha256=jH0xi_0PxiIBIVnXAGkfNoTVazTr7ol_P1XgBVDRkrU,19781
+mcp_dbutils/sqlite/server.py,sha256=EBKNKz_wTvChwg6BZlvZIBA1H5mmE2NiNEMOgu_CMy4,7373
+mcp_dbutils-0.23.0.dist-info/METADATA,sha256=8GKU0Z9Z-TQFAVJqjeouk5VRNrZaJfedBktVujVg3b4,8441
+mcp_dbutils-0.23.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+mcp_dbutils-0.23.0.dist-info/entry_points.txt,sha256=XTjt0QmYRgKOJQT6skR9bp1EMUfIrgpHeZJPZ3CJffs,49
+mcp_dbutils-0.23.0.dist-info/licenses/LICENSE,sha256=1A_CwpWVlbjrKdVEYO77vYfnXlW7oxcilZ8FpA_BzCI,1065
+mcp_dbutils-0.23.0.dist-info/RECORD,,