[CosmosDB] Upgrade API Version 2024-05-15-preview (#7613)

* Added Vendored SDK * Made changes to history.rst and setup.py * adding capacityMode option for CLI commands * adding test recordings * Corrected vendored sdk * Add azure_mgmt_mongocluster vendor sdk. Update commands to use new mongo cluster sdk and update tests. * Distributed Query changes * Fixes and test recordings * Specification bugfix which required re-release of SDK * Added test recordings * Restoring test recordings for CI * Test recordings using new subscription * Changes to tests to make them generic * Pylint * More pylint --------- Co-authored-by: chandrasekhar gunturi <[email protected]> Co-authored-by: Oliver Towers <[email protected]>
Azure · Jun 26, 2024 · 8097e12 · 8097e12
1 parent f3ce427
commit 8097e12
Show file tree

Hide file tree

Showing 203 changed files with 101,317 additions and 138,532 deletions.
diff --git a/src/cosmosdb-preview/HISTORY.rst b/src/cosmosdb-preview/HISTORY.rst
@@ -2,6 +2,12 @@
 
 Release History
 ===============
+
+1.0.1
+* Add support for CapacityMode in Global Database Accounts
+* Add support for DistributedQuery flag in SqlDedicatedGateway
+
++++++++
 1.0.0
 * Add support for Per-Region Per-Partition Autoscale. '--enable-prpp-autoscale' parameter can be used during account create/update.
 * Add support for Restore with Time-To-Live Disabled. '--disable-ttl' parameter can be used during restore.

diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_client_factory.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_client_factory.py
@@ -10,6 +10,12 @@ def cf_cosmosdb_preview(cli_ctx, *_):
     return get_mgmt_service_client(cli_ctx, CosmosDBManagementClient)
 
 
+def cf_mongocluster_preview(cli_ctx, *_):
+    from azure.cli.core.commands.client_factory import get_mgmt_service_client
+    from azext_cosmosdb_preview.vendored_sdks.azure_mgmt_mongocluster import MongoClusterMgmtClient
+    return get_mgmt_service_client(cli_ctx, MongoClusterMgmtClient)
+
+
 def cf_service(cli_ctx, _):
     return cf_cosmosdb_preview(cli_ctx).service
 
@@ -82,5 +88,9 @@ def cf_data_transfer_job(cli_ctx, _):
     return cf_cosmosdb_preview(cli_ctx).data_transfer_jobs
 
 
-def cf_mongo_cluster_job(cli_ctx, _):
-    return cf_cosmosdb_preview(cli_ctx).mongo_clusters
+def cf_mongo_clusters(cli_ctx, _):
+    return cf_mongocluster_preview(cli_ctx).mongo_clusters
+
+
+def cf_mongo_cluster_firewall_rules(cli_ctx, _):
+    return cf_mongocluster_preview(cli_ctx).firewall_rules
diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py b/src/cosmosdb-preview/azext_cosmosdb_preview/_params.py
@@ -29,17 +29,22 @@
     CreatePhysicalPartitionIdListAction)
 
 from azext_cosmosdb_preview.vendored_sdks.azure_mgmt_cosmosdb.models import (
-    ContinuousTier, DefaultPriorityLevel
-)
+    DefaultConsistencyLevel,
+    DatabaseAccountKind,
+    ServerVersion,
+    NetworkAclBypass,
+    BackupPolicyType,
+    AnalyticalStorageSchemaType,
+    BackupStorageRedundancy,
+    CapacityMode,
+    ContinuousTier,
+    DefaultPriorityLevel)
 
 from azure.cli.core.util import shell_safe_json_parse
 
 from azure.cli.core.commands.parameters import (
     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)
-
 from azure.cli.command_modules.cosmosdb.actions import (
     CreateLocation, CreateDatabaseRestoreResource, UtcDatetimeAction)
 
@@ -275,6 +280,7 @@ def load_arguments(self, _):
         c.argument('service_name', options_list=['--name', '-n'], help="Service Name.")
         c.argument('instance_count', options_list=['--count', '-c'], help="Instance Count.")
         c.argument('instance_size', options_list=['--size'], help="Instance Size. Possible values are: Cosmos.D4s, Cosmos.D8s, Cosmos.D16s etc")
