Skip to content

Add Cosmos DB container copy job MCP tools #90

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
niteshvijay1995 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add Cosmos DB container copy job MCP tools #90

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 126 additions & 1 deletion src/AzureCosmosDB.MCP.Toolkit/Controllers/MCPProtocolController.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,18 @@ namespace AzureCosmosDB.MCP.Toolkit.Controllers;
public class MCPProtocolController : ControllerBase
{
private readonly CosmosDbToolsService _cosmosDbTools;
private readonly CopyJobService _copyJobService;
private readonly AuthenticationService _authService;
private readonly ILogger<MCPProtocolController> _logger;

public MCPProtocolController(
CosmosDbToolsService cosmosDbTools,
CosmosDbToolsService cosmosDbTools,
CopyJobService copyJobService,
AuthenticationService authService,
ILogger<MCPProtocolController> logger)
{
_cosmosDbTools = cosmosDbTools;
_copyJobService = copyJobService;
_authService = authService;
_logger = logger;
}
Expand Down Expand Up @@ -232,6 +235,93 @@ public async Task<IActionResult> HandleMCPRequest([FromBody] JsonElement request
},
required = new string[] { "databaseId", "containerId", "searchText", "vectorProperty", "selectProperties", "topN" }
}
},
// Copy Job tools (ARM management plane)
new {
name = "create_copy_job",
description = "Creates a Cosmos DB container copy job. Supports NoSQL, Cassandra, and Mongo copy types, as well as backup/restore to Azure Blob Storage.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" },
jobName = new { type = "string", description = "Unique name for the copy job" },
jobProperties = new { type = "string", description = "JSON string describing copy source/destination. Example for NoSQL: {\"jobType\":\"NoSqlRUToNoSqlRU\",\"source\":{\"databaseName\":\"srcDb\",\"containerName\":\"srcContainer\"},\"destination\":{\"databaseName\":\"destDb\",\"containerName\":\"destContainer\"}}" },
mode = new { type = "string", description = "Copy mode: 'Offline' (default) or 'Online' (continuous)" },
workerCount = new { type = "integer", description = "Number of worker threads (default determined by service)" }
},
required = new string[] { "subscriptionId", "jobName", "jobProperties" }
}
},
new {
name = "get_copy_job",
description = "Gets the status and details of a specific Cosmos DB copy job.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" },
jobName = new { type = "string", description = "Name of the copy job to retrieve" }
},
required = new string[] { "subscriptionId", "jobName" }
}
},
new {
name = "list_copy_jobs",
description = "Lists all Cosmos DB copy jobs for the current account.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" }
},
required = new string[] { "subscriptionId" }
}
},
new {
name = "cancel_copy_job",
description = "Cancels a running Cosmos DB copy job.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" },
jobName = new { type = "string", description = "Name of the copy job to cancel" }
},
required = new string[] { "subscriptionId", "jobName" }
}
},
new {
name = "pause_copy_job",
description = "Pauses a running Cosmos DB copy job.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" },
jobName = new { type = "string", description = "Name of the copy job to pause" }
},
required = new string[] { "subscriptionId", "jobName" }
}
},
new {
name = "resume_copy_job",
description = "Resumes a paused Cosmos DB copy job.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" },
jobName = new { type = "string", description = "Name of the copy job to resume" }
},
required = new string[] { "subscriptionId", "jobName" }
}
},
new {
name = "complete_copy_job",
description = "Completes an Online Cosmos DB copy job, flushing remaining changes from source to destination.",
inputSchema = new {
type = "object",
properties = new {
subscriptionId = new { type = "string", description = "Azure subscription ID" },
jobName = new { type = "string", description = "Name of the Online copy job to complete" }
},
required = new string[] { "subscriptionId", "jobName" }
}
}
}
}
Expand Down Expand Up @@ -395,6 +485,41 @@ private async Task<object> ExecuteTool(string toolName, Dictionary<string, objec
GetStringArg(args, "selectProperties"),
GetRequiredIntArg(args, "topN"),
cancellationToken),
// Copy Job tools
"create_copy_job" => await _copyJobService.CreateCopyJob(
GetStringArg(args, "subscriptionId"),
GetStringArg(args, "jobName"),
GetStringArg(args, "jobProperties"),
GetStringArg(args, "mode"),
args.ContainsKey("workerCount") ? GetIntArg(args, "workerCount") : null,
cancellationToken),
"get_copy_job" => await _copyJobService.GetCopyJob(
GetStringArg(args, "subscriptionId"),
GetStringArg(args, "jobName"),
cancellationToken),
"list_copy_jobs" => await _copyJobService.ListCopyJobs(
GetStringArg(args, "subscriptionId"),
cancellationToken),
"cancel_copy_job" => await _copyJobService.CopyJobAction(
GetStringArg(args, "subscriptionId"),
GetStringArg(args, "jobName"),
"cancel",
cancellationToken),
"pause_copy_job" => await _copyJobService.CopyJobAction(
GetStringArg(args, "subscriptionId"),
GetStringArg(args, "jobName"),
"pause",
cancellationToken),
"resume_copy_job" => await _copyJobService.CopyJobAction(
GetStringArg(args, "subscriptionId"),
GetStringArg(args, "jobName"),
"resume",
cancellationToken),
"complete_copy_job" => await _copyJobService.CopyJobAction(
GetStringArg(args, "subscriptionId"),
GetStringArg(args, "jobName"),
"complete",
cancellationToken),
_ => throw new ArgumentException($"Unknown tool: {toolName}")
};
}
Expand Down
2 changes: 2 additions & 0 deletions src/AzureCosmosDB.MCP.Toolkit/Program.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -232,6 +232,8 @@
// Register services for dependency injection
builder.Services.AddScoped<AzureCosmosDB.MCP.Toolkit.Services.CosmosDbToolsService>();
builder.Services.AddScoped<AzureCosmosDB.MCP.Toolkit.Services.AuthenticationService>();
builder.Services.AddHttpClient();
builder.Services.AddScoped<AzureCosmosDB.MCP.Toolkit.Services.CopyJobService>();

// Configure forwarded headers for proxy scenarios
builder.Services.Configure<ForwardedHeadersOptions>(options =>
Expand Down
Loading