mcp-dbutils 0.16.1__py3-none-any.whl → 0.18.0__py3-none-any.whl

mcp_dbutils/mysql/handler.py

@@ -31,6 +31,35 @@ 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:
@@ -38,7 +67,7 @@ class MySQLHandler(ConnectionHandler):
             conn = mysql.connector.connect(**conn_params)
             with conn.cursor(dictionary=True) as cur:  # NOSONAR
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         TABLE_NAME as table_name,
                         TABLE_COMMENT as description
                     FROM information_schema.tables
@@ -69,7 +98,7 @@ class MySQLHandler(ConnectionHandler):
             with conn.cursor(dictionary=True) as cur:  # NOSONAR
                 # Get column information
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         COLUMN_NAME as column_name,
                         DATA_TYPE as data_type,
                         IS_NULLABLE as is_nullable,
@@ -119,24 +148,30 @@ class MySQLHandler(ConnectionHandler):
             self.log("debug", f"Executing query: {sql}")
 
             with conn.cursor(dictionary=True) as cur:  # NOSONAR
-                # Start read-only transaction
-                cur.execute("SET TRANSACTION READ ONLY")
+                # 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)
@@ -151,11 +186,14 @@ class MySQLHandler(ConnectionHandler):
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
             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 
+                    SELECT
                         TABLE_COMMENT as table_comment
-                    FROM information_schema.tables 
+                    FROM information_schema.tables
                     WHERE TABLE_SCHEMA = %s AND TABLE_NAME = %s
                 """, (self.config.database, table_name))
                 table_info = cur.fetchone()
@@ -163,7 +201,7 @@ class MySQLHandler(ConnectionHandler):
 
                 # Get column information
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         COLUMN_NAME as column_name,
                         DATA_TYPE as data_type,
                         COLUMN_DEFAULT as column_default,
@@ -184,14 +222,14 @@ class MySQLHandler(ConnectionHandler):
                     f"Comment: {table_comment or 'No comment'}\n",
                     COLUMNS_HEADER
                 ]
-                
+
                 for col in columns:
                     col_info = [
                         f"  {col['column_name']} ({col['data_type']})",
                         f"    Nullable: {col['is_nullable']}",
                         f"    Default: {col['column_default'] or 'None'}"
                     ]
-                    
+
                     if col['character_maximum_length']:
                         col_info.append(f"    Max Length: {col['character_maximum_length']}")
                     if col['numeric_precision']:
@@ -200,12 +238,12 @@ class MySQLHandler(ConnectionHandler):
                         col_info.append(f"    Scale: {col['numeric_scale']}")
                     if col['column_comment']:
                         col_info.append(f"    Comment: {col['column_comment']}")
-                        
+
                     description.extend(col_info)
                     description.append("")  # Empty line between columns
-                
+
                 return "\n".join(description)
-                
+
         except mysql.connector.Error as e:
             error_msg = f"Failed to get table description: {str(e)}"
             self.stats.record_error(e.__class__.__name__)
@@ -227,7 +265,7 @@ class MySQLHandler(ConnectionHandler):
                 if result:
                     return result['Create Table']
                 return f"Failed to get DDL for table {table_name}"
-                
+
         except mysql.connector.Error as e:
             error_msg = f"Failed to get table DDL: {str(e)}"
             self.stats.record_error(e.__class__.__name__)
@@ -243,9 +281,12 @@ class MySQLHandler(ConnectionHandler):
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
             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 
+                    SELECT
                         INDEX_NAME as index_name,
                         COLUMN_NAME as column_name,
                         NON_UNIQUE as non_unique,
@@ -264,7 +305,7 @@ class MySQLHandler(ConnectionHandler):
                 current_index = None
                 formatted_indexes = []
                 index_info = []
-                
+
                 for idx in indexes:
                     if current_index != idx['index_name']:
                         if index_info:
@@ -279,14 +320,14 @@ class MySQLHandler(ConnectionHandler):
                         ]
                         if idx['index_comment']:
                             index_info.insert(1, f"Comment: {idx['index_comment']}")
-                    
+
                     index_info.append(f"  - {idx['column_name']}")
 
                 if index_info:
                     formatted_indexes.extend(index_info)
 
                 return "\n".join(formatted_indexes)
-                
+
         except mysql.connector.Error as e:
             error_msg = f"Failed to get index information: {str(e)}"
             self.stats.record_error(e.__class__.__name__)
@@ -302,9 +343,12 @@ class MySQLHandler(ConnectionHandler):
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
             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 
+                    SELECT
                         TABLE_ROWS as table_rows,
                         AVG_ROW_LENGTH as avg_row_length,
                         DATA_LENGTH as data_length,
@@ -320,7 +364,7 @@ class MySQLHandler(ConnectionHandler):
 
                 # Get column statistics
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         COLUMN_NAME as column_name,
                         DATA_TYPE as data_type,
                         COLUMN_TYPE as column_type
@@ -367,20 +411,23 @@ class MySQLHandler(ConnectionHandler):
             conn_params = self.config.get_connection_params()
             conn = mysql.connector.connect(**conn_params)
             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 
+                    SELECT
                         k.CONSTRAINT_NAME as constraint_name,
                         t.CONSTRAINT_TYPE as constraint_type,
                         k.COLUMN_NAME as column_name,
                         k.REFERENCED_TABLE_NAME as referenced_table_name,
                         k.REFERENCED_COLUMN_NAME as referenced_column_name
                     FROM information_schema.key_column_usage k
-                    JOIN information_schema.table_constraints t 
+                    JOIN information_schema.table_constraints t
                         ON k.CONSTRAINT_NAME = t.CONSTRAINT_NAME
                         AND k.TABLE_SCHEMA = t.TABLE_SCHEMA
                         AND k.TABLE_NAME = t.TABLE_NAME
-                    WHERE k.TABLE_SCHEMA = %s 
+                    WHERE k.TABLE_SCHEMA = %s
                         AND k.TABLE_NAME = %s
                     ORDER BY t.CONSTRAINT_TYPE, k.CONSTRAINT_NAME, k.ORDINAL_POSITION
                 """, (self.config.database, table_name))
