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

mcp_dbutils/config.py

@@ -2,18 +2,111 @@
 
 import os
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Literal
+from typing import Any, Dict, List, Literal, Optional, Set, Union
 
 import yaml
 
 # Supported connection types
 ConnectionType = Literal['sqlite', 'postgres', 'mysql']
 
+# Supported write operations
+WriteOperationType = Literal['INSERT', 'UPDATE', 'DELETE']
+
+# Default policy for tables not explicitly listed in write_permissions
+DefaultPolicyType = Literal['read_only', 'allow_all']
+
+class WritePermissions:
+    """Write permissions configuration"""
+
+    def __init__(self, config: Optional[Dict[str, Any]] = None):
+        """Initialize write permissions
+
+        Args:
+            config: Write permissions configuration dictionary
+        """
+        self.tables: Dict[str, Set[WriteOperationType]] = {}
+        self.default_policy: DefaultPolicyType = 'read_only'
+
+        if config:
+            # Parse table permissions
+            if 'tables' in config and isinstance(config['tables'], dict):
+                for table_name, table_config in config['tables'].items():
+                    operations: Set[WriteOperationType] = set()
+
+                    if isinstance(table_config, dict) and 'operations' in table_config:
+                        ops = table_config['operations']
+                        if isinstance(ops, list):
+                            for op in ops:
+                                if op in ('INSERT', 'UPDATE', 'DELETE'):
+                                    operations.add(op)  # type: ignore
+
+                    # If no operations specified, allow all
+                    if not operations:
+                        operations = {'INSERT', 'UPDATE', 'DELETE'}  # type: ignore
+
+                    self.tables[table_name] = operations
+
+            # Parse default policy
+            if 'default_policy' in config:
+                policy = config['default_policy']
+                if policy in ('read_only', 'allow_all'):
+                    self.default_policy = policy  # type: ignore
+
+    def can_write_to_table(self, table_name: str) -> bool:
+        """Check if writing to the table is allowed
+
+        Args:
+            table_name: Name of the table
+
+        Returns:
+            True if writing to the table is allowed, False otherwise
+        """
+        # If table is explicitly listed, it's writable
+        if table_name in self.tables:
+            return True
+
+        # Otherwise, check default policy
+        return self.default_policy == 'allow_all'
+
+    def allowed_operations(self, table_name: str) -> Set[WriteOperationType]:
+        """Get allowed operations for a table
+
+        Args:
+            table_name: Name of the table
+
+        Returns:
+            Set of allowed operations
+        """
+        # If table is explicitly listed, return its allowed operations
+        if table_name in self.tables:
+            return self.tables[table_name]
+
+        # Otherwise, check default policy
+        if self.default_policy == 'allow_all':
+            return {'INSERT', 'UPDATE', 'DELETE'}  # type: ignore
+
+        # Default to empty set (no operations allowed)
+        return set()  # type: ignore
+
+    def is_operation_allowed(self, table_name: str, operation: WriteOperationType) -> bool:
+        """Check if an operation is allowed on a table
+
+        Args:
+            table_name: Name of the table
+            operation: Operation type
+
+        Returns:
+            True if the operation is allowed, False otherwise
+        """
+        return operation in self.allowed_operations(table_name)
+
 class ConnectionConfig(ABC):
     """Base class for connection configuration"""
 
     debug: bool = False
     type: ConnectionType  # Connection type
+    writable: bool = False  # Whether write operations are allowed
+    write_permissions: Optional[WritePermissions] = None  # Write permissions configuration
 
     @abstractmethod
     def get_connection_params(self) -> Dict[str, Any]:
@@ -50,6 +143,15 @@ class ConnectionConfig(ABC):
             if db_type not in ('sqlite', 'postgres', 'mysql'):
                 raise ValueError(f"Invalid type value in database configuration {conn_name}: {db_type}")
 
+            # Validate write permissions if writable is true
+            if db_config.get('writable', False):
+                if not isinstance(db_config.get('writable'), bool):
+                    raise ValueError(f"Invalid writable value in database configuration {conn_name}: {db_config['writable']}")
+
+                # Validate write_permissions if present
+                if 'write_permissions' in db_config and not isinstance(db_config['write_permissions'], dict):
+                    raise ValueError(f"Invalid write_permissions in database configuration {conn_name}: {db_config['write_permissions']}")
+
         return connections
 
     @classmethod

mcp_dbutils/mysql/config.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from typing import Any, Dict, Literal, Optional
 from urllib.parse import parse_qs, urlparse
 
