Merge pull request Azure#1 from pjohari-ms/capacitymodechangesonpr7613

Capacity Mode changes

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -38,7 +38,7 @@
     tags_type, get_resource_name_completion_list, name_type, get_enum_type, get_three_state_flag, get_location_type)
 
 from azure.mgmt.cosmosdb.models import (
-    DefaultConsistencyLevel, DatabaseAccountKind, ServerVersion, NetworkAclBypass, BackupPolicyType, AnalyticalStorageSchemaType, BackupStorageRedundancy)
+    DefaultConsistencyLevel, DatabaseAccountKind, ServerVersion, NetworkAclBypass, BackupPolicyType, AnalyticalStorageSchemaType, BackupStorageRedundancy, CapacityMode, CapacityModeTransitionStatus)
 
 from azure.cli.command_modules.cosmosdb.actions import (
     CreateLocation, CreateDatabaseRestoreResource, UtcDatetimeAction)
@@ -343,6 +343,7 @@ def load_arguments(self, _):
             c.argument('default_priority_level', arg_type=get_enum_type(DefaultPriorityLevel), help="Default Priority Level of Request if not specified.", is_preview=True)
             c.argument('enable_prpp_autoscale', arg_type=get_three_state_flag(), help="Enable or disable PerRegionPerPartitionAutoscale.", is_preview=True)
             c.argument('enable_partition_merge', arg_type=get_three_state_flag(), help="Flag to enable partition merge on the account.")
+            c.argument('capacity_mode', options_list=['--capacity-mode'], arg_type=get_enum_type(CapacityMode), help="CapacityMode of the account.", is_preview=True)
 
     with self.argument_context('cosmosdb update') as c:
         c.argument('key_uri', help="The URI of the key vault", is_preview=True)

src/cosmosdb-preview/azext_cosmosdb_preview/custom.py

@@ -55,10 +55,6 @@
     ManagedServiceIdentity,
     AnalyticalStorageConfiguration,
     ManagedServiceIdentityUserAssignedIdentity,
-    MongoCluster,
-    NodeGroupSpec,
-    NodeKind,
-    FirewallRule,
     CosmosCassandraDataTransferDataSourceSink,
     CosmosSqlDataTransferDataSourceSink,
     CosmosMongoDataTransferDataSourceSink
