Merge pull request Azure#2957 from NonStatic2014/bugfix

Support paginated response in GET-AzureRmMlWebService
solankisamir · Sep 20, 2016 · 35d5b90 · 35d5b90
2 parents 163e96c + c5864d0
commit 35d5b90
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 20 deletions.
diff --git a/src/ResourceManager/MachineLearning/Commands.MachineLearning/Cmdlets/GetAzureMLWebService.cs b/src/ResourceManager/MachineLearning/Commands.MachineLearning/Cmdlets/GetAzureMLWebService.cs
@@ -17,6 +17,7 @@
 using System.Management.Automation;
 using System.Threading.Tasks;
 using Microsoft.Azure.Management.MachineLearning.WebServices.Models;
+using Microsoft.Azure.Commands.ResourceManager.Common;
 
 namespace Microsoft.Azure.Commands.MachineLearning
 {
@@ -48,25 +49,27 @@ protected override void RunCmdlet()
             }
             else
             {
-                IList<WebService> services;
-                if (!string.IsNullOrWhiteSpace(this.ResourceGroupName))
+                // This is a collection of web services get call, so determine which flavor it is
+                Func<Task<ResponseWithContinuation<WebService[]>>> getFirstServicesPageCall = () => this.WebServicesClient.GetAzureMlWebServicesBySubscriptionAsync(null, this.CancellationToken);
+                Func<string, Task<ResponseWithContinuation<WebService[]>>> getNextPageCall = nextLink => this.WebServicesClient.GetAzureMlWebServicesBySubscriptionAsync(nextLink, this.CancellationToken);
+                if (this.ResourceGroupName != null)
                 {
-                    services = this.WebServicesClient.GetAzureMlWebServicesBySubscriptionAndGroupAsync(
-                                                        this.ResourceGroupName,
-                                                        null,
-                                                        this.CancellationToken).Result;
-                }
-                else
-                {
-                    services = this.WebServicesClient.GetAzureMlWebServicesBySubscriptionAsync(
-                                                        null,
-                                                        this.CancellationToken).Result;
+                    // This is a call for resource retrieval within a resource group
+                    getFirstServicesPageCall = () => this.WebServicesClient.GetAzureMlWebServicesBySubscriptionAndGroupAsync(this.ResourceGroupName, null, this.CancellationToken);
+                    getNextPageCall = nextLink => this.WebServicesClient.GetAzureMlWebServicesBySubscriptionAndGroupAsync(this.ResourceGroupName, nextLink, this.CancellationToken);
                 }
 
-                foreach (var service in services)
-                {
-                    this.WriteObject(service, true);
-                }
+                PaginatedResponseHelper.ForEach<WebService>(
+                    getFirstPage: getFirstServicesPageCall,
+                    getNextPage: getNextPageCall,
+                    cancellationToken: this.CancellationToken,
+                    action: webServices =>
+                    {
+                        foreach (var service in webServices)
+                        {
+                            this.WriteObject(service, true);
+                        }
+                    });
             }
         }
     }

diff --git a/src/ResourceManager/MachineLearning/Commands.MachineLearning/Utilities/WebServicesClient.cs b/src/ResourceManager/MachineLearning/Commands.MachineLearning/Utilities/WebServicesClient.cs
@@ -24,6 +24,7 @@
 using Microsoft.Azure.Management.MachineLearning.WebServices.Models;
 using APIClient = Microsoft.Azure.Management.MachineLearning.
                     WebServices.AzureMLWebServicesManagementClient;
+using Microsoft.Azure.Commands.ResourceManager.Common;
 
 namespace Microsoft.Azure.Commands.MachineLearning.Utilities
 {
@@ -95,7 +96,7 @@ public WebServiceKeys GetAzureMlWebServiceKeys(
             return this.apiClient.WebServices.ListKeys(resourceGroupName, webServiceName);
         }
 
-        public async Task<IList<WebService>> 
+        public async Task<ResponseWithContinuation<WebService[]>> 
                         GetAzureMlWebServicesBySubscriptionAndGroupAsync(
                                                     string resourceGroupName, 
                                                     string nextLink, 
@@ -108,10 +109,14 @@ public async Task<IList<WebService>>
                                                                         skipToken, 
                                                                         cancellationTokenParam).ConfigureAwait(false);
 
-            return paginatedResponse.Value;
+            return new ResponseWithContinuation<WebService[]>
+            {
+                Value = paginatedResponse.Value.ToArray(),
+                NextLink = paginatedResponse.NextLink
+            };
         }
 
-        public async Task<IList<WebService>> 
+        public async Task<ResponseWithContinuation<WebService[]>> 
                         GetAzureMlWebServicesBySubscriptionAsync(
                                                     string nextLink, 
                                                     CancellationToken? cancellationToken)
@@ -124,7 +129,11 @@ await this.apiClient.WebServices.ListAsync(
                                                         skipToken, 
                                                         cancellationTokenParam).ConfigureAwait(false);
 
-            return paginatedResponse.Value;
+            return new ResponseWithContinuation<WebService[]>
+            {
+                Value = paginatedResponse.Value.ToArray(),
+                NextLink = paginatedResponse.NextLink
+            };
         }
 
         private static string GetSkipTokenFromLink(string nextLink)