Merge remote-tracking branch 'upstream/main' into storage-preview-rem…

…ove-account-migration

.azure-pipelines/templates/azdev_setup.yml

@@ -17,6 +17,8 @@ steps:
       pip install azdev
       azdev --version
       azdev setup -c $CLI_REPO_PATH -r $CLI_EXT_REPO_PATH --debug
+      # Installing setuptools with a version higher than 70.0.0 will not generate metadata.json
+      pip install setuptools==70.0.0
       pip list -v
       az --version
     displayName: 'azdev setup'

src/aks-preview/azcli_aks_live_test/scripts/setup_venv.sh

@@ -67,14 +67,6 @@ installAZAKSTOOLFromLocal(){
     pip show az-aks-tool
 }
 
-# need to be executed in a venv
-installAZAKSTOOL(){
-    wheel_file="az_aks_tool-latest-py3-none-any.whl"
-    wheel_url="https://akspreview.blob.core.windows.net/azakstool/${wheel_file}"
-    curl -sLO "${wheel_url}"
-    installAZAKSTOOLFromLocal "${wheel_file}"
-}
-
 # need to be executed in a venv with kusto related modules installed
 removeKustoPTHFile(){
     pushd azEnv/lib/python"${PYTHON_VERSION}"/site-packages
@@ -149,13 +141,8 @@ if [[ -n ${setup_option} ]]; then
         installBuildTools
     elif [[ ${setup_option} == "setup-tool" ]]; then
         echo "Start to setup az-aks-tool!"
-        local_setup=${3:-"n"}
-        if [[ ${local_setup} == "y" ]]; then
-            wheel_file=${4:-"/az_aks_tool-latest-py3-none-any.whl"}
-            installAZAKSTOOLFromLocal "${wheel_file}"
-        else
-            installAZAKSTOOL
-        fi
+        wheel_file=${3:-$(find / -type f -name "az_aks_tool*" | head -n 1)}
+        installAZAKSTOOLFromLocal "${wheel_file}"
         removeKustoPTHFile
     elif [[ ${setup_option} == "setup-az" ]]; then
         echo "Start to setup azure-cli!"

src/aks-preview/azcli_aks_live_test/vsts-azcli-aks-live-test.yaml

@@ -59,7 +59,7 @@ jobs:
       condition: succeeded()
       displayName: "Start Container"
     - bash: |
-        docker exec $(ContainerName) /opt/scripts/setup_venv.sh setup-tool y y /az_aks_tool-latest-py3-none-any.whl
+        docker exec $(ContainerName) /opt/scripts/setup_venv.sh setup-tool y
         docker exec $(ContainerName) /opt/scripts/setup_venv.sh setup-az n azure-cli/ azure-cli-extensions/
       condition: succeeded()
       displayName: "Set up Virtual Environment"

src/aks-preview/azcli_aks_live_test/vsts-azcli-aks-unit-test.yaml

@@ -59,7 +59,7 @@ jobs:
       condition: succeeded()
       displayName: "Start Container"
     - bash: |
-        docker exec $(ContainerName) /opt/scripts/setup_venv.sh setup-tool y y /az_aks_tool-latest-py3-none-any.whl
+        docker exec $(ContainerName) /opt/scripts/setup_venv.sh setup-tool y
         docker exec $(ContainerName) /opt/scripts/setup_venv.sh setup-az n azure-cli/ azure-cli-extensions/
       condition: succeeded()
       displayName: "Set up Virtual Environment"

src/aks-preview/azext_aks_preview/_format.py

@@ -181,7 +181,7 @@ def flatten_version_table(release_info):
 def _custom_functions(preview_versions):
     class CustomFunctions(functions.Functions):  # pylint: disable=too-few-public-methods
 
-        @ functions.signature({'types': ['array']})
+        @functions.signature({'types': ['array']})
         def _func_sort_versions(self, versions):
             """Custom JMESPath `sort_versions` function that sorts an array of strings as software versions"""
             try:
@@ -190,7 +190,7 @@ def _func_sort_versions(self, versions):
             except (TypeError, ValueError):
                 return versions
 
-        @ functions.signature({'types': ['array']})
+        @functions.signature({'types': ['array']})
         def _func_set_preview_array(self, versions):
             """Custom JMESPath `set_preview_array` function that suffixes preview version"""
             try:
@@ -200,7 +200,7 @@ def _func_set_preview_array(self, versions):
             except (TypeError, ValueError):
                 return versions
 
-        @ functions.signature({'types': ['string']})
+        @functions.signature({'types': ['string']})
         def _func_set_preview(self, version):
             """Custom JMESPath `set_preview` function that suffixes preview version"""
             try:
@@ -210,7 +210,7 @@ def _func_set_preview(self, version):
             except (TypeError, ValueError):
                 return version
 
-        @ functions.signature({'types': ['object']})
+        @functions.signature({'types': ['object']})
         def _func_pprint_labels(self, labels):
             """Custom JMESPath `pprint_labels` function that pretty print labels"""
             if not labels:

src/aks-preview/azext_aks_preview/azext_metadata.json

@@ -1,4 +1,5 @@
 {
     "azext.minCliCoreVersion": "2.56.0",
-    "azext.isPreview": true
+    "azext.isPreview": true,
+    "name": "aks-preview"
 }

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.

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

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)

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)