Merge pull request #194 from AzureCosmosDB/copilot/fix-193

Add ReplaceIfExists setting for Azure Table API upsert functionality

Author: markjbrown

Directory.Packages.props

@@ -21,6 +21,7 @@
     <PackageVersion Include="Microsoft.Extensions.Configuration" Version="8.0.0" />
     <PackageVersion Include="Microsoft.Extensions.Configuration.Abstractions" Version="8.0.0" />
     <PackageVersion Include="Microsoft.Extensions.Configuration.Binder" Version="8.0.2" />
+    <PackageVersion Include="Microsoft.Extensions.Configuration.Json" Version="8.0.0" />
     <PackageVersion Include="Microsoft.Extensions.Configuration.UserSecrets" Version="8.0.1" />
     <PackageVersion Include="Microsoft.Extensions.Hosting" Version="8.0.1" />
     <PackageVersion Include="Microsoft.Extensions.Logging.Abstractions" Version="8.0.2" />

Extensions/AzureTableAPI/Cosmos.DataTransfer.AzureTableAPIExtension.UnitTests/AzureTableAPIDataSinkExtensionTests.cs

@@ -0,0 +1,91 @@
+using Microsoft.VisualStudio.TestTools.UnitTesting;
+using Cosmos.DataTransfer.AzureTableAPIExtension.Settings;
+using Microsoft.Extensions.Configuration;
+using System.Text.Json;
+
+namespace Cosmos.DataTransfer.AzureTableAPIExtension.UnitTests
+{
+    [TestClass]
+    public class AzureTableAPIDataSinkExtensionTests
+    {
+        [TestMethod]
+        public void AzureTableAPIDataSinkSettings_WriteMode_DefaultsToCreate()
+        {
+            var settings = new AzureTableAPIDataSinkSettings();
+            
+            Assert.AreEqual(EntityWriteMode.Create, settings.WriteMode ?? EntityWriteMode.Replace, "WriteMode should default to Create");
+        }
+
+        [TestMethod]
+        public void AzureTableAPIDataSinkSettings_WriteMode_CanBeSetToCreate()
+        {
+            var settings = new AzureTableAPIDataSinkSettings()
+            {
+                WriteMode = EntityWriteMode.Create
+            };
+            
+            Assert.AreEqual(EntityWriteMode.Create, settings.WriteMode ?? EntityWriteMode.Replace, "WriteMode should be settable to Create");
+        }
+
+        [TestMethod]
+        public void AzureTableAPIDataSinkSettings_WriteMode_CanBeSetToReplace()
+        {
+            var settings = new AzureTableAPIDataSinkSettings()
+            {
+                WriteMode = EntityWriteMode.Replace
+            };
+            
+            Assert.AreEqual(EntityWriteMode.Replace, settings.WriteMode ?? EntityWriteMode.Create, "WriteMode should be settable to Replace");
+        }
+
+        [TestMethod]
+        public void AzureTableAPIDataSinkSettings_WriteMode_CanBeSetToMerge()
+        {
+            var settings = new AzureTableAPIDataSinkSettings()
+            {
+                WriteMode = EntityWriteMode.Merge
+            };
+            
+            Assert.AreEqual(EntityWriteMode.Merge, settings.WriteMode ?? EntityWriteMode.Create, "WriteMode should be settable to Merge");
+        }
+
+        [TestMethod]
+        public void AzureTableAPIDataSinkSettings_WriteMode_CanBeNull()
+        {
+            var settings = new AzureTableAPIDataSinkSettings()
+            {
+                WriteMode = null
+            };
+            
+            Assert.IsNull(settings.WriteMode, "WriteMode should be settable to null");
+        }
+
+        [TestMethod]
+        public void AzureTableAPIDataSinkSettings_WriteMode_JsonSerializationSupported()
+        {
+            // Test JSON to enum conversion for Create
+            var jsonCreate = """{"WriteMode": "Create"}""";
+            var configCreate = new ConfigurationBuilder()
+                .AddJsonStream(new MemoryStream(System.Text.Encoding.UTF8.GetBytes(jsonCreate)))
+                .Build();
+            var settingsCreate = configCreate.Get<AzureTableAPIDataSinkSettings>();
+            Assert.AreEqual(EntityWriteMode.Create, settingsCreate?.WriteMode, "WriteMode should be deserialized from JSON string 'Create'");
+
+            // Test JSON to enum conversion for Replace
+            var jsonReplace = """{"WriteMode": "Replace"}""";
+            var configReplace = new ConfigurationBuilder()
+                .AddJsonStream(new MemoryStream(System.Text.Encoding.UTF8.GetBytes(jsonReplace)))
+                .Build();
+            var settingsReplace = configReplace.Get<AzureTableAPIDataSinkSettings>();
+            Assert.AreEqual(EntityWriteMode.Replace, settingsReplace?.WriteMode, "WriteMode should be deserialized from JSON string 'Replace'");
+
+            // Test JSON to enum conversion for Merge
+            var jsonMerge = """{"WriteMode": "Merge"}""";
+            var configMerge = new ConfigurationBuilder()
+                .AddJsonStream(new MemoryStream(System.Text.Encoding.UTF8.GetBytes(jsonMerge)))
+                .Build();
+            var settingsMerge = configMerge.Get<AzureTableAPIDataSinkSettings>();
+            Assert.AreEqual(EntityWriteMode.Merge, settingsMerge?.WriteMode, "WriteMode should be deserialized from JSON string 'Merge'");
+        }
+    }
+}

