Add option --non_live_versions to set_num_instances command

This commit introduced a new flag to enable SetNumInstancesCommand to
be able to set the number of instances for all non-live versions for
a given service or for all deployed services.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=222826003

java/google/registry/tools/AppEngineAdminApiModule.java

@@ -0,0 +1,40 @@
+// Copyright 2018 The Nomulus Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.tools;
+
+import com.google.api.client.googleapis.extensions.appengine.auth.oauth2.AppIdentityCredential;
+import com.google.api.client.googleapis.util.Utils;
+import com.google.api.services.appengine.v1.Appengine;
+import dagger.Module;
+import dagger.Provides;
+import google.registry.config.CredentialModule.AppEngineAdminApiCredential;
+import google.registry.config.RegistryConfig.Config;
+import javax.inject.Singleton;
+
+/** Module providing the instance of {@link Appengine} to access App Engine Admin Api. */
+@Module
+public abstract class AppEngineAdminApiModule {
+
+  @Provides
+  @Singleton
+  public static Appengine provideAppengine(
+      @AppEngineAdminApiCredential AppIdentityCredential appIdentityCredential,
+      @Config("projectId") String projectId) {
+    return new Appengine.Builder(
+            Utils.getDefaultTransport(), Utils.getDefaultJsonFactory(), appIdentityCredential)
+        .setApplicationName(projectId)
+        .build();
+  }
+}

java/google/registry/tools/AppEngineConnection.java

@@ -66,10 +66,21 @@ class AppEngineConnection {
   }
 
   enum Service {
-    DEFAULT,
-    TOOLS,
-    BACKEND,
-    PUBAPI
+    DEFAULT("default"),
+    TOOLS("tools"),
+    BACKEND("backend"),
+    PUBAPI("pubapi");
+
+    private final String serviceId;
+
+    Service(String serviceId) {
+      this.serviceId = serviceId;
+    }
+
+    /** Returns the actual service id in App Engine. */
+    String getServiceId() {
+      return serviceId;
+    }
   }
 
   /** Returns a copy of this connection that talks to a different service. */

java/google/registry/tools/BUILD

