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

mcp_server_motherduck/server.py

@@ -1,149 +1,242 @@
+"""
+FastMCP Server for MotherDuck and DuckDB.
+
+This module creates and configures the FastMCP server with all tools.
+"""
+
+import json
 import logging
-from pydantic import AnyUrl
-from typing import Literal
-import mcp.types as types
-from mcp.server import NotificationOptions, Server
-from mcp.server.models import InitializationOptions
+from pathlib import Path
+
+from fastmcp import FastMCP
+from fastmcp.utilities.types import Image
+from mcp.types import Icon
+
 from .configs import SERVER_VERSION
 from .database import DatabaseClient
-from .prompt import PROMPT_TEMPLATE
-
+from .instructions import get_instructions
+from .tools.execute_query import execute_query as execute_query_fn
+from .tools.list_columns import list_columns as list_columns_fn
+from .tools.list_databases import list_databases as list_databases_fn
+from .tools.list_tables import list_tables as list_tables_fn
+from .tools.switch_database_connection import (
+    switch_database_connection as switch_database_connection_fn,
+)
 
 logger = logging.getLogger("mcp_server_motherduck")
 
+# Server icon - embedded as data URI from local file
+ASSETS_DIR = Path(__file__).parent / "assets"
+ICON_PATH = ASSETS_DIR / "duck_feet_square.png"
+
 
-def build_application(
+def create_mcp_server(
     db_path: str,
     motherduck_token: str | None = None,
     home_dir: str | None = None,
     saas_mode: bool = False,
     read_only: bool = False,
-):
-    logger.info("Starting MotherDuck MCP Server")
-    server = Server("mcp-server-motherduck")
+    ephemeral_connections: bool = True,
+    max_rows: int = 1024,
+    max_chars: int = 50000,
+    query_timeout: int = -1,
+    init_sql: str | None = None,
+    allow_switch_databases: bool = False,
+) -> FastMCP:
+    """
+    Create and configure the FastMCP server.
+
+    Args:
+        db_path: Path to database (local file, :memory:, md:, or s3://)
+        motherduck_token: MotherDuck authentication token
+        home_dir: Home directory for DuckDB
+        saas_mode: Enable MotherDuck SaaS mode
+        read_only: Enable read-only mode
+        ephemeral_connections: Use temporary connections for read-only local files
+        max_rows: Maximum rows to return from queries
+        max_chars: Maximum characters in query results
+        query_timeout: Query timeout in seconds (-1 to disable)
+        init_sql: SQL file path or string to execute on startup
+        allow_switch_databases: Enable the switch_database_connection tool
+
+    Returns:
+        Configured FastMCP server instance
+    """
+    # Create database client
     db_client = DatabaseClient(
         db_path=db_path,
         motherduck_token=motherduck_token,
         home_dir=home_dir,
         saas_mode=saas_mode,
         read_only=read_only,
+        ephemeral_connections=ephemeral_connections,
+        max_rows=max_rows,
+        max_chars=max_chars,
+        query_timeout=query_timeout,
+        init_sql=init_sql,
     )
 
-    logger.info("Registering handlers")
+    # Get instructions with connection context
+    instructions = get_instructions(
+        read_only=read_only,
+        saas_mode=saas_mode,
+        db_path=db_path,
+        allow_switch_databases=allow_switch_databases,
+    )
 
-    @server.list_resources()
-    async def handle_list_resources() -> list[types.Resource]:
-        """
-        List available note resources.
-        Each note is exposed as a resource with a custom note:// URI scheme.
-        """
-        logger.info("No resources available to list")
-        return []
+    # Create server icon from local file
+    icons = []
+    if ICON_PATH.exists():
+        img = Image(path=str(ICON_PATH))
+        icons.append(Icon(src=img.to_data_uri(), mimeType="image/png"))
+
+    # Create FastMCP server with icon
+    mcp = FastMCP(
+        name="mcp-server-motherduck",
+        instructions=instructions,
+        version=SERVER_VERSION,
+        icons=icons if icons else None,
+    )
 
