add stage tool

Author: BohuTANG

README.md

@@ -6,11 +6,20 @@ An MCP server for Databend database interactions.
 
 ## What You Can Do
 
+### Database Operations
 - **execute_sql** - Execute SQL queries with timeout protection and safe mode security
 - **show_databases** - List all databases
 - **show_tables** - List tables in a database (with optional filter)
 - **describe_table** - Get table schema information
 
+### Stage Management
+- **show_stages** - List all available Databend stages
+- **list_stage_files** - List files in a specific stage (supports @stage_name format)
+- **create_stage** - Create a new stage with connection support
+
+### Connection Management
+- **show_connections** - List all available Databend connections
+
 ## Security Features
 
 ### MCP Safe Mode (Enabled by Default)
@@ -93,11 +102,21 @@ Add to your MCP client configuration (e.g., Claude Desktop, Windsurf):
 ### Step 4: Start Using
 
 Once configured, you can ask your AI assistant to:
+
+**Database Operations:**
 - "Show me all databases"
-- "List tables in the sales database"
+- "List tables in the sales database" 
 - "Describe the users table structure"
 - "Run this SQL query: SELECT * FROM products LIMIT 10"
 
+**Stage Management:**
+- "Show me all stages"
+- "List files in @my_stage"
+- "Create a stage named my_s3_stage with URL s3://my-bucket using connection my_connection"
+
+**Connection Management:**
+- "Show all connections"
+
 ## Development
 
 ```bash

mcp_databend/server.py

@@ -204,6 +204,73 @@ def describe_table(table: str, database: Optional[str] = None):
     return execute_sql(sql)
 
 
+@mcp.tool()
+def show_stages():
+    """List available Databend stages (safe operation, not affected by MCP safe mode)"""
+    logger.info("Listing all stages")
+    return _execute_sql("SHOW STAGES")
+
+
+@mcp.tool()
+def list_stage_files(stage_name: str, path: Optional[str] = None):
+    """
+    List files in a Databend stage (safe operation, not affected by MCP safe mode)
+    Args:
+        stage_name: The stage name (with @ prefix)
+        path: Optional path within the stage
+
+    Returns:
+        Dictionary containing either query results or error information
+    """
+    if not stage_name.startswith('@'):
+        stage_name = f"@{stage_name}"
+    
+    if path:
+        stage_path = f"{stage_name}/{path.strip('/')}"
+    else:
+        stage_path = stage_name
+    
+    logger.info(f"Listing files in stage '{stage_path}'")
+    sql = f"LIST {stage_path}"
+    return _execute_sql(sql)
+
+
+@mcp.tool()
+def show_connections():
+    """List available Databend connections (safe operation, not affected by MCP safe mode)"""
+    logger.info("Listing all connections")
+    return _execute_sql("SHOW CONNECTIONS")
+
+
+@mcp.tool()
+async def create_stage(name: str, url: str, connection_name: Optional[str] = None, **kwargs) -> dict:
+    """
+    Create a Databend stage with connection
+    Args:
+        name: The stage name
+        url: The stage URL (e.g., 's3://bucket-name')
+        connection_name: Optional connection name to use
+        **kwargs: Additional stage options
+
+    Returns:
+        Dictionary containing either query results or error information
+    """
+    logger.info(f"Creating stage '{name}' with URL '{url}'")
+    
+    sql_parts = [f"CREATE STAGE {name}", f"URL = '{url}'"]
+    
+    if connection_name:
+        sql_parts.append(f"CONNECTION = (CONNECTION_NAME = '{connection_name}')")
+    
+    # Add any additional options from kwargs
+    for key, value in kwargs.items():
+        if key not in ['name', 'url', 'connection_name']:
+            sql_parts.append(f"{key.upper()} = '{value}'")
+    
+    sql = " ".join(sql_parts)
+    return _execute_sql(sql)
+
+
 def main():
     """Main entry point for the MCP server."""
     try:

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "mcp-databend"
-version = "0.1.7"
+version = "0.1.8"
 description = "An MCP server for Databend."
 readme = "README.md"
 requires-python = ">=3.12"