@@ -68,6 +68,8 @@ java_library(
         "//third_party/objectify:objectify-v4_1",
         "@com_beust_jcommander",
         "@com_google_api_client",
+        "@com_google_api_client_appengine",
+        "@com_google_apis_google_api_services_appengine",
         "@com_google_apis_google_api_services_bigquery",
         "@com_google_apis_google_api_services_dns",
         "@com_google_appengine_api_1_0_sdk",

java/google/registry/tools/RegistryToolComponent.java

@@ -46,6 +46,7 @@ import javax.inject.Singleton;
 @Singleton
 @Component(
     modules = {
+      AppEngineAdminApiModule.class,
       AppEngineServiceUtilsModule.class,
       // TODO(b/36866706): Find a way to replace this with a command-line friendly version
       AuthModule.class,

java/google/registry/tools/SetNumInstancesCommand.java

@@ -15,11 +15,28 @@
 package google.registry.tools;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.collect.ImmutableSet.toImmutableSet;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
+import com.google.api.services.appengine.v1.Appengine;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.MultimapBuilder;
+import com.google.common.collect.Multimaps;
+import com.google.common.collect.Sets;
 import com.google.common.flogger.FluentLogger;
+import google.registry.config.RegistryConfig.Config;
+import google.registry.tools.AppEngineConnection.Service;
 import google.registry.util.AppEngineServiceUtils;
+import java.io.IOException;
+import java.util.AbstractMap.SimpleEntry;
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Stream;
 import javax.inject.Inject;
 
 /** A command to set the number of instances for an App Engine service. */
@@ -32,35 +49,147 @@ final class SetNumInstancesCommand implements CommandWithRemoteApi {
 
   private static final FluentLogger logger = FluentLogger.forEnclosingClass();
 
+  private static final ImmutableSet<String> ALL_VALID_SERVICES =
+      Arrays.stream(Service.values()).map(Service::name).collect(toImmutableSet());
+
+  private static final ImmutableSet<String> ALL_DEPLOYED_SERVICE_IDS =
+      Arrays.stream(Service.values()).map(Service::getServiceId).collect(toImmutableSet());
+
+  // TODO(b/119629679): Use List<Service> after upgrading jcommander to latest version.
   @Parameter(
-      names = "--service",
-      description = "Name of the App Engine service, e.g., default.",
-      required = true)
-  private String service;
+      names = "--services",
+      description =
+          "Comma-delimited list of App Engine services to set. "
+              + "Allowed values: [DEFAULT, TOOLS, BACKEND, PUBAPI]")
+  private List<String> services = ImmutableList.of();
 
   @Parameter(
-      names = "--version",
-      description = "Name of the service's version, e.g., canary.",
-      required = true)
-  private String version;
+      names = "--versions",
+      description =
+          "Comma-delimited list of App Engine versions to set, e.g., canary. "
+              + "Cannot be set if --non_live_versions is set.")
+  private List<String> versions = ImmutableList.of();
 
   @Parameter(
-      names = "--numInstances",
-      description = "The new number of instances for the version.",
+      names = "--num_instances",
+      description =
+          "The new number of instances for the given versions "
+              + "or for all non-live versions if --non_live_versions is set.",
       required = true)
   private Long numInstances;
 
+  @Parameter(
+      names = "--non_live_versions",
+      description = "Whether to set number of instances for all non-live versions.",
+      arity = 1)
+  private Boolean nonLiveVersions = false;
+
   @Inject AppEngineServiceUtils appEngineServiceUtils;
+  @Inject Appengine appengine;
+
+  @Inject
+  @Config("projectId")
+  String projectId;
 
   @Override
   public void run() throws Exception {
-    checkArgument(!service.isEmpty(), "Service must be specified");
-    checkArgument(!version.isEmpty(), "Version must be specified");
-    checkArgument(numInstances > 0, "Number of instances must be greater than zero");
+    Set<String> invalidServiceIds =
+        Sets.difference(ImmutableSet.copyOf(services), ALL_VALID_SERVICES);
+    checkArgument(invalidServiceIds.isEmpty(), "Invalid service(s): %s", invalidServiceIds);
 
+    Set<String> serviceIds =
+        services.stream()
+            .map(service -> Service.valueOf(service).getServiceId())
+            .collect(toImmutableSet());
+
+    if (nonLiveVersions) {
+      checkArgument(versions.isEmpty(), "--versions cannot be set if --non_live_versions is set");
+
+      serviceIds = serviceIds.isEmpty() ? ALL_DEPLOYED_SERVICE_IDS : serviceIds;
+      Multimap<String, String> allLiveVersionsMap = getAllLiveVersionsMap(serviceIds);
+      Multimap<String, String> manualScalingVersionsMap = getManualScalingVersionsMap(serviceIds);
+
+      // Set number of instances for versions which are manual scaling and non-live
+      manualScalingVersionsMap.forEach(
+          (serviceId, versionId) -> {
+            if (!allLiveVersionsMap.containsEntry(serviceId, versionId)) {
+              setNumInstances(serviceId, versionId, numInstances);
+            }
+          });
+    } else {
+      checkArgument(!serviceIds.isEmpty(), "Service must be specified");
+      checkArgument(!versions.isEmpty(), "Version must be specified");
+      checkArgument(numInstances > 0, "Number of instances must be greater than zero");
+
+      Multimap<String, String> manualScalingVersionsMap = getManualScalingVersionsMap(serviceIds);
+
+      for (String serviceId : serviceIds) {
+        for (String versionId : versions) {
+          checkArgument(
+              manualScalingVersionsMap.containsEntry(serviceId, versionId),
+              "Version %s of service %s is not managed through manual scaling",
+              versionId,
+              serviceId);
+          setNumInstances(serviceId, versionId, numInstances);
+        }
+      }
+    }
+  }
+
+  private void setNumInstances(String service, String version, long numInstances) {
     appEngineServiceUtils.setNumInstances(service, version, numInstances);
     logger.atInfo().log(
         "Successfully set version %s of service %s to %d instances.",
         version, service, numInstances);
   }
+
+  private Multimap<String, String> getAllLiveVersionsMap(Set<String> services) {
+    try {
+      return Stream.of(appengine.apps().services().list(projectId).execute().getServices())
+          .flatMap(Collection::stream)
+          .filter(service -> services.contains(service.getId()))
+          .flatMap(
+              service ->
+                  // getAllocations returns only live versions or null
+                  Stream.of(service.getSplit().getAllocations())
+                      .flatMap(
+                          allocations ->
+                              allocations.keySet().stream()
+                                  .map(versionId -> new SimpleEntry<>(service.getId(), versionId))))
+          .collect(
+              Multimaps.toMultimap(
+                  SimpleEntry::getKey,
+                  SimpleEntry::getValue,
+                  MultimapBuilder.treeKeys().arrayListValues()::build));
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
+  private Multimap<String, String> getManualScalingVersionsMap(Set<String> services) {
+    return services.stream()
+        .flatMap(
+            serviceId -> {
+              try {
+                return Stream.of(
+                        appengine
+                            .apps()
+                            .services()
+                            .versions()
+                            .list(projectId, serviceId)
+                            .execute()
+                            .getVersions())
+                    .flatMap(Collection::stream)
+                    .filter(version -> version.getManualScaling() != null)
+                    .map(version -> new SimpleEntry<>(serviceId, version.getId()));
+              } catch (IOException e) {
+                throw new RuntimeException(e);
+              }
+            })
+        .collect(
+            Multimaps.toMultimap(
+                SimpleEntry::getKey,
+                SimpleEntry::getValue,
+                MultimapBuilder.treeKeys().arrayListValues()::build));
+  }
 }