Refactor communication with AWS for KeyValueStore updates

Directory.Packages.props

@@ -13,8 +13,6 @@
     <PackageVersion Include="Amazon.Lambda.S3Events" Version="3.1.0" />
     <PackageVersion Include="Amazon.Lambda.Serialization.SystemTextJson" Version="2.4.4" />
     <PackageVersion Include="Amazon.Lambda.SQSEvents" Version="2.2.0" />
-    <PackageVersion Include="AWSSDK.CloudFront" Version="4.0.0.10" />
-    <PackageVersion Include="AWSSDK.CloudFrontKeyValueStore" Version="4.0.0.9" />
     <PackageVersion Include="AWSSDK.Core" Version="4.0.0.2" />
     <PackageVersion Include="AWSSDK.SQS" Version="4.0.0.1" />
     <PackageVersion Include="AWSSDK.S3" Version="4.0.0.1" />

src/tooling/Elastic.Documentation.Tooling/ExternalCommands/ExternalCommandExecutor.cs

@@ -12,6 +12,7 @@ namespace Elastic.Documentation.Tooling.ExternalCommands;
 public abstract class ExternalCommandExecutor(DiagnosticsCollector collector, IDirectoryInfo workingDirectory)
 {
 	protected IDirectoryInfo WorkingDirectory => workingDirectory;
+	protected DiagnosticsCollector Collector => collector;
 	protected void ExecIn(Dictionary<string, string> environmentVars, string binary, params string[] args)
 	{
 		var arguments = new ExecArguments(binary, args)
@@ -77,13 +78,13 @@ string[] CaptureOutput()
 	}
 
 
-	protected string Capture(string binary, params string[] args) => Capture(false, binary, args);
-
-	protected string Capture(bool muteExceptions, string binary, params string[] args)
+	protected string Capture(string binary, params string[] args) => Capture(false, 10, binary, args);
+	protected string Capture(bool muteExceptions, string binary, params string[] args) => Capture(muteExceptions, 10, binary, args);
+	protected string Capture(bool muteExceptions, int attempts, string binary, params string[] args)
 	{
 		// Try 10 times to capture the output of the command, if it fails, we'll throw an exception on the last try
 		Exception? e = null;
-		for (var i = 0; i <= 9; i++)
+		for (var i = 1; i <= attempts; i++)
 		{
 			try
 			{

src/tooling/docs-assembler/Cli/DeployCommands.cs

@@ -6,9 +6,6 @@
 using System.IO.Abstractions;
 using System.Text.Json;
 using Actions.Core.Services;
-using Amazon.CloudFront;
-using Amazon.CloudFrontKeyValueStore;
-using Amazon.CloudFrontKeyValueStore.Model;
 using Amazon.S3;
 using Amazon.S3.Transfer;
 using ConsoleAppFramework;
@@ -20,12 +17,6 @@
 
 namespace Documentation.Assembler.Cli;
 
-internal enum KvsOperation
-{
-	Puts,
-	Deletes
-}
-
 internal sealed class DeployCommands(ILoggerFactory logger, ICoreService githubActionsService)
 {
 	[SuppressMessage("Usage", "CA2254:Template should be a static expression")]
@@ -144,74 +135,11 @@ public async Task<int> UpdateRedirects(
 		}
 
 		var kvsName = $"elastic-docs-v3-{environment}-redirects-kvs";
+		var cloudFrontClient = new AwsCloudFrontKeyValueStoreProxy(collector, new FileSystem().DirectoryInfo.New(Directory.GetCurrentDirectory()));
 
-		var cfClient = new AmazonCloudFrontClient();
-		var kvsClient = new AmazonCloudFrontKeyValueStoreClient();
-
-		ConsoleApp.Log("Describing KVS");
-		var describeResponse = await cfClient.DescribeKeyValueStoreAsync(new Amazon.CloudFront.Model.DescribeKeyValueStoreRequest { Name = kvsName }, ctx);
-
-		var kvsArn = describeResponse.KeyValueStore.ARN;
-		var eTag = describeResponse.ETag;
-		var existingRedirects = new HashSet<string>();
-
-		var listKeysRequest = new ListKeysRequest { KvsARN = kvsArn };
-		ListKeysResponse listKeysResponse;
-
-		do
-		{
-			listKeysResponse = await kvsClient.ListKeysAsync(listKeysRequest, ctx);
-			foreach (var item in listKeysResponse.Items)
-				_ = existingRedirects.Add(item.Key);
-			listKeysRequest.NextToken = listKeysResponse.NextToken;
-		}
-		while (!string.IsNullOrEmpty(listKeysResponse.NextToken));
-
-		var toPut = sourcedRedirects
-			.Select(kvp => new PutKeyRequestListItem { Key = kvp.Key, Value = kvp.Value });
-		var toDelete = existingRedirects
-			.Except(sourcedRedirects.Keys)
-			.Select(k => new DeleteKeyRequestListItem { Key = k });
-
-		ConsoleApp.Log("Updating redirects in KVS");
-		const int batchSize = 50;
-
-		eTag = await ProcessBatchUpdatesAsync(kvsClient, kvsArn, eTag, toPut, batchSize, KvsOperation.Puts, ctx);
-		_ = await ProcessBatchUpdatesAsync(kvsClient, kvsArn, eTag, toDelete, batchSize, KvsOperation.Deletes, ctx);
+		cloudFrontClient.UpdateRedirects(kvsName, sourcedRedirects);
 
 		await collector.StopAsync(ctx);
 		return collector.Errors;
 	}
-
-	private static async Task<string> ProcessBatchUpdatesAsync(
-		IAmazonCloudFrontKeyValueStore kvsClient,
-		string kvsArn,
-		string eTag,
-		IEnumerable<object> items,
-		int batchSize,
-		KvsOperation operation,
-		Cancel ctx)
-	{
-		var enumerable = items.ToList();
-		for (var i = 0; i < enumerable.Count; i += batchSize)
-		{
-			var batch = enumerable.Skip(i).Take(batchSize);
-			var updateRequest = new UpdateKeysRequest
-			{
-				KvsARN = kvsArn,
-				IfMatch = eTag
-			};
-
-			if (operation is KvsOperation.Puts)
-				updateRequest.Puts = batch.Cast<PutKeyRequestListItem>().ToList();
-			else if (operation is KvsOperation.Deletes)
-				updateRequest.Deletes = batch.Cast<DeleteKeyRequestListItem>().ToList();
-
-			var update = await kvsClient.UpdateKeysAsync(updateRequest, ctx);
-			eTag = update.ETag;
-		}
-
-		return eTag;
-	}
-
 }

src/tooling/docs-assembler/Deploying/AwsCloudFrontKeyValueStoreProxy.cs

@@ -0,0 +1,127 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System.IO.Abstractions;
+using System.Text.Json;
+using ConsoleAppFramework;
+using Documentation.Assembler.Deploying.Serialization;
+using Elastic.Documentation.Diagnostics;
+using Elastic.Documentation.Tooling.ExternalCommands;
+
+namespace Documentation.Assembler.Deploying;
+
+internal enum KvsOperation
+{
+	Puts,
+	Deletes
+}
+
+public class AwsCloudFrontKeyValueStoreProxy(DiagnosticsCollector collector, IDirectoryInfo workingDirectory) : ExternalCommandExecutor(collector, workingDirectory)
+{
+	public void UpdateRedirects(string kvsName, IReadOnlyDictionary<string, string> sourcedRedirects)
+	{
+		var (kvsArn, eTag) = DescribeKeyValueStore(kvsName);
+		if (string.IsNullOrEmpty(kvsArn) || string.IsNullOrEmpty(eTag))
+			return;
+
+		var existingRedirects = ListAllKeys(kvsArn);
+
+		var toPut = sourcedRedirects
+			.Select(kvp => new PutKeyRequestListItem { Key = kvp.Key, Value = kvp.Value });
+		var toDelete = existingRedirects
+			.Except(sourcedRedirects.Keys)
+			.Select(k => new DeleteKeyRequestListItem { Key = k });
+
+		eTag = ProcessBatchUpdates(kvsArn, eTag, toPut, KvsOperation.Puts);
+		_ = ProcessBatchUpdates(kvsArn, eTag, toDelete, KvsOperation.Deletes);
+	}
+
+	private (string? Arn, string? ETag) DescribeKeyValueStore(string kvsName)
+	{
+		ConsoleApp.Log("Describing KeyValueStore");
+		try
+		{
+			var json = Capture("aws", "cloudfront", "describe-key-value-store", "--name", kvsName, "|", "jq", "-c");
+			var describeResponse = JsonSerializer.Deserialize<DescribeKeyValueStoreResponse>(json, AwsCloudFrontKeyValueStoreJsonContext.Default.DescribeKeyValueStoreResponse);
+			if (describeResponse?.ETag is not null && describeResponse.KeyValueStore is { ARN.Length: > 0 })
+				return (describeResponse.KeyValueStore.ARN, describeResponse.ETag);
+
+			Collector.EmitError("", "Could not deserialize the DescribeKeyValueStoreResponse");
+			return (null, null);
+		}
+		catch (Exception e)
+		{
+			Collector.EmitError("", "An error occurred while describing the KeyValueStore", e);
+			return (null, null);
+		}
+	}
+
+	private HashSet<string> ListAllKeys(string kvsArn)
+	{
+		ConsoleApp.Log("Acquiring existing redirects");
+		var allKeys = new HashSet<string>();
+		string[] baseArgs = ["cloudfront-keyvaluestore", "list-keys", "--kvs-arn", kvsArn];
+		string? nextToken = null;
+		try
+		{
+			do
+			{
+				var json = Capture("aws", [.. baseArgs, .. nextToken is not null ? (string[])["--starting-token", nextToken] : [], "|", "jq", "-c"]);
+				var response = JsonSerializer.Deserialize<ListKeysResponse>(json, AwsCloudFrontKeyValueStoreJsonContext.Default.ListKeysResponse);
+
+				if (response?.Items != null)
+				{
+					foreach (var item in response.Items)
+						_ = allKeys.Add(item.Key);
+				}
+
+				nextToken = response?.NextToken;
+			} while (!string.IsNullOrEmpty(nextToken));
+		}
+		catch (Exception e)
+		{
+			Collector.EmitError("", "An error occurred while acquiring existing redirects in the KeyValueStore", e);
+			return [];
+		}
+		return allKeys;
+	}
+
+
+	private string ProcessBatchUpdates(
+		string kvsArn,
+		string eTag,
+		IEnumerable<object> items,
+		KvsOperation operation)
+	{
+		const int batchSize = 50;
+		ConsoleApp.Log($"Processing {items.Count()} items in batches of {batchSize} for {operation} update operation.");
+		try
+		{
+			foreach (var batch in items.Chunk(batchSize))
+			{
+				var payload = operation switch
+				{
+					KvsOperation.Puts => JsonSerializer.Serialize(batch.Cast<PutKeyRequestListItem>().ToList(),
+						AwsCloudFrontKeyValueStoreJsonContext.Default.ListPutKeyRequestListItem),
+					KvsOperation.Deletes => JsonSerializer.Serialize(batch.Cast<DeleteKeyRequestListItem>().ToList(),
+						AwsCloudFrontKeyValueStoreJsonContext.Default.ListDeleteKeyRequestListItem),
+					_ => string.Empty
+				};
+				var responseJson = Capture(false, 1, "aws", "cloudfront-keyvaluestore", "update-keys", "--kvs-arn", kvsArn, "--if-match", eTag,
+					$"--{operation.ToString().ToLowerInvariant()}", "--payload", payload, "|", "jq", "-c");
+				var updateResponse = JsonSerializer.Deserialize<UpdateKeysResponse>(responseJson, AwsCloudFrontKeyValueStoreJsonContext.Default.UpdateKeysResponse);
+
+				if (string.IsNullOrEmpty(updateResponse?.ETag))
+					throw new Exception("Failed to get new ETag after update operation.");
+
+				eTag = updateResponse.ETag;
+			}
+		}
+		catch (Exception e)
+		{
+			Collector.EmitError("", $"An error occurred while performing a {operation} update to the KeyValueStore", e);
+		}
+		return eTag;
+	}
+}

src/tooling/docs-assembler/Deploying/Serialization/AwsCloudFrontKeyValueStoreModels.cs

@@ -0,0 +1,37 @@
+// Licensed to Elasticsearch B.V under one or more agreements.
+// Elasticsearch B.V licenses this file to you under the Apache 2.0 License.
+// See the LICENSE file in the project root for more information
+
+using System.Text.Json.Serialization;
+
+namespace Documentation.Assembler.Deploying.Serialization;
+
+public record DescribeKeyValueStoreResponse([property: JsonPropertyName("ETag")] string ETag, [property: JsonPropertyName("KeyValueStore")] KeyValueStore KeyValueStore);
+public record KeyValueStore([property: JsonPropertyName("ARN")] string ARN);
+
+public record ListKeysResponse([property: JsonPropertyName("NextToken")] string? NextToken, [property: JsonPropertyName("Items")] List<KeyItem> Items);
+public record KeyItem([property: JsonPropertyName("Key")] string Key);
+
+public record UpdateKeysResponse([property: JsonPropertyName("ETag")] string ETag);
+
+public record PutKeyRequestListItem
+{
+	[JsonPropertyName("Key")]
+	public required string Key { get; init; }
+	[JsonPropertyName("Value")]
+	public required string Value { get; init; }
+}
+
+public record DeleteKeyRequestListItem
+{
+	[JsonPropertyName("Key")]
+	public required string Key { get; init; }
+}
+
+[JsonSourceGenerationOptions(WriteIndented = false, UseStringEnumConverter = true)]
+[JsonSerializable(typeof(DescribeKeyValueStoreResponse))]
+[JsonSerializable(typeof(ListKeysResponse))]
+[JsonSerializable(typeof(UpdateKeysResponse))]
+[JsonSerializable(typeof(List<PutKeyRequestListItem>))]
+[JsonSerializable(typeof(List<DeleteKeyRequestListItem>))]
+internal sealed partial class AwsCloudFrontKeyValueStoreJsonContext : JsonSerializerContext;

src/tooling/docs-assembler/docs-assembler.csproj

@@ -17,8 +17,6 @@
   </PropertyGroup>
 
   <ItemGroup>
-    <PackageReference Include="AWSSDK.CloudFront" />
-    <PackageReference Include="AWSSDK.CloudFrontKeyValueStore" />
     <PackageReference Include="AWSSDK.S3"/>
     <PackageReference Include="ConsoleAppFramework.Abstractions"/>
     <PackageReference Include="ConsoleAppFramework" />