diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index 4e2d16a33e..4f326fbab5 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -65,6 +65,12 @@ # ServiceLabel: %tools-Advisor # ServiceOwners: @ankiga-MSFT +# PRLabel: %tools-DocumentDb +/tools/Azure.Mcp.Tools.DocumentDb/ @microsoft/azure-mcp + +# ServiceLabel: %tools-DocumentDb +# ServiceOwners: @microsoft/azure-mcp + # PRLabel: %tools-Aks /tools/Azure.Mcp.Tools.Aks/ @jongio @anuchandy @microsoft/azure-mcp diff --git a/Directory.Packages.props b/Directory.Packages.props index 8326ef89b8..aa7fbef669 100644 --- a/Directory.Packages.props +++ b/Directory.Packages.props @@ -82,6 +82,7 @@ + diff --git a/servers/Azure.Mcp.Server/Azure.Mcp.Server.slnx b/servers/Azure.Mcp.Server/Azure.Mcp.Server.slnx index feec36b074..f4312ba1ed 100644 --- a/servers/Azure.Mcp.Server/Azure.Mcp.Server.slnx +++ b/servers/Azure.Mcp.Server/Azure.Mcp.Server.slnx @@ -151,6 +151,13 @@ + + + + + + + diff --git a/servers/Azure.Mcp.Server/README.md b/servers/Azure.Mcp.Server/README.md index ee7d34214e..70e553234f 100644 --- a/servers/Azure.Mcp.Server/README.md +++ b/servers/Azure.Mcp.Server/README.md @@ -851,6 +851,18 @@ Check out the remote hosting [azd templates](https://github.com/microsoft/mcp/bl * "List my Advisor recommendations" +### 🗄️ Azure Cosmos DB for MongoDB (vCore) + +* "Connect to my DocumentDB instance" +* "List all databases in DocumentDB" +* "Show me collections in database 'mydb'" +* "Find documents in collection 'users' where status is active" +* "Count documents in collection 'orders'" +* "Insert a document into collection 'products'" +* "Update documents in collection 'inventory'" +* "Create an index on collection 'customers'" +* "Get statistics for database 'analytics'" + ### 🔎 Azure AI Search * "What indexes do I have in my Azure AI Search service 'mysvc'?" @@ -1042,6 +1054,7 @@ The Azure MCP Server provides tools for interacting with **41+ Azure service are - 📦 **Azure Container Apps** - Container hosting - 📦 **Azure Container Registry (ACR)** - Container registry management - 📊 **Azure Cosmos DB** - NoSQL database operations +- 🗄️ **Azure Cosmos DB for MongoDB (vCore)** - MongoDB vCore database operations - 🧮 **Azure Data Explorer** - Analytics queries and KQL - 🐬 **Azure Database for MySQL** - MySQL database management - 🐘 **Azure Database for PostgreSQL** - PostgreSQL database management diff --git a/servers/Azure.Mcp.Server/docs/azmcp-commands.md b/servers/Azure.Mcp.Server/docs/azmcp-commands.md index db9267d976..e90d27b6b5 100644 --- a/servers/Azure.Mcp.Server/docs/azmcp-commands.md +++ b/servers/Azure.Mcp.Server/docs/azmcp-commands.md @@ -257,6 +257,51 @@ azmcp server info azmcp advisor recommendations list --subscription ``` +### Azure Cosmos DB for MongoDB (vCore) Operations + +```bash +# Connection Management +azmcp documentdb connect --connection-string +azmcp documentdb disconnect +azmcp documentdb connection status + +# Database Operations +azmcp documentdb database list +azmcp documentdb database stats --db-name +azmcp documentdb database info --db-name +azmcp documentdb database drop --db-name + +# Collection Operations +azmcp documentdb collection stats --db-name --collection-name +azmcp documentdb collection rename --db-name --collection-name --new-collection-name +azmcp documentdb collection drop --db-name --collection-name +azmcp documentdb collection sample --db-name --collection-name --sample-size 10 + +# Document Operations +azmcp documentdb find documents --db-name --collection-name --query '{"status": "active"}' +azmcp documentdb count documents --db-name --collection-name --query '{}' +azmcp documentdb insert document --db-name --collection-name --document '{"name": "example"}' +azmcp documentdb insert many --db-name --collection-name --documents '[{"name": "doc1"}, {"name": "doc2"}]' +azmcp documentdb update document --db-name --collection-name --filter '{"_id": "123"}' --update '{"$set": {"status": "updated"}}' +azmcp documentdb update many --db-name --collection-name --filter '{"status": "pending"}' --update '{"$set": {"status": "processed"}}' +azmcp documentdb delete document --db-name --collection-name --filter '{"_id": "123"}' +azmcp documentdb delete many --db-name --collection-name --filter '{"status": "archived"}' +azmcp documentdb aggregate --db-name --collection-name --pipeline '[{"$group": {"_id": "$status", "count": {"$sum": 1}}}]' +azmcp documentdb find and modify --db-name --collection-name --query '{"_id": "123"}' --update '{"$set": {"modified": true}}' + +# Query Explanation +azmcp documentdb explain find --db-name --collection-name --query '{}' +azmcp documentdb explain count --db-name --collection-name --query '{}' +azmcp documentdb explain aggregate --db-name --collection-name --pipeline '[{"$match": {"status": "active"}}]' + +# Index Operations +azmcp documentdb index create --db-name --collection-name --keys '{"fieldName": 1}' +azmcp documentdb index list --db-name --collection-name +azmcp documentdb index drop --db-name --collection-name --index-name +azmcp documentdb index stats --db-name --collection-name +azmcp documentdb current ops +``` + ### Azure AI Search Operations ```bash diff --git a/servers/Azure.Mcp.Server/docs/documentdb-mcp-commands.md b/servers/Azure.Mcp.Server/docs/documentdb-mcp-commands.md new file mode 100644 index 0000000000..46c3eb0b14 --- /dev/null +++ b/servers/Azure.Mcp.Server/docs/documentdb-mcp-commands.md @@ -0,0 +1,821 @@ +# DocumentDB MCP Server Commands Documentation + +## Table of Contents + +1. [Connection Tools](#1-connection-tools) +2. [Database Tools](#2-database-tools) +3. [Collection Tools](#3-collection-tools) +4. [Document Tools](#4-document-tools) +5. [Index Tools](#5-index-tools) +6. [Workflow Tools](#6-workflow-tools) + +--- + +## 1. Connection Tools + +### 1.1 connect_mongodb + +**Name**: Connect to MongoDB +**Description**: Connect to a MongoDB instance with a connection string + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `connection_string` | `string` | Yes | MongoDB connection string (e.g., `mongodb://localhost:27017`) | +| `test_connection` | `boolean \| string` | No | Test the connection after connecting (default: `true`) | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted connection result with status information" +} +``` + +--- + +### 1.2 disconnect_mongodb + +**Name**: Disconnect from MongoDB +**Description**: Disconnect from the current MongoDB instance + +**Parameters**: None + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted disconnection result" +} +``` + +--- + +### 1.3 get_connection_status + +**Name**: Get Connection Status +**Description**: Get the current MongoDB connection status and details + +**Parameters**: None + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted connection status information" +} +``` + +--- + +## 2. Database Tools + +### 2.1 list_databases + +**Name**: List Databases +**Description**: List all databases in the DocumentDB instance + +**Parameters**: None + +**Returns**: + +```json +{ + "type": "text", + "text": "[\"admin\", \"local\", \"mydb\"]" +} +``` + +Return format: JSON array of database names + +--- + +### 2.2 db_stats + +**Name**: Database Statistics +**Description**: Get detailed statistics about a database's size and storage usage + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted database statistics (includes dataSize, storageSize, indexSize, etc.)" +} +``` + +--- + +### 2.3 get_db_info + +**Name**: Get Database Info +**Description**: Get database information including all collections and their document counts + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "database_name": "string", + "collections": [ + { + "name": "string", + "count": "number", + "error": "string (optional)" + } + ] + } +} +``` + +--- + +### 2.4 drop_database + +**Name**: Drop Database +**Description**: Drop a database and all its collections + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database to drop | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "success": "boolean", + "message": "string", + "data": "object" + } +} +``` + +--- + +## 3. Collection Tools + +### 3.1 collection_stats + +**Name**: Collection Statistics +**Description**: Get detailed statistics about a collection's size and storage usage + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted collection statistics" +} +``` + +--- + +### 3.2 rename_collection + +**Name**: Rename Collection +**Description**: Rename a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection to rename | +| `new_collection_name` | `string` | Yes | New name for the collection | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "message": "Collection renamed successfully" + } +} +``` + +--- + +### 3.3 drop_collection + +**Name**: Drop Collection +**Description**: Drop a collection from a database + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection to drop | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "message": "Collection dropped successfully" + } +} +``` + +--- + +### 3.4 sample_documents + +**Name**: Sample Documents +**Description**: Retrieve sample documents from a specific collection. Useful for understanding data schema and query generation + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `sample_size` | `number \| string` | No | Number of documents to sample (default: `10`) | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON array of sample documents" +} +``` + +--- + +## 4. Document Tools + +### 4.1 find_documents + +**Name**: Find Documents +**Description**: Find documents in a collection. Supports consolidated "options" object (limit, skip, sort, projection) + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection to query | +| `query` | `object \| string` | No | Query filter in MongoDB style (default: `{}`) | +| `options` | `object \| string` | No | Query options: `limit` (default 100), `skip` (default 0), `sort`, `projection` | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "documents": "Array", + "total_count": "number", + "returned_count": "number", + "has_more": "boolean", + "query": "object", + "applied_options": { + "limit": "number", + "skip": "number", + "sort": "object (optional)", + "projection": "object (optional)" + } + } +} +``` + +--- + +### 4.2 count_documents + +**Name**: Count Documents +**Description**: Count documents in a collection matching a query + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection to query | +| `query` | `object \| string` | No | Query filter in MongoDB style (default: `{}`) | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "count": "number", + "query": "object" + } +} +``` + +--- + +### 4.3 insert_document + +**Name**: Insert Document +**Description**: Insert a single document into a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `document` | `object` | Yes | Document to insert | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "inserted_id": "string", + "acknowledged": "boolean", + "inserted_count": 1 + } +} +``` + +--- + +### 4.4 insert_many + +**Name**: Insert Many Documents +**Description**: Insert multiple documents into a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `documents` | `Array \| string` | Yes | List of documents to insert | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "inserted_ids": "Array", + "acknowledged": "boolean", + "inserted_count": "number" + } +} +``` + +--- + +### 4.5 update_document + +**Name**: Update Single Document +**Description**: Update a document in a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `filter` | `object \| string` | Yes | Query filter to find the document | +| `update` | `object \| string` | Yes | Update operations ($set, $inc, etc.) | +| `upsert` | `boolean \| string` | No | Create document if it doesn't exist (default: `false`) | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "matched_count": "number", + "modified_count": "number", + "upserted_id": "string | null", + "acknowledged": "boolean" + } +} +``` + +--- + +### 4.6 update_many + +**Name**: Update Many Documents +**Description**: Update multiple documents in a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `filter` | `object \| string` | Yes | Query filter to find the documents | +| `update` | `object \| string` | Yes | Update operations ($set, $inc, etc.) | +| `upsert` | `boolean \| string` | No | Create document if it doesn't exist (default: `false`) | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "matched_count": "number", + "modified_count": "number", + "upserted_id": "string | null", + "acknowledged": "boolean" + } +} +``` + +--- + +### 4.7 delete_document + +**Name**: Delete Document +**Description**: Delete a document from a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `filter` | `object \| string` | Yes | Query filter to find the document | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "deleted_count": "number", + "acknowledged": "boolean" + } +} +``` + +--- + +### 4.8 delete_many + +**Name**: Delete Many Documents +**Description**: Delete multiple documents from a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `filter` | `object \| string` | Yes | Query filter to find the documents | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "deleted_count": "number", + "acknowledged": "boolean" + } +} +``` + +--- + +### 4.9 aggregate + +**Name**: Aggregate Pipeline +**Description**: Run an aggregation pipeline on a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `pipeline` | `Array \| string` | Yes | List of aggregation stages | +| `allow_disk_use` | `boolean \| string` | No | Allow pipeline stages to write to disk (default: `false`) | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "results": "Array", + "total_count": "number" + } +} +``` + +--- + +### 4.10 find_and_modify + +**Name**: Find And Modify Document +**Description**: Find one document by filter and apply update; returns the document BEFORE modification (or null if it doesn't exist) + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection to query | +| `query` | `object \| string` | Yes | Query filter in MongoDB style | +| `update` | `object \| string` | Yes | Update operations ($set, $inc, etc.) | +| `upsert` | `boolean \| string` | No | Create document if it does not exist (default: `false`) | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "matched": "boolean", + "upsertedId": "string | undefined", + "original_document": "object | null", + "query": "object", + "update": "object", + "upsert": "boolean" + } +} +``` + +--- + +### 4.11 explain_find_query + +**Name**: Explain Find Query +**Description**: Explain the execution plan with execution stats for a find query using consolidated options (sort, projection, limit, skip) + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `query` | `object \| string` | No | Query filter in MongoDB style (default: `{}`) | +| `options` | `object \| string` | No | Query options: `sort`, `projection`, `limit`, `skip` | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "options_applied": { + "sort": "object (optional)", + "projection": "object (optional)", + "limit": "number (optional)", + "skip": "number (optional)" + }, + "explain": "object (MongoDB explain output)" + } +} +``` + +--- + +### 4.12 explain_count_query + +**Name**: Explain Count Query +**Description**: Explain the execution plan with execution stats for count query on a given collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `query` | `object \| string` | No | Query filter in MongoDB style (default: `{}`) | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted explain output" +} +``` + +--- + +### 4.13 explain_aggregate_query + +**Name**: Explain Aggregate Query +**Description**: Explain the execution plan with execution stats for an aggregation query on a given collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `pipeline` | `Array \| string` | Yes | List of aggregation stages | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted explain output" +} +``` + +--- + +## 5. Index Tools + +### 5.1 create_index + +**Name**: Create Index +**Description**: Create an index on a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `keys` | `object \| string` | Yes | Dictionary defining the index (e.g., `{"field": 1}` for ascending) | +| `options` | `object \| string` | No | Index options (e.g., `{unique: true, name: 'idx'}`) (default: `{}`) | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "index_name": "string", + "keys": "object", + "options": "object" + } +} +``` + +--- + +### 5.2 list_indexes + +**Name**: List Indexes +**Description**: List all indexes on a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "indexes": "Array", + "count": "number" + } +} +``` + +--- + +### 5.3 drop_index + +**Name**: Drop Index +**Description**: Drop an index from a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | +| `index_name` | `string` | Yes | Name of the index to drop | + +**Returns**: + +```json +{ + "type": "text", + "text": { + "success": "boolean", + "message": "string", + "data": "object" + } +} +``` + +--- + +### 5.4 index_stats + +**Name**: Index Statistics +**Description**: Get statistics for indexes on a collection + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `db_name` | `string` | Yes | Name of the database | +| `collection_name` | `string` | Yes | Name of the collection | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON array of index statistics" +} +``` + +--- + +### 5.5 current_ops + +**Name**: Current Operations +**Description**: Get information about current MongoDB operations + +**Parameters**: + +| Parameter | Type | Required | Description | +|-----------|------|----------|-------------| +| `ops` | `object \| string \| null` | No | Optional filter to narrow down the operations returned | + +**Returns**: + +```json +{ + "type": "text", + "text": "JSON formatted current operations information" +} +``` + +--- + + +## Command Summary + +| Category | Command | Description | +|----------|---------|-------------| +| **Connection** | `connect_mongodb` | Connect to MongoDB | +| **Connection** | `disconnect_mongodb` | Disconnect | +| **Connection** | `get_connection_status` | Get connection status | +| **Database** | `list_databases` | List databases | +| **Database** | `db_stats` | Database statistics | +| **Database** | `get_db_info` | Database info | +| **Database** | `drop_database` | Drop database | +| **Collection** | `collection_stats` | Collection statistics | +| **Collection** | `rename_collection` | Rename collection | +| **Collection** | `drop_collection` | Drop collection | +| **Collection** | `sample_documents` | Sample documents | +| **Document** | `find_documents` | Find documents | +| **Document** | `count_documents` | Count documents | +| **Document** | `insert_document` | Insert single document | +| **Document** | `insert_many` | Insert multiple documents | +| **Document** | `update_document` | Update single document | +| **Document** | `update_many` | Update multiple documents | +| **Document** | `delete_document` | Delete single document | +| **Document** | `delete_many` | Delete multiple documents | +| **Document** | `aggregate` | Aggregation pipeline | +| **Document** | `find_and_modify` | Find and modify | +| **Document** | `explain_find_query` | Explain find query | +| **Document** | `explain_count_query` | Explain count query | +| **Document** | `explain_aggregate_query` | Explain aggregate query | +| **Index** | `create_index` | Create index | +| **Index** | `list_indexes` | List indexes | +| **Index** | `drop_index` | Drop index | +| **Index** | `index_stats` | Index statistics | +| **Index** | `current_ops` | Current operations | + +**Total**: 29 Commands \ No newline at end of file diff --git a/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md b/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md index 79e841aec9..ef6655f84c 100644 --- a/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md +++ b/servers/Azure.Mcp.Server/docs/e2eTestPrompts.md @@ -10,6 +10,70 @@ This file contains prompts used for end-to-end testing to ensure each tool is in | advisor_recommendations_list | Show me Advisor recommendations in the subscription | | advisor_recommendations_list | List all Advisor recommendations in the subscription | +## Azure Cosmos DB for MongoDB (vCore) + +| Tool Name | Test Prompt | +|:----------|:----------| +| documentdb_connect | Connect to DocumentDB with connection string | +| documentdb_connect | Connect to my DocumentDB instance using | +| documentdb_disconnect | Disconnect from DocumentDB | +| documentdb_disconnect | Close the DocumentDB connection | +| documentdb_connection_status | Show me the DocumentDB connection status | +| documentdb_connection_status | Is DocumentDB connected? | +| documentdb_database_list | List all databases in DocumentDB | +| documentdb_database_list | Show me all DocumentDB databases | +| documentdb_database_stats | Get statistics for database | +| documentdb_database_stats | Show me stats for DocumentDB database | +| documentdb_database_info | Get information about database | +| documentdb_database_info | Show me details of database | +| documentdb_database_drop | Drop database | +| documentdb_database_drop | Delete the database from DocumentDB | +| documentdb_collection_stats | Get statistics for collection in database | +| documentdb_collection_stats | Show me stats for collection in | +| documentdb_collection_rename | Rename collection to in database | +| documentdb_collection_rename | Rename collection to in DocumentDB database | +| documentdb_collection_drop | Drop collection from database | +| documentdb_collection_drop | Delete collection from database | +| documentdb_collection_sample | Sample 10 documents from collection in database | +| documentdb_collection_sample | Show me sample documents from collection in database | +| documentdb_find_documents | Find all documents in collection in database | +| documentdb_find_documents | Find documents where status equals active in collection in database | +| documentdb_find_documents | Query collection in database for documents | +| documentdb_count_documents | Count all documents in collection in database | +| documentdb_count_documents | How many documents are in collection in database | +| documentdb_insert_document | Insert a document into collection in database | +| documentdb_insert_document | Add a new document to collection in database | +| documentdb_insert_many | Insert multiple documents into collection in database | +| documentdb_insert_many | Bulk insert documents into collection in database | +| documentdb_update_document | Update a document in collection in database | +| documentdb_update_document | Update documents where id equals in collection in database | +| documentdb_update_many | Update multiple documents in collection in database | +| documentdb_update_many | Bulk update documents in collection in database | +| documentdb_delete_document | Delete a document from collection in database | +| documentdb_delete_document | Remove document where id equals from collection in database | +| documentdb_delete_many | Delete multiple documents from collection in database | +| documentdb_delete_many | Remove all documents where status equals inactive from collection in database | +| documentdb_aggregate | Run an aggregation pipeline on collection in database | +| documentdb_aggregate | Aggregate documents in collection in database | +| documentdb_find_and_modify | Find and modify a document in collection in database | +| documentdb_find_and_modify | Update and return a document from collection in database | +| documentdb_explain_find | Explain the query plan for finding documents in collection in database | +| documentdb_explain_find | Show me the execution plan for a find query on collection in database | +| documentdb_explain_count | Explain the query plan for counting documents in collection in database | +| documentdb_explain_count | Show me how the count query executes on collection in database | +| documentdb_explain_aggregate | Explain the aggregation pipeline on collection in database | +| documentdb_explain_aggregate | Show me the execution plan for aggregation on collection in database | +| documentdb_index_create | Create an index on collection in database | +| documentdb_index_create | Add a new index to collection in database | +| documentdb_index_list | List all indexes on collection in database | +| documentdb_index_list | Show me indexes for collection in database | +| documentdb_index_drop | Drop index from collection in database | +| documentdb_index_drop | Delete index from collection in database | +| documentdb_index_stats | Get index statistics for collection in database | +| documentdb_index_stats | Show me index stats for collection in database | +| documentdb_current_ops | Show current operations running in DocumentDB | +| documentdb_current_ops | What operations are currently running in DocumentDB | + ## Azure AI Search | Tool Name | Test Prompt | diff --git a/servers/Azure.Mcp.Server/src/Program.cs b/servers/Azure.Mcp.Server/src/Program.cs index c50d0fc156..00f381f364 100644 --- a/servers/Azure.Mcp.Server/src/Program.cs +++ b/servers/Azure.Mcp.Server/src/Program.cs @@ -98,6 +98,7 @@ private static IAreaSetup[] RegisterAreas() new Azure.Mcp.Tools.AzureMigrate.AzureMigrateSetup(), new Azure.Mcp.Tools.AzureTerraformBestPractices.AzureTerraformBestPracticesSetup(), new Azure.Mcp.Tools.Deploy.DeploySetup(), + new Azure.Mcp.Tools.DocumentDb.DocumentDbSetup(), new Azure.Mcp.Tools.EventGrid.EventGridSetup(), new Azure.Mcp.Tools.Acr.AcrSetup(), new Azure.Mcp.Tools.Advisor.AdvisorSetup(), diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/AssemblyInfo.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/AssemblyInfo.cs new file mode 100644 index 0000000000..51d7ca2f6e --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/AssemblyInfo.cs @@ -0,0 +1,7 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Runtime.CompilerServices; + +[assembly: InternalsVisibleTo("Azure.Mcp.Tools.DocumentDb.UnitTests")] +[assembly: InternalsVisibleTo("Azure.Mcp.Tools.DocumentDb.LiveTests")] diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Azure.Mcp.Tools.DocumentDb.csproj b/tools/Azure.Mcp.Tools.DocumentDb/src/Azure.Mcp.Tools.DocumentDb.csproj new file mode 100644 index 0000000000..0edac7ca58 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Azure.Mcp.Tools.DocumentDb.csproj @@ -0,0 +1,20 @@ + + + true + + + + + + + + + + + + + + + + + diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/BaseDocumentDbCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/BaseDocumentDbCommand.cs new file mode 100644 index 0000000000..2db38e3d98 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/BaseDocumentDbCommand.cs @@ -0,0 +1,17 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using System.Diagnostics.CodeAnalysis; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Tools.DocumentDb.Options; +using Microsoft.Extensions.Logging; + +namespace Azure.Mcp.Tools.DocumentDb.Commands; + +public abstract class BaseDocumentDbCommand< + [DynamicallyAccessedMembers(TrimAnnotations.CommandAnnotations)] TOptions>(ILogger> logger) + : GlobalCommand where TOptions : BaseDocumentDbOptions, new() +{ + protected readonly ILogger> _logger = logger; +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/CollectionStatsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/CollectionStatsCommand.cs new file mode 100644 index 0000000000..b1871dbb99 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/CollectionStatsCommand.cs @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Collection; + +public sealed class CollectionStatsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "b8c9d0e1-f2a3-4b8c-5d6e-7f8a9b0c1d2e"; + + public override string Name => "collection_stats"; + + public override string Description => "Get detailed statistics about a collection's size and storage usage"; + + public override string Title => "Collection Statistics"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + } + + protected override CollectionStatsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.GetCollectionStatsAsync(options.DbName!, options.CollectionName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbHelpers.SerializeBsonToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to get collection statistics for database: {DbName}, collection: {CollectionName}", options.DbName, options.CollectionName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/DropCollectionCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/DropCollectionCommand.cs new file mode 100644 index 0000000000..3177bddc65 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/DropCollectionCommand.cs @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Collection; + +public sealed class DropCollectionCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "d0e1f2a3-b4c5-4d0e-7f8a-9b0c1d2e3f4a"; + + public override string Name => "drop_collection"; + + public override string Description => "Drop a collection from a database"; + + public override string Title => "Drop Collection"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + } + + protected override DropCollectionOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.DropCollectionAsync(options.DbName!, options.CollectionName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to drop collection: {CollectionName} from database: {DbName}", options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/RenameCollectionCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/RenameCollectionCommand.cs new file mode 100644 index 0000000000..e14b3d3df3 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/RenameCollectionCommand.cs @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Collection; + +public sealed class RenameCollectionCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "c9d0e1f2-a3b4-4c9d-6e7f-8a9b0c1d2e3f"; + + public override string Name => "rename_collection"; + + public override string Description => "Rename a collection"; + + public override string Title => "Rename Collection"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.NewCollectionName); + } + + protected override RenameCollectionOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.NewCollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.NewCollectionName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.RenameCollectionAsync(options.DbName!, options.CollectionName!, options.NewCollectionName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to rename collection from {OldName} to {NewName} in database: {DbName}", options.CollectionName, options.NewCollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/SampleDocumentsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/SampleDocumentsCommand.cs new file mode 100644 index 0000000000..5d676efc2f --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Collection/SampleDocumentsCommand.cs @@ -0,0 +1,92 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Collection; + +public sealed class SampleDocumentsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "e1f2a3b4-c5d6-4e1f-8a9b-0c1d2e3f4a5b"; + + public override string Name => "sample_documents"; + + public override string Description => "Retrieve sample documents from a specific collection. Useful for understanding data schema and query generation"; + + public override string Title => "Sample Documents"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + + var dbNameOption = new Option("--db-name") + { + Description = "Database name", + Required = true + }; + command.Options.Add(dbNameOption); + + var collectionNameOption = new Option("--collection-name") + { + Description = "Collection name", + Required = true + }; + command.Options.Add(collectionNameOption); + + command.Options.Add(DocumentDbOptionDefinitions.SampleSize); + } + + protected override SampleDocumentsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.SampleSize = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.SampleSize.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.SampleDocumentsAsync(options.DbName!, options.CollectionName!, options.SampleSize, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result.Select(doc => DocumentDbHelpers.SerializeBsonToJson(doc)))); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to sample documents from collection: {CollectionName} in database: {DbName} with sample size: {SampleSize}", options.CollectionName, options.DbName, options.SampleSize); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/ConnectDocumentDbCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/ConnectDocumentDbCommand.cs new file mode 100644 index 0000000000..79c014dc7a --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/ConnectDocumentDbCommand.cs @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Commands; +using Microsoft.Mcp.Core.Models.Command; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Connection; + +public sealed class ConnectDocumentDbCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "a1b2c3d4-e5f6-4a1b-8c9d-0e1f2a3b4c5d"; + + public override string Name => "connect"; + + public override string Description => "Connect to an Azure Cosmos DB for MongoDB (vCore) instance with a connection string"; + + public override string Title => "Connect to DocumentDB"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.ConnectionString); + command.Options.Add(DocumentDbOptionDefinitions.TestConnection); + } + + protected override ConnectDocumentDbOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.ConnectionString = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.ConnectionString.Name); + options.TestConnection = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.TestConnection.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.ConnectAsync(options.ConnectionString!, options.TestConnection, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to connect to DocumentDB with connection string: {ConnectionString}", options.ConnectionString); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/DisconnectDocumentDbCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/DisconnectDocumentDbCommand.cs new file mode 100644 index 0000000000..c8a2a994ac --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/DisconnectDocumentDbCommand.cs @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Connection; + +public sealed class DisconnectDocumentDbCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "b2c3d4e5-f6a7-4b2c-9d0e-1f2a3b4c5d6e"; + + public override string Name => "disconnect"; + + public override string Description => "Disconnect from the current DocumentDB instance"; + + public override string Title => "Disconnect from DocumentDB"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = false + }; + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.DisconnectAsync(cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to disconnect from DocumentDB"); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/GetConnectionStatusCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/GetConnectionStatusCommand.cs new file mode 100644 index 0000000000..d6d8eb4529 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Connection/GetConnectionStatusCommand.cs @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Connection; + +public sealed class GetConnectionStatusCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "c3d4e5f6-a7b8-4c3d-0e1f-2a3b4c5d6e7f"; + + public override string Name => "get_connection_status"; + + public override string Description => "Get the current DocumentDB connection status and details"; + + public override string Title => "Get Connection Status"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = service.GetConnectionStatus(); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return await Task.FromResult(context.Response); + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to get DocumentDB connection status"); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/DbStatsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/DbStatsCommand.cs new file mode 100644 index 0000000000..708e59ed78 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/DbStatsCommand.cs @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Database; + +public sealed class DbStatsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "e5f6a7b8-c9d0-4e5f-2a3b-4c5d6e7f8a9b"; + + public override string Name => "db_stats"; + + public override string Description => "Get detailed statistics about a database's size and storage usage"; + + public override string Title => "Database Statistics"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + } + + protected override DbStatsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.GetDatabaseStatsAsync(options.DbName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbHelpers.SerializeBsonToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to get database statistics for database: {DbName}", options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/DropDatabaseCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/DropDatabaseCommand.cs new file mode 100644 index 0000000000..9729c0703f --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/DropDatabaseCommand.cs @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Database; + +public sealed class DropDatabaseCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "a7b8c9d0-e1f2-4a7b-4c5d-6e7f8a9b0c1d"; + + public override string Name => "drop_database"; + + public override string Description => "Drop a database and all its collections"; + + public override string Title => "Drop Database"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + } + + protected override DropDatabaseOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.DropDatabaseAsync(options.DbName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to drop database: {DbName}", options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/GetDbInfoCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/GetDbInfoCommand.cs new file mode 100644 index 0000000000..6a324e23d9 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/GetDbInfoCommand.cs @@ -0,0 +1,75 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Database; + +public sealed class GetDbInfoCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "f6a7b8c9-d0e1-4f6a-3b4c-5d6e7f8a9b0c"; + + public override string Name => "get_db_info"; + + public override string Description => "Get database information including all collections and their document counts"; + + public override string Title => "Get Database Info"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + } + + protected override GetDbInfoOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.GetDatabaseInfoAsync(options.DbName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to get database info for database: {DbName}", options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/ListDatabasesCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/ListDatabasesCommand.cs new file mode 100644 index 0000000000..888ad953c9 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Database/ListDatabasesCommand.cs @@ -0,0 +1,62 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Database; + +public sealed class ListDatabasesCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "d4e5f6a7-b8c9-4d4e-1f2a-3b4c5d6e7f8a"; + + public override string Name => "list_databases"; + + public override string Description => "List all databases in the DocumentDB instance"; + + public override string Title => "List Databases"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.ListDatabasesAsync(cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to list databases"); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/AggregateCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/AggregateCommand.cs new file mode 100644 index 0000000000..18e7aef58d --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/AggregateCommand.cs @@ -0,0 +1,88 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class AggregateCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "b0c1d2e3-f4a5-4b0c-7d8e-9f0a1b2c3d4e"; + + public override string Name => "aggregate"; + + public override string Description => "Run an aggregation pipeline on a collection"; + + public override string Title => "Aggregate Pipeline"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Pipeline); + command.Options.Add(DocumentDbOptionDefinitions.AllowDiskUse); + } + + protected override AggregateOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Pipeline = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Pipeline.Name); + options.AllowDiskUse = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.AllowDiskUse.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var pipeline = DocumentDbHelpers.ParseBsonDocumentList(options.Pipeline); + + if (pipeline == null || pipeline.Count == 0) + { + throw new ArgumentException("Invalid pipeline format or empty pipeline"); + } + + var result = await service.AggregateAsync(options.DbName!, options.CollectionName!, pipeline, options.AllowDiskUse, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to run aggregation pipeline on collection: {CollectionName}, database: {DbName}, allowDiskUse: {AllowDiskUse}", options.CollectionName, options.DbName, options.AllowDiskUse); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/CountDocumentsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/CountDocumentsCommand.cs new file mode 100644 index 0000000000..817bfa91ae --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/CountDocumentsCommand.cs @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + + + + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class CountDocumentsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "a3b4c5d6-e7f8-4a3b-0c1d-2e3f4a5b6c7d"; + + public override string Name => "count_documents"; + + public override string Description => "Count documents in a collection matching a query"; + + public override string Title => "Count Documents"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Query); + } + + protected override CountDocumentsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Query = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Query.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var query = DocumentDbHelpers.ParseBsonDocument(options.Query); + + var result = await service.CountDocumentsAsync(options.DbName!, options.CollectionName!, query, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to count documents in collection: {CollectionName}, database: {DbName}, query: {Query}", options.CollectionName, options.DbName, options.Query); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/DeleteDocumentCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/DeleteDocumentCommand.cs new file mode 100644 index 0000000000..db2443f9ca --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/DeleteDocumentCommand.cs @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + + + + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class DeleteDocumentCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "f8a9b0c1-d2e3-4f8a-5b6c-7d8e9f0a1b2c"; + + public override string Name => "delete_document"; + + public override string Description => "Delete a document from a collection"; + + public override string Title => "Delete Document"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Filter); + } + + protected override DeleteDocumentOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Filter = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Filter.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var filter = DocumentDbHelpers.ParseBsonDocument(options.Filter); + + if (filter == null) + { + throw new ArgumentException("Invalid filter format"); + } + + var result = await service.DeleteDocumentAsync(options.DbName!, options.CollectionName!, filter, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to delete document from collection: {CollectionName}, database: {DbName}, filter: {Filter}", options.CollectionName, options.DbName, options.Filter); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/DeleteManyCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/DeleteManyCommand.cs new file mode 100644 index 0000000000..cfb2d2258f --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/DeleteManyCommand.cs @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + + + + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class DeleteManyCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "a9b0c1d2-e3f4-4a9b-6c7d-8e9f0a1b2c3d"; + + public override string Name => "delete_many"; + + public override string Description => "Delete multiple documents from a collection"; + + public override string Title => "Delete Many Documents"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Filter); + } + + protected override DeleteManyOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Filter = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Filter.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var filter = DocumentDbHelpers.ParseBsonDocument(options.Filter); + + if (filter == null) + { + throw new ArgumentException("Invalid filter format"); + } + + var result = await service.DeleteManyAsync(options.DbName!, options.CollectionName!, filter, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to delete multiple documents from collection: {CollectionName}, database: {DbName}, filter: {Filter}", options.CollectionName, options.DbName, options.Filter); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainAggregateQueryCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainAggregateQueryCommand.cs new file mode 100644 index 0000000000..79d6d02b4b --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainAggregateQueryCommand.cs @@ -0,0 +1,86 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class ExplainAggregateQueryCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "f4a5b6c7-d8e9-4f4a-1b2c-3d4e5f6a7b8c"; + + public override string Name => "explain_aggregate_query"; + + public override string Description => "Explain the execution plan with execution stats for an aggregation query on a given collection"; + + public override string Title => "Explain Aggregate Query"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Pipeline); + } + + protected override ExplainAggregateQueryOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Pipeline = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Pipeline.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var pipeline = DocumentDbHelpers.ParseBsonDocumentList(options.Pipeline); + + if (pipeline == null || pipeline.Count == 0) + { + throw new ArgumentException("Invalid pipeline format or empty pipeline"); + } + + var result = await service.ExplainAggregateQueryAsync(options.DbName!, options.CollectionName!, pipeline, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to explain aggregate query on collection: {CollectionName}, database: {DbName}", options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainCountQueryCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainCountQueryCommand.cs new file mode 100644 index 0000000000..51bc926fe8 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainCountQueryCommand.cs @@ -0,0 +1,81 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class ExplainCountQueryCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "e3f4a5b6-c7d8-4e3f-0a1b-2c3d4e5f6a7b"; + + public override string Name => "explain_count_query"; + + public override string Description => "Explain the execution plan with execution stats for count query on a given collection"; + + public override string Title => "Explain Count Query"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Query); + } + + protected override ExplainCountQueryOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Query = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Query.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var query = DocumentDbHelpers.ParseBsonDocument(options.Query); + + var result = await service.ExplainCountQueryAsync(options.DbName!, options.CollectionName!, query, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to explain count query on collection: {CollectionName}, database: {DbName}, query: {Query}", options.CollectionName, options.DbName, options.Query); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainFindQueryCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainFindQueryCommand.cs new file mode 100644 index 0000000000..20d39d9b8c --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/ExplainFindQueryCommand.cs @@ -0,0 +1,84 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class ExplainFindQueryCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "d2e3f4a5-b6c7-4d2e-9f0a-1b2c3d4e5f6a"; + + public override string Name => "explain_find_query"; + + public override string Description => "Explain the execution plan with execution stats for a find query using consolidated options (sort, projection, limit, skip)"; + + public override string Title => "Explain Find Query"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Query); + command.Options.Add(DocumentDbOptionDefinitions.Options); + } + + protected override ExplainFindQueryOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Query = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Query.Name); + options.Options = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Options.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var query = DocumentDbHelpers.ParseBsonDocument(options.Query); + var queryOptions = string.IsNullOrWhiteSpace(options.Options) ? null : DocumentDbResponseHelper.DeserializeFromJson(options.Options); + + var result = await service.ExplainFindQueryAsync(options.DbName!, options.CollectionName!, query, queryOptions, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to explain find query on collection: {CollectionName}, database: {DbName}, query: {Query}", options.CollectionName, options.DbName, options.Query); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/FindAndModifyCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/FindAndModifyCommand.cs new file mode 100644 index 0000000000..ca41b0c65b --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/FindAndModifyCommand.cs @@ -0,0 +1,94 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + + + + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class FindAndModifyCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "c1d2e3f4-a5b6-4c1d-8e9f-0a1b2c3d4e5f"; + + public override string Name => "find_and_modify"; + + public override string Description => "Find one document by filter and apply update; returns the document BEFORE modification (or null if it doesn't exist)"; + + public override string Title => "Find And Modify Document"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Query); + command.Options.Add(DocumentDbOptionDefinitions.Update); + command.Options.Add(DocumentDbOptionDefinitions.Upsert); + } + + protected override FindAndModifyOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Query = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Query.Name); + options.Update = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Update.Name); + options.Upsert = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Upsert.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var query = DocumentDbHelpers.ParseBsonDocument(options.Query); + var update = DocumentDbHelpers.ParseBsonDocument(options.Update); + + if (query == null || update == null) + { + throw new ArgumentException("Invalid query or update format"); + } + + var result = await service.FindAndModifyAsync(options.DbName!, options.CollectionName!, query, update, options.Upsert, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to find and modify document in collection: {CollectionName}, database: {DbName}, query: {Query}, update: {Update}, upsert: {Upsert}", options.CollectionName, options.DbName, options.Query, options.Update, options.Upsert); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/FindDocumentsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/FindDocumentsCommand.cs new file mode 100644 index 0000000000..e7371195bb --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/FindDocumentsCommand.cs @@ -0,0 +1,97 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class FindDocumentsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "f2a3b4c5-d6e7-4f2a-9b0c-1d2e3f4a5b6c"; + + public override string Name => "find_documents"; + + public override string Description => "Find documents in a collection. Supports consolidated \"options\" object (limit, skip, sort, projection)"; + + public override string Title => "Find Documents"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + + var dbNameOption = new Option("--db-name") + { + Description = "Database name", + Required = true + }; + command.Options.Add(dbNameOption); + + var collectionNameOption = new Option("--collection-name") + { + Description = "Collection name", + Required = true + }; + command.Options.Add(collectionNameOption); + + command.Options.Add(DocumentDbOptionDefinitions.Query); + command.Options.Add(DocumentDbOptionDefinitions.Options); + } + + protected override FindDocumentsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Query = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Query.Name); + options.Options = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Options.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var query = DocumentDbHelpers.ParseBsonDocument(options.Query); + var queryOptions = string.IsNullOrWhiteSpace(options.Options) ? null : DocumentDbResponseHelper.DeserializeFromJson(options.Options); + + var result = await service.FindDocumentsAsync(options.DbName!, options.CollectionName!, query, queryOptions, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to find documents in collection: {CollectionName}, database: {DbName}, query: {Query}", options.CollectionName, options.DbName, options.Query); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/InsertDocumentCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/InsertDocumentCommand.cs new file mode 100644 index 0000000000..b9d2e80e6f --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/InsertDocumentCommand.cs @@ -0,0 +1,85 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class InsertDocumentCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "b4c5d6e7-f8a9-4b4c-1d2e-3f4a5b6c7d8e"; + + public override string Name => "insert_document"; + + public override string Description => "Insert a single document into a collection"; + + public override string Title => "Insert Document"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Document); + } + + protected override InsertDocumentOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Document = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Document.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var document = DocumentDbHelpers.ParseBsonDocument(options.Document); + if (document == null) + { + throw new ArgumentException("Invalid document format"); + } + + var result = await service.InsertDocumentAsync(options.DbName!, options.CollectionName!, document, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to insert document into collection: {CollectionName}, database: {DbName}", options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/InsertManyCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/InsertManyCommand.cs new file mode 100644 index 0000000000..44fa41452c --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/InsertManyCommand.cs @@ -0,0 +1,88 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + + + + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class InsertManyCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "c5d6e7f8-a9b0-4c5d-2e3f-4a5b6c7d8e9f"; + + public override string Name => "insert_many"; + + public override string Description => "Insert multiple documents into a collection"; + + public override string Title => "Insert Many Documents"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Documents); + } + + protected override InsertManyOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Documents = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Documents.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var documents = DocumentDbHelpers.ParseBsonDocumentList(options.Documents); + if (documents == null || documents.Count == 0) + { + throw new ArgumentException("Invalid documents format or empty list"); + } + + var result = await service.InsertManyAsync(options.DbName!, options.CollectionName!, documents, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to insert multiple documents into collection: {CollectionName}, database: {DbName}", options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/UpdateDocumentCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/UpdateDocumentCommand.cs new file mode 100644 index 0000000000..8b3d45c440 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/UpdateDocumentCommand.cs @@ -0,0 +1,91 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class UpdateDocumentCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "d6e7f8a9-b0c1-4d6e-3f4a-5b6c7d8e9f0a"; + + public override string Name => "update_document"; + + public override string Description => "Update a document in a collection"; + + public override string Title => "Update Single Document"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Filter); + command.Options.Add(DocumentDbOptionDefinitions.Update); + command.Options.Add(DocumentDbOptionDefinitions.Upsert); + } + + protected override UpdateDocumentOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Filter = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Filter.Name); + options.Update = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Update.Name); + options.Upsert = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Upsert.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var filter = DocumentDbHelpers.ParseBsonDocument(options.Filter); + var update = DocumentDbHelpers.ParseBsonDocument(options.Update); + + if (filter == null || update == null) + { + throw new ArgumentException("Invalid filter or update format"); + } + + var result = await service.UpdateDocumentAsync(options.DbName!, options.CollectionName!, filter, update, options.Upsert, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to update document in collection: {CollectionName}, database: {DbName}, filter: {Filter}, update: {Update}, upsert: {Upsert}", options.CollectionName, options.DbName, options.Filter, options.Update, options.Upsert); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/UpdateManyCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/UpdateManyCommand.cs new file mode 100644 index 0000000000..0c0637372f --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Document/UpdateManyCommand.cs @@ -0,0 +1,94 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + + + + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Document; + +public sealed class UpdateManyCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "e7f8a9b0-c1d2-4e7f-4a5b-6c7d8e9f0a1b"; + + public override string Name => "update_many"; + + public override string Description => "Update multiple documents in a collection"; + + public override string Title => "Update Many Documents"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Filter); + command.Options.Add(DocumentDbOptionDefinitions.Update); + command.Options.Add(DocumentDbOptionDefinitions.Upsert); + } + + protected override UpdateManyOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Filter = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Filter.Name); + options.Update = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Update.Name); + options.Upsert = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Upsert.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var filter = DocumentDbHelpers.ParseBsonDocument(options.Filter); + var update = DocumentDbHelpers.ParseBsonDocument(options.Update); + + if (filter == null || update == null) + { + throw new ArgumentException("Invalid filter or update format"); + } + + var result = await service.UpdateManyAsync(options.DbName!, options.CollectionName!, filter, update, options.Upsert, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to update multiple documents in collection: {CollectionName}, database: {DbName}, filter: {Filter}, update: {Update}, upsert: {Upsert}", options.CollectionName, options.DbName, options.Filter, options.Update, options.Upsert); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbHelpers.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbHelpers.cs new file mode 100644 index 0000000000..80d43f54fc --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbHelpers.cs @@ -0,0 +1,125 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using MongoDB.Bson; +using MongoDB.Bson.IO; +using MongoDB.Bson.Serialization; + +namespace Azure.Mcp.Tools.DocumentDb.Commands; + +internal static class DocumentDbHelpers +{ + public static BsonDocument? ParseBsonDocument(string? json) + { + if (string.IsNullOrWhiteSpace(json)) + return null; + + try + { + return BsonDocument.Parse(json); + } + catch + { + return null; + } + } + + public static BsonDocument? ParseBsonDocument(object? value) + { + if (value == null) + return null; + + if (value is string str) + return ParseBsonDocument(str); + + if (value is BsonDocument doc) + return doc; + + try + { + var json = DocumentDbResponseHelper.SerializeToJson(value); + return BsonDocument.Parse(json); + } + catch + { + return null; + } + } + + public static List? ParseBsonDocumentList(string? json) + { + if (string.IsNullOrWhiteSpace(json)) + return null; + + try + { + var bsonArray = BsonSerializer.Deserialize(json); + return bsonArray.Select(item => item.AsBsonDocument).ToList(); + } + catch + { + return null; + } + } + + public static List? ParseBsonDocumentList(object? value) + { + if (value == null) + return null; + + if (value is string str) + return ParseBsonDocumentList(str); + + if (value is List list) + return list; + + try + { + var json = DocumentDbResponseHelper.SerializeToJson(value); + return ParseBsonDocumentList(json); + } + catch + { + return null; + } + } + + public static bool ParseBoolean(string? value, bool defaultValue = false) + { + if (string.IsNullOrWhiteSpace(value)) + return defaultValue; + + if (bool.TryParse(value, out var result)) + return result; + + // Handle common string representations + return value.Trim().ToLowerInvariant() switch + { + "true" or "1" or "yes" => true, + "false" or "0" or "no" => false, + _ => defaultValue + }; + } + + public static int ParseInt(string? value, int defaultValue = 0) + { + if (string.IsNullOrWhiteSpace(value)) + return defaultValue; + + return int.TryParse(value, out var result) ? result : defaultValue; + } + + public static string SerializeBsonToJson(BsonDocument document) + { + var jsonWriterSettings = new JsonWriterSettings { OutputMode = JsonOutputMode.RelaxedExtendedJson }; + return document.ToJson(jsonWriterSettings); + } + + public static string SerializeBsonToJson(object obj) + { + if (obj is BsonDocument doc) + return SerializeBsonToJson(doc); + + return DocumentDbResponseHelper.SerializeToJson(obj); + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbJsonContext.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbJsonContext.cs new file mode 100644 index 0000000000..c615dedd82 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbJsonContext.cs @@ -0,0 +1,73 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.Text.Json.Serialization; +using MongoDB.Bson; + +namespace Azure.Mcp.Tools.DocumentDb.Commands; + +[JsonSourceGenerationOptions( + WriteIndented = false, + PropertyNamingPolicy = JsonKnownNamingPolicy.CamelCase, + DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull)] +[JsonSerializable(typeof(string))] +[JsonSerializable(typeof(object))] +[JsonSerializable(typeof(int))] +[JsonSerializable(typeof(bool))] +[JsonSerializable(typeof(long))] +[JsonSerializable(typeof(List))] +[JsonSerializable(typeof(Dictionary))] +[JsonSerializable(typeof(Dictionary))] +[JsonSerializable(typeof(List>))] +[JsonSerializable(typeof(List>))] +[JsonSerializable(typeof(System.Text.Json.JsonElement))] +internal partial class DocumentDbJsonContext : JsonSerializerContext; + +/// +/// Helper class for creating ResponseResult from JSON strings +/// +internal static class DocumentDbResponseHelper +{ + public static Microsoft.Mcp.Core.Models.Command.ResponseResult CreateFromJson(string json) + { + // Parse the JSON string to a JsonElement to get proper serialization + var element = System.Text.Json.JsonSerializer.Deserialize(json, DocumentDbJsonContext.Default.JsonElement); + return Microsoft.Mcp.Core.Models.Command.ResponseResult.Create(element, DocumentDbJsonContext.Default.JsonElement); + } + + public static string SerializeToJson(object value) + { + return value switch + { + // Handle BsonDocument by converting to JSON first + MongoDB.Bson.BsonDocument bsonDoc => bsonDoc.ToJson(new MongoDB.Bson.IO.JsonWriterSettings { OutputMode = MongoDB.Bson.IO.JsonOutputMode.RelaxedExtendedJson }), + List bsonList => "[" + string.Join(",", bsonList.Select(doc => doc.ToJson(new MongoDB.Bson.IO.JsonWriterSettings { OutputMode = MongoDB.Bson.IO.JsonOutputMode.RelaxedExtendedJson }))) + "]", + + // Handle standard types + Dictionary dict => System.Text.Json.JsonSerializer.Serialize(dict, DocumentDbJsonContext.Default.DictionaryStringObject), + List> list => System.Text.Json.JsonSerializer.Serialize(list, DocumentDbJsonContext.Default.ListDictionaryStringObject), + List strList => System.Text.Json.JsonSerializer.Serialize(strList, DocumentDbJsonContext.Default.ListString), + string str => System.Text.Json.JsonSerializer.Serialize(str, DocumentDbJsonContext.Default.String), + int i => System.Text.Json.JsonSerializer.Serialize(i, DocumentDbJsonContext.Default.Int32), + long l => System.Text.Json.JsonSerializer.Serialize(l, DocumentDbJsonContext.Default.Int64), + bool b => System.Text.Json.JsonSerializer.Serialize(b, DocumentDbJsonContext.Default.Boolean), + System.Text.Json.JsonElement element => System.Text.Json.JsonSerializer.Serialize(element, DocumentDbJsonContext.Default.JsonElement), + + // Handle IEnumerable (LINQ results) + System.Collections.Generic.IEnumerable enumStr => System.Text.Json.JsonSerializer.Serialize(enumStr.ToList(), DocumentDbJsonContext.Default.ListString), + + _ => throw new NotSupportedException($"Type {value.GetType().FullName} is not supported for AOT serialization. Please add it to DocumentDbJsonContext.") + }; + } + + public static T? DeserializeFromJson(string json) where T : class + { + // Only supports object type for AOT compatibility + if (typeof(T) == typeof(object)) + { + return System.Text.Json.JsonSerializer.Deserialize(json, DocumentDbJsonContext.Default.Object) as T; + } + + throw new NotSupportedException($"Type {typeof(T).Name} is not supported. Only 'object' type is AOT-compatible."); + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbOptionDefinitions.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbOptionDefinitions.cs new file mode 100644 index 0000000000..b9a938ebb8 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/DocumentDbOptionDefinitions.cs @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; + +namespace Azure.Mcp.Tools.DocumentDb.Commands; + +internal static class DocumentDbOptionDefinitions +{ + public static readonly Option ConnectionString = new("--connection-string") + { + Description = "DocumentDB connection string", + Required = true + }; + + public static readonly Option TestConnection = new("--test-connection") + { + Description = "Test connection after connecting", + DefaultValueFactory = _ => true + }; + + public static readonly Option DbName = new("--db-name") + { + Description = "Database name" + }; + + public static readonly Option CollectionName = new("--collection-name") + { + Description = "Collection name" + }; + + public static readonly Option NewCollectionName = new("--new-collection-name") + { + Description = "New collection name" + }; + + public static readonly Option SampleSize = new("--sample-size") + { + Description = "Number of documents to sample", + DefaultValueFactory = _ => 10 + }; + + public static readonly Option Query = new("--query") + { + Description = "Query filter in JSON format" + }; + + public static readonly Option Options = new("--options") + { + Description = "Query options" + }; + + public static readonly Option Document = new("--document") + { + Description = "Document to insert" + }; + + public static readonly Option Documents = new("--documents") + { + Description = "Documents to insert" + }; + + public static readonly Option Filter = new("--filter") + { + Description = "Filter for update/delete" + }; + + public static readonly Option Update = new("--update") + { + Description = "Update operations" + }; + + public static readonly Option Upsert = new("--upsert") + { + Description = "Create document if it doesn't exist", + DefaultValueFactory = _ => false + }; + + public static readonly Option Pipeline = new("--pipeline") + { + Description = "Aggregation pipeline" + }; + + public static readonly Option AllowDiskUse = new("--allow-disk-use") + { + Description = "Allow pipeline stages to write to disk", + DefaultValueFactory = _ => false + }; + + public static readonly Option Keys = new("--keys") + { + Description = "Index keys" + }; + + public static readonly Option IndexName = new("--index-name") + { + Description = "Index name" + }; + + public static readonly Option Ops = new("--ops") + { + Description = "Filter for current operations" + }; +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/CreateIndexCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/CreateIndexCommand.cs new file mode 100644 index 0000000000..c4475d8967 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/CreateIndexCommand.cs @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Index; + +public sealed class CreateIndexCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "a5b6c7d8-e9f0-4a5b-2c3d-4e5f6a7b8c9d"; + + public override string Name => "create_index"; + + public override string Description => "Create an index on a collection"; + + public override string Title => "Create Index"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.Keys); + command.Options.Add(DocumentDbOptionDefinitions.Options); + } + + protected override CreateIndexOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.Keys = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Keys.Name); + options.Options = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Options.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var keys = DocumentDbHelpers.ParseBsonDocument(options.Keys); + if (keys == null) + { + throw new ArgumentException("Invalid keys format"); + } + + var indexOptions = DocumentDbHelpers.ParseBsonDocument(options.Options); + + var result = await service.CreateIndexAsync(options.DbName!, options.CollectionName!, keys, indexOptions, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to create index on collection: {CollectionName}, database: {DbName}, keys: {Keys}", options.CollectionName, options.DbName, options.Keys); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/CurrentOpsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/CurrentOpsCommand.cs new file mode 100644 index 0000000000..3c815c15bd --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/CurrentOpsCommand.cs @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Index; + +public sealed class CurrentOpsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "e9f0a1b2-c3d4-4e9f-6a7b-8c9d0e1f2a3b"; + + public override string Name => "current_ops"; + + public override string Description => "Get information about current DocumentDB operations"; + + public override string Title => "Current Operations"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.Ops); + } + + protected override CurrentOpsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.Ops = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.Ops.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var filter = DocumentDbHelpers.ParseBsonDocument(options.Ops); + + var result = await service.GetCurrentOpsAsync(filter, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbHelpers.SerializeBsonToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to get current operations with filter: {Ops}", options.Ops); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/DropIndexCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/DropIndexCommand.cs new file mode 100644 index 0000000000..f008da7680 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/DropIndexCommand.cs @@ -0,0 +1,79 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Index; + +public sealed class DropIndexCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "c7d8e9f0-a1b2-4c7d-4e5f-6a7b8c9d0e1f"; + + public override string Name => "drop_index"; + + public override string Description => "Drop an index from a collection"; + + public override string Title => "Drop Index"; + + public override ToolMetadata Metadata => new() + { + Destructive = true, + ReadOnly = false + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + command.Options.Add(DocumentDbOptionDefinitions.IndexName); + } + + protected override DropIndexOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + options.IndexName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.IndexName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.DropIndexAsync(options.DbName!, options.CollectionName!, options.IndexName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to drop index: {IndexName} from collection: {CollectionName}, database: {DbName}", options.IndexName, options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/IndexStatsCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/IndexStatsCommand.cs new file mode 100644 index 0000000000..7f6fd34eaf --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/IndexStatsCommand.cs @@ -0,0 +1,77 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Index; + +public sealed class IndexStatsCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "d8e9f0a1-b2c3-4d8e-5f6a-7b8c9d0e1f2a"; + + public override string Name => "index_stats"; + + public override string Description => "Get statistics for indexes on a collection"; + + public override string Title => "Index Statistics"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + command.Options.Add(DocumentDbOptionDefinitions.DbName); + command.Options.Add(DocumentDbOptionDefinitions.CollectionName); + } + + protected override IndexStatsOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.GetIndexStatsAsync(options.DbName!, options.CollectionName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result.Select(doc => DocumentDbHelpers.SerializeBsonToJson(doc)))); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to get index statistics for collection: {CollectionName}, database: {DbName}", options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/ListIndexesCommand.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/ListIndexesCommand.cs new file mode 100644 index 0000000000..79c74179c8 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Commands/Index/ListIndexesCommand.cs @@ -0,0 +1,89 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using Azure.Mcp.Core.Commands; +using Azure.Mcp.Core.Extensions; +using Azure.Mcp.Tools.DocumentDb.Options; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb.Commands.Index; + +public sealed class ListIndexesCommand(ILogger logger) + : BaseDocumentDbCommand(logger) +{ + public override string Id => "b6c7d8e9-f0a1-4b6c-3d4e-5f6a7b8c9d0e"; + + public override string Name => "list_indexes"; + + public override string Description => "List all indexes on a collection"; + + public override string Title => "List Indexes"; + + public override ToolMetadata Metadata => new() + { + Destructive = false, + ReadOnly = true + }; + + protected override void RegisterOptions(Command command) + { + base.RegisterOptions(command); + + var dbNameOption = new Option("--db-name") + { + Description = "Database name", + Required = true + }; + command.Options.Add(dbNameOption); + + var collectionNameOption = new Option("--collection-name") + { + Description = "Collection name", + Required = true + }; + command.Options.Add(collectionNameOption); + } + + protected override ListIndexesOptions BindOptions(ParseResult parseResult) + { + var options = base.BindOptions(parseResult); + options.DbName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.DbName.Name); + options.CollectionName = parseResult.GetValueOrDefault(DocumentDbOptionDefinitions.CollectionName.Name); + return options; + } + + public override async Task ExecuteAsync( + CommandContext context, + ParseResult parseResult, + CancellationToken cancellationToken) + { + var options = BindOptions(parseResult); + + try + { + if (!Validate(parseResult.CommandResult, context.Response).IsValid) + { + return context.Response; + } + + var service = context.GetService(); + + var result = await service.ListIndexesAsync(options.DbName!, options.CollectionName!, cancellationToken); + + context.Response.Results = DocumentDbResponseHelper.CreateFromJson( + DocumentDbResponseHelper.SerializeToJson(result)); + + return context.Response; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to list indexes on collection: {CollectionName}, database: {DbName}", options.CollectionName, options.DbName); + HandleException(context, ex); + return context.Response; + } + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/DocumentDbSetup.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/DocumentDbSetup.cs new file mode 100644 index 0000000000..0a6f3e6586 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/DocumentDbSetup.cs @@ -0,0 +1,197 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using Azure.Mcp.Tools.DocumentDb.Commands.Collection; +using Azure.Mcp.Tools.DocumentDb.Commands.Connection; +using Azure.Mcp.Tools.DocumentDb.Commands.Database; +using Azure.Mcp.Tools.DocumentDb.Commands.Document; +using Azure.Mcp.Tools.DocumentDb.Commands.Index; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Mcp.Core.Areas; +using Microsoft.Mcp.Core.Commands; + +namespace Azure.Mcp.Tools.DocumentDb; + +public class DocumentDbSetup : IAreaSetup +{ + public string Name => "documentdb"; + public string Title => "Azure Cosmos DB for MongoDB (vCore)"; + + public void ConfigureServices(IServiceCollection services) + { + services.AddSingleton(); + + // Connection Commands + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + + // Database Commands + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + + // Collection Commands + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + + // Document Commands + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + + // Index Commands + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + services.AddSingleton(); + } + + public CommandGroup RegisterCommands(IServiceProvider serviceProvider) + { + // Create DocumentDB root command group + var documentDb = new CommandGroup( + Name, + "Azure Cosmos DB for MongoDB (vCore) operations - Manage databases, collections, and documents in MongoDB-compatible Azure Cosmos DB.", + Title); + + // Connection subgroup + var connection = new CommandGroup( + "connection", + "Connection management - Commands for connecting to and managing DocumentDB connections."); + documentDb.AddSubGroup(connection); + + connection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + connection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + connection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + + // Database subgroup + var database = new CommandGroup( + "database", + "Database operations - Commands for managing DocumentDB databases."); + documentDb.AddSubGroup(database); + + database.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + database.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + database.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + database.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + + // Collection subgroup + var collection = new CommandGroup( + "collection", + "Collection operations - Commands for managing DocumentDB collections."); + documentDb.AddSubGroup(collection); + + collection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + collection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + collection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + collection.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + + // Document subgroup + var document = new CommandGroup( + "document", + "Document operations - Commands for querying and manipulating documents in DocumentDB collections."); + documentDb.AddSubGroup(document); + + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + document.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + + // Index subgroup + var index = new CommandGroup( + "index", + "Index operations - Commands for managing indexes on DocumentDB collections."); + documentDb.AddSubGroup(index); + + index.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + index.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + index.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + index.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + index.AddCommand( + serviceProvider.GetRequiredService().Name, + serviceProvider.GetRequiredService()); + + return documentDb; + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/GlobalUsings.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/GlobalUsings.cs new file mode 100644 index 0000000000..b41cc886b4 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/GlobalUsings.cs @@ -0,0 +1,4 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +global using System.CommandLine; diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Options/BaseDocumentDbOptions.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Options/BaseDocumentDbOptions.cs new file mode 100644 index 0000000000..76756a211f --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Options/BaseDocumentDbOptions.cs @@ -0,0 +1,8 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using Azure.Mcp.Core.Options; + +namespace Azure.Mcp.Tools.DocumentDb.Options; + +public class BaseDocumentDbOptions : GlobalOptions; diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Options/DocumentDbCommandOptions.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Options/DocumentDbCommandOptions.cs new file mode 100644 index 0000000000..d2ff12cf91 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Options/DocumentDbCommandOptions.cs @@ -0,0 +1,194 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +namespace Azure.Mcp.Tools.DocumentDb.Options; + +// Connection Options +public class ConnectDocumentDbOptions : BaseDocumentDbOptions +{ + public string? ConnectionString { get; set; } + public bool TestConnection { get; set; } = true; +} + +public class DisconnectDocumentDbOptions : BaseDocumentDbOptions; + +public class GetConnectionStatusOptions : BaseDocumentDbOptions; + +// Database Options +public class ListDatabasesOptions : BaseDocumentDbOptions; + +public class DbStatsOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } +} + +public class GetDbInfoOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } +} + +public class DropDatabaseOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } +} + +// Collection Options +public class CollectionStatsOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } +} + +public class RenameCollectionOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? NewCollectionName { get; set; } +} + +public class DropCollectionOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } +} + +public class SampleDocumentsOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public int SampleSize { get; set; } = 10; +} + +// Document Options +public class FindDocumentsOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Query { get; set; } + public string? Options { get; set; } +} + +public class CountDocumentsOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Query { get; set; } +} + +public class InsertDocumentOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Document { get; set; } +} + +public class InsertManyOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Documents { get; set; } +} + +public class UpdateDocumentOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Filter { get; set; } + public string? Update { get; set; } + public bool Upsert { get; set; } +} + +public class UpdateManyOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Filter { get; set; } + public string? Update { get; set; } + public bool Upsert { get; set; } +} + +public class DeleteDocumentOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Filter { get; set; } +} + +public class DeleteManyOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Filter { get; set; } +} + +public class AggregateOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Pipeline { get; set; } + public bool AllowDiskUse { get; set; } +} + +public class FindAndModifyOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Query { get; set; } + public string? Update { get; set; } + public bool Upsert { get; set; } +} + +public class ExplainFindQueryOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Query { get; set; } + public string? Options { get; set; } +} + +public class ExplainCountQueryOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Query { get; set; } +} + +public class ExplainAggregateQueryOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Pipeline { get; set; } +} + +// Index Options +public class CreateIndexOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? Keys { get; set; } + public string? Options { get; set; } +} + +public class ListIndexesOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } +} + +public class DropIndexOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } + public string? IndexName { get; set; } +} + +public class IndexStatsOptions : BaseDocumentDbOptions +{ + public string? DbName { get; set; } + public string? CollectionName { get; set; } +} + +public class CurrentOpsOptions : BaseDocumentDbOptions +{ + public string? Ops { get; set; } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Services/DocumentDbService.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Services/DocumentDbService.cs new file mode 100644 index 0000000000..9b9acf9e31 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Services/DocumentDbService.cs @@ -0,0 +1,1128 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using Microsoft.Extensions.Logging; +using MongoDB.Bson; +using MongoDB.Driver; + +namespace Azure.Mcp.Tools.DocumentDb.Services; + +public class DocumentDbService : IDocumentDbService +{ + private readonly ILogger _logger; + private MongoClient? _client; + private string? _connectionString; + private bool _disposed; + + public DocumentDbService(ILogger logger) + { + _logger = logger ?? throw new ArgumentNullException(nameof(logger)); + } + + /// + /// Helper method to convert BsonDocument to JSON string for serialization + /// + private static string? BsonDocumentToJson(BsonDocument? doc) + { + return doc?.ToJson(new MongoDB.Bson.IO.JsonWriterSettings { OutputMode = MongoDB.Bson.IO.JsonOutputMode.RelaxedExtendedJson }); + } + + /// + /// Helper method to convert List of BsonDocument to List of JSON strings for serialization + /// + private static List BsonDocumentListToJson(List docs) + { + return docs.Select(doc => doc.ToJson(new MongoDB.Bson.IO.JsonWriterSettings { OutputMode = MongoDB.Bson.IO.JsonOutputMode.RelaxedExtendedJson })).ToList(); + } + + #region Connection Management + + public async Task ConnectAsync(string connectionString, bool testConnection = true, CancellationToken cancellationToken = default) + { + try + { + if (string.IsNullOrWhiteSpace(connectionString)) + { + return new Dictionary + { + ["success"] = false, + ["message"] = "Connection string cannot be empty", + ["data"] = null + }; + } + + // Disconnect any existing connection + if (_client != null) + { + await DisconnectAsync(cancellationToken); + } + + _connectionString = connectionString; + var settings = MongoClientSettings.FromConnectionString(connectionString); + settings.ServerSelectionTimeout = TimeSpan.FromSeconds(10); + _client = new MongoClient(settings); + + if (testConnection) + { + // Test the connection by listing databases + var databases = await _client.ListDatabaseNames(cancellationToken: cancellationToken).ToListAsync(cancellationToken); + _logger.LogInformation("Successfully connected to DocumentDB. Found {Count} databases", databases.Count); + + return new Dictionary + { + ["success"] = true, + ["message"] = "Connected successfully", + ["data"] = new Dictionary + { + ["databaseCount"] = databases.Count, + ["databases"] = databases + } + }; + } + + return new Dictionary + { + ["success"] = true, + ["message"] = "Connected successfully (not tested)", + ["data"] = null + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Failed to connect to DocumentDB"); + _client = null; + _connectionString = null; + return new Dictionary + { + ["success"] = false, + ["message"] = $"Connection failed: {ex.Message}", + ["data"] = null + }; + } + } + + public Task DisconnectAsync(CancellationToken cancellationToken = default) + { + try + { + if (_client == null) + { + return Task.FromResult(new Dictionary + { + ["success"] = true, + ["message"] = "No active connection" + }); + } + + _client = null; + _connectionString = null; + _logger.LogInformation("Disconnected from DocumentDB"); + return Task.FromResult(new Dictionary + { + ["success"] = true, + ["message"] = "Disconnected successfully" + }); + } + catch (Exception ex) + { + _logger.LogError(ex, "Error during disconnect"); + return Task.FromResult(new Dictionary + { + ["success"] = false, + ["message"] = $"Disconnect failed: {ex.Message}" + }); + } + } + + public object GetConnectionStatus() + { + if (_client == null) + { + return new Dictionary + { + ["isConnected"] = false, + ["connectionString"] = null, + ["details"] = null + }; + } + + var sanitizedConnectionString = SanitizeConnectionString(_connectionString); + return new Dictionary + { + ["isConnected"] = true, + ["connectionString"] = sanitizedConnectionString, + ["details"] = new Dictionary + { + ["status"] = "Connected" + } + }; + } + + private string? SanitizeConnectionString(string? connectionString) + { + if (string.IsNullOrEmpty(connectionString)) + return null; + + // Hide password from connection string + try + { + var uri = new Uri(connectionString); + if (!string.IsNullOrEmpty(uri.UserInfo)) + { + var sanitized = connectionString.Replace(uri.UserInfo, "***:***"); + return sanitized; + } + } + catch + { + // If parsing fails, just return a placeholder + return "mongodb://***"; + } + + return connectionString; + } + + #endregion + + #region Database Operations + + public async Task> ListDatabasesAsync(CancellationToken cancellationToken = default) + { + EnsureConnected(); + try + { + var databases = await _client!.ListDatabaseNames(cancellationToken: cancellationToken).ToListAsync(cancellationToken); + return databases; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error listing databases"); + throw new Exception($"Failed to list databases: {ex.Message}", ex); + } + } + + public async Task GetDatabaseStatsAsync(string databaseName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var command = new BsonDocument { { "dbStats", 1 } }; + var stats = await database.RunCommandAsync(command, cancellationToken: cancellationToken); + return stats; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error getting database stats for {DatabaseName}", databaseName); + throw new Exception($"Failed to get database stats: {ex.Message}", ex); + } + } + + public async Task GetDatabaseInfoAsync(string databaseName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collections = await database.ListCollectionNames(cancellationToken: cancellationToken).ToListAsync(cancellationToken); + + var collectionInfos = new List>(); + foreach (var collectionName in collections) + { + try + { + var collection = database.GetCollection(collectionName); + var count = await collection.CountDocumentsAsync(new BsonDocument(), cancellationToken: cancellationToken); + collectionInfos.Add(new Dictionary + { + ["name"] = collectionName, + ["count"] = count + }); + } + catch (Exception ex) + { + collectionInfos.Add(new Dictionary + { + ["name"] = collectionName, + ["error"] = ex.Message + }); + } + } + + return new Dictionary + { + ["database_name"] = databaseName, + ["collections"] = collectionInfos + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error getting database info for {DatabaseName}", databaseName); + throw new Exception($"Failed to get database info: {ex.Message}", ex); + } + } + + public async Task DropDatabaseAsync(string databaseName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + + try + { + await _client!.DropDatabaseAsync(databaseName, cancellationToken); + _logger.LogWarning("Dropped database {DatabaseName}", databaseName); + return new Dictionary + { + ["success"] = true, + ["message"] = $"Database '{databaseName}' dropped successfully", + ["data"] = new Dictionary + { + ["database_name"] = databaseName + } + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error dropping database {DatabaseName}", databaseName); + return new Dictionary + { + ["success"] = false, + ["message"] = $"Failed to drop database: {ex.Message}", + ["data"] = null + }; + } + } + + #endregion + + #region Collection Operations + + public async Task GetCollectionStatsAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var command = new BsonDocument { { "collStats", collectionName } }; + var stats = await database.RunCommandAsync(command, cancellationToken: cancellationToken); + return stats; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error getting collection stats for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to get collection stats: {ex.Message}", ex); + } + } + + public async Task RenameCollectionAsync(string databaseName, string oldName, string newName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(oldName, nameof(oldName)); + ValidateParameter(newName, nameof(newName)); + + try + { + var database = _client!.GetDatabase(databaseName); + await database.RenameCollectionAsync(oldName, newName, cancellationToken: cancellationToken); + _logger.LogInformation("Renamed collection {OldName} to {NewName} in database {DatabaseName}", oldName, newName, databaseName); + return new Dictionary + { + ["success"] = true, + ["message"] = $"Collection renamed from '{oldName}' to '{newName}' successfully" + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error renaming collection in {DatabaseName}", databaseName); + return new Dictionary + { + ["success"] = false, + ["message"] = $"Failed to rename collection: {ex.Message}" + }; + } + } + + public async Task DropCollectionAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + await database.DropCollectionAsync(collectionName, cancellationToken); + _logger.LogWarning("Dropped collection {CollectionName} from database {DatabaseName}", collectionName, databaseName); + return new Dictionary + { + ["success"] = true, + ["message"] = $"Collection '{collectionName}' dropped successfully" + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error dropping collection {CollectionName} from {DatabaseName}", collectionName, databaseName); + return new Dictionary + { + ["success"] = false, + ["message"] = $"Failed to drop collection: {ex.Message}" + }; + } + } + + public async Task> SampleDocumentsAsync(string databaseName, string collectionName, int sampleSize = 10, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var pipeline = new[] + { + new BsonDocument("$sample", new BsonDocument("size", sampleSize)) + }; + + var documents = await collection.Aggregate(pipeline, cancellationToken: cancellationToken).ToListAsync(cancellationToken); + return documents; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error sampling documents from {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to sample documents: {ex.Message}", ex); + } + } + + #endregion + + #region Document Operations + + public async Task FindDocumentsAsync(string databaseName, string collectionName, BsonDocument? query = null, object? options = null, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var filter = query ?? new BsonDocument(); + var findOptions = new FindOptions(); + + // Parse options if provided + int limit = 100; + int skip = 0; + BsonDocument? sort = null; + BsonDocument? projection = null; + + if (options is BsonDocument optionsDoc) + { + if (optionsDoc.Contains("limit")) + limit = optionsDoc["limit"].ToInt32(); + if (optionsDoc.Contains("skip")) + skip = optionsDoc["skip"].ToInt32(); + if (optionsDoc.Contains("sort") && optionsDoc["sort"].IsBsonDocument) + sort = optionsDoc["sort"].AsBsonDocument; + if (optionsDoc.Contains("projection") && optionsDoc["projection"].IsBsonDocument) + projection = optionsDoc["projection"].AsBsonDocument; + } + + findOptions.Limit = limit; + findOptions.Skip = skip; + if (sort != null) + findOptions.Sort = new SortDefinitionBuilder().Combine(sort); + if (projection != null) + findOptions.Projection = new ProjectionDefinitionBuilder().Combine(projection); + + var cursor = collection.Find(filter) + .Limit(limit) + .Skip(skip); + + if (sort != null) + cursor = cursor.Sort(sort); + if (projection != null) + cursor = cursor.Project(projection); + + var documents = await cursor.ToListAsync(cancellationToken); + var totalCount = await collection.CountDocumentsAsync(filter, cancellationToken: cancellationToken); + + return new Dictionary + { + ["documents"] = BsonDocumentListToJson(documents), + ["total_count"] = totalCount, + ["returned_count"] = documents.Count, + ["has_more"] = totalCount > (skip + documents.Count), + ["query"] = BsonDocumentToJson(filter), + ["applied_options"] = new Dictionary + { + ["limit"] = limit, + ["skip"] = skip, + ["sort"] = BsonDocumentToJson(sort), + ["projection"] = BsonDocumentToJson(projection) + } + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error finding documents in {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to find documents: {ex.Message}", ex); + } + } + + public async Task CountDocumentsAsync(string databaseName, string collectionName, BsonDocument? query = null, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var filter = query ?? new BsonDocument(); + var count = await collection.CountDocumentsAsync(filter, cancellationToken: cancellationToken); + + return new Dictionary + { + ["count"] = count, + ["query"] = BsonDocumentToJson(filter) + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error counting documents in {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to count documents: {ex.Message}", ex); + } + } + + public async Task InsertDocumentAsync(string databaseName, string collectionName, BsonDocument document, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(document); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + await collection.InsertOneAsync(document, cancellationToken: cancellationToken); + var insertedId = document["_id"].ToString(); + + _logger.LogInformation("Inserted document with ID {Id} into {DatabaseName}.{CollectionName}", insertedId, databaseName, collectionName); + + return new Dictionary + { + ["inserted_id"] = insertedId, + ["acknowledged"] = true, + ["inserted_count"] = 1 + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error inserting document into {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to insert document: {ex.Message}", ex); + } + } + + public async Task InsertManyAsync(string databaseName, string collectionName, List documents, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(documents); + + if (documents.Count == 0) + { + return new Dictionary + { + ["inserted_ids"] = Array.Empty(), + ["acknowledged"] = true, + ["inserted_count"] = 0 + }; + } + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + await collection.InsertManyAsync(documents, cancellationToken: cancellationToken); + var insertedIds = documents.Select(d => d["_id"].ToString()).ToList(); + + _logger.LogInformation("Inserted {Count} documents into {DatabaseName}.{CollectionName}", documents.Count, databaseName, collectionName); + + return new Dictionary + { + ["inserted_ids"] = insertedIds, + ["acknowledged"] = true, + ["inserted_count"] = documents.Count + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error inserting documents into {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to insert documents: {ex.Message}", ex); + } + } + + public async Task UpdateDocumentAsync(string databaseName, string collectionName, BsonDocument filter, BsonDocument update, bool upsert = false, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(filter); + ArgumentNullException.ThrowIfNull(update); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var options = new UpdateOptions { IsUpsert = upsert }; + var result = await collection.UpdateOneAsync(filter, update, options, cancellationToken); + + _logger.LogInformation("Updated document in {DatabaseName}.{CollectionName}. Matched: {Matched}, Modified: {Modified}", + databaseName, collectionName, result.MatchedCount, result.ModifiedCount); + + return new Dictionary + { + ["matched_count"] = result.MatchedCount, + ["modified_count"] = result.ModifiedCount, + ["upserted_id"] = result.UpsertedId?.ToString(), + ["acknowledged"] = result.IsAcknowledged + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error updating document in {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to update document: {ex.Message}", ex); + } + } + + public async Task UpdateManyAsync(string databaseName, string collectionName, BsonDocument filter, BsonDocument update, bool upsert = false, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(filter); + ArgumentNullException.ThrowIfNull(update); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var options = new UpdateOptions { IsUpsert = upsert }; + var result = await collection.UpdateManyAsync(filter, update, options, cancellationToken); + + _logger.LogInformation("Updated documents in {DatabaseName}.{CollectionName}. Matched: {Matched}, Modified: {Modified}", + databaseName, collectionName, result.MatchedCount, result.ModifiedCount); + + return new Dictionary + { + ["matched_count"] = result.MatchedCount, + ["modified_count"] = result.ModifiedCount, + ["upserted_id"] = result.UpsertedId?.ToString(), + ["acknowledged"] = result.IsAcknowledged + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error updating documents in {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to update documents: {ex.Message}", ex); + } + } + + public async Task DeleteDocumentAsync(string databaseName, string collectionName, BsonDocument filter, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(filter); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var result = await collection.DeleteOneAsync(filter, cancellationToken); + + _logger.LogInformation("Deleted {Count} document from {DatabaseName}.{CollectionName}", + result.DeletedCount, databaseName, collectionName); + + return new Dictionary + { + ["deleted_count"] = result.DeletedCount, + ["acknowledged"] = result.IsAcknowledged + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error deleting document from {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to delete document: {ex.Message}", ex); + } + } + + public async Task DeleteManyAsync(string databaseName, string collectionName, BsonDocument filter, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(filter); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var result = await collection.DeleteManyAsync(filter, cancellationToken); + + _logger.LogInformation("Deleted {Count} documents from {DatabaseName}.{CollectionName}", + result.DeletedCount, databaseName, collectionName); + + return new Dictionary + { + ["deleted_count"] = result.DeletedCount, + ["acknowledged"] = result.IsAcknowledged + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error deleting documents from {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to delete documents: {ex.Message}", ex); + } + } + + public async Task AggregateAsync(string databaseName, string collectionName, List pipeline, bool allowDiskUse = false, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(pipeline); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var options = new AggregateOptions { AllowDiskUse = allowDiskUse }; + var results = await collection.Aggregate(pipeline, options, cancellationToken: cancellationToken).ToListAsync(cancellationToken); + + return new Dictionary + { + ["results"] = BsonDocumentListToJson(results), + ["total_count"] = results.Count + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error executing aggregation pipeline in {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to execute aggregation: {ex.Message}", ex); + } + } + + public async Task FindAndModifyAsync(string databaseName, string collectionName, BsonDocument query, BsonDocument update, bool upsert = false, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(query); + ArgumentNullException.ThrowIfNull(update); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var options = new FindOneAndUpdateOptions + { + IsUpsert = upsert, + ReturnDocument = ReturnDocument.Before + }; + + var result = await collection.FindOneAndUpdateAsync(query, update, options, cancellationToken); + + return new Dictionary + { + ["matched"] = result != null, + ["upsertedId"] = result?["_id"]?.ToString(), + ["original_document"] = BsonDocumentToJson(result), + ["query"] = BsonDocumentToJson(query), + ["update"] = BsonDocumentToJson(update), + ["upsert"] = upsert + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error in find and modify for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to find and modify document: {ex.Message}", ex); + } + } + + public async Task ExplainFindQueryAsync(string databaseName, string collectionName, BsonDocument? query = null, object? options = null, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var filter = query ?? new BsonDocument(); + var findOptions = new FindOptions(); + + // Parse options + BsonDocument? sort = null; + BsonDocument? projection = null; + int? limit = null; + int? skip = null; + + if (options is BsonDocument optionsDoc) + { + if (optionsDoc.Contains("sort") && optionsDoc["sort"].IsBsonDocument) + sort = optionsDoc["sort"].AsBsonDocument; + if (optionsDoc.Contains("projection") && optionsDoc["projection"].IsBsonDocument) + projection = optionsDoc["projection"].AsBsonDocument; + if (optionsDoc.Contains("limit")) + limit = optionsDoc["limit"].ToInt32(); + if (optionsDoc.Contains("skip")) + skip = optionsDoc["skip"].ToInt32(); + } + + // Build the explain command + var findCommand = new BsonDocument + { + { "find", collectionName }, + { "filter", filter } + }; + + if (sort != null) + findCommand.Add("sort", sort); + if (projection != null) + findCommand.Add("projection", projection); + if (limit.HasValue) + findCommand.Add("limit", limit.Value); + if (skip.HasValue) + findCommand.Add("skip", skip.Value); + + var command = new BsonDocument + { + { "explain", findCommand }, + { "verbosity", "executionStats" } + }; + + var explain = await database.RunCommandAsync(command, cancellationToken: cancellationToken); + + return new Dictionary + { + ["options_applied"] = new Dictionary + { + ["sort"] = BsonDocumentToJson(sort), + ["projection"] = BsonDocumentToJson(projection), + ["limit"] = limit, + ["skip"] = skip + }, + ["explain"] = BsonDocumentToJson(explain) + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error explaining find query for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to explain find query: {ex.Message}", ex); + } + } + + public async Task ExplainCountQueryAsync(string databaseName, string collectionName, BsonDocument? query = null, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + + // Build explain command for count + var command = new BsonDocument + { + { "explain", new BsonDocument + { + { "count", collectionName }, + { "query", query ?? new BsonDocument() } + } + }, + { "verbosity", "executionStats" } + }; + + var explain = await database.RunCommandAsync(command, cancellationToken: cancellationToken); + return new Dictionary + { + ["explain"] = BsonDocumentToJson(explain) + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error explaining count query for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to explain count query: {ex.Message}", ex); + } + } + + public async Task ExplainAggregateQueryAsync(string databaseName, string collectionName, List pipeline, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(pipeline); + + try + { + var database = _client!.GetDatabase(databaseName); + + var command = new BsonDocument + { + { "explain", new BsonDocument + { + { "aggregate", collectionName }, + { "pipeline", new BsonArray(pipeline) }, + { "cursor", new BsonDocument() } + } + }, + { "verbosity", "executionStats" } + }; + + var explain = await database.RunCommandAsync(command, cancellationToken: cancellationToken); + return new Dictionary + { + ["explain"] = BsonDocumentToJson(explain) + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error explaining aggregate query for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to explain aggregate query: {ex.Message}", ex); + } + } + + #endregion + + #region Index Operations + + public async Task CreateIndexAsync(string databaseName, string collectionName, BsonDocument keys, BsonDocument? options = null, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ArgumentNullException.ThrowIfNull(keys); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var indexKeysDefinition = new BsonDocumentIndexKeysDefinition(keys); + var createIndexOptions = new CreateIndexOptions(); + + if (options != null) + { + if (options.Contains("unique")) + createIndexOptions.Unique = options["unique"].AsBoolean; + if (options.Contains("name")) + createIndexOptions.Name = options["name"].AsString; + if (options.Contains("sparse")) + createIndexOptions.Sparse = options["sparse"].AsBoolean; + if (options.Contains("expireAfterSeconds")) + createIndexOptions.ExpireAfter = TimeSpan.FromSeconds(options["expireAfterSeconds"].ToInt32()); + } + + var model = new CreateIndexModel(indexKeysDefinition, createIndexOptions); + var indexName = await collection.Indexes.CreateOneAsync(model, cancellationToken: cancellationToken); + + _logger.LogInformation("Created index {IndexName} on {DatabaseName}.{CollectionName}", indexName, databaseName, collectionName); + + return new Dictionary + { + ["index_name"] = indexName, + ["keys"] = BsonDocumentToJson(keys), + ["options"] = BsonDocumentToJson(options ?? new BsonDocument()) + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error creating index on {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to create index: {ex.Message}", ex); + } + } + + public async Task ListIndexesAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var indexes = await collection.Indexes.List(cancellationToken: cancellationToken).ToListAsync(cancellationToken); + + return new Dictionary + { + ["indexes"] = BsonDocumentListToJson(indexes), + ["count"] = indexes.Count + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error listing indexes for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to list indexes: {ex.Message}", ex); + } + } + + public async Task DropIndexAsync(string databaseName, string collectionName, string indexName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + ValidateParameter(indexName, nameof(indexName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + await collection.Indexes.DropOneAsync(indexName, cancellationToken); + + _logger.LogWarning("Dropped index {IndexName} from {DatabaseName}.{CollectionName}", indexName, databaseName, collectionName); + + return new Dictionary + { + ["success"] = true, + ["message"] = $"Index '{indexName}' dropped successfully", + ["data"] = new Dictionary + { + ["index_name"] = indexName + } + }; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error dropping index {IndexName} from {DatabaseName}.{CollectionName}", indexName, databaseName, collectionName); + return new Dictionary + { + ["success"] = false, + ["message"] = $"Failed to drop index: {ex.Message}", + ["data"] = null + }; + } + } + + public async Task> GetIndexStatsAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default) + { + EnsureConnected(); + ValidateParameter(databaseName, nameof(databaseName)); + ValidateParameter(collectionName, nameof(collectionName)); + + try + { + var database = _client!.GetDatabase(databaseName); + var collection = database.GetCollection(collectionName); + + var pipeline = new[] + { + new BsonDocument("$indexStats", new BsonDocument()) + }; + + var stats = await collection.Aggregate(pipeline, cancellationToken: cancellationToken).ToListAsync(cancellationToken); + return stats; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error getting index stats for {DatabaseName}.{CollectionName}", databaseName, collectionName); + throw new Exception($"Failed to get index stats: {ex.Message}", ex); + } + } + + public async Task GetCurrentOpsAsync(BsonDocument? filter = null, CancellationToken cancellationToken = default) + { + EnsureConnected(); + + try + { + var adminDb = _client!.GetDatabase("admin"); + var command = new BsonDocument("currentOp", 1); + + if (filter != null && filter.ElementCount > 0) + { + foreach (var element in filter) + { + command.Add(element); + } + } + + var result = await adminDb.RunCommandAsync(command, cancellationToken: cancellationToken); + return result; + } + catch (Exception ex) + { + _logger.LogError(ex, "Error getting current operations"); + throw new Exception($"Failed to get current operations: {ex.Message}", ex); + } + } + + #endregion + + #region Helper Methods + + private void EnsureConnected() + { + if (_client == null) + { + throw new InvalidOperationException("Not connected to DocumentDB. Please call ConnectAsync first."); + } + } + + private void ValidateParameter(string? value, string paramName) + { + if (string.IsNullOrWhiteSpace(value)) + { + throw new ArgumentException($"{paramName} cannot be null or empty", paramName); + } + } + + #endregion + + #region IDisposable + + public void Dispose() + { + if (_disposed) + return; + + _client = null; + _connectionString = null; + _disposed = true; + + GC.SuppressFinalize(this); + } + + #endregion +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/src/Services/IDocumentDbService.cs b/tools/Azure.Mcp.Tools.DocumentDb/src/Services/IDocumentDbService.cs new file mode 100644 index 0000000000..3383626944 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/src/Services/IDocumentDbService.cs @@ -0,0 +1,49 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using MongoDB.Bson; +using MongoDB.Driver; + +namespace Azure.Mcp.Tools.DocumentDb.Services; + +public interface IDocumentDbService : IDisposable +{ + // Connection Management + Task ConnectAsync(string connectionString, bool testConnection = true, CancellationToken cancellationToken = default); + Task DisconnectAsync(CancellationToken cancellationToken = default); + object GetConnectionStatus(); + + // Database Operations + Task> ListDatabasesAsync(CancellationToken cancellationToken = default); + Task GetDatabaseStatsAsync(string databaseName, CancellationToken cancellationToken = default); + Task GetDatabaseInfoAsync(string databaseName, CancellationToken cancellationToken = default); + Task DropDatabaseAsync(string databaseName, CancellationToken cancellationToken = default); + + // Collection Operations + Task GetCollectionStatsAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default); + Task RenameCollectionAsync(string databaseName, string oldName, string newName, CancellationToken cancellationToken = default); + Task DropCollectionAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default); + Task> SampleDocumentsAsync(string databaseName, string collectionName, int sampleSize = 10, CancellationToken cancellationToken = default); + + // Document Operations + Task FindDocumentsAsync(string databaseName, string collectionName, BsonDocument? query = null, object? options = null, CancellationToken cancellationToken = default); + Task CountDocumentsAsync(string databaseName, string collectionName, BsonDocument? query = null, CancellationToken cancellationToken = default); + Task InsertDocumentAsync(string databaseName, string collectionName, BsonDocument document, CancellationToken cancellationToken = default); + Task InsertManyAsync(string databaseName, string collectionName, List documents, CancellationToken cancellationToken = default); + Task UpdateDocumentAsync(string databaseName, string collectionName, BsonDocument filter, BsonDocument update, bool upsert = false, CancellationToken cancellationToken = default); + Task UpdateManyAsync(string databaseName, string collectionName, BsonDocument filter, BsonDocument update, bool upsert = false, CancellationToken cancellationToken = default); + Task DeleteDocumentAsync(string databaseName, string collectionName, BsonDocument filter, CancellationToken cancellationToken = default); + Task DeleteManyAsync(string databaseName, string collectionName, BsonDocument filter, CancellationToken cancellationToken = default); + Task AggregateAsync(string databaseName, string collectionName, List pipeline, bool allowDiskUse = false, CancellationToken cancellationToken = default); + Task FindAndModifyAsync(string databaseName, string collectionName, BsonDocument query, BsonDocument update, bool upsert = false, CancellationToken cancellationToken = default); + Task ExplainFindQueryAsync(string databaseName, string collectionName, BsonDocument? query = null, object? options = null, CancellationToken cancellationToken = default); + Task ExplainCountQueryAsync(string databaseName, string collectionName, BsonDocument? query = null, CancellationToken cancellationToken = default); + Task ExplainAggregateQueryAsync(string databaseName, string collectionName, List pipeline, CancellationToken cancellationToken = default); + + // Index Operations + Task CreateIndexAsync(string databaseName, string collectionName, BsonDocument keys, BsonDocument? options = null, CancellationToken cancellationToken = default); + Task ListIndexesAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default); + Task DropIndexAsync(string databaseName, string collectionName, string indexName, CancellationToken cancellationToken = default); + Task> GetIndexStatsAsync(string databaseName, string collectionName, CancellationToken cancellationToken = default); + Task GetCurrentOpsAsync(BsonDocument? filter = null, CancellationToken cancellationToken = default); +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Azure.Mcp.Tools.DocumentDb.UnitTests.csproj b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Azure.Mcp.Tools.DocumentDb.UnitTests.csproj new file mode 100644 index 0000000000..ba9891c40c --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Azure.Mcp.Tools.DocumentDb.UnitTests.csproj @@ -0,0 +1,17 @@ + + + true + Exe + + + + + + + + + + + + + diff --git a/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Collection/SampleDocumentsCommandTests.cs b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Collection/SampleDocumentsCommandTests.cs new file mode 100644 index 0000000000..ceef780051 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Collection/SampleDocumentsCommandTests.cs @@ -0,0 +1,237 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using System.Net; +using System.Text.Json; +using Azure.Mcp.Tools.DocumentDb.Commands.Collection; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using MongoDB.Bson; +using NSubstitute; +using NSubstitute.ExceptionExtensions; +using Xunit; + +namespace Azure.Mcp.Tools.DocumentDb.UnitTests.Collection; + +public class SampleDocumentsCommandTests +{ + private readonly IServiceProvider _serviceProvider; + private readonly IDocumentDbService _documentDbService; + private readonly ILogger _logger; + private readonly SampleDocumentsCommand _command; + private readonly CommandContext _context; + private readonly Command _commandDefinition; + + public SampleDocumentsCommandTests() + { + _documentDbService = Substitute.For(); + _logger = Substitute.For>(); + _command = new(_logger); + _commandDefinition = _command.GetCommand(); + _serviceProvider = new ServiceCollection() + .AddSingleton(_documentDbService) + .BuildServiceProvider(); + _context = new(_serviceProvider); + } + + [Fact] + public async Task ExecuteAsync_ReturnsSampleDocuments_WhenDocumentsExist() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var sampleSize = 5; + var expectedDocuments = new List + { + new() { { "_id", ObjectId.GenerateNewId() }, { "name", "doc1" } }, + new() { { "_id", ObjectId.GenerateNewId() }, { "name", "doc2" } } + }; + + _documentDbService.SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Is(sampleSize), + Arg.Any()) + .Returns(expectedDocuments); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--sample-size", sampleSize.ToString() + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_ReturnsEmptyList_WhenNoDocumentsExist() + { + // Arrange + var dbName = "testdb"; + var collectionName = "emptycollection"; + + _documentDbService.SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any()) + .Returns([]); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_UsesDefaultSampleSize_WhenNotProvided() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var defaultSampleSize = 10; + + _documentDbService.SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Is(defaultSampleSize), + Arg.Any()) + .Returns([]); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + await _documentDbService.Received(1).SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Is(defaultSampleSize), + Arg.Any()); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenCollectionNotFound() + { + // Arrange + var dbName = "testdb"; + var collectionName = "nonexistent"; + var expectedError = "Collection not found"; + + _documentDbService.SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenDatabaseNotFound() + { + // Arrange + var dbName = "nonexistentdb"; + var collectionName = "testcollection"; + var expectedError = "Database not found"; + + _documentDbService.SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Theory] + [InlineData("--db-name", "testdb")] + [InlineData("--collection-name", "testcollection")] + public async Task ExecuteAsync_Returns400_WhenRequiredParametersAreMissing(params string[] args) + { + // Arrange & Act + var response = await _command.ExecuteAsync(_context, _commandDefinition.Parse(args), TestContext.Current.CancellationToken); + + // Assert + Assert.Equal(HttpStatusCode.BadRequest, response.Status); + Assert.Contains("required", response.Message.ToLower()); + } + + [Theory] + [InlineData(1)] + [InlineData(100)] + [InlineData(1000)] + public async Task ExecuteAsync_HandlesVariousSampleSizes(int sampleSize) + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + + _documentDbService.SampleDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Is(sampleSize), + Arg.Any()) + .Returns([]); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--sample-size", sampleSize.ToString() + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Connection/ConnectDocumentDbCommandTests.cs b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Connection/ConnectDocumentDbCommandTests.cs new file mode 100644 index 0000000000..0c4aac4215 --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Connection/ConnectDocumentDbCommandTests.cs @@ -0,0 +1,160 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using System.Net; +using System.Text.Json; +using Azure.Mcp.Tools.DocumentDb.Commands.Connection; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using NSubstitute; +using NSubstitute.ExceptionExtensions; +using Xunit; + +namespace Azure.Mcp.Tools.DocumentDb.UnitTests.Connection; + +public class ConnectDocumentDbCommandTests +{ + private readonly IServiceProvider _serviceProvider; + private readonly IDocumentDbService _documentDbService; + private readonly ILogger _logger; + private readonly ConnectDocumentDbCommand _command; + private readonly CommandContext _context; + private readonly Command _commandDefinition; + + public ConnectDocumentDbCommandTests() + { + _documentDbService = Substitute.For(); + _logger = Substitute.For>(); + _command = new(_logger); + _commandDefinition = _command.GetCommand(); + _serviceProvider = new ServiceCollection() + .AddSingleton(_documentDbService) + .BuildServiceProvider(); + _context = new(_serviceProvider); + } + + [Fact] + public async Task ExecuteAsync_ReturnsSuccess_WhenConnectionSucceeds() + { + // Arrange + var connectionString = "mongodb://localhost:27017"; + var expectedResult = new Dictionary + { + ["success"] = true, + ["message"] = "Connected successfully", + ["data"] = new Dictionary + { + ["databaseCount"] = 2, + ["databases"] = new List { "test", "admin" } + } + }; + + _documentDbService.ConnectAsync( + Arg.Is(connectionString), + Arg.Is(true), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--connection-string", connectionString + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_ReturnsSuccess_WhenConnectionSucceedsWithoutTestFlag() + { + // Arrange + var connectionString = "mongodb://localhost:27017"; + var expectedResult = new Dictionary + { + ["success"] = true, + ["message"] = "Connected successfully (not tested)", + ["data"] = null + }; + + _documentDbService.ConnectAsync( + Arg.Is(connectionString), + Arg.Is(false), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--connection-string", connectionString, + "--test-connection", "false" + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenConnectionFails() + { + // Arrange + var connectionString = "mongodb://invalid:27017"; + var expectedError = "Failed to connect to DocumentDB"; + + _documentDbService.ConnectAsync( + Arg.Is(connectionString), + Arg.Is(true), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--connection-string", connectionString + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Fact] + public async Task ExecuteAsync_Returns400_WhenConnectionStringIsMissing() + { + // Arrange & Act + var response = await _command.ExecuteAsync(_context, _commandDefinition.Parse([]), TestContext.Current.CancellationToken); + + // Assert + Assert.Equal(HttpStatusCode.BadRequest, response.Status); + Assert.Contains("required", response.Message.ToLower()); + } + + [Theory] + [InlineData("")] + [InlineData(" ")] + public async Task ExecuteAsync_Returns400_WhenConnectionStringIsEmpty(string connectionString) + { + // Arrange + var args = _commandDefinition.Parse([ + "--connection-string", connectionString + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.Equal(HttpStatusCode.BadRequest, response.Status); + Assert.Contains("required", response.Message.ToLower()); + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Database/ListDatabasesCommandTests.cs b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Database/ListDatabasesCommandTests.cs new file mode 100644 index 0000000000..6b3fefd5bf --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Database/ListDatabasesCommandTests.cs @@ -0,0 +1,121 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using System.Net; +using System.Text.Json; +using Azure.Mcp.Tools.DocumentDb.Commands.Database; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using NSubstitute; +using NSubstitute.ExceptionExtensions; +using Xunit; + +namespace Azure.Mcp.Tools.DocumentDb.UnitTests.Database; + +public class ListDatabasesCommandTests +{ + private readonly IServiceProvider _serviceProvider; + private readonly IDocumentDbService _documentDbService; + private readonly ILogger _logger; + private readonly ListDatabasesCommand _command; + private readonly CommandContext _context; + private readonly Command _commandDefinition; + + public ListDatabasesCommandTests() + { + _documentDbService = Substitute.For(); + _logger = Substitute.For>(); + _command = new(_logger); + _commandDefinition = _command.GetCommand(); + _serviceProvider = new ServiceCollection() + .AddSingleton(_documentDbService) + .BuildServiceProvider(); + _context = new(_serviceProvider); + } + + [Fact] + public async Task ExecuteAsync_ReturnsDatabases_WhenDatabasesExist() + { + // Arrange + var expectedDatabases = new List { "database1", "database2", "admin" }; + _documentDbService.ListDatabasesAsync(Arg.Any()) + .Returns(expectedDatabases); + + var args = _commandDefinition.Parse([]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + + // Verify the service was called + await _documentDbService.Received(1).ListDatabasesAsync(Arg.Any()); + } + + [Fact] + public async Task ExecuteAsync_ReturnsEmpty_WhenNoDatabasesExist() + { + // Arrange + _documentDbService.ListDatabasesAsync(Arg.Any()) + .Returns([]); + + var args = _commandDefinition.Parse([]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + + // Verify the service was called + await _documentDbService.Received(1).ListDatabasesAsync(Arg.Any()); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenServiceThrowsException() + { + // Arrange + var expectedError = "Failed to connect to DocumentDB"; + + _documentDbService.ListDatabasesAsync(Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenServiceThrowsTimeoutException() + { + // Arrange + var expectedError = "Connection timeout"; + + _documentDbService.ListDatabasesAsync(Arg.Any()) + .ThrowsAsync(new TimeoutException(expectedError)); + + var args = _commandDefinition.Parse([]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Document/FindDocumentsCommandTests.cs b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Document/FindDocumentsCommandTests.cs new file mode 100644 index 0000000000..105a56044d --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Document/FindDocumentsCommandTests.cs @@ -0,0 +1,345 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using System.Net; +using System.Text.Json; +using Azure.Mcp.Tools.DocumentDb.Commands.Document; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using MongoDB.Bson; +using NSubstitute; +using NSubstitute.ExceptionExtensions; +using Xunit; + +namespace Azure.Mcp.Tools.DocumentDb.UnitTests.Document; + +public class FindDocumentsCommandTests +{ + private readonly IServiceProvider _serviceProvider; + private readonly IDocumentDbService _documentDbService; + private readonly ILogger _logger; + private readonly FindDocumentsCommand _command; + private readonly CommandContext _context; + private readonly Command _commandDefinition; + + public FindDocumentsCommandTests() + { + _documentDbService = Substitute.For(); + _logger = Substitute.For>(); + _command = new(_logger); + _commandDefinition = _command.GetCommand(); + _serviceProvider = new ServiceCollection() + .AddSingleton(_documentDbService) + .BuildServiceProvider(); + _context = new(_serviceProvider); + } + + [Fact] + public async Task ExecuteAsync_FindsDocuments_WhenQueryIsProvided() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var query = "{\"status\": \"active\"}"; + var expectedResult = new Dictionary + { + ["documents"] = new List { "{\"_id\":\"1\",\"status\":\"active\"}" }, + ["total_count"] = 1L, + ["returned_count"] = 1, + ["has_more"] = false, + ["query"] = "{\"status\":\"active\"}", + ["applied_options"] = new Dictionary + { + ["limit"] = 100, + ["skip"] = 0, + ["sort"] = null, + ["projection"] = null + } + }; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--query", query + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_FindsAllDocuments_WhenNoQueryProvided() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedResult = new Dictionary + { + ["documents"] = new List { "{\"_id\":\"1\"}", "{\"_id\":\"2\"}" }, + ["total_count"] = 2L, + ["returned_count"] = 2, + ["has_more"] = false, + ["query"] = "{}", + ["applied_options"] = new Dictionary + { + ["limit"] = 100, + ["skip"] = 0, + ["sort"] = null, + ["projection"] = null + } + }; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Is(x => x == null), + Arg.Any(), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_FindsDocuments_WhenOptionsProvided() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var query = "{\"status\": \"active\"}"; + var options = "{\"limit\": 10, \"skip\": 5}"; + var expectedResult = new Dictionary + { + ["documents"] = new List { "{\"_id\":\"1\"}" }, + ["total_count"] = 15L, + ["returned_count"] = 1, + ["has_more"] = true, + ["query"] = "{\"status\":\"active\"}", + ["applied_options"] = new Dictionary + { + ["limit"] = 10, + ["skip"] = 5, + ["sort"] = null, + ["projection"] = null + } + }; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--query", query, + "--options", options + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenInvalidJsonInQuery() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var invalidQuery = "{invalid json}"; + var expectedError = "Invalid JSON in query"; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--query", invalidQuery + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenCollectionNotFound() + { + // Arrange + var dbName = "testdb"; + var collectionName = "nonexistent"; + var expectedError = "Collection not found"; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Fact] + public async Task ExecuteAsync_ReturnsEmptyResult_WhenNoDocumentsMatch() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var query = "{\"status\": \"nonexistent\"}"; + var expectedResult = new Dictionary + { + ["documents"] = new List(), + ["total_count"] = 0L, + ["returned_count"] = 0, + ["has_more"] = false, + ["query"] = "{\"status\":\"nonexistent\"}", + ["applied_options"] = new Dictionary + { + ["limit"] = 100, + ["skip"] = 0, + ["sort"] = null, + ["projection"] = null + } + }; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--query", query + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Theory] + [InlineData("--db-name", "testdb")] + [InlineData("--collection-name", "testcollection")] + public async Task ExecuteAsync_Returns400_WhenRequiredParametersAreMissing(params string[] args) + { + // Arrange & Act + var response = await _command.ExecuteAsync(_context, _commandDefinition.Parse(args), TestContext.Current.CancellationToken); + + // Assert + Assert.Equal(HttpStatusCode.BadRequest, response.Status); + Assert.Contains("required", response.Message.ToLower()); + } + + [Theory] + [InlineData("{\"name\": \"test\"}")] + [InlineData("{\"age\": {\"$gt\": 18}}")] + [InlineData("{\"$and\": [{\"status\": \"active\"}, {\"verified\": true}]}")] + public async Task ExecuteAsync_HandlesVariousQueryFormats(string query) + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedResult = new Dictionary + { + ["documents"] = new List(), + ["total_count"] = 0L, + ["returned_count"] = 0, + ["has_more"] = false, + ["query"] = query, + ["applied_options"] = new Dictionary + { + ["limit"] = 100, + ["skip"] = 0, + ["sort"] = null, + ["projection"] = null + } + }; + + _documentDbService.FindDocumentsAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any(), + Arg.Any(), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName, + "--query", query + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + } +} diff --git a/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Index/ListIndexesCommandTests.cs b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Index/ListIndexesCommandTests.cs new file mode 100644 index 0000000000..b872fd946d --- /dev/null +++ b/tools/Azure.Mcp.Tools.DocumentDb/tests/Azure.Mcp.Tools.DocumentDb.UnitTests/Index/ListIndexesCommandTests.cs @@ -0,0 +1,278 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +using System.CommandLine; +using System.Net; +using System.Text.Json; +using Azure.Mcp.Tools.DocumentDb.Commands.Index; +using Azure.Mcp.Tools.DocumentDb.Services; +using Microsoft.Extensions.DependencyInjection; +using Microsoft.Extensions.Logging; +using Microsoft.Mcp.Core.Models.Command; +using NSubstitute; +using NSubstitute.ExceptionExtensions; +using Xunit; + +namespace Azure.Mcp.Tools.DocumentDb.UnitTests.Index; + +public class ListIndexesCommandTests +{ + private readonly IServiceProvider _serviceProvider; + private readonly IDocumentDbService _documentDbService; + private readonly ILogger _logger; + private readonly ListIndexesCommand _command; + private readonly CommandContext _context; + private readonly Command _commandDefinition; + + public ListIndexesCommandTests() + { + _documentDbService = Substitute.For(); + _logger = Substitute.For>(); + _command = new(_logger); + _commandDefinition = _command.GetCommand(); + _serviceProvider = new ServiceCollection() + .AddSingleton(_documentDbService) + .BuildServiceProvider(); + _context = new(_serviceProvider); + } + + [Fact] + public async Task ExecuteAsync_ReturnsIndexes_WhenIndexesExist() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedResult = new Dictionary + { + ["indexes"] = new List + { + "{\"name\":\"_id_\",\"key\":{\"_id\":1}}", + "{\"name\":\"status_1\",\"key\":{\"status\":1}}" + }, + ["count"] = 2 + }; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_ReturnsDefaultIndex_WhenNoCustomIndexesExist() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedResult = new Dictionary + { + ["indexes"] = new List + { + "{\"name\":\"_id_\",\"key\":{\"_id\":1}}" + }, + ["count"] = 1 + }; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenCollectionNotFound() + { + // Arrange + var dbName = "testdb"; + var collectionName = "nonexistent"; + var expectedError = "Collection not found"; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenDatabaseNotFound() + { + // Arrange + var dbName = "nonexistentdb"; + var collectionName = "testcollection"; + var expectedError = "Database not found"; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .ThrowsAsync(new Exception(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } + + [Theory] + [InlineData("--db-name", "testdb")] + [InlineData("--collection-name", "testcollection")] + public async Task ExecuteAsync_Returns400_WhenRequiredParametersAreMissing(params string[] args) + { + // Arrange & Act + var response = await _command.ExecuteAsync(_context, _commandDefinition.Parse(args), TestContext.Current.CancellationToken); + + // Assert + Assert.Equal(HttpStatusCode.BadRequest, response.Status); + Assert.Contains("required", response.Message.ToLower()); + } + + [Fact] + public async Task ExecuteAsync_ReturnsCompoundIndexes_WhenTheyExist() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedResult = new Dictionary + { + ["indexes"] = new List + { + "{\"name\":\"_id_\",\"key\":{\"_id\":1}}", + "{\"name\":\"name_1_age_1\",\"key\":{\"name\":1,\"age\":1}}" + }, + ["count"] = 2 + }; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_ReturnsTextIndexes_WhenTheyExist() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedResult = new Dictionary + { + ["indexes"] = new List + { + "{\"name\":\"_id_\",\"key\":{\"_id\":1}}", + "{\"name\":\"description_text\",\"key\":{\"description\":\"text\"}}" + }, + ["count"] = 2 + }; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .Returns(expectedResult); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.OK, response.Status); + Assert.NotNull(response.Results); + } + + [Fact] + public async Task ExecuteAsync_Returns500_WhenServiceThrowsUnauthorizedException() + { + // Arrange + var dbName = "testdb"; + var collectionName = "testcollection"; + var expectedError = "Unauthorized access"; + + _documentDbService.ListIndexesAsync( + Arg.Is(dbName), + Arg.Is(collectionName), + Arg.Any()) + .ThrowsAsync(new UnauthorizedAccessException(expectedError)); + + var args = _commandDefinition.Parse([ + "--db-name", dbName, + "--collection-name", collectionName + ]); + + // Act + var response = await _command.ExecuteAsync(_context, args, TestContext.Current.CancellationToken); + + // Assert + Assert.NotNull(response); + Assert.Equal(HttpStatusCode.InternalServerError, response.Status); + Assert.StartsWith(expectedError, response.Message); + } +}