-from ..config import ConnectionConfig
+from ..config import ConnectionConfig, WritePermissions
 
 
 @dataclass
@@ -16,30 +16,30 @@ class SSLConfig:
 
 def parse_url(url: str) -> Dict[str, Any]:
     """Parse MySQL URL into connection parameters
-    
+
     Args:
         url: URL (e.g. mysql://host:port/dbname?ssl-mode=verify_identity)
-        
+
     Returns:
         Dictionary of connection parameters including SSL settings
     """
     if not url.startswith('mysql://'):
         raise ValueError("Invalid MySQL URL format")
-    
+
     if '@' in url:
         raise ValueError("URL should not contain credentials. Please provide username and password separately.")
-    
+
     # Parse URL and query parameters
     parsed = urlparse(url)
     query_params = parse_qs(parsed.query)
-    
+
     params = {
         'host': parsed.hostname or 'localhost',
         'port': str(parsed.port or 3306),
         'database': parsed.path.lstrip('/') if parsed.path else '',
         'charset': query_params.get('charset', ['utf8mb4'])[0]
     }
-    
+
     if not params['database']:
         raise ValueError("MySQL database name must be specified in URL")
 
@@ -50,17 +50,17 @@ def parse_url(url: str) -> Dict[str, Any]:
         if mode not in ['disabled', 'required', 'verify_ca', 'verify_identity']:
             raise ValueError(f"Invalid ssl-mode: {mode}")
         ssl_params['mode'] = mode
-        
+
     if 'ssl-ca' in query_params:
         ssl_params['ca'] = query_params['ssl-ca'][0]
     if 'ssl-cert' in query_params:
         ssl_params['cert'] = query_params['ssl-cert'][0]
     if 'ssl-key' in query_params:
         ssl_params['key'] = query_params['ssl-key'][0]
-        
+
     if ssl_params:
         params['ssl'] = SSLConfig(**ssl_params)
-        
+
     return params
 
 @dataclass
@@ -75,18 +75,20 @@ class MySQLConfig(ConnectionConfig):
     type: Literal['mysql'] = 'mysql'
     url: Optional[str] = None
     ssl: Optional[SSLConfig] = None
+    writable: bool = False  # Whether write operations are allowed
+    write_permissions: Optional[WritePermissions] = None  # Write permissions configuration
 
     @classmethod
     def _validate_connection_config(cls, configs: dict, db_name: str) -> dict:
         """验证连接配置是否有效
-        
+
         Args:
             configs: 配置字典
             db_name: 连接名称
-            
+
         Returns:
             dict: 数据库配置
-            
+
         Raises:
             ValueError: 如果配置无效
         """
@@ -107,17 +109,17 @@ 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: 配置对象
         """
@@ -135,18 +137,18 @@ class MySQLConfig(ConnectionConfig):
             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配置无效
         """
@@ -156,10 +158,10 @@ class MySQLConfig(ConnectionConfig):
             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'],
@@ -171,30 +173,30 @@ class MySQLConfig(ConnectionConfig):
             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'),
@@ -212,7 +214,7 @@ class MySQLConfig(ConnectionConfig):
             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)
 
@@ -221,26 +223,35 @@ class MySQLConfig(ConnectionConfig):
             config = cls._create_config_from_url(db_config, local_host)
         else:
             config = cls._create_config_from_params(db_config, local_host)
-            
+
+        # Parse write permissions
+        config.writable = db_config.get('writable', False)
+        if config.writable and 'write_permissions' in db_config:
+            config.write_permissions = WritePermissions(db_config['write_permissions'])
+
         config.debug = cls.get_debug_mode()
         return config
 
     @classmethod
-    def from_url(cls, url: str, user: str, password: str, 
-                local_host: Optional[str] = None) -> 'MySQLConfig':
+    def from_url(cls, url: str, user: str, password: str,
+                local_host: Optional[str] = None,
+                writable: bool = False,
+                write_permissions: Optional[Dict[str, Any]] = None) -> 'MySQLConfig':
         """Create configuration from URL and credentials
-        
+
         Args:
             url: URL (mysql://host:port/dbname)
             user: Username for connection
             password: Password for connection
             local_host: Optional local host address
-            
+            writable: Whether write operations are allowed
+            write_permissions: Write permissions configuration
+
         Raises:
             ValueError: If URL format is invalid or required parameters are missing
         """
         params = parse_url(url)
-        
+
         config = cls(
             database=params['database'],
             user=user,
@@ -250,8 +261,14 @@ class MySQLConfig(ConnectionConfig):
             charset=params['charset'],
             local_host=local_host,
             url=url,
-            ssl=params.get('ssl')
+            ssl=params.get('ssl'),
+            writable=writable
         )