-    @server.read_resource()
-    async def handle_read_resource(uri: AnyUrl) -> str:
-        """
-        Read a specific note's content by its URI.
-        The note name is extracted from the URI host component.
+    # Define query tool annotations (dynamic based on read_only flag)
+    query_annotations = {
+        "readOnlyHint": read_only,
+        "destructiveHint": not read_only,
+        "openWorldHint": False,
+    }
+
+    # Catalog tool annotations (always read-only)
+    catalog_annotations = {
+        "readOnlyHint": True,
+        "destructiveHint": False,
+        "openWorldHint": False,
+    }
+
+    # Switch database annotations (open world - can connect to any database)
+    switch_db_annotations = {
+        "readOnlyHint": False,
+        "destructiveHint": False,
+        "openWorldHint": True,
+    }
+
+    # Register query tool
+    @mcp.tool(
+        name="execute_query",
+        title="Execute Query",
+        description="Execute a SQL query on the DuckDB or MotherDuck database. Unqualified table names resolve to current_database() and current_schema() automatically. Fully qualified names (database.schema.table) are only needed when multiple DuckDB databases are attached or when connected to MotherDuck.",
+        annotations=query_annotations,
+    )
+    def execute_query(sql: str) -> str:
         """
-        logger.info(f"Reading resource: {uri}")
-        raise ValueError(f"Unsupported URI scheme: {uri.scheme}")
+        Execute a SQL query on the DuckDB or MotherDuck database.
 
-    @server.list_prompts()
-    async def handle_list_prompts() -> list[types.Prompt]:
+        Args:
+            sql: SQL query to execute (DuckDB SQL dialect)
+
+        Returns:
+            JSON string with query results
+
+        Raises:
+            ValueError: If the query fails
         """
-        List available prompts.
-        Each prompt can have optional arguments to customize its behavior.
+        result = execute_query_fn(sql, db_client)
+        if not result.get("success", True):
+            # Raise exception so FastMCP marks as isError=True
+            raise ValueError(json.dumps(result, indent=2, default=str))
+        return json.dumps(result, indent=2, default=str)
+
+    # Register list_databases tool
+    @mcp.tool(
+        name="list_databases",
+        title="List Databases",
+        description="List all databases available in the connection. Useful when multiple DuckDB databases are attached or when connected to MotherDuck.",
+        annotations=catalog_annotations,
+    )
+    def list_databases_tool() -> str:
         """
-        logger.info("Listing prompts")
-        # TODO: Check where and how this is used, and how to optimize this.
-        # Check postgres and sqlite servers.
-        return [
-            types.Prompt(
-                name="duckdb-motherduck-initial-prompt",
-                description="A prompt to initialize a connection to duckdb or motherduck and start working with it",
-            )
-        ]
+        List all databases available in the connection.
 
-    @server.get_prompt()
-    async def handle_get_prompt(
-        name: str, arguments: dict[str, str] | None
-    ) -> types.GetPromptResult:
+        Returns:
+            JSON string with database list
         """
-        Generate a prompt by combining arguments with server state.
-        The prompt includes all current notes and can be customized via arguments.
+        result = list_databases_fn(db_client)
+        return json.dumps(result, indent=2, default=str)
+
+    # Register list_tables tool
+    @mcp.tool(
+        name="list_tables",
+        title="List Tables",
+        description="List all tables and views in a database with their comments. If database is not specified, uses the current database.",
+        annotations=catalog_annotations,
+    )
+    def list_tables(database: str | None = None, schema: str | None = None) -> str:
         """
-        logger.info(f"Getting prompt: {name}::{arguments}")
-        # TODO: Check where and how this is used, and how to optimize this.
-        # Check postgres and sqlite servers.
-        if name != "duckdb-motherduck-initial-prompt":
-            raise ValueError(f"Unknown prompt: {name}")
-
-        return types.GetPromptResult(
-            description="Initial prompt for interacting with DuckDB/MotherDuck",
-            messages=[
-                types.PromptMessage(
-                    role="user",
-                    content=types.TextContent(type="text", text=PROMPT_TEMPLATE),
-                )
-            ],
-        )
+        List all tables and views in a database.
 
