Fixed a bug to continue visting NextPageLink when listing key vaults from ARM API (Azure#17618)

* Fixed a bug to continue visting NextPageLink when listing key vaults from ARM API

* refactor code

* minor fix

Author: BethanyZhou

src/KeyVault/KeyVault/ChangeLog.md

@@ -18,7 +18,8 @@
         - Additional information about change #1
 -->
 ## Upcoming Release
-
+* Fixed a bug to continue visting `NextPageLink` when listing key vaults from ARM API
+ 
 ## Version 4.3.0
 * `New-AzKeyVaultManagedHsm`: supported specifying how long a deleted managed hsm is retained by `SoftDeleteRetentionInDays` and enabling purge protection by `EnablePurgeProtection`
 * `Update-AzKeyVaultManagedHsm`: supported enabling purge protection by `EnablePurgeProtection`

src/KeyVault/KeyVault/Commands/GetAzureKeyVault.cs

@@ -14,6 +14,8 @@
 
 using System;
 using System.Collections;
+using System.Collections.Generic;
+using System.Linq;
 using System.Management.Automation;
 using Microsoft.Azure.Commands.KeyVault.Helpers;
 using Microsoft.Azure.Commands.KeyVault.Models;
@@ -103,7 +105,13 @@ public override void ExecuteCmdlet()
             switch (ParameterSetName)
             {
                 case GetVaultParameterSet:
-                    ResourceGroupName = string.IsNullOrWhiteSpace(ResourceGroupName) ? GetResourceGroupName(VaultName) : ResourceGroupName;
+                    List<PSKeyVaultIdentityItem> vaults = null;
+                    
+                    if (string.IsNullOrWhiteSpace(ResourceGroupName))
+                    {
+                        vaults = ListVaults(ResourceGroupName, Tag);
+                        ResourceGroupName = vaults?.FirstOrDefault(r => r.VaultName.Equals(VaultName, StringComparison.OrdinalIgnoreCase))?.ResourceGroupName;
+                    }
 
                     if (ShouldGetByName(ResourceGroupName, VaultName))
                     {
@@ -115,7 +123,7 @@ public override void ExecuteCmdlet()
                     }
                     else
                     {
-                        WriteObject(TopLevelWildcardFilter(ResourceGroupName, VaultName, ListVaults(ResourceGroupName, Tag)), true);
+                        WriteObject(TopLevelWildcardFilter(ResourceGroupName, VaultName, vaults ?? ListVaults(ResourceGroupName, Tag)), true);
                     }
 
                     break;

src/KeyVault/KeyVault/Models/KeyVaultManagementCmdletBase.cs

@@ -135,21 +135,12 @@ protected List<PSKeyVaultIdentityItem> ListVaults(string resourceGroupName, Hash
                 return vaults;
             }
 
-            IEnumerable<PSKeyVaultIdentityItem> listResult;
             var resourceType = resourceTypeName.Equals(ResourceTypeName.Hsm) ?
                 KeyVaultManagementClient.ManagedHsmResourceType : KeyVaultManagementClient.VaultsResourceType;
-            if (ShouldListByResourceGroup(resourceGroupName, null))
-            {
-                listResult = ListByResourceGroup(resourceGroupName,
-                    new Rest.Azure.OData.ODataQuery<GenericResourceFilter>(
-                        r => r.ResourceType == resourceType));
-            }
-            else
-            {
-                listResult = ListResources(
-                    new Rest.Azure.OData.ODataQuery<GenericResourceFilter>(
-                        r => r.ResourceType == resourceType));
-            }
+
+            IEnumerable<PSKeyVaultIdentityItem> listResult = ListPagable(resourceGroupName,
+                new Rest.Azure.OData.ODataQuery<GenericResourceFilter>(
+                    r => r.ResourceType == resourceType));
 
             if (listResult != null)
             {
@@ -161,22 +152,40 @@ protected List<PSKeyVaultIdentityItem> ListVaults(string resourceGroupName, Hash
             return vaults;
         }
 
-        public virtual IEnumerable<PSKeyVaultIdentityItem> ListResources(Rest.Azure.OData.ODataQuery<GenericResourceFilter> filter = null, ulong first = ulong.MaxValue, ulong skip = ulong.MinValue)
+        IEnumerable<PSKeyVaultIdentityItem> ListPagable(string resourceGroupName, Rest.Azure.OData.ODataQuery<GenericResourceFilter> filter = null)
         {
-            IResourceManagementClient armClient = ResourceClient;
+            string nextPageLink = null;
+            var results = new List<PSKeyVaultIdentityItem>();
+            do
+            {
+                var response = ShouldListByResourceGroup(resourceGroupName, null) ?
+                    ListByResourceGroup(resourceGroupName, filter, nextPageLink) :
+                    ListResources(filter, nextPageLink);
+                results.AddRange(response.Select(r => new PSKeyVaultIdentityItem(r)));
+                nextPageLink = response?.NextPageLink;
+            } while (!string.IsNullOrEmpty(nextPageLink));
+            return results;
+        }
 
-            return new GenericPageEnumerable<GenericResource>(() => armClient.Resources.List(filter), armClient.Resources.ListNext, first, skip).Select(r => new PSKeyVaultIdentityItem(r));
+        public virtual Rest.Azure.IPage<GenericResource> ListResources(
+            Rest.Azure.OData.ODataQuery<GenericResourceFilter> filter = null, 
+            string NextPageLink = null)
+        {
+            IResourceManagementClient armClient = ResourceClient;
+            return string.IsNullOrEmpty(NextPageLink) ? 
+                armClient.Resources.List(filter) : 
+                armClient.Resources.ListNext(NextPageLink);
         }
 
-        private IEnumerable<PSKeyVaultIdentityItem> ListByResourceGroup(
+        private Rest.Azure.IPage<GenericResource> ListByResourceGroup(
             string resourceGroupName,
-            Rest.Azure.OData.ODataQuery<GenericResourceFilter> filter,
-            ulong first = ulong.MaxValue,
-            ulong skip = ulong.MinValue)
+            Rest.Azure.OData.ODataQuery<GenericResourceFilter> filter = null,
+            string NextPageLink = null)
         {
             IResourceManagementClient armClient = ResourceClient;
-
-            return new GenericPageEnumerable<GenericResource>(() => armClient.ResourceGroups.ListResources(resourceGroupName, filter), armClient.ResourceGroups.ListResourcesNext, first, skip).Select(r => new PSKeyVaultIdentityItem(r));
+            return string.IsNullOrEmpty(NextPageLink) ? 
+                armClient.ResourceGroups.ListResources(resourceGroupName, filter) :
+                armClient.ResourceGroups.ListResourcesNext(NextPageLink);
         }
 
         protected string GetResourceGroupName(string name, bool isHsm = false)