@@ -404,7 +451,7 @@ class MySQLHandler(ConnectionHandler):
                             f"\n{con['constraint_type']} Constraint: {con['constraint_name']}",
                             COLUMNS_HEADER
                         ]
-                    
+
                     col_info = f"  - {con['column_name']}"
                     if con['referenced_table_name']:
                         col_info += f" -> {con['referenced_table_name']}.{con['referenced_column_name']}"
@@ -446,7 +493,7 @@ class MySQLHandler(ConnectionHandler):
                 ]
                 for row in explain_result:
                     output.append(str(row['EXPLAIN']))
-                
+
                 output.extend([
                     "\nActual Plan (ANALYZE):",
                     "----------------------"
@@ -464,6 +511,26 @@ class MySQLHandler(ConnectionHandler):
             if conn:
                 conn.close()
 
+    async def test_connection(self) -> bool:
+        """Test database connection
+
+        Returns:
+            bool: True if connection is successful, False otherwise
+        """
+        conn = None
+        try:
+            conn_params = self.config.get_connection_params()
+            conn = mysql.connector.connect(**conn_params)
+            with conn.cursor() as cur:
+                cur.execute("SELECT 1")
+                return True
+        except mysql.connector.Error as e:
+            self.log("error", f"Connection test failed: {str(e)}")
+            return False
+        finally:
+            if conn:
+                conn.close()
+
     async def cleanup(self):
         """Cleanup resources"""
         # Log final stats before cleanup

mcp_dbutils/postgres/handler.py

@@ -165,7 +165,7 @@ class PostgreSQLHandler(ConnectionHandler):
                         (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass,
                         'pg_class'
                     ) as table_comment
-                    FROM information_schema.tables 
+                    FROM information_schema.tables
                     WHERE table_name = %s
                 """, (table_name,))
                 table_info = cur.fetchone()
@@ -173,7 +173,7 @@ class PostgreSQLHandler(ConnectionHandler):
 
                 # 获取列信息
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         column_name,
                         data_type,
                         column_default,
@@ -197,14 +197,14 @@ class PostgreSQLHandler(ConnectionHandler):
                     f"Comment: {table_comment or 'No comment'}\n",
                     COLUMNS_HEADER
                 ]
-                
+
                 for col in columns:
                     col_info = [
                         f"  {col[0]} ({col[1]})",
                         f"    Nullable: {col[3]}",
                         f"    Default: {col[2] or 'None'}"
                     ]
-                    
+
                     if col[4]:  # character_maximum_length
                         col_info.append(f"    Max Length: {col[4]}")
                     if col[5]:  # numeric_precision
@@ -213,12 +213,12 @@ class PostgreSQLHandler(ConnectionHandler):
                         col_info.append(f"    Scale: {col[6]}")
                     if col[7]:  # column_comment
                         col_info.append(f"    Comment: {col[7]}")
-                        
+
                     description.extend(col_info)
                     description.append("")  # Empty line between columns
-                
+
                 return "\n".join(description)
-                
+
         except psycopg2.Error as e:
             error_msg = f"Failed to get index information: [Code: {e.pgcode}] {e.pgerror or str(e)}"
             self.stats.record_error(e.__class__.__name__)
@@ -236,7 +236,7 @@ class PostgreSQLHandler(ConnectionHandler):
             with conn.cursor() as cur:
                 # 获取列定义
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         column_name,
                         data_type,
                         column_default,
@@ -252,7 +252,7 @@ class PostgreSQLHandler(ConnectionHandler):
 
                 # 获取约束
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         conname as constraint_name,
                         pg_get_constraintdef(c.oid) as constraint_def
                     FROM pg_constraint c
@@ -263,12 +263,12 @@ class PostgreSQLHandler(ConnectionHandler):
 
                 # 构建CREATE TABLE语句
                 ddl = [f"CREATE TABLE {table_name} ("]
-                
+
                 # 添加列定义
                 column_defs = []
                 for col in columns:
                     col_def = [f"    {col[0]} {col[1]}"]
-                    
+
                     if col[4]:  # character_maximum_length
                         col_def[0] = f"{col_def[0]}({col[4]})"
                     elif col[5]:  # numeric_precision
@@ -276,24 +276,24 @@ class PostgreSQLHandler(ConnectionHandler):
                             col_def[0] = f"{col_def[0]}({col[5]},{col[6]})"
                         else:
                             col_def[0] = f"{col_def[0]}({col[5]})"
-                            
+
                     if col[2]:  # default
                         col_def.append(f"DEFAULT {col[2]}")
                     if col[3] == 'NO':  # not null
                         col_def.append("NOT NULL")
-                        
+
                     column_defs.append(" ".join(col_def))
-                
+
                 # 添加约束定义
                 for con in constraints:
                     column_defs.append(f"    CONSTRAINT {con[0]} {con[1]}")
-                
+
                 ddl.append(",\n".join(column_defs))
                 ddl.append(");")
-                
+
                 # 添加注释
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         c.column_name,
                         col_description(
                             (quote_ident(table_schema) || '.' || quote_ident(table_name))::regclass,
@@ -307,15 +307,15 @@ class PostgreSQLHandler(ConnectionHandler):
                     WHERE c.table_name = %s
                 """, (table_name,))
                 comments = cur.fetchall()
-                
+
                 for comment in comments:
                     if comment[2]:  # table comment
                         ddl.append(f"\nCOMMENT ON TABLE {table_name} IS '{comment[2]}';")
                     if comment[1]:  # column comment
                         ddl.append(f"COMMENT ON COLUMN {table_name}.{comment[0]} IS '{comment[1]}';")
-                
+
                 return "\n".join(ddl)
-                
+
         except psycopg2.Error as e:
             error_msg = f"Failed to get table DDL: [Code: {e.pgcode}] {e.pgerror or str(e)}"
             self.stats.record_error(e.__class__.__name__)
@@ -336,7 +336,7 @@ class PostgreSQLHandler(ConnectionHandler):
                     SELECT
                         i.relname as index_name,
                         a.attname as column_name,
-                        CASE 
+                        CASE
                             WHEN ix.indisprimary THEN 'PRIMARY KEY'
                             WHEN ix.indisunique THEN 'UNIQUE'
                             ELSE 'INDEX'
@@ -362,7 +362,7 @@ class PostgreSQLHandler(ConnectionHandler):
                 current_index = None
                 formatted_indexes = []
                 index_info = []
-                
+
                 for idx in indexes:
                     if current_index != idx[0]:
                         if index_info:
@@ -377,14 +377,14 @@ class PostgreSQLHandler(ConnectionHandler):
                         ]
                         if idx[5]:  # index comment
                             index_info.insert(1, f"Comment: {idx[5]}")