-    @server.list_tools()
-    async def handle_list_tools() -> list[types.Tool]:
-        """
-        List available tools.
-        Each tool specifies its arguments using JSON Schema validation.
+        Args:
+            database: Database name to list tables from (defaults to current database)
+            schema: Optional schema name to filter by
+
+        Returns:
+            JSON string with table/view list
         """
-        logger.info("Listing tools")
-        return [
-            types.Tool(
-                name="query",
-                description="Use this to execute a query on the MotherDuck or DuckDB database",
-                inputSchema={
-                    "type": "object",
-                    "properties": {
-                        "query": {
-                            "type": "string",
-                            "description": "SQL query to execute that is a dialect of DuckDB SQL",
-                        },
-                    },
-                    "required": ["query"],
-                },
-            ),
-        ]
-
-    @server.call_tool()
-    async def handle_tool_call(
-        name: str, arguments: dict | None
-    ) -> list[types.TextContent | types.ImageContent | types.EmbeddedResource]:
+        result = list_tables_fn(db_client, database, schema)
+        return json.dumps(result, indent=2, default=str)
+
+    # Register list_columns tool
+    @mcp.tool(
+        name="list_columns",
+        title="List Columns",
+        description="List all columns of a table or view with their types and comments. If database/schema are not specified, uses the current database/schema.",
+        annotations=catalog_annotations,
+    )
+    def list_columns(table: str, database: str | None = None, schema: str | None = None) -> str:
         """
-        Handle tool execution requests.
-        Tools can modify server state and notify clients of changes.
+        List all columns of a table or view.
+
+        Args:
+            table: Table or view name
+            database: Database name (defaults to current database)
+            schema: Schema name (defaults to current schema)
+
+        Returns:
+            JSON string with column list
         """
-        logger.info(f"Calling tool: {name}::{arguments}")
-        try:
-            if name == "query":
-                if arguments is None:
-                    return [
-                        types.TextContent(type="text", text="Error: No query provided")
-                    ]
-                tool_response = db_client.query(arguments["query"])
-                return [types.TextContent(type="text", text=str(tool_response))]
-
-            return [types.TextContent(type="text", text=f"Unsupported tool: {name}")]
-
-        except Exception as e:
-            logger.error(f"Error executing tool {name}: {e}")
-            raise ValueError(f"Error executing tool {name}: {str(e)}")
-
-    initialization_options = InitializationOptions(
-        server_name="motherduck",
-        server_version=SERVER_VERSION,
-        capabilities=server.get_capabilities(
-            notification_options=NotificationOptions(),
-            experimental_capabilities={},
-        ),
-    )
+        result = list_columns_fn(table, db_client, database, schema)
+        return json.dumps(result, indent=2, default=str)
+
+    # Conditionally register switch_database_connection tool
+    if allow_switch_databases:
+        # Store server's read_only setting for switch_database_connection
+        server_read_only_mode = read_only
+
+        @mcp.tool(
+            name="switch_database_connection",
+            title="Switch Database Connection",
+            description="Switch to a different database connection. For local files, use absolute paths only. The new connection respects the server's read-only/read-write mode. For local files, the file must exist unless create_if_not_exists=True (requires read-write mode).",
+            annotations=switch_db_annotations,
+        )
+        def switch_database_connection(path: str, create_if_not_exists: bool = False) -> str:
+            """
+            Switch to a different primary database.
+
+            Args:
+                path: Database path. For local files, must be an absolute path.
+                      Also accepts :memory:, md:database_name, or s3:// paths.
+                create_if_not_exists: If True, create the database file if it doesn't exist.
+                                   Only works in read-write mode.
+
+            Returns:
+                JSON string with result
+            """
+            result = switch_database_connection_fn(
+                path=path,
+                db_client=db_client,
+                server_read_only=server_read_only_mode,
+                create_if_not_exists=create_if_not_exists,
+            )
+            return json.dumps(result, indent=2, default=str)
+
+    logger.info(f"FastMCP server created with {len(mcp._tool_manager._tools)} tools")
 
