fix(provider/google): Add pagination support for retrieving image res…

…ources. (#1549)
spinnaker · Mar 31, 2017 · dbffbbf · dbffbbf
1 parent a5f4659
commit dbffbbf
Show file tree

Hide file tree

Showing 6 changed files with 127 additions and 45 deletions.
diff --git a/...ver-google/src/main/groovy/com/netflix/spinnaker/clouddriver/google/deploy/GCEUtil.groovy b/...ver-google/src/main/groovy/com/netflix/spinnaker/clouddriver/google/deploy/GCEUtil.groovy
@@ -82,32 +82,51 @@ class GCEUtil {
 
     def imageProjects = [projectName] + credentials?.imageProjects + baseImageProjects - null
     def sourceImage = null
+    // We want predictable iteration order that matches the order of insertion.
+    def imageProjectToNextPageTokenMap = new LinkedHashMap<String, String>()
 
-    def imageListBatch = buildBatchRequest(compute, clouddriverUserAgentApplicationName)
-    def imageListCallback = new JsonBatchCallback<ImageList>() {
-      @Override
-      void onFailure(GoogleJsonError e, HttpHeaders responseHeaders) throws IOException {
-        updateStatusAndThrowNotFoundException("Error locating $imageName in these projects: $imageProjects: $e.message", task, phase)
-      }
+    // This will ensure that each image project is queried.
+    imageProjects.each { imageProjectToNextPageTokenMap[it] = null }
 
-      @Override
-      void onSuccess(ImageList imageList, HttpHeaders responseHeaders) throws IOException {
-        // No need to look through these images if the requested image was already found.
-        if (!sourceImage) {
-          for (def image : imageList.items) {
-            if (image.name == imageName) {
-              sourceImage = image
-            }
+    while (!sourceImage && imageProjectToNextPageTokenMap) {
+      def imageListBatch = buildBatchRequest(compute, clouddriverUserAgentApplicationName)
+      def imageListCallback = new JsonBatchCallback<ImageList>() {
+        @Override
+        void onFailure(GoogleJsonError e, HttpHeaders responseHeaders) throws IOException {
+          updateStatusAndThrowNotFoundException("Error locating $imageName in these projects: $imageProjects: $e.message", task, phase)
+        }
+
+        @Override
+        void onSuccess(ImageList imageList, HttpHeaders responseHeaders) throws IOException {
+          // No need to look through these images if the requested image was already found.
+          if (!sourceImage) {
+            sourceImage = imageList.items.find { it.name == imageName }
+          }
+
+          // selfLinks look like this: https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images
+          def selfLinkTokens = imageList.getSelfLink().tokenize("/")
+          def imageProject = selfLinkTokens[selfLinkTokens.size() - 3]
+
+          if (imageList.nextPageToken) {
+            imageProjectToNextPageTokenMap[imageProject] = imageList.nextPageToken
+          } else {
+            imageProjectToNextPageTokenMap.remove(imageProject)
           }
         }
       }
-    }
 
-    for (imageProject in imageProjects) {
-      compute.images().list(imageProject).queue(imageListBatch, imageListCallback)
-    }
+      imageProjectToNextPageTokenMap.each { imageProject, pageToken ->
+        def imagesList = compute.images().list(imageProject)
+
+        if (pageToken) {
+          imagesList = imagesList.setPageToken(pageToken)
+        }
+
+        imagesList.queue(imageListBatch, imageListCallback)
+      }
 
-    executor.timeExecuteBatch( imageListBatch, "findImage")
+      executor.timeExecuteBatch(imageListBatch, "findImage")
+    }
 
     if (sourceImage) {
       return sourceImage

diff --git a/...vy/com/netflix/spinnaker/clouddriver/google/provider/agent/GoogleImageCachingAgent.groovy b/...vy/com/netflix/spinnaker/clouddriver/google/provider/agent/GoogleImageCachingAgent.groovy
@@ -20,6 +20,7 @@ import com.fasterxml.jackson.databind.ObjectMapper
 import com.google.api.client.googleapis.batch.BatchRequest
 import com.google.api.client.googleapis.batch.json.JsonBatchCallback
 import com.google.api.client.http.HttpHeaders
+import com.google.api.services.compute.Compute
 import com.google.api.services.compute.model.Image
 import com.google.api.services.compute.model.ImageList
 import com.netflix.servo.util.VisibleForTesting
@@ -72,19 +73,33 @@ class GoogleImageCachingAgent extends AbstractGoogleCachingAgent {
 
   List<Image> loadImages() {
     List<Image> imageList = []
+    List<String> allImageProjects = [project] + imageProjects + baseImageProjects - null
 
-    BatchRequest imageRequest = buildBatchRequest()
+    // We want predictable iteration order that matches the order of insertion.
+    LinkedHashMap<String, String> imageProjectToNextPageTokenMap = new LinkedHashMap<>()
 
-    compute.images().list(project).queue(imageRequest, new AllImagesCallback(imageList: imageList))
-    imageProjects.each {
-      compute.images().list(it).queue(imageRequest, new AllImagesCallback(imageList: imageList))
-    }
-    baseImageProjects.each {
-      compute.images().list(it).queue(imageRequest, new AllImagesCallback<ImageList>(imageList: imageList))
+    // This will ensure that each image project is queried.
+    allImageProjects.each { imageProjectToNextPageTokenMap[it] = null }
+
+    while (imageProjectToNextPageTokenMap) {
+      BatchRequest imageListBatch = buildBatchRequest()
+      AllImagesCallback<ImageList> imageListCallback =
+        new AllImagesCallback(imageProjectToNextPageTokenMap: imageProjectToNextPageTokenMap, imageList: imageList)
+
+      imageProjectToNextPageTokenMap.each { imageProject, pageToken ->
+        Compute.Images.List imagesList = compute.images().list(imageProject)
+
+        if (pageToken) {
+          imagesList = imagesList.setPageToken(pageToken)
+        }
+
+        imagesList.queue(imageListBatch, imageListCallback)
+      }
+
+      executeIfRequestsAreQueued(imageListBatch, "ImageCaching.image")
     }
-    executeIfRequestsAreQueued(imageRequest, "ImageCaching.image")
 
-    imageList
+    return imageList
   }
 
   private CacheResult buildCacheResult(ProviderCache _, List<Image> imageList) {
@@ -113,6 +128,7 @@ class GoogleImageCachingAgent extends AbstractGoogleCachingAgent {
 
   class AllImagesCallback<ImageList> extends JsonBatchCallback<ImageList> implements FailureLogger {
 
+    LinkedHashMap<String, String> imageProjectToNextPageTokenMap
     List<Image> imageList
 
     @Override
@@ -121,6 +137,16 @@ class GoogleImageCachingAgent extends AbstractGoogleCachingAgent {
       if (nonDeprecatedImages) {
         imageList.addAll(nonDeprecatedImages)
       }
+
+      // selfLinks look like this: https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images
+      def selfLinkTokens = imageListResult.getSelfLink().tokenize("/")
+      def imageProject = selfLinkTokens[selfLinkTokens.size() - 3]
+
+      if (imageListResult.nextPageToken) {
+        imageProjectToNextPageTokenMap[imageProject] = imageListResult.nextPageToken
+      } else {
+        imageProjectToNextPageTokenMap.remove(imageProject)
+      }
     }
   }
 }
diff --git a/...google/src/test/groovy/com/netflix/spinnaker/clouddriver/google/deploy/GCEUtilSpec.groovy b/...google/src/test/groovy/com/netflix/spinnaker/clouddriver/google/deploy/GCEUtilSpec.groovy
@@ -100,8 +100,10 @@ class GCEUtilSpec extends Specification {
 
       batchMock.demand.size { return 1 }
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([soughtImage])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [soughtImage]
+        )
         callback.onSuccess(imageList, null)
       }
 
@@ -152,8 +154,10 @@ class GCEUtilSpec extends Specification {
 
       batchMock.demand.size { return 1 }
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([soughtImage])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [soughtImage]
+        )
         callback.onSuccess(imageList, null)
       }
 
@@ -202,8 +206,20 @@ class GCEUtilSpec extends Specification {
 
       batchMock.demand.size { return 1 }
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([new Image(name: IMAGE_NAME + "-WRONG")])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [new Image(name: IMAGE_NAME + "-WRONG")]
+        )
+        callback.onSuccess(imageList, null)
+        imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/centos-cloud/global/images",
+          items: [new Image(name: IMAGE_NAME + "-WRONG")]
+        )
+        callback.onSuccess(imageList, null)
+        imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images",
+          items: [new Image(name: IMAGE_NAME + "-WRONG")]
+        )
         callback.onSuccess(imageList, null)
       }
 

diff --git a/...pinnaker/clouddriver/google/deploy/ops/CreateGoogleInstanceAtomicOperationUnitSpec.groovy b/...pinnaker/clouddriver/google/deploy/ops/CreateGoogleInstanceAtomicOperationUnitSpec.groovy
@@ -85,8 +85,10 @@ class CreateGoogleInstanceAtomicOperationUnitSpec extends Specification implemen
 
       batchMock.demand.size { return 1 }
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([new Image(name: IMAGE)])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [new Image(name: IMAGE)]
+        )
         callback.onSuccess(imageList, null)
       }
 

diff --git a/...innaker/clouddriver/google/deploy/ops/UpsertGoogleImageTagsAtomicOperationUnitSpec.groovy b/...innaker/clouddriver/google/deploy/ops/UpsertGoogleImageTagsAtomicOperationUnitSpec.groovy
@@ -82,8 +82,10 @@ class UpsertGoogleImageTagsAtomicOperationUnitSpec extends Specification impleme
 
       batchMock.demand.size { return 1 }
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([new Image(name: IMAGE_NAME, selfLink: IMAGE_SELF_LINK)])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [new Image(name: IMAGE_NAME, selfLink: IMAGE_SELF_LINK)]
+        )
         callback.onSuccess(imageList, null)
       }
 