-                    
+
                     index_info.append(f"  - {idx[1]}")
 
                 if index_info:
                     formatted_indexes.extend(index_info)
 
                 return "\n".join(formatted_indexes)
-                
+
         except psycopg2.Error as e:
             error_msg = f"Failed to get index information: [Code: {e.pgcode}] {e.pgerror or str(e)}"
             self.stats.record_error(e.__class__.__name__)
@@ -402,7 +402,7 @@ class PostgreSQLHandler(ConnectionHandler):
             with conn.cursor() as cur:
                 # Get table statistics
                 cur.execute("""
-                    SELECT 
+                    SELECT
                         c.reltuples::bigint as row_estimate,
                         pg_size_pretty(pg_total_relation_size(c.oid)) as total_size,
                         pg_size_pretty(pg_table_size(c.oid)) as table_size,
@@ -428,8 +428,8 @@ class PostgreSQLHandler(ConnectionHandler):
                         s.n_distinct as distinct_values,
                         pg_column_size(a.attname::text) as approx_width
                     FROM pg_stats s
-                    JOIN pg_attribute a ON a.attrelid = %s::regclass 
-                        AND a.attnum > 0 
+                    JOIN pg_attribute a ON a.attrelid = %s::regclass
+                        AND a.attnum > 0
                         AND a.attname = s.attname
                     WHERE s.schemaname = 'public'
                     AND s.tablename = %s
@@ -522,10 +522,10 @@ class PostgreSQLHandler(ConnectionHandler):
 
                     if con[4]:  # is_deferrable
                         output.append(f"    Deferrable: {'Deferred' if con[5] else 'Immediate'}")
-                    
+
                     if con[6]:  # comment
                         output.append(f"    Comment: {con[6]}")
-                    
+
                     output.append("")  # Empty line between constraints
 
                 return "\n".join(output)
@@ -567,7 +567,7 @@ class PostgreSQLHandler(ConnectionHandler):
                     "----------------"
                 ]
                 output.extend(line[0] for line in regular_plan)
-                
+
                 output.extend([
                     "\nActual Plan (ANALYZE):",
                     "----------------------"
@@ -584,6 +584,26 @@ class PostgreSQLHandler(ConnectionHandler):
             if conn:
                 conn.close()
 
+    async def test_connection(self) -> bool:
+        """Test database connection
+
+        Returns:
+            bool: True if connection is successful, False otherwise
+        """
+        conn = None
+        try:
+            conn_params = self.config.get_connection_params()
+            conn = psycopg2.connect(**conn_params)
+            with conn.cursor() as cur:
+                cur.execute("SELECT 1")
+                return True
+        except psycopg2.Error as e:
+            self.log("error", f"Connection test failed: [Code: {e.pgcode}] {e.pgerror or str(e)}")
+            return False
+        finally:
+            if conn:
+                conn.close()
+
     async def cleanup(self):
         """Cleanup resources"""
         # Log final stats before cleanup