-    return server, initialization_options
+    return mcp

mcp_server_motherduck/tools/__init__.py

@@ -0,0 +1,19 @@
+"""
+MCP Tools for MotherDuck/DuckDB server.
+
+Each tool is defined in its own module and exported here.
+"""
+
+from .execute_query import execute_query
+from .list_columns import list_columns
+from .list_databases import list_databases
+from .list_tables import list_tables
+from .switch_database_connection import switch_database_connection
+
+__all__ = [
+    "execute_query",
+    "list_databases",
+    "list_tables",
+    "list_columns",
+    "switch_database_connection",
+]

mcp_server_motherduck/tools/execute_query.py

@@ -0,0 +1,21 @@
+"""
+Execute Query tool - Execute SQL queries against DuckDB/MotherDuck databases.
+"""
+
+from typing import Any
+
+DESCRIPTION = "Execute a SQL query on the DuckDB or MotherDuck database."
+
+
+def execute_query(sql: str, db_client: Any) -> dict[str, Any]:
+    """
+    Execute a SQL query on the DuckDB or MotherDuck database.
+
+    Args:
+        sql: SQL query to execute (DuckDB SQL dialect)
+        db_client: DatabaseClient instance (injected by server)
+
+    Returns:
+        JSON-serializable dict with query results or error
+    """
+    return db_client.query(sql)

mcp_server_motherduck/tools/list_columns.py

@@ -0,0 +1,99 @@
+"""
+List columns tool - List all columns of a table or view.
+"""
+
+from typing import Any
+
+DESCRIPTION = (
+    "List all columns of a table or view with their types and comments. "
+    "If database/schema are not specified, uses the current database/schema."
+)
+
+
+def list_columns(
+    table: str,
+    db_client: Any,
+    database: str | None = None,
+    schema: str | None = None,
+) -> dict[str, Any]:
+    """
+    List all columns of a table or view.
+
+    Args:
+        table: Table or view name
+        db_client: DatabaseClient instance (injected by server)
+        database: Database name (defaults to current database)
+        schema: Schema name (defaults to current schema)
+
+    Returns:
+        JSON-serializable dict with column list or error
+    """
+    try:
+        # Get current database if not specified
+        if database is None:
+            _, _, db_rows = db_client.execute_raw("SELECT current_database()")
+            database = db_rows[0][0]
+
+        # Get current schema if not specified
+        if schema is None:
+            _, _, schema_rows = db_client.execute_raw("SELECT current_schema()")
+            schema = schema_rows[0][0]
+
+        # Query columns using DuckDB system function
+        sql = f"""
+            SELECT
+                column_name as name,
+                data_type as type,
+                is_nullable = 'YES' as nullable,
+                comment
+            FROM duckdb_columns()
+            WHERE database_name = '{database}'
+              AND schema_name = '{schema}'
+              AND table_name = '{table}'
+            ORDER BY column_index
+        """
+
+        _, _, rows = db_client.execute_raw(sql)
+
+        # Transform results
+        columns = [
+            {
+                "name": row[0],
+                "type": row[1],
+                "nullable": bool(row[2]),
+                "comment": row[3] if row[3] else None,
+            }
+            for row in rows
+        ]
+
+        # Determine if it's a view or table
+        object_type = "table"
+        try:
+            _, _, view_rows = db_client.execute_raw(f"""
+                SELECT 1 FROM duckdb_views()
+                WHERE database_name = '{database}'
+                  AND schema_name = '{schema}'
+                  AND view_name = '{table}'
+                LIMIT 1
+            """)
+            if view_rows:
+                object_type = "view"
+        except Exception:
+            pass  # Assume table if check fails
+
+        return {
+            "success": True,
+            "database": database,
+            "schema": schema,
+            "table": table,
+            "objectType": object_type,
+            "columns": columns,
+            "columnCount": len(columns),
+        }
+
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "errorType": type(e).__name__,
+        }