@@ -152,8 +154,10 @@ class UpsertGoogleImageTagsAtomicOperationUnitSpec extends Specification impleme
 
       batchMock.demand.size { return 1 }
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([new Image(name: IMAGE_NAME, selfLink: IMAGE_SELF_LINK, labels: LABELS)])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [new Image(name: IMAGE_NAME, selfLink: IMAGE_SELF_LINK, labels: LABELS)]
+        )
         callback.onSuccess(imageList, null)
       }
 
@@ -219,8 +223,20 @@ class UpsertGoogleImageTagsAtomicOperationUnitSpec extends Specification impleme
       }
 
       batchMock.demand.execute {
-        def imageList = new ImageList()
-        imageList.setItems([new Image(name: IMAGE_NAME + "-WRONG")])
+        def imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/$PROJECT_NAME/global/images",
+          items: [new Image(name: IMAGE_NAME + "-WRONG")]
+        )
+        callback.onSuccess(imageList, null)
+        imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/centos-cloud/global/images",
+          items: [new Image(name: IMAGE_NAME + "-WRONG")]
+        )
+        callback.onSuccess(imageList, null)
+        imageList = new ImageList(
+          selfLink: "https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images",
+          items: [new Image(name: IMAGE_NAME + "-WRONG")]
+        )
         callback.onSuccess(imageList, null)
       }
 

