Fix SQL container throughput update to preserve existing throughput buckets

src/CosmosDB/CosmosDB.Test/ScenarioTests/SqlOperationsTests.cs

@@ -74,9 +74,23 @@ public void TestSqlThroughputCmdlets()
 
         [Fact]
         [Trait(Category.AcceptanceType, Category.CheckIn)]
-        public void TestSqlThroughputBucketsCmdlets()
+        public void TestSqlThroughputBucketsManualCmdlets()
         {
-            TestRunner.RunTestScript("Test-SqlThroughputBucketsCmdlets");
+            TestRunner.RunTestScript("Test-SqlThroughputBucketsCmdlets-ManualContainer");
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestSqlThroughputBucketsAutoscaleCmdlets()
+        {
+            TestRunner.RunTestScript("Test-SqlThroughputBucketsCmdlets-AutoscaleContainer");
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void TestSqlThroughputBucketsMigrationCmdlets()
+        {
+            TestRunner.RunTestScript("Test-SqlThroughputBucketsCmdlets-Migration");
         }
 
         [Fact]

src/CosmosDB/CosmosDB/ChangeLog.md

@@ -19,6 +19,7 @@
 -->
 
 ## Upcoming Release
+* Added support for preserving existing throughput buckets during SQL container throughput updates.
 
 ## Version 2.0.7
 * Added support for throughput buckets.

src/CosmosDB/CosmosDB/Helpers/ThroughputHelper.cs

@@ -31,33 +31,23 @@ public static ThroughputSettingsUpdateParameters CreateThroughputSettingsObject(
             AutoscaleMaxThroughput = autoscaleMaxThroughput;
 
             ThroughputHelper.ThroughputValidation(validateBothPresent:true);
+            List<ThroughputBucketResource> buckets = ToBucketList(throughputBucketsObject);
             ThroughputSettingsUpdateParameters throughputSettingsUpdateParameters = new ThroughputSettingsUpdateParameters();
+
             if (throughput != null)
             {
                 throughputSettingsUpdateParameters.Resource = new ThroughputSettingsResource
                 {
                     Throughput = throughput.Value,
-                    ThroughputBuckets = throughputBucketsObject != null && throughputBucketsObject.Length > 0
-                        ? new List<ThroughputBucketResource>(throughputBucketsObject.Select(t => new ThroughputBucketResource
-                        {
-                            Id = t.Id,
-                            MaxThroughputPercentage = t.MaxThroughputPercentage
-                        }))
-                        : null
+                    ThroughputBuckets = buckets
                 };
             }
             else if (autoscaleMaxThroughput != null)
             {
                 throughputSettingsUpdateParameters.Resource = new ThroughputSettingsResource
                 {
                     AutoscaleSettings = new AutoscaleSettingsResource { MaxThroughput = autoscaleMaxThroughput.Value },
-                    ThroughputBuckets = throughputBucketsObject != null && throughputBucketsObject.Length > 0
-                        ? new List<ThroughputBucketResource>(throughputBucketsObject.Select(t => new ThroughputBucketResource
-                        {
-                            Id = t.Id,
-                            MaxThroughputPercentage = t.MaxThroughputPercentage
-                        }))
-                        : null
+                    ThroughputBuckets = buckets
                 };
             }
 
@@ -97,5 +87,22 @@ private static void ThroughputValidation(bool validateBothPresent = false)
             return;
         }
 
+        private static List<ThroughputBucketResource> ToBucketList(PSThroughputBucket[] throughputBucketsObject)
+        {
+            if (throughputBucketsObject == null)
+            {
+                return null;
+            }
+            if (throughputBucketsObject.Length == 0)
+            {
+                return new List<ThroughputBucketResource>();
+            }
+            return new List<ThroughputBucketResource>(throughputBucketsObject.Select(t => new ThroughputBucketResource
+            {
+                Id = t.Id,
+                MaxThroughputPercentage = t.MaxThroughputPercentage
+            }));
+        }
+
     }
 }

src/CosmosDB/CosmosDB/SQL/UpdateAzCosmosDBSqlContainerThroughput.cs

@@ -13,6 +13,7 @@
 // ----------------------------------------------------------------------------------
 
 using System.Management.Automation;
+using System.Linq;
 using Microsoft.Azure.Commands.CosmosDB.Models;
 using Microsoft.Azure.Commands.CosmosDB.Helpers;
 using Microsoft.Azure.Management.Internal.Resources.Utilities.Models;
@@ -65,5 +66,19 @@ public override void CallSDKMethod(ThroughputSettingsUpdateParameters throughput
             }
         }
 
+        protected override PSThroughputBucket[] GetExistingThroughputBuckets()
+        {
+            // Read current throughput to preserve existing buckets when user omits the parameter
+            ThroughputSettingsGetResults current = CosmosDBManagementClient.SqlResources
+                .GetSqlContainerThroughputWithHttpMessagesAsync(ResourceGroupName, AccountName, DatabaseName, Name)
+                .GetAwaiter().GetResult().Body;
+
+            System.Collections.Generic.IList<ThroughputBucketResource> buckets = current?.Resource?.ThroughputBuckets;
+            if (buckets == null)
+            {
+                return null;
+            }
+            return buckets.Select(b => new PSThroughputBucket(b)).ToArray();
+        }
     }
 }

src/CosmosDB/CosmosDB/UpdateAzCosmosDBThroughput.cs

@@ -55,14 +55,21 @@ public override void ExecuteCmdlet()
                 PopulateFromInputObject();
             }
 
-            ThroughputSettingsUpdateParameters throughputSettingsUpdateParameters = ThroughputHelper.CreateThroughputSettingsObject(Throughput, AutoscaleMaxThroughput, ThroughputBucketsObject);
+            // Determine effective throughput buckets:
+            // - If user omitted the parameter (null), fetch and preserve existing buckets via hook
+            // - If user passed empty array, clear buckets by sending an empty list
+            // - If user passed values, replace with those values
+            PSThroughputBucket[] effectiveThroughputBuckets = ThroughputBucketsObject ?? GetExistingThroughputBuckets();
+
+            ThroughputSettingsUpdateParameters throughputSettingsUpdateParameters = ThroughputHelper.CreateThroughputSettingsObject(Throughput, AutoscaleMaxThroughput, effectiveThroughputBuckets);
 
             CallSDKMethod(throughputSettingsUpdateParameters);
         }
 
         public virtual void PopulateFromParentObject() { }
         public virtual void PopulateFromInputObject() { }
         public virtual void CallSDKMethod(ThroughputSettingsUpdateParameters throughputSettingsUpdateParameters) { }
+        protected virtual PSThroughputBucket[] GetExistingThroughputBuckets() { return null; }
 
     }
 }