mcp_server_motherduck/tools/list_databases.py

@@ -0,0 +1,52 @@
+"""
+List databases tool - Show all databases available in the connection.
+"""
+
+from typing import Any
+
+DESCRIPTION = "List all databases with their names and types."
+
+
+def list_databases(db_client: Any) -> dict[str, Any]:
+    """
+    List all databases available in the connection.
+
+    For MotherDuck: Uses MD_ALL_DATABASES() to list all databases.
+    For local DuckDB: Uses duckdb_databases() system function.
+
+    Excludes internal databases: 'system' and 'temp'.
+
+    Args:
+        db_client: DatabaseClient instance (injected by server)
+
+    Returns:
+        JSON-serializable dict with database list or error
+    """
+    try:
+        # Try MotherDuck function first (works for MotherDuck connections)
+        try:
+            _, _, rows = db_client.execute_raw(
+                "SELECT alias, type FROM MD_ALL_DATABASES() "
+                "WHERE alias IS NOT NULL AND alias NOT IN ('system', 'temp')"
+            )
+            databases = [{"name": row[0], "type": row[1]} for row in rows]
+        except Exception:
+            # Fall back to DuckDB system function (works for local DuckDB)
+            _, _, rows = db_client.execute_raw(
+                "SELECT database_name, type FROM duckdb_databases() "
+                "WHERE database_name NOT IN ('system', 'temp')"
+            )
+            databases = [{"name": row[0], "type": row[1]} for row in rows]
+
+        return {
+            "success": True,
+            "databases": databases,
+            "databaseCount": len(databases),
+        }
+
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "errorType": type(e).__name__,
+        }

mcp_server_motherduck/tools/list_tables.py

@@ -0,0 +1,91 @@
+"""
+List tables tool - List all tables and views in a database.
+"""
+
+from typing import Any
+
+DESCRIPTION = (
+    "List all tables and views in a database with their comments. "
+    "If database is not specified, uses the current database."
+)
+
+
+def list_tables(
+    db_client: Any,
+    database: str | None = None,
+    schema: str | None = None,
+) -> dict[str, Any]:
+    """
+    List all tables and views in a database.
+
+    Args:
+        db_client: DatabaseClient instance (injected by server)
+        database: Database name to list tables from (defaults to current database)
+        schema: Optional schema name to filter by (defaults to all schemas)
+
+    Returns:
+        JSON-serializable dict with table/view list or error
+    """
+    try:
+        # Get current database if not specified
+        if database is None:
+            _, _, db_rows = db_client.execute_raw("SELECT current_database()")
+            database = db_rows[0][0]
+
+        # Build schema filter
+        schema_filter = f"AND schema_name = '{schema}'" if schema else ""
+
+        # Query tables and views using DuckDB system functions
+        sql = f"""
+            SELECT
+                schema_name as schema,
+                table_name as name,
+                'table' as type,
+                comment
+            FROM duckdb_tables()
+            WHERE database_name = '{database}' {schema_filter}
+
+            UNION ALL
+
+            SELECT
+                schema_name as schema,
+                view_name as name,
+                'view' as type,
+                comment
+            FROM duckdb_views()
+            WHERE database_name = '{database}' {schema_filter}
+
+            ORDER BY schema, type, name
+        """
+
+        _, _, rows = db_client.execute_raw(sql)
+
+        # Transform results
+        tables = [
+            {
+                "schema": row[0],
+                "name": row[1],
+                "type": row[2],
+                "comment": row[3] if row[3] else None,
+            }
+            for row in rows
+        ]
+
+        table_count = sum(1 for t in tables if t["type"] == "table")
+        view_count = sum(1 for t in tables if t["type"] == "view")
+
+        return {
+            "success": True,
+            "database": database,
+            "schema": schema or "all",
+            "tables": tables,
+            "tableCount": table_count,
+            "viewCount": view_count,
+        }
+
+    except Exception as e:
+        return {
+            "success": False,
+            "error": str(e),
+            "errorType": type(e).__name__,
+        }