{dls} az dls account: Gen1 account management sdk upgrade (#30094)

* Revert "{Packaging} Vendor `azure-mgmt-datalake-store` (#29275)"

This reverts commit acec6c3.

* Update package versions

* Working tests test_dls_file_access_mgmt and test_dls_file_mgmt

* Tests passing apart from 2 update commands

* Fix update commands

* Fix style violations

* Updating recordings with live test. Disabled network rules tests

* Fix style, linter errors after removing network rules commands. Update recordings

Author: akharit

src/azure-cli-core/azure/cli/core/profiles/_shared.py

@@ -101,7 +101,7 @@ class ResourceType(Enum):  # pylint: disable=too-few-public-methods
     MGMT_CONTAINERINSTANCE = ('azure.mgmt.containerinstance', None)
     MGMT_COSMOSDB = ('azure.mgmt.cosmosdb', None)
     MGMT_DATALAKE_ANALYTICS = ('azure.cli.command_modules.dla.vendored_sdks.azure_mgmt_datalake_analytics', None)
-    MGMT_DATALAKE_STORE = ('azure.cli.command_modules.dls.vendored_sdks.azure_mgmt_datalake_store', None)
+    MGMT_DATALAKE_STORE = ('azure.mgmt.datalake.store', None)
     MGMT_DATAMIGRATION = ('azure.mgmt.datamigration', None)
     MGMT_EVENTGRID = ('azure.mgmt.eventgrid', None)
     MGMT_MAPS = ('azure.mgmt.maps', None)

src/azure-cli/azure/cli/command_modules/dls/_client_factory.py

@@ -6,25 +6,19 @@
 
 def cf_dls_account(cli_ctx, _):
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from .vendored_sdks.azure_mgmt_datalake_store import DataLakeStoreAccountManagementClient
+    from azure.mgmt.datalake.store import DataLakeStoreAccountManagementClient
     return get_mgmt_service_client(cli_ctx, DataLakeStoreAccountManagementClient).accounts
 
 
 def cf_dls_account_firewall(cli_ctx, _):
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from .vendored_sdks.azure_mgmt_datalake_store import DataLakeStoreAccountManagementClient
+    from azure.mgmt.datalake.store import DataLakeStoreAccountManagementClient
     return get_mgmt_service_client(cli_ctx, DataLakeStoreAccountManagementClient).firewall_rules
 
 
-def cf_dls_account_virtual_network(cli_ctx, _):
-    from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from .vendored_sdks.azure_mgmt_datalake_store import DataLakeStoreAccountManagementClient
-    return get_mgmt_service_client(cli_ctx, DataLakeStoreAccountManagementClient).virtual_network_rules
-
-
 def cf_dls_account_trusted_provider(cli_ctx, _):
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
-    from .vendored_sdks.azure_mgmt_datalake_store import DataLakeStoreAccountManagementClient
+    from azure.mgmt.datalake.store import DataLakeStoreAccountManagementClient
     return get_mgmt_service_client(cli_ctx, DataLakeStoreAccountManagementClient).trusted_id_providers
 
 

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

@@ -93,6 +93,16 @@
 helps['dls account firewall update'] = """
 type: command
 short-summary: Updates a firewall rule in a Data Lake Store account.
+parameters:
+  - name: --end-ip-address
+    type: string
+    short-summary: 'The end of the valid ip range for the firewall rule. This can be either ipv4 or ipv6. Start and End should be in the same protocol.'
+  - name: --start-ip-address
+    type: string
+    short-summary: 'The start of the valid ip range for the firewall rule. This can be either ipv4 or ipv6. Start and End should be in the same protocol.'
+  - name: --firewall-rule-name
+    type: string
+    short-summary: 'The name of the firewall rule.'
 """
 
 helps['dls account list'] = """
@@ -104,43 +114,6 @@
     crafted: true
 """
 
-helps['dls account network-rule'] = """
-type: group
-short-summary: Manage Data Lake Store account virtual network rules.
-"""
-
-helps['dls account network-rule create'] = """
-type: command
-short-summary: Creates a virtual network rule in a Data Lake Store account.
-parameters:
-  - name: --subnet
-    type: string
-    short-summary: 'The subnet name or id for the virtual network rule.'
-  - name: --vnet-name
-    type: string
-    short-summary: 'The name of the virtual network rule.'
-"""
-
-helps['dls account network-rule delete'] = """
-type: command
-short-summary: Deletes a virtual network rule in a Data Lake Store account.
-"""
-
-helps['dls account network-rule list'] = """
-type: command
-short-summary: Lists virtual network rules in a Data Lake Store account.
-"""
-
-helps['dls account network-rule show'] = """
-type: command
-short-summary: Get the details of a virtual network rule in a Data Lake Store account.
-"""
-
-helps['dls account network-rule update'] = """
-type: command
-short-summary: Updates a virtual network rule in a Data Lake Store account.
-"""
-
 helps['dls account show'] = """
 type: command
 short-summary: Get the details of a Data Lake Store account.
@@ -155,6 +128,30 @@
 short-summary: Manage Data Lake Store account trusted identity providers.
 """
 
+helps['dls account trusted-provider create'] = """
+type: command
+short-summary:  Creates or updates the specified trusted identity provider.
+parameters:
+  - name: --id-provider
+    type: string
+    short-summary: 'The URL of this trusted identity provider.'
+  - name: --trusted-id-provider-name
+    type: string
+    short-summary: 'The name of the trusted identity provider. This is used for differentiation of providers in the account.'
+"""
+
+helps['dls account trusted-provider update'] = """
+type: command
+short-summary:   Updates the specified trusted identity provider.
+parameters:
+  - name: --id-provider
+    type: string
+    short-summary: 'The URL of this trusted identity provider.'
+  - name: --trusted-id-provider-name
+    type: string
+    short-summary: 'The name of the trusted identity provider. This is used for differentiation of providers in the account.'
+"""
+
 helps['dls account update'] = """
 type: command
 short-summary: Updates a Data Lake Store account.

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

@@ -9,13 +9,13 @@
     tags_type, get_resource_name_completion_list, resource_group_name_type, get_enum_type)
 
 from azure.cli.command_modules.dls._validators import validate_resource_group_name
-from .vendored_sdks.azure_mgmt_datalake_store.models import (
+from azure.mgmt.datalake.store.models import (
     FirewallState,
     TrustedIdProviderState,
     TierType,
     FirewallAllowAzureIpsState)
 
-from .vendored_sdks.azure_mgmt_datalake_store.models import EncryptionConfigType
+from azure.mgmt.datalake.store.models import EncryptionConfigType
 
 
 # pylint: disable=line-too-long, too-many-statements

src/azure-cli/azure/cli/command_modules/dls/_validators.py

@@ -29,7 +29,7 @@ def _get_resource_group_from_account_name(client, account_name):
 
 # COMMAND NAMESPACE VALIDATORS
 def validate_resource_group_name(cmd, ns):
-    from .vendored_sdks.azure_mgmt_datalake_store import DataLakeStoreAccountManagementClient
+    from azure.mgmt.datalake.store import DataLakeStoreAccountManagementClient
     if not ns.resource_group_name:
         try:
             account_name = ns.name

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

@@ -8,36 +8,27 @@
 from azure.cli.command_modules.dls._client_factory import (
     cf_dls_account,
     cf_dls_account_firewall,
-    cf_dls_account_virtual_network,
     cf_dls_account_trusted_provider)
 
 
 # pylint: disable=too-many-statements
 def load_command_table(self, _):
-    from ._validators import (
-        validate_subnet
-    )
-    adls_format_path = 'azure.cli.command_modules.dls.vendored_sdks.azure_mgmt_datalake_store.operations.{}#{}.{{}}'
+    adls_format_path = 'azure.mgmt.datalake.store.operations#{}.{{}}'
 
     dls_custom = CliCommandType(operations_tmpl='azure.cli.command_modules.dls.custom#{}')
 
     dls_account_sdk = CliCommandType(
-        operations_tmpl=adls_format_path.format('accounts_operations', 'AccountsOperations'),
+        operations_tmpl=adls_format_path.format('AccountsOperations'),
         client_factory=cf_dls_account
     )
 
     dls_firewall_sdk = CliCommandType(
-        operations_tmpl=adls_format_path.format('firewall_rules_operations', 'FirewallRulesOperations'),
-        client_factory=cf_dls_account
-    )
-
-    dls_virtual_network_sdk = CliCommandType(
-        operations_tmpl=adls_format_path.format('virtual_network_rules_operations', 'VirtualNetworkRulesOperations'),
+        operations_tmpl=adls_format_path.format('FirewallRulesOperations'),
         client_factory=cf_dls_account
     )
 
     dls_provider_sdk = CliCommandType(
-        operations_tmpl=adls_format_path.format('trusted_id_providers_operations', 'TrustedIdProvidersOperations'),
+        operations_tmpl=adls_format_path.format('TrustedIdProvidersOperations'),
         client_factory=cf_dls_account_trusted_provider
     )
 
@@ -46,32 +37,22 @@ def load_command_table(self, _):
         g.custom_command('create', 'create_adls_account')
         g.custom_command('update', 'update_adls_account')
         g.custom_command('list', 'list_adls_account')
-        g.command('delete', 'delete')
+        g.command('delete', 'begin_delete')
         g.show_command('show', 'get')
         g.command('enable-key-vault', 'enable_key_vault')
 
     # account firewall operations
     with self.command_group('dls account firewall', dls_firewall_sdk, client_factory=cf_dls_account_firewall) as g:
         g.custom_command('create', 'add_adls_firewall_rule')
-        g.command('update', 'update')
-        g.command('list', 'list_by_account')
-        g.show_command('show', 'get')
-        g.command('delete', 'delete')
-
-    # account virtual network rule operations
-    with self.command_group('dls account network-rule',
-                            dls_virtual_network_sdk,
-                            client_factory=cf_dls_account_virtual_network) as g:
-        g.custom_command('create', 'add_adls_virtual_network_rule', validator=validate_subnet)
-        g.generic_update_command('update')
+        g.custom_command('update', 'update_adls_firewall_rule')
         g.command('list', 'list_by_account')
         g.show_command('show', 'get')
         g.command('delete', 'delete')
 
     # account trusted id provider operations
-    with self.command_group('dls account trusted-provider', dls_provider_sdk) as g:
-        g.command('create', 'create_or_update')
-        g.command('update', 'update')
+    with self.command_group('dls account trusted-provider', dls_provider_sdk, client_factory=cf_dls_account_trusted_provider) as g:
+        g.custom_command('create', 'add_trusted_provider_rule')
+        g.custom_command('update', 'update_trusted_provider_rule')
         g.command('list', 'list_by_account')
         g.show_command('show', 'get')
         g.command('delete', 'delete')

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

@@ -6,7 +6,7 @@
 from knack.log import get_logger
 from knack.util import CLIError
 
-from .vendored_sdks.azure_mgmt_datalake_store.models import (
+from azure.mgmt.datalake.store.models import (
     UpdateDataLakeStoreAccountParameters,
     CreateDataLakeStoreAccountParameters,
     EncryptionConfigType,
@@ -22,7 +22,11 @@
 from azure.cli.command_modules.dls._client_factory import (cf_dls_filesystem)
 from azure.cli.core.commands.client_factory import get_mgmt_service_client
 from azure.cli.core.profiles import ResourceType
-
+from azure.mgmt.datalake.store.models import (
+    CreateOrUpdateFirewallRuleParameters,
+    UpdateFirewallRuleParameters,
+    CreateOrUpdateTrustedIdProviderParameters,
+    UpdateTrustedIdProviderParameters)
 
 logger = get_logger(__name__)
 
@@ -78,8 +82,7 @@ def create_adls_account(cmd, client, resource_group_name, account_name, location
         create_params.encryption_state = EncryptionState.disabled
         create_params.identity = None
         create_params.encryption_config = None
-
-    return client.create(resource_group_name, account_name, create_params).result()
+    return client.begin_create(resource_group_name, account_name, create_params)
 
 
 def update_adls_account(client, account_name, resource_group_name, tags=None, default_group=None, firewall_state=None,
@@ -97,7 +100,7 @@ def update_adls_account(client, account_name, resource_group_name, tags=None, de
         update_params.encryption_config = UpdateEncryptionConfig(
             key_vault_meta_info=UpdateKeyVaultMetaInfo(encryption_key_version=key_version))
 
-    return client.update(resource_group_name, account_name, update_params).result()
+    return client.begin_update(resource_group_name, account_name, update_params)
 # endregion
 
 
@@ -111,21 +114,49 @@ def add_adls_firewall_rule(client,
     return client.create_or_update(resource_group_name,
                                    account_name,
                                    firewall_rule_name,
-                                   start_ip_address,
-                                   end_ip_address)
+                                   CreateOrUpdateFirewallRuleParameters(
+                                       start_ip_address=start_ip_address,
+                                       end_ip_address=end_ip_address))
+
+
+def update_adls_firewall_rule(client,
+                              account_name,
+                              firewall_rule_name,
+                              resource_group_name,
+                              start_ip_address=None,
+                              end_ip_address=None):
+    return client.update(resource_group_name,
+                         account_name,
+                         firewall_rule_name,
+                         UpdateFirewallRuleParameters(
+                             start_ip_address=start_ip_address,
+                             end_ip_address=end_ip_address))
+
 # endregion
 
 
-# region virtual network
-def add_adls_virtual_network_rule(client,
-                                  account_name,
-                                  virtual_network_rule_name,
-                                  subnet,
-                                  resource_group_name):
-    return client.create_or_update(resource_group_name,
-                                   account_name,
-                                   virtual_network_rule_name,
-                                   subnet)
+# region trusted provider
+def add_trusted_provider_rule(client,
+                              account_name,
+                              trusted_id_provider_name,
+                              id_provider,
+                              resource_group_name):
+    return client.create_or_update(resource_group_name=resource_group_name,
+                                   account_name=account_name,
+                                   trusted_id_provider_name=trusted_id_provider_name,
+                                   parameters=CreateOrUpdateTrustedIdProviderParameters(id_provider=id_provider))
+
+
+def update_trusted_provider_rule(client,
+                                 account_name,
+                                 resource_group_name,
+                                 trusted_id_provider_name,
+                                 id_provider=None):
+    return client.update(resource_group_name=resource_group_name,
+                         account_name=account_name,
+                         trusted_id_provider_name=trusted_id_provider_name,
+                         parameters=UpdateTrustedIdProviderParameters(
+                             id_provider=id_provider))
 # endregion
 
 
@@ -216,16 +247,10 @@ def test_adls_item(cmd, account_name, path):
 def preview_adls_item(cmd, account_name, path, length=None, offset=0, force=False):
     client = cf_dls_filesystem(cmd.cli_ctx, account_name)
     if length:
-        try:
-            length = long(length)
-        except NameError:
-            length = int(length)
+        length = int(length)
 
     if offset:
-        try:
-            offset = long(offset)
-        except NameError:
-            offset = int(offset)
+        offset = int(offset)
 
     if not length or length <= 0:
         length = client.info(path)['length'] - offset
@@ -258,10 +283,7 @@ def set_adls_item_expiry(cmd, account_name, path, expiration_time):
         raise CLIError('The specified path does not exist or is not a file. Please ensure the path points to a file and it exists. Path supplied: {}'.format(path))
 
     expiration_time = float(expiration_time)
-    try:
-        expiration_time = long(expiration_time)
-    except NameError:
-        expiration_time = int(expiration_time)
+    expiration_time = int(expiration_time)
     client.set_expiry(path, ExpiryOptionType.absolute.value, expiration_time)