Improve status and list commands (#8102)

Azure · Apr 16, 2024 · 8ac3c6b · 8ac3c6b
1 parent de61b0c
commit 8ac3c6b
Show file tree

Hide file tree

Showing 5 changed files with 117 additions and 69 deletions.
diff --git a/tools/secret-management/Azure.Sdk.Tools.SecretManagement.Cli/Commands/ListCommand.cs b/tools/secret-management/Azure.Sdk.Tools.SecretManagement.Cli/Commands/ListCommand.cs
@@ -1,7 +1,6 @@
 using System.CommandLine.Invocation;
-using System.Text.Json;
+using System.Text;
 using Azure.Sdk.Tools.SecretRotation.Configuration;
-using Azure.Sdk.Tools.SecretRotation.Core;
 
 namespace Azure.Sdk.Tools.SecretManagement.Cli.Commands;
 
@@ -16,7 +15,21 @@ protected override Task HandleCommandAsync(ILogger logger, RotationConfiguration
     {
         foreach (PlanConfiguration plan in rotationConfiguration.PlanConfigurations)
         {
-            Console.WriteLine($"name: {plan.Name} - tags: {string.Join(", ", plan.Tags)}");
+            logger.LogInformation(plan.Name);
+
+            if (logger.IsEnabled(LogLevel.Debug))
+            {
+                var builder = new StringBuilder();
+
+                builder.AppendLine($"  Tags: {string.Join(", ", plan.Tags)}");
+                builder.AppendLine($"  Rotation Period: {plan.RotationPeriod}");
+                builder.AppendLine($"  Rotation Threshold: {plan.RotationThreshold}");
+                builder.AppendLine($"  Warning Threshold: {plan.WarningThreshold}");
+                builder.AppendLine($"  Revoke After Period: {plan.RevokeAfterPeriod}");
+                builder.AppendLine($"  Store Types: {string.Join(", ", plan.StoreConfigurations.Select(s => s.Type).Distinct() )}");
+
+                logger.LogDebug(builder.ToString());
+            }
         }
 
         return Task.CompletedTask;

diff --git a/tools/secret-management/Azure.Sdk.Tools.SecretManagement.Cli/Commands/StatusCommand.cs b/tools/secret-management/Azure.Sdk.Tools.SecretManagement.Cli/Commands/StatusCommand.cs
@@ -1,8 +1,5 @@
 using System.CommandLine.Invocation;
-using System.Reflection;
 using System.Text;
-using System.Text.Json;
-using System.Text.Json.Serialization;
 using Azure.Sdk.Tools.SecretRotation.Configuration;
 using Azure.Sdk.Tools.SecretRotation.Core;
 
@@ -18,93 +15,80 @@ protected override async Task HandleCommandAsync(ILogger logger, RotationConfigu
         InvocationContext invocationContext)
     {
         var timeProvider = new TimeProvider();
-        IEnumerable<RotationPlan> plans = rotationConfiguration.GetAllRotationPlans(logger, timeProvider);
+        RotationPlan[] plans = rotationConfiguration.GetAllRotationPlans(logger, timeProvider).ToArray();
 
-        List<(RotationPlan Plan, RotationPlanStatus Status)> statuses = new();
+        logger.LogInformation($"Getting status for {plans.Length} plans");
 
-        foreach (RotationPlan plan in plans)
-        {
-            logger.LogInformation($"Getting status for plan '{plan.Name}'");
-            RotationPlanStatus status = await plan.GetStatusAsync();
-
-            if (logger.IsEnabled(LogLevel.Debug))
-            {
-                var builder = new StringBuilder();
-
-                builder.AppendLine($"  Plan:");
-                builder.AppendLine($"    RotationPeriod: {plan.RotationPeriod}");
-                builder.AppendLine($"    RotationThreshold: {plan.RotationThreshold}");
-                builder.AppendLine($"    RevokeAfterPeriod: {plan.RevokeAfterPeriod}");
-
-                builder.AppendLine($"  Status:");
-                builder.AppendLine($"    ExpirationDate: {status.ExpirationDate}");
-                builder.AppendLine($"    State: {status.State}");
-                builder.AppendLine($"    RequiresRevocation: {status.RequiresRevocation}");
-                builder.AppendLine($"    Exception: {status.Exception?.Message}");
-
-                logger.LogDebug(builder.ToString());
-            }
+        (RotationPlan Plan, RotationPlanStatus Status)[] statuses = await plans
+            .Select(async plan => {
+                logger.LogDebug($"Getting status for plan '{plan.Name}'.");
+                return (plan, await plan.GetStatusAsync());
+            })
+            .LimitConcurrencyAsync(10);
 
 
-            statuses.Add((plan, status));
-        }
-
         var plansBuyState = statuses.GroupBy(x => x.Status.State)
             .ToDictionary(x => x.Key, x => x.ToArray());
 
-        var statusBuilder = new StringBuilder();
 
-        void AppendStatusSection(RotationState state, string header)
+        void LogStatusSection(RotationState state, string header)
         {
-            if (!plansBuyState.TryGetValue(RotationState.Expired, out var matchingPlans))
+            if (!plansBuyState.TryGetValue(state, out var matchingPlans))
             {
                 return;
             }
 
-            statusBuilder.AppendLine();
-            statusBuilder.AppendLine(header);
+            logger.LogInformation($"\n{header}");
+
             foreach ((RotationPlan plan, RotationPlanStatus status) in matchingPlans)
             {
-                foreach (string line in GetPlanStatusLine(plan, status).Split("\n"))
+                var builder = new StringBuilder();
+                var debugBuilder = new StringBuilder();
+
+                builder.Append($"  {plan.Name} - ");
+                DateTimeOffset? expirationDate = status.ExpirationDate;
+                if (expirationDate.HasValue)
+                {
+                    builder.AppendLine($"{expirationDate} ({FormatTimeSpan(expirationDate.Value.Subtract(DateTimeOffset.UtcNow))})");
+                }
+                else
                 {
-                    statusBuilder.Append("  ");
-                    statusBuilder.AppendLine(line);
+                    builder.AppendLine("no expiration date");
                 }
+
+                debugBuilder.AppendLine($"    Plan:");
+                debugBuilder.AppendLine($"      Rotation Period: {plan.RotationPeriod}");
+                debugBuilder.AppendLine($"      Rotation Threshold: {plan.RotationThreshold}");
+                debugBuilder.AppendLine($"      Warning Threshold: {plan.WarningThreshold}");
+                debugBuilder.AppendLine($"      Revoke After Period: {plan.RevokeAfterPeriod}");
+                debugBuilder.AppendLine($"    Status:");
+                debugBuilder.AppendLine($"      Expiration Date: {status.ExpirationDate}");
+                debugBuilder.AppendLine($"      State: {status.State}");
+                debugBuilder.AppendLine($"      Requires Revocation: {status.RequiresRevocation}");
+
+                if (status.Exception != null)
+                {
+                    builder.AppendLine($"    Exception:");
+                    builder.AppendLine($"      {status.Exception.Message}");
+                }
+
+                logger.LogInformation(builder.ToString());
+                logger.LogDebug(debugBuilder.ToString());
             }
         }
 
-        AppendStatusSection(RotationState.Expired, "Expired:");
-        AppendStatusSection(RotationState.Warning, "Expiring:");
-        AppendStatusSection(RotationState.Rotate, "Should Rotate:");
-        AppendStatusSection(RotationState.UpToDate, "Up-to-date:");
-        AppendStatusSection(RotationState.Error, "Error reading plan status:");
-
-        logger.LogInformation(statusBuilder.ToString());
+        LogStatusSection(RotationState.Expired, "Expired:");
+        LogStatusSection(RotationState.Warning, "Expiring:");
+        LogStatusSection(RotationState.Rotate, "Should Rotate:");
+        LogStatusSection(RotationState.UpToDate, "Up-to-date:");
+        LogStatusSection(RotationState.Error, "Error reading plan status:");
 
         if (statuses.Any(x => x.Status.State is RotationState.Expired or RotationState.Warning))
         {
             invocationContext.ExitCode = 1;
         }
     }
 
-    private static string GetPlanStatusLine(RotationPlan plan, RotationPlanStatus status)
-    {
-        if (status.Exception != null)
-        {
-            return $"{plan.Name}:\n  {status.Exception.Message}";
-        }
-
-        DateTimeOffset? expirationDate = status.ExpirationDate;
-
-        DateTimeOffset now = DateTimeOffset.UtcNow;
-
-        string expiration = expirationDate.HasValue
-            ? $"{FormatTimeSpan(expirationDate.Value.Subtract(now))}"
-            : "No expiration date";
-
-        return $"{plan.Name} - {expiration} / ({FormatTimeSpan(plan.RotationPeriod)} @ {FormatTimeSpan(plan.RotationThreshold)})";
-    }
-
     private static string FormatTimeSpan(TimeSpan timeSpan)
     {
         if (timeSpan == TimeSpan.Zero)

diff --git a/tools/secret-management/Azure.Sdk.Tools.SecretManagement.Cli/SimplerConsoleFormatter.cs b/tools/secret-management/Azure.Sdk.Tools.SecretManagement.Cli/SimplerConsoleFormatter.cs
@@ -1,4 +1,4 @@
-using Microsoft.Extensions.Logging.Abstractions;
+using Microsoft.Extensions.Logging.Abstractions;
 using Microsoft.Extensions.Logging.Console;
 using Microsoft.Extensions.Options;
 
@@ -83,7 +83,7 @@ private DateTimeOffset GetCurrentDateTime()
         return logLevel switch
         {
             LogLevel.Trace => "trce: ",
-            LogLevel.Debug => "dbug: ",
+            LogLevel.Debug => null,
             LogLevel.Information => null,
             LogLevel.Warning => "warn: ",
             LogLevel.Error => "fail: ",

diff --git a/tools/secret-management/Azure.Sdk.Tools.SecretRotation.Core/RotationException.cs b/tools/secret-management/Azure.Sdk.Tools.SecretRotation.Core/RotationException.cs
@@ -1,4 +1,4 @@
-namespace Azure.Sdk.Tools.SecretRotation.Core;
+namespace Azure.Sdk.Tools.SecretRotation.Core;
 
 public class RotationException : Exception
 {

diff --git a/tools/secret-management/Azure.Sdk.Tools.SecretRotation.Core/TaskExtensions.cs b/tools/secret-management/Azure.Sdk.Tools.SecretRotation.Core/TaskExtensions.cs
@@ -0,0 +1,51 @@
+namespace Azure.Sdk.Tools.SecretRotation.Core;
+
+public static class TaskExtensions
+{
+    public static async Task<T[]> LimitConcurrencyAsync<T>(this IEnumerable<Task<T>> tasks, int concurrencyLimit = 1, CancellationToken cancellationToken = default)
+    {
+        if (concurrencyLimit == int.MaxValue)
+        {
+            return await Task.WhenAll(tasks);
+        }
+
+        var results = new List<T>();
+
+        if (concurrencyLimit == 1)
+        {
+            foreach (var task in tasks)
+            {
+                results.Add(await task);
+            }
+
+            return results.ToArray();
+        }
+
+        var pending = new List<Task<T>>();
+
+        foreach (var task in tasks)
+        {
+            pending.Add(task);
+
+            if (cancellationToken.IsCancellationRequested)
+            {
+                break;
+            }
+
+            if (pending.Count < concurrencyLimit) continue;
+
+            var completed = await Task.WhenAny(pending);
+            pending.Remove(completed);
+            results.Add(await completed);
+        }
+
+        results.AddRange(await Task.WhenAll(pending));
+
+        return results.ToArray();
+    }
+
+    public static Task<TResult[]> LimitConcurrencyAsync<TSource, TResult>(this IEnumerable<TSource> source, Func<TSource, Task<TResult>> taskFactory, int concurrencyLimit = 1, CancellationToken cancellationToken = default)
+    {
+        return LimitConcurrencyAsync(source.Select(taskFactory), concurrencyLimit, cancellationToken);
+    }
+}