[CosmosDB] az cosmosdb fleet: Update GA contract (#9395)

* Changed version number

* Generated python SDK

* Add test and update documentation for IsDefaultBucket

* Successful recordings

* Update and record throughput bucketing test

* Garnet issues fixed

* Fleet contract updated

* Flake 8

* Live_only tests due to missing role assignments

* Mongo cluster tests working

* Corrected live_only

* Corrected HISTORY entries

* Remaining recordings

* Bug fix

* Remaining role issue - live_only

* Correct bucketing test recording

* Updated versioning

---------

Co-authored-by: Achint Agrawal <[email protected]>

Author: pjohari-ms

src/cosmosdb-preview/HISTORY.rst

@@ -2,6 +2,10 @@
 Release History
 ===============
 
+1.6.2
+* Added throughput bucketing.
++++++++
+
 1.6.1
 +++++
 * Fix SQL container throughput update to preserve existing throughput buckets when not explicitly specified.

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -452,7 +452,8 @@ def validate_fleetspace_body(cmd, ns):
         if not isinstance(tp_config, dict):
             raise InvalidArgumentValueError('Missing or invalid "throughputPoolConfiguration" in properties.')
 
-        for field in ['minThroughput', 'maxThroughput', 'serviceTier', 'dataRegions']:
+        # Check for minThroughput and maxThroughput in throughputPoolConfiguration
+        for field in ['minThroughput', 'maxThroughput']:
             if field not in tp_config:
                 raise InvalidArgumentValueError(f'Missing "{field}" in throughputPoolConfiguration.')
 
@@ -462,11 +463,14 @@ def validate_fleetspace_body(cmd, ns):
         if not isinstance(tp_config['maxThroughput'], int) or tp_config['maxThroughput'] <= 0:
             raise InvalidArgumentValueError('"maxThroughput" must be a positive integer.')
 
-        if not isinstance(tp_config['serviceTier'], str):
-            raise InvalidArgumentValueError('"serviceTier" must be a string.')
+        # Check for serviceTier and dataRegions at base properties level
+        if 'serviceTier' in props:
+            if not isinstance(props['serviceTier'], str):
+                raise InvalidArgumentValueError('"serviceTier" must be a string.')
 
-        if not isinstance(tp_config['dataRegions'], list) or not all(isinstance(r, str) for r in tp_config['dataRegions']):
-            raise InvalidArgumentValueError('"dataRegions" must be a list of strings.')
+        if 'dataRegions' in props:
+            if not isinstance(props['dataRegions'], list) or not all(isinstance(r, str) for r in props['dataRegions']):
+                raise InvalidArgumentValueError('"dataRegions" must be a list of strings.')
 
         ns.fleetspace_body = body
 

src/cosmosdb-preview/azext_cosmosdb_preview/_validators.py

@@ -3305,23 +3305,33 @@ def cli_cosmosdb_fleetspace_create(client,
 
     """Creates an Azure Cosmos DB Fleetspace."""
 
+    # Extract service_tier and data_regions from base level (mandatory for create)
+    service_tier = fleetspace_body['properties'].get('serviceTier')
+    data_regions = fleetspace_body['properties'].get('dataRegions')
+
+    if not service_tier:
+        raise CLIError('Missing required field "serviceTier" in properties.')
+
+    if not data_regions:
+        raise CLIError('Missing required field "dataRegions" in properties.')
+
     throughput_pool_config = FleetspacePropertiesThroughputPoolConfiguration(
         min_throughput=fleetspace_body['properties']['throughputPoolConfiguration']['minThroughput'],
-        max_throughput=fleetspace_body['properties']['throughputPoolConfiguration']['maxThroughput'],
-        service_tier=fleetspace_body['properties']['throughputPoolConfiguration']['serviceTier'],
-        data_regions=fleetspace_body['properties']['throughputPoolConfiguration']['dataRegions']
+        max_throughput=fleetspace_body['properties']['throughputPoolConfiguration']['maxThroughput']
     )
 
-    fleetspace_body = FleetspaceResource(
+    fleetspace_resource = FleetspaceResource(
         fleetspace_api_kind="NoSQL",
-        throughput_pool_configuration=throughput_pool_config
+        throughput_pool_configuration=throughput_pool_config,
+        service_tier=service_tier,
+        data_regions=data_regions
     )
 
     return client.begin_create(
         resource_group_name=resource_group_name,
         fleet_name=fleet_name,
         fleetspace_name=fleetspace_name,
-        body=fleetspace_body
+        body=fleetspace_resource
     )
 
 
@@ -3333,22 +3343,27 @@ def cli_cosmosdb_fleetspace_update(client,
 
     """Updates an existing Azure Cosmos DB Fleetspace."""
 
+    # Extract service_tier and data_regions from base level (optional for update)
+    service_tier = fleetspace_body['properties'].get('serviceTier')
+    data_regions = fleetspace_body['properties'].get('dataRegions')
+
     throughput_pool_config = FleetspacePropertiesThroughputPoolConfiguration(
         min_throughput=fleetspace_body['properties']['throughputPoolConfiguration']['minThroughput'],
-        max_throughput=fleetspace_body['properties']['throughputPoolConfiguration']['maxThroughput'],
-        service_tier=fleetspace_body['properties']['throughputPoolConfiguration']['serviceTier']
+        max_throughput=fleetspace_body['properties']['throughputPoolConfiguration']['maxThroughput']
     )
 
-    fleetspace_body = FleetspaceResource(
+    fleetspace_resource = FleetspaceResource(
         fleetspace_api_kind="NoSQL",
-        throughput_pool_configuration=throughput_pool_config
+        throughput_pool_configuration=throughput_pool_config,
+        service_tier=service_tier,
+        data_regions=data_regions
     )
 
     return client.begin_update(
         resource_group_name=resource_group_name,
         fleet_name=fleet_name,
         fleetspace_name=fleetspace_name,
-        body=fleetspace_body
+        body=fleetspace_resource
     )