diff --git a/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs b/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs new file mode 100644 index 00000000..236e8832 --- /dev/null +++ b/CosmosDBShell.Tests/CommandTests/BatchCommandTests.cs @@ -0,0 +1,457 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.CommandTests; + +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Commands; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; +using Azure.Data.Cosmos.Shell.States; +using Azure.Data.Cosmos.Shell.Util; +using Microsoft.Azure.Cosmos; + +/// +/// Unit tests for . These cover the pure parsing and +/// validation logic, which can be exercised without a live Cosmos DB connection. +/// +public class BatchCommandTests +{ + [Fact] + public async Task Status_WithoutActiveBatch_ReturnsInactive() + { + using var shell = ShellInterpreter.CreateInstance(); + var command = new BatchCommand { Subcommand = "status" }; + + var state = await command.ExecuteAsync(shell, new CommandState(), "batch status", CancellationToken.None); + + var json = Assert.IsType(state.Result!.ConvertShellObject(Azure.Data.Cosmos.Shell.Parser.DataType.Json)); + Assert.False(json.GetProperty("active").GetBoolean()); + Assert.NotNull(state.RenderUser); + } + + [Fact] + public async Task Show_WithoutActiveBatch_ReturnsEmptyArray() + { + using var shell = ShellInterpreter.CreateInstance(); + var command = new BatchCommand { Subcommand = "show" }; + + var state = await command.ExecuteAsync(shell, new CommandState(), "batch show", CancellationToken.None); + + var json = Assert.IsType(state.Result!.ConvertShellObject(Azure.Data.Cosmos.Shell.Parser.DataType.Json)); + Assert.Equal(JsonValueKind.Array, json.ValueKind); + Assert.Empty(json.EnumerateArray()); + } + + [Fact] + public async Task Status_WithActiveBatch_ReturnsBatchDetails() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.CurrentBatch = CreatePendingBatch(); + shell.CurrentBatch.Operations.AddRange(BatchOperationParser.Parse( + "batch", + "[{\"op\":\"create\",\"item\":{\"id\":\"1\"}},{\"op\":\"delete\",\"id\":\"2\"}]")); + var command = new BatchCommand { Subcommand = "status" }; + + var state = await command.ExecuteAsync(shell, new CommandState(), "batch status", CancellationToken.None); + + var json = Assert.IsType(state.Result!.ConvertShellObject(Azure.Data.Cosmos.Shell.Parser.DataType.Json)); + Assert.True(json.GetProperty("active").GetBoolean()); + Assert.Equal("TestDatabase", json.GetProperty("database").GetString()); + Assert.Equal("TestContainer", json.GetProperty("container").GetString()); + Assert.Equal("tenant-1", json.GetProperty("partitionKey").GetString()); + Assert.Equal(2, json.GetProperty("operationCount").GetInt32()); + Assert.Equal("1", json.GetProperty("operations")[0].GetProperty("id").GetString()); + Assert.Equal("delete", json.GetProperty("operations")[1].GetProperty("op").GetString()); + Assert.NotNull(state.RenderUser); + } + + [Fact] + public async Task Show_WithActiveBatch_ReturnsOriginalOperations() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.CurrentBatch = CreatePendingBatch(); + shell.CurrentBatch.Operations.AddRange(BatchOperationParser.Parse( + "batch", + "[{\"op\":\"create\",\"item\":{\"id\":\"1\",\"name\":\"Ada\"}},{\"op\":\"delete\",\"id\":\"2\"}]")); + var command = new BatchCommand { Subcommand = "show" }; + + var state = await command.ExecuteAsync(shell, new CommandState(), "batch show", CancellationToken.None); + + var json = Assert.IsType(state.Result!.ConvertShellObject(Azure.Data.Cosmos.Shell.Parser.DataType.Json)); + Assert.Equal(2, json.GetArrayLength()); + Assert.Equal("Ada", json[0].GetProperty("item").GetProperty("name").GetString()); + Assert.Equal("2", json[1].GetProperty("id").GetString()); + Assert.Null(state.RenderUser); + } + + [Fact] + public async Task Add_WithData_QueuesOperations() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.CurrentBatch = CreatePendingBatch(); + var command = new BatchCommand + { + Subcommand = "add", + Data = "[{\"op\":\"create\",\"item\":{\"id\":\"1\"}},{\"op\":\"delete\",\"id\":\"2\"}]", + }; + + var state = await command.ExecuteAsync(shell, new CommandState(), "batch add", CancellationToken.None); + + Assert.Equal(2, shell.CurrentBatch.Operations.Count); + Assert.NotNull(state.RenderUser); + } + + [Fact] + public async Task Add_WithPipedJson_QueuesOperation() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.CurrentBatch = CreatePendingBatch(); + var command = new BatchCommand { Subcommand = "add" }; + var input = new CommandState + { + Result = new ShellJson(JsonSerializer.SerializeToElement(new + { + op = "delete", + id = "piped-item", + })), + }; + + await command.ExecuteAsync(shell, input, "batch add", CancellationToken.None); + + var operation = Assert.Single(shell.CurrentBatch.Operations); + Assert.Equal("piped-item", operation.Id); + } + + [Theory] + [InlineData("cancel")] + [InlineData("abort")] + public async Task Cancel_WithActiveBatch_ClearsBatch(string subcommand) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.CurrentBatch = CreatePendingBatch(); + var command = new BatchCommand { Subcommand = subcommand }; + + var state = await command.ExecuteAsync(shell, new CommandState(), $"batch {subcommand}", CancellationToken.None); + + Assert.Null(shell.CurrentBatch); + Assert.NotNull(state.RenderUser); + } + + [Theory] + [InlineData("add")] + [InlineData("cancel")] + public async Task StatefulCommand_WithoutActiveBatch_Throws(string subcommand) + { + using var shell = ShellInterpreter.CreateInstance(); + var command = new BatchCommand + { + Subcommand = subcommand, + Data = "{\"op\":\"delete\",\"id\":\"1\"}", + }; + + var exception = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), $"batch {subcommand}", CancellationToken.None)); + + Assert.Equal(MessageService.GetString("command-batch-error-not_active"), exception.Message); + } + + [Theory] + [InlineData("")] + [InlineData("unknown")] + public async Task InvalidSubcommand_Throws(string subcommand) + { + using var shell = ShellInterpreter.CreateInstance(); + var command = new BatchCommand { Subcommand = subcommand }; + + await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "batch", CancellationToken.None)); + } + + [Theory] + [InlineData("run")] + [InlineData("begin")] + [InlineData("execute")] + [InlineData("exec")] + [InlineData("commit")] + public async Task ConnectedCommand_WhenDisconnected_Throws(string subcommand) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DisconnectedState(); + var command = new BatchCommand { Subcommand = subcommand }; + + await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), $"batch {subcommand}", CancellationToken.None)); + } + + [Fact] + public async Task Begin_WhenBatchAlreadyActive_Throws() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new ConnectedState(CreateTestClient()); + shell.CurrentBatch = CreatePendingBatch(); + var command = new BatchCommand { Subcommand = "begin", PartitionKeyArgument = "tenant-1" }; + + var exception = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "batch begin", CancellationToken.None)); + + Assert.Equal(MessageService.GetString("command-batch-error-already_active"), exception.Message); + } + + [Theory] + [InlineData("run")] + [InlineData("begin")] + public async Task StartCommand_WithoutPartitionKey_Throws(string subcommand) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new ConnectedState(CreateTestClient()); + var command = new BatchCommand { Subcommand = subcommand }; + + var exception = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), $"batch {subcommand}", CancellationToken.None)); + + Assert.Equal(MessageService.GetString("command-batch-error-missing_pk"), exception.Message); + } + + [Fact] + public async Task Execute_WithEmptyActiveBatch_Throws() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new ConnectedState(CreateTestClient()); + shell.CurrentBatch = CreatePendingBatch(); + var command = new BatchCommand { Subcommand = "execute" }; + + var exception = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "batch execute", CancellationToken.None)); + + Assert.Equal(MessageService.GetString("command-batch-error-empty"), exception.Message); + } + + [Fact] + public async Task Add_WhenBatchWouldExceedMaximum_Throws() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.CurrentBatch = CreatePendingBatch(); + var existingOperation = BatchOperationParser.Parse("batch", "{\"op\":\"delete\",\"id\":\"1\"}")[0]; + shell.CurrentBatch.Operations.AddRange(Enumerable.Repeat(existingOperation, BatchExecutor.MaxOperations)); + var command = new BatchCommand + { + Subcommand = "add", + Data = "{\"op\":\"delete\",\"id\":\"101\"}", + }; + + await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "batch add", CancellationToken.None)); + } + + [Fact] + public async Task Execute_WithoutOperations_Throws() + { + await Assert.ThrowsAsync( + () => BatchExecutor.ExecuteAsync("batch", null!, default, [], CancellationToken.None)); + } + + [Fact] + public async Task Execute_OverMaximumOperations_Throws() + { + var operation = BatchOperationParser.Parse("batch", "{\"op\":\"delete\",\"id\":\"1\"}")[0]; + var operations = Enumerable.Repeat(operation, BatchExecutor.MaxOperations + 1).ToArray(); + + await Assert.ThrowsAsync( + () => BatchExecutor.ExecuteAsync("batch", null!, default, operations, CancellationToken.None)); + } + + [Theory] + [InlineData(1, "1 operation")] + [InlineData(2, "2 operations")] + public void BatchMessages_PluralizeNumericOperationCounts(int count, string expected) + { + var message = MessageService.GetArgsString("command-batch-cancelled", "count", count); + + Assert.Contains(expected, message); + } + + [Fact] + public void Parse_Array_ReturnsAllOperationsInOrder() + { + var specs = BatchOperationParser.Parse( + "batch", + "[{\"op\":\"create\",\"item\":{\"id\":\"1\"}},{\"op\":\"delete\",\"id\":\"2\"}]"); + + Assert.Equal(2, specs.Count); + Assert.Equal(BatchOperationKind.Create, specs[0].Kind); + Assert.Equal(BatchOperationKind.Delete, specs[1].Kind); + Assert.Equal("2", specs[1].Id); + } + + [Fact] + public void Parse_SingleObject_ReturnsOneOperation() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"upsert\",\"item\":{\"id\":\"7\"}}"); + + Assert.Single(specs); + Assert.Equal(BatchOperationKind.Upsert, specs[0].Kind); + Assert.NotNull(specs[0].Item); + } + + [Fact] + public void Parse_Create_ExtractsIdFromItem() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"create\",\"item\":{\"id\":\"abc\",\"name\":\"x\"}}"); + + Assert.Equal("abc", specs[0].Id); + } + + [Fact] + public void Parse_Create_AllowsMissingId() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"create\",\"item\":{\"name\":\"x\"}}"); + + Assert.Null(specs[0].Id); + Assert.Equal(BatchOperationKind.Create, specs[0].Kind); + } + + [Fact] + public void Parse_Replace_UsesExplicitIdOverItemId() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"replace\",\"id\":\"explicit\",\"item\":{\"id\":\"inner\"}}"); + + Assert.Equal("explicit", specs[0].Id); + Assert.Equal(BatchOperationKind.Replace, specs[0].Kind); + } + + [Fact] + public void Parse_Replace_FallsBackToItemId() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"replace\",\"item\":{\"id\":\"inner\"}}"); + + Assert.Equal("inner", specs[0].Id); + } + + [Fact] + public void Parse_Replace_MissingId_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"replace\",\"item\":{\"name\":\"x\"}}")); + } + + [Fact] + public void Parse_Delete_MissingId_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"delete\"}")); + } + + [Fact] + public void Parse_Patch_BuildsPatchOperations() + { + var specs = BatchOperationParser.Parse( + "batch", + "{\"op\":\"patch\",\"id\":\"1\",\"operations\":[{\"op\":\"set\",\"path\":\"/name\",\"value\":\"x\"},{\"op\":\"incr\",\"path\":\"/n\",\"value\":2}]}"); + + Assert.Equal(BatchOperationKind.Patch, specs[0].Kind); + Assert.Equal("1", specs[0].Id); + Assert.NotNull(specs[0].PatchOperations); + Assert.Equal(2, specs[0].PatchOperations!.Count); + } + + [Fact] + public void Parse_Patch_MissingOperations_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"patch\",\"id\":\"1\",\"operations\":[]}")); + } + + [Fact] + public void Parse_Patch_InvalidOperationEntry_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"patch\",\"id\":\"1\",\"operations\":[{\"op\":\"set\"}]}")); + } + + [Fact] + public void Parse_MissingItem_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"create\"}")); + } + + [Fact] + public void Parse_ItemNotObject_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"create\",\"item\":\"not-an-object\"}")); + } + + [Fact] + public void Parse_MissingOp_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"item\":{\"id\":\"1\"}}")); + } + + [Fact] + public void Parse_UnsupportedOp_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{\"op\":\"merge\",\"item\":{\"id\":\"1\"}}")); + } + + [Fact] + public void Parse_InvalidJson_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "{not json")); + } + + [Fact] + public void Parse_NonObjectJson_Throws() + { + Assert.Throws(() => + BatchOperationParser.Parse("batch", "42")); + } + + [Fact] + public void Parse_NumericId_IsPreservedAsString() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"delete\",\"id\":5}"); + + Assert.Equal("5", specs[0].Id); + } + + [Fact] + public void Parse_ClonedItem_SurvivesSourceDocumentDisposal() + { + var specs = BatchOperationParser.Parse("batch", "{\"op\":\"create\",\"item\":{\"id\":\"1\",\"name\":\"Ada\"}}"); + + // The parser disposes the source JsonDocument internally; the cloned item must remain valid. + Assert.Equal("Ada", specs[0].Item!.Value.GetProperty("name").GetString()); + } + + [Fact] + public void Parse_RawOperation_PreservesOriginalOperationJson() + { + var specs = BatchOperationParser.Parse( + "batch", + "[{\"op\":\"create\",\"item\":{\"id\":\"1\"}},{\"op\":\"patch\",\"id\":\"1\",\"operations\":[{\"op\":\"set\",\"path\":\"/n\",\"value\":1}]}]"); + + Assert.Equal("create", specs[0].RawOperation.GetProperty("op").GetString()); + Assert.Equal("1", specs[0].RawOperation.GetProperty("item").GetProperty("id").GetString()); + Assert.Equal("patch", specs[1].RawOperation.GetProperty("op").GetString()); + Assert.Equal(1, specs[1].RawOperation.GetProperty("operations").GetArrayLength()); + } + + private static PendingBatchState CreatePendingBatch() + { + return new PendingBatchState("TestDatabase", "TestContainer", "tenant-1", new PartitionKey("tenant-1")); + } + + private static CosmosClient CreateTestClient() + { + return new CosmosClient( + "https://localhost:8081", + Convert.ToBase64String(new byte[64]), + new CosmosClientOptions { ConnectionMode = ConnectionMode.Gateway }); + } +} diff --git a/CosmosDBShell.Tests/Integration/BatchOperationTests.cs b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs new file mode 100644 index 00000000..758ba7cd --- /dev/null +++ b/CosmosDBShell.Tests/Integration/BatchOperationTests.cs @@ -0,0 +1,192 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.Integration; + +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Core; + +using Xunit; + +public class BatchOperationTests : EmulatorFixtureTestBase +{ + public BatchOperationTests(EmulatorDatabaseFixture fixture) + : base(fixture) + { + } + + public override async ValueTask InitializeAsync() + { + await base.InitializeAsync(); + + await ExecuteAsync("cd"); + await ExecuteAsync($"cd {Fixture.DatabaseName}"); + } + + [Fact] + public async Task BatchRun_MultipleCreates_CommitsAtomically() + { + await CreateBatchContainerAsync(); + const string pk = "tenant-success"; + var json = "[" + + "{\"op\":\"create\",\"item\":{\"id\":\"a\",\"pk\":\"" + pk + "\"}}," + + "{\"op\":\"create\",\"item\":{\"id\":\"b\",\"pk\":\"" + pk + "\"}}]"; + + var output = await ExecuteWithOutputAsync($"batch run '{json}' --partition-key {pk}"); + var root = JsonDocument.Parse(output).RootElement; + + Assert.True(root.GetProperty("success").GetBoolean()); + Assert.Equal(2, root.GetProperty("operationCount").GetInt32()); + + var itemA = JsonDocument.Parse(await ExecuteWithOutputAsync($"print a {pk}")).RootElement; + Assert.Equal("a", itemA.GetProperty("id").GetString()); + var itemB = JsonDocument.Parse(await ExecuteWithOutputAsync($"print b {pk}")).RootElement; + Assert.Equal("b", itemB.GetProperty("id").GetString()); + } + + [Fact] + public async Task BatchRun_CreateThenPatch_AppliesBoth() + { + await CreateBatchContainerAsync(); + const string pk = "tenant-patch"; + var json = "[" + + "{\"op\":\"create\",\"item\":{\"id\":\"c1\",\"pk\":\"" + pk + "\",\"status\":\"new\"}}," + + "{\"op\":\"patch\",\"id\":\"c1\",\"operations\":[{\"op\":\"set\",\"path\":\"/status\",\"value\":\"done\"}]}]"; + + var output = await ExecuteWithOutputAsync($"batch run '{json}' --partition-key {pk}"); + Assert.True(JsonDocument.Parse(output).RootElement.GetProperty("success").GetBoolean()); + + var item = JsonDocument.Parse(await ExecuteWithOutputAsync($"print c1 {pk}")).RootElement; + Assert.Equal("done", item.GetProperty("status").GetString()); + } + + [Fact] + public async Task BatchRun_FailingOperation_RollsBackEntireBatch() + { + await CreateBatchContainerAsync(); + const string pk = "tenant-rollback"; + + // Seed an existing item so the second create conflicts. + await ExecuteAsync($"mkitem '{{\"id\":\"existing\",\"pk\":\"{pk}\"}}'"); + + var json = "[" + + "{\"op\":\"create\",\"item\":{\"id\":\"fresh\",\"pk\":\"" + pk + "\"}}," + + "{\"op\":\"create\",\"item\":{\"id\":\"existing\",\"pk\":\"" + pk + "\"}}]"; + + var outputFile = CreateTempFile(); + Shell.StdOutRedirect = outputFile; + try + { + var batchState = await ExecuteAsync($"batch run '{json}' --partition-key {pk}"); + Assert.True(batchState.IsError); + Assert.Equal(ShellExitCode.GeneralFailure, batchState.ExitCode); + + var output = await File.ReadAllTextAsync(outputFile, TestContext.Current.CancellationToken); + var root = JsonDocument.Parse(output).RootElement; + Assert.False(root.GetProperty("success").GetBoolean()); + } + finally + { + Shell.StdOutRedirect = null; + } + + // The first operation must have been rolled back: 'fresh' should not exist. + var printState = await ExecuteAsync($"print fresh {pk}"); + Assert.True(printState.IsError); + } + + [Fact] + public async Task StatefulBatch_BeginAddExecute_CommitsQueuedOperations() + { + await CreateBatchContainerAsync(); + const string pk = "tenant-stateful"; + + await ExecuteAsync($"batch begin --partition-key {pk}"); + await ExecuteAsync($"batch add '{{\"op\":\"create\",\"item\":{{\"id\":\"s1\",\"pk\":\"{pk}\"}}}}'"); + await ExecuteAsync($"batch add '{{\"op\":\"patch\",\"id\":\"s1\",\"operations\":[{{\"op\":\"set\",\"path\":\"/status\",\"value\":\"done\"}}]}}'"); + + var statusOutput = await ExecuteWithOutputAsync("batch status"); + var status = JsonDocument.Parse(statusOutput).RootElement; + Assert.True(status.GetProperty("active").GetBoolean()); + Assert.Equal(2, status.GetProperty("operationCount").GetInt32()); + + var execOutput = await ExecuteWithOutputAsync("batch execute"); + Assert.True(JsonDocument.Parse(execOutput).RootElement.GetProperty("success").GetBoolean()); + + var item = JsonDocument.Parse(await ExecuteWithOutputAsync($"print s1 {pk}")).RootElement; + Assert.Equal("done", item.GetProperty("status").GetString()); + + var afterStatus = JsonDocument.Parse(await ExecuteWithOutputAsync("batch status")).RootElement; + Assert.False(afterStatus.GetProperty("active").GetBoolean()); + } + + [Fact] + public async Task StatefulBatch_Cancel_DiscardsQueuedOperations() + { + await CreateBatchContainerAsync(); + const string pk = "tenant-cancel"; + + await ExecuteAsync($"batch begin --partition-key {pk}"); + await ExecuteAsync($"batch add '{{\"op\":\"create\",\"item\":{{\"id\":\"x1\",\"pk\":\"{pk}\"}}}}'"); + await ExecuteAsync("batch cancel"); + + var status = JsonDocument.Parse(await ExecuteWithOutputAsync("batch status")).RootElement; + Assert.False(status.GetProperty("active").GetBoolean()); + + var printState = await ExecuteAsync($"print x1 {pk}"); + Assert.True(printState.IsError); + } + + [Fact] + public async Task StatefulBatch_Show_PrintsQueuedOperationsAsJsonArray() + { + await CreateBatchContainerAsync(); + const string pk = "tenant-show"; + + await ExecuteAsync($"batch begin --partition-key {pk}"); + await ExecuteAsync($"batch add '{{\"op\":\"create\",\"item\":{{\"id\":\"sh1\",\"pk\":\"{pk}\"}}}}'"); + await ExecuteAsync($"batch add '{{\"op\":\"delete\",\"id\":\"sh2\"}}'"); + + var output = await ExecuteWithOutputAsync("batch show"); + var array = JsonDocument.Parse(output).RootElement; + + Assert.Equal(JsonValueKind.Array, array.ValueKind); + Assert.Equal(2, array.GetArrayLength()); + Assert.Equal("create", array[0].GetProperty("op").GetString()); + Assert.Equal("sh1", array[0].GetProperty("item").GetProperty("id").GetString()); + Assert.Equal("delete", array[1].GetProperty("op").GetString()); + Assert.Equal("sh2", array[1].GetProperty("id").GetString()); + + await ExecuteAsync("batch cancel"); + } + + [Fact] + public async Task BatchAdd_WithoutBegin_ReturnsError() + { + await CreateBatchContainerAsync(); + + var state = await ExecuteAsync("batch add '{\"op\":\"create\",\"item\":{\"id\":\"z\",\"pk\":\"p\"}}'"); + + Assert.True(state.IsError); + } + + [Fact] + public async Task BatchExecute_WithoutBegin_ReturnsError() + { + await CreateBatchContainerAsync(); + + var state = await ExecuteAsync("batch execute"); + + Assert.True(state.IsError); + } + + private async Task CreateBatchContainerAsync() + { + var name = $"batch_{Guid.NewGuid():N}"; + var state = await ExecuteAsync($"mkcon {name} /pk"); + Assert.False(state.IsError, FormatError(state)); + await ExecuteAsync($"cd {name}"); + return name; + } +} diff --git a/CosmosDBShell.Tests/McpResponseFactoryTests.cs b/CosmosDBShell.Tests/McpResponseFactoryTests.cs index 4a9088bd..ba6d4d2c 100644 --- a/CosmosDBShell.Tests/McpResponseFactoryTests.cs +++ b/CosmosDBShell.Tests/McpResponseFactoryTests.cs @@ -81,6 +81,27 @@ public void CreateSuccess_WhenErrorMessageMissing_UsesFallbackError() Assert.Equal("Command execution failed.", document.RootElement.GetProperty("error").GetString()); } + [Fact] + public void CreateSuccess_StructuredError_IncludesErrorAndResult() + { + var state = new StructuredErrorCommandState( + new CommandException("batch", "Batch failed."), + new ShellJson(JsonSerializer.SerializeToElement(new + { + success = false, + statusCode = 409, + }))); + + var result = McpResponseFactory.CreateSuccess(state, new ConnectedState(null!)); + var text = Assert.IsType(Assert.Single(result.Content)).Text; + + using var document = JsonDocument.Parse(text); + Assert.True(result.IsError); + Assert.Equal("Batch failed.", document.RootElement.GetProperty("error").GetString()); + Assert.False(document.RootElement.GetProperty("result").GetProperty("success").GetBoolean()); + Assert.Equal(409, document.RootElement.GetProperty("result").GetProperty("statusCode").GetInt32()); + } + [Fact] public void CreateError_WrapsMessageWithCurrentLocation() { diff --git a/CosmosDBShell.Tests/Runtime/ShellExitCodeTests.cs b/CosmosDBShell.Tests/Runtime/ShellExitCodeTests.cs index 99c215c9..4468e8e7 100644 --- a/CosmosDBShell.Tests/Runtime/ShellExitCodeTests.cs +++ b/CosmosDBShell.Tests/Runtime/ShellExitCodeTests.cs @@ -231,4 +231,20 @@ public void FromCommandState_ErrorWithGenericException_ReturnsGeneralFailure() Assert.Equal(ShellExitCode.GeneralFailure, ShellExitCode.FromCommandState(state)); } + + [Theory] + [InlineData(HttpStatusCode.NotFound, ShellExitCode.NotFound)] + [InlineData(HttpStatusCode.TooManyRequests, ShellExitCode.Throttled)] + public void FromCommandState_StructuredErrorWithHttpStatus_ReturnsClassifiedExitCode(HttpStatusCode statusCode, int expected) + { + var exception = new CommandException( + "batch", + "Batch failed.", + new RequestFailedException((int)statusCode, "Batch failed.")); + var state = new StructuredErrorCommandState( + exception, + new ShellJson(JsonSerializer.SerializeToElement(new { success = false }))); + + Assert.Equal(expected, ShellExitCode.FromCommandState(state)); + } } diff --git a/CosmosDBShell.Tests/Shell/CosmosShellPromptTests.cs b/CosmosDBShell.Tests/Shell/CosmosShellPromptTests.cs index 52f311cb..1a08916c 100644 --- a/CosmosDBShell.Tests/Shell/CosmosShellPromptTests.cs +++ b/CosmosDBShell.Tests/Shell/CosmosShellPromptTests.cs @@ -54,4 +54,17 @@ public void GetPromptString_InDatabaseState_EscapesMarkupCharactersOnce() Assert.Contains(Markup.Escape("Db[Name]"), prompt); Assert.DoesNotContain(Markup.Escape(Markup.Escape("Db[Name]")), prompt); } + + [Fact] + public void GetPromptString_WithActiveBatch_EscapesIndicatorOnce() + { + var shell = ShellInterpreter.CreateInstance(); + shell.State = new DisconnectedState(); + shell.CurrentBatch = new PendingBatchState("TestDatabase", "TestContainer", "tenant-1", new PartitionKey("tenant-1")); + + var prompt = new CosmosShellPrompt(shell).GetPromptString(); + + Assert.Contains(Markup.Escape("[batch:0]"), prompt); + Assert.DoesNotContain(Markup.Escape(Markup.Escape("[batch:0]")), prompt); + } } diff --git a/CosmosDBShell.Tests/Shell/ExecuteCommandExceptionTests.cs b/CosmosDBShell.Tests/Shell/ExecuteCommandExceptionTests.cs index 426663a5..739d8355 100644 --- a/CosmosDBShell.Tests/Shell/ExecuteCommandExceptionTests.cs +++ b/CosmosDBShell.Tests/Shell/ExecuteCommandExceptionTests.cs @@ -4,7 +4,9 @@ namespace CosmosShell.Tests.Shell; +using System.Text.Json; using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; public class ExecuteCommandExceptionTests { @@ -13,6 +15,39 @@ private ShellInterpreter CreateInterpreter() return new ShellInterpreter(); } + [Fact] + public void PrintState_StructuredErrorInMachineMode_WritesSingleEnvelopeToStderr() + { + using var interpreter = CreateInterpreter(); + var stdoutFile = Path.GetTempFileName(); + var stderrFile = Path.GetTempFileName(); + interpreter.Options = new Program.CosmosShellOptions { Output = "json" }; + interpreter.StdOutRedirect = stdoutFile; + interpreter.ErrOutRedirect = stderrFile; + try + { + var state = new StructuredErrorCommandState( + new CommandException("batch", "Batch failed."), + new ShellJson(JsonSerializer.SerializeToElement(new { success = false, statusCode = 409 }))); + + interpreter.PrintState(state); + + Assert.Empty(File.ReadAllText(stdoutFile)); + using var document = JsonDocument.Parse(File.ReadAllText(stderrFile)); + Assert.Equal("error", document.RootElement.GetProperty("status").GetString()); + Assert.Equal("Batch failed.", document.RootElement.GetProperty("error").GetString()); + Assert.False(document.RootElement.GetProperty("result").GetProperty("success").GetBoolean()); + Assert.Equal(409, document.RootElement.GetProperty("result").GetProperty("statusCode").GetInt32()); + } + finally + { + interpreter.StdOutRedirect = null; + interpreter.ErrOutRedirect = null; + File.Delete(stdoutFile); + File.Delete(stderrFile); + } + } + [Fact] public async Task ExecuteCommandAsync_ShellException_ReturnsErrorState() { diff --git a/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs b/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs index dec8b5b0..473145d5 100644 --- a/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs +++ b/CosmosDBShell.Tests/ToolOperationsCallToolTests.cs @@ -230,6 +230,78 @@ public async Task CallTool_MissingRequiredParameter_ReturnsError() } } + [Theory] + [InlineData("query", "query")] + [InlineData("batch", "subcommand")] + public async Task CallTool_NullRequiredParameter_ReturnsMissingParameterError(string command, string parameter) + { + var tool = CreateToolOperations(); + var arguments = new Dictionary + { + [parameter] = Json("null"), + }; + + var result = await tool.CallToolHandler(CallContext(command, arguments), CancellationToken.None); + + var (isError, root, document) = ReadResult(result); + using (document) + { + Assert.True(isError); + Assert.Contains("Missing required parameter", root.GetProperty("error").GetString()); + } + } + + [Theory] + [InlineData("begin")] + [InlineData("add")] + [InlineData("execute")] + [InlineData("exec")] + [InlineData("commit")] + [InlineData("cancel")] + [InlineData("abort")] + [InlineData("status")] + [InlineData("show")] + public async Task CallTool_StatefulBatchSubcommand_ReturnsError(string subcommand) + { + var tool = CreateToolOperations(); + var arguments = new Dictionary + { + ["subcommand"] = Json($"\"{subcommand}\""), + }; + + var result = await tool.CallToolHandler(CallContext("batch", arguments), CancellationToken.None); + + var (isError, root, document) = ReadResult(result); + using (document) + { + Assert.True(isError); + var error = root.GetProperty("error").GetString(); + Assert.Contains("only the stateless 'batch run'", error); + Assert.Contains("manually in the shell", error); + } + } + + [Fact] + public async Task CallTool_WhitespaceBatchSubcommand_ReturnsMissingSubcommandError() + { + var tool = CreateToolOperations(); + var arguments = new Dictionary + { + ["subcommand"] = Json("\" \""), + }; + + var result = await tool.CallToolHandler(CallContext("batch", arguments), CancellationToken.None); + + var (isError, root, document) = ReadResult(result); + using (document) + { + Assert.True(isError); + var error = root.GetProperty("error").GetString(); + Assert.Contains("subcommand", error, StringComparison.OrdinalIgnoreCase); + Assert.DoesNotContain("only the stateless 'batch run'", error); + } + } + [Fact] public async Task CallTool_InvalidValueType_ReturnsSanitizedError() { diff --git a/CosmosDBShell.Tests/ToolOperationsTests.cs b/CosmosDBShell.Tests/ToolOperationsTests.cs index 79aa33b3..26abae04 100644 --- a/CosmosDBShell.Tests/ToolOperationsTests.cs +++ b/CosmosDBShell.Tests/ToolOperationsTests.cs @@ -170,6 +170,17 @@ public void GetTool_DoesNotAppendWarningForUnrestrictedCommands() Assert.DoesNotContain("cannot be invoked through MCP", tool.Description, StringComparison.OrdinalIgnoreCase); } + [Fact] + public void GetTool_BatchDescription_OnlyOffersStatelessRunThroughMcp() + { + var factory = new CommandRunner().Commands["batch"]; + + var tool = ToolOperations.GetTool(factory); + + Assert.Contains("MCP supports only the one-shot 'run' subcommand", tool.Description); + Assert.Contains("available only in the interactive shell", tool.Description); + } + [Fact] public void GetTool_MapsReadOnlyAnnotationHints() { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchCommand.cs new file mode 100644 index 00000000..b00351ed --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchCommand.cs @@ -0,0 +1,360 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Commands; + +using System.Globalization; +using System.Text.Json; +using System.Text.Json.Nodes; +using Azure.Data.Cosmos.Shell.Mcp; +using Azure.Data.Cosmos.Shell.Parser; +using Azure.Data.Cosmos.Shell.Util; +using global::Azure.Data.Cosmos.Shell.Core; +using global::Azure.Data.Cosmos.Shell.States; +using Spectre.Console; + +[CosmosCommand("batch")] +[CosmosExample("batch run '[{\"op\":\"create\",\"item\":{\"id\":\"1\",\"pk\":\"a\"}},{\"op\":\"delete\",\"id\":\"2\"}]' --partition-key a", Description = "Atomically apply multiple operations in a single transaction")] +[CosmosExample("batch begin --partition-key a", Description = "Start a stateful batch for partition key 'a'")] +[CosmosExample("batch add '{\"op\":\"upsert\",\"item\":{\"id\":\"3\",\"pk\":\"a\"}}'", Description = "Queue an operation onto the active batch")] +[CosmosExample("batch add '{\"op\":\"patch\",\"id\":\"3\",\"operations\":[{\"op\":\"set\",\"path\":\"/status\",\"value\":\"done\"}]}'", Description = "Queue a patch operation onto the active batch")] +[CosmosExample("batch execute", Description = "Execute the queued operations atomically")] +[CosmosExample("batch status", Description = "Show the active batch and its queued operations")] +[CosmosExample("batch show", Description = "Print the queued operations as a JSON array")] +[CosmosExample("batch cancel", Description = "Discard the active batch")] +#pragma warning disable SA1118 // Parameter should not span multiple lines +[McpAnnotation( + Title = "Batch", + Description = @" +Executes multiple write operations against a single partition key as one atomic Cosmos DB transactional batch. + +Subcommands: +- 'run --partition-key ' parses a JSON array of operations and executes them atomically in one call. +- MCP supports only the one-shot 'run' subcommand. Stateful subcommands ('begin', 'add', 'execute', 'cancel', 'status', and 'show') are available only in the interactive shell. + +Each operation is a JSON object: +- {""op"":""create"",""item"":{...}} +- {""op"":""upsert"",""item"":{...}} +- {""op"":""replace"",""id"":""1"",""item"":{...}} +- {""op"":""delete"",""id"":""3""} +- {""op"":""patch"",""id"":""1"",""operations"":[{""op"":""set"",""path"":""/name"",""value"":""x""}]} + +All operations must share the same partition key. A batch holds at least 1 and at most 100 operations. If any operation fails, the whole batch is rolled back. +", + ReadOnly = false)] +#pragma warning restore SA1118 // Parameter should not span multiple lines +internal class BatchCommand : CosmosCommand +{ + [CosmosParameter("subcommand", RequiredErrorKey = "command-batch-error-missing_subcommand")] + public string Subcommand { get; init; } = string.Empty; + + [CosmosParameter("data", IsRequired = false)] + public string? Data { get; init; } + + [CosmosOption("partition-key", "pk")] + public string? PartitionKeyArgument { get; init; } + + [CosmosOption("database", "db")] + public string? Database { get; init; } + + [CosmosOption("container", "con")] + public string? Container { get; init; } + + public override async Task ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + var subcommand = this.Subcommand.Trim().ToLowerInvariant(); + + return subcommand switch + { + "run" => await this.RunAsync(shell, commandState, token), + "begin" => await this.BeginAsync(shell, token), + "add" => this.Add(shell, commandState), + "execute" or "exec" or "commit" => await this.ExecuteBatchAsync(shell, token), + "cancel" or "abort" => Cancel(shell), + "status" => Status(shell), + "show" => Show(shell), + "" => throw new CommandException("batch", MessageService.GetString("command-batch-error-missing_subcommand")), + _ => throw new CommandException( + "batch", + MessageService.GetArgsString("command-batch-error-invalid_subcommand", "subcommand", subcommand)), + }; + } + + private static PartitionKey ParsePartitionKey(string rawValue) + { + try + { + return CreatePartitionKeyFromArgument(rawValue); + } + catch (JsonException ex) + { + throw new CommandException("batch", MessageService.GetString("command-batch-error-invalid_pk_json"), ex); + } + } + + private static CommandState Cancel(ShellInterpreter shell) + { + var batch = shell.CurrentBatch + ?? throw new CommandException("batch", MessageService.GetString("command-batch-error-not_active")); + + var count = batch.Operations.Count; + shell.CurrentBatch = null; + var message = MessageService.GetArgsString( + "command-batch-cancelled", + "count", + count); + return new CommandState { RenderUser = () => ShellInterpreter.WriteLine(message) }; + } + + private static CommandState Status(ShellInterpreter shell) + { + var batch = shell.CurrentBatch; + JsonObject root; + Action renderUser; + if (batch is null) + { + root = new JsonObject { ["active"] = false }; + renderUser = () => ShellInterpreter.WriteLine(MessageService.GetString("command-batch-status-inactive")); + } + else + { + var operations = new JsonArray(); + foreach (var operation in batch.Operations) + { + var node = new JsonObject { ["op"] = operation.Kind.ToString().ToLowerInvariant() }; + if (operation.Id is { } id) + { + node["id"] = id; + } + + operations.Add(node); + } + + root = new JsonObject + { + ["active"] = true, + ["database"] = batch.DatabaseName, + ["container"] = batch.ContainerName, + ["partitionKey"] = batch.PartitionKeyArgument, + ["operationCount"] = batch.Operations.Count, + ["operations"] = operations, + }; + renderUser = () => RenderStatus(batch); + } + + using var document = JsonDocument.Parse(root.ToJsonString()); + return new CommandState + { + Result = new ShellJson(document.RootElement.Clone()), + RenderUser = renderUser, + }; + } + + private static void RenderStatus(PendingBatchState batch) + { + var details = new Table().HideHeaders(); + details.AddColumn(string.Empty); + details.AddColumn(string.Empty); + void AddDetail(string labelKey, string value) => + details.AddRow( + Theme.FormatHelpName(Markup.Escape(MessageService.GetString(labelKey))), + Theme.FormatTableValue(Markup.Escape(value))); + + AddDetail("command-batch-status-target", $"{batch.DatabaseName}/{batch.ContainerName}"); + AddDetail("command-batch-status-partition-key", batch.PartitionKeyArgument); + AddDetail("command-batch-status-operation-count", batch.Operations.Count.ToString(CultureInfo.InvariantCulture)); + AnsiConsole.Write(details); + + if (batch.Operations.Count == 0) + { + return; + } + + var operations = new Table(); + operations.AddColumn(Theme.FormatSectionHeader(MessageService.GetString("command-batch-status-column-index"))); + operations.AddColumn(Theme.FormatSectionHeader(MessageService.GetString("command-batch-status-column-operation"))); + operations.AddColumn(Theme.FormatSectionHeader(MessageService.GetString("command-batch-status-column-id"))); + for (var index = 0; index < batch.Operations.Count; index++) + { + var operation = batch.Operations[index]; + operations.AddRow( + Theme.FormatTableValue((index + 1).ToString(CultureInfo.InvariantCulture)), + Theme.FormatTableValue(Markup.Escape(operation.Kind.ToString().ToLowerInvariant())), + Theme.FormatTableValue(Markup.Escape(operation.Id ?? string.Empty))); + } + + AnsiConsole.Write(operations); + } + + private static CommandState Show(ShellInterpreter shell) + { + var operations = new JsonArray(); + if (shell.CurrentBatch is { } batch) + { + foreach (var operation in batch.Operations) + { + operations.Add(JsonNode.Parse(operation.RawOperation.GetRawText())); + } + } + + using var document = JsonDocument.Parse(operations.ToJsonString()); + return new CommandState { Result = new ShellJson(document.RootElement.Clone()) }; + } + + private async Task RunAsync(ShellInterpreter shell, CommandState commandState, CancellationToken token) + { + if (shell.State is not ConnectedState connectedState) + { + throw new NotConnectedException("batch"); + } + + if (string.IsNullOrWhiteSpace(this.PartitionKeyArgument)) + { + throw new CommandException("batch", MessageService.GetString("command-batch-error-missing_pk")); + } + + var json = this.ResolveData(commandState); + var specs = BatchOperationParser.Parse("batch", json); + + var (_, _, container) = await ResolveContainerAsync( + connectedState.Client, + shell.State, + this.Database, + this.Container, + "batch", + token); + + var partitionKey = ParsePartitionKey(this.PartitionKeyArgument); + try + { + return await BatchExecutor.ExecuteAsync("batch", container, partitionKey, specs, token); + } + catch (CosmosException ce) + { + throw new CommandException( + "batch", + MessageService.GetArgsString( + "command-batch-error-execution_failed", + "status", + ce.StatusCode.ToString(), + "message", + CommandException.GetDisplayMessage(ce)), + ce); + } + } + + private async Task BeginAsync(ShellInterpreter shell, CancellationToken token) + { + if (shell.State is not ConnectedState connectedState) + { + throw new NotConnectedException("batch"); + } + + if (shell.CurrentBatch is not null) + { + throw new CommandException("batch", MessageService.GetString("command-batch-error-already_active")); + } + + if (string.IsNullOrWhiteSpace(this.PartitionKeyArgument)) + { + throw new CommandException("batch", MessageService.GetString("command-batch-error-missing_pk")); + } + + var (databaseName, containerName, _) = await ResolveContainerAsync( + connectedState.Client, + shell.State, + this.Database, + this.Container, + "batch", + token); + + var partitionKey = ParsePartitionKey(this.PartitionKeyArgument); + shell.CurrentBatch = new PendingBatchState(databaseName!, containerName!, this.PartitionKeyArgument, partitionKey); + + var message = MessageService.GetArgsString( + "command-batch-begun", + "database", + databaseName!, + "container", + containerName!); + return new CommandState { RenderUser = () => ShellInterpreter.WriteLine(message) }; + } + + private CommandState Add(ShellInterpreter shell, CommandState commandState) + { + var batch = shell.CurrentBatch + ?? throw new CommandException("batch", MessageService.GetString("command-batch-error-not_active")); + + var json = this.ResolveData(commandState); + var specs = BatchOperationParser.Parse("batch", json); + + if (batch.Operations.Count + specs.Count > BatchExecutor.MaxOperations) + { + throw new CommandException( + "batch", + MessageService.GetArgsString( + "command-batch-error-too_many", + "count", + batch.Operations.Count + specs.Count)); + } + + batch.Operations.AddRange(specs); + var message = MessageService.GetArgsString( + "command-batch-added", + "count", + specs.Count, + "total", + batch.Operations.Count); + return new CommandState { RenderUser = () => ShellInterpreter.WriteLine(message) }; + } + + private async Task ExecuteBatchAsync(ShellInterpreter shell, CancellationToken token) + { + if (shell.State is not ConnectedState connectedState) + { + throw new NotConnectedException("batch"); + } + + var batch = shell.CurrentBatch + ?? throw new CommandException("batch", MessageService.GetString("command-batch-error-not_active")); + + if (batch.Operations.Count == 0) + { + throw new CommandException("batch", MessageService.GetString("command-batch-error-empty")); + } + + var container = connectedState.Client.GetDatabase(batch.DatabaseName).GetContainer(batch.ContainerName); + + try + { + var result = await BatchExecutor.ExecuteAsync("batch", container, batch.PartitionKey, batch.Operations, token); + shell.CurrentBatch = null; + return result; + } + catch (CosmosException ce) + { + throw new CommandException( + "batch", + MessageService.GetArgsString( + "command-batch-error-execution_failed", + "status", + ce.StatusCode.ToString(), + "message", + CommandException.GetDisplayMessage(ce)), + ce); + } + } + + private string ResolveData(CommandState commandState) + { + var evaluatedResult = commandState.Result?.ConvertShellObject(DataType.Text); + var json = this.Data ?? (evaluatedResult as string); + if (string.IsNullOrWhiteSpace(json)) + { + throw new CommandException("batch", MessageService.GetString("command-batch-error-missing_data")); + } + + return json; + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs new file mode 100644 index 00000000..67e109d1 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchExecutor.cs @@ -0,0 +1,141 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Commands; + +using System.Globalization; +using System.Text.Json; +using System.Text.Json.Nodes; +using Azure.Data.Cosmos.Shell.Parser; +using Azure.Data.Cosmos.Shell.Util; +using global::Azure; +using global::Azure.Data.Cosmos.Shell.Core; + +internal static class BatchExecutor +{ + internal const int MaxOperations = 100; + + public static async Task ExecuteAsync( + string commandName, + Container container, + PartitionKey partitionKey, + IReadOnlyList operations, + CancellationToken token) + { + if (operations.Count == 0) + { + throw new CommandException(commandName, MessageService.GetString("command-batch-error-empty")); + } + + if (operations.Count > MaxOperations) + { + throw new CommandException( + commandName, + MessageService.GetArgsString("command-batch-error-too_many", "count", operations.Count)); + } + + var batch = container.CreateTransactionalBatch(partitionKey); + foreach (var operation in operations) + { + switch (operation.Kind) + { + case BatchOperationKind.Create: + batch.CreateItem(operation.Item!.Value); + break; + + case BatchOperationKind.Upsert: + batch.UpsertItem(operation.Item!.Value); + break; + + case BatchOperationKind.Replace: + batch.ReplaceItem(operation.Id!, operation.Item!.Value); + break; + + case BatchOperationKind.Delete: + batch.DeleteItem(operation.Id!); + break; + + case BatchOperationKind.Patch: + batch.PatchItem(operation.Id!, operation.PatchOperations!); + break; + } + } + + using var response = await batch.ExecuteAsync(token); + + var summary = BuildSummary(operations, response); + + if (response.IsSuccessStatusCode) + { + var successMessage = MessageService.GetArgsString( + "command-batch-success", + "count", + operations.Count, + "charge", + response.RequestCharge.ToString("F2", CultureInfo.InvariantCulture)); + return CreateResultState(summary, successMessage); + } + + var errorMessage = MessageService.GetArgsString( + "command-batch-error-failed", + "status", + ((int)response.StatusCode).ToString(CultureInfo.InvariantCulture), + "charge", + response.RequestCharge.ToString("F2", CultureInfo.InvariantCulture)); + var errorState = new StructuredErrorCommandState( + new CommandException( + commandName, + errorMessage, + new RequestFailedException((int)response.StatusCode, errorMessage)), + new ShellJson(summary)); + errorState.RenderUser = () => ShellInterpreter.WriteLine(errorMessage); + return errorState; + } + + private static CommandState CreateResultState(JsonElement summary, string message) + { + var state = new CommandState { Result = new ShellJson(summary) }; + state.RenderUser = () => ShellInterpreter.WriteLine(message); + return state; + } + + private static JsonElement BuildSummary(IReadOnlyList operations, TransactionalBatchResponse response) + { + var operationsArray = new JsonArray(); + for (var i = 0; i < response.Count; i++) + { + var result = response[i]; + var node = new JsonObject + { + ["index"] = i, + ["op"] = operations[i].Kind.ToString().ToLowerInvariant(), + ["statusCode"] = (int)result.StatusCode, + }; + + if (operations[i].Id is { } id) + { + node["id"] = id; + } + + if (!string.IsNullOrEmpty(result.ETag)) + { + node["etag"] = result.ETag; + } + + operationsArray.Add(node); + } + + var root = new JsonObject + { + ["success"] = response.IsSuccessStatusCode, + ["statusCode"] = (int)response.StatusCode, + ["requestCharge"] = response.RequestCharge, + ["operationCount"] = operations.Count, + ["operations"] = operationsArray, + }; + + using var document = JsonDocument.Parse(root.ToJsonString()); + return document.RootElement.Clone(); + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchOperationParser.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchOperationParser.cs new file mode 100644 index 00000000..b320ba51 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BatchOperationParser.cs @@ -0,0 +1,198 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Commands; + +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Util; +using global::Azure.Data.Cosmos.Shell.Core; + +internal static class BatchOperationParser +{ + public static List Parse(string commandName, string json) + { + JsonDocument document; + try + { + document = JsonDocument.Parse(json); + } + catch (JsonException ex) + { + throw new CommandException( + commandName, + MessageService.GetArgsString("command-batch-error-invalid_json", "message", ex.Message), + ex); + } + + using (document) + { + var root = document.RootElement; + var specs = new List(); + + switch (root.ValueKind) + { + case JsonValueKind.Array: + foreach (var element in root.EnumerateArray()) + { + specs.Add(ParseOne(commandName, element)); + } + + break; + + case JsonValueKind.Object: + specs.Add(ParseOne(commandName, root)); + break; + + default: + throw new CommandException(commandName, MessageService.GetString("command-batch-error-not_object")); + } + + return specs; + } + } + + private static BatchOperationSpec ParseOne(string commandName, JsonElement element) + { + if (element.ValueKind != JsonValueKind.Object) + { + throw new CommandException(commandName, MessageService.GetString("command-batch-error-not_object")); + } + + if (!element.TryGetProperty("op", out var opElement) || opElement.ValueKind != JsonValueKind.String) + { + throw new CommandException(commandName, MessageService.GetString("command-batch-error-missing_op")); + } + + var op = opElement.GetString()!.Trim().ToLowerInvariant(); + var raw = element.Clone(); + + switch (op) + { + case "create": + return new BatchOperationSpec { Kind = BatchOperationKind.Create, Item = RequireItem(commandName, element, op), Id = ExtractItemId(element), RawOperation = raw }; + + case "upsert": + return new BatchOperationSpec { Kind = BatchOperationKind.Upsert, Item = RequireItem(commandName, element, op), Id = ExtractItemId(element), RawOperation = raw }; + + case "replace": + var item = RequireItem(commandName, element, op); + var replaceId = ExtractExplicitId(element) ?? ExtractItemId(element); + if (string.IsNullOrEmpty(replaceId)) + { + throw new CommandException(commandName, MessageService.GetArgsString("command-batch-error-missing_id", "op", op)); + } + + return new BatchOperationSpec { Kind = BatchOperationKind.Replace, Item = item, Id = replaceId, RawOperation = raw }; + + case "delete": + var deleteId = ExtractExplicitId(element); + if (string.IsNullOrEmpty(deleteId)) + { + throw new CommandException(commandName, MessageService.GetArgsString("command-batch-error-missing_id", "op", op)); + } + + return new BatchOperationSpec { Kind = BatchOperationKind.Delete, Id = deleteId, RawOperation = raw }; + + case "patch": + var patchId = ExtractExplicitId(element); + if (string.IsNullOrEmpty(patchId)) + { + throw new CommandException(commandName, MessageService.GetArgsString("command-batch-error-missing_id", "op", op)); + } + + return new BatchOperationSpec { Kind = BatchOperationKind.Patch, Id = patchId, PatchOperations = ParsePatchOperations(commandName, element), RawOperation = raw }; + + default: + throw new CommandException( + commandName, + MessageService.GetArgsString("command-batch-error-unsupported_op", "op", op)); + } + } + + private static JsonElement RequireItem(string commandName, JsonElement element, string op) + { + if (!element.TryGetProperty("item", out var itemElement)) + { + throw new CommandException(commandName, MessageService.GetArgsString("command-batch-error-missing_item", "op", op)); + } + + if (itemElement.ValueKind != JsonValueKind.Object) + { + throw new CommandException(commandName, MessageService.GetArgsString("command-batch-error-invalid_item", "op", op)); + } + + return itemElement.Clone(); + } + + private static string? ExtractExplicitId(JsonElement element) + { + if (element.TryGetProperty("id", out var idElement)) + { + return idElement.ValueKind switch + { + JsonValueKind.String => idElement.GetString(), + JsonValueKind.Number => idElement.GetRawText(), + _ => null, + }; + } + + return null; + } + + private static string? ExtractItemId(JsonElement element) + { + if (element.TryGetProperty("item", out var itemElement) + && itemElement.ValueKind == JsonValueKind.Object + && itemElement.TryGetProperty("id", out var idElement)) + { + return idElement.ValueKind switch + { + JsonValueKind.String => idElement.GetString(), + JsonValueKind.Number => idElement.GetRawText(), + _ => null, + }; + } + + return null; + } + + private static List ParsePatchOperations(string commandName, JsonElement element) + { + if (!element.TryGetProperty("operations", out var operationsElement) || operationsElement.ValueKind != JsonValueKind.Array) + { + throw new CommandException(commandName, MessageService.GetString("command-batch-error-missing_patch_ops")); + } + + var operations = new List(); + foreach (var patchElement in operationsElement.EnumerateArray()) + { + if (patchElement.ValueKind != JsonValueKind.Object + || !patchElement.TryGetProperty("op", out var patchOp) + || patchOp.ValueKind != JsonValueKind.String + || !patchElement.TryGetProperty("path", out var patchPath) + || patchPath.ValueKind != JsonValueKind.String) + { + throw new CommandException(commandName, MessageService.GetString("command-batch-error-invalid_patch_op")); + } + + var value = patchElement.TryGetProperty("value", out var valueElement) + ? valueElement.GetRawText() + : null; + + operations.Add(PatchOperationFactory.Build( + commandName, + patchOp.GetString()!, + patchPath.GetString()!, + value, + "command-batch-error-unsupported_patch_op")); + } + + if (operations.Count == 0) + { + throw new CommandException(commandName, MessageService.GetString("command-batch-error-missing_patch_ops")); + } + + return operations; + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs index cb0c5fff..cc684d5e 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchCommand.cs @@ -4,7 +4,6 @@ namespace Azure.Data.Cosmos.Shell.Commands; -using System.Globalization; using System.Text.Json; using Azure.Data.Cosmos.Shell.Parser; using Azure.Data.Cosmos.Shell.Util; @@ -66,8 +65,8 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co throw new CommandException("patch", MessageService.GetString("command-patch-error-missing_op")); } - var op = NormalizeOperation(this.Op); - if (!IsSupportedOperation(op)) + var op = PatchOperationFactory.Normalize(this.Op); + if (!PatchOperationFactory.IsSupported(op)) { throw new CommandException( "patch", @@ -94,7 +93,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co "patch", token); - var operation = BuildPatchOperation(op, this.Path!, this.Value); + var operation = PatchOperationFactory.Build("patch", op, this.Path!, this.Value); var requestOptions = string.IsNullOrEmpty(this.ETag) ? null @@ -163,132 +162,4 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co ce); } } - - private static PatchOperation BuildPatchOperation(string opRaw, string path, string? value) - { - var op = NormalizeOperation(opRaw); - - switch (op) - { - case "remove": - if (value is not null) - { - throw new CommandException( - "patch", - MessageService.GetString("command-patch-error-unexpected_value_for_remove")); - } - - return PatchOperation.Remove(path); - - case "add": - return PatchOperation.Add(path, ParseValue(op, value)); - - case "set": - return PatchOperation.Set(path, ParseValue(op, value)); - - case "replace": - return PatchOperation.Replace(path, ParseValue(op, value)); - - case "incr": - case "increment": - return BuildIncrementOperation(path, value); - - default: - throw new CommandException( - "patch", - MessageService.GetString( - "command-patch-error-unsupported_op", - new Dictionary { { "op", op } })); - } - } - - private static string NormalizeOperation(string op) => op.Trim().ToLowerInvariant(); - - private static bool IsSupportedOperation(string op) - { - return op is "set" or "add" or "replace" or "remove" or "incr" or "increment"; - } - - private static PatchOperation BuildIncrementOperation(string path, string? rawValue) - { - if (rawValue == null) - { - throw new CommandException( - "patch", - MessageService.GetString( - "command-patch-error-missing_value_for_op", - new Dictionary { { "op", "incr" } })); - } - - var trimmed = rawValue.Trim(); - if (long.TryParse(trimmed, NumberStyles.Integer, CultureInfo.InvariantCulture, out var intValue)) - { - return PatchOperation.Increment(path, intValue); - } - - if (double.TryParse(trimmed, NumberStyles.Float, CultureInfo.InvariantCulture, out var doubleValue)) - { - return PatchOperation.Increment(path, doubleValue); - } - - throw new CommandException( - "patch", - MessageService.GetString("command-patch-error-increment_number")); - } - - private static object? ParseValue(string op, string? rawValue) - { - if (rawValue == null) - { - throw new CommandException( - "patch", - MessageService.GetString( - "command-patch-error-missing_value_for_op", - new Dictionary { { "op", op } })); - } - - var trimmed = rawValue.Trim(); - if (LooksLikeJsonLiteral(trimmed)) - { - try - { - using var doc = JsonDocument.Parse(trimmed); - return JsonSerializer.Deserialize(doc.RootElement.GetRawText()); - } - catch (JsonException) - { - // Fall through and treat as plain string. - } - } - - return rawValue; - } - - private static bool LooksLikeJsonLiteral(string trimmed) - { - if (trimmed.Length == 0) - { - return false; - } - - var first = trimmed[0]; - if (first == '{' || first == '[' || first == '"') - { - return true; - } - - if (first == '-' || char.IsDigit(first)) - { - return true; - } - - if (string.Equals(trimmed, "true", StringComparison.Ordinal) - || string.Equals(trimmed, "false", StringComparison.Ordinal) - || string.Equals(trimmed, "null", StringComparison.Ordinal)) - { - return true; - } - - return false; - } } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchOperationFactory.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchOperationFactory.cs new file mode 100644 index 00000000..f2c920c9 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PatchOperationFactory.cs @@ -0,0 +1,146 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Commands; + +using System.Globalization; +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Util; +using global::Azure.Data.Cosmos.Shell.Core; + +internal static class PatchOperationFactory +{ + public static string Normalize(string op) => op.Trim().ToLowerInvariant(); + + public static bool IsSupported(string op) + { + return op is "set" or "add" or "replace" or "remove" or "incr" or "increment"; + } + + public static PatchOperation Build( + string commandName, + string opRaw, + string path, + string? value, + string unsupportedOpMessageKey = "command-patch-error-unsupported_op") + { + var op = Normalize(opRaw); + + switch (op) + { + case "remove": + if (value is not null) + { + throw new CommandException( + commandName, + MessageService.GetString("command-patch-error-unexpected_value_for_remove")); + } + + return PatchOperation.Remove(path); + + case "add": + return PatchOperation.Add(path, ParseValue(commandName, op, value)); + + case "set": + return PatchOperation.Set(path, ParseValue(commandName, op, value)); + + case "replace": + return PatchOperation.Replace(path, ParseValue(commandName, op, value)); + + case "incr": + case "increment": + return BuildIncrementOperation(commandName, path, value); + + default: + throw new CommandException( + commandName, + MessageService.GetString( + unsupportedOpMessageKey, + new Dictionary { { "op", op } })); + } + } + + private static PatchOperation BuildIncrementOperation(string commandName, string path, string? rawValue) + { + if (rawValue == null) + { + throw new CommandException( + commandName, + MessageService.GetString( + "command-patch-error-missing_value_for_op", + new Dictionary { { "op", "incr" } })); + } + + var trimmed = rawValue.Trim(); + if (long.TryParse(trimmed, NumberStyles.Integer, CultureInfo.InvariantCulture, out var intValue)) + { + return PatchOperation.Increment(path, intValue); + } + + if (double.TryParse(trimmed, NumberStyles.Float, CultureInfo.InvariantCulture, out var doubleValue)) + { + return PatchOperation.Increment(path, doubleValue); + } + + throw new CommandException( + commandName, + MessageService.GetString("command-patch-error-increment_number")); + } + + private static object? ParseValue(string commandName, string op, string? rawValue) + { + if (rawValue == null) + { + throw new CommandException( + commandName, + MessageService.GetString( + "command-patch-error-missing_value_for_op", + new Dictionary { { "op", op } })); + } + + var trimmed = rawValue.Trim(); + if (LooksLikeJsonLiteral(trimmed)) + { + try + { + using var doc = JsonDocument.Parse(trimmed); + return JsonSerializer.Deserialize(doc.RootElement.GetRawText()); + } + catch (JsonException) + { + // Fall through and treat as plain string. + } + } + + return rawValue; + } + + private static bool LooksLikeJsonLiteral(string trimmed) + { + if (trimmed.Length == 0) + { + return false; + } + + var first = trimmed[0]; + if (first == '{' || first == '[' || first == '"') + { + return true; + } + + if (first == '-' || char.IsDigit(first)) + { + return true; + } + + if (string.Equals(trimmed, "true", StringComparison.Ordinal) + || string.Equals(trimmed, "false", StringComparison.Ordinal) + || string.Equals(trimmed, "null", StringComparison.Ordinal)) + { + return true; + } + + return false; + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/BatchOperationSpec.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/BatchOperationSpec.cs new file mode 100644 index 00000000..87df9d10 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/BatchOperationSpec.cs @@ -0,0 +1,29 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +using System.Text.Json; + +internal enum BatchOperationKind +{ + Create, + Upsert, + Replace, + Delete, + Patch, +} + +internal sealed class BatchOperationSpec +{ + public BatchOperationKind Kind { get; init; } + + public string? Id { get; init; } + + public JsonElement? Item { get; init; } + + public IReadOnlyList? PatchOperations { get; init; } + + public JsonElement RawOperation { get; init; } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosShellPrompt.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosShellPrompt.cs index 83f656fd..8bf54177 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosShellPrompt.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosShellPrompt.cs @@ -20,6 +20,7 @@ internal class CosmosShellPrompt(ShellInterpreter shell) : ILineEditorPrompt, IS private Markup prompt = new(string.Empty); private State? oldState; private ThemeOptions? oldTheme; + private string? oldBatchSignature; internal static string PromptMarker => SupportsFancyPromptMarker() ? FancyPromptMarker : AsciiPromptMarker; @@ -39,10 +40,12 @@ internal class CosmosShellPrompt(ShellInterpreter shell) : ILineEditorPrompt, IS return (new Markup(Theme.FormatMuted("...")), 1); } - if (this.oldState != this.shell.State || !ReferenceEquals(this.oldTheme, Theme.Current)) + var batchSignature = this.GetBatchSignature(); + if (this.oldState != this.shell.State || !ReferenceEquals(this.oldTheme, Theme.Current) || this.oldBatchSignature != batchSignature) { this.oldState = this.shell.State; this.oldTheme = Theme.Current; + this.oldBatchSignature = batchSignature; this.prompt = new Markup(this.GetPromptString()); } @@ -54,8 +57,21 @@ public string GetPromptString() this.oldState ??= this.shell.State; this.oldTheme ??= Theme.Current; #pragma warning disable VSTHRD002 // Synchronously waiting - required by ILineEditorPrompt interface - return this.oldState.AcceptAsync(this, null, default).Result ?? string.Empty; + var basePrompt = this.oldState.AcceptAsync(this, null, default).Result ?? string.Empty; #pragma warning restore VSTHRD002 + var batch = this.shell.CurrentBatch; + if (batch is not null) + { + basePrompt += " " + Theme.FormatMuted($"[batch:{batch.Operations.Count}]"); + } + + return basePrompt; + } + + private string GetBatchSignature() + { + var batch = this.shell.CurrentBatch; + return batch is null ? string.Empty : batch.Operations.Count.ToString(System.Globalization.CultureInfo.InvariantCulture); } Task IStateVisitor.VisitConnectedStateAsync(ConnectedState state, object? data, CancellationToken token) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/PendingBatchState.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/PendingBatchState.cs new file mode 100644 index 00000000..34041b00 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/PendingBatchState.cs @@ -0,0 +1,26 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +internal sealed class PendingBatchState +{ + public PendingBatchState(string databaseName, string containerName, string partitionKeyArgument, PartitionKey partitionKey) + { + this.DatabaseName = databaseName; + this.ContainerName = containerName; + this.PartitionKeyArgument = partitionKeyArgument; + this.PartitionKey = partitionKey; + } + + public string DatabaseName { get; } + + public string ContainerName { get; } + + public string PartitionKeyArgument { get; } + + public PartitionKey PartitionKey { get; } + + public List Operations { get; } = []; +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index 4826518f..bd9c77f0 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -179,6 +179,8 @@ internal static char CSVSeparator internal int? McpPort { get; set; } + internal PendingBatchState? CurrentBatch { get; set; } + internal Queue VariableContainers { get; } = new(); /// @@ -1459,6 +1461,7 @@ internal void Connect(CosmosClient client, ArmCosmosContext? armContext = null) { this.State?.Dispose(); this.State = new ConnectedState(client, armContext); + this.CurrentBatch = null; CosmosCompleteCommand.ClearDatabases(); CosmosCompleteCommand.ClearContainers(); this.Diagnostics?.LogConnect(client.Endpoint, client.ClientOptions.ConnectionMode); @@ -1523,6 +1526,7 @@ internal void Disconnect() { this.State?.Dispose(); this.State = new DisconnectedState(); + this.CurrentBatch = null; } internal void PrintCommand(string cmdString) @@ -1575,6 +1579,12 @@ internal CommandState PrintState(CommandState state, bool markAsRendered = false return state; } + if (inMachineMode && state is StructuredErrorCommandState structuredError) + { + this.WriteMachineError(structuredError.Exception.Message, structuredError.Result); + return state; + } + string? output; if (state.Result?.DataType == Parser.DataType.Json) @@ -2106,14 +2116,20 @@ internal static bool IsIncompleteInput(string text) // redirection (`ErrOutRedirect` / `2>` / `2>>`) so scripts that redirect // stderr still capture errors in --quiet / --output json modes; otherwise // the object is written to the process stderr. - private void WriteMachineError(string errorMessage) + private void WriteMachineError(string errorMessage, ShellObject? result = null) { - var errObj = new + var error = new Dictionary { - status = "error", - error = errorMessage, + ["status"] = "error", + ["error"] = errorMessage, }; - var json = JsonSerializer.Serialize(errObj); + + if (result?.ConvertShellObject(Parser.DataType.Json) is JsonElement resultElement) + { + error["result"] = resultElement; + } + + var json = JsonSerializer.Serialize(error); if (this.ErrOutRedirect != null) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/StructuredErrorCommandState.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/StructuredErrorCommandState.cs new file mode 100644 index 00000000..4895d3ba --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/StructuredErrorCommandState.cs @@ -0,0 +1,16 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +using Azure.Data.Cosmos.Shell.Parser; + +internal sealed class StructuredErrorCommandState : ErrorCommandState +{ + public StructuredErrorCommandState(Exception exception, ShellObject result) + : base(exception) + { + this.Result = result; + } +} \ No newline at end of file diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs index c7c4249e..7a197539 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/McpResponseFactory.cs @@ -87,6 +87,12 @@ private static JsonObject CreateSuccessPayload(CommandState commandState) { payload["error"] = GetErrorPayloadMessage(commandState); + var errorResultNode = CreateResultNode(commandState); + if (errorResultNode != null) + { + payload["result"] = errorResultNode; + } + return payload; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index 0ec8a145..6f39542f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs @@ -417,7 +417,6 @@ private async ValueTask OnCallToolsAsync( var parameter = command.Parameters.FirstOrDefault(a => MatchesArgumentName(a.Name, par.Key)); if (parameter != null) { - suppliedParameters.Add(parameter.Name[0]); var bindError = this.BindMember( cmd, parameter.PropertyInfo, @@ -431,6 +430,13 @@ private async ValueTask OnCallToolsAsync( return bindError; } + var boundValue = parameter.PropertyInfo.GetValue(cmd); + if (boundValue != null + && (boundValue is not string stringValue || !string.IsNullOrWhiteSpace(stringValue))) + { + suppliedParameters.Add(parameter.Name[0]); + } + continue; } @@ -458,6 +464,15 @@ private async ValueTask OnCallToolsAsync( return McpResponseFactory.CreateError(missingMessage, ShellInterpreter.Instance.State); } + var batchSubcommand = (cmd as BatchCommand)?.Subcommand?.Trim(); + if (!string.IsNullOrEmpty(batchSubcommand) + && !string.Equals(batchSubcommand, "run", StringComparison.OrdinalIgnoreCase)) + { + const string errorMessage = "MCP supports only the stateless 'batch run' subcommand. Run stateful batch commands manually in the shell."; + this.logger?.LogWarning(errorMessage); + return McpResponseFactory.CreateError(errorMessage, ShellInterpreter.Instance.State); + } + if (RequiresConfirmation(command)) { var server = parameters.Server; diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index d0a5ae95..5bf85ad9 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -265,6 +265,54 @@ command-patch-error-failed = Failed to patch item: { $status } - { $message } command-patch-error-not_found = Item '{ $id }' not found. command-patch-error-etag_mismatch = Item '{ $id }' was modified since it was last read (ETag mismatch). +command-batch-description = Executes multiple write operations against a single partition key as one atomic transactional batch, either in a single call (run) or as a stateful batch (begin, add, execute, cancel, status, show). +command-batch-description-subcommand = The action to perform: run, begin, add, execute, cancel, status, or show. +command-batch-description-data = The batch operations as a JSON array, or a single operation as a JSON object. +command-batch-description-partition-key = The partition key shared by every operation in the batch. +command-batch-description-database = The database containing the target container. +command-batch-description-container = The target container for the batch. +command-batch-success = Batch of { $count } { $count -> + [one] operation + *[other] operations +} committed (RU charge: { $charge }) +command-batch-begun = Started a batch on { $database }/{ $container }. Add operations with 'batch add' and commit with 'batch execute'. +command-batch-added = Added { $count } { $count -> + [one] operation + *[other] operations +} ({ $total } pending). +command-batch-cancelled = Discarded the pending batch ({ $count } { $count -> + [one] operation + *[other] operations +}). +command-batch-status-inactive = No batch is currently active. +command-batch-status-target = Target +command-batch-status-partition-key = Partition key +command-batch-status-operation-count = Operations +command-batch-status-column-index = # +command-batch-status-column-operation = Operation +command-batch-status-column-id = Item ID +command-batch-error-missing_subcommand = Missing subcommand. Use one of: run, begin, add, execute, cancel, status, show. +command-batch-error-invalid_subcommand = Unknown subcommand '{ $subcommand }'. Use one of: run, begin, add, execute, cancel, status, show. +command-batch-error-missing_pk = A partition key is required. Specify it with --partition-key. +command-batch-error-invalid_pk_json = Partition key must be a JSON scalar value or a JSON array of values for hierarchical partition keys. +command-batch-error-missing_data = Batch operations are required. Provide a JSON array of operations, or a single operation object. +command-batch-error-invalid_json = The batch operations are not valid JSON: { $message } +command-batch-error-not_object = Each batch operation must be a JSON object, and the batch itself must be a JSON object or an array of objects. +command-batch-error-missing_op = Each batch operation requires an 'op' string field. Supported: create, upsert, replace, delete, patch. +command-batch-error-unsupported_op = Unsupported batch operation '{ $op }'. Supported: create, upsert, replace, delete, patch. +command-batch-error-missing_item = Operation '{ $op }' requires an 'item' object. +command-batch-error-invalid_item = The 'item' for operation '{ $op }' must be a JSON object. +command-batch-error-missing_id = Operation '{ $op }' requires an 'id'. +command-batch-error-missing_patch_ops = A patch operation requires a non-empty 'operations' array. +command-batch-error-invalid_patch_op = Each entry in 'operations' requires 'op' and 'path' string fields. +command-batch-error-unsupported_patch_op = Unsupported patch operation '{ $op }' in a batch patch entry. Supported operations: set, add, replace, remove, incr. +command-batch-error-empty = The batch has no operations. Add at least one operation before executing. +command-batch-error-too_many = A transactional batch supports at most 100 operations, but { $count } were provided. +command-batch-error-already_active = A batch is already in progress. Run 'batch execute' or 'batch cancel' first. +command-batch-error-not_active = No batch is in progress. Start one with 'batch begin'. +command-batch-error-failed = Batch failed with status { $status } and was rolled back (RU charge: { $charge }) +command-batch-error-execution_failed = Failed to execute batch: { $status } - { $message } + command-export-description = Exports items from a container to a JSON Lines, JSON array, or CSV file. command-export-description-file = Destination file path. command-export-description-database = The database to read from. diff --git a/README.md b/README.md index 21011f6d..fdbf0c76 100644 --- a/README.md +++ b/README.md @@ -10,6 +10,7 @@ A terminal-native shell for Azure Cosmos DB — navigate databases like a filesy - Navigate with `ls` and `cd` (Account -> Databases -> Containers -> Items) - Inspect the current location with `pwd` - Create, query, replace, patch, delete: `mkdb`, `mkcon`, `mkitem`, `query`, `replace`, `patch`, `rm` +- Atomic multi-operation transactions on a single partition key: `batch` - Bulk roundtrip with `import` / `export` for JSON Lines and JSON array files, plus CSV import/export (CSV import coerces values to strings; `--partition-key` nests a CSV column under a nested partition key path) - Manage container indexing policies with `index` (`show`, `add`, `remove`, `set`) - Inspect container/database/account configuration and usage statistics with `info` (partition key, throughput, policies, indexing policy summary, document count, storage size, regions; `--partitions` and `--detailed` for distribution analysis) diff --git a/docs/commands.md b/docs/commands.md index 163d843e..36bc710e 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -294,6 +294,97 @@ patch set order-42 customer-7 /name "Ada Lovelace" --etag="" - `remove` with a `value` argument is rejected up front. - `incr` with a non-numeric value is rejected up front. +### batch + +Execute multiple write operations against a single partition key as one atomic Cosmos DB transactional batch. Either run a batch in a single call, or build one up statefully across several commands. Every operation in a batch must share the same partition key, execution requires between 1 and 100 operations, and if any operation fails the entire batch is rolled back. A pending stateful batch may be empty until operations are added. + +```text +Usage: batch subcommand [data] [--partition-key ] [-database ] [-container ] + +Arguments: + subcommand The action to perform: run, begin, add, execute, cancel, status, or show + [data] Batch operations as a JSON array, or a single operation as a JSON object (Optional) + +Options: + --partition-key, --pk + The partition key shared by every operation in the batch. Required for `run` and `begin`; the stateful `add`, `execute`, `cancel`, `status`, and `show` subcommands use the active batch's partition key. + -database, -db + Override database name (Optional) + -container, -con + Override container name (Optional) +``` + +#### Subcommands + +|Subcommand|Description| +|-|-| +|`run --partition-key `|Parse a JSON array of operations and execute them atomically in a single call. Also reads piped input.| +|`begin --partition-key `|Start a stateful batch bound to a partition key, database, and container.| +|`add `|Queue one operation (JSON object) or several (JSON array) onto the active batch.| +|`execute` (`exec`, `commit`)|Commit the queued operations atomically and clear the active batch.| +|`cancel` (`abort`)|Discard the active batch without executing it.| +|`status`|Report the active batch target, partition key, and a compact list of queued operations.| +|`show`|Print the queued operations as a JSON array (the same shape accepted by `run`/`add`).| + +When a stateful batch is active the prompt shows a `[batch:N]` indicator, where `N` is the number of queued operations. + +In interactive user output, `status` uses a compact table while `show` always prints the full queued-operation JSON. Use `--output json`, redirection, or a pipeline to obtain the structured JSON result from `status`. + +#### Operation schema + +Each operation is a JSON object with an `op` field: + +|Operation|Shape| +|-|-| +|`create`|`{"op":"create","item":{...}}`| +|`upsert`|`{"op":"upsert","item":{...}}`| +|`replace`|`{"op":"replace","id":"1","item":{...}}` (the `id` is optional when `item.id` is present)| +|`delete`|`{"op":"delete","id":"3"}`| +|`patch`|`{"op":"patch","id":"1","operations":[{"op":"set","path":"/name","value":"x"}]}`| + +Patch sub-operations use the same `op`/`path`/`value` shape and semantics as the [`patch`](#patch) command (`set`, `add`, `replace`, `remove`, `incr`). Patch values are typed JSON values, for example `"value":"active"`, `"value":42`, or `"value":true`. + +#### Result + +In interactive user output, `batch run` and `batch execute` print a concise outcome with the operation count and request charge. Their full result is emitted as a JSON summary when using `--output json`, redirection, a pipeline, or MCP: + +```json +{ + "success": true, + "statusCode": 200, + "requestCharge": 12.34, + "operationCount": 2, + "operations": [ + { "index": 0, "op": "create", "statusCode": 201, "id": "a", "etag": "..." }, + { "index": 1, "op": "delete", "statusCode": 204, "id": "b" } + ] +} +``` + +When the batch fails, `success` is `false`, the failing operation reports its own status code, and the remaining operations report `424` (Failed Dependency) because the transaction was rolled back. + +#### Examples + +```bash +batch run '[{"op":"create","item":{"id":"1","pk":"a"}},{"op":"delete","id":"2"}]' --partition-key a +echo '[{"op":"upsert","item":{"id":"3","pk":"a"}}]' | batch run --partition-key a + +batch begin --partition-key a +batch add '{"op":"upsert","item":{"id":"3","pk":"a"}}' +batch add '{"op":"patch","id":"3","operations":[{"op":"set","path":"/status","value":"done"}]}' +batch status +batch show # prints the queued operations as a JSON array +batch execute +``` + +#### Errors + +- Missing `--partition-key` for `run` or `begin` is rejected up front. +- `add`, `execute`, or `cancel` with no active batch: `No batch is in progress. Start one with 'batch begin'.` +- `begin` while a batch is already active: `A batch is already in progress. Run 'batch execute' or 'batch cancel' first.` +- More than 100 operations is rejected before any call to Cosmos DB. +- A transactional failure prints a one-line message, returns a result summary with `success` set to `false` and per-operation status codes, and rolls back every operation. + ### rm Remove items from container. diff --git a/docs/mcp.md b/docs/mcp.md index 7a089d67..10b89a25 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -42,6 +42,8 @@ The MCP server runs locally with your user permissions. Connected clients can ex Server-side programming commands — stored procedures (`sproc`), user-defined functions (`udf`), and triggers (`trigger`) — are restricted from MCP. Run those commands manually in the shell. +Transactional batches invoked through MCP must use the one-shot `batch run` subcommand. Stateful batch subcommands (`begin`, `add`, `execute`, `cancel`, `status`, and `show`, including their aliases) are restricted to the interactive shell because MCP tool calls share no client-specific batch state. + ### Destructive Command Confirmation Destructive commands (`delete`, `rm`, `rmcon`, `rmdb`) are gated behind an explicit user confirmation. When a client invokes one, the server sends an MCP elicitation prompt describing the exact command line before anything runs: @@ -96,10 +98,10 @@ Both representations are always byte-for-byte equivalent. | Field | When present | Description | | ----- | ------------ | ----------- | -| `result` | Successful commands that produce output | The command result as JSON (objects, arrays, or a scalar). Text-only results are represented as a JSON string. | +| `result` | Commands that produce output | The command result as JSON (objects, arrays, or a scalar). Text-only results are represented as a JSON string. Failed transactional batches include their per-operation summary here alongside `error`. | | `outputText` | CSV output commands with non-empty text | The CSV rendering of the result. Omitted when the CSV output is empty or whitespace. | | `error` | Failed commands | The error message. | | `currentLocation` | Always | The shell's current navigation path (for example `/MyDatabase/MyContainer`), or `null` when disconnected. | -Successful results set `result` (and optionally `outputText`); failed results set `error` and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls. +Successful results set `result` (and optionally `outputText`); failed results set `error`, may also include a structured `result`, and mark the tool result as an error. `currentLocation` is always included so a client can track navigation state across calls.