Add transactional batch support (batch command)

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<string> 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;
+    }
+}

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<TextContentBlock>(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()
     {

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));
+    }
 }

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);
+    }
 }

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()
     {

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<string, JsonElement>
+        {
+            [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<string, JsonElement>
+        {
+            ["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<string, JsonElement>
+        {
+            ["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()
     {