[Batch] Upgrade api-version to v2022-06-01, adding 'batch account network-profile' and 'batch account network-profile network-rule'

src/azure-cli/azure/cli/command_modules/batch/_help.py

@@ -74,6 +74,103 @@
 long-summary: List the endpoints that a Batch Compute Node under this Batch Account may call as part of Batch service administration. If you are deploying a Pool inside of a virtual network that you specify, you must make sure your network allows outbound access to these endpoints. Failure to allow access to these endpoints may cause Batch to mark the affected nodes as unusable. For more information about creating a pool inside of a virtual network, see https://docs.microsoft.com/azure/batch/batch-virtual-network."
 """
 
+helps['batch account identity'] = """
+type: group
+short-summary: Manage identities of a batch account.
+"""
+
+helps['batch account identity assign'] = """
+type: command
+short-summary: Add managed identities to an existing batch account.
+examples:
+  - name: Add a system assigned managed identity to an existing batch account.
+    text: >
+        az batch account identity assign --name MyBatchAccount --resource-group MyResourceGroup --system-assigned
+  - name: Add a user assigned managed identity to an existing batch account.
+    text: >
+        az batch account identity assign --name MyBatchAccount --resource-group MyResourceGroup --user-assigned MyAssignedId
+"""
+
+helps['batch account identity remove'] = """
+type: command
+short-summary: Remove managed identities from an existing batch account.
+examples:
+  - name: Remove a system assigned managed identity from an existing batch account.
+    text: >
+        az batch account identity remove --name MyBatchAccount --resource-group MyResourceGroup --system-assigned
+  - name: Remove a user assigned managed identity from an existing batch account.
+    text: >
+        az batch account identity remove --name MyBatchAccount --resource-group MyResourceGroup --user-assigned MyAssignedId
+  - name: Remove all user assigned managed identities from an existing batch account.
+    text: >
+        az batch account identity remove --name MyBatchAccount --resource-group MyResourceGroup --user-assigned
+"""
+
+helps['batch account identity show'] = """
+type: command
+short-summary: Display managed identities of a batch account.
+examples:
+  - name: Display managed identities of a batch account.
+    text: |
+        az batch account identity show --name MyBatchAccount --resource-group MyResourceGroup
+"""
+
+helps['batch account network-profile'] = """
+type: group
+short-summary: Manage Batch account Network profiles.
+"""
+
+helps['batch account network-profile show'] = """
+type: command
+short-summary: Get information about the Network profile for Batch account.
+examples:
+  - name: Show the network-profile for both BatchAccount and NodeManagement
+    text: >
+        az batch account network-profile show -g MyResourceGroup -n MyBatchAccount
+"""
+
+helps['batch account network-profile set'] = """
+type: command
+short-summary: Set the Network profile for Batch account.
+examples:
+  - name: Set the BatchAccount network-profile to the Allow
+    text: >
+        az batch account network-profile set -g MyResourceGroup -n MyBatchAccount --profile BatchAccount --default-action Allow
+"""
+
+helps['batch account network-profile network-rule'] = """
+type: group
+short-summary: Manage Batch account Network rules in Network Profile.
+"""
+
+helps['batch account network-profile network-rule list'] = """
+type: command
+short-summary: List the Network rules from a Network Profile.
+examples:
+  - name: List the Batch Accounts network profile
+    text: >
+        az batch account network-profile network-rule list -g MyResourceGroup -n MyBatchAccount
+"""
+
+helps['batch account network-profile network-rule add'] = """
+type: command
+short-summary: Add a Network rule from a Network Profile.
+examples:
+  - name: Add ip address to BatchAccount network rule
+    text: >
+        az batch account network-profile network-rule add -g MyResourceGroup -n MyBatchAccount --profile BatchAccount --ip-address 1.2.3.4
+"""
+
+helps['batch account network-profile network-rule delete'] = """
+type: command
+short-summary: Delete a Network rule from a Network Profile.
+examples:
+  - name: Delete ip address from BatchAccount network rule
+    text: >
+        az batch account network-profile network-rule delete -g MyResourceGroup -n MyBatchAccount --profile BatchAccount --ip-address 1.2.3.4
+"""
+
+
 helps['batch application'] = """
 type: group
 short-summary: Manage Batch applications.