+
+        # Parse write permissions
+        if writable and write_permissions:
+            config.write_permissions = WritePermissions(write_permissions)
+
         config.debug = cls.get_debug_mode()
         return config
 
@@ -266,7 +283,7 @@ class MySQLConfig(ConnectionConfig):
             'charset': self.charset,
             'use_unicode': True
         }
-        
+
         # Add SSL parameters if configured
         if self.ssl:
             params['ssl_mode'] = self.ssl.mode
@@ -276,7 +293,7 @@ class MySQLConfig(ConnectionConfig):
                 params['ssl_cert'] = self.ssl.cert
             if self.ssl.key:
                 params['ssl_key'] = self.ssl.key
-                
+
         return {k: v for k, v in params.items() if v is not None}
 
     def get_masked_connection_info(self) -> Dict[str, Any]:

mcp_dbutils/mysql/handler.py

@@ -181,6 +181,66 @@ class MySQLHandler(ConnectionHandler):
             if conn:
                 conn.close()
 
+    async def _execute_write_query(self, sql: str) -> str:
+        """Execute SQL write query
+
+        Args:
+            sql: SQL write query (INSERT, UPDATE, DELETE)
+
+        Returns:
+            str: Execution result
+
+        Raises:
+            ConnectionHandlerError: If query execution fails
+        """
+        conn = None
+        try:
+            # Check if the query is a write operation
+            sql_upper = sql.strip().upper()
+            is_insert = sql_upper.startswith("INSERT")
+            is_update = sql_upper.startswith("UPDATE")
+            is_delete = sql_upper.startswith("DELETE")
+            is_transaction = sql_upper.startswith(("BEGIN", "COMMIT", "ROLLBACK", "START TRANSACTION"))
+
+            if not (is_insert or is_update or is_delete or is_transaction):
+                raise ConnectionHandlerError("Only INSERT, UPDATE, DELETE, and transaction statements are allowed for write operations")
+
+            conn_params = self.config.get_connection_params()
+            conn = mysql.connector.connect(**conn_params)
+            self.log("debug", f"Executing write operation: {sql}")
+
+            with conn.cursor() as cur:
+                try:
+                    # Execute the write operation
+                    cur.execute(sql)
+
+                    # Get number of affected rows
+                    affected_rows = cur.rowcount
+
+                    # Commit the transaction if not in a transaction block
+                    if not is_transaction:
+                        conn.commit()
+
+                    self.log("debug", f"Write operation executed successfully, affected {affected_rows} rows")
+
+                    # Return result
+                    if is_transaction:
+                        return f"Transaction operation executed successfully"
+                    else:
+                        return f"Write operation executed successfully. {affected_rows} row{'s' if affected_rows != 1 else ''} affected."
+                except mysql.connector.Error as e:
+                    # Rollback on error
+                    if not is_transaction:
+                        conn.rollback()
+                    self.log("error", f"Write operation error: {str(e)}")
+                    raise ConnectionHandlerError(str(e))
+        except mysql.connector.Error as e:
+            error_msg = f"[{self.db_type}] Write operation failed: {str(e)}"
+            raise ConnectionHandlerError(error_msg)
+        finally:
+            if conn:
+                conn.close()
+
     async def get_table_description(self, table_name: str) -> str:
         """Get detailed table description"""
         conn = None

mcp_dbutils/postgres/config.py

@@ -3,7 +3,7 @@ from dataclasses import dataclass
 from typing import Any, Dict, Literal, Optional
 from urllib.parse import parse_qs, urlparse
 
-from ..config import ConnectionConfig
+from ..config import ConnectionConfig, WritePermissions
 
 
 @dataclass
@@ -16,29 +16,29 @@ class SSLConfig:
 
 def parse_url(url: str) -> Dict[str, Any]:
     """Parse PostgreSQL URL into connection parameters
-    
+
     Args:
         url: URL (e.g. postgresql://host:port/dbname?sslmode=verify-full)
-        
+
     Returns:
         Dictionary of connection parameters including SSL settings
     """
     if not url.startswith('postgresql://'):
         raise ValueError("Invalid PostgreSQL URL format")
-    
+
     if '@' in url:
         raise ValueError("URL should not contain credentials. Please provide username and password separately.")
-    
+
     # Parse URL and query parameters
     parsed = urlparse(url)
     query_params = parse_qs(parsed.query)