Extensions/AzureTableAPI/Cosmos.DataTransfer.AzureTableAPIExtension.UnitTests/Cosmos.DataTransfer.AzureTableAPIExtension.UnitTests.csproj

@@ -10,6 +10,9 @@
 
   <ItemGroup>
     <PackageReference Include="Azure.Data.Tables" />
+    <PackageReference Include="Microsoft.Extensions.Configuration" />
+    <PackageReference Include="Microsoft.Extensions.Configuration.Binder" />
+    <PackageReference Include="Microsoft.Extensions.Configuration.Json" />
     <PackageReference Include="Microsoft.NET.Test.Sdk" />
     <PackageReference Include="MSTest.TestAdapter" />
     <PackageReference Include="MSTest.TestFramework" />

Extensions/AzureTableAPI/Cosmos.DataTransfer.AzureTableAPIExtension/AzureTableAPIDataSinkExtension.cs

@@ -48,6 +48,7 @@ public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfigurati
             await tableClient.CreateIfNotExistsAsync().ConfigureAwait(false);
 
             var maxConcurrency = settings.MaxConcurrentEntityWrites ?? 10;
+            var writeMode = settings.WriteMode ?? EntityWriteMode.Create;
 
             logger.LogInformation("Writing data to Azure Table Storage with a maximum of {MaxConcurrency} concurrent writes.", maxConcurrency);
 
@@ -60,7 +61,7 @@ await Parallel.ForEachAsync<IDataItem>(dataItems,
                 try
                 {
                     var entity = item.ToTableEntity(settings.PartitionKeyFieldName, settings.RowKeyFieldName);
-                    await AddEntityWithRetryAsync(tableClient, entity, ct);
+                    await AddEntityWithRetryAsync(tableClient, entity, writeMode, ct);
                 }
                 catch (Exception ex)
                 {
@@ -81,17 +82,31 @@ public IEnumerable<IDataExtensionSettings> GetSettings()
         /// </summary>
         /// <param name="tableClient"></param>
         /// <param name="entity"></param>
+        /// <param name="writeMode"></param>
         /// <param name="cancellationToken"></param>
         /// <returns></returns>
-        private static async Task AddEntityWithRetryAsync(TableClient tableClient, TableEntity entity, CancellationToken cancellationToken)
+        private static async Task AddEntityWithRetryAsync(TableClient tableClient, TableEntity entity, EntityWriteMode writeMode, CancellationToken cancellationToken)
         {
             var retryPolicy = Policy
                 .Handle<RequestFailedException>(ex => TransientStatusCodes.Contains(ex.Status))
                 .WaitAndRetryAsync(3, retryAttempt => TimeSpan.FromSeconds(Math.Pow(2, retryAttempt)));
 
             await retryPolicy.ExecuteAsync(async () =>
             {
-                await tableClient.AddEntityAsync(entity, cancellationToken);
+                switch (writeMode)
+                {
+                    case EntityWriteMode.Create:
+                        await tableClient.AddEntityAsync(entity, cancellationToken);
+                        break;
+                    case EntityWriteMode.Replace:
+                        await tableClient.UpsertEntityAsync(entity, TableUpdateMode.Replace, cancellationToken);
+                        break;
+                    case EntityWriteMode.Merge:
+                        await tableClient.UpsertEntityAsync(entity, TableUpdateMode.Merge, cancellationToken);
+                        break;
+                    default:
+                        throw new ArgumentOutOfRangeException(nameof(writeMode), writeMode, "Unsupported EntityWriteMode");
+                }
             });
         }
     }

Extensions/AzureTableAPI/Cosmos.DataTransfer.AzureTableAPIExtension/Settings/AzureTableAPIDataSinkSettings.cs

@@ -7,5 +7,13 @@ public class AzureTableAPIDataSinkSettings : AzureTableAPISettingsBase
         /// This setting is used to control the number of concurrent writes to the Azure Table API.
         /// </summary>
         public int? MaxConcurrentEntityWrites { get; set; }
+        
+        /// <summary>
+        /// Specifies the behavior when writing entities to the table.
+        /// Create: Adds new entities only, fails if entity already exists (default).
+        /// Replace: Upserts entities, completely replacing existing ones.
+        /// Merge: Upserts entities, merging properties with existing ones.
+        /// </summary>
+        public EntityWriteMode? WriteMode { get; set; } = EntityWriteMode.Create;
     }
 }