@@ -113,7 +109,7 @@ def cli_cosmosdb_mongocluster_firewall_rule_create(client,
 
     '''Creates an Azure Cosmos DB Mongo Cluster Firewall rule'''
 
-    firewall_rule = FirewallRule(start_ip_address=start_ip_address, end_ip_address=end_ip_address)
+    firewall_rule = None
 
     return client.begin_create_or_update_firewall_rule(resource_group_name, cluster_name, rule_name, firewall_rule)
 
@@ -135,7 +131,7 @@ def cli_cosmosdb_mongocluster_firewall_rule_update(client,
     if end_ip_address is None:
         end_ip_address = mongo_cluster_firewallRule.endIpAddress
 
-    firewall_rule = FirewallRule(start_ip_address=start_ip_address, end_ip_address=end_ip_address)
+    firewall_rule = None
 
     return client.begin_create_or_update_firewall_rule(resource_group_name, cluster_name, rule_name, firewall_rule)
 
@@ -179,23 +175,10 @@ def cli_cosmosdb_mongocluster_create(client,
     if ((administrator_login is None and administrator_login_password is not None) or (administrator_login is not None and administrator_login_password is None)):
         raise InvalidArgumentValueError('Both(administrator_login and administrator_login_password) Mongo Cluster admin user parameters must be provided together')
 
-    node_group_spec = NodeGroupSpec(
-        sku=shard_node_tier,
-        disk_size_gb=shard_node_disk_size_gb,
-        enable_ha=shard_node_ha,
-        kind=NodeKind.SHARD.value,
-        node_count=shard_node_count
-    )
+    node_group_spec = None
 
     node_group_specs = [node_group_spec]
-    mongodb_cluster = MongoCluster(
-        location=location,
-        tags=tags,
-        create_mode=CreateMode.DEFAULT.value,
-        administrator_login=administrator_login,
-        administrator_login_password=administrator_login_password,
-        server_version=server_version,
-        node_group_specs=node_group_specs)
+    mongodb_cluster = None
 
     return client.begin_create_or_update(resource_group_name, cluster_name, mongodb_cluster)
 
@@ -239,23 +222,10 @@ def cli_cosmosdb_mongocluster_update(client,
     if shard_node_ha is None:
         shard_node_ha = mongo_cluster_resource.node_group_specs[0].enable_ha
 
-    node_group_spec = NodeGroupSpec(
-        sku=shard_node_tier,
-        disk_size_gb=shard_node_disk_size_gb,
-        enable_ha=shard_node_ha,
-        kind=NodeKind.SHARD.value,
-        node_count=None,
-    )
+    node_group_spec = None
 
     node_group_specs = [node_group_spec]
-    mongodb_cluster = MongoCluster(
-        location=location,
-        tags=tags,
-        create_mode=CreateMode.DEFAULT.value,
-        administrator_login=administrator_login,
-        administrator_login_password=administrator_login_password,
-        server_version=server_version,
-        node_group_specs=node_group_specs)
+    mongodb_cluster = None
 
     return client.begin_create_or_update(resource_group_name, cluster_name, mongodb_cluster)
 
@@ -781,7 +751,8 @@ def cli_cosmosdb_create(cmd,
                         enable_priority_based_execution=None,
                         default_priority_level=None,
                         enable_prpp_autoscale=None,
-                        enable_partition_merge=None):
+                        enable_partition_merge=None,
+                        capacity_mode=None):
     """Create a new Azure Cosmos DB database account."""
 
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
@@ -839,7 +810,8 @@ def cli_cosmosdb_create(cmd,
                                     enable_priority_based_execution=enable_priority_based_execution,
                                     default_priority_level=default_priority_level,
                                     enable_prpp_autoscale=enable_prpp_autoscale,
-                                    enable_partition_merge=enable_partition_merge)
+                                    enable_partition_merge=enable_partition_merge,
+                                    capacity_mode=capacity_mode)
 
 
 # pylint: disable=too-many-branches
@@ -875,7 +847,8 @@ def cli_cosmosdb_update(client,
                         enable_priority_based_execution=None,
                         default_priority_level=None,
                         enable_prpp_autoscale=None,
-                        enable_partition_merge=None):
+                        enable_partition_merge=None,
+                        capacity_mode=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
 
@@ -966,7 +939,8 @@ def cli_cosmosdb_update(client,
         enable_priority_based_execution=enable_priority_based_execution,
         default_priority_level=default_priority_level,
         enable_per_region_per_partition_autoscale=enable_prpp_autoscale,
-        enable_partition_merge=enable_partition_merge)
+        enable_partition_merge=enable_partition_merge,
+        capacity_mode=capacity_mode)
 
     async_docdb_update = client.begin_update(resource_group_name, account_name, params)
     docdb_account = async_docdb_update.result()
@@ -1176,7 +1150,8 @@ def _create_database_account(client,
                              default_priority_level=None,
                              enable_prpp_autoscale=None,
                              disable_ttl=None,
-                             enable_partition_merge=None):
+                             enable_partition_merge=None,
+                             capacity_mode=None):
     consistency_policy = None
     if default_consistency_level is not None:
         consistency_policy = ConsistencyPolicy(default_consistency_level=default_consistency_level,
@@ -1315,7 +1290,8 @@ def _create_database_account(client,
         enable_priority_based_execution=enable_priority_based_execution,
         default_priority_level=default_priority_level,
         enable_per_region_per_partition_autoscale=enable_prpp_autoscale,
-        enable_partition_merge=enable_partition_merge
+        enable_partition_merge=enable_partition_merge,
+        capacity_mode=capacity_mode
     )
 
     async_docdb_create = client.begin_create_or_update(resource_group_name, account_name, params)