awslabs.documentdb-mcp-server 0.0.1__py3-none-any.whl

awslabs/documentdb_mcp_server/db_management_tools.py

@@ -0,0 +1,176 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Database management tools for DocumentDB MCP Server."""
+
+from awslabs.documentdb_mcp_server.config import serverConfig
+from awslabs.documentdb_mcp_server.connection_tools import DocumentDBConnection
+from loguru import logger
+from pydantic import Field
+from typing import Annotated, Any, Dict, List
+
+
+async def list_databases(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+) -> Dict[str, Any]:
+    """List all available databases in the DocumentDB cluster.
+
+    This tool returns the names of all accessible databases in the connected cluster.
+
+    Returns:
+        Dict[str, Any]: List of database names
+    """
+    try:
+        client = DocumentDBConnection.get_connection(connection_id)
+        databases = client.list_database_names()
+        logger.info(f'Found {len(databases)} databases')
+        return {'databases': databases, 'count': len(databases)}
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error listing databases: {str(e)}')
+        raise ValueError(f'Failed to list databases: {str(e)}')
+
+
+async def create_collection(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection to create')],
+) -> Dict[str, Any]:
+    """Create a new collection in a DocumentDB database.
+
+    This tool creates a new collection in the specified database.
+
+    Returns:
+        Dict[str, Any]: Status of collection creation
+    """
+    # Check if server is in read-only mode
+    if serverConfig.read_only_mode:
+        logger.warning('Create collection operation denied: Server is in read-only mode')
+        raise ValueError('Operation not permitted: Server is configured in read-only mode')
+
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+
+        # Check if collection already exists
+        existing_collections = db.list_collection_names()
+        if collection in existing_collections:
+            return {
+                'success': False,
+                'message': f"Collection '{collection}' already exists in database '{database}'",
+            }
+
+        # Create the collection
+        db.create_collection(collection)
+
+        logger.info(f"Created collection '{collection}' in database '{database}'")
+        return {'success': True, 'message': f"Collection '{collection}' created successfully"}
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error creating collection: {str(e)}')
+        raise ValueError(f'Failed to create collection: {str(e)}')
+
+
+async def list_collections(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+) -> List[str]:
+    """List collections in a DocumentDB database.
+
+    This tool returns the names of all collections in a specified database.
+
+    Returns:
+        List[str]: List of collection names
+    """
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+        db = client[database]
+        collections = db.list_collection_names()
+        logger.info(f"Found {len(collections)} collections in database '{database}'")
+        return collections
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error listing collections: {str(e)}')
+        raise ValueError(f'Failed to list collections: {str(e)}')
+
+
+async def drop_collection(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection to drop')],
+) -> Dict[str, Any]:
+    """Drop a collection from a DocumentDB database.
+
+    This tool completely removes a collection and all its documents from the specified database.
+    This operation cannot be undone, so use it with caution.
+
+    Returns:
+        Dict[str, Any]: Status of the drop operation
+    """
+    # Check if server is in read-only mode
+    if serverConfig.read_only_mode:
+        logger.warning('Drop collection operation denied: Server is in read-only mode')
+        raise ValueError('Operation not permitted: Server is configured in read-only mode')
+
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+
+        # Check if collection exists
+        existing_collections = db.list_collection_names()
+        if collection not in existing_collections:
+            return {
+                'success': False,
+                'message': f"Collection '{collection}' does not exist in database '{database}'",
+            }
+
+        # Drop the collection
+        db.drop_collection(collection)
+
+        logger.info(f"Dropped collection '{collection}' from database '{database}'")
+        return {'success': True, 'message': f"Collection '{collection}' dropped successfully"}
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error dropping collection: {str(e)}')
+        raise ValueError(f'Failed to drop collection: {str(e)}')

awslabs/documentdb_mcp_server/query_tools.py

@@ -0,0 +1,121 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Query tools for DocumentDB MCP Server."""
+
+from awslabs.documentdb_mcp_server.connection_tools import DocumentDBConnection
+from loguru import logger
+from pydantic import Field
+from typing import Annotated, Any, Dict, List, Optional
+
+
+async def find(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection')],
+    query: Annotated[
+        Dict[str, Any], Field(description='Query filter (e.g., {"name": "example"})')
+    ],
+    projection: Annotated[
+        Optional[Dict[str, Any]],
+        Field(description='Fields to include/exclude (e.g., {"_id": 0, "name": 1})'),
+    ] = None,
+    limit: Annotated[
+        int, Field(description='Maximum number of documents to return (default: 10)')
+    ] = 10,
+) -> List[Dict[str, Any]]:
+    """Run a query against a DocumentDB collection.
+
+    This tool queries documents from a specified collection based on a filter.
+
+    Returns:
+        List[Dict[str, Any]]: List of matching documents
+    """
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+        coll = db[collection]
+
+        result = list(coll.find(query, projection).limit(limit))
+
+        # Convert ObjectId to string for JSON serialization
+        for doc in result:
+            if '_id' in doc and not isinstance(doc['_id'], str):
+                doc['_id'] = str(doc['_id'])
+
+        logger.info(f'Query returned {len(result)} documents')
+        return result
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error querying DocumentDB: {str(e)}')
+        raise ValueError(f'Failed to query DocumentDB: {str(e)}')
+
+
+async def aggregate(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection')],
+    pipeline: Annotated[
+        List[Dict[str, Any]], Field(description='DocumentDB aggregation pipeline')
+    ],
+    limit: Annotated[
+        int, Field(description='Maximum number of documents to return (default: 10)')
+    ] = 10,
+) -> List[Dict[str, Any]]:
+    """Run an aggregation pipeline against a DocumentDB collection.
+
+    This tool executes a DocumentDB aggregation pipeline on a specified collection.
+
+    Returns:
+        List[Dict[str, Any]]: List of aggregation results
+    """
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+        coll = db[collection]
+
+        # Add limit stage if not already in pipeline
+        if limit > 0 and not any('$limit' in stage for stage in pipeline):
+            pipeline.append({'$limit': limit})
+
+        result = list(coll.aggregate(pipeline))
+
+        # Convert ObjectId to string for JSON serialization
+        for doc in result:
+            if '_id' in doc and not isinstance(doc['_id'], str):
+                doc['_id'] = str(doc['_id'])
+
+        logger.info(f'Aggregation returned {len(result)} results')
+        return result
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error running aggregation in DocumentDB: {str(e)}')
+        raise ValueError(f'Failed to run aggregation: {str(e)}')

