diff --git a/CHANGELOG.md b/CHANGELOG.md index 0421584f..3cf773c4 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,10 +11,22 @@ ### Fixes - Preserve `?` in unquoted command arguments so scalar indexing paths such as `index add /city/?` and `index remove /city/?` reach Cosmos DB unchanged. ([#260](https://github.com/Azure/CosmosDBShell/issues/260)) +- Long-running commands and scripts against a local emulator no longer fail with a timeout and disconnect after 10 seconds while the emulator is still responding. The 10-second limit now applies only when nothing makes progress, and it is paused while confirmation prompts and external editors wait for input. +- Preserve partial bulk write counts through pipelines on emulator timeouts, report consumed inactivity cancellations as failures, and keep active bulk selection, export, and piped removal scans from timing out during local processing. +- Prevent scripts from rendering timed-out partial results as successful, pause emulator inactivity deadlines during file/theme editing and change-feed polling intervals, and cancel stalled change-feed requests. +- Keep Ctrl+C effective for startup commands and redirected scripts, and exclude local external tools such as `jq` from emulator inactivity deadlines. +- Report interrupted startup commands as failures, retain structured partial write results, and prevent MCP prompt refreshes from canceling queued startup work. +- Retain settled import counts and RUs on cancellation, prevent interrupted startup scripts from writing partial success output, and keep intentional watch stops successful through pipelines, blocks, and scripts. +- Stop follow-up startup commands after an intentional watch stop in redirected stdin, and align timed-out import result error messages and exit codes with the command failure. +- Keep import progress counts and observed RUs current when writes finish while the next input record is still being read. +- Keep bulk completion counts current during slow input reads, allow cancellation while loading resume journals, and recognize active local validation, removal scans, and successful journal recording without treating failed requests as activity. +- Stop interrupted startup stdin reads, confirmation prompts, and watch expressions before subsequent work; keep machine-readable stdout clean on Ctrl+C and exclude external command expressions from emulator inactivity deadlines. - Database and container creation on key, emulator, and static-token connections now report `created: false` and `outcome: "existing"` when an existing resource is reused, instead of always reporting `created: true`. Human-readable output says "Using existing database/container" in that case. ([#261](https://github.com/Azure/CosmosDBShell/issues/261)) ### Improvements +- Show periodic progress for `import`, `bulk`, and pattern-based `rm`, including elapsed time, counts, and observed RU usage, even while requests are waiting. In a terminal, a single progress line is updated in place and removed when the command finishes. Structured output and MCP responses remain unchanged. + ## 1.1.291 — 2026-10-26 ### Breaking changes diff --git a/CosmosDBShell.Tests/CommandTests/BulkCommandTests.cs b/CosmosDBShell.Tests/CommandTests/BulkCommandTests.cs index 706c9b0a..ce484148 100644 --- a/CosmosDBShell.Tests/CommandTests/BulkCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/BulkCommandTests.cs @@ -394,6 +394,183 @@ public async Task WriteClassifiesServiceErrors(int status, string expected) } // Execution engine. + [Theory] + [InlineData("success", false)] + [InlineData("record", false)] + [InlineData("write", false)] + [InlineData("cancel", false)] + [InlineData("source", false)] + [InlineData("record-and-source", false)] + [InlineData("write", true)] + [InlineData("source", true)] + [InlineData("record-and-source", true)] + public async Task PendingSourceReadSettlesCompletedWritesAndFinishesBeforeDisposal(string failureKind, bool applicationException) + { + using var timeout = Timeout(); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(timeout.Token); + var readStarted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var releaseRead = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var settled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var writes = new[] + { + new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), + new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), + }; + Exception expected = applicationException ? new ApplicationException(failureKind) : new IOException(failureKind); + Exception sourceFailure = applicationException ? new ApplicationException("source") : new IOException("source"); + var readFinished = false; + var disposed = false; + var progressUpdates = 0; + var summary = new BulkSummary(); + var records = new List(); + async IAsyncEnumerable OperationsAsync() + { + try + { + yield return Operation(0); + yield return Operation(1); + readStarted.SetResult(); + try + { + // Deliberately ignore cancellation until released to check disposal ordering. + await releaseRead.Task; + if (failureKind is "source" or "record-and-source") + { + throw sourceFailure; + } + } + finally + { + readFinished = true; + } + } + finally + { + disposed = true; + Assert.True(readFinished); + } + } + + var run = BulkExecutor.ExecuteAsync( + OperationsAsync(), + (operation, _) => writes[(int)operation.Index].Task, + outcome => + { + records.Add(outcome); + if (outcome.Status != BulkOutcome.Started) + { + settled.TrySetResult(); + if (failureKind is "record" or "record-and-source") + { + throw expected; + } + } + + return Task.CompletedTask; + }, + new Dictionary(), + summary, 3, null, true, false, cancellation.Token, + progressChanged: () => progressUpdates++); + try + { + await readStarted.Task.WaitAsync(timeout.Token); + if (failureKind == "write") + { + writes[0].SetException(expected); + } + else + { + writes[0].SetResult(Succeeded(0)); + } + + await settled.Task.WaitAsync(timeout.Token); + Assert.Equal(1, summary.Processed); + Assert.Equal(1, progressUpdates); + Assert.Equal(failureKind == "write" ? 0 : 1, summary.Succeeded); + Assert.Equal(failureKind == "write" ? 0 : 1, summary.RequestCharge); + if (failureKind == "cancel") + { + await cancellation.CancelAsync(); + } + + Assert.False(run.IsCompleted); + Assert.False(disposed); + } + finally + { + releaseRead.TrySetResult(); + writes[0].TrySetResult(Succeeded(0)); + writes[1].TrySetResult(Succeeded(1)); + } + + if (failureKind == "cancel") + { + await Assert.ThrowsAnyAsync(() => run.WaitAsync(timeout.Token)); + } + else if (failureKind != "success") + { + var actual = await Assert.ThrowsAnyAsync(() => run.WaitAsync(timeout.Token)); + Assert.Same(failureKind == "source" ? sourceFailure : expected, actual); + } + else + { + await run.WaitAsync(timeout.Token); + } + + Assert.True(readFinished); + Assert.True(disposed); + Assert.Equal(2, summary.Processed); + Assert.Equal(2, progressUpdates); + Assert.Equal(2, records.Count(outcome => outcome.Status != BulkOutcome.Started)); + } + + [Fact] + public async Task FinalDrainSettlesWritesInCompletionOrder() + { + using var timeout = Timeout(); + var gates = new[] + { + new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), + new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), + }; + var started = new[] { new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously), new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously) }; + var secondSettled = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var summary = new BulkSummary(); + var run = BulkExecutor.ExecuteAsync( + SourceAsync(2, timeout.Token), + (operation, _) => + { + started[(int)operation.Index].SetResult(); + return gates[(int)operation.Index].Task; + }, + IgnoreAsync, + new Dictionary(), + summary, + 3, + null, + false, + false, + timeout.Token, + progressChanged: () => secondSettled.TrySetResult()); + try + { + await Task.WhenAll(started.Select(gate => gate.Task)).WaitAsync(timeout.Token); + gates[1].SetResult(Succeeded(1)); + + await secondSettled.Task.WaitAsync(timeout.Token); + Assert.Equal(1, summary.Succeeded); + Assert.False(run.IsCompleted); + } + finally + { + gates[0].TrySetResult(Succeeded(0)); + gates[1].TrySetResult(Succeeded(1)); + } + + await run.WaitAsync(timeout.Token); + Assert.Equal(2, summary.Succeeded); + } + [Theory] [InlineData(1)] [InlineData(3)] @@ -564,6 +741,7 @@ public async Task ExhaustedRuBudgetDoesNotPreventLocalJournalSkips(string status var summary = new BulkSummary { RequestCharge = 2 }; var writes = new List(); var records = new List(); + var processed = new List(); await BulkExecutor.ExecuteAsync( SourceAsync(3, TestContext.Current.CancellationToken), @@ -577,9 +755,11 @@ await BulkExecutor.ExecuteAsync( records.Add(outcome); return Task.CompletedTask; }, - previous, summary, 1, 1, true, false, TestContext.Current.CancellationToken); + previous, summary, 1, 1, true, false, TestContext.Current.CancellationToken, + progressChanged: () => processed.Add(summary.Processed)); Assert.Empty(writes); + Assert.Equal([1L, 2L, 3L], processed); Assert.Equal(status == BulkOutcome.Succeeded ? 3 : 0, records.Count); Assert.All(records, outcome => Assert.Equal(BulkOutcome.Succeeded, outcome.Status)); Assert.Equal(3, summary.Skipped); @@ -1180,6 +1360,165 @@ public async Task JournalValidatesRecordIdentityAgainstCurrentOperationsBeforeAn } // End-to-end command behavior against a substituted container. + [Theory] + [InlineData(HttpStatusCode.OK, false)] + [InlineData(HttpStatusCode.Conflict, true)] + [InlineData(HttpStatusCode.RequestTimeout, true)] + [InlineData(HttpStatusCode.ServiceUnavailable, true)] + public async Task EmulatorActivityRefreshesForSettledSuccessesButNotWriteErrors(HttpStatusCode status, bool timesOut) + { + using var fixture = new Fixture(localEmulator: true); + fixture.Shell.Options = new Program.CosmosShellOptions { Output = "json" }; + fixture.Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(1); + fixture.Container.DeleteItemAsync( + Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(async call => + { + await Task.Delay(TimeSpan.FromMilliseconds(200), call.ArgAt(3)); + if (status != HttpStatusCode.OK) + { + throw new CosmosException("write failed", status, 0, "activity", 1); + } + + var response = Substitute.For>(); + response.StatusCode.Returns(status); + response.RequestCharge.Returns(1); + return response; + }); + var operations = string.Join(",", Enumerable.Range(0, 10).Select(index => Operation(index).Json)); + + var state = await fixture.Shell.ExecuteCommandAsync( + $"bulk run '[{operations}]' --db db --con items --concurrency 1 --continue-on-error --yes", + TestContext.Current.CancellationToken); + + Assert.Equal(timesOut, state.IsError); + if (timesOut) + { + var error = Assert.IsType(state); + Assert.Equal(ShellExitCode.ConnectionError, error.ExitCode); + Assert.Contains("timed out", error.Exception.Message, StringComparison.OrdinalIgnoreCase); + Assert.IsType(fixture.Shell.State); + Assert.True(fixture.Writes().Count() < 10); + } + else + { + Assert.IsType(fixture.Shell.State); + Assert.Equal(10, Assert.IsType(state.Result).Value.GetProperty("succeeded").GetInt64()); + } + } + + [Fact] + public async Task CanceledJournalReplayStopsBeforeParsingAndReleasesFile() + { + using var fixture = new Fixture(); + using var cancellation = Timeout(); + var path = fixture.Directory.PathFor("canceled-journal.jsonl"); + var hash = BulkSpool.ComputeHash([Operation(0).Json]); + using (var journal = BulkJournal.Open(path, "https://unit-test.documents.azure.com/", "rid", hash, 1)) + { + await journal.RecordAsync(Succeeded(0)); + } + + await File.AppendAllTextAsync(path, "invalid record\n", TestContext.Current.CancellationToken); + await fixture.RunAsync("bulk begin --db db --con items"); + await fixture.RunAsync($"bulk add '{Operation(0).Json}'"); + using var response = new ResponseMessage(HttpStatusCode.OK) + { + Content = new CancelOnDisposeStream( + Encoding.UTF8.GetBytes("""{"_rid":"rid","partitionKey":{"paths":["/pk"]}}"""), cancellation), + }; + fixture.Container.ReadContainerStreamAsync(Arg.Any(), Arg.Any()) + .Returns(Task.FromResult(response)); + var command = new BulkCommand { Subcommand = "execute", Journal = path, Yes = true }; + + await Assert.ThrowsAnyAsync(() => + command.ExecuteAsync(fixture.Shell, new CommandState(), "bulk execute", cancellation.Token)); + + Assert.Empty(fixture.Writes()); + using var reopened = new FileStream(path, FileMode.Open, FileAccess.ReadWrite, FileShare.None); + Assert.True(reopened.Length > 0); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task EmulatorTimeout_PreservesPartialResultsAndDisconnects(bool piped) + { + using var fixture = new Fixture(localEmulator: true); + fixture.Shell.Options = new Program.CosmosShellOptions { Output = "json" }; + fixture.Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromMilliseconds(100); + fixture.Container.DeleteItemAsync("1", Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(async call => + { + await Task.Delay(System.Threading.Timeout.InfiniteTimeSpan, call.ArgAt(3)); + return await Task.FromException>(new InvalidOperationException("The stalled write should be canceled.")); + }); + var command = $$"""bulk run '[{{Operation(0).Json}},{{Operation(1).Json}},{{Operation(2).Json}}]' --db db --con items --concurrency 1 --yes"""; + + var state = await fixture.Shell.ExecuteCommandAsync( + command + (piped ? " | echo" : string.Empty), TestContext.Current.CancellationToken); + + var error = Assert.IsType(state); + Assert.True(error.ExitCode == ShellExitCode.ConnectionError, error.Exception.ToString()); + Assert.Contains("timed out", error.Exception.Message, StringComparison.OrdinalIgnoreCase); + Assert.IsType(fixture.Shell.State); + var result = Assert.IsType(error.Result).Value; + Assert.Equal(1, result.GetProperty("succeeded").GetInt64()); + Assert.Equal(1, result.GetProperty("uncertain").GetInt64()); + Assert.Equal(2, result.GetProperty("attempted").GetInt64()); + Assert.True(result.GetProperty("resultIncomplete").GetBoolean()); + Assert.Equal(7, state.RequestCharge); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task WhereSelection_ReportsActivityWhileSpoolingOnePage(bool save) + { + using var fixture = new Fixture(localEmulator: true); + fixture.Shell.Options = new Program.CosmosShellOptions { Output = "json" }; + fixture.Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(1); + var rows = Enumerable.Range(0, 10) + .Select(index => JsonSerializer.SerializeToElement(new { id = index.ToString(CultureInfo.InvariantCulture), __pk0 = "t" })) + .ToArray(); + using var iterator = Substitute.For>(); + var read = false; + iterator.HasMoreResults.Returns(_ => !read); + iterator.ReadNextAsync(Arg.Any()).Returns(call => + { + read = true; + var page = Substitute.For>(); + page.GetEnumerator().Returns(_ => SlowRows(call.Arg()).GetEnumerator()); + return Task.FromResult(page); + }); + fixture.Container.GetItemQueryIterator( + Arg.Any(), Arg.Any(), Arg.Any()).Returns(iterator); + var savePath = fixture.Directory.PathFor("selection.jsonl"); + var command = "bulk delete --where true --dry-run --db db --con items" + + (save ? $" --save '{savePath}'" : string.Empty); + var watch = System.Diagnostics.Stopwatch.StartNew(); + + var state = await fixture.Shell.ExecuteCommandAsync(command, TestContext.Current.CancellationToken); + + Assert.False(state.IsError); + Assert.True(watch.Elapsed > fixture.Shell.LocalEmulatorInactivityTimeout); + Assert.IsType(fixture.Shell.State); + Assert.Equal(rows.Length, Assert.IsType(state.Result).Value.GetProperty("operationCount").GetInt64()); + if (save) + { + Assert.Equal(rows.Length, (await File.ReadAllLinesAsync(savePath, TestContext.Current.CancellationToken)).Length); + } + + IEnumerable SlowRows(CancellationToken token) + { + foreach (var row in rows) + { + Assert.False(token.WaitHandle.WaitOne(TimeSpan.FromMilliseconds(200)), "Active selection timed out."); + yield return row; + } + } + } + [Fact] public async Task RunExecutesMixedOperationsAcrossPartitions() { @@ -1932,6 +2271,18 @@ private static async IAsyncEnumerable SourceAsync(int count, [Enu } } + private sealed class CancelOnDisposeStream(byte[] content, CancellationTokenSource cancellation) : MemoryStream(content) + { + protected override void Dispose(bool disposing) + { + base.Dispose(disposing); + if (disposing) + { + cancellation.Cancel(); + } + } + } + private sealed class FailingSpoolFileStream(string path) : FileStream(path, new FileStreamOptions { Mode = FileMode.CreateNew, @@ -2009,11 +2360,11 @@ private sealed class Fixture : IDisposable private string? failingPatchId; private HttpStatusCode failingStatus = HttpStatusCode.OK; - public Fixture() + public Fixture(bool localEmulator = false) { var client = Substitute.For(); var database = Substitute.For(); - client.Endpoint.Returns(new Uri("https://unit-test.documents.azure.com")); + client.Endpoint.Returns(new Uri(localEmulator ? "https://localhost:8081" : "https://unit-test.documents.azure.com")); client.GetDatabase("db").Returns(database); database.GetContainer("items").Returns(this.Container); this.Shell.State = new ConnectedState(client); diff --git a/CosmosDBShell.Tests/CommandTests/ConnectCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ConnectCommandTests.cs index af70823c..b6acffa0 100644 --- a/CosmosDBShell.Tests/CommandTests/ConnectCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ConnectCommandTests.cs @@ -36,6 +36,23 @@ public void CreateClientOptions_Emulator_UsesShortRequestTimeout() Assert.Equal(TimeSpan.FromSeconds(5), options.RequestTimeout); } + [Fact] + public void LocalEmulatorInactivityTimeout_IsInitializedOnSharedAndNewInstances() + { + using var shell = ShellInterpreter.CreateInstance(); + + Assert.Equal(TimeSpan.FromSeconds(10), ShellInterpreter.Instance.LocalEmulatorInactivityTimeout); + Assert.Equal(TimeSpan.FromSeconds(10), shell.LocalEmulatorInactivityTimeout); + } + + [Fact] + public void ReportOperationProgress_WithoutRunningCommand_IsNoOp() + { + using var shell = ShellInterpreter.CreateInstance(); + + shell.ReportOperationProgress(); + } + [Fact] public void ConnectivityFailure_LocalEmulator_DisconnectsShell() { diff --git a/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs index 2f31b78f..cc565f77 100644 --- a/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ExportCommandTests.cs @@ -122,6 +122,54 @@ public async Task WriteCsvAsync_WritesHeaderUnionAndRows() Assert.Equal("\"2\",\"\",\"Seattle\"", lines[2]); } + [Fact] + public async Task WriteCsvAsync_ReportsEachRowOfBothPasses() + { + var items = ToAsyncEnumerableAsync( + JsonSerializer.SerializeToElement(new { id = "1" }), + JsonSerializer.SerializeToElement(new { id = "2" }), + JsonSerializer.SerializeToElement(new { id = "3" })); + using var writer = new StringWriter(); + var rows = 0; + + var count = await ExportCommand.WriteCsvAsync(items, writer, ',', CancellationToken.None, onRowWritten: () => rows++); + + Assert.Equal(3, count); + Assert.Equal(6, rows); + } + + [Theory] + [InlineData((int)ExportFormat.Json, 1)] + [InlineData((int)ExportFormat.Csv, 2)] + public async Task WriteFileAsync_ReportsActivityDuringEveryPass(int format, int passes) + { + var directory = Directory.CreateTempSubdirectory("cosmos-export-progress-"); + var rows = 0; + try + { + var count = await ExportCommand.WriteFileAsync( + ItemsAsync(), (ExportFormat)format, Path.Join(directory.FullName, "items"), false, + TestContext.Current.CancellationToken, () => rows++); + + Assert.Equal(3, count); + Assert.Equal(3 * passes, rows); + } + finally + { + directory.Delete(recursive: true); + } + + async IAsyncEnumerable ItemsAsync() + { + for (var index = 0; index < 3; index++) + { + Assert.Equal(index, rows); + yield return JsonSerializer.SerializeToElement(index); + await Task.Yield(); + } + } + } + [Fact] public async Task WriteCsvAsync_EscapesSeparatorsQuotesAndNewlines() { diff --git a/CosmosDBShell.Tests/CommandTests/ImportConcurrencyTests.cs b/CosmosDBShell.Tests/CommandTests/ImportConcurrencyTests.cs index d41e3bf9..3483ccdd 100644 --- a/CosmosDBShell.Tests/CommandTests/ImportConcurrencyTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ImportConcurrencyTests.cs @@ -16,6 +16,147 @@ namespace CosmosShell.Tests.CommandTests; public class ImportConcurrencyTests { + [Theory] + [InlineData(0)] + [InlineData(1)] + [InlineData(2)] + public async Task SlowInput_RecordsSettledWritesAndSettlesReadBeforeDisposal(int failure) + { + using var timeout = CreateTimeout(); + var container = Substitute.For(); + var write = new TaskCompletionSource>(TaskCreationOptions.RunContinuationsAsynchronously); + ConfigureWrites(container, (_, _) => write.Task); + var reading = NewGate(); + var releaseRead = NewGate(); + var recorded = NewGate(); + var disposed = false; + var progress = new ImportCommand.ImportProgress(); + async IAsyncEnumerable<(int LineNumber, JsonElement Item)> SlowSource() + { + try + { + yield return (1, JsonSerializer.SerializeToElement(new { id = "first" })); + reading.TrySetResult(); + await releaseRead.Task.WaitAsync(timeout.Token); + } + finally + { + disposed = true; + } + } + + var import = ImportCommand.WriteItemsAsync( + SlowSource(), container, ImportMode.Insert, 2, false, [], new CommandState(), timeout.Token, + progress: progress, + progressChanged: () => recorded.TrySetResult()); + try + { + await reading.Task.WaitAsync(timeout.Token); + if (failure == 1) + { + write.SetException(new InvalidOperationException("write failed")); + } + else if (failure == 2) + { + write.SetCanceled(timeout.Token); + } + else + { + write.SetResult(Response(HttpStatusCode.Created, 2.5)); + await recorded.Task.WaitAsync(timeout.Token); + Assert.Equal(1, progress.Success); + Assert.Equal(0, progress.Failed); + Assert.Equal(2.5, progress.Charge); + } + + Assert.False(import.IsCompleted); + Assert.False(disposed); + } + finally + { + write.TrySetResult(Response(HttpStatusCode.Created, 2.5)); + releaseRead.TrySetResult(); + } + + if (failure == 1) + { + var error = await Assert.ThrowsAsync(() => import.WaitAsync(timeout.Token)); + Assert.Equal("write failed", error.Message); + } + else if (failure == 2) + { + await Assert.ThrowsAnyAsync(() => import.WaitAsync(timeout.Token)); + } + else + { + Assert.Equal((1, 0, 2.5), await import.WaitAsync(timeout.Token)); + } + + Assert.True(disposed); + } + + [Fact] + public async Task FinalDrain_RecordsWritesThatFinishBehindASlowerWrite() + { + using var timeout = CreateTimeout(); + var container = Substitute.For(); + var gates = new[] { NewGate(), NewGate() }; + var started = new[] { NewGate(), NewGate() }; + ConfigureWrites(container, async (item, token) => + { + var index = item.GetProperty("id").GetInt32(); + started[index].TrySetResult(); + await gates[index].Task.WaitAsync(token); + return Response(HttpStatusCode.Created, 1); + }); + using var reader = NewReader(2); + var progress = new ImportCommand.ImportProgress(); + var secondRecorded = NewGate(); + var import = ImportCommand.WriteItemsAsync( + ImportCommand.EnumerateJsonLinesAsync(reader, timeout.Token), + container, ImportMode.Insert, 3, false, [], new CommandState(), timeout.Token, + progress: progress, + progressChanged: () => secondRecorded.TrySetResult()); + try + { + await Task.WhenAll(started.Select(gate => gate.Task)).WaitAsync(timeout.Token); + gates[1].TrySetResult(); + + await secondRecorded.Task.WaitAsync(timeout.Token); + Assert.Equal(1, progress.Success); + Assert.False(import.IsCompleted); + } + finally + { + gates[0].TrySetResult(); + gates[1].TrySetResult(); + } + + Assert.Equal((2, 0, 2.0), await import.WaitAsync(timeout.Token)); + } + + [Fact] + public async Task ProgressUpdatesIncludeEverySettledWriteAndObservedCharge() + { + using var timeout = CreateTimeout(); + var container = Substitute.For(); + ConfigureWrites(container, (item, _) => Task.FromResult( + Response(item.GetProperty("id").GetInt32() == 1 ? HttpStatusCode.Conflict : HttpStatusCode.Created, 2.5))); + using var reader = NewReader(3); + var progress = new ImportCommand.ImportProgress(); + var snapshots = new List<(int Success, int Failed, double Charge)>(); + + await ImportCommand.WriteItemsAsync( + ImportCommand.EnumerateJsonLinesAsync(reader, timeout.Token), + container, ImportMode.Insert, 3, true, [], new CommandState(), timeout.Token, + progress: progress, + progressChanged: () => snapshots.Add((progress.Success, progress.Failed, progress.Charge))); + + Assert.Equal(3, snapshots.Count); + Assert.Equal((2, 1, 7.5), snapshots[^1]); + Assert.Equal([1, 2, 3], snapshots.Select(snapshot => snapshot.Success + snapshot.Failed)); + } + [Theory] [InlineData(1)] [InlineData(3)] diff --git a/CosmosDBShell.Tests/CommandTests/ImportFailureResultTests.cs b/CosmosDBShell.Tests/CommandTests/ImportFailureResultTests.cs index 71fcdbe7..323da8f6 100644 --- a/CosmosDBShell.Tests/CommandTests/ImportFailureResultTests.cs +++ b/CosmosDBShell.Tests/CommandTests/ImportFailureResultTests.cs @@ -15,10 +15,133 @@ namespace CosmosShell.Tests.CommandTests; using Azure.Data.Cosmos.Shell.Util; using Microsoft.Azure.Cosmos; using NSubstitute; +using CosmosShell.Tests.Integration; [Collection(ConsoleOutputTestCollection.Name)] public class ImportFailureResultTests { + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task EmulatorTimeout_SynchronizesStructuredClassificationAndPreservesSettledWrites(bool block) + { + var file = Path.GetTempFileName(); + var errors = Path.GetTempFileName(); + var output = Path.GetTempFileName(); + using var shell = CreateShell(out var container); + Assert.IsAssignableFrom(shell.State).Client!.Endpoint.Returns(new Uri("https://localhost:8081")); + shell.Options = new Program.CosmosShellOptions { Output = "json", Quiet = true }; + shell.LocalEmulatorInactivityTimeout = TimeSpan.FromMilliseconds(100); + shell.ErrOutRedirect = errors; + shell.StdOutRedirect = output; + var calls = 0; + container.CreateItemStreamAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(async call => + { + if (++calls > 1) + { + await Task.Delay(Timeout.InfiniteTimeSpan, call.Arg()); + } + + return Response(2); + }); + try + { + await File.WriteAllTextAsync(file, """[{"id":"first","pk":"x"},{"id":"second","pk":"x"}]""", TestContext.Current.CancellationToken); + var command = $"import {JsonSerializer.Serialize(file)} --concurrency=1"; + if (block) + { + command = $"{{ {command}; echo downstream }}"; + } + + var state = Assert.IsType(await shell.ExecuteCommandAsync(command, TestContext.Current.CancellationToken)); + + Assert.Equal(ShellExitCode.ConnectionError, state.ExitCode); + Assert.IsType(shell.State); + var result = Assert.IsType(state.Result).Value; + Assert.Equal(ShellExitCode.ConnectionError, result.GetProperty("exitCode").GetInt32()); + Assert.Equal(state.Exception.Message, result.GetProperty("error").GetString()); + Assert.Contains("timed out", state.Exception.Message, StringComparison.OrdinalIgnoreCase); + Assert.Equal(state.Exception.Message, IntegrationTestBase.GetErrorMessage(state)); + Assert.Equal(1, result.GetProperty("imported").GetInt32()); + Assert.Equal(0, result.GetProperty("failed").GetInt32()); + Assert.Equal(2, result.GetProperty("requestCharge").GetDouble()); + Assert.Equal(3.25, state.RequestCharge); + Assert.Empty(await File.ReadAllTextAsync(output, TestContext.Current.CancellationToken)); + using var error = JsonDocument.Parse(await File.ReadAllTextAsync(errors, TestContext.Current.CancellationToken)); + Assert.True(JsonElement.DeepEquals(result, error.RootElement.GetProperty("result"))); + var mcp = McpResponseFactory.CreateSuccess(state, shell.State); + Assert.True(mcp.IsError); + Assert.True(JsonElement.DeepEquals(result, Assert.IsType(mcp.StructuredContent).GetProperty("result"))); + } + finally + { + shell.StdOutRedirect = null; + shell.ErrOutRedirect = null; + File.Delete(file); + File.Delete(errors); + File.Delete(output); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task StartupCancellation_PreservesSettledWritesAndRendersPartialResult(bool cancelInFlight) + { + var file = Path.GetTempFileName(); + var errors = Path.GetTempFileName(); + var output = Path.GetTempFileName(); + using var shell = CreateShell(out var container); + shell.Options = new Program.CosmosShellOptions { Output = "json", Quiet = true }; + shell.ErrOutRedirect = errors; + shell.StdOutRedirect = output; + var calls = 0; + container.CreateItemStreamAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(call => + { + calls++; + if (calls == (cancelInFlight ? 2 : 1)) + { + shell.CancelPrompt(); + if (cancelInFlight) + { + return Task.FromCanceled(call.Arg()); + } + } + + return Task.FromResult(Response(2)); + }); + try + { + await File.WriteAllTextAsync(file, """[{"id":"first","pk":"x"},{"id":"second","pk":"x"},{"id":"third","pk":"x"}]""", TestContext.Current.CancellationToken); + + var (state, interrupted) = await Program.ExecuteStartupCommandAsync(shell, $"import {JsonSerializer.Serialize(file)} --concurrency=1"); + + Assert.IsType(state); + Assert.True(interrupted); + Assert.True(state.IsError); + var result = Assert.IsType(state.Result).Value; + Assert.Equal(1, result.GetProperty("imported").GetInt32()); + Assert.Equal(0, result.GetProperty("failed").GetInt32()); + Assert.Equal(2, result.GetProperty("requestCharge").GetDouble()); + Assert.Equal(3.25, state.RequestCharge); + Assert.Equal(cancelInFlight ? 2 : 1, calls); + Assert.Empty(await File.ReadAllTextAsync(output, TestContext.Current.CancellationToken)); + using var error = JsonDocument.Parse(await File.ReadAllTextAsync(errors, TestContext.Current.CancellationToken)); + Assert.Equal("error", error.RootElement.GetProperty("status").GetString()); + Assert.True(JsonElement.DeepEquals(result, error.RootElement.GetProperty("result"))); + } + finally + { + shell.StdOutRedirect = null; + shell.ErrOutRedirect = null; + File.Delete(file); + File.Delete(errors); + File.Delete(output); + } + } + [Theory] [InlineData(429, ShellExitCode.Throttled, false, false)] [InlineData(429, ShellExitCode.Throttled, true, false)] diff --git a/CosmosDBShell.Tests/CommandTests/IngestionProgressTests.cs b/CosmosDBShell.Tests/CommandTests/IngestionProgressTests.cs new file mode 100644 index 00000000..5883cb20 --- /dev/null +++ b/CosmosDBShell.Tests/CommandTests/IngestionProgressTests.cs @@ -0,0 +1,307 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.CommandTests; + +using Azure.Data.Cosmos.Shell.Commands; +using Azure.Data.Cosmos.Shell.Core; + +[Collection(ConsoleOutputTestCollection.Name)] +public class IngestionProgressTests +{ + [Theory] + [InlineData(null, null, false, false, false, true, true)] + [InlineData("user", null, false, false, false, true, true)] + [InlineData("table", null, false, false, false, true, true)] + [InlineData("json", null, false, false, false, true, false)] + [InlineData("csv", null, false, false, false, true, false)] + [InlineData(null, "json", false, false, false, true, false)] + [InlineData(null, "csv", false, false, false, true, false)] + [InlineData("json", "user", false, false, false, true, false)] + [InlineData("user", null, true, false, false, true, false)] + [InlineData("user", null, false, true, false, true, false)] + [InlineData("user", null, false, false, true, true, false)] + [InlineData("user", null, false, false, false, false, false)] + [InlineData(null, null, false, false, false, false, false)] + public void OutputPolicy( + string? globalFormat, string? localFormat, bool quiet, bool mcp, bool redirected, bool terminal, bool expected) + { + // The documented environment default would otherwise change the cases without a format. + var previousFormat = Environment.GetEnvironmentVariable("COSMOSDB_SHELL_FORMAT"); + Environment.SetEnvironmentVariable("COSMOSDB_SHELL_FORMAT", null); + try + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = globalFormat, Quiet = quiet }; + shell.IsMcpInvocation = mcp; + shell.StdOutRedirect = redirected ? "result.json" : null; + shell.IsOutputTerminal = () => terminal; + var state = new CommandState(); + state.SetFormat(localFormat); + + Assert.Equal(expected, IngestionProgress.IsEnabled(shell, state)); + } + finally + { + Environment.SetEnvironmentVariable("COSMOSDB_SHELL_FORMAT", previousFormat); + } + } + + [Theory] + [InlineData("command-import-progress", "3 items processed", "2 imported", "1 failed", "items/s")] + [InlineData("command-bulk-progress", "3 of 10 operations processed", "2 succeeded", "1 failed", "operations/s")] + [InlineData("command-import-validation-progress", "3 items validated offline", "no writes sent", "Dry run:", "items/s")] + [InlineData("command-rm-progress", "3 items scanned", "2 deleted", "Remove:", "RU charge")] + [InlineData("command-rm-dry-run-progress", "3 items scanned", "2 would be deleted", "Dry run:", "RU charge")] + public async Task ReportsWhileWorkIsStalledAndStopsBeforeReturning( + string key, string count, string succeeded, string failed, string rate) + { + using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + timeout.CancelAfter(TimeSpan.FromSeconds(15)); + using var shell = CreateTerminalShell(); + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var output = new CompletingWriter(completion, completeOnWrite: true); + var original = Console.Out; + try + { + Console.SetOut(output); + var reporter = new IngestionProgress(shell, new CommandState(), key, 10); + reporter.Update(3, 2, 1, 7.5); + + await reporter.RunAsync(() => completion.Task.WaitAsync(timeout.Token)); + + var text = output.ToString(); + Assert.Contains(count, text); + Assert.Contains(succeeded, text); + Assert.Contains(failed, text); + Assert.Contains("elapsed", text); + Assert.Contains(rate, text); + Assert.DoesNotContain("{ $", text); + if (key != "command-import-validation-progress") + { + Assert.Contains("7.50", text); + } + } + finally + { + completion.TrySetResult(); + Console.SetOut(original); + } + } + + [Fact] + public async Task FastWorkIsSilentAndFailuresAndCancellationPropagate() + { + using var shell = CreateTerminalShell(); + using var output = new StringWriter(); + var original = Console.Out; + try + { + Console.SetOut(output); + var reporter = new IngestionProgress(shell, new CommandState(), "command-import-progress"); + await reporter.RunAsync(() => Task.CompletedTask); + var failure = new IOException("input failed"); + Assert.Same(failure, await Assert.ThrowsAsync( + () => reporter.RunAsync(() => Task.FromException(failure)))); + await Assert.ThrowsAnyAsync( + () => reporter.RunAsync(() => Task.FromCanceled(new CancellationToken(true)))); + Assert.Empty(output.ToString()); + } + finally + { + Console.SetOut(original); + } + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ReportingFailure_DoesNotChangeOperationOutcome(bool operationFails) + { + using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + timeout.CancelAfter(TimeSpan.FromSeconds(15)); + using var shell = CreateTerminalShell(); + var writeAttempted = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var output = new FailingWriter(writeAttempted); + var original = Console.Out; + var operationFailure = new InvalidOperationException("write failed"); + try + { + Console.SetOut(output); + var reporter = new IngestionProgress(shell, new CommandState(), "command-import-progress"); + + async Task OperationAsync() + { + await writeAttempted.Task.WaitAsync(timeout.Token); + if (operationFails) + { + throw operationFailure; + } + } + + if (operationFails) + { + Assert.Same(operationFailure, await Assert.ThrowsAsync(() => reporter.RunAsync(OperationAsync))); + } + else + { + await reporter.RunAsync(OperationAsync); + } + + Assert.True(writeAttempted.Task.IsCompleted); + } + finally + { + writeAttempted.TrySetResult(); + Console.SetOut(original); + } + } + [Fact] + public async Task Terminal_RedrawsOneLineInPlaceAndClearsItBeforeReturning() + { + using var timeout = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + timeout.CancelAfter(TimeSpan.FromSeconds(15)); + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = "user" }; + shell.IsOutputTerminal = () => true; + shell.ConsoleWidth = () => 1000; + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var output = new CompletingWriter(completion, completeOnWrite: true); + var original = Console.Out; + try + { + Console.SetOut(output); + var reporter = new IngestionProgress(shell, new CommandState(), "command-import-progress"); + reporter.Update(3, 2, 1, 7.5); + + await reporter.RunAsync(() => completion.Task.WaitAsync(timeout.Token)); + + var text = output.ToString(); + Assert.StartsWith("\rImport: 3 items processed", text); + Assert.DoesNotContain('\n', text); + Assert.Equal(0, output.Lines); + var drawn = text.IndexOf('\r', 1); + Assert.Equal("\r" + new string(' ', drawn - 1) + "\r", text[drawn..]); + } + finally + { + completion.TrySetResult(); + Console.SetOut(original); + } + } + + [Fact] + public void TransientLine_PadsShorterTextAndClearsOnce() + { + using var shell = ShellInterpreter.CreateInstance(); + using var output = new StringWriter(); + var original = Console.Out; + try + { + Console.SetOut(output); + shell.WriteTransientLine("abcdef"); + shell.WriteTransientLine("xy"); + shell.ClearTransientLine(); + shell.ClearTransientLine(); + + // The padding already blanks the longer text, so clearing covers only "xy". + Assert.Equal("\rabcdef\rxy \r \r", output.ToString()); + } + finally + { + Console.SetOut(original); + } + } + + [Fact] + public void TransientLine_UsesCompactTextWhenFullTextDoesNotFitAndTruncatesAsLastResort() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.ConsoleWidth = () => 12; + using var output = new StringWriter(); + var original = Console.Out; + try + { + Console.SetOut(output); + shell.WriteTransientLine("full status that is too long", "compact"); + shell.WriteTransientLine("full status that is too long", "compact but still long"); + shell.WriteTransientLine("fits"); + + Assert.Equal("\rcompact\rcompact but\rfits ", output.ToString()); + } + finally + { + Console.SetOut(original); + } + } + + [Fact] + public async Task ReportsDuringSynchronousWorkBeforeTheFirstAwait() + { + using var shell = CreateTerminalShell(); + var drawn = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + using var output = new CompletingWriter(drawn, completeOnWrite: true); + var original = Console.Out; + try + { + Console.SetOut(output); + var reporter = new IngestionProgress(shell, new CommandState(), "command-import-progress"); + + await reporter.RunAsync(() => + { + // Blocks the calling thread like a long synchronous scan of piped input. + Assert.True(drawn.Task.Wait(TimeSpan.FromSeconds(15), TestContext.Current.CancellationToken), "No progress was drawn during synchronous work."); + return Task.CompletedTask; + }); + + Assert.Contains("Import:", output.ToString()); + } + finally + { + drawn.TrySetResult(); + Console.SetOut(original); + } + } + + private static ShellInterpreter CreateTerminalShell() + { + var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = "user" }; + shell.IsOutputTerminal = () => true; + + // The real console width would select compact messages on narrow terminals. + shell.ConsoleWidth = () => 1000; + return shell; + } + + private sealed class FailingWriter(TaskCompletionSource attempted) : StringWriter + { + public override void Write(string? value) + { + attempted.TrySetResult(); + throw new IOException("console unavailable"); + } + } + + private sealed class CompletingWriter(TaskCompletionSource completion, bool completeOnWrite = false) : StringWriter + { + internal int Lines { get; private set; } + + public override void Write(string? value) + { + base.Write(value); + if (completeOnWrite) + { + completion.TrySetResult(); + } + } + + public override void WriteLine(string? value) + { + base.WriteLine(value); + this.Lines++; + completion.TrySetResult(); + } + } +} diff --git a/CosmosDBShell.Tests/CommandTests/RmCommandTests.cs b/CosmosDBShell.Tests/CommandTests/RmCommandTests.cs index 4c07885e..80ac6bd8 100644 --- a/CosmosDBShell.Tests/CommandTests/RmCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/RmCommandTests.cs @@ -11,6 +11,7 @@ namespace CosmosShell.Tests.CommandTests; using Azure.Data.Cosmos.Shell.States; using Azure.Data.Cosmos.Shell.Util; using Microsoft.Azure.Cosmos; +using NSubstitute; public class RmCommandTests { @@ -20,6 +21,92 @@ private static CosmosClient CreateTestClient() return new CosmosClient(connectionString); } + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task PipedScan_ObservesCancellationBeforeFiltering(bool dryRun, bool singleItem) + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + var database = Substitute.For(); + var container = Substitute.For(); + client.GetDatabase("db").Returns(database); + database.GetContainer("items").Returns(container); + var databaseResponse = Substitute.For(); + databaseResponse.StatusCode.Returns(System.Net.HttpStatusCode.OK); + database.ReadAsync(Arg.Any(), Arg.Any()).Returns(databaseResponse); + var containerResponse = Substitute.For(); + containerResponse.StatusCode.Returns(System.Net.HttpStatusCode.OK); + containerResponse.Resource.Returns(new ContainerProperties("items", "/pk")); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + var reads = 0; + container.ReadContainerAsync(Arg.Any(), Arg.Any()).Returns(_ => + { + if (++reads == 2) + { + cancellation.Cancel(); + } + + return Task.FromResult(containerResponse); + }); + shell.State = new ConnectedState(client); + using var input = JsonDocument.Parse(singleItem ? """{"id":"excluded","pk":"a"}""" : """[{"id":"excluded","pk":"a"}]"""); + var state = new CommandState + { + Result = new ShellJson(input.RootElement.Clone()), + }; + var command = new RmCommand { Pattern = "match-*", Key = "id", DryRun = dryRun, Database = "db", Container = "items" }; + + await Assert.ThrowsAnyAsync( + () => command.ExecuteAsync(shell, state, "rm", cancellation.Token)); + await container.DidNotReceiveWithAnyArgs().DeleteItemAsync(default!, default, default, TestContext.Current.CancellationToken); + } + + [Fact] + public async Task PatternScan_StopsBeforeNextRowWhenDeletionCancelsOperation() + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + var database = Substitute.For(); + var container = Substitute.For(); + client.GetDatabase("db").Returns(database); + database.GetContainer("items").Returns(container); + var databaseResponse = Substitute.For(); + databaseResponse.StatusCode.Returns(System.Net.HttpStatusCode.OK); + database.ReadAsync(Arg.Any(), Arg.Any()).Returns(databaseResponse); + var metadata = Substitute.For(); + metadata.StatusCode.Returns(System.Net.HttpStatusCode.OK); + metadata.Resource.Returns(new ContainerProperties("items", "/pk")); + container.ReadContainerAsync(Arg.Any(), Arg.Any()).Returns(metadata); + using var iterator = Substitute.For(); + iterator.HasMoreResults.Returns(true, false); + container.GetItemQueryStreamIterator(Arg.Any(), Arg.Any(), Arg.Any()).Returns(iterator); + using var response = new ResponseMessage(System.Net.HttpStatusCode.OK) + { + Content = new MemoryStream(System.Text.Encoding.UTF8.GetBytes( + """{"Documents":[{"id":"first","pk":"x"},{"id":"second","pk":"x"},{"id":"third","pk":"x"}]}""")), + }; + iterator.ReadNextAsync(Arg.Any()).Returns(response); + using var cancellation = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + var deleted = Substitute.For>(); + container.DeleteItemAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(_ => + { + cancellation.Cancel(); + return Task.FromResult(deleted); + }); + shell.State = new ContainerState("items", "db", client); + var command = new RmCommand { Pattern = "*", PartitionKeyArgument = "\"x\"" }; + + await Assert.ThrowsAnyAsync( + () => command.ExecuteAsync(shell, new CommandState(), "rm *", cancellation.Token)); + + await container.Received(1).DeleteItemAsync("first", Arg.Any(), Arg.Any(), Arg.Any()); + Assert.Single(container.ReceivedCalls(), call => call.GetMethodInfo().Name == nameof(Container.DeleteItemAsync)); + } + [Fact] public async Task ExecuteAsync_NoPatternAndNoPipeInput_ThrowsCommandException() { diff --git a/CosmosDBShell.Tests/CommandTests/WatchCommandTests.cs b/CosmosDBShell.Tests/CommandTests/WatchCommandTests.cs index d890fa23..2dc9671e 100644 --- a/CosmosDBShell.Tests/CommandTests/WatchCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/WatchCommandTests.cs @@ -6,7 +6,12 @@ namespace CosmosShell.Tests.CommandTests; 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 System.Net; +using System.Text; +using Microsoft.Azure.Cosmos; +using NSubstitute; /// /// Unit tests for . Covers the pure helpers that @@ -14,6 +19,7 @@ namespace CosmosShell.Tests.CommandTests; /// plus the option-validation and not-connected guards, all of which can be /// exercised without a live Cosmos DB connection. /// +[Collection(ConsoleOutputTestCollection.Name)] public class WatchCommandTests { [Fact] @@ -118,4 +124,99 @@ public async Task ExecuteAsync_Disconnected_ThrowsNotConnected() await Assert.ThrowsAsync( () => command.ExecuteAsync(shell, new CommandState(), "watch", TestContext.Current.CancellationToken)); } + + [Theory] + [InlineData(false, false, false)] + [InlineData(true, false, false)] + [InlineData(true, true, false)] + [InlineData(true, true, true)] + [InlineData(true, true, true, "{0} | echo downstream")] + [InlineData(true, true, true, "{{ {0}; echo downstream }}")] + [InlineData(true, true, true, "$capture = ({0}); echo downstream")] + [InlineData(true, true, true, "echo ({0})")] + [InlineData(true, true, true, "def stop {{ {0} }}; $capture = (stop)")] + [InlineData(true, false, false, "$capture = ({0})")] + public async Task PollingDelay_IsPausedButStalledRequestsTimeOut(bool stalled, bool callerCancels, bool startup, string commandPattern = "{0}") + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + var database = Substitute.For(); + var container = Substitute.For(); + client.Endpoint.Returns(new Uri("https://localhost:8081")); + client.GetDatabase("db").Returns(database); + database.GetContainer("items").Returns(container); + var databaseResponse = Substitute.For(); + databaseResponse.StatusCode.Returns(HttpStatusCode.OK); + database.ReadAsync(Arg.Any(), Arg.Any()).Returns(databaseResponse); + var containerResponse = Substitute.For(); + containerResponse.StatusCode.Returns(HttpStatusCode.OK); + container.ReadContainerAsync(Arg.Any(), Arg.Any()).Returns(containerResponse); + shell.State = new ConnectedState(client); + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + shell.LocalEmulatorInactivityTimeout = callerCancels ? TimeSpan.FromSeconds(30) : TimeSpan.FromMilliseconds(500); + using var caller = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + if (callerCancels && !startup) + { + caller.CancelAfter(TimeSpan.FromMilliseconds(50)); + } + using var iterator = Substitute.For(); + container.GetChangeFeedStreamIterator( + Arg.Any(), Arg.Any(), Arg.Any()).Returns(iterator); + var polls = 0; + iterator.ReadNextAsync(Arg.Any()).Returns(async call => + { + if (stalled) + { + if (startup) + { + shell.CancelPrompt(); + } + + await Task.Delay(System.Threading.Timeout.InfiniteTimeSpan, call.Arg()); + } + + return ++polls == 1 + ? new ResponseMessage(HttpStatusCode.NotModified) + : new ResponseMessage(HttpStatusCode.OK) + { + Content = new MemoryStream(Encoding.UTF8.GetBytes("""{"Documents":[{"id":"1"}]}""")), + }; + }); + + var watch = System.Diagnostics.Stopwatch.StartNew(); + shell.SetVariable("capture", new ShellText("unchanged")); + var command = string.Format(commandPattern, "watch --db db --con items --interval=1 --max=1"); + CommandState state; + if (startup) + { + var startupResult = await Program.ExecuteStartupCommandAsync(shell, command); + state = startupResult.State; + Assert.True(startupResult.Interrupted); + } + else + { + state = await shell.ExecuteCommandAsync(command, caller.Token); + } + + if (callerCancels) + { + Assert.False(state.IsError); + Assert.True(state.UserCancellationHandled); + Assert.Equal("unchanged", Assert.IsType(shell.GetVariable("capture")).Text); + Assert.IsType(shell.State); + } + else if (stalled) + { + Assert.True(state.IsError); + Assert.Equal(ShellExitCode.ConnectionError, state.ExitCode); + Assert.IsType(shell.State); + } + else + { + Assert.False(state.IsError); + Assert.True(watch.Elapsed > shell.LocalEmulatorInactivityTimeout); + Assert.Equal(2, polls); + Assert.IsType(shell.State); + } + } } diff --git a/CosmosDBShell.Tests/Integration/AzureIngestionBenchmarkTests.cs b/CosmosDBShell.Tests/Integration/AzureIngestionBenchmarkTests.cs new file mode 100644 index 00000000..5bc5c865 --- /dev/null +++ b/CosmosDBShell.Tests/Integration/AzureIngestionBenchmarkTests.cs @@ -0,0 +1,196 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.Integration; + +using System.Diagnostics; +using System.Text; +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; +using Azure.Data.Cosmos.Shell.Util; +using Azure.Identity; +using Microsoft.Azure.Cosmos; + +/// +/// Opt-in ingestion benchmark against a real Azure Cosmos DB account, using the same workload and +/// matrix as . It signs in with the Azure CLI identity, needs a +/// data-plane write role plus ARM rights to create and delete databases, and incurs provisioned +/// throughput charges. It creates one temporary database and container and always deletes them. +/// +[Trait("Category", "LiveAzure")] +public class AzureIngestionBenchmarkTests : IntegrationTestBase +{ + private const string EndpointVariable = "COSMOSDB_SHELL_AZURE_BENCHMARK_ENDPOINT"; + private const string SubscriptionVariable = "COSMOSDB_SHELL_AZURE_BENCHMARK_SUBSCRIPTION"; + private const string ResourceGroupVariable = "COSMOSDB_SHELL_AZURE_BENCHMARK_RESOURCE_GROUP"; + private const string OutputVariable = "COSMOSDB_SHELL_AZURE_BENCHMARK_CSV"; + private const string ThroughputVariable = "COSMOSDB_SHELL_AZURE_BENCHMARK_RU"; + + [Fact] + public async Task MeasureImportAndBulkOnAzure() + { + var endpoint = Environment.GetEnvironmentVariable(EndpointVariable); + var subscription = Environment.GetEnvironmentVariable(SubscriptionVariable); + var resourceGroup = Environment.GetEnvironmentVariable(ResourceGroupVariable); + var output = Environment.GetEnvironmentVariable(OutputVariable); + Assert.SkipWhen( + string.IsNullOrWhiteSpace(endpoint) || string.IsNullOrWhiteSpace(subscription) + || string.IsNullOrWhiteSpace(resourceGroup) || string.IsNullOrWhiteSpace(output), + $"Set {EndpointVariable}, {SubscriptionVariable}, {ResourceGroupVariable}, and {OutputVariable} to run the opt-in Azure benchmark."); + var throughput = int.TryParse(Environment.GetEnvironmentVariable(ThroughputVariable), out var ru) ? ru : 4000; + + var token = TestContext.Current.CancellationToken; + Shell.Options = new Program.CosmosShellOptions { Output = "json" }; + Shell.IsInteractiveSession = () => false; + await Shell.ConnectAsync( + endpoint!, + credentialMethod: CredentialMethod.AzureCli, + subscriptionId: subscription, + resourceGroupName: resourceGroup, + token: token); + + var databaseName = $"IngestionBenchmark_{Guid.NewGuid():N}"; + var inputPath = Path.GetTempFileName(); + + // Set before mkdb: a failed or canceled create can still have been accepted by Azure. + var creationAttempted = false; + var creationCompleted = false; + Exception? failure = null; + try + { + creationAttempted = true; + + // Cancellation must count as failure here: a canceled ARM create may still complete, so + // only a create that returned without error or cancellation is treated as completed. + var database = await Shell.ExecuteCommandAsync($"mkdb {databaseName} --scale m --ru 400", token, cancellationIsError: true); + Assert.False(database.IsError, FormatError(database)); + token.ThrowIfCancellationRequested(); + creationCompleted = true; + var container = await this.RunScriptAsync($"mkcon items /pk --database {databaseName} --scale m --ru {throughput}"); + Assert.False(container.IsError, FormatError(container)); + var navigation = await this.RunScriptAsync($"cd /{databaseName}/items"); + Assert.False(navigation.IsError, FormatError(navigation)); + + using var verifier = new CosmosClient(endpoint, new AzureCliCredential()); + var items = verifier.GetContainer(databaseName, "items"); + using var results = new StreamWriter(Path.GetFullPath(output!), append: false, new UTF8Encoding(false)) { AutoFlush = true }; + await results.WriteLineAsync(IngestionBenchmarkWorkload.CsvHeader); + foreach (var size in IngestionBenchmarkWorkload.Sizes) + { + foreach (var command in IngestionBenchmarkWorkload.Commands) + { + foreach (var concurrency in IngestionBenchmarkWorkload.Concurrencies) + { + // Each case writes uniquely prefixed IDs into the shared container. + var caseId = $"{command}-{size}-{concurrency}"; + await IngestionBenchmarkWorkload.WriteInputAsync(inputPath, size, command == "bulk", caseId, token); + var text = command == "import" + ? $"import {ShellQuotedPath(inputPath)} --concurrency={concurrency}" + : $"bulk run {ShellQuotedPath(inputPath)} --concurrency={concurrency} --yes"; + var watch = Stopwatch.StartNew(); + var state = await Shell.RunCommandAsync(new CommandState(), text, token); + watch.Stop(); + Assert.False(state.IsError, FormatError(state)); + var result = Assert.IsType(state.Result).Value; + Assert.Equal(size, result.GetProperty(command == "import" ? "imported" : "succeeded").GetInt64()); + Assert.Equal(0, result.GetProperty("failed").GetInt64()); + + var query = new QueryDefinition("SELECT VALUE COUNT(1) FROM c WHERE c.caseId = @caseId").WithParameter("@caseId", caseId); + using var iterator = items.GetItemQueryIterator(query); + long verified = 0; + while (iterator.HasMoreResults) + { + verified += (await iterator.ReadNextAsync(token)).Sum(); + } + + Assert.Equal(size, verified); + await results.WriteLineAsync(IngestionBenchmarkWorkload.FormatRow( + command, size, new FileInfo(inputPath).Length, concurrency, watch.Elapsed, result, verified)); + } + } + } + } + catch (Exception ex) + { + failure = ex; + throw; + } + finally + { + File.Delete(inputPath); + if (creationAttempted) + { + var cleanupProblem = await this.DeleteDatabaseAsync(databaseName, creationCompleted); + if (cleanupProblem is not null) + { + // Keep the benchmark's own failure, if any, as the inner exception. + throw new InvalidOperationException(cleanupProblem, failure); + } + } + } + } + + /// + /// Deletes the benchmark database so provisioned throughput stops billing, and returns a problem + /// description when deletion cannot be confirmed. A completed create is deleted with retries. After + /// a failed or canceled create, the outcome is unknown: the database is watched and deleted if it + /// appears, but absence is never reported as success because the create could still complete. + /// Cancellation of cleanup commands counts as failure so incomplete cleanup is not hidden. + /// + private async Task DeleteDatabaseAsync(string databaseName, bool creationCompleted) + { + var attempts = creationCompleted ? 12 : 30; + var notFound = MessageService.GetString( + "command-rmdb-error-database_not_found", + new Dictionary { { "db", databaseName } }); + using var cleanupTimeout = new CancellationTokenSource(TimeSpan.FromMinutes(10)); + var token = cleanupTimeout.Token; + var lastError = "no attempt completed"; + var manualCheck = $"Check the account for {databaseName} and delete it manually to stop throughput charges."; + + // rmdb is rejected inside a container, so return to the account root first. + var root = await Shell.ExecuteCommandAsync("cd /", token, cancellationIsError: true); + if (root.IsError) + { + return $"Cleanup could not navigate to the account root ({FormatError(root)}). {manualCheck}"; + } + + for (var attempt = 1; attempt <= attempts; attempt++) + { + var probe = await Shell.ExecuteCommandAsync($"rmdb {databaseName} --dry-run", token, cancellationIsError: true); + if (probe is ErrorCommandState { Exception.Message: var probeMessage } && probeMessage == notFound) + { + if (creationCompleted) + { + // The database existed, so not found means it is already deleted. + return null; + } + + lastError = "not found yet; the canceled or failed create may still complete"; + } + else if (probe.IsError) + { + lastError = FormatError(probe); + } + else + { + var delete = await Shell.ExecuteCommandAsync($"rmdb {databaseName} true", token, cancellationIsError: true); + if (!delete.IsError) + { + return null; + } + + // A create that is still running can reject the delete; retry after it settles. + lastError = FormatError(delete); + } + + await Task.Delay(TimeSpan.FromSeconds(10), token); + } + + return creationCompleted + ? $"Cleanup could not confirm that {databaseName} was deleted ({lastError}). {manualCheck}" + : $"Cleanup is unconfirmed: database creation did not complete, and {databaseName} did not appear within {attempts * 10} seconds ({lastError}). {manualCheck}"; + } +} diff --git a/CosmosDBShell.Tests/Integration/IngestionBenchmarkTests.cs b/CosmosDBShell.Tests/Integration/IngestionBenchmarkTests.cs new file mode 100644 index 00000000..3368fe1f --- /dev/null +++ b/CosmosDBShell.Tests/Integration/IngestionBenchmarkTests.cs @@ -0,0 +1,85 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.Integration; + +using System.Diagnostics; +using System.Globalization; +using System.Text; +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; +using Microsoft.Azure.Cosmos.Linq; + +public class IngestionBenchmarkTests : ConnectedEmulatorTestBase +{ + private const string OutputVariable = "COSMOSDB_SHELL_INGESTION_BENCHMARK_CSV"; + + public override async ValueTask InitializeAsync() + { + Assert.SkipWhen( + string.IsNullOrWhiteSpace(Environment.GetEnvironmentVariable(OutputVariable)), + $"Set {OutputVariable} to a results CSV path to run the opt-in emulator benchmark."); + await base.InitializeAsync(); + } + + [Fact] + public async Task MeasureImportAndBulk() + { + var token = TestContext.Current.CancellationToken; + var outputPath = Path.GetFullPath(Environment.GetEnvironmentVariable(OutputVariable)!); + using var results = new StreamWriter(outputPath, append: false, new UTF8Encoding(false)) { AutoFlush = true }; + await results.WriteLineAsync(IngestionBenchmarkWorkload.CsvHeader); + var databaseName = $"IngestionBenchmark_{Guid.NewGuid():N}"; + CreatedDatabases.Add(databaseName); + var database = (await this.CosmosClient.CreateDatabaseAsync(databaseName, cancellationToken: token)).Database; + Shell.Options = new Program.CosmosShellOptions { Output = "json" }; + Shell.IsInteractiveSession = () => false; + var inputPath = Path.GetTempFileName(); + try + { + foreach (var size in new[] { 1_000, 10_000, 100_000 }) + { + foreach (var command in new[] { "import", "bulk" }) + { + await IngestionBenchmarkWorkload.WriteInputAsync(inputPath, size, command == "bulk", caseId: null, token); + foreach (var concurrency in new[] { 1, 16, 32 }) + { + var containerName = $"items-{command}-{size}-{concurrency}"; + var container = (await database.CreateContainerAsync( + containerName, "/pk", throughput: 10_000, cancellationToken: token)).Container; + try + { + var navigation = await this.ExecuteAsync($"cd /{databaseName}/{containerName}"); + Assert.False(navigation.IsError, FormatError(navigation)); + var text = command == "import" + ? $"import {ShellQuotedPath(inputPath)} --concurrency={concurrency}" + : $"bulk run {ShellQuotedPath(inputPath)} --concurrency={concurrency} --yes"; + var watch = Stopwatch.StartNew(); + // Use the parsed command runner so timings stay comparable with docs/ingestion.md. + var state = await Shell.RunCommandAsync(new CommandState(), text, token); + watch.Stop(); + Assert.False(state.IsError, FormatError(state)); + var result = Assert.IsType(state.Result).Value; + Assert.Equal(size, result.GetProperty(command == "import" ? "imported" : "succeeded").GetInt64()); + Assert.Equal(0, result.GetProperty("failed").GetInt64()); + var verified = (await container.GetItemLinqQueryable().CountAsync(token)).Resource; + Assert.Equal(size, verified); + await results.WriteLineAsync(IngestionBenchmarkWorkload.FormatRow( + command, size, new FileInfo(inputPath).Length, concurrency, watch.Elapsed, result, verified)); + } + finally + { + await container.DeleteContainerAsync(cancellationToken: token); + } + } + } + } + } + finally + { + File.Delete(inputPath); + } + } +} diff --git a/CosmosDBShell.Tests/Integration/IngestionBenchmarkWorkload.cs b/CosmosDBShell.Tests/Integration/IngestionBenchmarkWorkload.cs new file mode 100644 index 00000000..3a9d67c4 --- /dev/null +++ b/CosmosDBShell.Tests/Integration/IngestionBenchmarkWorkload.cs @@ -0,0 +1,65 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.Integration; + +using System.Globalization; +using System.Text; +using System.Text.Json; + +/// +/// The shared ingestion benchmark workload: 1 KiB JSON items spread evenly across 1,000 +/// partition-key values, and the CSV row format for results. +/// +internal static class IngestionBenchmarkWorkload +{ + internal const int ItemBytes = 1024; + + internal const string CsvHeader = "command,items,itemBytes,inputBytes,concurrency,elapsedSeconds,itemsPerSecond,requestCharge,verifiedItems"; + + internal static readonly int[] Sizes = [1_000, 10_000, 100_000]; + + internal static readonly string[] Commands = ["import", "bulk"]; + + internal static readonly int[] Concurrencies = [1, 16, 32]; + + /// + /// Writes a JSON Lines input file. When is set, IDs are prefixed with + /// it and items carry a caseId property, so several cases can share one container. + /// + internal static async Task WriteInputAsync(string path, int size, bool bulk, string? caseId, CancellationToken token) + { + using var writer = new StreamWriter(path, append: false, new UTF8Encoding(false)) { NewLine = "\n" }; + string? payload = null; + for (var index = 0; index < size; index++) + { + token.ThrowIfCancellationRequested(); + var id = caseId is null + ? index.ToString("D8", CultureInfo.InvariantCulture) + : $"{caseId}-{index.ToString("D8", CultureInfo.InvariantCulture)}"; + var pk = $"p{index % 1000:D4}"; + payload ??= new string('x', ItemBytes - Encoding.UTF8.GetByteCount(Serialize(id, pk, caseId, string.Empty))); + var item = Serialize(id, pk, caseId, payload); + Assert.Equal(ItemBytes, Encoding.UTF8.GetByteCount(item)); + await writer.WriteLineAsync(bulk ? $$"""{"op":"create","item":{{item}}}""" : item); + } + } + + internal static string FormatRow(string command, int size, long inputBytes, int concurrency, TimeSpan elapsed, JsonElement result, long verified) => + string.Join( + ",", + command, + size, + ItemBytes, + inputBytes, + concurrency, + elapsed.TotalSeconds.ToString("F3", CultureInfo.InvariantCulture), + (size / elapsed.TotalSeconds).ToString("F1", CultureInfo.InvariantCulture), + result.GetProperty("requestCharge").GetDouble().ToString("F2", CultureInfo.InvariantCulture), + verified); + + private static string Serialize(string id, string pk, string? caseId, string payload) => caseId is null + ? JsonSerializer.Serialize(new { id, pk, payload }) + : JsonSerializer.Serialize(new { id, pk, caseId, payload }); +} diff --git a/CosmosDBShell.Tests/Integration/IntegrationTestBase.cs b/CosmosDBShell.Tests/Integration/IntegrationTestBase.cs index acbad524..07ec956d 100644 --- a/CosmosDBShell.Tests/Integration/IntegrationTestBase.cs +++ b/CosmosDBShell.Tests/Integration/IntegrationTestBase.cs @@ -39,7 +39,7 @@ internal static string GetText(CommandState state) internal static string GetErrorMessage(CommandState state) { - return Assert.IsType(state).Exception.Message; + return Assert.IsAssignableFrom(state).Exception.Message; } internal static string FormatError(CommandState state) @@ -76,6 +76,12 @@ internal void SetVariable(string name, ShellObject value) /// internal static string ShellPath(string path) => path.Replace('\\', '/'); + /// + /// Formats a file-system path as a single-quoted shell argument, doubling embedded + /// apostrophes as the lexer requires (for example in a user profile named O'Brien). + /// + internal static string ShellQuotedPath(string path) => "'" + ShellPath(path).Replace("'", "''") + "'"; + public virtual void Dispose() { foreach (var file in tempFiles) diff --git a/CosmosDBShell.Tests/Integration/LocalEmulatorInactivityTimeoutTests.cs b/CosmosDBShell.Tests/Integration/LocalEmulatorInactivityTimeoutTests.cs new file mode 100644 index 00000000..241cdb8c --- /dev/null +++ b/CosmosDBShell.Tests/Integration/LocalEmulatorInactivityTimeoutTests.cs @@ -0,0 +1,134 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.Integration; + +using System.Diagnostics; +using System.Globalization; +using System.Text.Json; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.States; +using Azure.Data.Cosmos.Shell.Util; +using Microsoft.Azure.Cosmos; +using Microsoft.Azure.Cosmos.Linq; + +public class LocalEmulatorInactivityTimeoutTests : ConnectedEmulatorTestBase +{ + private string databaseName = null!; + private Container container = null!; + private string inputPath = null!; + + public override async ValueTask InitializeAsync() + { + await base.InitializeAsync(); + var token = TestContext.Current.CancellationToken; + this.databaseName = $"InactivityTimeout_{Guid.NewGuid():N}"; + CreatedDatabases.Add(this.databaseName); + var database = (await this.CosmosClient.CreateDatabaseAsync(this.databaseName, cancellationToken: token)).Database; + this.container = (await database.CreateContainerAsync("items", "/pk", cancellationToken: token)).Container; + var options = ShellInterpreter.CreateClientOptions(ConnectionMode.Gateway, isEmulator: true); + options.ServerCertificateCustomValidationCallback = (_, _, _) => true; + options.CustomHandlers.Add(new OperationProgressHandler(Shell)); + options.CustomHandlers.Add(new PacedRequestHandler()); + Shell.Connect(new CosmosClient(ParsedDocDBConnectionString.BuildEmulatorConnectionString(EmulatorEndpoint), options)); + Shell.Options = new Program.CosmosShellOptions { Output = "json" }; + var navigation = await this.ExecuteAsync($"cd /{this.databaseName}/items"); + Assert.False(navigation.IsError, FormatError(navigation)); + this.inputPath = Path.GetTempFileName(); + await File.WriteAllLinesAsync( + this.inputPath, + Enumerable.Range(0, 300).Select(index => JsonSerializer.Serialize(new + { + id = index.ToString(CultureInfo.InvariantCulture), + pk = $"p{index % 10}", + })), + token); + } + + public override void Dispose() + { + if (this.inputPath is not null) + { + File.Delete(this.inputPath); + } + + base.Dispose(); + } + + [Fact] + public async Task Import_LongerThanTimeoutButMakingProgress_Completes() + { + Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(1); + + var watch = Stopwatch.StartNew(); + var state = await this.ExecuteAsync($"import {ShellQuotedPath(this.inputPath)} --concurrency=1"); + watch.Stop(); + + Assert.False(state.IsError, FormatError(state)); + Assert.True(watch.Elapsed > Shell.LocalEmulatorInactivityTimeout, $"Import finished in {watch.Elapsed}; the test did not exceed the timeout."); + Assert.Equal(300, GetJson(state).GetProperty("imported").GetInt32()); + Assert.IsNotType(Shell.State); + var count = await this.container.GetItemLinqQueryable().CountAsync(TestContext.Current.CancellationToken); + Assert.Equal(300, count.Resource); + } + + [Fact] + public async Task RmWildcard_LongerThanTimeoutButMakingProgress_Completes() + { + var import = await this.ExecuteAsync($"import {ShellQuotedPath(this.inputPath)}"); + Assert.False(import.IsError, FormatError(import)); + Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(1); + + var watch = Stopwatch.StartNew(); + var state = await this.ExecuteAsync("rm *"); + watch.Stop(); + + Assert.False(state.IsError, FormatError(state)); + Assert.True(watch.Elapsed > Shell.LocalEmulatorInactivityTimeout, $"rm finished in {watch.Elapsed}; the test did not exceed the timeout."); + Assert.IsNotType(Shell.State); + var count = await this.container.GetItemLinqQueryable().CountAsync(TestContext.Current.CancellationToken); + Assert.Equal(0, count.Resource); + } + + [Fact] + public async Task Script_LongerThanTimeout_CountsEveryEmulatorResponseAsProgress() + { + Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(1); + var script = string.Join( + Environment.NewLine, + Enumerable.Range(0, 200).Select(index => $$"""mkitem '{"id":"script-{{index}}","pk":"s{{index % 5}}"}'""")); + + var watch = Stopwatch.StartNew(); + var state = await this.ExecuteAsync(script); + watch.Stop(); + + Assert.False(state.IsError, FormatError(state)); + Assert.True(watch.Elapsed > Shell.LocalEmulatorInactivityTimeout, $"Script finished in {watch.Elapsed}; the test did not exceed the timeout."); + Assert.IsNotType(Shell.State); + var count = await this.container.GetItemLinqQueryable().CountAsync(TestContext.Current.CancellationToken); + Assert.Equal(200, count.Resource); + } + + [Fact] + public async Task Command_WithoutProgress_StillTimesOut() + { + Shell.LocalEmulatorInactivityTimeout = TimeSpan.FromMilliseconds(1); + + var state = await this.ExecuteAsync($"import {ShellQuotedPath(this.inputPath)} --concurrency=1"); + + Assert.True(state.IsError); + Assert.Contains("timed out", GetErrorMessage(state), StringComparison.OrdinalIgnoreCase); + Assert.IsType(Shell.State); + } + + private sealed class PacedRequestHandler : RequestHandler + { + public override async Task SendAsync(RequestMessage request, CancellationToken cancellationToken) + { + // Make the elapsed-time assertions independent of emulator or CI machine speed. + await Task.Delay(TimeSpan.FromMilliseconds(20), cancellationToken); + return await base.SendAsync(request, cancellationToken); + } + } +} diff --git a/CosmosDBShell.Tests/Shell/OperationInactivityTimeoutTests.cs b/CosmosDBShell.Tests/Shell/OperationInactivityTimeoutTests.cs new file mode 100644 index 00000000..ea75a55c --- /dev/null +++ b/CosmosDBShell.Tests/Shell/OperationInactivityTimeoutTests.cs @@ -0,0 +1,398 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.Shell; + +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Commands; +using Azure.Data.Cosmos.Shell.Parser; +using Azure.Data.Cosmos.Shell.States; +using Microsoft.Azure.Cosmos; +using NSubstitute; +using System.Text.Json; + +public class OperationInactivityTimeoutTests +{ + private static readonly TimeSpan Timeout = TimeSpan.FromSeconds(10); + + [Fact] + public void WithoutProgress_CancelsAfterTimeout() + { + var clock = new ManualTimeProvider(); + using var source = new CancellationTokenSource(global::System.Threading.Timeout.InfiniteTimeSpan, clock); + _ = new OperationInactivityTimeout(source, Timeout); + + clock.Advance(Timeout - TimeSpan.FromMilliseconds(1)); + Assert.False(source.IsCancellationRequested); + + clock.Advance(TimeSpan.FromMilliseconds(1)); + Assert.True(source.IsCancellationRequested); + } + + [Fact] + public void Progress_RestartsTimeout() + { + var clock = new ManualTimeProvider(); + using var source = new CancellationTokenSource(global::System.Threading.Timeout.InfiniteTimeSpan, clock); + var deadline = new OperationInactivityTimeout(source, Timeout); + + // Three gaps of 0.9x the timeout exceed it in total, but each restarts the period. + for (var i = 0; i < 3; i++) + { + clock.Advance(Timeout * 0.9); + deadline.ReportProgress(); + } + + Assert.False(source.IsCancellationRequested); + clock.Advance(Timeout); + Assert.True(source.IsCancellationRequested); + } + + [Fact] + public void NestedPauses_SuspendUntilLastResume_AndIgnoreProgress() + { + var clock = new ManualTimeProvider(); + using var source = new CancellationTokenSource(global::System.Threading.Timeout.InfiniteTimeSpan, clock); + var deadline = new OperationInactivityTimeout(source, Timeout); + + var outer = deadline.Pause(); + var inner = deadline.Pause(); + deadline.ReportProgress(); + clock.Advance(Timeout * 3); + Assert.False(source.IsCancellationRequested); + + inner.Dispose(); + inner.Dispose(); + clock.Advance(Timeout * 3); + Assert.False(source.IsCancellationRequested); + + outer.Dispose(); + clock.Advance(Timeout - TimeSpan.FromMilliseconds(1)); + Assert.False(source.IsCancellationRequested); + clock.Advance(TimeSpan.FromMilliseconds(1)); + Assert.True(source.IsCancellationRequested); + } + + [Fact] + public void ProgressAndPause_AfterSourceDisposed_DoNotThrow() + { + var source = new CancellationTokenSource(); + var deadline = new OperationInactivityTimeout(source, Timeout); + source.Dispose(); + + deadline.ReportProgress(); + deadline.Pause().Dispose(); + } + + [Fact] + public void ShellWithoutRunningCommand_PauseIsNoOp() + { + using var shell = ShellInterpreter.CreateInstance(); + + shell.PauseOperationTimeout().Dispose(); + shell.ReportOperationProgress(); + } + + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task ConsumedCancellation_OnlyInternalTimeoutFailsAndDisconnects(bool callerCancels) + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + client.Endpoint.Returns(new Uri("https://localhost:8081")); + shell.State = new ConnectedState(client); + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + shell.LocalEmulatorInactivityTimeout = callerCancels ? TimeSpan.FromSeconds(30) : TimeSpan.FromMilliseconds(50); + using var caller = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + if (callerCancels) + { + caller.CancelAfter(TimeSpan.FromMilliseconds(50)); + } + + Assert.True(CommandFactory.TryCreateFactory(typeof(ConsumedCancellationCommand), out var factory)); + shell.App.Commands["consumed-cancellation-test"] = factory; + + var state = await shell.ExecuteCommandAsync("consumed-cancellation-test", caller.Token); + + if (callerCancels) + { + Assert.False(state.IsError); + Assert.IsType(shell.State); + } + else + { + var error = Assert.IsType(state); + Assert.Equal(ShellExitCode.ConnectionError, error.ExitCode); + Assert.Contains("timed out", error.Exception.Message, StringComparison.OrdinalIgnoreCase); + Assert.IsType(shell.State); + } + } + + [CosmosCommand("consumed-cancellation-test")] + internal sealed class ConsumedCancellationCommand : CosmosCommand + { + public override async Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + try + { + await Task.Delay(global::System.Threading.Timeout.InfiniteTimeSpan, token); + } + catch (OperationCanceledException) when (token.IsCancellationRequested) + { + commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { marker = "partial-success-marker" })); + return commandState; + } + + throw new InvalidOperationException("The stalled command should be canceled."); + } + } + + [CosmosCommand("local-processing-test", External = true)] + internal sealed class LocalProcessingCommand : CosmosCommand + { + public override async Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + await Task.Delay(TimeSpan.FromSeconds(1), token); + return commandState; + } + } + + [Theory] + [InlineData(true)] + [InlineData(false)] + public async Task Import_LocalValidationCountsAsActivityButUnavailableResponsesDoNot(bool localValidation) + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + client.Endpoint.Returns(new Uri("https://localhost:8081")); + var database = Substitute.For(); + var container = Substitute.For(); + client.GetDatabase("db").Returns(database); + database.GetContainer("items").Returns(container); + using var unavailable = Substitute.For(); + unavailable.StatusCode.Returns(System.Net.HttpStatusCode.ServiceUnavailable); + container.CreateItemStreamAsync(Arg.Any(), Arg.Any(), Arg.Any(), Arg.Any()) + .Returns(_ => Task.FromResult(unavailable)); + shell.State = new ConnectedState(client); + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + + // A manual clock keeps this deterministic: each item advances time by 0.6x the limit, so + // the run exceeds the limit in total but never between two items that report activity. + shell.LocalEmulatorTimeProvider = new ManualTimeProvider(); + shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(10); + Assert.True(CommandFactory.TryCreateFactory(typeof(ImportValidationProgressCommand), out var factory)); + shell.App.Commands["import-validation-progress-test"] = factory; + + var state = await shell.ExecuteCommandAsync( + "import-validation-progress-test" + (localValidation ? " --local" : ""), TestContext.Current.CancellationToken); + + if (localValidation) + { + Assert.False(state.IsError, state is ErrorCommandState error ? error.Exception.ToString() : state.GetType().Name); + Assert.IsType(shell.State); + Assert.Equal(ImportValidationProgressCommand.ItemCount, Assert.IsType(state.Result).Value.GetProperty("failed").GetInt32()); + } + else + { + Assert.True(state.IsError); + Assert.Equal(ShellExitCode.ConnectionError, state.ExitCode); + Assert.IsType(shell.State); + } + } + + [CosmosCommand("import-validation-progress-test")] + internal sealed class ImportValidationProgressCommand : CosmosCommand + { + internal const int ItemCount = 20; + + [CosmosOption("local")] + public bool? Local { get; set; } + + public override async Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + var container = ((ConnectedState)shell.State).Client.GetDatabase("db").GetContainer("items"); + var clock = (ManualTimeProvider)shell.LocalEmulatorTimeProvider!; + var step = shell.LocalEmulatorInactivityTimeout * 0.6; + async IAsyncEnumerable<(int LineNumber, JsonElement Item)> Items() + { + for (var index = 0; index < ItemCount; index++) + { + await Task.Yield(); + clock.Advance(step); + yield return (index + 1, JsonSerializer.SerializeToElement(new + { + id = index.ToString(System.Globalization.CultureInfo.InvariantCulture), + pk = this.Local == true ? JsonSerializer.SerializeToElement(Array.Empty()) : JsonSerializer.SerializeToElement("key"), + })); + } + } + + var result = await ImportCommand.WriteItemsAsync( + Items(), container, ImportMode.Insert, 1, true, ["pk"], commandState, token, warningShell: shell); + commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { failed = result.Failed })); + return commandState; + } + } + + [Theory] + [InlineData("echo data | local-processing-test")] + [InlineData("$result = (local-processing-test)")] + [InlineData("echo (local-processing-test)")] + public async Task ExternalCommand_LongerThanDeadline_DoesNotTimeOut(string command) + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + client.Endpoint.Returns(new Uri("https://localhost:8081")); + shell.State = new ConnectedState(client); + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + shell.LocalEmulatorInactivityTimeout = TimeSpan.FromMilliseconds(500); + Assert.True(CommandFactory.TryCreateFactory(typeof(LocalProcessingCommand), out var factory)); + shell.App.Commands["local-processing-test"] = factory; + + var state = await shell.ExecuteCommandAsync(command, TestContext.Current.CancellationToken); + + Assert.False(state.IsError); + Assert.IsType(shell.State); + } + + [Theory] + [InlineData(false, false)] + [InlineData(false, true)] + [InlineData(true, false)] + [InlineData(true, true)] + public async Task ConsumedTimeout_DoesNotRenderPartialSuccessInsideScriptOrBlock(bool script, bool callerCancels) + { + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + client.Endpoint.Returns(new Uri("https://localhost:8081")); + shell.State = new ConnectedState(client); + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + shell.LocalEmulatorInactivityTimeout = callerCancels ? TimeSpan.FromSeconds(30) : TimeSpan.FromMilliseconds(50); + using var caller = CancellationTokenSource.CreateLinkedTokenSource(TestContext.Current.CancellationToken); + if (callerCancels) + { + caller.CancelAfter(TimeSpan.FromMilliseconds(50)); + } + + Assert.True(CommandFactory.TryCreateFactory(typeof(ConsumedCancellationCommand), out var factory)); + shell.App.Commands["consumed-cancellation-test"] = factory; + var directory = Directory.CreateTempSubdirectory("cosmos-timeout-render-"); + var scriptPath = Path.Join(Environment.CurrentDirectory, $"timeout-{Guid.NewGuid():N}.csh"); + try + { + var output = Path.Join(directory.FullName, "result.json"); + var statement = $"consumed-cancellation-test > '{output.Replace("'", "''")}'"; + string command; + if (script) + { + await File.WriteAllTextAsync(scriptPath, statement, TestContext.Current.CancellationToken); + command = Path.GetFileName(scriptPath); + } + else + { + command = $"{{ {statement} }}"; + } + + var state = await shell.ExecuteCommandAsync(command, caller.Token); + + Assert.Equal(!callerCancels, state.IsError); + if (!callerCancels) + { + Assert.Equal(ShellExitCode.ConnectionError, state.ExitCode); + Assert.IsType(shell.State); + Assert.False(File.Exists(output), "Partial success was written before reporting the timeout."); + } + else + { + Assert.IsType(shell.State); + } + } + finally + { + File.Delete(scriptPath); + directory.Delete(recursive: true); + } + } + + [Fact] + public void EmulatorTimeout_KeepsScriptFramesOutermostAndClassifiesAsConnectionError() + { + var canceled = new OperationCanceledException("The operation was canceled."); + var inner = new PositionalException("inner.csh", canceled, 3, 1, "bulk run plan.jsonl --yes"); + var outer = new PositionalException("outer.csh", inner, 7, 5, "inner.csh"); + + var timeout = ShellInterpreter.CreateEmulatorTimeoutException(outer); + + var frames = PositionalException.GetSourceTrace(timeout); + Assert.Equal(["inner.csh", "outer.csh"], frames.Select(frame => frame.FileName)); + Assert.Equal([3, 7], frames.Select(frame => frame.Line)); + Assert.IsType(timeout); + Assert.Contains("timed out", timeout.Message, StringComparison.OrdinalIgnoreCase); + Assert.Equal(ShellExitCode.ConnectionError, ShellExitCode.FromException(timeout)); + Assert.Same(canceled, timeout.InnerException!.InnerException!.InnerException!.InnerException); + } + + [Fact] + public void EmulatorTimeout_WithoutScriptFrame_ClassifiesAsConnectionError() + { + var timeout = ShellInterpreter.CreateEmulatorTimeoutException(new OperationCanceledException("The operation was canceled.")); + + Assert.IsType(timeout); + Assert.IsType(timeout.InnerException); + Assert.Contains("timed out", timeout.Message, StringComparison.OrdinalIgnoreCase); + Assert.Equal(ShellExitCode.ConnectionError, ShellExitCode.FromException(timeout)); + } + + /// + /// A deterministic clock: timers fire only when moves time past them. + /// + private sealed class ManualTimeProvider : TimeProvider + { + private readonly List timers = []; + private TimeSpan now; + + public override ITimer CreateTimer(TimerCallback callback, object? state, TimeSpan dueTime, TimeSpan period) + { + var timer = new ManualTimer(this, callback, state); + timer.Change(dueTime, period); + this.timers.Add(timer); + return timer; + } + + public void Advance(TimeSpan delta) + { + this.now += delta; + foreach (var timer in this.timers.ToArray().Where(timer => timer.DueAt is { } due && due <= this.now)) + { + timer.DueAt = null; + timer.Fire(); + } + } + + private sealed class ManualTimer(ManualTimeProvider owner, TimerCallback callback, object? state) : ITimer + { + public TimeSpan? DueAt { get; set; } + + public bool Change(TimeSpan dueTime, TimeSpan period) + { + this.DueAt = dueTime == global::System.Threading.Timeout.InfiniteTimeSpan ? null : owner.now + dueTime; + return true; + } + + public void Fire() => callback(state); + + public void Dispose() => this.DueAt = null; + + public ValueTask DisposeAsync() + { + this.Dispose(); + return ValueTask.CompletedTask; + } + } + } +} diff --git a/CosmosDBShell.Tests/UtilTest/ExternalEditorTests.cs b/CosmosDBShell.Tests/UtilTest/ExternalEditorTests.cs index 5ce533fe..c6759fa6 100644 --- a/CosmosDBShell.Tests/UtilTest/ExternalEditorTests.cs +++ b/CosmosDBShell.Tests/UtilTest/ExternalEditorTests.cs @@ -5,6 +5,10 @@ namespace CosmosShell.Tests.UtilTest; using Azure.Data.Cosmos.Shell.Util; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.States; +using Microsoft.Azure.Cosmos; +using NSubstitute; // Several tests below mutate process-wide environment variables (VISUAL/EDITOR). // xUnit parallelizes across test classes, so this collection disables @@ -18,6 +22,51 @@ public sealed class ExternalEditorTestCollection [Collection(ExternalEditorTestCollection.Name)] public class ExternalEditorTests { + [Theory] + [InlineData(false)] + [InlineData(true)] + public async Task FileAndThemeEditing_PauseTheExecutingShellDeadline(bool theme) + { + Assert.SkipUnless(OperatingSystem.IsWindows(), "Uses the Windows PowerShell fake editor."); + var directory = Directory.CreateTempSubdirectory("cosmos-editor-timeout-"); + var visual = Environment.GetEnvironmentVariable("VISUAL"); + var savedTheme = Theme.Current; + try + { + var script = Path.Join(directory.FullName, "editor.ps1"); + await File.WriteAllTextAsync(script, "param($path)\nStart-Sleep -Milliseconds 1500\nexit 0\n", TestContext.Current.CancellationToken); + var executable = Path.Join(Environment.GetFolderPath(Environment.SpecialFolder.System), "WindowsPowerShell", "v1.0", "powershell.exe"); + Environment.SetEnvironmentVariable("VISUAL", $"\"{executable}\" -NoProfile -NonInteractive -ExecutionPolicy Bypass -File \"{script}\""); + var path = Path.Join(directory.FullName, theme ? "test.toml" : "test.csh"); + if (theme) + { + ThemeFile.Save("editor-timeout-test", Theme.Current, path); + } + + using var shell = ShellInterpreter.CreateInstance(); + using var client = Substitute.For(); + client.Endpoint.Returns(new Uri("https://localhost:8081")); + shell.State = new ConnectedState(client); + shell.IsInputRedirected = () => false; + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + shell.LocalEmulatorInactivityTimeout = TimeSpan.FromSeconds(1); + var command = $"{(theme ? "theme edit" : "edit")} '{path.Replace("'", "''")}'"; + var watch = System.Diagnostics.Stopwatch.StartNew(); + + var state = await shell.ExecuteCommandAsync(command, TestContext.Current.CancellationToken); + + Assert.False(state.IsError, state is ErrorCommandState error ? error.Exception.ToString() : string.Empty); + Assert.True(watch.Elapsed > shell.LocalEmulatorInactivityTimeout); + Assert.IsType(shell.State); + } + finally + { + Environment.SetEnvironmentVariable("VISUAL", visual); + Theme.Apply(savedTheme); + directory.Delete(recursive: true); + } + } + [Fact] public void Resolve_ExplicitEditor_TakesPrecedence() { diff --git a/CosmosDBShell.Tests/UtilTest/StartupOptionTests.cs b/CosmosDBShell.Tests/UtilTest/StartupOptionTests.cs index e09cea50..72588142 100644 --- a/CosmosDBShell.Tests/UtilTest/StartupOptionTests.cs +++ b/CosmosDBShell.Tests/UtilTest/StartupOptionTests.cs @@ -5,9 +5,488 @@ namespace CosmosShell.Tests.UtilTest; using Azure.Data.Cosmos.Shell.Util; +using Azure.Data.Cosmos.Shell.Commands; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Parser; +using System.Reflection; +using System.Text.Json; +[Collection(ConsoleOutputTestCollection.Name)] public class StartupOptionTests { + [Theory] + [InlineData(ConsoleKey.C, true, true)] + [InlineData(ConsoleKey.N, false, false)] + [InlineData(ConsoleKey.Escape, false, false)] + public async Task Confirmation_InterruptStopsFollowingCommandsButDeclineDoesNot(ConsoleKey key, bool control, bool interrupted) + { + using var shell = ShellInterpreter.CreateInstance(); + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupConfirmationCommand), out var factory)); + shell.App.Commands["startup-confirm-test"] = factory; + StartupConfirmationCommand.Key = new ConsoleKeyInfo( + key == ConsoleKey.N ? 'n' : '\0', key, false, false, control); + var output = Path.GetFullPath($"startup-following-{Guid.NewGuid():N}.txt"); + try + { + var (state, wasInterrupted) = await Program.ExecuteStartupCommandAsync( + shell, $"{{ startup-confirm-test; echo downstream > {JsonSerializer.Serialize(output)} }}"); + + Assert.Equal(interrupted, wasInterrupted); + Assert.Equal(interrupted, state.IsError); + Assert.Equal(!interrupted, File.Exists(output)); + Assert.False((await shell.ExecuteCommandAsync("echo next", TestContext.Current.CancellationToken)).IsError); + } + finally + { + File.Delete(output); + } + } + + [CosmosCommand("startup-confirm-test")] + internal sealed class StartupConfirmationCommand : CosmosCommand + { + internal static ConsoleKeyInfo Key { get; set; } + + public override Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + ShellInterpreter.Confirm("yes_char", shell, () => Key); + return Task.FromResult(commandState); + } + } + + [Theory] + [InlineData("json", false, true, false, false)] + [InlineData("csv", false, true, false, false)] + [InlineData("user", false, false, false, false)] + [InlineData("table", false, true, true, false)] + [InlineData("user", true, true, false, false)] + [InlineData("user", false, true, false, true)] + public void ConsoleCancellation_KeepsMachineOutputClean( + string format, bool quiet, bool terminal, bool redirected, bool markerExpected) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = format, Quiet = quiet }; + shell.IsOutputTerminal = () => terminal; + shell.StdOutRedirect = redirected ? "not-written.txt" : null; + using var cancellation = ShellInterpreter.UserCancellationTokenSource; + using var output = new StringWriter(); + using var errors = new StringWriter(); + var savedOutput = Console.Out; + var savedError = Console.Error; + try + { + Console.SetOut(output); + Console.SetError(errors); + RaiseConsoleCancellation(shell); + + Assert.True(cancellation.IsCancellationRequested); + Assert.Equal(markerExpected ? "^C" + Environment.NewLine : "", output.ToString()); + Assert.Empty(errors.ToString()); + } + finally + { + Console.SetOut(savedOutput); + Console.SetError(savedError); + } + } + + [Theory] + [InlineData("user", null, "startup-format-test --format=json", false)] + [InlineData("user", null, "startup-format-test -f csv", false)] + [InlineData("user", null, "startup-format-test --format=user", true)] + [InlineData("user", null, "startup-format-test --format=table", true)] + [InlineData("user", null, "startup-format-test", true)] + [InlineData("json", null, "startup-format-test --format=user", false)] + [InlineData("csv", null, "startup-format-test --format=user", false)] + [InlineData(null, "json", "startup-format-test --format=user", true)] + [InlineData(null, "csv", "startup-format-test --format=table", true)] + [InlineData("user", null, "echo (startup-format-test --format=json)", false)] + [InlineData("user", null, "$captured = (startup-format-test --format=csv)", false)] + [InlineData("user", null, "startup-format-outer-test --format=json --inner-format=user", false)] + [InlineData("user", null, "startup-format-outer-test --format=user --inner-format=csv", false)] + public async Task ConsoleCancellation_UsesScopedCommandFormat( + string? sessionFormat, string? environmentFormat, string command, bool markerExpected) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = sessionFormat }; + shell.IsOutputTerminal = () => true; + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupFormatCancellationCommand), out var factory)); + shell.App.Commands["startup-format-test"] = factory; + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupOuterFormatCancellationCommand), out var outerFactory)); + shell.App.Commands["startup-format-outer-test"] = outerFactory; + using var output = new StringWriter(); + using var errors = new StringWriter(); + var savedOutput = Console.Out; + var savedError = Console.Error; + var savedEnvironmentFormat = Environment.GetEnvironmentVariable("COSMOSDB_SHELL_FORMAT"); + try + { + Environment.SetEnvironmentVariable("COSMOSDB_SHELL_FORMAT", environmentFormat); + Console.SetOut(output); + Console.SetError(errors); + + var (state, interrupted) = await Program.ExecuteStartupCommandAsync(shell, command); + + Assert.False(state.IsError); + Assert.True(interrupted); + Assert.True(state.UserCancellationHandled); + Assert.Equal(markerExpected ? "^C" + Environment.NewLine : "", output.ToString()); + Assert.Empty(errors.ToString()); + + output.GetStringBuilder().Clear(); + var (nextState, nextInterrupted) = await Program.ExecuteStartupCommandAsync( + shell, "startup-format-test --format=user"); + + Assert.False(nextState.IsError); + Assert.True(nextInterrupted); + Assert.Equal(shell.IsMachineMode ? "" : "^C" + Environment.NewLine, output.ToString()); + Assert.Empty(errors.ToString()); + } + finally + { + Environment.SetEnvironmentVariable("COSMOSDB_SHELL_FORMAT", savedEnvironmentFormat); + Console.SetOut(savedOutput); + Console.SetError(savedError); + } + } + + [CosmosCommand("startup-format-test")] + internal sealed class StartupFormatCancellationCommand : CosmosCommand + { + [CosmosOption("format", "f")] + public string? SelectedFormat { get; init; } + + public override async Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + await Task.Run(() => RaiseConsoleCancellation(shell), token); + return new CommandState { UserCancellationHandled = true }; + } + } + + [CosmosCommand("startup-format-outer-test")] + internal sealed class StartupOuterFormatCancellationCommand : CosmosCommand + { + [CosmosOption("format", "f")] + public string? SelectedFormat { get; init; } + + [CosmosOption("inner-format")] + public string? InnerFormat { get; init; } + + public override Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + => shell.ExecuteCosmosCommandAsync( + new StartupFormatCancellationCommand { SelectedFormat = this.InnerFormat }, + commandState, commandText, token); + } + + [Fact] + public async Task RedirectedInput_CancellationAbortsBlockingReadWithoutExecutingLateScript() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Quiet = true }; + using var input = new BlockingStartupReader(); + var output = Path.GetFullPath($"startup-stdin-{Guid.NewGuid():N}.txt"); + input.Script = $"echo should-not-run > {JsonSerializer.Serialize(output)}"; + try + { + var startup = Program.ExecuteStartupInputAsync(shell, input); + await input.Started.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + + RaiseConsoleCancellation(shell); + var (state, interrupted) = await startup.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + + Assert.True(interrupted); + Assert.True(Assert.IsType(state).IsError); + Assert.True(input.ReadToken.IsCancellationRequested); + input.Release.Set(); + await input.Finished.Task.WaitAsync(TimeSpan.FromSeconds(10), TestContext.Current.CancellationToken); + Assert.False(File.Exists(output)); + Assert.False((await shell.ExecuteCommandAsync("echo next", TestContext.Current.CancellationToken)).IsError); + } + finally + { + input.Release.Set(); + File.Delete(output); + } + } + + [Theory] + [InlineData("")] + [InlineData(" \r\n")] + public async Task RedirectedInput_EmptyInputPreservesInteractiveStartup(string text) + { + using var shell = ShellInterpreter.CreateInstance(); + using var input = new StringReader(text); + + var (state, interrupted) = await Program.ExecuteStartupInputAsync(shell, input); + + Assert.Null(state); + Assert.False(interrupted); + } + + [Fact] + public async Task RedirectedInput_HandledCancellationRemainsSuccessfulAndStopsScript() + { + using var shell = ShellInterpreter.CreateInstance(); + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupHandledCancellationCommand), out var factory)); + shell.App.Commands["startup-handled-test"] = factory; + using var input = new StringReader("startup-handled-test\necho downstream"); + + var (state, interrupted) = await Program.ExecuteStartupInputAsync(shell, input); + + Assert.True(interrupted); + Assert.NotNull(state); + Assert.False(state.IsError); + Assert.True(state.UserCancellationHandled); + Assert.Equal("handled", Assert.IsType(state.Result).Text); + } + + private static void RaiseConsoleCancellation(ShellInterpreter shell) + { + var args = (ConsoleCancelEventArgs)Activator.CreateInstance( + typeof(ConsoleCancelEventArgs), BindingFlags.Instance | BindingFlags.NonPublic, + binder: null, args: [ConsoleSpecialKey.ControlC], culture: null)!; + typeof(ShellInterpreter).GetMethod("Console_CancelKeyPress", BindingFlags.Instance | BindingFlags.NonPublic)! + .Invoke(shell, [null, args]); + Assert.True(args.Cancel); + } + + private sealed class BlockingStartupReader : TextReader + { + internal TaskCompletionSource Started { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + internal TaskCompletionSource Finished { get; } = new(TaskCreationOptions.RunContinuationsAsynchronously); + internal ManualResetEventSlim Release { get; } = new(); + internal string Script { get; set; } = ""; + internal CancellationToken ReadToken { get; private set; } + + public override Task ReadToEndAsync(CancellationToken cancellationToken) + { + this.ReadToken = cancellationToken; + this.Started.SetResult(); + this.Release.Wait(); + this.Finished.SetResult(); + return Task.FromResult(this.Script); + } + + protected override void Dispose(bool disposing) + { + this.Release.Set(); + base.Dispose(disposing); + } + } + + [Theory] + [InlineData("-c")] + [InlineData("-k")] + [InlineData("stdin")] + public async Task StartupCommands_ObserveConsoleCancellation(string mode) + { + using var shell = ShellInterpreter.CreateInstance(); + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupCancellationCommand), out var factory)); + shell.App.Commands["startup-cancel-test"] = factory; + string command; + if (mode == "stdin") + { + command = "startup-cancel-test\n"; + } + else + { + var (root, map) = Program.BuildRootCommand(); + var parsed = root.Parse(Program.NormalizeArguments([mode, "startup-cancel-test"])); + Assert.Empty(parsed.Errors); + command = parsed.GetValue(mode == "-c" ? map.ExecuteAndQuit : map.ExecuteAndContinue)!; + } + + var (state, interrupted) = await Program.ExecuteStartupCommandAsync(shell, command); + + var error = Assert.IsType(state); + Assert.True(interrupted); + Assert.NotEqual(ShellExitCode.Success, error.ExitCode); + Assert.True(Assert.IsAssignableFrom(error.Exception.GetBaseException()).CancellationToken.IsCancellationRequested); + } + + [CosmosCommand("startup-cancel-test")] + internal sealed class StartupCancellationCommand : CosmosCommand + { + public override Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + shell.CancelPrompt(); + commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(token.IsCancellationRequested)); + return Task.FromResult(commandState); + } + } + + [Theory] + [InlineData("startup-handled-test | echo downstream")] + [InlineData("{ startup-handled-test; echo downstream }")] + [InlineData("loop { startup-handled-test }")] + [InlineData("while true { startup-handled-test }")] + [InlineData("do { startup-handled-test } while true")] + [InlineData("for $item in [1,2] { startup-handled-test }")] + [InlineData("def stop { startup-handled-test; echo downstream }; stop")] + [InlineData("script")] + public async Task HandledCancellation_StopsEnclosingRunnersSuccessfully(string command) + { + using var shell = ShellInterpreter.CreateInstance(); + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupHandledCancellationCommand), out var factory)); + shell.App.Commands["startup-handled-test"] = factory; + var script = $"startup-cancel-{Guid.NewGuid():N}.csh"; + try + { + if (command == "script") + { + await File.WriteAllTextAsync(script, "startup-handled-test\necho downstream", TestContext.Current.CancellationToken); + command = script; + } + + var (state, interrupted) = await Program.ExecuteStartupCommandAsync(shell, command); + + Assert.False(state.IsError); + Assert.True(interrupted); + Assert.True(state.UserCancellationHandled); + Assert.Equal("handled", Assert.IsType(state.Result).Text); + } + finally + { + File.Delete(script); + } + } + + [CosmosCommand("startup-handled-test")] + internal sealed class StartupHandledCancellationCommand : CosmosCommand + { + public override Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + shell.CancelPrompt(); + commandState.UserCancellationHandled = true; + commandState.Result = new ShellText("handled"); + return Task.FromResult(commandState); + } + } + + [Theory] + [InlineData(false, false)] + [InlineData(true, false)] + [InlineData(false, true)] + [InlineData(true, true)] + public async Task UnhandledCancellation_DoesNotRedirectPartialSuccess(bool script, bool pipeline) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = "json", Quiet = true }; + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupCancellationCommand), out var factory)); + shell.App.Commands["startup-cancel-test"] = factory; + var output = Path.GetTempFileName(); + var file = $"startup-cancel-{Guid.NewGuid():N}.csh"; + const string sentinel = "previous complete result"; + try + { + await File.WriteAllTextAsync(output, sentinel, TestContext.Current.CancellationToken); + var inner = "startup-cancel-test" + (pipeline ? " | echo downstream" : ""); + var command = $"{{ {inner} > {JsonSerializer.Serialize(output)} }}"; + if (script) + { + await File.WriteAllTextAsync(file, $"{inner} > {JsonSerializer.Serialize(output)}", TestContext.Current.CancellationToken); + command = file; + } + + var (state, interrupted) = await Program.ExecuteStartupCommandAsync(shell, command); + + Assert.True(state.IsError); + Assert.True(interrupted); + Assert.Equal(sentinel, await File.ReadAllTextAsync(output, TestContext.Current.CancellationToken)); + Assert.False((await shell.ExecuteCommandAsync("echo next", TestContext.Current.CancellationToken)).IsError); + } + finally + { + File.Delete(output); + File.Delete(file); + } + } + + [Fact] + public async Task StartupCancellation_RendersStructuredPartialResults() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.Options = new Program.CosmosShellOptions { Output = "json" }; + Assert.True(CommandFactory.TryCreateFactory(typeof(StartupPartialCancellationCommand), out var factory)); + shell.App.Commands["startup-partial-test"] = factory; + using var output = new StringWriter(); + using var errors = new StringWriter(); + var savedOutput = Console.Out; + var savedError = Console.Error; + try + { + Console.SetOut(output); + Console.SetError(errors); + + var (state, interrupted) = await Program.ExecuteStartupCommandAsync(shell, "startup-partial-test"); + + Assert.IsType(state); + Assert.True(interrupted); + Assert.True(state.IsError); + Assert.Empty(output.ToString()); + using var payload = JsonDocument.Parse(errors.ToString()); + Assert.Equal("error", payload.RootElement.GetProperty("status").GetString()); + var result = payload.RootElement.GetProperty("result"); + Assert.Equal(1, result.GetProperty("succeeded").GetInt32()); + Assert.Equal(1, result.GetProperty("uncertain").GetInt32()); + Assert.True(result.GetProperty("resultIncomplete").GetBoolean()); + } + finally + { + Console.SetOut(savedOutput); + Console.SetError(savedError); + } + } + + [CosmosCommand("startup-partial-test")] + internal sealed class StartupPartialCancellationCommand : CosmosCommand + { + public override Task ExecuteAsync( + ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) + { + shell.CancelPrompt(); + return Task.FromResult(new StructuredErrorCommandState( + new OperationCanceledException(token), + new ShellJson(JsonSerializer.SerializeToElement(new { succeeded = 1, uncertain = 1, resultIncomplete = true })))); + } + } + + [Fact] + public async Task PromptRefresh_DoesNotCancelQueuedStartupWork() + { + using var shell = ShellInterpreter.CreateInstance(); + var release = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var holder = shell.RunSerializedAsync(async () => + { + await release.Task.WaitAsync(TestContext.Current.CancellationToken); + shell.RefreshPrompt(); + return 0; + }, TestContext.Current.CancellationToken); + try + { + var queued = Program.ExecuteStartupCommandAsync(shell, "echo queued"); + Assert.False(queued.IsCompleted); + release.TrySetResult(); + await holder; + + var (state, interrupted) = await queued; + + Assert.False(interrupted); + Assert.False(state.IsError); + Assert.Equal("queued", Assert.IsType(state.Result).Text); + } + finally + { + release.TrySetResult(); + await holder; + } + } + [Theory] [InlineData("--output", "JSON", "json")] [InlineData("-o", "csv", "csv")] diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BucketCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BucketCommand.cs index 8756412e..263127b6 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BucketCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BucketCommand.cs @@ -139,7 +139,7 @@ private static bool ConfirmWrite(ShellInterpreter shell, bool yes, string summar } AnsiConsole.MarkupLine(MessageService.GetArgsString(summaryKey, args)); - return ShellInterpreter.Confirm("command-bucket-confirm"); + return ShellInterpreter.Confirm("command-bucket-confirm", shell); } private static CommandState BuildShowResult(ConnectedState state, ThroughputBucketsView view) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkCommand.cs index 7b694911..2884502f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkCommand.cs @@ -4,7 +4,6 @@ namespace Azure.Data.Cosmos.Shell.Commands; -using System.Diagnostics; using System.Globalization; using System.Net; using System.Runtime.CompilerServices; @@ -63,8 +62,6 @@ internal sealed class BulkCommand : CosmosCommand private const int StatusPreviewCount = 100; - private static readonly TimeSpan ProgressInterval = TimeSpan.FromSeconds(2); - private static readonly Dictionary AllowedOptions = new() { ["run"] = ["partition-key", "database", "container", "concurrency", "max-ru", "dry-run", "yes", "continue-on-error", "journal", "retry-uncertain"], @@ -688,7 +685,7 @@ private void Confirm(ShellInterpreter shell, long count, Target target) count, "target", $"{target.DatabaseName}/{target.ContainerName}")); - if (!ShellInterpreter.Confirm("command-bulk-confirm")) + if (!ShellInterpreter.Confirm("command-bulk-confirm", shell)) { throw Error("command-bulk-error-declined"); } @@ -733,6 +730,7 @@ private async Task RunAsync(ShellInterpreter shell, CommandState c { var operation = BulkOperationNormalizer.Normalize(element, target.PartitionKeyPaths, defaultKey, spool.Count); await spool.AddAsync(operation.Json); + shell.ReportOperationProgress(); } if (spool.Count == 0) @@ -773,8 +771,10 @@ private async Task RunSelectionAsync(ShellInterpreter shell, strin var page = await iterator.ReadNextAsync(token); summary.RequestCharge += page.RequestCharge; + shell.ReportOperationProgress(); foreach (var row in page) { + token.ThrowIfCancellationRequested(); if (this.MaxItems is { } max && spool.Count >= max) { summary.SelectionLimited = true; @@ -788,6 +788,8 @@ private async Task RunSelectionAsync(ShellInterpreter shell, strin { await save.WriteLineAsync(operation.Json); } + + shell.ReportOperationProgress(); } if (this.MaxItems is { } limit && spool.Count >= limit && iterator.HasMoreResults) @@ -927,10 +929,14 @@ private async Task ExecuteOperationsAsync( this.Confirm(shell, count, target); using var journal = this.Journal is null ? null - : BulkJournal.Open(this.Journal, target.Endpoint, target.ContainerRid, hash, count); + : BulkJournal.Open(this.Journal, target.Endpoint, target.ContainerRid, hash, count, () => + { + token.ThrowIfCancellationRequested(); + shell.ReportOperationProgress(); + }); if (journal is not null) { - await journal.ValidateAsync(operations, token); + await journal.ValidateAsync(operations, token, shell.ReportOperationProgress); } var previous = new Dictionary(); @@ -964,7 +970,23 @@ private async Task ExecuteOperationsAsync( } } - var progress = Stopwatch.StartNew(); + var progress = new IngestionProgress(shell, new CommandState(), "command-bulk-progress", count); + var reportedSkips = summary.Skipped; + var reportedSuccesses = summary.Succeeded; + void UpdateProgress() + { + progress.Update(summary.Processed, summary.Succeeded, summary.Failed, summary.RequestCharge); + + // Only successful writes and local skips restart inactivity, not synthetic SDK errors. + if (summary.Skipped != reportedSkips || summary.Succeeded != reportedSuccesses) + { + reportedSkips = summary.Skipped; + reportedSuccesses = summary.Succeeded; + shell.ReportOperationProgress(); + } + } + + UpdateProgress(); async Task RecordAsync(BulkOutcome outcome) { if (memory != null) @@ -979,26 +1001,11 @@ async Task RecordAsync(BulkOutcome outcome) { await journal.RecordAsync(outcome); } - - if (outcome.Status != BulkOutcome.Started && progress.Elapsed >= ProgressInterval && !shell.IsMachineMode) - { - progress.Restart(); - ShellInterpreter.WriteLine(MessageService.GetArgsString( - "command-bulk-progress", - "processed", - summary.Processed, - "total", - count, - "failed", - summary.Failed, - "charge", - summary.RequestCharge.ToString("F2", CultureInfo.InvariantCulture))); - } } try { - await BulkExecutor.ExecuteAsync( + await progress.RunAsync(() => BulkExecutor.ExecuteAsync( operations, (operation, cancellation) => WriteAsync(target.Container, operation, cancellation), RecordAsync, @@ -1008,7 +1015,8 @@ await BulkExecutor.ExecuteAsync( this.MaxRu, this.ContinueOnError == true, this.RetryUncertain == true, - token); + token, + UpdateProgress)); } catch (Exception ex) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkExecutor.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkExecutor.cs index 983eebbd..9cdf2651 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkExecutor.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkExecutor.cs @@ -18,7 +18,8 @@ public static async Task ExecuteAsync( double? maxRu, bool continueOnError, bool retryUncertain, - CancellationToken token) + CancellationToken token, + Action? progressChanged = null) { ArgumentOutOfRangeException.ThrowIfLessThan(concurrency, 1); var pending = new List<(BulkOperation Operation, PartitionKey Key, Task Task, Guid AttemptId)>(); @@ -57,6 +58,7 @@ async Task SettleAsync(int index) blockedItems.Add((operation.Id, key)); } + progressChanged?.Invoke(); await record(result); } @@ -73,8 +75,50 @@ async Task SettleCompletedAsync() try { - await foreach (var operation in source.WithCancellation(token)) + await using var items = source.GetAsyncEnumerator(token); + async ValueTask ReadNextAsync() { + var next = items.MoveNextAsync(); + if (next.IsCompleted || pending.Count == 0) + { + return await next; + } + + var read = next.AsTask(); + try + { + while (!read.IsCompleted && pending.Count > 0) + { +#pragma warning disable VSTHRD003 // All pending tasks were started by this invocation. + await Task.WhenAny(pending.Select(entry => (Task)entry.Task).Append(read)).WaitAsync(token); +#pragma warning restore VSTHRD003 + await SettleCompletedAsync(); + } + + return await read; + } + catch (Exception ex) + { + failure ??= ExceptionDispatchInfo.Capture(ex); + throw; + } + finally + { + // Finish the read before disposing its enumerator, retaining the original failure. + try + { + await read; + } + catch (Exception ex) + { + failure ??= ExceptionDispatchInfo.Capture(ex); + } + } + } + + while (await ReadNextAsync()) + { + var operation = items.Current; var key = BulkOperationNormalizer.ParseKey(operation.PartitionKey); await SettleCompletedAsync(); @@ -102,6 +146,7 @@ async Task SettleCompletedAsync() { summary.Skipped++; summary.Processed++; + progressChanged?.Invoke(); await record(old); continue; } @@ -112,6 +157,7 @@ async Task SettleCompletedAsync() summary.Processed++; summary.AddFailure(old with { Status = BulkOutcome.Uncertain }); blockedItems.Add((operation.Id, key)); + progressChanged?.Invoke(); continue; } } @@ -121,6 +167,7 @@ async Task SettleCompletedAsync() summary.Skipped++; summary.Processed++; summary.ResultIncomplete = true; + progressChanged?.Invoke(); continue; } @@ -145,11 +192,15 @@ async Task SettleCompletedAsync() finally { // Observe every started write even when reading, cancellation, or journal I/O fails. + // Settle in completion order so progress reflects writes that finish behind a slow one. while (pending.Count > 0) { try { - await SettleAsync(0); +#pragma warning disable VSTHRD003 // All pending tasks were started by this invocation. + var completed = await Task.WhenAny(pending.Select(entry => entry.Task)); +#pragma warning restore VSTHRD003 + await SettleAsync(pending.FindIndex(entry => entry.Task == completed)); } catch (Exception ex) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkJournal.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkJournal.cs index 053c1a1c..db9e9798 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkJournal.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/BulkJournal.cs @@ -32,7 +32,7 @@ internal BulkJournal(FileStream stream, Dictionary previous) public IReadOnlyDictionary Previous { get; } - public static BulkJournal Open(string path, string endpoint, string containerRid, string hash, long operationCount) + public static BulkJournal Open(string path, string endpoint, string containerRid, string hash, long operationCount, Action? onRecordRead = null) { var fullPath = Path.GetFullPath(path); var options = new FileStreamOptions @@ -73,7 +73,7 @@ public static BulkJournal Open(string path, string endpoint, string containerRid throw Invalid(null); } - ReadExisting(stream, header, previous, completeLength < stream.Length); + ReadExisting(stream, header, previous, completeLength < stream.Length, onRecordRead); stream.SetLength(completeLength); stream.Seek(0, SeekOrigin.End); } @@ -105,7 +105,7 @@ public async Task RecordAsync(BulkOutcome outcome) } } - public async Task ValidateAsync(IAsyncEnumerable operations, CancellationToken token) + public async Task ValidateAsync(IAsyncEnumerable operations, CancellationToken token, Action? onOperationValidated = null) { if (this.Previous.Count == 0) { @@ -115,6 +115,7 @@ public async Task ValidateAsync(IAsyncEnumerable operations, Canc var validated = 0; await foreach (var operation in operations.WithCancellation(token)) { + onOperationValidated?.Invoke(); if (this.Previous.TryGetValue(operation.Index, out var outcome)) { if (outcome.Op != operation.Op || outcome.Id != operation.Id || @@ -173,7 +174,7 @@ private static long FindCompleteLength(FileStream stream) return 0; } - private static void ReadExisting(FileStream stream, Header expected, Dictionary previous, bool tornTail) + private static void ReadExisting(FileStream stream, Header expected, Dictionary previous, bool tornTail, Action? onRecordRead) { stream.Position = 0; using var reader = new StreamReader(stream, new UTF8Encoding(false), false, 65536, leaveOpen: true); @@ -187,6 +188,7 @@ private static void ReadExisting(FileStream stream, Header expected, Dictionary< while (reader.ReadLine() is { } line) { + onRecordRead?.Invoke(); if (tornTail && reader.EndOfStream) { break; diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/EditCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/EditCommand.cs index 59db19e7..eca67c47 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/EditCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/EditCommand.cs @@ -31,7 +31,7 @@ public override Task ExecuteAsync(ShellInterpreter shell, CommandS throw new CommandException("edit", MessageService.GetString("command-edit-missing-path")); } - if (Console.IsInputRedirected || !string.IsNullOrEmpty(shell.CurrentScriptFileName)) + if (shell.IsInputRedirected() || !string.IsNullOrEmpty(shell.CurrentScriptFileName)) { var message = MessageService.GetString("command-edit-not-interactive"); AnsiConsole.MarkupLine(Theme.FormatError(message)); @@ -78,6 +78,7 @@ public override Task ExecuteAsync(ShellInterpreter shell, CommandS int exitCode; try { + using var pause = shell.PauseOperationTimeout(); using var process = System.Diagnostics.Process.Start(new System.Diagnostics.ProcessStartInfo { FileName = editor.FileName, diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs index e741be72..56330885 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ExportCommand.cs @@ -92,7 +92,7 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co var max = ResultLimit.ResolveMaxItemCount(this.Max, defaultMaxItemCount: null); var format = ResolveFormat(filePath); - var (count, charge) = await ExecuteExportAsync(container, query, max, format, filePath, this.Force == true, token); + var (count, charge) = await ExecuteExportAsync(shell, container, query, max, format, filePath, this.Force == true, token); resultState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { @@ -151,14 +151,16 @@ internal static string SerializeJsonLine(JsonElement element) /// The items to write. /// The destination stream. /// Cancellation token. + /// Invoked after each item is written. /// The number of items written. - internal static async Task WriteArrayAsync(IAsyncEnumerable items, Stream stream, CancellationToken token) + internal static async Task WriteArrayAsync(IAsyncEnumerable items, Stream stream, CancellationToken token, Action? onRowWritten = null) { var count = 0; await using var writer = new Utf8JsonWriter(stream, new JsonWriterOptions { Indented = true, IndentSize = 4, NewLine = "\n" }); writer.WriteStartArray(); await foreach (var item in items.WithCancellation(token)) { + token.ThrowIfCancellationRequested(); item.WriteTo(writer); if (writer.BytesPending >= 64 * 1024) { @@ -166,6 +168,7 @@ internal static async Task WriteArrayAsync(IAsyncEnumerable it } count++; + onRowWritten?.Invoke(); } writer.WriteEndArray(); @@ -185,8 +188,9 @@ internal static async Task WriteArrayAsync(IAsyncEnumerable it /// The field separator. /// Cancellation token. /// The header for non-object rows; defaults to the environment setting. + /// Invoked after each row is written in both the spooling and final passes. /// The number of data rows written. - internal static async Task WriteCsvAsync(IAsyncEnumerable items, TextWriter writer, char separator, CancellationToken token, string? scalarColumnName = null) + internal static async Task WriteCsvAsync(IAsyncEnumerable items, TextWriter writer, char separator, CancellationToken token, string? scalarColumnName = null, Action? onRowWritten = null) { scalarColumnName ??= Environment.GetEnvironmentVariable(ScalarColumnEnvironmentVariable) ?? string.Empty; var spoolOptions = new FileStreamOptions @@ -209,7 +213,9 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item var hasScalarRows = false; await foreach (var item in items.WithCancellation(token)) { + token.ThrowIfCancellationRequested(); await spoolWriter.WriteLineAsync(SerializeJsonLine(item).AsMemory(), token); + onRowWritten?.Invoke(); if (item.ValueKind != JsonValueKind.Object) { hasScalarRows = true; @@ -287,6 +293,7 @@ internal static async Task WriteCsvAsync(IAsyncEnumerable item await writer.WriteLineAsync(sb.ToString().AsMemory(), token); count++; + onRowWritten?.Invoke(); } await writer.FlushAsync(token); @@ -299,6 +306,7 @@ private static string GetScalarCsvValue(JsonElement item) } private static async Task<(int Count, double Charge)> ExecuteExportAsync( + ShellInterpreter shell, Container container, string query, int? max, @@ -318,7 +326,20 @@ private static string GetScalarCsvValue(JsonElement item) var totalCharge = 0.0; using var iterator = container.GetItemQueryIterator(query, requestOptions: options); var count = await WriteFileAsync( - EnumerateAsync(iterator, max, charge => totalCharge += charge, token), format, filePath, overwrite, token); + EnumerateAsync( + iterator, + max, + charge => + { + totalCharge += charge; + shell.ReportOperationProgress(); + }, + token), + format, + filePath, + overwrite, + token, + shell.ReportOperationProgress); return (count, totalCharge); } catch (CosmosException ce) @@ -340,7 +361,8 @@ internal static async Task WriteFileAsync( ExportFormat format, string filePath, bool overwrite, - CancellationToken token) + CancellationToken token, + Action? onRowWritten = null) { var destination = Path.GetFullPath(filePath); @@ -379,11 +401,11 @@ internal static async Task WriteFileAsync( { using var writer = new StreamWriter(stream, new UTF8Encoding(encoderShouldEmitUTF8Identifier: false)); writer.NewLine = "\n"; - count = await WriteCsvAsync(items, writer, ShellInterpreter.CSVSeparator, token); + count = await WriteCsvAsync(items, writer, ShellInterpreter.CSVSeparator, token, onRowWritten: onRowWritten); } else { - count = await WriteArrayAsync(items, stream, token); + count = await WriteArrayAsync(items, stream, token, onRowWritten); } } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs index 701e5fc1..3794565a 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ImportCommand.cs @@ -732,10 +732,19 @@ public override async Task ExecuteAsync(ShellInterpreter shell, Co var partitionKeySegments = ParsePartitionKeySegments(this.PartitionKey); var progress = new ImportProgress(); + var reporter = new IngestionProgress( + shell, + selectedFormat, + dryRun ? "command-import-validation-progress" : "command-import-progress"); + void UpdateProgress() => reporter.Update( + dryRun ? progress.Validated : progress.Success + progress.Failed, + progress.Success, + progress.Failed, + progress.Charge); Exception? failure = null; try { - await ExecuteImportAsync( + await reporter.RunAsync(() => ExecuteImportAsync( shell, filePath, format, @@ -748,7 +757,8 @@ await ExecuteImportAsync( concurrency, commandState, progress, - token); + token, + UpdateProgress)); } catch (DecoderFallbackException ex) { @@ -757,6 +767,10 @@ await ExecuteImportAsync( MessageService.GetArgsString("command-import-error-invalid_document", "message", ex.Message), ex); } + catch (OperationCanceledException ex) when (token.IsCancellationRequested) + { + failure = ex; + } catch (Exception ex) when (ex is CommandException or JsonException or IOException or UnauthorizedAccessException) { failure = ex; @@ -872,7 +886,8 @@ private static async Task ExecuteImportAsync( int concurrency, CommandState commandState, ImportProgress progress, - CancellationToken token) + CancellationToken token, + Action progressChanged) { FileStream? stream = null; IAsyncEnumerable<(int LineNumber, JsonElement Item)> source; @@ -894,6 +909,8 @@ private static async Task ExecuteImportAsync( { ValidateDryRunItem(item.Item, item.LineNumber); progress.Validated++; + progressChanged(); + shell.ReportOperationProgress(); } return; @@ -910,7 +927,8 @@ await WriteItemsAsync( commandState, token, warningShell: shell, - progress: progress); + progress: progress, + progressChanged: progressChanged); } finally { @@ -971,7 +989,8 @@ private static void ValidateDryRunItem(JsonElement item, int lineNumber) CommandState commandState, CancellationToken token, ShellInterpreter? warningShell = null, - ImportProgress? progress = null) + ImportProgress? progress = null, + Action? progressChanged = null) { ArgumentOutOfRangeException.ThrowIfLessThan(concurrency, 1); progress ??= new ImportProgress(); @@ -991,6 +1010,8 @@ void Record((double Charge, Exception? Error) result) progress.Failed++; progress.FirstFailure ??= result.Error; } + + progressChanged?.Invoke(); } async Task DrainCompletedAsync() @@ -1006,6 +1027,32 @@ async Task DrainCompletedAsync() } await using var items = source.GetAsyncEnumerator(token); + async ValueTask ReadNextAsync() + { + var next = items.MoveNextAsync(); + if (next.IsCompleted || pending.Count == 0) + { + return await next; + } + + var read = next.AsTask(); + try + { + while (!read.IsCompleted && pending.Count > 0) + { + await Task.WhenAny(pending.Cast().Append(read)); + await DrainCompletedAsync(); + } + + return await read; + } + finally + { + // Settle the read before disposing its enumerator, even when a write fails. + await read; + } + } + try { while (true) @@ -1023,7 +1070,7 @@ async Task DrainCompletedAsync() } token.ThrowIfCancellationRequested(); - if (!await items.MoveNextAsync()) + if (!await ReadNextAsync()) { break; } @@ -1058,18 +1105,35 @@ async Task DrainCompletedAsync() items.Current.LineNumber, items.Current.Item, partitionKeyPaths, - token)); + token, + warningShell)); } } finally { try { - // Drain every started write before returning or releasing the source, even on parse errors. - var results = await Task.WhenAll(pending); - foreach (var result in results) + // Drain every started write before returning or releasing the source, even on parse + // errors. Record each write as it completes so progress does not wait for the slowest one. + Task<(double Charge, Exception? Error)>? firstUnsuccessful = null; + while (pending.Count > 0) + { + var completed = await Task.WhenAny(pending); + pending.Remove(completed); + if (completed.IsCompletedSuccessfully) + { + Record(await completed); + } + else if (firstUnsuccessful is null || (firstUnsuccessful.IsCanceled && completed.IsFaulted)) + { + firstUnsuccessful = completed; + } + } + + if (firstUnsuccessful != null) { - Record(result); + // Matches Task.WhenAll: a fault takes precedence over cancellation. + await firstUnsuccessful; } } finally @@ -1091,7 +1155,8 @@ async Task DrainCompletedAsync() int lineNumber, JsonElement item, IReadOnlyList partitionKeyPaths, - CancellationToken token) + CancellationToken token, + ShellInterpreter? shell) { try { @@ -1125,6 +1190,7 @@ async Task DrainCompletedAsync() } catch (Exception ex) when (ex is JsonException or InvalidOperationException) { + shell?.ReportOperationProgress(); return (0, new CommandException( "import", MessageService.GetArgsString( diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/IngestionProgress.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/IngestionProgress.cs new file mode 100644 index 00000000..3fb88240 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/IngestionProgress.cs @@ -0,0 +1,125 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Commands; + +using System.Diagnostics; +using System.Globalization; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.Util; + +internal sealed class IngestionProgress(ShellInterpreter shell, CommandState state, string messageKey, long? total = null) +{ + private readonly object sync = new(); + private long processed; + private long succeeded; + private long failed; + private double charge; + + internal static bool IsEnabled(ShellInterpreter shell, CommandState state) + { + var format = state.OutputFormatExplicitlySet ? state.OutputFormat : shell.DefaultOutputFormat; + return !shell.IsMachineMode && !shell.IsMcpInvocation + && string.IsNullOrEmpty(shell.StdOutRedirect) + && shell.IsOutputTerminal() + && !OutputPolicy.IsStructuredFormat(format); + } + + internal void Update(long processed, long succeeded, long failed, double charge) + { + // Display only: failed writes, including SDK-synthesized timeouts, must not count as + // emulator activity. Real responses restart the inactivity timeout in OperationProgressHandler. + lock (this.sync) + { + this.processed = processed; + this.succeeded = succeeded; + this.failed = failed; + this.charge = charge; + } + } + + internal async Task RunAsync(Func execute) + { + if (!IsEnabled(shell, state)) + { + await execute(); + return; + } + + // Start reporting before invoking the work so a long synchronous stretch before its first + // await is still covered; the timer continuations run on the thread pool. + var completion = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + var reporting = this.ReportSafelyAsync(completion.Task); + try + { + await execute(); + } + finally + { + completion.TrySetResult(); + await reporting; + try + { + shell.ClearTransientLine(); + } + catch (Exception ex) when (ex is not OutOfMemoryException) + { + // Like rendering, cleanup must never replace the operation's own outcome. + Trace.TraceWarning("Clearing the ingestion progress line failed ({0}).", ex.GetType().Name); + } + } + } + + private async Task ReportSafelyAsync(Task completed) + { + try + { + await this.ReportUntilCompleteAsync(completed); + } + catch (Exception ex) when (ex is not OutOfMemoryException) + { + // Progress is informational: a formatting or output failure must not turn a + // successful ingestion into a failed one. Stop reporting and let the work finish. + Trace.TraceWarning("Ingestion progress reporting stopped ({0}).", ex.GetType().Name); + } + } + + private async Task ReportUntilCompleteAsync(Task completed) + { + var elapsed = Stopwatch.StartNew(); + using var timer = new PeriodicTimer(TimeSpan.FromSeconds(2)); + var tick = timer.WaitForNextTickAsync().AsTask(); + while (await Task.WhenAny(completed, tick).ConfigureAwait(false) != completed) + { + await tick.ConfigureAwait(false); + string message; + string compactMessage; + lock (this.sync) + { + object[] args = + [ + "processed", + this.processed, + "succeeded", + this.succeeded, + "failed", + this.failed, + "total", + total ?? 0, + "elapsed", + elapsed.Elapsed.TotalSeconds.ToString("F1", CultureInfo.InvariantCulture), + "rate", + (this.processed / elapsed.Elapsed.TotalSeconds).ToString("F1", CultureInfo.InvariantCulture), + "charge", + this.charge.ToString("F2", CultureInfo.InvariantCulture), + ]; + message = MessageService.GetArgsString(messageKey, args); + compactMessage = MessageService.GetArgsString(messageKey + "-compact", args); + } + + shell.WriteTransientLine(message, compactMessage); + tick = timer.WaitForNextTickAsync().AsTask(); + } + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs index 6b7ae64c..ae505f93 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs @@ -235,6 +235,7 @@ private async Task ListContainerItemsAsync(ConnectedState state, S while (feedIterator.HasMoreResults) { using var response = await feedIterator.ReadNextAsync(token); + interpreter.ReportOperationProgress(); AccumulateRequestCharge(returnState, response.Headers.RequestCharge); JsonDocument queryDocument; try diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs index a6cdd7b4..2b93948b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs @@ -746,6 +746,7 @@ internal async Task ExecuteQueryAsync(Container container, ShellIn } using var response = await feedIterator.ReadNextAsync(token); + shell.ReportOperationProgress(); var pageRequestCharge = response.Headers.RequestCharge; if (!response.IsSuccessStatusCode) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs index 1f7ec48e..9dee444b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs @@ -189,8 +189,26 @@ private async Task RemoveItemsFromContainerAsync(ConnectedState state, var matchKeyPropertyNames = string.IsNullOrEmpty(this.Key) ? partitionKeyPropertyNames : [this.Key]; var totalCount = 0; + var scannedCount = 0; double totalCharge = 0; bool dryRun = this.DryRun == true; + var progress = new IngestionProgress( + this.shell, + new CommandState(), + dryRun ? "command-rm-dry-run-progress" : "command-rm-progress"); + + void ReportProgress() => progress.Update(scannedCount, totalCount, 0, totalCharge); + + void RecordDelete((bool Counted, double RequestCharge) result) + { + totalCharge += result.RequestCharge; + if (result.Counted) + { + totalCount++; + } + + ReportProgress(); + } bool ShouldDelete(JsonElement element) { @@ -230,159 +248,161 @@ bool TryGetPipedItemPartitionKey(JsonElement item, out PartitionKey itemPartitio } } - // Process pipe input if available - if (hasPipeInput && commandState.Result is ShellJson jsonResult) + async Task RemoveMatchesAsync() { - // Handle items from pipe - check if it's an array or object with items property - JsonElement itemsArray; - - if (jsonResult.Value.ValueKind == JsonValueKind.Array) - { - itemsArray = jsonResult.Value; - } - else if (jsonResult.Value.TryGetProperty("items", out var itemsProp) && itemsProp.ValueKind == JsonValueKind.Array) - { - itemsArray = itemsProp; - } - else + // Process pipe input if available + if (hasPipeInput && commandState.Result is ShellJson jsonResult) { - // Single item - treat as array of one - itemsArray = jsonResult.Value; - } + // Handle items from pipe - check if it's an array or object with items property + JsonElement itemsArray; - if (itemsArray.ValueKind == JsonValueKind.Array) - { - foreach (var item in itemsArray.EnumerateArray()) + if (jsonResult.Value.ValueKind == JsonValueKind.Array) { - // Filter before converting the key so an excluded item with an invalid - // partition-key value cannot block deletion of matching items. - if (item.TryGetProperty("id", out var idElement) - && idElement.GetString() is { } id - && ShouldDelete(item) - && TryGetPipedItemPartitionKey(item, out var itemPartitionKey)) + itemsArray = jsonResult.Value; + } + else if (jsonResult.Value.TryGetProperty("items", out var itemsProp) && itemsProp.ValueKind == JsonValueKind.Array) + { + itemsArray = itemsProp; + } + else + { + // Single item - treat as array of one + itemsArray = jsonResult.Value; + } + + if (itemsArray.ValueKind == JsonValueKind.Array) + { + foreach (var item in itemsArray.EnumerateArray()) { - var deleteResult = await TryDeleteAsync(id, itemPartitionKey); - totalCharge += deleteResult.RequestCharge; - if (deleteResult.Counted) + token.ThrowIfCancellationRequested(); + scannedCount++; + this.shell.ReportOperationProgress(); + + // Filter before converting the key so an excluded item with an invalid + // partition-key value cannot block deletion of matching items. + if (item.TryGetProperty("id", out var idElement) + && idElement.GetString() is { } id + && ShouldDelete(item) + && TryGetPipedItemPartitionKey(item, out var itemPartitionKey)) { - totalCount++; + RecordDelete(await TryDeleteAsync(id, itemPartitionKey)); + } + else + { + ReportProgress(); } } } - } - else - { - // Single item from pipe - if (itemsArray.TryGetProperty("id", out var idElement) - && idElement.GetString() is { } id - && ShouldDelete(itemsArray) - && TryGetPipedItemPartitionKey(itemsArray, out var itemPartitionKey)) + else { - var deleteResult = await TryDeleteAsync(id, itemPartitionKey); - totalCharge += deleteResult.RequestCharge; - if (deleteResult.Counted) + // Single item from pipe + token.ThrowIfCancellationRequested(); + scannedCount++; + this.shell.ReportOperationProgress(); + if (itemsArray.TryGetProperty("id", out var idElement) + && idElement.GetString() is { } id + && ShouldDelete(itemsArray) + && TryGetPipedItemPartitionKey(itemsArray, out var itemPartitionKey)) { - totalCount++; + RecordDelete(await TryDeleteAsync(id, itemPartitionKey)); } } } - } - else if (this.matcher != null) - { - var needsDeletionPlan = !dryRun - && this.partitionKey == null - && partitionKeyPropertyNames.Length > 0; - var queryOptions = this.partitionKey is PartitionKey scope - ? new QueryRequestOptions { PartitionKey = scope } - : null; - - using var feedIterator = container.GetItemQueryStreamIterator("SELECT * FROM c", requestOptions: queryOptions); - await using var planStream = needsDeletionPlan ? this.CreateDeletionPlanStream() : null; - await using var planWriter = planStream is null - ? null - : new StreamWriter(planStream, new System.Text.UTF8Encoding(false), 4096, leaveOpen: true); - var undefinedPartitionKeyCount = 0; - - while (feedIterator.HasMoreResults) + else if (this.matcher != null) { - token.ThrowIfCancellationRequested(); - using var response = await feedIterator.ReadNextAsync(token); + var needsDeletionPlan = !dryRun + && this.partitionKey == null + && partitionKeyPropertyNames.Length > 0; + var queryOptions = this.partitionKey is PartitionKey scope + ? new QueryRequestOptions { PartitionKey = scope } + : null; + + using var feedIterator = container.GetItemQueryStreamIterator("SELECT * FROM c", requestOptions: queryOptions); + await using var planStream = needsDeletionPlan ? this.CreateDeletionPlanStream() : null; + await using var planWriter = planStream is null + ? null + : new StreamWriter(planStream, new System.Text.UTF8Encoding(false), 4096, leaveOpen: true); + var undefinedPartitionKeyCount = 0; + + while (feedIterator.HasMoreResults) + { + token.ThrowIfCancellationRequested(); + using var response = await feedIterator.ReadNextAsync(token); - // The scan pages that locate matching items consume RUs regardless of whether - // any item is ultimately deleted (including in --dry-run), so include each - // page's request charge from the response headers. - totalCharge += response.Headers.RequestCharge; + // The scan pages that locate matching items consume RUs regardless of whether + // any item is ultimately deleted (including in --dry-run), so include each + // page's request charge from the response headers. + totalCharge += response.Headers.RequestCharge; - using var streamReader = new StreamReader(response.Content); - using var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync(token)); + using var streamReader = new StreamReader(response.Content); + using var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync(token)); - foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) - { - if (!element.TryGetProperty(IdPropertyName, out var idElement) - || idElement.ValueKind != JsonValueKind.String - || idElement.GetString() is not string id - || !ShouldDelete(element)) + foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray()) { - continue; - } + token.ThrowIfCancellationRequested(); + scannedCount++; + this.shell.ReportOperationProgress(); + ReportProgress(); + if (!element.TryGetProperty(IdPropertyName, out var idElement) + || idElement.ValueKind != JsonValueKind.String + || idElement.GetString() is not string id + || !ShouldDelete(element)) + { + continue; + } - TryCreatePartitionKeyFromItem( - element, - partitionKeyPropertyNames, - out var hasUndefinedComponent, - out var itemPartitionKey); + TryCreatePartitionKeyFromItem( + element, + partitionKeyPropertyNames, + out var hasUndefinedComponent, + out var itemPartitionKey); - if (planWriter != null) - { - if (hasUndefinedComponent) + if (planWriter != null) { - undefinedPartitionKeyCount++; + if (hasUndefinedComponent) + { + undefinedPartitionKeyCount++; + } + + var partitionKeyJson = itemPartitionKey.Equals(PartitionKey.None) + ? "{}" + : itemPartitionKey.ToString(); + var candidate = new DeletionCandidate(id, partitionKeyJson); + await planWriter.WriteLineAsync(JsonSerializer.Serialize(candidate).AsMemory(), token); } - - var partitionKeyJson = itemPartitionKey.Equals(PartitionKey.None) - ? "{}" - : itemPartitionKey.ToString(); - var candidate = new DeletionCandidate(id, partitionKeyJson); - await planWriter.WriteLineAsync(JsonSerializer.Serialize(candidate).AsMemory(), token); - } - else - { - var deleteResult = await TryDeleteAsync(id, itemPartitionKey); - totalCharge += deleteResult.RequestCharge; - if (deleteResult.Counted) + else { - totalCount++; + RecordDelete(await TryDeleteAsync(id, itemPartitionKey)); } } + + ReportProgress(); } - } - token.ThrowIfCancellationRequested(); - if (planWriter != null) - { - await planWriter.FlushAsync(token); - planStream!.Position = 0; + token.ThrowIfCancellationRequested(); + if (planWriter != null) + { + await planWriter.FlushAsync(token); + planStream!.Position = 0; - AddUndefinedPartitionKeyWarning(commandState, undefinedPartitionKeyCount); - this.shell!.RenderMutationWarnings(commandState); + AddUndefinedPartitionKeyWarning(commandState, undefinedPartitionKeyCount); + this.shell!.RenderMutationWarnings(commandState); - using var planReader = new StreamReader(planStream, System.Text.Encoding.UTF8, false, 4096, leaveOpen: true); - while (await planReader.ReadLineAsync(token) is { } line) - { - token.ThrowIfCancellationRequested(); - var candidate = JsonSerializer.Deserialize(line) - ?? throw new JsonException("Unable to read a delete candidate from the temporary plan."); - var itemPartitionKey = CreatePartitionKeyFromArgument(candidate.PartitionKeyJson); - var deleteResult = await TryDeleteAsync(candidate.Id, itemPartitionKey); - totalCharge += deleteResult.RequestCharge; - if (deleteResult.Counted) + using var planReader = new StreamReader(planStream, System.Text.Encoding.UTF8, false, 4096, leaveOpen: true); + while (await planReader.ReadLineAsync(token) is { } line) { - totalCount++; + token.ThrowIfCancellationRequested(); + var candidate = JsonSerializer.Deserialize(line) + ?? throw new JsonException("Unable to read a delete candidate from the temporary plan."); + var itemPartitionKey = CreatePartitionKeyFromArgument(candidate.PartitionKeyJson); + RecordDelete(await TryDeleteAsync(candidate.Id, itemPartitionKey)); } } } } + await progress.RunAsync(RemoveMatchesAsync); + string renderMessage = totalCount > 0 ? MessageService.GetString( dryRun ? "command-rm-dry-run-plan" : "command-rm-deleted_items", diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmContainerCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmContainerCommand.cs index 615e93e3..98d5a768 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmContainerCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmContainerCommand.cs @@ -50,7 +50,7 @@ async Task IStateVisitor.VisitConnectedSta { if (!string.IsNullOrEmpty(this.Database)) { - return await this.RemoveContainerAsync(state, this.Database, token); + return await this.RemoveContainerAsync(state, shell, this.Database, token); } throw new NotInDatabaseException("rmcon"); @@ -60,23 +60,23 @@ async Task IStateVisitor.VisitDatabaseStat { if (!string.IsNullOrEmpty(this.Database)) { - return await this.RemoveContainerAsync(state, this.Database, token); + return await this.RemoveContainerAsync(state, shell, this.Database, token); } - return await this.RemoveContainerAsync(state, state.DatabaseName, token); + return await this.RemoveContainerAsync(state, shell, state.DatabaseName, token); } async Task IStateVisitor.VisitContainerStateAsync(ContainerState state, ShellInterpreter shell, CancellationToken token) { if (!string.IsNullOrEmpty(this.Database)) { - return await this.RemoveContainerAsync(state, this.Database, token); + return await this.RemoveContainerAsync(state, shell, this.Database, token); } throw new NotInContainerException("rmcon"); } - private async Task RemoveContainerAsync(ConnectedState state, string databaseName, CancellationToken token) + private async Task RemoveContainerAsync(ConnectedState state, ShellInterpreter shell, string databaseName, CancellationToken token) { // Validate database exists await ValidateDatabaseExistsAsync(state, databaseName, "rmcon", token); @@ -98,7 +98,7 @@ private async Task RemoveContainerAsync(ConnectedState state, string d return 0; } - if (this.Force == true || ShellInterpreter.Confirm("command-rmcon-confirm_container_deletion")) + if (this.Force == true || ShellInterpreter.Confirm("command-rmcon-confirm_container_deletion", shell)) { await CosmosResourceFacade.DeleteContainerAsync(state, databaseName, containerName, token); CosmosCompleteCommand.ClearContainers(); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmDbCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmDbCommand.cs index a6095cd0..2d7767f4 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmDbCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmDbCommand.cs @@ -92,7 +92,7 @@ private async Task RemoveDatabaseAsync(ConnectedState state, ShellInte return 0; } - if (this.Force is true || ShellInterpreter.Confirm("command-rmdb-confirm_db_deletion")) + if (this.Force is true || ShellInterpreter.Confirm("command-rmdb-confirm_db_deletion", shell)) { await CosmosResourceFacade.DeleteDatabaseAsync(state, databaseName, token); UpdateStateAfterDelete(shell, state.Client, state.ArmContext, databaseName); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs index f23cb148..d428dd26 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/SprocCommand.cs @@ -373,7 +373,7 @@ private async Task CreateAsync(Container container, ShellInterpret } var seed = existingBody ?? DefaultStoredProcedureBody(); - var edited = await this.LaunchEditorAsync(seed, name, token); + var edited = await this.LaunchEditorAsync(shell, seed, name, token); if (!string.IsNullOrWhiteSpace(edited)) { @@ -384,7 +384,7 @@ private async Task CreateAsync(Container container, ShellInterpret ShellInterpreter.WriteLine(); } - if (string.IsNullOrWhiteSpace(edited) || !ShellInterpreter.Confirm("command-sproc-create-confirm")) + if (string.IsNullOrWhiteSpace(edited) || !ShellInterpreter.Confirm("command-sproc-create-confirm", shell)) { commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "sproc", id = name, created = false })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString("command-sproc-create-discarded", "name", name)); @@ -525,7 +525,7 @@ private async Task EditAsync(Container container, ShellInterpreter throw NotFound(name, ex); } - var newBody = await this.LaunchEditorAsync(existingBody, name, token); + var newBody = await this.LaunchEditorAsync(shell, existingBody, name, token); if (string.Equals(newBody, existingBody, StringComparison.Ordinal)) { @@ -553,7 +553,7 @@ private async Task EditAsync(Container container, ShellInterpreter /// external editor, waits for the editor to close, and returns the edited /// contents. The temporary file is always removed. /// - private async Task LaunchEditorAsync(string initialBody, string name, CancellationToken token) + private async Task LaunchEditorAsync(ShellInterpreter shell, string initialBody, string name, CancellationToken token) { var editor = ExternalEditor.Resolve(null); if (editor is null) @@ -568,6 +568,7 @@ private async Task LaunchEditorAsync(string initialBody, string name, Ca try { await File.WriteAllTextAsync(tempPath, initialBody, token); + using var pause = shell.PauseOperationTimeout(); AnsiConsole.MarkupLine(Theme.FormatMuted(MessageService.GetArgsString( "command-sproc-edit-launching", diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThemeCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThemeCommand.cs index e9200ea6..cb869cc3 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThemeCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThemeCommand.cs @@ -56,7 +56,7 @@ public override Task ExecuteAsync(ShellInterpreter shell, CommandS "load" => Task.FromResult(this.RunLoad(commandState)), "validate" => Task.FromResult(this.RunValidate(commandState)), "save" => Task.FromResult(this.RunSave(commandState)), - "edit" => Task.FromResult(this.RunEdit(commandState)), + "edit" => Task.FromResult(this.RunEdit(shell, commandState)), "open" => Task.FromResult(this.RunOpen(commandState)), "reload" => Task.FromResult(this.RunReload(commandState)), _ => Task.FromResult(this.RunUnknownAction(commandState, action)), @@ -586,7 +586,7 @@ private CommandState RunOpen(CommandState commandState) return commandState; } - private CommandState RunEdit(CommandState commandState) + private CommandState RunEdit(ShellInterpreter shell, CommandState commandState) { // Resolve target. Empty -> active theme; otherwise treat as theme name // or path (same rules as load/validate). @@ -689,6 +689,7 @@ private CommandState RunEdit(CommandState commandState) int exitCode; try { + using var pause = shell.PauseOperationTimeout(); using var process = System.Diagnostics.Process.Start(new System.Diagnostics.ProcessStartInfo { FileName = editor.FileName, diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs index 2289c6e3..118d365a 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs @@ -260,7 +260,7 @@ private static bool ConfirmWrite(ShellInterpreter shell, bool yes, string databa string modeLabel = MessageService.GetString(isAutoscale ? "command-throughput-mode-autoscale" : "command-throughput-mode-manual"); string ruText = ru.ToString(CultureInfo.InvariantCulture); AnsiConsole.MarkupLine(MessageService.GetArgsString("command-throughput-confirm_summary", "resource", Markup.Escape(resourceName), "mode", modeLabel, "ru", ruText)); - return ShellInterpreter.Confirm("command-throughput-confirm"); + return ShellInterpreter.Confirm("command-throughput-confirm", shell); } private async Task ExecuteOnScopeAsync(ConnectedState state, ShellInterpreter shell, string databaseName, string? containerName, CancellationToken token) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs index e13c0863..aea734b3 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/TriggerCommand.cs @@ -342,7 +342,7 @@ private async Task CreateAsync(Container container, ShellInterpret } var seed = existingBody ?? DefaultTriggerBody(); - var edited = await this.LaunchEditorAsync(seed, name, token); + var edited = await this.LaunchEditorAsync(shell, seed, name, token); if (!string.IsNullOrWhiteSpace(edited)) { @@ -353,7 +353,7 @@ private async Task CreateAsync(Container container, ShellInterpret ShellInterpreter.WriteLine(); } - if (string.IsNullOrWhiteSpace(edited) || !ShellInterpreter.Confirm("command-trigger-create-confirm")) + if (string.IsNullOrWhiteSpace(edited) || !ShellInterpreter.Confirm("command-trigger-create-confirm", shell)) { commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "trigger", id = name, created = false })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString("command-trigger-create-discarded", "name", name)); @@ -470,7 +470,7 @@ private async Task EditAsync(Container container, ShellInterpreter } var existingBody = existing.Body ?? string.Empty; - var newBody = await this.LaunchEditorAsync(existingBody, name, token); + var newBody = await this.LaunchEditorAsync(shell, existingBody, name, token); if (string.Equals(newBody, existingBody, StringComparison.Ordinal)) { @@ -504,7 +504,7 @@ private async Task EditAsync(Container container, ShellInterpreter /// external editor, waits for the editor to close, and returns the edited /// contents. The temporary file is always removed. /// - private async Task LaunchEditorAsync(string initialBody, string name, CancellationToken token) + private async Task LaunchEditorAsync(ShellInterpreter shell, string initialBody, string name, CancellationToken token) { var editor = ExternalEditor.Resolve(null); if (editor is null) @@ -519,6 +519,7 @@ private async Task LaunchEditorAsync(string initialBody, string name, Ca try { await File.WriteAllTextAsync(tempPath, initialBody, token); + using var pause = shell.PauseOperationTimeout(); AnsiConsole.MarkupLine(Theme.FormatMuted(MessageService.GetArgsString( "command-trigger-edit-launching", diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs index d5d2384d..af9f5a90 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/UdfCommand.cs @@ -281,7 +281,7 @@ private async Task CreateAsync(Container container, ShellInterpret } var seed = existingBody ?? DefaultUserDefinedFunctionBody(); - var edited = await this.LaunchEditorAsync(seed, name, token); + var edited = await this.LaunchEditorAsync(shell, seed, name, token); if (!string.IsNullOrWhiteSpace(edited)) { @@ -292,7 +292,7 @@ private async Task CreateAsync(Container container, ShellInterpret ShellInterpreter.WriteLine(); } - if (string.IsNullOrWhiteSpace(edited) || !ShellInterpreter.Confirm("command-udf-create-confirm")) + if (string.IsNullOrWhiteSpace(edited) || !ShellInterpreter.Confirm("command-udf-create-confirm", shell)) { commandState.Result = new ShellJson(JsonSerializer.SerializeToElement(new { type = "udf", id = name, created = false })); commandState.RenderUser = () => ShellInterpreter.WriteLine(MessageService.GetArgsString("command-udf-create-discarded", "name", name)); @@ -395,7 +395,7 @@ private async Task EditAsync(Container container, ShellInterpreter throw NotFound(name, ex); } - var newBody = await this.LaunchEditorAsync(existingBody, name, token); + var newBody = await this.LaunchEditorAsync(shell, existingBody, name, token); if (string.Equals(newBody, existingBody, StringComparison.Ordinal)) { @@ -423,7 +423,7 @@ private async Task EditAsync(Container container, ShellInterpreter /// external editor, waits for the editor to close, and returns the edited /// contents. The temporary file is always removed. /// - private async Task LaunchEditorAsync(string initialBody, string name, CancellationToken token) + private async Task LaunchEditorAsync(ShellInterpreter shell, string initialBody, string name, CancellationToken token) { var editor = ExternalEditor.Resolve(null); if (editor is null) @@ -438,6 +438,7 @@ private async Task LaunchEditorAsync(string initialBody, string name, Ca try { await File.WriteAllTextAsync(tempPath, initialBody, token); + using var pause = shell.PauseOperationTimeout(); AnsiConsole.MarkupLine(Theme.FormatMuted(MessageService.GetArgsString( "command-udf-edit-launching", diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs index bf11ff49..939e39c8 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/WatchCommand.cs @@ -63,12 +63,7 @@ public async override Task ExecuteAsync(ShellInterpreter shell, Co "watch", token); - // A tail is long-running and must not be killed by the per-command - // timeout. Switch to the interruptible (no-timeout) cancellation source - // so the loop runs until the user presses Ctrl+C or --max is reached. - var watchToken = ShellInterpreter.UserCancellationTokenSource.Token; - - return await this.WatchAsync(shell, container, watchToken); + return await this.WatchAsync(shell, container, token); } /// @@ -178,7 +173,11 @@ private async Task WatchAsync(ShellInterpreter shell, Container co if (response.StatusCode == HttpStatusCode.NotModified) { // Caught up with the feed; wait before polling for more. - await Task.Delay(pollInterval, token); + using (shell.PauseOperationTimeout()) + { + await Task.Delay(pollInterval, token); + } + continue; } @@ -195,6 +194,7 @@ private async Task WatchAsync(ShellInterpreter shell, Container co var limitReached = false; foreach (var element in documents) { + token.ThrowIfCancellationRequested(); if (!redirected) { AnsiConsole.MarkupLine(JsonOutputHighlighter.BuildMarkup(element)); @@ -202,6 +202,7 @@ private async Task WatchAsync(ShellInterpreter shell, Container co collected?.Add(element); count++; + shell.ReportOperationProgress(); if (max.HasValue && count >= max.Value) { @@ -216,14 +217,14 @@ private async Task WatchAsync(ShellInterpreter shell, Container co } } } - catch (OperationCanceledException) + catch (OperationCanceledException) when (token.IsCancellationRequested) { // Ctrl+C: stop tailing and report what was seen so far. } AnsiConsole.MarkupLine(MessageService.GetArgsString("command-watch-stopped", "count", count.ToString())); - var result = new CommandState(); + var result = new CommandState { UserCancellationHandled = token.IsCancellationRequested }; result.SetFormat(this.OutputFormat); if (collected != null) { diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs index 7a1568c9..29b1716a 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CommandState.cs @@ -67,6 +67,8 @@ public OutputFormat OutputFormat internal bool OutputRendered { get; set; } + internal bool UserCancellationHandled { get; set; } + /// /// Gets or sets an optional delegate that renders the result for a human on an /// interactive terminal. When set and the effective format is , @@ -296,6 +298,11 @@ internal string GenerateOutputText() } } + internal sealed class HandledCancellationException(CommandState state, CancellationToken token) : OperationCanceledException(token) + { + internal CommandState State { get; } = state; + } + internal sealed class FailureException : Exception { public FailureException(CommandState state) diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/OperationInactivityTimeout.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/OperationInactivityTimeout.cs new file mode 100644 index 00000000..afe70645 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/OperationInactivityTimeout.cs @@ -0,0 +1,134 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +/// +/// Cancels a command after a period without progress. Progress restarts the period, and +/// pauses (for example while waiting for user input) suspend it until every pause ends. +/// +internal sealed class OperationInactivityTimeout : IDisposable +{ + internal static readonly IDisposable NoPause = new EmptyScope(); + + private readonly object sync = new(); + private readonly CancellationTokenSource source; + private readonly CancellationToken callerToken; + private readonly TimeSpan timeout; + private readonly ITimer? timer; + private int pauseCount; + + /// The source canceled when the timeout expires. + /// The period without progress that cancels . + /// The caller's token; its cancellation is not reported as a timeout. + /// A clock for tests. When , the source's own timer is used. + internal OperationInactivityTimeout(CancellationTokenSource source, TimeSpan timeout, CancellationToken callerToken = default, TimeProvider? timeProvider = null) + { + this.source = source; + this.callerToken = callerToken; + this.timeout = timeout; + this.timer = timeProvider?.CreateTimer( + static state => ((OperationInactivityTimeout)state!).Expire(), + this, + Timeout.InfiniteTimeSpan, + Timeout.InfiniteTimeSpan); + this.Restart(timeout); + } + + internal void ThrowIfTimedOut() + { + if (!this.callerToken.IsCancellationRequested) + { + this.source.Token.ThrowIfCancellationRequested(); + } + } + + internal void ReportProgress() + { + lock (this.sync) + { + if (this.pauseCount == 0) + { + this.Restart(this.timeout); + } + } + } + + internal IDisposable Pause() + { + lock (this.sync) + { + if (this.pauseCount++ == 0) + { + this.Restart(Timeout.InfiniteTimeSpan); + } + } + + return new PauseScope(this); + } + + private void Resume() + { + lock (this.sync) + { + if (--this.pauseCount == 0) + { + this.Restart(this.timeout); + } + } + } + + public void Dispose() => this.timer?.Dispose(); + + private void Expire() + { + try + { + this.source.Cancel(); + } + catch (ObjectDisposedException) + { + // The command finished just before the timer fired. + } + } + + private void Restart(TimeSpan delay) + { + try + { + if (this.timer is not null) + { + this.timer.Change(delay, Timeout.InfiniteTimeSpan); + } + else + { + this.source.CancelAfter(delay); + } + } + catch (ObjectDisposedException) + { + // The command finished while a late completion was reporting progress. + } + } + + private sealed class EmptyScope : IDisposable + { + public void Dispose() + { + } + } + + private sealed class PauseScope(OperationInactivityTimeout owner) : IDisposable + { + private int disposed; + + public void Dispose() + { + if (Interlocked.Exchange(ref this.disposed, 1) == 0) + { + owner.Resume(); + } + } + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/OperationProgressHandler.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/OperationProgressHandler.cs new file mode 100644 index 00000000..a18d445f --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/OperationProgressHandler.cs @@ -0,0 +1,26 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +using System.Net; + +/// +/// Treats every response the local emulator actually sends as progress for the running +/// command's inactivity timeout. Request timeouts and unavailability do not count, because +/// the SDK synthesizes them when the emulator does not answer. +/// +internal sealed class OperationProgressHandler(ShellInterpreter shell) : RequestHandler +{ + public override async Task SendAsync(RequestMessage request, CancellationToken cancellationToken) + { + var response = await base.SendAsync(request, cancellationToken); + if (response.StatusCode is not (HttpStatusCode.RequestTimeout or HttpStatusCode.ServiceUnavailable)) + { + shell.ReportOperationProgress(); + } + + return response; + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index b6c9257c..614562ab 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -45,6 +45,9 @@ public partial class ShellInterpreter : IDisposable private const int OptionalArmDiscoveryTimeoutSeconds = 3; + // A constant, not a static TimeSpan field: Instance is constructed before later static initializers run. + private const int LocalEmulatorOperationTimeoutSeconds = 10; + internal const int MaximumCallDepth = 64; private const string EncodedHistoryLinePrefix = "CosmosDBShellHistoryV1:"; @@ -56,8 +59,6 @@ public partial class ShellInterpreter : IDisposable private static readonly TimeSpan HistoryFileOpenRetryDelay = TimeSpan.FromMilliseconds(25); - private static readonly TimeSpan LocalEmulatorOperationTimeout = TimeSpan.FromSeconds(10); - private static CancellationTokenSource? currentTokenSource; private readonly string cfgPath; @@ -70,6 +71,8 @@ public partial class ShellInterpreter : IDisposable private readonly object historyLock = new(); + private readonly object transientLineLock = new(); + private readonly List pendingHistoryEntries = []; private readonly object lineEditorLock = new(); @@ -78,6 +81,13 @@ public partial class ShellInterpreter : IDisposable private long stateVersion; + private OperationInactivityTimeout? localEmulatorDeadline; + private CancellationToken cancellationErrorToken; + + private int activeInterruptOutputFormat = -1; + + private int transientLineLength; + private int callDepth; private TokenCredential? activeCredential; @@ -283,9 +293,32 @@ internal long SessionRequestChargeGeneration internal string WelcomeMarkerFile => this.welcomeMarkerFile; + /// + /// Gets or sets how long a command connected to a local emulator may run without reporting + /// progress before it is canceled. Each call restarts it. + /// + internal TimeSpan LocalEmulatorInactivityTimeout { get; set; } = TimeSpan.FromSeconds(LocalEmulatorOperationTimeoutSeconds); + + /// + /// Gets or sets the clock for the local-emulator inactivity timeout. Tests set a manual clock; + /// when , the timeout uses the system clock. + /// + internal TimeProvider? LocalEmulatorTimeProvider { get; set; } + internal Func IsInteractiveSession { get; set; } = static () => !Console.IsInputRedirected && !Console.IsOutputRedirected; + /// + /// Gets or sets a value indicating whether standard output is a terminal that can redraw a + /// line in place. Progress falls back to one line per update when it is not. + /// + internal Func IsOutputTerminal { get; set; } = static () => !Console.IsOutputRedirected; + + /// + /// Gets or sets the console width used to fit in-place status lines; 0 means unknown. + /// + internal Func ConsoleWidth { get; set; } = GetConsoleWidth; + internal Func IsInputRedirected { get; set; } = static () => Console.IsInputRedirected; // RadLine detects terminal capabilities itself when no console is supplied. @@ -500,14 +533,19 @@ public static void Write(string message) /// Prompts the user for confirmation with a yes/no question. /// /// The message to display to the user. + /// The interpreter executing the command, whose inactivity timeout is paused while waiting for the answer. Defaults to . /// true if the user confirms; otherwise, false. - public static bool Confirm(string message) + public static bool Confirm(string message, ShellInterpreter? shell = null) + => Confirm(message, shell ?? Instance, static () => Console.ReadKey(intercept: true)); + + internal static bool Confirm(string message, ShellInterpreter shell, Func readKey) { + using var pause = shell.PauseOperationTimeout(); var yes = char.ToUpper(MessageService.GetString("yes_char")[0]); var no = char.ToUpper(MessageService.GetString("no_char")[0]); // Take over Ctrl+C handling for the lifetime of the prompt so it cancels the - // question (returns false) instead of being swallowed by the global cancel-key + // operation instead of being swallowed by the global cancel-key // handler, which would leave this blocking ReadKey loop spinning forever. var restoreControlC = TrySetTreatControlCAsInput(true, out var originalTreatControlCAsInput); try @@ -519,7 +557,7 @@ public static bool Confirm(string message) ConsoleKeyInfo key; try { - key = Console.ReadKey(intercept: true); + key = readKey(); } catch (InvalidOperationException) { @@ -531,8 +569,8 @@ public static bool Confirm(string message) if (key.Key == ConsoleKey.C && key.Modifiers.HasFlag(ConsoleModifiers.Control)) { - WriteLine("^C"); - return false; + shell.CancelPrompt(); + throw new OperationCanceledException(); } if (key.Key == ConsoleKey.Escape) @@ -604,6 +642,11 @@ public void CancelPrompt() } } + this.RefreshPrompt(); + } + + internal void RefreshPrompt() + { this.editorCancelTokenSource.Cancel(); this.editorCancelTokenSource = new CancellationTokenSource(); } @@ -613,20 +656,33 @@ public void CancelPrompt() /// /// The command string to execute. /// A cancellation token to observe while waiting for the task to complete. + /// Whether unhandled caller cancellation should be reported as a failure. /// A representing the result of the command execution. - public async Task ExecuteCommandAsync(string command, CancellationToken token) + public async Task ExecuteCommandAsync(string command, CancellationToken token, bool cancellationIsError = false) { try { - return await this.RunSerializedAsync(() => this.ExecuteCommandCoreAsync(command, token), token); + return await this.RunSerializedAsync(() => this.ExecuteCommandCoreAsync(command, token, cancellationIsError), token); } - catch (OperationCanceledException) when (token.IsCancellationRequested) + catch (OperationCanceledException e) when (token.IsCancellationRequested) { this.Diagnostics?.LogCancelled(0, command); - return new CommandState(); + return cancellationIsError ? this.ReportCancellation(e, new CommandState(), command) : new CommandState(); } } + private CommandState ReportCancellation(OperationCanceledException exception, CommandState state, string command) + { + var error = new ErrorCommandState(new ShellException(MessageService.GetString("runtime-error-canceled"), exception)) + { + RequestCharge = state.RequestCharge, + }; + error.ShareWarningsFrom(state); + this.RenderWarnings(error, !string.IsNullOrEmpty(this.StdOutRedirect), this.IsMachineMode, deferToErrorEnvelope: true); + this.ReportExecutionError(error.Exception, command, error.Warnings); + return error; + } + /// /// Callers must not already hold the gate; nested work runs through the ungated core methods. /// @@ -643,7 +699,7 @@ internal async Task RunSerializedAsync(Func> operation, Cancellati } } - private async Task ExecuteCommandCoreAsync(string command, CancellationToken token) + private async Task ExecuteCommandCoreAsync(string command, CancellationToken token, bool cancellationIsError) { using var activity = TracingBootstrap.StartCommandActivity("cosmosdbshell.command"); var isLocalEmulatorOperation = this.State is ConnectedState connectedState @@ -665,20 +721,84 @@ private async Task ExecuteCommandCoreAsync(string command, Cancell CommandState? result = null; var wasCancelled = false; using var operationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(token); - if (isLocalEmulatorOperation) + var previousLocalEmulatorDeadline = this.localEmulatorDeadline; + var previousCancellationErrorToken = this.cancellationErrorToken; + this.cancellationErrorToken = cancellationIsError ? token : default; + using var deadline = isLocalEmulatorOperation + ? new OperationInactivityTimeout(operationTokenSource, this.LocalEmulatorInactivityTimeout, token, this.LocalEmulatorTimeProvider) + : null; + if (deadline is not null) { - operationTokenSource.CancelAfter(LocalEmulatorOperationTimeout); + Volatile.Write(ref this.localEmulatorDeadline, deadline); } try { + CommandState PrintStructuredError(StructuredErrorCommandState structuredError) + { + // A command can turn the inactivity cancellation into a structured partial result + // (bulk does); still report it as an emulator timeout and disconnect. + var timedOut = isLocalEmulatorOperation && operationTokenSource.IsCancellationRequested && !token.IsCancellationRequested; + if (timedOut) + { + structuredError.Exception = CreateEmulatorTimeoutException(structuredError.Exception); + if (structuredError.Result is ShellJson json + && json.Value.ValueKind == JsonValueKind.Object + && json.Value.TryGetProperty("exitCode", out _)) + { + var payload = JsonSerializer.SerializeToNode(json.Value)!.AsObject(); + payload["exitCode"] = structuredError.ExitCode; + if (payload.ContainsKey("error")) + { + payload["error"] = structuredError.Exception.Message; + } + + structuredError.Result = new ShellJson(JsonSerializer.SerializeToElement(payload)); + } + } + + var printed = this.PrintState(structuredError); + if (timedOut) + { + // PrintState already reports positional and redirected structured errors. + if (!this.IsMachineMode + && this.ErrOutRedirect == null + && structuredError.Exception is not Parser.PositionalException) + { + this.ReportExecutionError(structuredError.Exception); + } + + this.Disconnect(); + } + + return printed; + } + try { state = await this.RunCommandAsync(state, command, operationTokenSource.Token); + if (state is not StructuredErrorCommandState) + { + // Some commands consume cancellation and return partial results as a normal state. + Volatile.Read(ref this.localEmulatorDeadline)?.ThrowIfTimedOut(); + } } - catch (OperationCanceledException) when (token.IsCancellationRequested) + catch (CommandState.HandledCancellationException e) when (token.IsCancellationRequested) + { + wasCancelled = true; + e.State.ShareWarningsFrom(state); + result = this.PrintState(e.State); + return result; + } + catch (OperationCanceledException e) when (token.IsCancellationRequested) { wasCancelled = true; + if (cancellationIsError) + { + result = this.ReportCancellation(e, state, command); + return result; + } + result = new CommandState(); result.ShareWarningsFrom(state); this.RenderWarnings(result, !string.IsNullOrEmpty(this.StdOutRedirect), this.IsMachineMode, flushConsoleWarnings: true); @@ -686,9 +806,8 @@ private async Task ExecuteCommandCoreAsync(string command, Cancell } catch (OperationCanceledException e) when (isLocalEmulatorOperation && operationTokenSource.IsCancellationRequested) { - var shellException = new ShellException( - CommandException.GetDisplayMessage(System.Net.HttpStatusCode.RequestTimeout, e.Message), - e); + var message = CommandException.GetDisplayMessage(System.Net.HttpStatusCode.RequestTimeout, e.Message); + var shellException = new ShellException(message, new TimeoutException(message, e)); result = new ErrorCommandState(shellException); result.ShareWarningsFrom(state); this.RenderWarnings(result, !string.IsNullOrEmpty(this.StdOutRedirect), this.IsMachineMode, deferToErrorEnvelope: true); @@ -716,7 +835,7 @@ private async Task ExecuteCommandCoreAsync(string command, Cancell structuredError.Exception = e; } - result = this.PrintState(structuredError); + result = PrintStructuredError(structuredError); return result; } @@ -739,6 +858,24 @@ private async Task ExecuteCommandCoreAsync(string command, Cancell if (token.IsCancellationRequested) { wasCancelled = true; + if (state is StructuredErrorCommandState canceledStructuredError) + { + result = PrintStructuredError(canceledStructuredError); + return result; + } + + if (state.UserCancellationHandled) + { + result = this.PrintState(state); + return result; + } + + if (cancellationIsError) + { + result = this.ReportCancellation(new OperationCanceledException(token), state, command); + return result; + } + result = state; this.RenderWarnings(result, !string.IsNullOrEmpty(this.StdOutRedirect), this.IsMachineMode, flushConsoleWarnings: true); return result; @@ -755,11 +892,15 @@ private async Task ExecuteCommandCoreAsync(string command, Cancell return result; } - result = this.PrintState(state); + result = state is StructuredErrorCommandState returnedStructuredError + ? PrintStructuredError(returnedStructuredError) + : this.PrintState(state); return result; } finally { + Volatile.Write(ref this.localEmulatorDeadline, previousLocalEmulatorDeadline); + this.cancellationErrorToken = previousCancellationErrorToken; this.StdOutRedirect = savedStdOut; this.AppendOutRedirection = savedAppendOut; this.ErrOutRedirect = savedErrOut; @@ -1177,20 +1318,36 @@ internal Task ExecuteCosmosCommandAsync( CosmosCommand command, CommandState commandState, string commandText, - CancellationToken token) - => this.ExecuteCosmosCommandCoreAsync(command, commandState, commandText, token); + CancellationToken token, + bool isExternal = false) + => this.ExecuteCosmosCommandCoreAsync(command, commandState, commandText, token, isExternal); private async Task ExecuteCosmosCommandCoreAsync( CosmosCommand command, CommandState commandState, string commandText, - CancellationToken token) + CancellationToken token, + bool isExternal) { + using var localWorkPause = isExternal ? this.PauseOperationTimeout() : null; + // CommandState is intentionally reused by pipelines and expressions. Clear the // previous command's charge so an uninstrumented command cannot count it again. commandState.RequestCharge = null; var generation = this.SessionRequestChargeGeneration; using var requestChargeScope = RequestChargeContext.Begin(); + var previousInterruptFormat = Volatile.Read(ref this.activeInterruptOutputFormat); + var formatProperty = command.GetType().GetProperties().FirstOrDefault(property => + property.GetCustomAttribute()?.Names.Contains("format", StringComparer.OrdinalIgnoreCase) == true); + var effectiveFormat = OutputFormats.TryParse(formatProperty?.GetValue(command) as string, out var selectedFormat) + ? selectedFormat + : commandState.OutputFormatExplicitlySet ? commandState.OutputFormat : this.DefaultOutputFormat; + + // Nested command output must not add terminal markers to an enclosing structured result. + var interruptFormat = previousInterruptFormat >= 0 && OutputPolicy.IsStructuredFormat((OutputFormat)previousInterruptFormat) + ? previousInterruptFormat + : (int)effectiveFormat; + Volatile.Write(ref this.activeInterruptOutputFormat, interruptFormat); try { var result = await command.ExecuteAsync(this, commandState, commandText, token); @@ -1214,6 +1371,10 @@ private async Task ExecuteCosmosCommandCoreAsync( throw; } + finally + { + Volatile.Write(ref this.activeInterruptOutputFormat, previousInterruptFormat); + } } internal void RecordRequestCharge(CommandState commandState, long generation) @@ -1325,13 +1486,18 @@ internal async Task ConnectAsync(string connectionString, string? loginHint = nu WriteLine(MessageService.GetString("shell-connect-key-auth")); var keyMode = mode ?? (isEmulator ? ConnectionMode.Gateway : ConnectionMode.Direct); var keyOptions = CreateClientOptions(keyMode, isEmulator); + if (isEmulator) + { + keyOptions.CustomHandlers.Add(new OperationProgressHandler(this)); + } + client = new CosmosClient(connectionString, keyOptions); AccountProperties keyProps; using var operationTokenSource = CancellationTokenSource.CreateLinkedTokenSource(token); if (isEmulator) { - operationTokenSource.CancelAfter(LocalEmulatorOperationTimeout); + operationTokenSource.CancelAfter(TimeSpan.FromSeconds(LocalEmulatorOperationTimeoutSeconds)); } try @@ -1548,6 +1714,24 @@ private static async Task ReadAccountAsync(CosmosClient clien return await client.ReadAccountAsync().WaitAsync(token); } + /// + /// Re-labels an exception caused by the local-emulator inactivity timeout. Script location + /// frames stay outermost so diagnostics still identify the failing statement, and the cause + /// becomes a so the exit code reports a connection error. + /// + /// The exception produced when the inactivity timeout canceled the command. + /// The equivalent timeout exception. + internal static Exception CreateEmulatorTimeoutException(Exception exception) + { + if (exception is Parser.PositionalException { InnerException: { } inner } frame) + { + return new Parser.PositionalException(frame.FileName, CreateEmulatorTimeoutException(inner), frame.Line, frame.Column, frame.LineText); + } + + var message = CommandException.GetDisplayMessage(System.Net.HttpStatusCode.RequestTimeout, exception.Message); + return new ShellException(message, new TimeoutException(message, exception)); + } + internal static string GetLocalEmulatorConnectionFailureMessage(Uri endpoint) { var alternate = BuildAlternateEmulatorUri(endpoint); @@ -1948,6 +2132,76 @@ internal void DisconnectLocalEmulatorAfterConnectivityFailure(Exception exceptio } } + /// + /// Restarts the local-emulator inactivity timeout for the running command, so long-running + /// work is canceled only after it stops making progress. Has no effect for other connections. + /// + internal void ReportOperationProgress() => Volatile.Read(ref this.localEmulatorDeadline)?.ReportProgress(); + + /// + /// Suspends the local-emulator inactivity timeout while waiting for the user, such as at a + /// confirmation prompt or in an external editor. Dispose the result to resume it. + /// + /// A scope that resumes the timeout when disposed. + internal IDisposable PauseOperationTimeout() => + Volatile.Read(ref this.localEmulatorDeadline)?.Pause() ?? OperationInactivityTimeout.NoPause; + + /// + /// Draws a status line that replaces the previous one in place, such as command progress. + /// When the text does not fit the console width, the compact form is used instead, and is + /// truncated if necessary, so the line never wraps. + /// + /// The status text, without a line break. + /// A shorter form that keeps the key values, used on narrow consoles. + internal void WriteTransientLine(string text, string? compactText = null) + { + lock (this.transientLineLock) + { + var width = this.ConsoleWidth(); + if (width > 1 && text.Length >= width) + { + text = compactText ?? text; + if (text.Length >= width) + { + text = text[..(width - 1)]; + } + } + + Write("\r" + text + new string(' ', Math.Max(0, this.transientLineLength - text.Length))); + this.transientLineLength = text.Length; + } + } + + /// + /// Erases the status line drawn by , if any, so subsequent + /// output starts on a clean line. + /// + internal void ClearTransientLine() + { + lock (this.transientLineLock) + { + if (this.transientLineLength == 0) + { + return; + } + + Write("\r" + new string(' ', this.transientLineLength) + "\r"); + this.transientLineLength = 0; + } + } + + private static int GetConsoleWidth() + { + try + { + return Console.WindowWidth; + } + catch (Exception ex) when (ex is IOException or InvalidOperationException or PlatformNotSupportedException) + { + return 0; + } + } + internal void PrintCommand(string cmdString) { // Print the shell prompt similar to how it appears when typing command @@ -1993,6 +2247,15 @@ internal void PrintCommand(string cmdString) /// The rendered state, or an error state when rendering failed. internal CommandState PrintState(CommandState state, bool markAsRendered = false, bool deferMachineModeWarnings = false) { + if (deferMachineModeWarnings && state is not StructuredErrorCommandState) + { + Volatile.Read(ref this.localEmulatorDeadline)?.ThrowIfTimedOut(); + if (!state.UserCancellationHandled) + { + this.cancellationErrorToken.ThrowIfCancellationRequested(); + } + } + var result = this.PrintStateCore(state, markAsRendered); if (!ReferenceEquals(result, state)) { @@ -2322,53 +2585,72 @@ private void RenderWarnings( warnings.AddRange(state.TakeDeferredWarnings()); } - foreach (var warning in warnings) + // A warning raised mid-command must not be appended to an in-place progress line. Hold the + // progress-line lock until the warnings are written so a heartbeat cannot redraw in between. + var holdsProgressLine = false; + try { - if (warning.IsRendered - && (!renderUpdates || string.Equals(warning.LastRenderedMessage, warning.Message, StringComparison.Ordinal))) + if (warnings.Count > 0) { - continue; + Monitor.Enter(this.transientLineLock, ref holdsProgressLine); + this.ClearTransientLine(); } - if (machineModeWarningRecords - && (!warning.IsRendered || (renderUpdates && warning.EmittedAsMachineModeWarningRecord))) - { - Console.Error.WriteLine(JsonSerializer.Serialize(new - { - status = renderUpdates && warning.EmittedAsMachineModeWarningRecord ? "warning-update" : "warning", - warning = warning.Message, - })); - warning.EmittedAsMachineModeWarningRecord = true; - } - else if (warning.DeferredToConsole) + foreach (var warning in warnings) { - // Held until the top-level command finalizes; see DeferMachineModeConsoleWarnings. - if (!flushConsoleWarnings) + if (warning.IsRendered + && (!renderUpdates || string.Equals(warning.LastRenderedMessage, warning.Message, StringComparison.Ordinal))) { continue; } - Console.Error.WriteLine(warning.Message); - } - else if (this.ErrOutRedirect != null) - { - this.WriteErrorRedirect(warning.Message + Environment.NewLine); - this.ErrorRedirectWarnings.HasWarning = true; - } - else if (!redirected && !inMachineMode) - { - AnsiConsole.MarkupLine(Theme.FormatWarning(Markup.Escape(warning.Message))); + if (machineModeWarningRecords + && (!warning.IsRendered || (renderUpdates && warning.EmittedAsMachineModeWarningRecord))) + { + Console.Error.WriteLine(JsonSerializer.Serialize(new + { + status = renderUpdates && warning.EmittedAsMachineModeWarningRecord ? "warning-update" : "warning", + warning = warning.Message, + })); + warning.EmittedAsMachineModeWarningRecord = true; + } + else if (warning.DeferredToConsole) + { + // Held until the top-level command finalizes; see DeferMachineModeConsoleWarnings. + if (!flushConsoleWarnings) + { + continue; + } + + Console.Error.WriteLine(warning.Message); + } + else if (this.ErrOutRedirect != null) + { + this.WriteErrorRedirect(warning.Message + Environment.NewLine); + this.ErrorRedirectWarnings.HasWarning = true; + } + else if (!redirected && !inMachineMode) + { + AnsiConsole.MarkupLine(Theme.FormatWarning(Markup.Escape(warning.Message))); + } + else + { + Console.Error.WriteLine(warning.Message); + } + + warning.IsRendered = true; + warning.LastRenderedMessage = warning.Message; + warning.RenderedErrorRedirect = warning.DeferredToConsole || this.ErrOutRedirect == null + ? null + : Path.GetFullPath(this.ErrOutRedirect); } - else + } + finally + { + if (holdsProgressLine) { - Console.Error.WriteLine(warning.Message); + Monitor.Exit(this.transientLineLock); } - - warning.IsRendered = true; - warning.LastRenderedMessage = warning.Message; - warning.RenderedErrorRedirect = warning.DeferredToConsole || this.ErrOutRedirect == null - ? null - : Path.GetFullPath(this.ErrOutRedirect); } } @@ -2623,8 +2905,17 @@ private void Console_CancelKeyPress(object? sender, ConsoleCancelEventArgs e) this.cosmosShellPrompt.InContinuation = false; } + var activeFormat = Volatile.Read(ref this.activeInterruptOutputFormat); + var effectiveFormat = activeFormat >= 0 ? (OutputFormat)activeFormat : this.DefaultOutputFormat; + var showMarker = !this.IsMachineMode + && !OutputPolicy.IsStructuredFormat(effectiveFormat) + && this.IsOutputTerminal() + && string.IsNullOrEmpty(this.StdOutRedirect); this.CancelPrompt(); - WriteLine("̂C"); + if (showMarker) + { + WriteLine("^C"); + } } private void SaveHistory() diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs index e4deb029..10c79ec3 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Mcp/ToolOperations.cs @@ -795,8 +795,8 @@ internal async Task ExecuteToolAsync( bulk.ConfirmationApproved = confirmedVersion.HasValue; } - var response = await shell.ExecuteCosmosCommandAsync(cmd, commandState, command.CommandName, cancellationToken); - shell.CancelPrompt(); + var response = await shell.ExecuteCosmosCommandAsync(cmd, commandState, command.CommandName, cancellationToken, command.IsExternal); + shell.RefreshPrompt(); return McpResponseFactory.CreateSuccess(response, shell.State); } finally diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs index 5714f596..9a623a61 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Expression/CommandExpression.cs @@ -115,6 +115,11 @@ public override async Task EvaluateAsync(ShellInterpreter interpret throw failure; } + if (resultState.UserCancellationHandled) + { + throw new CommandState.HandledCancellationException(resultState, cancellationToken); + } + // Return the result from the command state if (resultState.Result != null) { @@ -170,7 +175,7 @@ internal async Task ExecuteCommandAsync(ShellInterpreter shell, Co if (shell.App.Commands.TryGetValue(this.Name, out var factory)) { var cmd = await this.CreateCommandAsync(factory, shell, commandState, token); - return await shell.ExecuteCosmosCommandAsync(cmd, commandState, string.Empty, token); + return await shell.ExecuteCosmosCommandAsync(cmd, commandState, string.Empty, token, factory.IsExternal); } // Check for script files diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/BlockStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/BlockStatement.cs index 1952d5d5..76f9460f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/BlockStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/BlockStatement.cs @@ -111,7 +111,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman return commandState; } - if (commandState.IsError) + if (commandState.IsError || commandState.UserCancellationHandled) { return commandState; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs index 17a4b012..8658d992 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/CommandStatement.cs @@ -191,7 +191,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman } var cmd = await this.CreateCommandAsync(factory, shell, commandState, token); - var result = await shell.ExecuteCosmosCommandAsync(cmd, commandState, string.Empty, token); + var result = await shell.ExecuteCosmosCommandAsync(cmd, commandState, string.Empty, token, factory.IsExternal); if (result is ErrorCommandState error && shell.CurrentScriptFileName is { } sourceName && shell.CurrentScriptContent is { } sourceText) { var (line, column, lineText) = PositionalErrorHelper.GetLineAndColumn(sourceText, this.Start); @@ -453,7 +453,7 @@ public async Task RunScriptAsync(ShellInterpreter shell, CommandSt { // Run the parsed statements currentState = await statement.RunAsync(shell, currentState, token); - if (currentState.IsError) + if (currentState.IsError || currentState.UserCancellationHandled) { break; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/DoWhileStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/DoWhileStatement.cs index a170f606..b98d432f 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/DoWhileStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/DoWhileStatement.cs @@ -92,7 +92,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman throw; } - if (commandState.IsError || commandState.ReturnFunc) + if (commandState.IsError || commandState.UserCancellationHandled || commandState.ReturnFunc) { return commandState; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/ForStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/ForStatement.cs index 693544e4..003b3c8d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/ForStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/ForStatement.cs @@ -136,7 +136,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman throw; } - if (commandState.IsError || commandState.ReturnFunc) + if (commandState.IsError || commandState.UserCancellationHandled || commandState.ReturnFunc) { return commandState; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/LoopStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/LoopStatement.cs index 9bcb25ae..5d5c3c3d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/LoopStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/LoopStatement.cs @@ -95,7 +95,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman throw; } - if (commandState.IsError || commandState.ReturnFunc) + if (commandState.IsError || commandState.UserCancellationHandled || commandState.ReturnFunc) { return commandState; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/PipeStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/PipeStatement.cs index 4d130d0f..63cf5c93 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/PipeStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/PipeStatement.cs @@ -84,16 +84,16 @@ public override async Task RunAsync(ShellInterpreter shell, Comman { foreach (var statement in this.Statements) { - token.ThrowIfCancellationRequested(); - if (commandState.IsError || commandState.BreakBlock || commandState.ContinueBlock || commandState.ReturnFunc) + if (commandState.IsError || commandState.UserCancellationHandled || commandState.BreakBlock || commandState.ContinueBlock || commandState.ReturnFunc) { return commandState; } + token.ThrowIfCancellationRequested(); commandState = await statement.RunAsync(shell, commandState, token); } - if (commandState.IsError || commandState.BreakBlock || commandState.ContinueBlock || commandState.ReturnFunc) + if (commandState.IsError || commandState.UserCancellationHandled || commandState.BreakBlock || commandState.ContinueBlock || commandState.ReturnFunc) { return commandState; } diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/WhileStatement.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/WhileStatement.cs index f9ab75b6..e1b2a23b 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/WhileStatement.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Parser/Statement/WhileStatement.cs @@ -88,7 +88,7 @@ public override async Task RunAsync(ShellInterpreter shell, Comman throw; } - if (commandState.IsError || commandState.ReturnFunc) + if (commandState.IsError || commandState.UserCancellationHandled || commandState.ReturnFunc) { return commandState; } diff --git a/CosmosDBShell/Program.cs b/CosmosDBShell/Program.cs index 1f8e47db..38801031 100644 --- a/CosmosDBShell/Program.cs +++ b/CosmosDBShell/Program.cs @@ -18,6 +18,36 @@ internal class Program { private const int DefaultMcpPort = 6128; + internal static async Task<(CommandState State, bool Interrupted)> ExecuteStartupCommandAsync(ShellInterpreter shell, string command) + { + using var cancellation = ShellInterpreter.UserCancellationTokenSource; + var state = await shell.ExecuteCommandAsync(command, cancellation.Token, cancellationIsError: true); + return (state, cancellation.IsCancellationRequested); + } + + internal static async Task<(CommandState? State, bool Interrupted)> ExecuteStartupInputAsync(ShellInterpreter shell, TextReader input) + { + using var cancellation = ShellInterpreter.UserCancellationTokenSource; + var token = cancellation.Token; + try + { + // Console.In can perform synchronous reads even through its async API. + var script = await Task.Run(() => input.ReadToEndAsync(token), token).WaitAsync(token); + token.ThrowIfCancellationRequested(); + if (string.IsNullOrWhiteSpace(script)) + { + return (null, false); + } + + var state = await shell.ExecuteCommandAsync(script + Environment.NewLine, token, cancellationIsError: true); + return (state, cancellation.IsCancellationRequested); + } + catch (OperationCanceledException) when (token.IsCancellationRequested) + { + return (new ErrorCommandState(new OperationCanceledException(token)), true); + } + } + public static async Task Main(string[] args) { // Normalize argv first so that any --lsp/--stdio token that is part @@ -441,39 +471,35 @@ await ShellInterpreter.Instance.ConnectAsync( if (Console.IsInputRedirected) { - // Read entire stdin (script) - var script = await Console.In.ReadToEndAsync() + Environment.NewLine; - if (!string.IsNullOrWhiteSpace(script)) + var (state, interrupted) = await ExecuteStartupInputAsync(ShellInterpreter.Instance, Console.In); + if (state?.IsError == true) { - var state = await ShellInterpreter.Instance.ExecuteCommandAsync(script, default); - if (state.IsError) - { - Environment.ExitCode = ShellExitCode.FromCommandState(state); - if (executeAndContinueCommand is null) - { - await StopHostAsync(host, hostTask); + Environment.ExitCode = ShellExitCode.FromCommandState(state); + } - return; - } - } + if (interrupted || (state?.IsError == true && executeAndContinueCommand is null)) + { + await StopHostAsync(host, hostTask); - // If user only wants to execute piped script then quit (unless -k / ExecuteAndContinue) - if (executeAndContinueCommand is null && executeAndQuitCommand is null) - { - await StopHostAsync(host, hostTask); + return; + } - return; - } + // If user only wants to execute piped script then quit (unless -k / ExecuteAndContinue) + if (state is not null && executeAndContinueCommand is null && executeAndQuitCommand is null) + { + await StopHostAsync(host, hostTask); + + return; } } if (explicitCommand is not null) { - var state = await ShellInterpreter.Instance.ExecuteCommandAsync(explicitCommand, default); + var (state, interrupted) = await ExecuteStartupCommandAsync(ShellInterpreter.Instance, explicitCommand); if (state.IsError) { Environment.ExitCode = ShellExitCode.FromCommandState(state); - if (executeAndContinueCommand is null) + if (executeAndContinueCommand is null || interrupted) { await StopHostAsync(host, hostTask); diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index 1c0ca28c..60e46fb0 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -258,6 +258,10 @@ command-rm-dry-run-plan = Dry run: Would delete { $count } { $count -> [one] item *[other] items }. No changes were made. +command-rm-progress = Remove: { $processed } items scanned, { $succeeded } deleted; elapsed { $elapsed } s (observed RU charge: { $charge }) +command-rm-progress-compact = Remove: { $succeeded } deleted, { $processed } scanned, { $elapsed } s, { $charge } RU +command-rm-dry-run-progress = Dry run: { $processed } items scanned, { $succeeded } would be deleted; elapsed { $elapsed } s (observed RU charge: { $charge }) +command-rm-dry-run-progress-compact = Dry run: { $succeeded } to delete, { $processed } scanned, { $elapsed } s, { $charge } RU command-rm-error-no_filter = Filter is missing. command-rm-error-invalid_pk_json = Invalid partition key JSON. Use { "{" }{ "}" }, a JSON scalar, or an array of one to three string, finite number, boolean, null, or { "{" }{ "}" } components. Non-empty objects and nested arrays are not supported. command-rm-error-partition_key_components = The partition key must specify all { $expected } components of the container's partition key, but { $actual } { $actual -> @@ -497,7 +501,8 @@ command-bulk-status-succeeded = Succeeded command-bulk-status-failed = Failed command-bulk-status-uncertain = Uncertain command-bulk-status-more = ...and { $count } more -command-bulk-progress = { $processed } of { $total } operations processed, { $failed } failed (RU charge: { $charge }) +command-bulk-progress = Bulk: { $processed } of { $total } operations processed, { $succeeded } succeeded, { $failed } failed; elapsed { $elapsed } s, { $rate } operations/s (observed RU charge: { $charge }) +command-bulk-progress-compact = Bulk: { $processed }/{ $total }, { $failed } failed, { $elapsed } s, { $rate }/s, { $charge } RU command-bulk-success = Bulk completed { $count } { $count -> [one] operation *[other] operations @@ -573,6 +578,10 @@ command-import-description-concurrency = Maximum in-flight writes (default: 16). command-import-description-partition-key = For CSV import, the partition key path. Nested paths (e.g. /address/city) place the matching column under that path. command-import-description-continue-on-error = Continue importing after individual item write failures. Otherwise stop scheduling writes on the first observed failure and wait for in-flight writes. Parse or validation errors still abort the import. command-import-description-dry-run = Validate file syntax and item IDs offline without writing items or accessing the target. Does not check target partition keys, permissions, conflicts, or other service constraints. +command-import-progress = Import: { $processed } items processed, { $succeeded } imported, { $failed } failed; elapsed { $elapsed } s, { $rate } items/s (observed RU charge: { $charge }) +command-import-progress-compact = Import: { $succeeded } ok, { $failed } failed, { $elapsed } s, { $rate }/s, { $charge } RU +command-import-validation-progress = Dry run: { $processed } items validated offline; elapsed { $elapsed } s, { $rate } items/s; no writes sent +command-import-validation-progress-compact = Dry run: { $processed } validated, { $elapsed } s, { $rate }/s command-import-success = Imported { $count } { $count -> [one] item *[other] items diff --git a/docs/commands.md b/docs/commands.md index 76ddafa3..7e78d8e1 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -669,6 +669,8 @@ batch execute Execute many **independent** write operations with bounded concurrency. `bulk` uses the same operation JSON and subcommands as [`batch`](#batch), with different execution semantics: +During write execution, human-readable output shows progress about every two seconds, including processed/total operations, successes, failures, elapsed seconds, average processing rate, and observed RU charge. The progress line is updated in place and removed when the command finishes. Updates continue while writes are waiting; unchanged counts do not by themselves indicate a stalled job. Validation, selection, confirmation, and journal verification happen before this execution progress starts. Progress appears only when standard output is a terminal; quiet mode, structured output, shell or operating-system output redirection, operating-system pipes, and MCP calls do not emit it. Within an in-shell `|` pipeline, an earlier stage still shows progress on the terminal; it is erased before the next stage runs and is never part of the piped data. See [ingestion sizing and limitations](ingestion.md) for planning guidance and emulator measurements. + | | `batch` | `bulk` | |-|-|-| | Partition keys | One per batch (`--partition-key`) | Per operation `partitionKey`, or `--partition-key` as a default | @@ -825,6 +827,8 @@ Examples: - `rm test-* --dry-run` reports how many items match without deleting anything. - `rm tmp-* --key=id --partition-key=customer-42 --dry-run` counts matching items only in the `customer-42` logical partition. +Pattern-based and piped removals show progress about every two seconds in human-readable output: items scanned, items deleted (or that would be deleted in a dry run), elapsed seconds, and observed RU charge. The progress line is updated in place and removed when the command finishes. Unscoped deletes scan the whole container before deleting, so the scanned count stops rising while the deleted count catches up. Items are deleted one at a time. Progress appears only when standard output is a terminal; quiet mode, structured output, output redirection, operating-system pipes, and MCP calls suppress it without changing the result. + #### Partition-key scoping `--partition-key` (alias `--pk`) restricts both the dry run and the deletion to a single logical partition. It uses the same typed parsing as `patch` and `batch`: JSON scalars keep their type (`--pk=7` targets the number `7`, `--pk='"7"'` targets the string `"7"`), `{}` targets an undefined value, and hierarchical keys are passed as a JSON array (`--pk='["tenant-1",{}]'`). The value must supply every component of the container's partition key; a prefix of a hierarchical key is rejected. Piped items outside the partition are skipped. @@ -889,6 +893,8 @@ Normal output shows a readable summary with the number of items written and the Successful imports display a readable summary in normal output mode. `--format=json`, `--format=csv`, and `--format=table` select the command-result presentation; they do not change how the input file is read. When `--format` is omitted, the result inherits the shell's output settings. +Human-readable output shows progress about every two seconds: completed item counts, imported/failed counts, elapsed seconds, average processing rate, and observed write RU charge. The progress line is updated in place and removed when the command finishes. Updates continue while reads or writes are waiting. Dry runs report offline validation counts instead. Because input is streamed without a counting pass, import does not report a total, percentage, or ETA. Progress appears only when standard output is a terminal; quiet mode, structured output, output redirection, operating-system pipes, and MCP calls suppress it without changing the final result. See [ingestion sizing and limitations](ingestion.md) for planning guidance and emulator measurements. + Invalid result formats are rejected before accessing the target container or writing items. Diagnostics for JSON object sequences and CSV records use physical start-line numbers; JSON arrays use 1-based item indexes. Non-object roots in an object sequence are rejected before their contents are loaded. Bulk-load items from a JSON, JSON Lines, or CSV file into a container. A `.csv` extension selects CSV. For every other file, import detects the structure from its contents: a top-level `[` selects an array of items (the format the Azure Databases extension for VS Code imports and exports); otherwise it reads a single object or a sequence of objects. Sequences support JSON Lines, multiple objects on one line, and objects spanning several lines, without commas between objects or an enclosing array. No `--legacy` flag is needed, regardless of the filename extension. Malformed JSON, non-object items, and trailing garbage still fail import. Default mode is `insert`; pass `--mode=upsert` to replace items that already exist. For CSV, the header row defines property names and every value is imported as a string; the CSV separator follows `COSMOSDB_SHELL_CSVSEP` (default `;`). Empty headers with empty cells are ignored; a value under an empty header fails import with its line and column instead of being silently discarded. All formats are read incrementally rather than loading the complete file into memory. CSV supports quoted separators, escaped quotes, and multiline fields; malformed records abort the import with their physical start line. When an imported item omits a configured partition-key path, Cosmos DB infers an undefined value. Interactively and when stderr is redirected, import displays a notice before the first affected write, then reports one final bounded warning counted from writes scheduled. Machine mode with console stderr emits the pre-write notice and final count as JSON Lines warning and `warning-update` records. Invalid partition-key values are reported as per-item failures; `--continue-on-error` continues with subsequent rows. Earlier writes are not rolled back, so use `--dry-run` first when the entire file must be validated before any writes. @@ -967,6 +973,8 @@ Options: By default `watch` starts from now and follows the container, printing each change as highlighted JSON until you press Ctrl+C. Use `--from-beginning` to replay existing items first, `--partition-key` to scope the feed to one partition, and `--max` to stop automatically after a number of changes. Use `--interval` to change how long the shell waits between polls once it has caught up (default 1 second; values below 0.1 are clamped to avoid hammering the container). The change feed surfaces creates and updates (not deletes). This command is interactive and streaming, so it is not exposed over MCP. +On a local emulator, intentional waits between polls do not count toward the 10-second inactivity timeout. A stalled request still times out and disconnects; Ctrl+C remains a normal way to stop watching. + ```bash watch watch --from-beginning diff --git a/docs/ingestion.md b/docs/ingestion.md new file mode 100644 index 00000000..7783dc08 --- /dev/null +++ b/docs/ingestion.md @@ -0,0 +1,149 @@ +# Ingestion sizing, progress, and limitations + +Use [`import`](commands.md#import) for files of items (JSON arrays, JSON object sequences/JSON Lines, or CSV). Use [`bulk run`](commands.md#bulk) for operation plans, mixed mutations, or journal-based recovery. Neither command is a transaction: completed writes are not rolled back. + +## Choosing a workload size + +There is no single recommended maximum file size or completion time. Item size and indexed properties, available RU/s, partition-key distribution, network latency, client resources, and competing traffic all matter. A file's byte count alone is not enough to predict its write cost. + +1. Start with **1,000 representative items** on a disposable container with the intended partition key, indexing policy, and throughput. This is a pilot size, not a product limit. +2. Keep the default **concurrency of 16** initially. Compare 32 only if the pilot shows spare capacity; increasing concurrency does not increase RU/s or fix a hot partition. Use 1 for sequential import ordering, accepting lower throughput. +3. Measure elapsed time, successful writes, failures, and observed RU charge. Then try 10,000 items before committing to a much larger job. The [Azure measurements](#azure-measurements-west-europe) below show one reference configuration; repeat them with your own item shape, indexing, throughput, and region before quoting completion times. +4. Choose chunks around the duration you can afford to restart or reconcile. An initial estimate is `items per chunk = measured items/second * acceptable seconds per chunk`; remeasure when item shape or target settings change. +5. Budget for other workloads. `items * observed RU/item / available RU/s` is only an optimistic throughput-based duration estimate, not a completion-time guarantee. Retries, skew, client limits, and network overhead can make a job slower. + +For example: + +```text +import pilot.jsonl --dry-run +import pilot.jsonl --concurrency=16 + +bulk run operations.jsonl --dry-run +bulk run operations.jsonl --concurrency=16 --journal operations.journal --yes +``` + +`import --dry-run` is offline syntax/ID validation, not a service acceptance or performance test. `bulk --dry-run` validates operations against target metadata; query-driven bulk selection can consume RUs even in a dry run. + +## Reading progress + +Human-readable import and bulk execution report about every two seconds, even if no writes have completed since the previous update. A single progress line is updated in place and removed when the command finishes, leaving the final summary: + +- Import shows completed items, imported/failed counts, elapsed seconds, average completed items/second, and observed write RU charge. Offline validation shows validated items instead. +- Bulk execution shows processed/total operations, succeeded/failed counts, elapsed seconds, average processed operations/second, and observed RU charge. Processed operations can include journal skips, so this is not necessarily a new-write rate. Failures include uncertain outcomes; use the final result/journal to distinguish them. +- Bulk execution progress starts **after** validation/selection, confirmation, and journal verification. These preparation steps can add time before the first update. +- These are cumulative averages, not instantaneous rates. Observed RU charge covers recorded responses, not a billing estimate or a count of internal SDK retries. Unchanged counts can mean input work, pending requests, or SDK retry waits; they do not prove a hang or identify throttling. +- Streamed imports do not pre-count records, so they have no percentage or ETA. Bulk knows its operation count. +- Progress appears only when standard output is a terminal. Quiet mode, structured output, shell or operating-system output redirection, operating-system pipes, and MCP calls do not emit it. In an in-shell `|` pipeline, an earlier stage still shows progress on the terminal, but the line is erased before the next stage runs and never becomes part of the piped data. Final structured results retain their existing fields. + +## Limits and recovery + +- **Emulator inactivity timeout:** commands connected to `localhost` or `127.0.0.1` are canceled, and the shell disconnects, after **10 seconds without progress**. Every response the emulator returns counts as progress, as does local import/bulk work such as validating and spooling operations. Long imports, deletes, queries, exports, and scripts therefore complete as long as the emulator keeps answering. The timer is paused while the shell waits for you at a confirmation prompt or in an external editor, during intentional `watch` polling intervals, and while local external tools such as `jq` run. A stalled or unreachable emulator still times out. Earlier writes remain committed. This is an emulator-only safeguard, not an input-size limit or an Azure service limit. +- **Item size:** the NoSQL service limit is **2 MB per item**, measured as the UTF-8 JSON representation, not 2 MB per input file. Import dry runs do not enforce this service limit. See [service limits](https://learn.microsoft.com/azure/cosmos-db/concepts-limits#per-item-limits). +- **File versus in-memory jobs:** import streams its input. File-based `bulk run` validates into a temporary operation spool before writing, so allow temporary-disk capacity as well as input-file space. Stateful `bulk begin`/`add` jobs are limited to **10,000 operations and 64 MiB of canonical JSON**; use `bulk run ` for larger plans. +- **Partial failure:** by default, new scheduling stops after an observed write failure, but started writes settle and can succeed afterward. `--continue-on-error` allows more writes; it does not make malformed input safe. Always inspect the final result. +- **Retry safety:** import has no resumable journal. Replaying inserts can conflict; upsert can replace existing items and must not be used merely to hide uncertainty. Reconcile completed writes before retrying. For bulk recovery, keep the same operation plan and journal; do not blindly retry uncertain writes or assume a new journal prevents replay. +- **Not covered by these measurements:** larger documents, CSV/array parsing differences, skewed keys, large indexes, mixed mutations, ETag conflicts, journal overhead, injected failures, memory ceilings, multi-region accounts, throttled (429) workloads, and clients far from the account region. No supported maximum file size is established by these matrices. + +## Azure measurements (West Europe) + +These are measurements against a real Azure Cosmos DB account for one reference configuration. They show the shape of the trade-offs; they are not service guarantees. + +Configuration: + +- Account `mikkrg` (provisioned throughput), single region **West Europe**, Session consistency, no availability zones. The client ran in Europe (about 55–60 ms HTTPS round trip to the account endpoint) on the same Windows 11 host as the emulator runs. +- Entra ID sign-in (Azure CLI identity); the shell's default **Direct** connection mode for token-based connections. Cosmos .NET SDK 3.63.2; .NET 10. +- One container with partition key `/pk`, default indexing, and **4,000 RU/s manual** throughput. All 18 cases wrote into this container; each case used its own ID prefix and a `caseId` property for verification, and items stayed exactly 1,024 bytes. +- Same workload, matrix, timing boundaries, and verification as the emulator runs: insert/create only, 1,000 evenly used partition-key values, one sample per combination, run sequentially on 2026-10-29. + +### Results (2026-10-29) + +All 18 cases completed with matching item counts (verified by count query) and zero reported write failures. Times are measured seconds. + +| Items | Concurrency | Import (seconds) | Bulk create (seconds) | Rate (items/s) | +|---:|---:|---:|---:|---:| +| 1,000 | 1 | 24.99 | 23.54 | 40–43 | +| 1,000 | 16 | 6.55 | 3.21 | 153–312 | +| 1,000 | 32 | 4.00 | 3.59 | 250–278 | +| 10,000 | 1 | 224.03 | 225.47 | 44–45 | +| 10,000 | 16 | 27.17 | 25.65 | 368–390 | +| 10,000 | 32 | 26.44 | 25.43 | 378–393 | +| 100,000 | 1 | 2,189.28 (36.5 min) | 2,209.65 (36.8 min) | 45–46 | +| 100,000 | 16 | 256.65 (4.3 min) | 255.23 (4.3 min) | 390–392 | +| 100,000 | 32 | 256.22 (4.3 min) | 255.95 (4.3 min) | 390–391 | + +What the numbers show: + +- **Cost per item:** each 1 KiB insert cost **7.0 RU** with default indexing (700,000 RU for 100,000 items; bulk adds 2 RU for target metadata). Plan RU budgets from measured charges, not from item count alone. +- **Concurrency 1 is latency-bound:** about 45 items/s, roughly 22 ms per write. Sequential ingestion of 100,000 items took about 37 minutes regardless of command. Use it only when file order matters. +- **The default concurrency (16) was about 8.5× faster** than sequential: about 390 items/s, or roughly 1.4 million items per hour, so 100,000 items took about 4.3 minutes. Increasing to 32 gave no further gain. +- **The plateau was not RU throttling.** Azure Monitor recorded no 429 (throttled) responses during the run, and peak normalized RU consumption was 83%. About 390 items/s × 7 RU is roughly 2,730 RU/s, below the provisioned 4,000 RU/s. The limiting factor (for example per-partition write latency or client-side limits) was not isolated, so raising RU/s alone is not guaranteed to raise this rate. Measure before provisioning for a target duration. +- **Small jobs are dominated by fixed overhead.** At 1,000 items, connection warm-up and metadata reads make results noisy; use 10,000 items or more to estimate sustained rates. +- `import` and `bulk create` performed the same for this workload. Choose between them by feature (file formats versus journals and mixed operations), not speed. + +### Reproducing the Azure matrix + +The opt-in `AzureIngestionBenchmarkTests` signs in with the Azure CLI identity (`az login`), which needs a data-plane write role (for example *Cosmos DB Built-in Data Contributor*) and Azure Resource Manager rights to create and delete databases. It creates one uniquely named database with a container at the configured manual throughput, and deletes the database at the end. **It incurs provisioned-throughput charges for the duration of the run** (about 1 hour 40 minutes for the full matrix in this configuration; 4,000 RU/s in West Europe was about $0.32 per hour at the time of measurement). Run it only against an account you are authorized to use for testing. + +```powershell +$env:COSMOSDB_SHELL_AZURE_BENCHMARK_ENDPOINT = 'https://.documents.azure.com:443/' +$env:COSMOSDB_SHELL_AZURE_BENCHMARK_SUBSCRIPTION = '' +$env:COSMOSDB_SHELL_AZURE_BENCHMARK_RESOURCE_GROUP = '' +$env:COSMOSDB_SHELL_AZURE_BENCHMARK_RU = '4000' # optional; defaults to 4000 +$env:COSMOSDB_SHELL_AZURE_BENCHMARK_CSV = Join-Path $PWD 'azure-ingestion-benchmark.csv' +dotnet test CosmosDBShell.Tests\CosmosDBShell.Tests.csproj --filter 'FullyQualifiedName~AzureIngestionBenchmarkTests' +``` + +Without these variables the test is skipped. If a run is interrupted, delete any remaining `IngestionBenchmark_*` database. + +## Local emulator measurements + +These are **local command-path measurements, not Azure throughput promises**. + +Configuration: + +- Windows 11, 8 logical processors, approximately 16 GiB RAM; shared development host. +- Windows Cosmos DB Emulator 2.14.28.0; Cosmos .NET SDK 3.63.2; .NET 10. +- Gateway connection to the local emulator. Fresh container per case, partition key `/pk`, default indexing, configured manual throughput of 10,000 RU/s. +- Insert/create only, unique IDs, 1,000 evenly used partition-key values. Each item is exactly 1,024 UTF-8 JSON bytes; JSON Lines adds one LF per item. Bulk input also includes the operation envelope. +- One sample per combination, sequentially run; no journal, no terminal progress rendering, no deliberate failure injection, and no controlled cold/warm-cache comparison. +- Timing includes command parsing, metadata access, input parsing, and writes (plus bulk plan validation/spooling). Input generation, container provisioning, independent count verification, and cleanup are outside the timed interval. The parsed command runner is invoked directly, outside the emulator inactivity timeout; SDK request/retry behavior is unchanged. +- Each completed case requires zero reported write failures and an independent count matching the requested number of items. + +### Results (2026-10-29) + +All 18 cases completed with matching item counts and zero reported write failures. Times below are measured seconds, rounded to two decimal places; they are not extrapolations. + +| Items | Concurrency | Import (seconds) | Bulk create (seconds) | +|---:|---:|---:|---:| +| 1,000 | 1 | 15.80 | 15.55 | +| 1,000 | 16 | 5.19 | 5.28 | +| 1,000 | 32 | 5.18 | 5.15 | +| 10,000 | 1 | 155.42 | 155.34 | +| 10,000 | 16 | 51.81 | 51.68 | +| 10,000 | 32 | 51.50 | 51.49 | +| 100,000 | 1 | 1,562.95 | 1,556.85 | +| 100,000 | 16 | 518.49 | 519.02 | +| 100,000 | 32 | 517.23 | 534.51 | + +At 100,000 items, input files were 102,500,000 bytes for import (about 97.75 MiB) and 104,800,000 bytes for bulk (about 99.95 MiB). Import reported 6,620 / 66,200 / 662,000 write RUs for the three sizes; bulk reported 2 additional RUs per case for target metadata. Verification and provisioning RUs are not included in these command results. + +On this host, concurrency 16 was about three times faster than 1; 32 showed no consistent improvement over 16. This supports keeping 16 as the initial setting, not advertising 32 as automatically faster. A 100,000-item file completed through the command runner, but that is not a maximum-size recommendation. Do not copy these local rates into Azure customer SLAs. + +### Reproducing the matrix + +The opt-in test creates a uniquely named database on **the local emulator only** and deletes its containers and database afterward. It performs 666,000 writes across 18 cases, so allow substantial time and do not run it alongside other performance tests. Existing emulator data is not reset. + +From the repository root in PowerShell: + +```powershell +$env:COSMOSDB_SHELL_INGESTION_BENCHMARK_CSV = Join-Path $PWD 'ingestion-benchmark.csv' +try { + dotnet test CosmosDBShell.Tests\CosmosDBShell.Tests.csproj --filter 'FullyQualifiedName~Integration.IngestionBenchmarkTests' +} finally { + Remove-Item Env:\COSMOSDB_SHELL_INGESTION_BENCHMARK_CSV +} +``` + +The CSV is overwritten at the chosen path and flushed after every verified case. It records input bytes, concurrency, elapsed seconds, observed RU charge, rate, and verified item count. Without the environment variable, the benchmark is skipped. + +These local measurements cannot establish Azure service capacity; see [Azure measurements](#azure-measurements-west-europe) for a real-account reference, and repeat representative tests with your own indexing, throughput, region, partition distribution, and item sizes. diff --git a/docs/mcp.md b/docs/mcp.md index 816f7a79..70f76a89 100644 --- a/docs/mcp.md +++ b/docs/mcp.md @@ -133,6 +133,8 @@ For deterministic ARM routing in multi-subscription environments, start the shel MCP `import` uses the same bounded concurrency as the CLI: up to 16 writes in flight by default, configurable with the `concurrency` argument. Its schema requires a positive integer (`minimum: 1`). Use `concurrency: 1` for sequential, file-order writes. In-flight writes may complete after an error; imports are not transactional. See [import](commands.md#import) for error and cancellation behavior. +The CLI's periodic ingestion messages are terminal-only: MCP `import` and `bulk` calls do not emit them or add progress fields to their final results. See [ingestion sizing and limitations](ingestion.md) before planning large jobs. + ### Bulk Operations The `bulk` tool exposes the stateless `run`, `patch`, and `delete` subcommands, with the same operation schema, selection, journal, and reporting semantics as the [CLI command](commands.md#bulk). Stateful subcommands (`begin`, `add`, `execute`, `cancel`, `status`, and `show`) are restricted to the interactive shell, like stateful batches. Writes always require MCP elicitation, even with `yes: true`; once approved, no second terminal prompt is shown. `dry-run: true` validates or selects without database writes, but selection reads consume RUs. Only dry runs without `save` skip elicitation: saving a plan creates a server-local file and requires confirmation. diff --git a/docs/navigation.md b/docs/navigation.md index 1556caba..e1e71bad 100644 --- a/docs/navigation.md +++ b/docs/navigation.md @@ -335,6 +335,10 @@ or `$LASTEXITCODE`): These values are a public contract. See the [CI/CD guide](ci.md#exit-code-contract) for install steps, auth patterns, and scripted failure handling. +Ctrl+C interrupts commands started with `-c`, `-k`, or redirected stdin. An interrupted startup command returns a nonzero exit code and stops any following startup command; structured partial write results, including settled import counts and observed RUs, remain visible. Canceled partial results are not written as successful redirected output. Successful writes are not rolled back, and canceled in-flight writes may require reconciliation before retrying. Stopping `watch` with Ctrl+C is intentional and remains successful; enclosing pipelines, blocks, and scripts stop without running subsequent work. An interrupted stdin script also skips any separately supplied `-c` or `-k` command, even when `watch` handled the interruption successfully. + +Interruption also applies while stdin is still being read and while a confirmation prompt is waiting. At a confirmation, `n` or Escape declines normally; Ctrl+C interrupts the operation. Stopping a `watch` expression stops its enclosing assignment or command before it changes a variable or performs subsequent work. Interrupt markers are not written to machine-readable or redirected stdout. + Scripted `bulk` writes require `--yes`; use `--dry-run` to preview first. Bulk operations are not transactional, and an exhausted observed RU budget returns exit code `6`. Supply `--journal` on the original run and reuse that same journal to resume; without it, reconcile completed writes before retrying. See [bulk operations](commands.md#bulk) for recovery and partial-success semantics. ### Environment Variables diff --git a/l10n/CosmosDBShell.json b/l10n/CosmosDBShell.json index 901f2d3e..7a7cde70 100644 --- a/l10n/CosmosDBShell.json +++ b/l10n/CosmosDBShell.json @@ -164,7 +164,8 @@ "command-bulk-example-6": "Queue an operation onto the active bulk job", "command-bulk-example-7": "Execute the queued operations", "command-bulk-mcp-help": "Available through MCP for the stateless run, patch, and delete subcommands. Writes and dry runs that save a plan require user confirmation; dry runs without --save skip confirmation and send no writes.", - "command-bulk-progress": "{0} of {1} operations processed, {2} failed (RU charge: {3})", + "command-bulk-progress": "Bulk: {0} of {1} operations processed, {2} succeeded, {3} failed; elapsed {4} s, {5} operations/s (observed RU charge: {6})", + "command-bulk-progress-compact": "Bulk: {0}/{1}, {2} failed, {3} s, {4}/s, {5} RU", "command-bulk-status-failed": "Failed", "command-bulk-status-inactive": "No bulk job is currently active.", "command-bulk-status-more": "...and {0} more", @@ -528,6 +529,8 @@ "command-import-example-6": "Keep importing after individual item write failures", "command-import-example-7": "Validate the file without writing any items", "command-import-example-8": "Import with up to 32 concurrent writes", + "command-import-progress": "Import: {0} items processed, {1} imported, {2} failed; elapsed {3} s, {4} items/s (observed RU charge: {5})", + "command-import-progress-compact": "Import: {0} ok, {1} failed, {2} s, {3}/s, {4} RU", "command-import-stopped": "Import stopped after {0} items (RU charge: {1}).", "command-import-success": "Imported {0} {1} (RU charge: {2})", "command-import-success-partial": "Imported {0} {1}, {2} failed (RU charge: {3})", @@ -535,6 +538,8 @@ "command-import-success-partial.__p1.other": "items", "command-import-success.__p1.one": "item", "command-import-success.__p1.other": "items", + "command-import-validation-progress": "Dry run: {0} items validated offline; elapsed {1} s, {2} items/s; no writes sent", + "command-import-validation-progress-compact": "Dry run: {0} validated, {1} s, {2}/s", "command-index-description": "Manages the indexing policy of a container via show, add, remove, and set subcommands.", "command-index-description-automatic": "Whether automatic indexing is enabled (true or false).", "command-index-description-container": "The container to read/update the indexing policy for", @@ -813,6 +818,8 @@ "command-rm-dry-run-plan": "Dry run: Would delete {0} {1}. No changes were made.", "command-rm-dry-run-plan.__p1.one": "item", "command-rm-dry-run-plan.__p1.other": "items", + "command-rm-dry-run-progress": "Dry run: {0} items scanned, {1} would be deleted; elapsed {2} s (observed RU charge: {3})", + "command-rm-dry-run-progress-compact": "Dry run: {0} to delete, {1} scanned, {2} s, {3} RU", "command-rm-error-etag_empty": "The --etag option requires a non-empty value.", "command-rm-error-etag_mismatch": "Item \u0027{0}\u0027 was modified since it was last read (ETag mismatch). The item was not deleted.", "command-rm-error-etag_requires_exact_id": "The --etag option can only be used to delete a single item by exact id (--key=id without wildcards or piped input).", @@ -831,6 +838,8 @@ "command-rm-example-6": "Preview deleting one item by id within a single logical partition, including its current ETag", "command-rm-example-7": "Delete one item only if it has not changed since the dry run", "command-rm-no-matches": "No items matched the pattern \u0027{0}\u0027 for key \u0027{1}\u0027", + "command-rm-progress": "Remove: {0} items scanned, {1} deleted; elapsed {2} s (observed RU charge: {3})", + "command-rm-progress-compact": "Remove: {0} deleted, {1} scanned, {2} s, {3} RU", "command-rm-warning-missing-partition-key": "Warning: Cannot delete item with id \u0027{0}\u0027 - missing partition key \u0027{1}\u0027", "command-rmcon-confirm_container_deletion": "Are you sure you want to delete this container", "command-rmcon-deleted_container": "Deleted container {0}",