mcp-server-motherduck 0.7.2__py3-none-any.whl → 1.0.0__py3-none-any.whl

mcp_server_motherduck/database.py

@@ -1,15 +1,39 @@
+import json
+import logging
 import os
+import re
+import threading
+from typing import Any, Literal, Optional
+
 import duckdb
-from typing import Literal, Optional
-import io
-from contextlib import redirect_stdout
-from tabulate import tabulate
-import logging
+
 from .configs import SERVER_VERSION
 
 logger = logging.getLogger("mcp_server_motherduck")
 
 
+def _is_read_scaling_connection(conn: duckdb.DuckDBPyConnection) -> bool:
+    """
+    Check if a MotherDuck connection is using read-scaling.
+
+    Read-scaling connections have a duckling ID ending with .rs.{number}
+    e.g., "my_database.rs.3", "app_db.rs.0"
+
+    Read-write connections end with .rw
+    e.g., "my_database.rw", "app_db.rw"
+    """
+    try:
+        # __md_duckling_id() is a table function, must use FROM clause
+        result = conn.execute("SELECT * FROM __md_duckling_id()").fetchone()
+        if result and result[0]:
+            duckling_id = result[0]
+            # Check if duckling ID ends with .rs.{number}
+            return bool(re.search(r"\.rs\.\d+$", duckling_id))
+        return False
+    except Exception:
+        return False
+
+
 class DatabaseClient:
     def __init__(
         self,
@@ -18,8 +42,18 @@ class DatabaseClient:
         home_dir: str | None = None,
         saas_mode: bool = False,
         read_only: bool = False,
+        ephemeral_connections: bool = True,
+        max_rows: int = 1024,
+        max_chars: int = 50000,
+        query_timeout: int = -1,
+        init_sql: str | None = None,
     ):
         self._read_only = read_only
+        self._ephemeral_connections = ephemeral_connections
+        self._max_rows = max_rows
+        self._max_chars = max_chars
+        self._query_timeout = query_timeout
+        self._init_sql = init_sql
         self.db_path, self.db_type = self._resolve_db_path_type(
             db_path, motherduck_token, saas_mode
         )
@@ -36,23 +70,30 @@ class DatabaseClient:
 
         logger.info(f"🔌 Connecting to {self.db_type} database")
 
-        # S3 databases don't support read-only mode
-        if self.db_type == "s3" and self._read_only:
-            raise ValueError("Read-only mode is not supported for S3 databases")
+        # Read-only handling for local DuckDB files (not in-memory)
+        is_local_file = self.db_type == "duckdb" and self.db_path != ":memory:"
 
-        if self.db_type == "duckdb" and self._read_only:
-            # check that we can connect, issue a `select 1` and then close + return None
+        if is_local_file and self._read_only:
+            # For read-only local DuckDB files, use short-lived connections by default
+            # to allow concurrent access from other processes
             try:
                 conn = duckdb.connect(
                     self.db_path,
-                    config={
-                        "custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"
-                    },
+                    config={"custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"},
                     read_only=self._read_only,
                 )
                 conn.execute("SELECT 1")
-                conn.close()
-                return None
+
+                if self._ephemeral_connections:
+                    # Default: close connection for concurrent access
+                    conn.close()
+                    return None
+                else:
+                    # User requested persistent connection via --no-ephemeral-connections
+                    logger.info("Using persistent read-only connection")
+                    # Execute init SQL
+                    self._execute_init_sql(conn)
+                    return conn
             except Exception as e:
                 logger.error(f"❌ Read-only check failed: {e}")
                 raise
@@ -60,27 +101,27 @@ class DatabaseClient:
         # Check if this is an S3 path
         if self.db_type == "s3":
             # For S3, we need to create an in-memory connection and attach the S3 database
-            conn = duckdb.connect(':memory:')
-            
+            conn = duckdb.connect(":memory:")
+
             # Install and load the httpfs extension for S3 support
             import io