awslabs/documentdb_mcp_server/server.py

@@ -0,0 +1,159 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""AWS Labs DocumentDB MCP Server implementation for querying AWS DocumentDB."""
+
+import argparse
+from awslabs.documentdb_mcp_server.analytic_tools import (
+    analyze_schema,
+    count_documents,
+    explain_operation,
+    get_collection_stats,
+    get_database_stats,
+)
+from awslabs.documentdb_mcp_server.config import serverConfig
+from awslabs.documentdb_mcp_server.connection_tools import (
+    DocumentDBConnection,
+    connect,
+    disconnect,
+)
+from awslabs.documentdb_mcp_server.db_management_tools import (
+    create_collection,
+    drop_collection,
+    list_collections,
+    list_databases,
+)
+from awslabs.documentdb_mcp_server.query_tools import aggregate, find
+from awslabs.documentdb_mcp_server.write_tools import delete, insert, update
+from loguru import logger
+from mcp.server.fastmcp import FastMCP
+
+
+# Create the FastMCP server
+mcp = FastMCP(
+    'awslabs.documentdb-mcp-server',
+    instructions="""DocumentDB MCP Server provides tools to connect to and query AWS DocumentDB databases.
+
+    Usage pattern:
+    1. First use the `connect` tool to establish a connection and get a connection_id
+    2. Use the connection_id with other tools to perform operations
+    3. When finished, use the `disconnect` tool to release resources
+
+    Server Configuration:
+    - The server can be configured in read-only mode, which blocks write operations
+      while still allowing read operations.""",
+    dependencies=[
+        'pydantic',
+        'loguru',
+        'pymongo',
+    ],
+)
+
+
+# Register all tools
+
+# Connection tools
+mcp.tool(name='connect')(connect)
+mcp.tool(name='disconnect')(disconnect)
+
+# Query tools
+mcp.tool(name='find')(find)
+mcp.tool(name='aggregate')(aggregate)
+
+# Write tools
+mcp.tool(name='insert')(insert)
+mcp.tool(name='update')(update)
+mcp.tool(name='delete')(delete)
+
+# Database management tools
+mcp.tool(name='listDatabases')(list_databases)
+mcp.tool(name='createCollection')(create_collection)
+mcp.tool(name='listCollections')(list_collections)
+mcp.tool(name='dropCollection')(drop_collection)
+
+# Analytic tools
+mcp.tool(name='countDocuments')(count_documents)
+mcp.tool(name='getDatabaseStats')(get_database_stats)
+mcp.tool(name='getCollectionStats')(get_collection_stats)
+mcp.tool(name='analyzeSchema')(analyze_schema)
+mcp.tool(name='explainOperation')(explain_operation)
+
+
+def main():
+    """Run the MCP server with CLI argument support."""
+    parser = argparse.ArgumentParser(
+        description='An AWS Labs Model Context Protocol (MCP) server for DocumentDB'
+    )
+    parser.add_argument('--sse', action='store_true', help='Use SSE transport')
+    parser.add_argument('--port', type=int, default=8888, help='Port to run the server on')
+    parser.add_argument('--host', type=str, default='127.0.0.1', help='Host to bind the server to')
+    parser.add_argument(
+        '--log-level',
+        type=str,
+        default='INFO',
+        choices=['TRACE', 'DEBUG', 'INFO', 'SUCCESS', 'WARNING', 'ERROR', 'CRITICAL'],
+        help='Set the logging level',
+    )
+    parser.add_argument(
+        '--connection-timeout',
+        type=int,
+        default=30,
+        help='Idle connection timeout in minutes (default: 30)',
+    )
+    parser.add_argument(
+        '--allow-write',
+        action='store_true',
+        help='Allow write operations (insert, update, delete). By default, the server runs in read-only mode.',
+    )
+
+    args = parser.parse_args()
+
+    # Configure logging
+    logger.remove()
+    logger.add(
+        lambda msg: print(msg),
+        level=args.log_level,
+        format='<green>{time:YYYY-MM-DD HH:mm:ss.SSS}</green> | <level>{level: <8}</level> | <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan> - <level>{message}</level>',
+    )
+
+    logger.info(f'Starting DocumentDB MCP Server on {args.host}:{args.port}')
+    logger.info(f'Log level: {args.log_level}')
+
+    # Set connection timeout
+    DocumentDBConnection._idle_timeout = args.connection_timeout
+    logger.info(f'Idle connection timeout: {args.connection_timeout} minutes')
+
+    # Configure read-only mode
+    serverConfig.read_only_mode = not args.allow_write
+    if serverConfig.read_only_mode:
+        logger.warning('Server is running in READ-ONLY mode. Write operations will be blocked.')
+    else:
+        logger.info('Server is running with WRITE operations ENABLED. Database can be modified.')
+
+    try:
+        # Run server with appropriate transport
+        if args.sse:
+            mcp.settings.port = args.port
+            mcp.settings.host = args.host
+            mcp.run(transport='sse')
+        else:
+            mcp.settings.port = args.port
+            mcp.settings.host = args.host
+            mcp.run()
+    except Exception as e:
+        logger.critical(f'Failed to start server: {str(e)}')
+    finally:
+        # Close all DB connections
+        DocumentDBConnection.close_all_connections()
+
+
+if __name__ == '__main__':
+    main()

awslabs/documentdb_mcp_server/write_tools.py

@@ -0,0 +1,202 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+# with the License. A copy of the License is located at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# or in the 'license' file accompanying this file. This file is distributed on an 'AS IS' BASIS, WITHOUT WARRANTIES
+# OR CONDITIONS OF ANY KIND, express or implied. See the License for the specific language governing permissions
+# and limitations under the License.
+
+"""Write tools for DocumentDB MCP Server."""
+
+from awslabs.documentdb_mcp_server.config import serverConfig
+from awslabs.documentdb_mcp_server.connection_tools import DocumentDBConnection
+from loguru import logger
+from pydantic import Field
+from typing import Annotated, Any, Dict, List, Union
+
+
+async def insert(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection')],
+    documents: Annotated[
+        Union[Dict[str, Any], List[Dict[str, Any]]],
+        Field(description='Document or list of documents to insert'),
+    ],
+) -> Dict[str, Any]:
+    """Insert one or more documents into a DocumentDB collection.
+
+    This tool inserts new documents into a specified collection.
+
+    Returns:
+        Dict[str, Any]: Insert operation results including document IDs
+    """
+    # Check if server is in read-only mode
+    if serverConfig.read_only_mode:
+        logger.warning('Insert operation denied: Server is in read-only mode')
+        raise ValueError('Operation not permitted: Server is configured in read-only mode')
+
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+        coll = db[collection]
+
+        # Handle single document or multiple documents
+        if isinstance(documents, dict):
+            result = coll.insert_one(documents)
+            inserted_ids = [str(result.inserted_id)]
+            count = 1
+        else:
+            result = coll.insert_many(documents)
+            inserted_ids = [str(id) for id in result.inserted_ids]
+            count = len(inserted_ids)
+
+        logger.info(f'Inserted {count} documents')
+        return {'success': True, 'inserted_count': count, 'inserted_ids': inserted_ids}
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error inserting into DocumentDB: {str(e)}')
+        raise ValueError(f'Failed to insert documents: {str(e)}')
+
+
+async def update(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection')],
+    filter: Annotated[Dict[str, Any], Field(description='Filter to select documents to update')],
+    update: Annotated[
+        Dict[str, Any],
+        Field(
+            description='Update operations to apply. It should either include DocumentDB operators like $set, or an entire document if the entire document needs to be replaced.'
+        ),
+    ],
+    upsert: Annotated[
+        bool,
+        Field(
+            description='Whether to create a new document if no match is found (default: False)'
+        ),
+    ] = False,
+    many: Annotated[
+        bool, Field(description='Whether to update multiple documents (default: False)')
+    ] = False,
+) -> Dict[str, Any]:
+    """Update documents in a DocumentDB collection.
+
+    This tool updates existing documents that match a specified filter.
+
+    Returns:
+        Dict[str, Any]: Update operation results
+    """
+    # Check if server is in read-only mode
+    if serverConfig.read_only_mode:
+        logger.warning('Update operation denied: Server is in read-only mode')
+        raise ValueError('Operation not permitted: Server is configured in read-only mode')
+
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+        coll = db[collection]
+
+        # If the update doesn't have any operators, then it's a replace
+        if not any(key.startswith('$') for key in update.keys()):
+            result = coll.replace_one(filter, update, upsert=upsert)
+            matched = result.matched_count
+            modified = result.modified_count
+        # If the update needs to update multiple documents
+        elif many:
+            result = coll.update_many(filter, update, upsert=upsert)
+            matched = result.matched_count
+            modified = result.modified_count
+        # Else only a single document needs to be updated
+        else:
+            result = coll.update_one(filter, update, upsert=upsert)
+            matched = result.matched_count
+            modified = result.modified_count
+
+        upserted_id = str(result.upserted_id) if result.upserted_id else None
+
+        logger.info(f'Updated {modified} documents (matched {matched})')
+        return {
+            'success': True,
+            'matched_count': matched,
+            'modified_count': modified,
+            'upserted_id': upserted_id,
+        }
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error updating DocumentDB: {str(e)}')
+        raise ValueError(f'Failed to update documents: {str(e)}')
+
+
+async def delete(
+    connection_id: Annotated[
+        str, Field(description='The connection ID returned by the connect tool')
+    ],
+    database: Annotated[str, Field(description='Name of the database')],
+    collection: Annotated[str, Field(description='Name of the collection')],
+    filter: Annotated[Dict[str, Any], Field(description='Filter to select documents to delete')],
+    many: Annotated[
+        bool, Field(description='Whether to delete multiple documents (default: False)')
+    ] = False,
+) -> Dict[str, Any]:
+    """Delete documents from a DocumentDB collection.
+
+    This tool deletes documents that match a specified filter.
+
+    Returns:
+        Dict[str, Any]: Delete operation results
+    """
+    # Check if server is in read-only mode
+    if serverConfig.read_only_mode:
+        logger.warning('Delete operation denied: Server is in read-only mode')
+        raise ValueError('Operation not permitted: Server is configured in read-only mode')
+
+    try:
+        # Get connection
+        if connection_id not in DocumentDBConnection._connections:
+            raise ValueError(f'Connection ID {connection_id} not found. You must connect first.')
+
+        connection_info = DocumentDBConnection._connections[connection_id]
+        client = connection_info.client
+
+        db = client[database]
+        coll = db[collection]
+
+        if many:
+            result = coll.delete_many(filter)
+            deleted = result.deleted_count
+        else:
+            result = coll.delete_one(filter)
+            deleted = result.deleted_count
+
+        logger.info(f'Deleted {deleted} documents')
+        return {'success': True, 'deleted_count': deleted}
+    except ValueError as e:
+        logger.error(f'Connection error: {str(e)}')
+        raise ValueError(str(e))
+    except Exception as e:
+        logger.error(f'Error deleting from DocumentDB: {str(e)}')
+        raise ValueError(f'Failed to delete documents: {str(e)}')