src/azure-cli/azure/cli/command_modules/batch/_params.py

@@ -3,6 +3,7 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
+from enum import Enum
 from argcomplete.completers import FilesCompleter
 
 from knack.arguments import CLIArgumentType
@@ -11,7 +12,8 @@
     AccountKeyType,
     KeySource,
     PublicNetworkAccessType,
-    ResourceIdentityType)
+    ResourceIdentityType,
+    EndpointAccessDefaultAction)
 from azure.batch.models import ComputeNodeDeallocationOption
 
 from azure.cli.core.commands.parameters import (
@@ -41,6 +43,16 @@
     validate_pool_resize_parameters)
 
 
+class NetworkProfile(Enum):
+    BatchAccount = "BatchAccount"
+    NodeManagement = "NodeManagement"
+
+
+profile_param_type = CLIArgumentType(
+    help="Network profile to set.",
+    arg_type=get_enum_type(NetworkProfile))
+
+
 # pylint: disable=line-too-long, too-many-statements
 def load_arguments(self, _):
     batch_name_type = CLIArgumentType(
@@ -71,15 +83,70 @@ def load_arguments(self, _):
         c.argument('encryption_key_source', help='Part of the encryption configuration for the Batch account. Type of the key source. Can be either Microsoft.Batch or Microsoft.KeyVault', arg_type=get_enum_type(KeySource))
         c.argument('encryption_key_identifier', help='Part of the encryption configuration for the Batch account. '
                                                      'Full path to the versioned secret. Example https://mykeyvault.vault.azure.net/keys/testkey/6e34a81fef704045975661e297a4c053.')
-        c.argument('identity_type', help="The type of identity used for the Batch account. Possible values include: 'SystemAssigned', 'None'.", arg_type=get_enum_type(ResourceIdentityType))
+        c.argument('identity_type', help="The type of identity used for the Batch account. Possible values include: 'SystemAssigned', 'None'.", arg_type=get_enum_type(ResourceIdentityType), deprecate_info=c.deprecate(hide=True))
+        c.argument('mi_user_assigned', help='Resource ID of the user assigned identity for the batch services account.', arg_group='Identity')
+        c.argument('mi_system_assigned', help='Set the system managed identity on the batch services account.', arg_group='Identity')
         c.ignore('keyvault_url')
 
     with self.argument_context('batch account set') as c:
         c.argument('tags', tags_type)
         c.argument('storage_account', help='The storage account name or resource ID to be used for auto storage.', validator=storage_account_id)
         c.argument('encryption_key_source', help='Part of the encryption configuration for the Batch account. Type of the key source. Can be either Microsoft.Batch or Microsoft.KeyVault')
+        c.argument('public_network_access', help="The network access type for accessing Azure Batch account. Values can either be enabled or disabled.", arg_type=get_enum_type(PublicNetworkAccessType))
         c.argument('encryption_key_identifier', help='Part of the encryption configuration for the Batch account. Full path to the versioned secret. Example https://mykeyvault.vault.azure.net/keys/testkey/6e34a81fef704045975661e297a4c053.')
-        c.argument('identity_type', help="The type of identity used for the Batch account. Possible values include: 'SystemAssigned', 'None'.", arg_type=get_enum_type(ResourceIdentityType))
+        c.argument('identity_type', help="The type of identity used for the Batch account. Possible values include: 'SystemAssigned', 'None'.", arg_type=get_enum_type(ResourceIdentityType), deprecate_info=c.deprecate(hide=True))
+        c.argument('mi_system_assigned', help='Set the system managed identity on the batch services account.', arg_group='Identity', deprecate_info=c.deprecate(hide=True))
+        c.argument('mi_user_assigned', help='Resource ID of the user assigned identity for the batch services account.', arg_group='Identity', deprecate_info=c.deprecate(hide=True))
+
+    with self.argument_context('batch account identity assign') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+        c.argument('mi_system_assigned', options_list=['--system-assigned'], action='store_true',
+                   arg_group='Managed Identity', help='Provide this flag to use system assigned identity for batch accounts. '
+                   'Check out help for more examples')
+        c.argument('mi_user_assigned', options_list=['--user-assigned'],
+                   arg_group='Managed Identity', help='User Assigned Identity ids to be used for batch account. '
+                   'Check out help for more examples')
+
+    with self.argument_context('batch account identity remove') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+        c.argument('mi_system_assigned', options_list=['--system-assigned'], action='store_true',
+                   arg_group='Managed Identity', help='Provide this flag to use system assigned identity for batch accounts. '
+                   'Check out help for more examples')
+        c.argument('mi_user_assigned', options_list=['--user-assigned'], nargs='*',
+                   arg_group='Managed Identity', help='User Assigned Identity ids to be used for batch account. '
+                   'Check out help for more examples')
+
+    with self.argument_context('batch account identity show') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+
+    with self.argument_context('batch account network-profile show') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+
+    with self.argument_context('batch account network-profile set') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+        c.argument('profile', arg_type=profile_param_type)
+        c.argument('default_action', help="Default action for endpoint access. It is only applicable when publicNetworkAccess is enabled. Possible values include: 'Allow', 'Deny'", arg_type=get_enum_type(EndpointAccessDefaultAction))
+
+    with self.argument_context('batch account network-profile network-rule list') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+
+    with self.argument_context('batch account network-profile network-rule add') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+        c.argument('profile', arg_type=profile_param_type)
+        c.argument('ip_address', help='IPv4 address or CIDR range.')
+
+    with self.argument_context('batch account network-profile network-rule delete') as c:
+        c.argument('resource_group_name', resource_group_name_type, help='Name of the resource group. If not specified will display currently set account.')
+        c.argument('account_name', batch_name_type, options_list=('--name', '-n'), help='Name of the batch account to show. If not specified will display currently set account.')
+        c.argument('profile', arg_type=profile_param_type)
+        c.argument('ip_address', help='IPv4 address or CIDR range.')
 
     with self.argument_context('batch account keys renew') as c:
         c.argument('resource_group_name', resource_group_name_type,

src/azure-cli/azure/cli/command_modules/batch/commands.py

@@ -62,6 +62,20 @@ def get_data_factory(name):
         g.custom_command('keys renew', 'renew_accounts_keys', table_transformer=account_keys_renew_table_format)
         g.command('outbound-endpoints', 'list_outbound_network_dependencies_endpoints')
 
+    with self.command_group('batch account identity', get_mgmt_type('batch_account'), client_factory=get_mgmt_factory('batch_account')) as g:
+        g.custom_command('assign', 'assign_batch_identity')
+        g.custom_command('remove', 'remove_batch_identity', confirmation=True)
+        g.custom_show_command('show', 'show_batch_identity')
+
+    with self.command_group('batch account network-profile', get_mgmt_type('batch_account'), client_factory=get_mgmt_factory('batch_account')) as g:
+        g.custom_show_command('show', 'get_network_profile')
+        g.custom_command('set', 'update_network_profile')
+
+    with self.command_group('batch account network-profile network-rule', get_mgmt_type('batch_account'), client_factory=get_mgmt_factory('batch_account')) as g:
+        g.custom_show_command('list', 'list_network_rules')
+        g.custom_command('add', 'add_network_rule')
+        g.custom_command('delete', 'delete_network_rule', confirmation=True)
+
     with self.command_group('batch application', get_mgmt_type('application'), client_factory=get_mgmt_factory('application')) as g:
         g.command('list', 'list', table_transformer=application_list_table_format)
         g.show_command('show', 'get')