Merge pull request #4 from databendlabs/databend-python

feat: add local databend mode

Author: sundy-li

.gitignore

@@ -1,3 +1,6 @@
+# databend dir
+.databend
+
 # Byte-compiled / optimized / DLL files
 __pycache__/
 *.py[cod]

README.md

@@ -61,6 +61,9 @@ Get your connection string from [Databend documentation](https://docs.databend.c
 | **Databend Cloud** | `databend://user:pwd@host:443/database?warehouse=wh` |
 | **Self-hosted** | `databend://user:pwd@localhost:8000/database?sslmode=disable` |
 
+Or use local Databend by setting `LOCAL_MODE=true`, the metadata is stored in `.databend` directory:
+
+
 ### Step 2: Install
 
 ```bash
@@ -71,10 +74,16 @@ uv tool install mcp-databend
 
 #### Option A: Claude Code (CLI)
 
+- For Databend server:
 ```bash
 claude mcp add mcp-databend --env DATABEND_DSN='your-connection-string-here' -- uv tool run mcp-databend
 ```
 
+- For local Databend:
+```bash
+claude mcp add mcp-databend --env -- uv tool run mcp-databend
+```
+
 #### Option B: MCP Configuration (JSON)
 
 Add to your MCP client configuration (e.g., Claude Desktop, Windsurf):
@@ -105,7 +114,7 @@ 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"
 
@@ -127,4 +136,7 @@ uv sync
 
 # Run locally
 uv run python -m mcp_databend.main
+
+# Use modelcontextprotocol/inspector to debug
+npx @modelcontextprotocol/inspector -e LOCAL_MODE=1 uv run python -m mcp_databend.main
 ```

mcp_databend/env.py

@@ -11,6 +11,7 @@ class DatabendConfig:
     Required environment variables:
         DATABEND_DSN: The dsn connect string (defaults to: "databend://default:@127.0.0.1:8000/?sslmode=disable")
         SAFE_MODE: Enable safe mode to restrict dangerous SQL operations (defaults to: "true")
+        LOCAL_MODE: Enable local mode to use in-memory Databend (defaults to: "false")
     """
 
     def __init__(self):
@@ -29,6 +30,16 @@ def safe_mode(self) -> bool:
         """Get the safe mode setting."""
         return os.environ.get("SAFE_MODE", "true").lower() in ("true", "1", "yes", "on")
 
+    @property
+    def local_mode(self) -> bool:
+        """Get the local mode setting."""
+        return os.environ.get("LOCAL_MODE", "false").lower() in (
+            "true",
+            "1",
+            "yes",
+            "on",
+        )
+
 
 # Global instance placeholder for the singleton pattern
 _CONFIG_INSTANCE = None

mcp_databend/server.py

@@ -32,43 +32,62 @@
 # Initialize MCP server
 mcp = FastMCP(SERVER_NAME)
 
+# Global Databend client singleton
+_databend_client = None
+
+
+def get_global_databend_client():
+    """Get global Databend client instance (deprecated, use create_databend_client)."""
+    global _databend_client
+    if _databend_client is None:
+        _databend_client = create_databend_client()
+    return _databend_client
+
 
 def is_sql_safe(sql: str) -> tuple[bool, str]:
     """
     Check if SQL query is safe to execute in safe mode.
-    
+
     Args:
         sql: SQL query string to check
-        
+
     Returns:
         Tuple of (is_safe, reason) where is_safe is boolean and reason is error message if unsafe
     """
     sql_upper = sql.upper().strip()
-    
+
     # List of dangerous operations to block in safe mode
     dangerous_patterns = [
-        (r'\bDROP\s+', "DROP operations are not allowed in MCP safe mode"),
-        (r'\bDELETE\s+', "DELETE operations are not allowed in MCP safe mode"),
-        (r'\bTRUNCATE\s+', "TRUNCATE operations are not allowed in MCP safe mode"),
-        (r'\bALTER\s+', "ALTER operations are not allowed in MCP safe mode"),
-        (r'\bUPDATE\s+', "UPDATE operations are not allowed in MCP safe mode"),
-        (r'\bREVOKE\s+', "REVOKE operations are not allowed in MCP safe mode"),
+        (r"\bDROP\s+", "DROP operations are not allowed in MCP safe mode"),
+        (r"\bDELETE\s+", "DELETE operations are not allowed in MCP safe mode"),
+        (r"\bTRUNCATE\s+", "TRUNCATE operations are not allowed in MCP safe mode"),
+        (r"\bALTER\s+", "ALTER operations are not allowed in MCP safe mode"),
+        (r"\bUPDATE\s+", "UPDATE operations are not allowed in MCP safe mode"),
+        (r"\bREVOKE\s+", "REVOKE operations are not allowed in MCP safe mode"),
     ]
-    
+
     # Check each dangerous pattern
     for pattern, reason in dangerous_patterns:
         if re.search(pattern, sql_upper, re.IGNORECASE | re.DOTALL):
             return False, reason
-    
+
     return True, ""
 
 
 def create_databend_client():
     """Create and return a Databend client instance."""
     config = get_config()
-    from databend_driver import BlockingDatabendClient
 
-    return BlockingDatabendClient(config.dsn)
+    if config.local_mode:
+        # Use local in-memory Databend
+        import databend
+
+        return databend.SessionContext()
+    else:
+        # Use remote Databend server
+        from databend_driver import BlockingDatabendClient
+
+        return BlockingDatabendClient(config.dsn)
 
 
 def execute_databend_query(sql: str) -> list[dict] | dict:
@@ -81,20 +100,28 @@ def execute_databend_query(sql: str) -> list[dict] | dict:
     Returns:
         List of dictionaries containing query results or error dictionary
     """
-    client = create_databend_client()
-    conn = client.get_conn()
+    client = get_global_databend_client()
+    config = get_config()
 
     try:
-        cursor = conn.query_iter(sql)
-        column_names = [field.name for field in cursor.schema().fields()]
-        results = []
-
-        for row in cursor:
-            row_data = dict(zip(column_names, list(row.values())))
-            results.append(row_data)
-
-        logger.info(f"Query executed successfully, returned {len(results)} rows")
-        return results
+        if config.local_mode:
+            # Handle local in-memory Databend
+            result = client.sql(sql)
+            df = result.to_pandas()
+            return df.to_dict("records")
+        else:
+            # Handle remote Databend server
+            conn = client.get_conn()
+            cursor = conn.query_iter(sql)
+            column_names = [field.name for field in cursor.schema().fields()]
+            results = []
+
+            for row in cursor:
+                row_data = dict(zip(column_names, list(row.values())))
+                results.append(row_data)
+
+            logger.info(f"Query executed successfully, returned {len(results)} rows")
+            return results
 
     except Exception as err:
         error_msg = f"Error executing query: {str(err)}"
@@ -104,7 +131,7 @@ def execute_databend_query(sql: str) -> list[dict] | dict:
 
 def _execute_sql(sql: str) -> dict:
     logger.info(f"Executing SQL query: {sql}")
-    
+
     # Check safe mode configuration
     config = get_config()
     if config.safe_mode:
@@ -141,12 +168,13 @@ def _execute_sql(sql: str) -> dict:
         logger.error(error_msg)
         return {"status": "error", "message": error_msg}
 
+
 @mcp.tool()
 async def execute_sql(sql: str) -> dict:
     """
     Execute SQL query against Databend database with MCP safe mode protection.
-    
-    Safe mode (enabled by default) blocks dangerous operations like DROP, DELETE, 
+
+    Safe mode (enabled by default) blocks dangerous operations like DROP, DELETE,
     TRUNCATE, ALTER, UPDATE, and REVOKE. Set SAFE_MODE=false to disable.
 
     Args:
@@ -159,14 +187,14 @@ async def execute_sql(sql: str) -> dict:
 
 
 @mcp.tool()
-def show_databases():
+async def show_databases():
     """List available Databend databases (safe operation, not affected by MCP safe mode)"""
     logger.info("Listing all databases")
     return _execute_sql("SHOW DATABASES")
 
 
 @mcp.tool()
-def show_tables(database: Optional[str] = None, filter: Optional[str] = None):
+async def show_tables(database: Optional[str] = None, filter: Optional[str] = None):
     """
     List available Databend tables in a database (safe operation, not affected by MCP safe mode)
     Args:
@@ -186,7 +214,7 @@ def show_tables(database: Optional[str] = None, filter: Optional[str] = None):
 
 
 @mcp.tool()
-def describe_table(table: str, database: Optional[str] = None):
+async def describe_table(table: str, database: Optional[str] = None):
     """
     Describe a Databend table (safe operation, not affected by MCP safe mode)
     Args:
@@ -201,7 +229,7 @@ def describe_table(table: str, database: Optional[str] = None):
         table = f"{database}.{table}"
     logger.info(f"Describing table '{table}'")
     sql = f"DESCRIBE TABLE {table}"
-    return execute_sql(sql)
+    return _execute_sql(sql)
 
 
 @mcp.tool()
@@ -212,7 +240,7 @@ def show_stages():
 
 
 @mcp.tool()
-def list_stage_files(stage_name: str, path: Optional[str] = None):
+async 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:
@@ -222,28 +250,30 @@ def list_stage_files(stage_name: str, path: Optional[str] = None):
     Returns:
         Dictionary containing either query results or error information
     """
-    if not stage_name.startswith('@'):
+    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():
+async 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:
+async def create_stage(
+    name: str, url: str, connection_name: Optional[str] = None, **kwargs
+) -> dict:
     """
     Create a Databend stage with connection
     Args:
@@ -256,17 +286,17 @@ async def create_stage(name: str, url: str, connection_name: Optional[str] = Non
         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']:
+        if key not in ["name", "url", "connection_name"]:
             sql_parts.append(f"{key.upper()} = '{value}'")
-    
+
     sql = " ".join(sql_parts)
     return _execute_sql(sql)
 

pyproject.toml

@@ -21,8 +21,11 @@ classifiers = [
     "Topic :: Software Development :: Libraries :: Python Modules",
 ]
 dependencies = [
+    "databend>=1.2.810",
     "databend-driver>=0.27.3",
     "mcp>=1.9.0",
+    "pandas>=2.3.2",
+    "pyarrow>=21.0.0",
     "python-dotenv>=1.1.0",
 ]