-            from contextlib import redirect_stdout, redirect_stderr
-            
+            from contextlib import redirect_stderr, redirect_stdout
+
             null_file = io.StringIO()
             with redirect_stdout(null_file), redirect_stderr(null_file):
                 try:
                     conn.execute("INSTALL httpfs;")
-                except:
+                except Exception:
                     pass  # Extension might already be installed
                 conn.execute("LOAD httpfs;")
-            
+
             # Configure S3 credentials from environment variables using CREATE SECRET
-            aws_access_key = os.environ.get('AWS_ACCESS_KEY_ID')
-            aws_secret_key = os.environ.get('AWS_SECRET_ACCESS_KEY')
-            aws_region = os.environ.get('AWS_DEFAULT_REGION', 'us-east-1')
-            
-            
-            if aws_access_key and aws_secret_key:
+            aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
+            aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
+            aws_session_token = os.environ.get("AWS_SESSION_TOKEN")
+            aws_region = os.environ.get("AWS_DEFAULT_REGION", "us-east-1")
+
+            if aws_access_key and aws_secret_key and not aws_session_token:
                 # Use CREATE SECRET for better credential management
                 conn.execute(f"""
                     CREATE SECRET IF NOT EXISTS s3_secret (
@@ -90,7 +131,17 @@ class DatabaseClient:
                         REGION '{aws_region}'
                     );
                 """)
-            
+            elif aws_session_token:
+                # Use credential_chain provider to automatically fetch credentials
+                # This supports IAM roles, SSO, instance profiles, etc.
+                conn.execute(f"""
+                    CREATE SECRET IF NOT EXISTS s3_secret (
+                        TYPE S3,
+                        PROVIDER credential_chain,
+                        REGION '{aws_region}'
+                    );
+                """)
+
             # Attach the S3 database
             try:
                 # For S3, we always attach as READ_ONLY since S3 storage is typically read-only
@@ -98,7 +149,9 @@ class DatabaseClient:
                 conn.execute(f"ATTACH '{self.db_path}' AS s3db (READ_ONLY);")
                 # Use the attached database
                 conn.execute("USE s3db;")
-                logger.info(f"✅ Successfully connected to {self.db_type} database (attached as read-only)")
+                logger.info(
+                    f"✅ Successfully connected to {self.db_type} database (attached as read-only)"
+                )
             except Exception as e:
                 logger.error(f"Failed to attach S3 database: {e}")
                 # If the database doesn't exist and we're not in read-only mode, try to create it
@@ -114,19 +167,62 @@ class DatabaseClient:
                         raise
                 else:
                     raise
-                
+
+            # Execute init SQL
+            self._execute_init_sql(conn)
             return conn
 
+        # For MotherDuck, pass read_only flag; for in-memory it's not applicable
+        read_only_flag = self._read_only if self.db_type == "motherduck" else False
+
         conn = duckdb.connect(
             self.db_path,
             config={"custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"},
-            read_only=self._read_only,
+            read_only=read_only_flag,
         )
 
         logger.info(f"✅ Successfully connected to {self.db_type} database")
 
+        # For MotherDuck with --read-only flag, verify it's a read-scaling connection
+        if self.db_type == "motherduck" and self._read_only:
+            if not _is_read_scaling_connection(conn):
+                conn.close()
+                raise ValueError(
+                    "The --read-only flag with MotherDuck requires a read-scaling token. "
+                    "You appear to be using a read/write token. Please use a read-scaling token instead. "
+                    "See: https://motherduck.com/docs/key-tasks/authenticating-and-connecting-to-motherduck/"
+                )
+            logger.info("Verified read-scaling connection for --read-only mode")
+
+        # Execute init SQL
+        self._execute_init_sql(conn)
+
         return conn
 