Extensions/AzureTableAPI/Cosmos.DataTransfer.AzureTableAPIExtension/Settings/EntityWriteMode.cs

@@ -0,0 +1,26 @@
+namespace Cosmos.DataTransfer.AzureTableAPIExtension.Settings
+{
+    /// <summary>
+    /// Defines the behavior when writing entities to Azure Table API.
+    /// </summary>
+    public enum EntityWriteMode
+    {
+        /// <summary>
+        /// Creates new entities only. Fails if an entity with the same partition key and row key already exists.
+        /// Uses AddEntityAsync method.
+        /// </summary>
+        Create,
+        
+        /// <summary>
+        /// Replaces existing entities or creates new ones. Completely replaces the entity in the table.
+        /// Uses UpsertEntityAsync with TableUpdateMode.Replace.
+        /// </summary>
+        Replace,
+        
+        /// <summary>
+        /// Merges with existing entities or creates new ones. Merges the properties of the supplied entity with the entity in the table.
+        /// Uses UpsertEntityAsync with TableUpdateMode.Merge.
+        /// </summary>
+        Merge
+    }
+}

Extensions/AzureTableAPI/README.md

@@ -30,11 +30,15 @@ The following setting is supported for the Source:
 
 ### Additional Sink Settings
 
-The AzureTableAPI Source extension has an additional setting that can be configured for writing Table entities.
+The AzureTableAPI Sink extension has additional settings that can be configured for writing Table entities.
 
-The following setting is supported for the Sink:
+The following settings are supported for the Sink:
 
 - `MaxConcurrentEntityWrites` - The Maximum number of concurrent entity writes to the Azure Table API. This setting is used to control the number of concurrent writes to the Azure Table API.
+- `WriteMode` - Specifies the behavior when writing entities to the table. Options are:
+  - `Create` (default): Creates new entities only. Fails if an entity with the same partition key and row key already exists.
+  - `Replace`: Upserts entities, completely replacing existing ones if they exist.
+  - `Merge`: Upserts entities, merging properties with existing entities if they exist.
 
 ## Authentication Methods
 
@@ -47,6 +51,7 @@ The AzureTableAPI extension supports two authentication methods for connecting t
 
 ### Example RBAC Settings
 
+**Source Example (RBAC)**
 ```json
 {
   "UseRbacAuth": true,
@@ -59,6 +64,20 @@ The AzureTableAPI extension supports two authentication methods for connecting t
 }
 ```
 
+**Sink Example (RBAC)**
+```json
+{
+  "UseRbacAuth": true,
+  "AccountEndpoint": "https://<storage-account-name>.table.core.windows.net",
+  "Table": "DestinationTable1",
+  "PartitionKeyFieldName": "State",
+  "RowKeyFieldName": "id",
+  "EnableInteractiveCredentials": true,
+  "WriteMode": "Merge",
+  "MaxConcurrentEntityWrites": 10
+}
+```
+
 ### Example ConnectionString Source and Sink Settings Usage
 
 The following are a couple example `settings.json` files for configuring the AzureTableAPI Source and Sink extensions.
@@ -82,6 +101,8 @@ The following are a couple example `settings.json` files for configuring the Azu
   "ConnectionString": "DefaultEndpointsProtocol=https;AccountName=<storage-account-name>;AccountKey=<key>;EndpointSuffix=core.windows.net",
   "Table": "DestinationTable1",
   "PartitionKeyFieldName": "State",
-  "RowKeyFieldName":  "id"
+  "RowKeyFieldName": "id",
+  "WriteMode": "Replace",
+  "MaxConcurrentEntityWrites": 5
 }
 ```