From 9dde01306611c37d74c76fe69d2f87a9ae3d07f8 Mon Sep 17 00:00:00 2001
From: nickfelt <nickfelt@google.com>
Date: Fri, 23 Sep 2016 13:53:46 -0700
Subject: [PATCH] Clean up LoadAndResaveCommand

This refactors LoadAndResaveCommand a bit so that it's more straightfoward
and just switches on an enum to determine the resource type.  This has the
pro of actually making the command line flag friendlier anyway.

Also changed how it handles a non-existent (or non-active) resource so that
it throws an IAE rather than silently doing stageEntityChange(null, null).

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=134113671
---
 .../tools/javascrap/LoadAndResaveCommand.java | 60 +++++++++----------
 1 file changed, 29 insertions(+), 31 deletions(-)

diff --git a/java/google/registry/tools/javascrap/LoadAndResaveCommand.java b/java/google/registry/tools/javascrap/LoadAndResaveCommand.java
index ef25d6c88..7c9cbb10b 100644
--- a/java/google/registry/tools/javascrap/LoadAndResaveCommand.java
+++ b/java/google/registry/tools/javascrap/LoadAndResaveCommand.java
@@ -14,23 +14,20 @@
 
 package google.registry.tools.javascrap;
 
-import static com.google.common.collect.Maps.uniqueIndex;
-import static google.registry.model.EppResourceUtils.loadByForeignKey;
-import static google.registry.model.EppResourceUtils.loadDomainApplication;
+import static google.registry.model.index.ForeignKeyIndex.loadAndGetKey;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 import static org.joda.time.DateTimeZone.UTC;
 
 import com.beust.jcommander.Parameter;
 import com.beust.jcommander.Parameters;
-import com.google.common.base.Function;
+import com.googlecode.objectify.Key;
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactResource;
 import google.registry.model.domain.DomainApplication;
 import google.registry.model.domain.DomainResource;
 import google.registry.model.host.HostResource;
 import google.registry.tools.MutatingCommand;
-import java.util.Arrays;
-import java.util.Map;
 import org.joda.time.DateTime;
 
 /** A command to load and resave an entity, which triggers @OnSave changes. */
@@ -39,41 +36,42 @@ import org.joda.time.DateTime;
     commandDescription = "Load and resave an object, to trigger @OnSave changes")
 public final class LoadAndResaveCommand extends MutatingCommand {
 
+  private enum ResourceType { CONTACT, HOST, DOMAIN, APPLICATION }
+
   @Parameter(
     names = "--type",
-    description =
-        "Resource type (ContactResource, DomainApplication, DomainResource, HostResource).")
-  protected String type;
+    description = "Resource type.")
+  protected ResourceType type;
 
   @Parameter(
     names = "--id",
     description = "Foreign key of the resource, or application ID of the domain application.")
   protected String uniqueId;
 
-  private static final Map<String, Class<? extends EppResource>> CLASSES_BY_NAME =
-      uniqueIndex(
-          Arrays.<Class<? extends EppResource>>asList(
-              ContactResource.class,
-              DomainApplication.class,
-              DomainResource.class,
-              HostResource.class),
-          new Function<Class<?>, String>(){
-            @Override
-            public String apply(Class<?> clazz) {
-              return clazz.getSimpleName();
-            }});
-
   @Override
   protected void init() throws Exception {
-    Class<? extends EppResource> clazz = CLASSES_BY_NAME.get(type);
-    EppResource existing =
-        (clazz == DomainApplication.class)
-            ? loadDomainApplication(uniqueId, DateTime.now(UTC))
-            : loadByForeignKey(clazz, uniqueId, DateTime.now(UTC));
-    // Find the resource by foreign key, and then reload it directly, bypassing loadByUniqueId().
-    // We need to do a reload because otherwise stageEntityChange() can fail due to the implicit
-    // changes done when forwarding the resource to "now" in cloneProjectedAtTime().
-    EppResource resource = ofy().load().entity(existing).now();
+    Key<? extends EppResource> resourceKey = checkArgumentNotNull(
+        getResourceKey(type, uniqueId, DateTime.now(UTC)),
+        "Could not find active resource of type %s: %s", type, uniqueId);
+    // Load the resource directly to bypass running cloneProjectedAtTime() automatically, which can
+    // cause stageEntityChange() to fail due to implicit projection changes.
+    EppResource resource = ofy().load().key(resourceKey).now();
     stageEntityChange(resource, resource);
   }
+
+  private Key<? extends EppResource> getResourceKey(
+      ResourceType type, String uniqueId, DateTime now) {
+    switch (type) {
+      case CONTACT:
+        return loadAndGetKey(ContactResource.class, uniqueId, now);
+      case HOST:
+        return loadAndGetKey(HostResource.class, uniqueId, now);
+      case DOMAIN:
+        return loadAndGetKey(DomainResource.class, uniqueId, now);
+      case APPLICATION:
+        return Key.create(DomainApplication.class, uniqueId);
+      default:
+        throw new IllegalStateException("Unknown type: " + type);
+    }
+  }
 }