+    def _execute_init_sql(self, conn: duckdb.DuckDBPyConnection) -> None:
+        """Execute initialization SQL if provided."""
+        if not self._init_sql:
+            return
+
+        try:
+            # Check if init_sql is a file path
+            if os.path.isfile(self._init_sql):
+                logger.info(f"Loading init SQL from file: {self._init_sql}")
+                with open(self._init_sql) as f:
+                    sql_content = f.read()
+            else:
+                # Treat as raw SQL string
+                logger.info("Executing init SQL string")
+                sql_content = self._init_sql
+
+            # Execute the SQL
+            conn.execute(sql_content)
+            logger.info("Init SQL executed successfully")
+
+        except Exception as e:
+            logger.error(f"Failed to execute init SQL: {e}")
+            raise ValueError(f"Init SQL execution failed: {e}") from e
+
     def _resolve_db_path_type(
         self, db_path: str, motherduck_token: str | None = None, saas_mode: bool = False
     ) -> tuple[str, Literal["duckdb", "motherduck", "s3"]]:
@@ -134,7 +230,7 @@ class DatabaseClient:
         # Handle S3 paths
         if db_path.startswith("s3://"):
             return db_path, "s3"
-        
+
         # Handle MotherDuck paths
         if db_path.startswith("md:"):
             if motherduck_token:
@@ -150,17 +246,16 @@ class DatabaseClient:
                         f"{db_path}?motherduck_token={motherduck_token}",
                         "motherduck",
                     )
-            elif os.getenv("motherduck_token"):
-                logger.info(
-                    "Using MotherDuck token from env to connect to database `md:`"
-                )
+            elif os.getenv("motherduck_token") or os.getenv("MOTHERDUCK_TOKEN"):
+                token = os.getenv("motherduck_token") or os.getenv("MOTHERDUCK_TOKEN")
+                logger.info("Using MotherDuck token from env to connect to database `md:`")
                 return (
-                    f"{db_path}?motherduck_token={os.getenv('motherduck_token')}",
+                    f"{db_path}?motherduck_token={token}",
                     "motherduck",
                 )
             else:
                 raise ValueError(
-                    "Please set the `motherduck_token` as an environment variable or pass it as an argument with `--motherduck-token` when using `md:` as db_path."
+                    "Please set the `motherduck_token` or `MOTHERDUCK_TOKEN` as an environment variable or pass it as an argument with `--motherduck-token` when using `md:` as db_path."
                 )
 
         if db_path == ":memory:":
@@ -168,32 +263,176 @@ class DatabaseClient:
 
         return db_path, "duckdb"
 
-    def _execute(self, query: str) -> str:
+    def _execute(self, query: str) -> dict[str, Any]:
+        """Execute query and return JSON-serializable result."""
+        # Get connection to use
         if self.conn is None:
-            # open short lived readonly connection for local DuckDB, run query, close connection, return result
             conn = duckdb.connect(
                 self.db_path,
                 config={"custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"},
                 read_only=self._read_only,
             )
-            q = conn.execute(query)
         else:
-            q = self.conn.execute(query)
+            conn = self.conn
 
-        out = tabulate(
-            q.fetchall(),
-            headers=[d[0] + "\n" + str(d[1]) for d in q.description],
-            tablefmt="pretty",
-        )
+        try:
+            # Execute with or without timeout
+            if self._query_timeout > 0:
+                columns, column_types, rows, has_more_rows = self._execute_with_timeout(conn, query)
+            else:
+                columns, column_types, rows, has_more_rows = self._execute_direct(conn, query)
 
