Merge pull request #190 from AzureCosmosDB/copilot/fix-189

Fix InitClientEncryption configuration not being respected in CosmosDataSinkExtension

Author: markjbrown

Extensions/Cosmos/Cosmos.DataTransfer.CosmosExtension.UnitTests/CosmosSinkSettingsTests.cs

@@ -155,4 +155,58 @@ public void GetValidationErrors_WhenPartitionKeysInvalid_ReturnsErrors()
 
         Assert.AreEqual(1, validationErrors.Count(v => v.Contains(nameof(CosmosSinkSettings.PartitionKeyPaths))));
     }
+
+    [TestMethod]
+    public void GetValidationErrors_WhenInitClientEncryptionWithoutRbac_ReturnsError()
+    {
+        var settings = new CosmosSinkSettings
+        {
+            ConnectionString = "AccountEndpoint=https://localhost:8081/;AccountKey=",
+            Database = "db",
+            Container = "container",
+            UseRbacAuth = false,
+            InitClientEncryption = true,
+        };
+
+        var validationErrors = settings.GetValidationErrors();
+        LogErrors(validationErrors);
+
+        Assert.AreEqual(1, validationErrors.Count(v => v.Contains(nameof(CosmosSinkSettings.InitClientEncryption))));
+    }
+
+    [TestMethod]
+    public void GetValidationErrors_WhenInitClientEncryptionWithRbac_Succeeds()
+    {
+        var settings = new CosmosSinkSettings
+        {
+            UseRbacAuth = true,
+            AccountEndpoint = "https://localhost:8081/",
+            Database = "db",
+            Container = "container",
+            InitClientEncryption = true,
+        };
+
+        var validationErrors = settings.GetValidationErrors();
+        LogErrors(validationErrors);
+
+        Assert.AreEqual(0, validationErrors.Count());
+    }
+
+    [TestMethod]
+    public void GetValidationErrors_WhenRbacWithoutInitClientEncryption_Succeeds()
+    {
+        var settings = new CosmosSinkSettings
+        {
+            UseRbacAuth = true,
+            AccountEndpoint = "https://localhost:8081/",
+            Database = "db",
+            Container = "container",
+            InitClientEncryption = false,
+        };
+
+        var validationErrors = settings.GetValidationErrors();
+        LogErrors(validationErrors);
+
+        Assert.AreEqual(0, validationErrors.Count());
+    }
 }

Extensions/Cosmos/Cosmos.DataTransfer.CosmosExtension/CosmosDataSinkExtension.cs

@@ -141,9 +141,18 @@ public async Task WriteAsync(IAsyncEnumerable<IDataItem> dataItems, IConfigurati
 
             var client = CosmosExtensionServices.CreateClient(settings, DisplayName, dataSource.DisplayName);
 
-            Container container = settings.UseRbacAuth
-                ? await client.GetContainer(settings.Database, settings.Container).InitializeEncryptionAsync(cancellationToken)
-                : await CreateDatabaseAndContainerAsync(client, settings, logger, cancellationToken);
+            Container container;
+            if (settings.UseRbacAuth)
+            {
+                var cosmosContainer = client.GetContainer(settings.Database, settings.Container);
+                container = settings.InitClientEncryption
+                    ? await cosmosContainer.InitializeEncryptionAsync(cancellationToken)
+                    : cosmosContainer;
+            }
+            else
+            {
+                container = await CreateDatabaseAndContainerAsync(client, settings, logger, cancellationToken);
+            }
 
             await CosmosExtensionServices.VerifyContainerAccess(container, settings.Container, logger, cancellationToken);