Address request charge review findings

Author: mkrueger

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ListCommand.cs

@@ -236,7 +236,8 @@ private async Task<CommandState> ListContainerItemsAsync(ConnectedState state, S
             }
             catch (Exception ex) when (ex is not OperationCanceledException)
             {
-                return new ErrorCommandState(ex) { RequestCharge = returnState.RequestCharge };
+                RequestChargeContext.Record(returnState.RequestCharge ?? 0);
+                throw new CommandException("ls", ex);
             }
 
             using (queryDocument)

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/PrintCommand.cs

@@ -54,46 +54,38 @@ private async Task<CommandState> PrintItemAsync(Container container, Cancellatio
         try
         {
             using var response = await container.ReadItemStreamAsync(this.Id, new PartitionKey(this.PartitionKey), cancellationToken: token);
+            RequestChargeContext.Record(response.Headers.RequestCharge);
 
             if (response.IsSuccessStatusCode)
             {
-                commandState.RequestCharge = response.Headers.RequestCharge;
                 using var reader = new StreamReader(response.Content);
                 var content = await reader.ReadToEndAsync();
 
                 // Parse the content as JSON for structured output
-                var jsonDocument = System.Text.Json.JsonDocument.Parse(content);
-                commandState.Result = new ShellJson(jsonDocument.RootElement);
+                using var jsonDocument = System.Text.Json.JsonDocument.Parse(content);
+                commandState.Result = new ShellJson(jsonDocument.RootElement.Clone());
             }
             else if (response.StatusCode == System.Net.HttpStatusCode.NotFound)
             {
-                return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary<string, object>
+                throw new CommandException("print", MessageService.GetString("command-print-error-item_not_found", new Dictionary<string, object>
                 {
                     { "id", this.Id ?? "(null)" },
                     { "partitionKey", this.PartitionKey ?? "(null)" },
-                })))
-                {
-                    RequestCharge = response.Headers.RequestCharge,
-                };
+                }));
             }
             else
             {
-                return new ErrorCommandState(new CommandException("print", MessageService.GetString("command-print-error-request_failed", new Dictionary<string, object>
+                throw new CommandException("print", MessageService.GetString("command-print-error-request_failed", new Dictionary<string, object>
                 {
                     { "id", this.Id ?? "(null)" },
                     { "status", (int)response.StatusCode },
-                })))
-                {
-                    RequestCharge = response.Headers.RequestCharge,
-                };
+                }));
             }
         }
         catch (CosmosException ex)
         {
-            return new ErrorCommandState(new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex))
-            {
-                RequestCharge = ex.RequestCharge > 0 ? ex.RequestCharge : null,
-            };
+            RequestChargeContext.Record(ex.RequestCharge);
+            throw new CommandException("print", MessageService.GetArgsString("command-print-error-reading_item", "message", CommandException.GetDisplayMessage(ex)), ex);
         }
 
         return commandState;

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/QueryCommand.cs

@@ -634,14 +634,8 @@ private async Task<CommandState> ExecuteExplainAsync(Container container, ShellI
             double requestCharge = response?.Headers.RequestCharge ?? 0;
             if (response is not null && !response.IsSuccessStatusCode)
             {
-                try
-                {
-                    await this.ThrowIfRequestFailedAsync(response, shell);
-                }
-                catch (Exception ex) when (ex is not OperationCanceledException)
-                {
-                    return new ErrorCommandState(ex) { RequestCharge = requestCharge > 0 ? requestCharge : null };
-                }
+                RequestChargeContext.Record(requestCharge);
+                await this.ThrowIfRequestFailedAsync(response, shell);
             }
 
             var cumulative = response?.Diagnostics.GetQueryMetrics()?.CumulativeMetrics;
@@ -723,15 +717,8 @@ private async Task<CommandState> ExecuteQueryAsync(Container container, ShellInt
                 var pageRequestCharge = response.Headers.RequestCharge;
                 if (!response.IsSuccessStatusCode)
                 {
-                    try
-                    {
-                        await this.ThrowIfRequestFailedAsync(response, shell);
-                    }
-                    catch (Exception ex) when (ex is not OperationCanceledException)
-                    {
-                        var failedCharge = totalRequestCharge + pageRequestCharge;
-                        return new ErrorCommandState(ex) { RequestCharge = failedCharge > 0 ? failedCharge : null };
-                    }
+                    RequestChargeContext.Record(totalRequestCharge + pageRequestCharge);
+                    await this.ThrowIfRequestFailedAsync(response, shell);
                 }
 
                 if (response.Content == null)

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ReplaceCommand.cs

@@ -130,6 +130,7 @@ private static async Task<ReplaceSummary> ReplaceArrayAsync(Container container,
 
         if (failCount > 0)
         {
+            RequestChargeContext.Record(charge);
             throw new CommandException(
                 "replace",
                 MessageService.GetArgsString(

CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/RmCommand.cs

@@ -249,7 +249,7 @@ private async Task<ExitCode> RemoveItemsFromContainerAsync(ConnectedState state,
                 totalCharge += response.Headers.RequestCharge;
 
                 using var streamReader = new StreamReader(response.Content);
-                var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync());
+                using var queryDocument = JsonDocument.Parse(await streamReader.ReadToEndAsync());
 
                 foreach (var element in queryDocument.RootElement.GetProperty("Documents").EnumerateArray())
                 {

CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs

@@ -2003,7 +2003,7 @@ private void SetSessionMaxRequestCharge(ShellObject value)
         double requestChargeTotal;
         lock (this.sessionRequestChargeLock)
         {
-            if (maximum != this.sessionMaxRequestCharge)
+            if (Math.Abs(maximum - this.sessionMaxRequestCharge) > 1e-9)
             {
                 this.sessionRequestChargeWarningIssued = false;
             }