+        c.argument('dedicated_gateway_type', options_list=['--gateway-type'], arg_type=get_enum_type(['IntegratedCache', 'DistributedQuery']), help="Dedicated Gateway Type. Valid only for SqlDedicatedGateway service kind")
 
     with self.argument_context('cosmosdb service create') as c:
         c.argument('instance_size', options_list=['--size'], help="Instance Size. Possible values are: Cosmos.D4s, Cosmos.D8s, Cosmos.D16s etc")
@@ -307,7 +313,7 @@ def load_arguments(self, _):
         c.argument('databases_to_restore', nargs='+', action=CreateDatabaseRestoreResource, is_preview=True, arg_group='Restore')
         c.argument('gremlin_databases_to_restore', nargs='+', action=CreateGremlinDatabaseRestoreResource, is_preview=True, arg_group='Restore')
         c.argument('tables_to_restore', nargs='+', action=CreateTableRestoreResource, is_preview=True, arg_group='Restore')
-        c.argument('enable_partition_merge', arg_type=get_three_state_flag(), help="Flag to enable partition merge on the account.")
+        c.argument('enable_partition_merge', arg_type=get_three_state_flag(), is_preview=True, help="Flag to enable partition merge on the account.")
 
     for scope in ['cosmosdb create', 'cosmosdb update']:
         with self.argument_context(scope) as c:
@@ -343,6 +349,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)

diff --git a/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py b/src/cosmosdb-preview/azext_cosmosdb_preview/commands.py
@@ -24,7 +24,8 @@
     cf_restorable_table_resources,
     cf_restorable_database_accounts,
     cf_data_transfer_job,
-    cf_mongo_cluster_job
+    cf_mongo_clusters,
+    cf_mongo_cluster_firewall_rules
 )
 
 
@@ -276,18 +277,23 @@ def load_command_table(self, _):
 
     # Mongo cluster operations
     cosmosdb_mongocluster_sdk = CliCommandType(
-        operations_tmpl='azext_cosmosdb_preview.vendored_sdks.azure_mgmt_cosmosdb.operations.#MongoClustersOperations.{}',
-        client_factory=cf_mongo_cluster_job)
+        operations_tmpl='azext_cosmosdb_preview.vendored_sdks.azure_mgmt_mongocluster.operations.#MongoClustersOperations.{}',
+        client_factory=cf_mongo_clusters)
+
+    # Mongo cluster firewall rule operations
+    cosmosdb_mongocluster_firewall_rule_sdk = CliCommandType(
+        operations_tmpl='azext_cosmosdb_preview.vendored_sdks.azure_mgmt_mongocluster.operations.#FirewallRulesOperations.{}',
+        client_factory=cf_mongo_cluster_firewall_rules)
 
     # Mongo Cluster create operations
-    with self.command_group('cosmosdb mongocluster', cosmosdb_mongocluster_sdk, client_factory=cf_mongo_cluster_job, is_preview=True) as g:
+    with self.command_group('cosmosdb mongocluster', cosmosdb_mongocluster_sdk, client_factory=cf_mongo_clusters, is_preview=True) as g:
         g.custom_command('create', 'cli_cosmosdb_mongocluster_create', is_preview=True)
         g.custom_command('update', 'cli_cosmosdb_mongocluster_update', is_preview=True)
         g.custom_command('list', 'cli_cosmosdb_mongocluster_list', is_preview=True)
         g.custom_show_command('show', 'cli_cosmosdb_mongocluster_get', is_preview=True)
         g.custom_command('delete', 'cli_cosmosdb_mongocluster_delete', confirmation=True)
 
-    with self.command_group('cosmosdb mongocluster firewall rule', cosmosdb_mongocluster_sdk, client_factory=cf_mongo_cluster_job, is_preview=True) as g:
+    with self.command_group('cosmosdb mongocluster firewall rule', cosmosdb_mongocluster_firewall_rule_sdk, client_factory=cf_mongo_cluster_firewall_rules, is_preview=True) as g:
         g.custom_command('create', 'cli_cosmosdb_mongocluster_firewall_rule_create', is_preview=True)
         g.custom_command('update', 'cli_cosmosdb_mongocluster_firewall_rule_update', is_preview=True)
         g.custom_command('list', 'cli_cosmosdb_mongocluster_firewall_rule_list', is_preview=True)