-        if self.conn is None:
-            conn.close()
+            # Build result object
+            result: dict[str, Any] = {
+                "success": True,
+                "columns": columns,
+                "columnTypes": column_types,
+                "rows": rows,
+                "rowCount": len(rows),
+            }
+
+            # Add row truncation warning
+            if has_more_rows:
+                result["truncated"] = True
+                result["warning"] = (
+                    f"Results limited to {self._max_rows:,} rows. Query returned more data."
+                )
+
+            # Check character limit on JSON output
+            json_output = json.dumps(result, default=str)
+            if len(json_output) > self._max_chars:
+                # Progressively reduce rows until under limit
+                while rows and len(json_output) > self._max_chars:
+                    # Remove ~10% of rows each iteration
+                    remove_count = max(1, len(rows) // 10)
+                    rows = rows[:-remove_count]
+                    result["rows"] = rows
+                    result["rowCount"] = len(rows)
+                    result["truncated"] = True
+                    result["warning"] = (
+                        f"Results limited to {len(rows):,} rows due to "
+                        f"{self._max_chars // 1000}KB output size limit."
+                    )
+                    json_output = json.dumps(result, default=str)
+
+            return result
+
+        finally:
+            # Close connection if it was temporary
+            if self.conn is None:
+                conn.close()
+
+    def _execute_direct(
+        self, conn: duckdb.DuckDBPyConnection, query: str
+    ) -> tuple[list[str], list[str], list[list[Any]], bool]:
+        """Execute query without timeout - returns columns, types, rows, has_more."""
+        q = conn.execute(query)
 
-        return out
+        # Get column metadata
+        columns = [d[0] for d in q.description] if q.description else []
+        column_types = [str(d[1]) for d in q.description] if q.description else []
+
+        # Fetch rows (max_rows + 1 to detect truncation)
+        raw_rows = q.fetchmany(self._max_rows + 1)
+        has_more_rows = len(raw_rows) > self._max_rows
+        if has_more_rows:
+            raw_rows = raw_rows[: self._max_rows]
+
+        # Convert rows to JSON-serializable lists
+        rows = [list(row) for row in raw_rows]
+
+        return columns, column_types, rows, has_more_rows
+
+    def _execute_with_timeout(
+        self, conn: duckdb.DuckDBPyConnection, query: str
+    ) -> tuple[list[str], list[str], list[list[Any]], bool]:
+        """Execute query with timeout using threading.Timer and conn.interrupt()."""
+        timer = threading.Timer(self._query_timeout, conn.interrupt)
+        timer.start()
 
-    def query(self, query: str) -> str:
         try:
-            return self._execute(query)
+            return self._execute_direct(conn, query)
+        except duckdb.InterruptException:
+            raise ValueError(
+                f"Query execution timed out after {self._query_timeout} seconds. "
+                "Increase timeout with --query-timeout argument when starting the mcp server."
+            )
+        finally:
+            timer.cancel()
 
+    def query(self, query: str) -> dict[str, Any]:
+        """Execute a SQL query and return JSON-serializable result."""
+        try:
+            return self._execute(query)
+        except ValueError:
+            # Re-raise ValueError (timeout, etc.) as-is
+            raise
         except Exception as e:
-            raise ValueError(f"❌ Error executing query: {e}")
+            # Return error as structured response
+            return {
+                "success": False,
+                "error": str(e),
+                "errorType": type(e).__name__,
+            }
+
+    def execute_raw(self, query: str) -> tuple[list[str], list[str], list[list[Any]]]:
+        """
+        Execute a query and return raw results (columns, types, rows).
+        Used by catalog tools that need custom result formatting.
+        """
+        if self.conn is None:
+            conn = duckdb.connect(
+                self.db_path,
+                config={"custom_user_agent": f"mcp-server-motherduck/{SERVER_VERSION}"},
+                read_only=self._read_only,
+            )
+        else:
+            conn = self.conn
+
+        try:
+            q = conn.execute(query)
+            columns = [d[0] for d in q.description] if q.description else []
+            column_types = [str(d[1]) for d in q.description] if q.description else []
+            rows = [list(row) for row in q.fetchall()]
+            return columns, column_types, rows
+        finally:
+            if self.conn is None:
+                conn.close()
+
+    def switch_database(self, path: str, read_only: bool = True) -> None:
+        """
+        Switch to a different primary database.
+
+        Closes any existing connection and updates the database path.
+        The next query will connect to the new database.
+
+        Args:
+            path: New database path (local file, :memory:, or md:database_name)
+            read_only: Whether to connect in read-only mode
+        """
+        # Close existing connection if any
+        if self.conn is not None:
+            try:
+                self.conn.close()
+            except Exception:
+                pass  # Ignore close errors
+            self.conn = None
+
+        # Update database configuration
+        self.db_path = path
+        self._read_only = read_only
+
+        # Determine new database type
+        if path.startswith("md:") or path.startswith("motherduck:"):
+            self.db_type = "motherduck"
+        elif path.startswith("s3://"):
+            self.db_type = "s3"
+        elif path == ":memory:":
+            self.db_type = "memory"
+        else:
+            self.db_type = "duckdb"
+
+        # Re-initialize connection (will be None for read-only local DuckDB)
+        self.conn = self._initialize_connection()
+
+        logger.info(f"Switched to database: {path} (read_only={read_only})")

mcp_server_motherduck/instructions.py

@@ -0,0 +1,187 @@
+"""
+Server instructions for DuckDB/MotherDuck MCP Server.
+
+These instructions are sent to the client during initialization
+to provide context about how to use the server's capabilities.
+"""
+
+INSTRUCTIONS_BASE = """Execute SQL queries against DuckDB and MotherDuck databases using DuckDB SQL syntax.
+
+## Available Tools
+
+- `execute_query`: Execute SQL queries (DuckDB SQL dialect)
+- `list_databases`: List all available databases
+- `list_tables`: List tables and views in a database
+- `list_columns`: List columns of a table or view
+
+## DuckDB SQL Quick Reference
+
+**Name Qualification**
+- Format: `database.schema.table` or just `schema.table` or `table`
+- Default schema is `main`: `db.table` = `db.main.table`
+- Use fully qualified names when joining tables across different databases
+
+**Identifiers and Literals:**
+- Use double quotes (`"`) for identifiers with spaces/special characters or case-sensitivity
+- Use single quotes (`'`) for string literals
+
+**Flexible Query Structure:**
+- Queries can start with `FROM`: `FROM my_table WHERE condition;`
+- `SELECT` without `FROM` for expressions: `SELECT 1 + 1 AS result;`
+- Support for `CREATE TABLE AS` (CTAS): `CREATE TABLE new_table AS SELECT * FROM old_table;`
+
+**Advanced Column Selection:**
+- Exclude columns: `SELECT * EXCLUDE (sensitive_data) FROM users;`
+- Replace columns: `SELECT * REPLACE (UPPER(name) AS name) FROM users;`
+- Pattern matching: `SELECT COLUMNS('sales_.*') FROM sales_data;`
+
+**Grouping and Ordering Shortcuts:**
+- Group by all non-aggregated columns: `SELECT category, SUM(sales) FROM sales_data GROUP BY ALL;`
+- Order by all columns: `SELECT * FROM my_table ORDER BY ALL;`
+
+**Complex Data Types:**
+- Lists: `SELECT [1, 2, 3] AS my_list;`
+- Structs: `SELECT {'a': 1, 'b': 'text'} AS my_struct;`
+- Maps: `SELECT MAP([1,2],['one','two']) AS my_map;`
+- JSON: `json_col->>'key'` (returns text) or `data->'$.user.id'` (returns JSON)
+
+**Date/Time Operations:**
+- String to timestamp: `strptime('2023-07-23', '%Y-%m-%d')::TIMESTAMP`
+- Format timestamp: `strftime(NOW(), '%Y-%m-%d')`
+- Extract parts: `EXTRACT(YEAR FROM DATE '2023-07-23')`
+
+### Schema Exploration
+
+```sql
+-- List all databases
+SELECT database_name, type FROM duckdb_databases();
+
+-- For MotherDuck: List all databases (including shared)
+SELECT alias as database_name, type FROM MD_ALL_DATABASES();
+
+-- List tables in a database
+SELECT schema_name, table_name FROM duckdb_tables()
+WHERE database_name = 'your_db';
+
+-- Get column info
+SELECT column_name, data_type FROM duckdb_columns()
+WHERE database_name = 'your_db' AND table_name = 'your_table';
+
+-- Quick preview with statistics
+SUMMARIZE your_table;
+```
+
+### Query Best Practices
+
+- Filter early to reduce data volume before blocking operations
+- Use CTEs to break complex queries into manageable parts
+- Avoid unnecessary `ORDER BY` on intermediate results
+- Use `arg_max()` and `arg_min()` for "most recent" queries
+- Use `QUALIFY` for filtering window function results
+
+```sql
+-- Get top 2 products by sales in each category
+SELECT category, product_name, sales_amount
+FROM products
+QUALIFY ROW_NUMBER() OVER (PARTITION BY category ORDER BY sales_amount DESC) <= 2;
+```
+
+### Persisting In-Memory Data to File
+
+To save an in-memory database to a persistent file:
+
+```sql
+-- Attach a new file-based database
+ATTACH '/path/to/my_database.db' AS my_db;
+
+-- Copy all data from memory to the file
+COPY FROM DATABASE memory TO my_db;
+
+-- Optionally detach when done
+DETACH my_db;
+```
+"""
+
+
+def get_instructions(
+    read_only: bool = False,
+    saas_mode: bool = False,
+    db_path: str = ":memory:",
+    allow_switch_databases: bool = False,
+) -> str:
+    """
+    Get server instructions with connection context.
+
+    Args:
+        read_only: Whether the server is in read-only mode
+        saas_mode: Whether MotherDuck is in SaaS mode
+        db_path: The database path being used
+        allow_switch_databases: Whether database switching is enabled
+
+    Returns:
+        Instructions string with context header
+    """
+    context_lines = []
+
+    # Database info
+    if db_path == ":memory:":
+        context_lines.append("- **Database**: In-memory (data will not persist after session ends)")
+    elif db_path.startswith("md:"):
+        context_lines.append(f"- **Database**: MotherDuck cloud database (`{db_path}`)")
+    elif db_path.startswith("s3://"):
+        context_lines.append(
+            f"- **Database**: S3-hosted DuckDB file (`{db_path}`) - always read-only"
+        )
+    else:
+        context_lines.append(f"- **Database**: Local DuckDB file (`{db_path}`)")
+
+    # Access mode
+    if read_only:
+        context_lines.append(
+            "- **Access mode**: Read-only - CREATE, INSERT, UPDATE, DELETE, and DROP operations are disabled"
+        )
+    else:
+        context_lines.append("- **Access mode**: Read-write - all SQL operations are allowed")
+
+    # Security modes
+    if saas_mode:
+        context_lines.append(
+            "- **SaaS mode**: Enabled - local filesystem access is restricted for security"
+        )
+
+    # Available tools
+    tools = ["execute_query", "list_databases", "list_tables", "list_columns"]
+    if allow_switch_databases:
+        tools.append("switch_database_connection")
+    context_lines.append(f"- **Available tools**: {', '.join(tools)}")
+
+    # Implications for the agent
+    context_lines.append("")
+    context_lines.append("### Important Implications")
+
+    if db_path == ":memory:" and not read_only:
+        context_lines.append("- Data created in this session will be lost when the session ends")
+        context_lines.append(
+            "- To persist data, use ATTACH and COPY FROM DATABASE (see 'Persisting In-Memory Data to File' below)"
+        )
+
+    if read_only:
+        context_lines.append(
+            "- You can only query existing data; any attempt to modify data will fail"
+        )
+        context_lines.append("- Use this mode for safe data exploration and analysis")
+
+    if allow_switch_databases and not read_only:
+        context_lines.append(
+            "- You can switch to different databases using switch_database_connection"
+        )
+        context_lines.append(
+            "- To create a new database file, use create_if_not_exists=True (only in read-write mode)"
+        )
+    elif allow_switch_databases and read_only:
+        context_lines.append(
+            "- You can switch to different existing databases, but cannot create new ones"
+        )
+
+    context = "## Server Configuration\n\n" + "\n".join(context_lines) + "\n\n"
+    return context + INSTRUCTIONS_BASE