diff --git a/...om/netflix/spinnaker/clouddriver/google/provider/agent/GoogleImageCachingAgentSpec.groovy b/...om/netflix/spinnaker/clouddriver/google/provider/agent/GoogleImageCachingAgentSpec.groovy
@@ -28,16 +28,18 @@ class GoogleImageCachingAgentSpec extends Specification {
     setup:
       def imageList = new ArrayList<String>()
       def imagesCallback1 = new GoogleImageCachingAgent.AllImagesCallback(new GoogleImageCachingAgent())
+      imagesCallback1.imageProjectToNextPageTokenMap = [:]
       imagesCallback1.imageList = imageList
-      def imageListResult1 = new ImageList()
+      def imageListResult1 = new ImageList(selfLink: "https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images")
       imageListResult1.setItems([new Image(name: "backports-debian-7-wheezy-v20141108"),
                                  new Image(name: "debian-7-wheezy-v20141108"),
                                  new Image(name: "someos-8-something-v20141108"),
                                  new Image(name: "otheros-9-something-v20141108")])
 
       def imagesCallback2 = new GoogleImageCachingAgent.AllImagesCallback(new GoogleImageCachingAgent())
+      imagesCallback2.imageProjectToNextPageTokenMap = [:]
       imagesCallback2.imageList = imageList
-      def imageListResult2 = new ImageList()
+      def imageListResult2 = new ImageList(selfLink: "https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images")
       imageListResult2.setItems([new Image(name: "ubuntu-1404-trusty-v20141028"),
                                  buildImage("ubuntu-1404-trusty-v20141029", true),
                                  new Image(name: "ubuntu-1404-trusty-v20141031a")])
@@ -59,8 +61,9 @@ class GoogleImageCachingAgentSpec extends Specification {
     setup:
       def imageList = new ArrayList<String>()
       def imagesCallback = new GoogleImageCachingAgent.AllImagesCallback(new GoogleImageCachingAgent())
+      imagesCallback.imageProjectToNextPageTokenMap = [:]
       imagesCallback.imageList = imageList
-      def imageListResult = new ImageList()
+      def imageListResult = new ImageList(selfLink: "https://www.googleapis.com/compute/alpha/projects/ubuntu-os-cloud/global/images")
       imageListResult.setItems([buildImage("my-image-1", false),
                                 buildImage("my-image-2", true),
                                 buildImage("my-image-3", false),