[Batch] Update to API Version 2020-05-01

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

@@ -7,7 +7,9 @@
 
 from knack.arguments import CLIArgumentType
 
-from azure.mgmt.batch.models import AccountKeyType
+from azure.mgmt.batch.models import \
+    AccountKeyType, KeySource,\
+    PublicNetworkAccessType, ResourceIdentityType
 from azure.batch.models import ComputeNodeDeallocationOption
 
 from azure.cli.core.commands.parameters import \
@@ -51,11 +53,19 @@ def load_arguments(self, _):
         c.argument('tags', tags_type, help="Space-separated tags in 'key[=value]' format.")
         c.argument('storage_account', help='The storage account name or resource ID to be used for auto storage.', validator=storage_account_id)
         c.argument('keyvault', help='The KeyVault name or resource ID to be used for an account with a pool allocation mode of \'User Subscription\'.', validator=keyvault_id)
+        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_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.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('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))
 
     with self.argument_context('batch account keys renew') as c:
         c.argument('key_name', arg_type=get_enum_type(AccountKeyType))

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

@@ -12,15 +12,17 @@
 from msrest.exceptions import DeserializationError
 
 from azure.mgmt.batch import BatchManagementClient
-from azure.mgmt.batch.models import (BatchAccountCreateParameters,
+from azure.mgmt.batch.models import (BatchAccountCreateParameters, BatchAccountUpdateParameters,
                                      AutoStorageBaseProperties,
-                                     Application)
+                                     Application, EncryptionProperties,
+                                     KeyVaultProperties, BatchAccountIdentity)
 from azure.mgmt.batch.operations import (ApplicationPackageOperations)
 
 from azure.batch.models import (CertificateAddParameter, PoolStopResizeOptions, PoolResizeParameter,
                                 PoolResizeOptions, JobListOptions, JobListFromJobScheduleOptions,
                                 TaskAddParameter, TaskAddCollectionParameter, TaskConstraints,
-                                PoolUpdatePropertiesParameter, StartTask, AffinityInformation)
+                                PoolUpdatePropertiesParameter, StartTask, AffinityInformation,
+                                )
 
 from azure.cli.core.commands.client_factory import get_mgmt_service_client
 from azure.cli.core.profiles import get_sdk, ResourceType
@@ -73,12 +75,25 @@ def get_account(cmd, client, resource_group_name=None, account_name=None):
 @transfer_doc(AutoStorageBaseProperties)
 def create_account(client,
                    resource_group_name, account_name, location, tags=None, storage_account=None,
-                   keyvault=None, keyvault_url=None, no_wait=False):
+                   keyvault=None, keyvault_url=None, no_wait=False, public_network_access=None,
+                   encryption_key_source=None, encryption_key_identifier=None, identity_type=None):
     properties = AutoStorageBaseProperties(storage_account_id=storage_account) \
         if storage_account else None
+    identity = BatchAccountIdentity(type=identity_type) if identity_type else None
+    if encryption_key_source == "Microsoft.KeyVault" and not encryption_key_identifier:
+        raise ValueError("The --encryption-key-identifier property is required when "
+                         "--encryption-key-source is set to Microsoft.KeyVault")
+    encryption_key_identifier = KeyVaultProperties(key_identifier=encryption_key_identifier) \
+        if encryption_key_identifier else None
+    encryption = EncryptionProperties(
+        key_source=encryption_key_source,
+        encryption_key_identifier=encryption_key_identifier) if encryption_key_source else None
     parameters = BatchAccountCreateParameters(location=location,
                                               tags=tags,
-                                              auto_storage=properties)
+                                              auto_storage=properties,
+                                              public_network_access=public_network_access,
+                                              encryption=encryption,
+                                              identity=identity)
     if keyvault:
         parameters.key_vault_reference = {'id': keyvault, 'url': keyvault_url}
         parameters.pool_allocation_mode = 'UserSubscription'
@@ -89,13 +104,27 @@ def create_account(client,
 
 @transfer_doc(AutoStorageBaseProperties)
 def update_account(client, resource_group_name, account_name,
-                   tags=None, storage_account=None):
+                   tags=None, storage_account=None, encryption_key_source=None,
+                   encryption_key_identifier=None, identity_type=None):
     properties = AutoStorageBaseProperties(storage_account_id=storage_account) \
         if storage_account else None
+    if encryption_key_source == "Microsoft.KeyVault" and not encryption_key_identifier:
+        raise ValueError("The --encryption-key-identifier property is required when "
+                         "--encryption-key-source is set to Microsoft.KeyVault")
+    encryption_key_identifier = KeyVaultProperties(key_identifier=encryption_key_identifier) \
+        if encryption_key_identifier else None
+    encryption = EncryptionProperties(
+        key_source=encryption_key_source,
+        encryption_key_identifier=encryption_key_identifier) if encryption_key_source else None
+    identity = BatchAccountIdentity(type=identity_type) if identity_type else None
+    parameters = BatchAccountUpdateParameters(
+        tags=tags,
+        encryption=encryption,
+        identity=identity,
+        auto_storage=properties)
     return client.update(resource_group_name=resource_group_name,
                          account_name=account_name,
-                         tags=tags,
-                         auto_storage=properties)
+                         parameters=parameters)
 
 
 # pylint: disable=inconsistent-return-statements