-    
+
     params = {
         'host': parsed.hostname or 'localhost',
         'port': str(parsed.port or 5432),
         'dbname': parsed.path.lstrip('/') if parsed.path else '',
     }
-    
+
     if not params['dbname']:
         raise ValueError("PostgreSQL database name must be specified in URL")
 
@@ -49,17 +49,17 @@ def parse_url(url: str) -> Dict[str, Any]:
         if mode not in ['disable', 'require', 'verify-ca', 'verify-full']:
             raise ValueError(f"Invalid sslmode: {mode}")
         ssl_params['mode'] = mode
-        
+
     if 'sslcert' in query_params:
         ssl_params['cert'] = query_params['sslcert'][0]
     if 'sslkey' in query_params:
         ssl_params['key'] = query_params['sslkey'][0]
     if 'sslrootcert' in query_params:
         ssl_params['root'] = query_params['sslrootcert'][0]
-        
+
     if ssl_params:
         params['ssl'] = SSLConfig(**ssl_params)
-        
+
     return params
 
 @dataclass
@@ -73,6 +73,8 @@ class PostgreSQLConfig(ConnectionConfig):
     type: Literal['postgres'] = 'postgres'
     url: Optional[str] = None
     ssl: Optional[SSLConfig] = None
+    writable: bool = False  # Whether write operations are allowed
+    write_permissions: Optional[WritePermissions] = None  # Write permissions configuration
 
     @classmethod
     def from_yaml(cls, yaml_path: str, db_name: str, local_host: Optional[str] = None) -> 'PostgreSQLConfig':
@@ -123,24 +125,24 @@ class PostgreSQLConfig(ConnectionConfig):
                 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, 'disable', 'require', 'verify-ca', 'verify-full']:
                     raise ValueError(f"Invalid sslmode: {ssl_params.get('mode')}")
-                
+
                 ssl_config = SSLConfig(
                     mode=ssl_params.get('mode', 'disable'),
                     cert=ssl_params.get('cert'),
                     key=ssl_params.get('key'),
                     root=ssl_params.get('root')
                 )
-            
+
             config = cls(
                 dbname=db_config['dbname'],
                 user=db_config['user'],
@@ -150,25 +152,35 @@ class PostgreSQLConfig(ConnectionConfig):
                 local_host=local_host,
                 ssl=ssl_config
             )
+
+        # Parse write permissions
+        config.writable = db_config.get('writable', False)
+        if config.writable and 'write_permissions' in db_config:
+            config.write_permissions = WritePermissions(db_config['write_permissions'])
+
         config.debug = cls.get_debug_mode()
         return config
 
     @classmethod
-    def from_url(cls, url: str, user: str, password: str, 
-                local_host: Optional[str] = None) -> 'PostgreSQLConfig':
+    def from_url(cls, url: str, user: str, password: str,
+                local_host: Optional[str] = None,
+                writable: bool = False,
+                write_permissions: Optional[Dict[str, Any]] = None) -> 'PostgreSQLConfig':
         """Create configuration from URL and credentials
-        
+
         Args:
             url: URL (postgresql://host:port/dbname)
             user: Username for connection
             password: Password for connection
             local_host: Optional local host address
-            
+            writable: Whether write operations are allowed
+            write_permissions: Write permissions configuration
+
         Raises:
             ValueError: If URL format is invalid or required parameters are missing
         """
         params = parse_url(url)
-        
+
         config = cls(
             dbname=params['dbname'],
             user=user,
@@ -177,8 +189,14 @@ class PostgreSQLConfig(ConnectionConfig):
             port=params['port'],
             local_host=local_host,
             url=url,
-            ssl=params.get('ssl')
+            ssl=params.get('ssl'),
+            writable=writable
         )
+
+        # Parse write permissions
+        if writable and write_permissions:
+            config.write_permissions = WritePermissions(write_permissions)
+
         config.debug = cls.get_debug_mode()
         return config
 
@@ -191,7 +209,7 @@ class PostgreSQLConfig(ConnectionConfig):
             'host': self.local_host or self.host,
             'port': self.port
         }
-        
+
         # Add SSL parameters if configured
         if self.ssl:
             params['sslmode'] = self.ssl.mode
@@ -201,7 +219,7 @@ class PostgreSQLConfig(ConnectionConfig):
                 params['sslkey'] = self.ssl.key
             if self.ssl.root:
                 params['sslrootcert'] = self.ssl.root
-                
+
         return {k: v for k, v in params.items() if v}
 
     def get_masked_connection_info(self) -> Dict[str, Any]: