mcp-dbutils 0.16.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

mcp_dbutils/mysql/config.py

@@ -77,15 +77,19 @@ class MySQLConfig(ConnectionConfig):
     ssl: Optional[SSLConfig] = None
 
     @classmethod
-    def from_yaml(cls, yaml_path: str, db_name: str, local_host: Optional[str] = None) -> 'MySQLConfig':
-        """Create configuration from YAML file
-
+    def _validate_connection_config(cls, configs: dict, db_name: str) -> dict:
+        """验证连接配置是否有效
+        
         Args:
-            yaml_path: Path to YAML configuration file
-            db_name: Connection configuration name to use
-            local_host: Optional local host address
+            configs: 配置字典
+            db_name: 连接名称
+            
+        Returns:
+            dict: 数据库配置
+            
+        Raises:
+            ValueError: 如果配置无效
         """
-        configs = cls.load_yaml_config(yaml_path)
         if not db_name:
             raise ValueError("Connection name must be specified")
         if db_name not in configs:
@@ -103,57 +107,121 @@ class MySQLConfig(ConnectionConfig):
             raise ValueError("User must be specified in connection configuration")
         if not db_config.get('password'):
             raise ValueError("Password must be specified in connection configuration")
+            
+        return db_config
+        
+    @classmethod
+    def _create_config_from_url(cls, db_config: dict, local_host: Optional[str] = None) -> 'MySQLConfig':
+        """从URL创建配置
+        
+        Args:
+            db_config: 数据库配置
+            local_host: 可选的本地主机地址
+            
+        Returns:
+            MySQLConfig: 配置对象
+        """
+        # Parse URL for connection parameters
+        params = parse_url(db_config['url'])
+        config = cls(
+            database=params['database'],
+            user=db_config['user'],
+            password=db_config['password'],
+            host=params['host'],
+            port=params['port'],
+            charset=params['charset'],
+            local_host=local_host,
+            url=db_config['url'],
+            ssl=params.get('ssl')
+        )
+        return config
+        
+    @classmethod
+    def _create_config_from_params(cls, db_config: dict, local_host: Optional[str] = None) -> 'MySQLConfig':
+        """从参数创建配置
+        
+        Args:
+            db_config: 数据库配置
+            local_host: 可选的本地主机地址
+            
+        Returns:
+            MySQLConfig: 配置对象
+            
+        Raises:
+            ValueError: 如果缺少必需参数或SSL配置无效
+        """
+        if not db_config.get('database'):
+            raise ValueError("MySQL database name must be specified in configuration")
+        if not db_config.get('host'):
+            raise ValueError("Host must be specified in connection configuration")
+        if not db_config.get('port'):
+            raise ValueError("Port must be specified in connection configuration")
+        
+        # Parse SSL configuration if present
+        ssl_config = cls._parse_ssl_config(db_config)
+        
+        config = cls(
+            database=db_config['database'],
+            user=db_config['user'],
+            password=db_config['password'],
+            host=db_config['host'],
+            port=str(db_config['port']),
+            charset=db_config.get('charset', 'utf8mb4'),
+            local_host=local_host,
+            ssl=ssl_config
+        )
+        return config
+        
+    @classmethod
+    def _parse_ssl_config(cls, db_config: dict) -> Optional[SSLConfig]:
+        """解析SSL配置
+        
+        Args:
+            db_config: 数据库配置
+            
+        Returns:
+            Optional[SSLConfig]: SSL配置或None
+            
+        Raises:
+            ValueError: 如果SSL配置无效
+        """
+        if 'ssl' not in db_config:
+            return None
+            
+        ssl_params = db_config['ssl']
+        if not isinstance(ssl_params, dict):
+            raise ValueError("SSL configuration must be a dictionary")
+        
+        if ssl_params.get('mode') not in [None, 'disabled', 'required', 'verify_ca', 'verify_identity']:
+            raise ValueError(f"Invalid ssl-mode: {ssl_params.get('mode')}")
+        
+        return SSLConfig(
+            mode=ssl_params.get('mode', 'disabled'),
+            ca=ssl_params.get('ca'),
+            cert=ssl_params.get('cert'),
+            key=ssl_params.get('key')
+        )
 
-        # Get connection parameters
+    @classmethod
+    def from_yaml(cls, yaml_path: str, db_name: str, local_host: Optional[str] = None) -> 'MySQLConfig':
+        """Create configuration from YAML file
+
+        Args:
+            yaml_path: Path to YAML configuration file
+            db_name: Connection configuration name to use
+            local_host: Optional local host address
+        """
+        configs = cls.load_yaml_config(yaml_path)
+        
+        # Validate connection config
+        db_config = cls._validate_connection_config(configs, db_name)
+
+        # Create configuration based on URL or parameters
         if 'url' in db_config:
-            # Parse URL for connection parameters
-            params = parse_url(db_config['url'])
-            config = cls(
-                database=params['database'],
-                user=db_config['user'],
-                password=db_config['password'],
-                host=params['host'],
-                port=params['port'],
-                charset=params['charset'],
-                local_host=local_host,
-                url=db_config['url'],
-                ssl=params.get('ssl')
-            )
+            config = cls._create_config_from_url(db_config, local_host)
         else:
-            if not db_config.get('database'):
-                raise ValueError("MySQL database name must be specified in configuration")
-            if not db_config.get('host'):
-                raise ValueError("Host must be specified in connection configuration")
-            if not db_config.get('port'):
-                raise ValueError("Port must be specified in connection configuration")
-            
-            # Parse SSL configuration if present
-            ssl_config = None
-            if 'ssl' in db_config:
-                ssl_params = db_config['ssl']
-                if not isinstance(ssl_params, dict):
-                    raise ValueError("SSL configuration must be a dictionary")
-                
-                if ssl_params.get('mode') not in [None, 'disabled', 'required', 'verify_ca', 'verify_identity']:
-                    raise ValueError(f"Invalid ssl-mode: {ssl_params.get('mode')}")
-                
-                ssl_config = SSLConfig(
-                    mode=ssl_params.get('mode', 'disabled'),
-                    ca=ssl_params.get('ca'),
-                    cert=ssl_params.get('cert'),
-                    key=ssl_params.get('key')
-                )
+            config = cls._create_config_from_params(db_config, local_host)
             
-            config = cls(
-                database=db_config['database'],
-                user=db_config['user'],
-                password=db_config['password'],
-                host=db_config['host'],
-                port=str(db_config['port']),
-                charset=db_config.get('charset', 'utf8mb4'),
-                local_host=local_host,
-                ssl=ssl_config
-            )
         config.debug = cls.get_debug_mode()
         return config
 

mcp_dbutils/mysql/handler.py

@@ -31,12 +31,41 @@ class MySQLHandler(ConnectionHandler):
         self.log("debug", f"Configuring connection with parameters: {masked_params}")
         self.pool = None
 
+    async def _check_table_exists(self, cursor, table_name: str) -> None:
+        """检查表是否存在
+        
+        Args:
+            cursor: 数据库游标
+            table_name: 表名
+            
+        Raises:
+            ConnectionHandlerError: 如果表不存在
+        """
+        cursor.execute("""
+            SELECT COUNT(*) as count
+            FROM information_schema.tables
+            WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
+        """, (self.config.database, table_name))
+        table_exists = cursor.fetchone()
+        
+        # Handle different formats of cursor results (dict or tuple)
+        if not table_exists:
+            raise ConnectionHandlerError(f"Table '{self.config.database}.{table_name}' doesn't exist")
+        
+        # If fetchone returns a dictionary (dictionary=True was used)
+        if isinstance(table_exists, dict) and 'count' in table_exists and table_exists['count'] == 0:
+            raise ConnectionHandlerError(f"Table '{self.config.database}.{table_name}' doesn't exist")
+        
+        # If fetchone returns a tuple
+        if isinstance(table_exists, tuple) and table_exists[0] == 0:
+            raise ConnectionHandlerError(f"Table '{self.config.database}.{table_name}' doesn't exist")
+    
     async def get_tables(self) -> list[types.Resource]:
         """Get all table resources"""
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
                 cur.execute("""
                     SELECT 
                         TABLE_NAME as table_name,
