diff --git a/CosmosDBShell.Tests/CommandTests/FilterCommandTests.cs b/CosmosDBShell.Tests/CommandTests/FilterCommandTests.cs index 9d96e98..8195994 100644 --- a/CosmosDBShell.Tests/CommandTests/FilterCommandTests.cs +++ b/CosmosDBShell.Tests/CommandTests/FilterCommandTests.cs @@ -442,7 +442,7 @@ public async Task ExecuteAsync_HonorsCancellation() }; using var cts = new CancellationTokenSource(); - cts.Cancel(); + await cts.CancelAsync(); await Assert.ThrowsAnyAsync( () => command.ExecuteAsync(shell, state, string.Empty, cts.Token)); diff --git a/CosmosDBShell.Tests/CommandTests/ThroughputCommandTests.cs b/CosmosDBShell.Tests/CommandTests/ThroughputCommandTests.cs new file mode 100644 index 0000000..9664037 --- /dev/null +++ b/CosmosDBShell.Tests/CommandTests/ThroughputCommandTests.cs @@ -0,0 +1,161 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace CosmosShell.Tests.CommandTests; + +using Azure.Data.Cosmos.Shell.Commands; +using Azure.Data.Cosmos.Shell.Core; +using Azure.Data.Cosmos.Shell.States; +using Azure.Data.Cosmos.Shell.Util; +using Microsoft.Azure.Cosmos; + +/// +/// Offline unit tests for . These cover the not-connected, +/// wrong-scope, and argument-validation branches that execute before any network call. +/// +public class ThroughputCommandTests +{ + [Fact] + public async Task Disconnected_ThrowsNotConnected() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DisconnectedState(); + var command = new ThroughputCommand { Subcommand = "show" }; + + await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput show", CancellationToken.None)); + } + + [Fact] + public async Task Connected_NoDatabase_ThrowsNotInDatabase() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new ConnectedState(CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "show" }; + + await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput show", CancellationToken.None)); + } + + [Fact] + public async Task InvalidSubcommand_ThrowsCommandException() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "bogus" }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput bogus", CancellationToken.None)); + Assert.Equal( + MessageService.GetArgsString("command-throughput-error-invalid_subcommand", "subcommand", "bogus"), + ex.Message); + } + + [Theory] + [InlineData("set")] + [InlineData("manual")] + [InlineData("autoscale")] + public async Task Write_MissingRu_ThrowsCommandException(string subcommand) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = subcommand }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), $"throughput {subcommand}", CancellationToken.None)); + Assert.Equal(MessageService.GetString("command-throughput-error-missing_ru"), ex.Message); + } + + [Theory] + [InlineData(0)] + [InlineData(-100)] + public async Task Write_NonPositiveRu_ThrowsCommandException(int ru) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "set", Ru = ru }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), $"throughput set {ru}", CancellationToken.None)); + Assert.Equal( + MessageService.GetArgsString("command-throughput-error-invalid_ru", "ru", ru), + ex.Message); + } + + [Fact] + public async Task Show_WithRu_ThrowsCommandException() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "show", Ru = 4000 }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput show 4000", CancellationToken.None)); + Assert.Equal(MessageService.GetString("command-throughput-error-show_no_args"), ex.Message); + } + + [Theory] + [InlineData("set", 100)] + [InlineData("manual", 399)] + public async Task Manual_BelowMinimum_ThrowsCommandException(string subcommand, int ru) + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = subcommand, Ru = ru }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), $"throughput {subcommand} {ru}", CancellationToken.None)); + Assert.Equal( + MessageService.GetArgsString("command-throughput-error-manual_min", "ru", ru, "min", 400), + ex.Message); + } + + [Fact] + public async Task Manual_NotMultipleOf100_ThrowsCommandException() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "manual", Ru = 450 }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput manual 450", CancellationToken.None)); + Assert.Equal( + MessageService.GetArgsString("command-throughput-error-manual_increment", "ru", 450, "increment", 100), + ex.Message); + } + + [Fact] + public async Task Autoscale_BelowMinimum_ThrowsCommandException() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "autoscale", Ru = 400 }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput autoscale 400", CancellationToken.None)); + Assert.Equal( + MessageService.GetArgsString("command-throughput-error-autoscale_min", "ru", 400, "min", 1000), + ex.Message); + } + + [Fact] + public async Task Autoscale_NotMultipleOf1000_ThrowsCommandException() + { + using var shell = ShellInterpreter.CreateInstance(); + shell.State = new DatabaseState("TestDatabase", CreateTestClient()); + var command = new ThroughputCommand { Subcommand = "autoscale", Ru = 1500 }; + + var ex = await Assert.ThrowsAsync( + () => command.ExecuteAsync(shell, new CommandState(), "throughput autoscale 1500", CancellationToken.None)); + Assert.Equal( + MessageService.GetArgsString("command-throughput-error-autoscale_increment", "ru", 1500, "increment", 1000), + ex.Message); + } + + private static CosmosClient CreateTestClient() + { + var connectionString = ParsedDocDBConnectionString.BuildEmulatorConnectionString("https://localhost:8081/"); + return new CosmosClient(connectionString); + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs new file mode 100644 index 0000000..dce8e66 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Commands/ThroughputCommand.cs @@ -0,0 +1,306 @@ +//------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +//------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Commands; + +using System.Globalization; +using System.Text.Json; +using System.Text.Json.Nodes; +using System.Text.RegularExpressions; +using Azure.Data.Cosmos.Shell.Mcp; +using Azure.Data.Cosmos.Shell.Parser; +using Azure.Data.Cosmos.Shell.Util; +using global::Azure.Data.Cosmos.Shell.Core; +using global::Azure.Data.Cosmos.Shell.States; +using Spectre.Console; + +[CosmosCommand("throughput")] +[CosmosExample("throughput show", Description = "Display the current provisioned throughput (RU/s)")] +[CosmosExample("throughput set 4000", Description = "Set manual throughput to 4000 RU/s")] +[CosmosExample("throughput manual 4000", Description = "Switch to manual provisioning at 4000 RU/s")] +[CosmosExample("throughput autoscale 10000", Description = "Switch to autoscale with a maximum of 10000 RU/s")] +[CosmosExample("throughput set 4000 --yes", Description = "Set throughput without the confirmation prompt")] +#pragma warning disable SA1118 // Parameter should not span multiple lines +[McpAnnotation( + Title = "Throughput", + Description = @" +Views or changes the provisioned throughput (RU/s) of a Cosmos DB database or container through subcommands: +- 'show' returns the current throughput as JSON, including the mode (manual or autoscale), provisioned RU/s, autoscale maximum, and minimum. +- 'set ' sets manual throughput to the given RU/s (alias of 'manual'). +- 'manual ' switches to manual provisioning at the given RU/s. +- 'autoscale ' switches to autoscale with the given maximum RU/s. + +By default the command targets the current scope: the container when in a container, otherwise the database. Use --database and --container to target a specific resource.", + ReadOnly = false)] +#pragma warning restore SA1118 // Parameter should not span multiple lines +internal class ThroughputCommand : CosmosCommand, IStateVisitor +{ + // Matches the Cosmos DB RBAC "action denied" message returned for both data-plane + // (CosmosException 403/5302) and control-plane throughput writes. The principal id + // can be padded with whitespace inside the brackets. + private static readonly Regex RbacPermissionRegex = new( + @"principal \[\s*([^\]]+?)\s*\] does not have required RBAC permissions to perform action \[([^\]]+)\]", + RegexOptions.IgnoreCase | RegexOptions.Compiled); + + [CosmosParameter("subcommand", RequiredErrorKey = "command-throughput-error-missing_subcommand")] + public string Subcommand { get; init; } = string.Empty; + + [CosmosParameter("ru", IsRequired = false)] + public int? Ru { get; init; } + + [CosmosOption("database", "db")] + public string? Database { get; init; } + + [CosmosOption("container", "con")] + public string? Container { get; init; } + + [CosmosOption("yes", "y", "force")] + public bool? Yes { get; init; } + + public override Task ExecuteAsync(ShellInterpreter shell, CommandState commandState, string commandText, CancellationToken token) => + shell.State.AcceptAsync(this, shell, token); + + Task IStateVisitor.VisitDisconnectedStateAsync(DisconnectedState state, ShellInterpreter shell, CancellationToken token) + { + throw new NotConnectedException("throughput"); + } + + async Task IStateVisitor.VisitConnectedStateAsync(ConnectedState state, ShellInterpreter shell, CancellationToken token) + { + if (string.IsNullOrEmpty(this.Database)) + { + throw new NotInDatabaseException("throughput"); + } + + return await this.ExecuteOnScopeAsync(state, shell, this.Database, this.Container, token); + } + + async Task IStateVisitor.VisitDatabaseStateAsync(DatabaseState state, ShellInterpreter shell, CancellationToken token) + { + string databaseName = this.Database ?? state.DatabaseName; + return await this.ExecuteOnScopeAsync(state, shell, databaseName, this.Container, token); + } + + async Task IStateVisitor.VisitContainerStateAsync(ContainerState state, ShellInterpreter shell, CancellationToken token) + { + string databaseName = this.Database ?? state.DatabaseName; + string containerName = this.Container ?? state.ContainerName; + return await this.ExecuteOnScopeAsync(state, shell, databaseName, containerName, token); + } + + private static CommandState BuildResult(ShellInterpreter shell, ThroughputView view) + { + string mode = view.Availability == ThroughputAvailability.NotConfigured + ? "none" + : view.IsAutoscale ? "autoscale" : "manual"; + + var root = new JsonObject + { + ["scope"] = view.Scope, + ["resource"] = view.ResourceName, + ["mode"] = mode, + ["throughput"] = view.Throughput, + ["autoscaleMaxThroughput"] = view.AutoscaleMaxThroughput, + ["minThroughput"] = view.MinThroughput, + }; + + using var jsonDoc = JsonDocument.Parse(root.ToJsonString(new JsonSerializerOptions { WriteIndented = true })); + var result = new ShellJson(jsonDoc.RootElement.Clone()); + + // When output is redirected to a file, let the interpreter emit the JSON result + // so 'throughput show > out.json' honors the documented JSON contract instead of + // writing a console table. Interactive sessions still get the friendly table, and + // MCP/piping consume the structured Result regardless. + if (!string.IsNullOrEmpty(shell.StdOutRedirect)) + { + return new CommandState { Result = result }; + } + + var table = new Table().HideHeaders(); + table.AddColumn(string.Empty); + table.AddColumn(string.Empty); + void Row(string labelKey, string value) => + table.AddRow( + Theme.FormatHelpName(Markup.Escape(MessageService.GetString(labelKey))), + Theme.FormatTableValue(Markup.Escape(value))); + + Row("command-throughput-label-scope", MessageService.GetString($"command-throughput-scope-{view.Scope}")); + Row("command-throughput-label-resource", view.ResourceName); + Row("command-throughput-label-mode", MessageService.GetString($"command-throughput-mode-{mode}")); + if (view.Throughput.HasValue) + { + Row("command-throughput-label-throughput", view.Throughput.Value.ToString(CultureInfo.InvariantCulture)); + } + + if (view.AutoscaleMaxThroughput.HasValue) + { + Row("command-throughput-label-max", view.AutoscaleMaxThroughput.Value.ToString(CultureInfo.InvariantCulture)); + } + + if (view.MinThroughput.HasValue) + { + Row("command-throughput-label-min", view.MinThroughput.Value.ToString(CultureInfo.InvariantCulture)); + } + + AnsiConsole.Write(table); + + return new CommandState + { + Result = result, + IsPrinted = true, + }; + } + + private static bool TryGetRbacError(Exception e, out string principalId, out string action) + { + var match = RbacPermissionRegex.Match(e.Message ?? string.Empty); + if (match.Success) + { + principalId = match.Groups[1].Value.Trim(); + action = match.Groups[2].Value.Trim(); + return true; + } + + principalId = string.Empty; + action = string.Empty; + return false; + } + + // Returns true when the throughput write should proceed. Interactive sessions get a + // billing-impact confirmation prompt; --yes/--force, MCP, script, and piped-input + // contexts skip it so automation never blocks on Console input. + private static bool ConfirmWrite(ShellInterpreter shell, bool yes, string databaseName, string? containerName, bool isAutoscale, int ru) + { + if (yes || shell.McpPort.HasValue || !string.IsNullOrEmpty(shell.CurrentScriptFileName) || Console.IsInputRedirected) + { + return true; + } + + string resourceName = containerName ?? databaseName; + 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"); + } + + private async Task ExecuteOnScopeAsync(ConnectedState state, ShellInterpreter shell, string databaseName, string? containerName, CancellationToken token) + { + bool isWrite; + bool isAutoscale; + switch (this.Subcommand.Trim().ToLowerInvariant()) + { + case "show": + isWrite = false; + isAutoscale = false; + break; + case "set": + case "manual": + isWrite = true; + isAutoscale = false; + break; + case "autoscale": + isWrite = true; + isAutoscale = true; + break; + default: + throw new CommandException( + "throughput", + MessageService.GetArgsString("command-throughput-error-invalid_subcommand", "subcommand", this.Subcommand)); + } + + int ru = 0; + if (isWrite) + { + ru = this.RequireRu(isAutoscale); + } + else if (this.Ru.HasValue) + { + throw new CommandException("throughput", MessageService.GetString("command-throughput-error-show_no_args")); + } + + await ValidateContainerExistsAsync(state, databaseName, containerName, "throughput", token); + + if (!isWrite) + { + var current = await CosmosResourceFacade.GetThroughputAsync(state, databaseName, containerName, token); + return BuildResult(shell, current); + } + + if (!ConfirmWrite(shell, this.Yes == true, databaseName, containerName, isAutoscale, ru)) + { + ShellInterpreter.WriteLine(MessageService.GetString("command-throughput-cancelled")); + return new CommandState { IsPrinted = true }; + } + + ThroughputView view; + try + { + view = await CosmosResourceFacade.ReplaceThroughputAsync(state, databaseName, containerName, new ThroughputUpdate(isAutoscale, ru), token); + } + catch (ThroughputNotConfiguredException ex) + { + throw new CommandException( + "throughput", + MessageService.GetArgsString("command-throughput-error-not_configured", "resource", ex.ResourceName), + ex); + } + catch (ThroughputModeSwitchNotSupportedException ex) + { + string targetMode = MessageService.GetString(ex.TargetIsAutoscale ? "command-throughput-mode-autoscale" : "command-throughput-mode-manual"); + throw new CommandException( + "throughput", + MessageService.GetArgsString("command-throughput-error-mode_switch_unsupported", "resource", ex.ResourceName, "mode", targetMode), + ex); + } + catch (Exception ex) when (TryGetRbacError(ex, out var principalId, out var action)) + { + throw new CommandException( + "throughput", + MessageService.GetArgsString("command-throughput-error-rbac", "id", principalId, "permission", action), + ex); + } + + ShellInterpreter.WriteLine(MessageService.GetString("command-throughput-updated")); + return BuildResult(shell, view); + } + + private int RequireRu(bool isAutoscale) + { + if (!this.Ru.HasValue) + { + throw new CommandException("throughput", MessageService.GetString("command-throughput-error-missing_ru")); + } + + int value = this.Ru.Value; + if (value <= 0) + { + throw new CommandException( + "throughput", + MessageService.GetArgsString("command-throughput-error-invalid_ru", "ru", value)); + } + + // Cosmos DB requires manual RU/s in multiples of 100 (minimum 400) and autoscale + // maximum RU/s in multiples of 1000 (minimum 1000). Validate up front so the user + // gets a clean message instead of a raw server rejection. + int minimum = isAutoscale ? 1000 : 400; + int increment = isAutoscale ? 1000 : 100; + if (value < minimum) + { + string key = isAutoscale ? "command-throughput-error-autoscale_min" : "command-throughput-error-manual_min"; + throw new CommandException( + "throughput", + MessageService.GetArgsString(key, "ru", value, "min", minimum)); + } + + if (value % increment != 0) + { + string key = isAutoscale ? "command-throughput-error-autoscale_increment" : "command-throughput-error-manual_increment"; + throw new CommandException( + "throughput", + MessageService.GetArgsString(key, "ru", value, "increment", increment)); + } + + return value; + } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ArmCosmosResourceOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ArmCosmosResourceOperations.cs index fe5db11..5fb80c6 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ArmCosmosResourceOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ArmCosmosResourceOperations.cs @@ -169,4 +169,110 @@ public async Task ReplaceIndexingPolicyAsync(string databaseName, string var updated = response.Value.Data.Resource.IndexingPolicy; return CosmosResourceJson.IndentJson(CosmosArmResourceProvider.WriteArmModel(updated)); } + + public async Task GetThroughputAsync(string databaseName, string? containerName, CancellationToken token) + { + string scope = containerName is null ? "database" : "container"; + string resourceName = containerName ?? databaseName; + try + { + ThroughputSettingsResourceInfo info; + if (string.IsNullOrEmpty(containerName)) + { + var database = await CosmosArmResourceProvider.GetDatabaseAsync(context, databaseName, token); + var response = await database.GetCosmosDBSqlDatabaseThroughputSetting().GetAsync(token); + info = response.Value.Data.Resource; + } + else + { + var container = await CosmosArmResourceProvider.GetContainerAsync(context, databaseName, containerName, token); + var response = await container.GetCosmosDBSqlContainerThroughputSetting().GetAsync(token); + info = response.Value.Data.Resource; + } + + return BuildThroughputView(scope, resourceName, info); + } + catch (RequestFailedException ex) when (ex.Status == (int)HttpStatusCode.NotFound) + { + return new ThroughputView(scope, resourceName, false, null, null, null, ThroughputAvailability.NotConfigured, null); + } + } + + public async Task ReplaceThroughputAsync(string databaseName, string? containerName, ThroughputUpdate update, CancellationToken token) + { + string scope = containerName is null ? "database" : "container"; + string resourceName = containerName ?? databaseName; + var resourceInfo = update.IsAutoscale + ? new ThroughputSettingsResourceInfo { AutoscaleSettings = new AutoscaleSettingsResourceInfo(update.Throughput) } + : new ThroughputSettingsResourceInfo { Throughput = update.Throughput }; + var data = new ThroughputSettingsUpdateData(context.Account.Data.Location, resourceInfo); + try + { + ThroughputSettingsResourceInfo info; + if (string.IsNullOrEmpty(containerName)) + { + var database = await CosmosArmResourceProvider.GetDatabaseAsync(context, databaseName, token); + var setting = database.GetCosmosDBSqlDatabaseThroughputSetting(); + var current = await setting.GetAsync(token); + if (IsAutoscale(current.Value.Data.Resource) != update.IsAutoscale) + { + if (update.IsAutoscale) + { + await setting.MigrateSqlDatabaseToAutoscaleAsync(WaitUntil.Completed, token); + } + else + { + await setting.MigrateSqlDatabaseToManualThroughputAsync(WaitUntil.Completed, token); + } + } + + var operation = await setting.CreateOrUpdateAsync(WaitUntil.Completed, data, token); + info = operation.Value.Data.Resource; + } + else + { + var container = await CosmosArmResourceProvider.GetContainerAsync(context, databaseName, containerName, token); + var setting = container.GetCosmosDBSqlContainerThroughputSetting(); + var current = await setting.GetAsync(token); + if (IsAutoscale(current.Value.Data.Resource) != update.IsAutoscale) + { + if (update.IsAutoscale) + { + await setting.MigrateSqlContainerToAutoscaleAsync(WaitUntil.Completed, token); + } + else + { + await setting.MigrateSqlContainerToManualThroughputAsync(WaitUntil.Completed, token); + } + } + + var operation = await setting.CreateOrUpdateAsync(WaitUntil.Completed, data, token); + info = operation.Value.Data.Resource; + } + + return BuildThroughputView(scope, resourceName, info); + } + catch (RequestFailedException ex) when (ex.Status is (int)HttpStatusCode.NotFound or (int)HttpStatusCode.BadRequest) + { + throw new ThroughputNotConfiguredException(resourceName, ex); + } + } + + private static bool IsAutoscale(ThroughputSettingsResourceInfo info) => info.AutoscaleSettings?.MaxThroughput != null; + + private static ThroughputView BuildThroughputView(string scope, string resourceName, ThroughputSettingsResourceInfo info) + { + int? min = int.TryParse(info.MinimumThroughput, out var parsedMin) ? parsedMin : null; + int? autoscaleMax = info.AutoscaleSettings?.MaxThroughput; + bool isAutoscale = autoscaleMax.HasValue; + return new ThroughputView( + scope, + resourceName, + isAutoscale, + info.Throughput, + autoscaleMax, + min, + ThroughputAvailability.Available, + null); + } } \ No newline at end of file diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosResourceFacade.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosResourceFacade.cs index 31eb24b..63820de 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosResourceFacade.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/CosmosResourceFacade.cs @@ -82,6 +82,16 @@ public static Task ReplaceIndexingPolicyAsync(ConnectedState state, stri return For(state).ReplaceIndexingPolicyAsync(databaseName, containerName, indexPolicyJson, token); } + public static Task GetThroughputAsync(ConnectedState state, string databaseName, string? containerName, CancellationToken token) + { + return For(state).GetThroughputAsync(databaseName, containerName, token); + } + + public static Task ReplaceThroughputAsync(ConnectedState state, string databaseName, string? containerName, ThroughputUpdate update, CancellationToken token) + { + return For(state).ReplaceThroughputAsync(databaseName, containerName, update, token); + } + private static ICosmosResourceOperations For(ConnectedState state) { return state.ArmContext is { } armContext diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs index 6b17843..1859f3a 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/DataPlaneCosmosResourceOperations.cs @@ -207,6 +207,81 @@ public async Task ReplaceIndexingPolicyAsync(string databaseName, string return JsonConvert.SerializeObject(replaced.Resource?.IndexingPolicy ?? policy, Formatting.Indented); } + public async Task GetThroughputAsync(string databaseName, string? containerName, CancellationToken token) + { + string scope = containerName is null ? "database" : "container"; + string resourceName = containerName ?? databaseName; + try + { + var throughputResponse = string.IsNullOrEmpty(containerName) + ? await client.GetDatabase(databaseName).ReadThroughputAsync(new RequestOptions(), token) + : await client.GetDatabase(databaseName).GetContainer(containerName).ReadThroughputAsync(new RequestOptions(), token); + return BuildThroughputView(scope, resourceName, throughputResponse); + } + catch (CosmosException ex) when (ex.StatusCode is HttpStatusCode.NotFound or HttpStatusCode.BadRequest) + { + return new ThroughputView(scope, resourceName, false, null, null, null, ThroughputAvailability.NotConfigured, null); + } + } + + public async Task ReplaceThroughputAsync(string databaseName, string? containerName, ThroughputUpdate update, CancellationToken token) + { + string scope = containerName is null ? "database" : "container"; + string resourceName = containerName ?? databaseName; + + ThroughputResponse currentResponse; + try + { + currentResponse = string.IsNullOrEmpty(containerName) + ? await client.GetDatabase(databaseName).ReadThroughputAsync(new RequestOptions(), token) + : await client.GetDatabase(databaseName).GetContainer(containerName).ReadThroughputAsync(new RequestOptions(), token); + } + catch (CosmosException ex) when (ex.StatusCode is HttpStatusCode.NotFound or HttpStatusCode.BadRequest) + { + throw new ThroughputNotConfiguredException(resourceName, ex); + } + + // The data-plane SDK can only change the value within the current mode; it + // cannot migrate between manual and autoscale. Detect that up front and fail + // with a clear error instead of silently leaving the mode unchanged. + bool currentIsAutoscale = currentResponse.Resource?.AutoscaleMaxThroughput is not null; + if (currentIsAutoscale != update.IsAutoscale) + { + throw new ThroughputModeSwitchNotSupportedException(resourceName, update.IsAutoscale); + } + + var properties = update.IsAutoscale + ? ThroughputProperties.CreateAutoscaleThroughput(update.Throughput) + : ThroughputProperties.CreateManualThroughput(update.Throughput); + try + { + var throughputResponse = string.IsNullOrEmpty(containerName) + ? await client.GetDatabase(databaseName).ReplaceThroughputAsync(properties, cancellationToken: token) + : await client.GetDatabase(databaseName).GetContainer(containerName).ReplaceThroughputAsync(properties, cancellationToken: token); + return BuildThroughputView(scope, resourceName, throughputResponse); + } + catch (CosmosException ex) when (ex.StatusCode is HttpStatusCode.NotFound or HttpStatusCode.BadRequest) + { + throw new ThroughputNotConfiguredException(resourceName, ex); + } + } + + private static ThroughputView BuildThroughputView(string scope, string resourceName, ThroughputResponse response) + { + var resource = response.Resource; + int? autoscaleMax = resource?.AutoscaleMaxThroughput; + bool isAutoscale = autoscaleMax.HasValue; + return new ThroughputView( + scope, + resourceName, + isAutoscale, + resource?.Throughput, + autoscaleMax, + response.MinThroughput, + ThroughputAvailability.Available, + null); + } + private static ContainerProperties GetContainerPropertiesOrThrow(ContainerResponse response) { return response.Resource ?? throw new ShellException(MessageService.GetString("error-unable_to_read_container")); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ICosmosResourceOperations.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ICosmosResourceOperations.cs index 94a7451..d0c0b0d 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ICosmosResourceOperations.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ICosmosResourceOperations.cs @@ -37,4 +37,8 @@ Task CreateContainerAsync( Task GetIndexingPolicyJsonAsync(string databaseName, string containerName, CancellationToken token); Task ReplaceIndexingPolicyAsync(string databaseName, string containerName, string indexPolicyJson, CancellationToken token); + + Task GetThroughputAsync(string databaseName, string? containerName, CancellationToken token); + + Task ReplaceThroughputAsync(string databaseName, string? containerName, ThroughputUpdate update, CancellationToken token); } \ No newline at end of file diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs index ff68730..3d24509 100644 --- a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ShellInterpreter.cs @@ -230,23 +230,85 @@ public static bool Confirm(string message) var yes = char.ToUpper(MessageService.GetString("yes_char")[0]); var no = char.ToUpper(MessageService.GetString("no_char")[0]); - while (true) + // 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 + // handler, which would leave this blocking ReadKey loop spinning forever. + var restoreControlC = TrySetTreatControlCAsInput(true, out var originalTreatControlCAsInput); + try { - Console.Write($"{MessageService.GetString(message)} ({yes}/{no})?"); - var key = Console.ReadKey(); - WriteLine(); - if (char.ToUpper(key.KeyChar) == yes) + while (true) { - return true; - } + Console.Write($"{MessageService.GetString(message)} ({yes}/{no})?"); + + ConsoleKeyInfo key; + try + { + key = Console.ReadKey(intercept: true); + } + catch (InvalidOperationException) + { + // No interactive console available (e.g. redirected input). Treat as + // a declined prompt rather than throwing. + WriteLine(); + return false; + } + + if (key.Key == ConsoleKey.C && key.Modifiers.HasFlag(ConsoleModifiers.Control)) + { + WriteLine("^C"); + return false; + } + + if (key.Key == ConsoleKey.Escape) + { + WriteLine(); + return false; + } + + // intercept:true suppresses the echo, so mirror the keystroke ourselves. + Console.Write(key.KeyChar); + WriteLine(); + + if (char.ToUpper(key.KeyChar) == yes) + { + return true; + } - if (char.ToUpper(key.KeyChar) == no || key.Key == ConsoleKey.Escape) + if (char.ToUpper(key.KeyChar) == no) + { + return false; + } + } + } + finally + { + if (restoreControlC) { - return false; + TrySetTreatControlCAsInput(originalTreatControlCAsInput, out _); } } } + private static bool TrySetTreatControlCAsInput(bool value, out bool originalValue) + { + try + { + originalValue = Console.TreatControlCAsInput; + Console.TreatControlCAsInput = value; + return true; + } + catch (IOException) + { + originalValue = false; + return false; + } + catch (InvalidOperationException) + { + originalValue = false; + return false; + } + } + /// /// Cancels the current prompt operation, including any ongoing editor or command input. /// diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputModeSwitchNotSupportedException.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputModeSwitchNotSupportedException.cs new file mode 100644 index 0000000..6a2ccfe --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputModeSwitchNotSupportedException.cs @@ -0,0 +1,25 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ +namespace Azure.Data.Cosmos.Shell.Core; + +/// +/// Thrown when a throughput write would switch a resource between manual and +/// autoscale, but the active connection cannot perform that migration. The Cosmos +/// data-plane SDK can only change the value within the current mode; switching +/// modes requires a control-plane (ARM/AAD) connection, the Azure portal, CLI, or +/// PowerShell. +/// +internal sealed class ThroughputModeSwitchNotSupportedException : System.InvalidOperationException +{ + public ThroughputModeSwitchNotSupportedException(string resourceName, bool targetIsAutoscale, System.Exception? innerException = null) + : base($"Switching throughput mode for '{resourceName}' is not supported on this connection.", innerException) + { + this.ResourceName = resourceName; + this.TargetIsAutoscale = targetIsAutoscale; + } + + public string ResourceName { get; } + + public bool TargetIsAutoscale { get; } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputNotConfiguredException.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputNotConfiguredException.cs new file mode 100644 index 0000000..5eb9be6 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputNotConfiguredException.cs @@ -0,0 +1,21 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ +namespace Azure.Data.Cosmos.Shell.Core; + +/// +/// Thrown when a throughput operation targets a resource that has no provisioned +/// throughput (for example a container inside a shared-throughput database, or a +/// serverless account). New callers can match this specific type to translate to a +/// localized command error. +/// +internal sealed class ThroughputNotConfiguredException : System.InvalidOperationException +{ + public ThroughputNotConfiguredException(string resourceName, System.Exception? innerException = null) + : base($"Resource '{resourceName}' has no provisioned throughput.", innerException) + { + this.ResourceName = resourceName; + } + + public string ResourceName { get; } +} diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputUpdate.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputUpdate.cs new file mode 100644 index 0000000..96e4f9c --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputUpdate.cs @@ -0,0 +1,7 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +internal sealed record ThroughputUpdate(bool IsAutoscale, int Throughput); diff --git a/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputView.cs b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputView.cs new file mode 100644 index 0000000..9a03ab3 --- /dev/null +++ b/CosmosDBShell/Azure.Data.Cosmos.Shell.Core/ThroughputView.cs @@ -0,0 +1,15 @@ +// ------------------------------------------------------------ +// Copyright (c) Microsoft Corporation. All rights reserved. +// ------------------------------------------------------------ + +namespace Azure.Data.Cosmos.Shell.Core; + +internal sealed record ThroughputView( + string Scope, + string ResourceName, + bool IsAutoscale, + int? Throughput, + int? AutoscaleMaxThroughput, + int? MinThroughput, + ThroughputAvailability Availability, + string? ErrorMessage); diff --git a/CosmosDBShell/lang/en.ftl b/CosmosDBShell/lang/en.ftl index 640734f..7e1df5a 100644 --- a/CosmosDBShell/lang/en.ftl +++ b/CosmosDBShell/lang/en.ftl @@ -336,6 +336,50 @@ command-index-error-show_no_args = 'index show' does not take any arguments. Use command-index-error-invalid_policy = Invalid indexing policy JSON. Please provide a valid Cosmos DB indexing policy. command-index-error-no_policy = The container has no indexing policy configured. +command-throughput-description = Views or changes the provisioned throughput (RU/s) of a database or container via show, set, manual, and autoscale subcommands. +command-throughput-description-subcommand = The action to perform: show, set, manual, or autoscale. +command-throughput-description-ru = The throughput in RU/s to provision (manual RU/s for set/manual, maximum RU/s for autoscale). +command-throughput-description-database = The database to target, or that contains the target container. +command-throughput-description-container = The container to read/update the throughput for. +command-throughput-description-yes = Skip the confirmation prompt before applying a throughput change. +command-throughput-updated = Throughput updated successfully. +command-throughput-confirm_summary = About to set { $mode } throughput to { $ru } RU/s on '{ $resource }'. This may affect your bill. +command-throughput-confirm = Apply this throughput change +command-throughput-cancelled = Throughput change cancelled. +command-throughput-label-scope = Scope +command-throughput-label-resource = Resource +command-throughput-label-mode = Mode +command-throughput-label-throughput = Throughput (RU/s) +command-throughput-label-max = Max throughput (RU/s) +command-throughput-label-min = Min throughput (RU/s) +command-throughput-scope-database = Database +command-throughput-scope-container = Container +command-throughput-mode-autoscale = Autoscale +command-throughput-mode-manual = Manual +command-throughput-mode-none = Not configured +command-throughput-error-missing_subcommand = Missing subcommand. Use one of: show, set, manual, autoscale. +command-throughput-error-invalid_subcommand = Unknown subcommand '{ $subcommand }'. Use one of: show, set, manual, autoscale. +command-throughput-error-missing_ru = No throughput value provided. Specify the RU/s, for example: throughput set 4000. +command-throughput-error-invalid_ru = Invalid throughput value '{ $ru }'. Provide a positive number of RU/s. +command-throughput-error-manual_min = Manual throughput must be at least { $min } RU/s. '{ $ru }' is too low. +command-throughput-error-manual_increment = Manual throughput must be a multiple of { $increment } RU/s. '{ $ru }' is not. +command-throughput-error-autoscale_min = Autoscale maximum throughput must be at least { $min } RU/s. '{ $ru }' is too low. +command-throughput-error-autoscale_increment = Autoscale maximum throughput must be a multiple of { $increment } RU/s. '{ $ru }' is not. +command-throughput-error-show_no_args = 'throughput show' does not take any arguments. Use 'throughput show' to display the current throughput. +command-throughput-error-not_configured = Resource '{ $resource }' has no provisioned throughput to change. It may be serverless or use shared database throughput. +command-throughput-error-rbac = + You do not have permission to change throughput on the selected account. + + Required action: '{ $permission }' + Principal id: '{ $id }' + + Learn more: https://aka.ms/cosmos-native-rbac + +command-throughput-error-mode_switch_unsupported = + Switching '{ $resource }' to { $mode } throughput is not supported on this connection. + + The Cosmos data-plane SDK can only change the value within the current mode. To switch between manual and autoscale, connect with an Azure AD (token) credential, or use the Azure portal, Azure CLI, or PowerShell. + command-ls-description = List resources in the current context. command-ls-description-filter = The filter pattern. command-ls-description-max = Maximum number of items returned when listing container items. Defaults to 100 if omitted. Use 0 or a negative value for no limit. diff --git a/docs/commands.md b/docs/commands.md index 1ea5dca..c700b18 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -571,6 +571,56 @@ index set --mode=consistent --automatic=true index set '{"indexingMode":"consistent","automatic":true,"includedPaths":[{"path":"/*"}],"excludedPaths":[]}' ``` +### throughput + +View or change the provisioned throughput (RU/s) of a database or container through subcommands. + +```text +Usage: throughput subcommand [ru] [-database ] [-container ] + +Arguments: + subcommand show, set, manual, or autoscale + [ru] Throughput in RU/s (manual RU/s for set/manual, maximum RU/s for autoscale) + +Options: + -database, -db + Override database name (Optional) + -container, -con + Override container name (Optional) + -yes, -y, -force + Skip the confirmation prompt before applying a change (Optional) +``` + +By default the command targets the current scope: the container when in a container, otherwise the database. Use `--database` and `--container` to target a specific resource. + +#### Subcommands + +|Subcommand|Behavior| +|-|-| +|`show`|Reads and returns the current throughput as JSON, including the mode (`manual`, `autoscale`, or `none`), provisioned RU/s, autoscale maximum, and minimum.| +|`set `|Sets manual throughput to the given RU/s. Alias of `manual`.| +|`manual `|Switches to manual provisioning at the given RU/s.| +|`autoscale `|Switches to autoscale with the given maximum RU/s.| + +Throughput changes apply to the resource's own provisioned throughput. Containers inside a shared-throughput database, and serverless accounts, have no dedicated throughput to change. + +Throughput values are validated before the request is sent: manual RU/s must be at least 400 and a multiple of 100, and autoscale maximum RU/s must be at least 1000 and a multiple of 1000. + +Switching between `manual` and `autoscale` is a mode migration. Over an Azure AD (token) connection this is performed automatically. Over a key-based (data-plane) connection the SDK cannot migrate modes, so a mode switch is rejected with guidance to use a token connection, the Azure portal, Azure CLI, or PowerShell; changing the RU/s value within the current mode still works. + +Write operations (`set`, `manual`, `autoscale`) ask for confirmation before applying, because throughput changes can affect your bill. Pass `--yes` (`-y`/`--force`) to skip the prompt. The prompt is also skipped automatically in non-interactive contexts (MCP, script execution, or piped input). + +#### Examples + +```bash +throughput show +throughput set 4000 +throughput manual 4000 +throughput autoscale 10000 +throughput set 4000 --yes +throughput show --database MyDatabase --container MyContainer +``` + ## Utilities ### az