@@ -66,7 +95,7 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
                 # Get column information
                 cur.execute("""
                     SELECT 
@@ -118,25 +147,31 @@ class MySQLHandler(ConnectionHandler):
             conn = mysql.connector.connect(**conn_params)
             self.log("debug", f"Executing query: {sql}")
 
-            with conn.cursor(dictionary=True) as cur:
-                # Start read-only transaction
-                cur.execute("SET TRANSACTION READ ONLY")
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
+                # Check if the query is a SELECT statement
+                sql_upper = sql.strip().upper()
+                is_select = sql_upper.startswith("SELECT")
+                
+                # Only set read-only transaction for SELECT statements
+                if is_select:
+                    cur.execute("SET TRANSACTION READ ONLY")
                 try:
                     cur.execute(sql)
-                    results = cur.fetchall()
+                    if not is_select:
+                        conn.commit()
+                    results = cur.fetchall() if is_select else []
+                    if cur.description is None:  # DDL statements
+                        return "Query executed successfully"
                     columns = [desc[0] for desc in cur.description]
-
-                    result_text = str({
-                        'type': self.db_type,
-                        'columns': columns,
-                        'rows': results,
-                        'row_count': len(results)
+                    return str({
+                        "columns": columns,
+                        "rows": results
                     })
-
-                    self.log("debug", f"Query completed, returned {len(results)} rows")
-                    return result_text
+                except mysql.connector.Error as e:
+                    self.log("error", f"Query error: {str(e)}")
+                    raise ConnectionHandlerError(str(e))
                 finally:
-                    cur.execute("ROLLBACK")
+                    cur.close()
         except mysql.connector.Error as e:
             error_msg = f"[{self.db_type}] Query execution failed: {str(e)}"
             raise ConnectionHandlerError(error_msg)
@@ -150,7 +185,10 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
+                # Check if table exists
+                await self._check_table_exists(cur, table_name)
+                
                 # Get table information and comment
                 cur.execute("""
                     SELECT 
@@ -220,7 +258,7 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
                 # MySQL provides a SHOW CREATE TABLE statement
                 cur.execute(f"SHOW CREATE TABLE {table_name}")
                 result = cur.fetchone()
@@ -242,7 +280,10 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
+                # Check if table exists
+                await self._check_table_exists(cur, table_name)
+
                 # Get index information
                 cur.execute("""
                     SELECT 
@@ -301,7 +342,10 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
+                # Check if table exists
+                await self._check_table_exists(cur, table_name)
+                
                 # Get table statistics
                 cur.execute("""
                     SELECT 
@@ -366,7 +410,10 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
+                # Check if table exists
+                await self._check_table_exists(cur, table_name)
+                
                 # Get constraint information
                 cur.execute("""
                     SELECT 
@@ -429,7 +476,7 @@ class MySQLHandler(ConnectionHandler):
         try:
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR
                 # Get EXPLAIN output
                 cur.execute(f"EXPLAIN FORMAT=TREE {sql}")
                 explain_result = cur.fetchall()

mcp_dbutils/mysql/server.py

@@ -49,7 +49,7 @@ class MySQLServer(ConnectionServer):
         """列出所有表资源"""
         try:
             conn = self.pool.get_connection()
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR - dictionary参数是正确的，用于返回字典格式的结果
                 cur.execute("""
                     SELECT 
                         table_name,
@@ -78,7 +78,7 @@ class MySQLServer(ConnectionServer):
         try:
             table_name = uri.split('/')[-2]
             conn = self.pool.get_connection()
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR - dictionary参数是正确的，用于返回字典格式的结果
                 # 获取列信息
                 cur.execute("""
                     SELECT 
@@ -170,7 +170,7 @@ class MySQLServer(ConnectionServer):
                 conn = self.pool.get_connection()
 
             self.log("info", f"执行查询: {sql}")
-            with conn.cursor(dictionary=True) as cur:
+            with conn.cursor(dictionary=True) as cur:  # NOSONAR - dictionary参数是正确的，用于返回字典格式的结果
                 # 设置只读事务
                 cur.execute("SET TRANSACTION READ ONLY")
                 try:

mcp_dbutils/sqlite/handler.py

@@ -1,6 +1,7 @@
 """SQLite connection handler implementation"""
 
 import sqlite3
+import time
 
 import mcp.types as types
 
@@ -78,28 +79,51 @@ class SQLiteHandler(ConnectionHandler):
     async def _execute_query(self, sql: str) -> str:
         """Execute SQL query"""
         try:
-            # Only allow SELECT statements
-            if not sql.strip().upper().startswith("SELECT"):
-                raise ConnectionHandlerError("cannot execute DELETE statement")
-            
-            with sqlite3.connect(self.config.path) as conn:
-                conn.row_factory = sqlite3.Row
-                cur = conn.cursor()
-                self.log("debug", f"Executing query: {sql}")
-                
-                cur.execute(sql)
-                results = cur.fetchall()
-                rows = [dict(row) for row in results]
+            # Check if the query is a DDL statement
+            sql_upper = sql.strip().upper()
+            is_ddl = sql_upper.startswith(("CREATE", "DROP", "ALTER", "TRUNCATE"))
+            is_dml = sql_upper.startswith(("INSERT", "UPDATE", "DELETE"))
+            is_select = sql_upper.startswith("SELECT")
 
-                result_text = str({
-                    'type': self.db_type,
-                    'columns': list(rows[0].keys()) if rows else [],
-                    'rows': rows,
-                    'row_count': len(rows)
-                })
+            if not (is_select or is_ddl or is_dml):
+                raise ConnectionHandlerError("Only SELECT, DDL, and DML statements are allowed")
+
+            conn = sqlite3.connect(self.config.path)
+            cur = conn.cursor()
 
-                self.log("debug", f"Query completed, returned {len(rows)} rows")
-                return result_text
+            try:
+                start_time = time.time()
+                cur.execute(sql)
+                conn.commit()
+                end_time = time.time()
+                elapsed_ms = (end_time - start_time) * 1000
+                self.log("debug", f"Query executed in {elapsed_ms:.2f}ms")
+                
+                if is_select:
+                    # Get column names
+                    columns = [description[0] for description in cur.description]
+                    # Fetch results and convert to dictionaries
+                    results = []
+                    for row in cur.fetchall():
+                        # Convert each row to a dictionary
+                        row_dict = {}
+                        for i, col_name in enumerate(columns):
+                            row_dict[col_name] = row[i]
+                        results.append(row_dict)
+                    
+                    return str({
+                        "columns": columns,
+                        "rows": results
+                    })
+                else:
+                    # For DDL/DML statements
+                    return "Query executed successfully"
+            except sqlite3.Error as e:
+                self.log("error", f"Query error: {str(e)}")
+                raise ConnectionHandlerError(str(e))
+            finally:
+                cur.close()
+                conn.close()
         except sqlite3.Error as e:
             error_msg = f"[{self.db_type}] Query execution failed: {str(e)}"
             raise ConnectionHandlerError(error_msg)
@@ -142,7 +166,7 @@ class SQLiteHandler(ConnectionHandler):
             with sqlite3.connect(self.config.path) as conn:
                 cur = conn.cursor()
                 # SQLite provides the complete CREATE statement
-                cur.execute(f"SELECT sql FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
+                cur.execute("SELECT sql FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
                 result = cur.fetchone()
                 
                 if not result:
@@ -151,7 +175,7 @@ class SQLiteHandler(ConnectionHandler):
                 ddl = result[0]
                 
                 # Get indexes
-                cur.execute(f"SELECT sql FROM sqlite_master WHERE type='index' AND tbl_name=?", (table_name,))
+                cur.execute("SELECT sql FROM sqlite_master WHERE type='index' AND tbl_name=?", (table_name,))
                 indexes = cur.fetchall()
                 
                 # Add index definitions
@@ -173,6 +197,12 @@ class SQLiteHandler(ConnectionHandler):
         try:
             with sqlite3.connect(self.config.path) as conn:
                 cur = conn.cursor()
+                
+                # Check if table exists
+                cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
+                if not cur.fetchone():
+                    raise ConnectionHandlerError(f"Table '{table_name}' doesn't exist")
+                
                 # 获取索引列表
                 cur.execute(f"PRAGMA index_list({table_name})")
                 indexes = cur.fetchall()
@@ -221,6 +251,11 @@ class SQLiteHandler(ConnectionHandler):
             with sqlite3.connect(self.config.path) as conn:
                 cur = conn.cursor()
                 
+                # Check if table exists
+                cur.execute("SELECT name FROM sqlite_master WHERE type='table' AND name=?", (table_name,))
+                if not cur.fetchone():
+                    raise ConnectionHandlerError(f"Table '{table_name}' doesn't exist")
+                
                 # Get basic table information
                 cur.execute(f"PRAGMA table_info({table_name})")
                 columns = cur.fetchall()
@@ -234,9 +269,9 @@ class SQLiteHandler(ConnectionHandler):
                 indexes = cur.fetchall()
                 
                 # Get page count and size
-                cur.execute(f"PRAGMA page_count")
+                cur.execute("PRAGMA page_count")
                 page_count = cur.fetchone()[0]
-                cur.execute(f"PRAGMA page_size")
+                cur.execute("PRAGMA page_size")
                 page_size = cur.fetchone()[0]
                 
                 # Calculate total size
@@ -382,6 +417,13 @@ class SQLiteHandler(ConnectionHandler):
             with sqlite3.connect(self.config.path) as conn:
                 cur = conn.cursor()
                 
+                # Check if the query is valid by preparing it
+                try:
+                    # Use prepare to validate the query without executing it
+                    conn.execute(f"EXPLAIN {sql}")
+                except sqlite3.Error as e:
+                    raise ConnectionHandlerError(f"Failed to explain query: {str(e)}")
+                
                 # Get EXPLAIN output
                 cur.execute(f"EXPLAIN QUERY PLAN {sql}")
                 plan = cur.fetchall()

mcp_dbutils/sqlite/server.py

@@ -54,7 +54,7 @@ class SQLiteServer(ConnectionServer):
             # 使用默认连接
             conn = self._get_connection()
 
-            with closing(conn) as connection:
+            with closing(conn) as _:
                 cursor = conn.execute(
                     "SELECT name FROM sqlite_master WHERE type='table'"
                 )
@@ -155,7 +155,7 @@ class SQLiteServer(ConnectionServer):
                 # 使用默认连接
                 conn = self._get_connection()
 
-            with closing(conn) as connection:
+            with closing(conn) as _:
                 self.log("info", f"执行查询: {sql}")
                 cursor = conn.execute(sql)
                 results = cursor.fetchall()