From 0f174d9ce0f90beec7a30b26d6bb42f464908a0a Mon Sep 17 00:00:00 2001
From: Shicong Huang <shicong@google.com>
Date: Wed, 20 May 2020 14:24:45 -0400
Subject: [PATCH] Add all existing entities to VKeyTranslatorFactory (#595)

EntityClasses.ALL_CLASSES has all of our registered entities so
we can use it to initialize VKeyTranslatorFactory.classRegistry
to avoid adding them one by one.

Also, this PR changed to use Key.getKind() to get the kind of
the entity to solve the problem that when the entity class
is an inner class, its kind should still be the class name
instead of OuterClass$InnerClass.
---
 .../registry/model/ofy/ObjectifyService.java  |  4 +--
 .../translators/VKeyTranslatorFactory.java    | 36 +++++++++----------
 2 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/core/src/main/java/google/registry/model/ofy/ObjectifyService.java b/core/src/main/java/google/registry/model/ofy/ObjectifyService.java
index d33ffcfd5..d45d3903d 100644
--- a/core/src/main/java/google/registry/model/ofy/ObjectifyService.java
+++ b/core/src/main/java/google/registry/model/ofy/ObjectifyService.java
@@ -36,8 +36,6 @@ import com.googlecode.objectify.impl.translate.opt.joda.MoneyStringTranslatorFac
 import google.registry.config.RegistryEnvironment;
 import google.registry.model.EntityClasses;
 import google.registry.model.ImmutableObject;
-import google.registry.model.contact.ContactResource;
-import google.registry.model.host.HostResource;
 import google.registry.model.translators.BloomFilterOfStringTranslatorFactory;
 import google.registry.model.translators.CidrAddressBlockTranslatorFactory;
 import google.registry.model.translators.CommitLogRevisionsTranslatorFactory;
@@ -131,7 +129,7 @@ public class ObjectifyService {
             new InetAddressTranslatorFactory(),
             new MoneyStringTranslatorFactory(),
             new ReadableInstantUtcTranslatorFactory(),
-            new VKeyTranslatorFactory(HostResource.class, ContactResource.class),
+            new VKeyTranslatorFactory(),
             new UpdateAutoTimestampTranslatorFactory())) {
       factory().getTranslators().add(translatorFactory);
     }
diff --git a/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java b/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java
index a60f45b88..c7bf520b5 100644
--- a/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java
+++ b/core/src/main/java/google/registry/model/translators/VKeyTranslatorFactory.java
@@ -16,13 +16,12 @@ package google.registry.model.translators;
 
 import static com.google.common.base.Functions.identity;
 import static com.google.common.collect.ImmutableMap.toImmutableMap;
+import static google.registry.model.EntityClasses.ALL_CLASSES;
 
 import com.google.appengine.api.datastore.Key;
-import com.google.common.base.Splitter;
 import com.google.common.collect.ImmutableMap;
+import com.googlecode.objectify.annotation.EntitySubclass;
 import google.registry.persistence.VKey;
-import java.util.List;
-import java.util.stream.Stream;
 
 /**
  * Translator factory for VKey.
@@ -34,21 +33,16 @@ public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey,
 
   // Class registry allowing us to restore the original class object from the unqualified class
   // name, which is all the datastore key gives us.
-  private final ImmutableMap<String, Class> classRegistry;
+  // Note that entities annotated with @EntitySubclass are removed because they share the same
+  // kind of the key with their parent class.
+  private static final ImmutableMap<String, Class> CLASS_REGISTRY =
+      ALL_CLASSES.stream()
+          .filter(clazz -> !clazz.isAnnotationPresent(EntitySubclass.class))
+          .collect(toImmutableMap(com.googlecode.objectify.Key::getKind, identity()));
+  ;
 
-  public VKeyTranslatorFactory(Class... refClasses) {
+  public VKeyTranslatorFactory() {
     super(VKey.class);
-
-    // Store a registry of all classes by their unqualified name.
-    classRegistry =
-        Stream.of(refClasses)
-            .collect(
-                toImmutableMap(
-                    clazz -> {
-                      List<String> nameComponent = Splitter.on('.').splitToList(clazz.getName());
-                      return nameComponent.get(nameComponent.size() - 1);
-                    },
-                    identity()));
   }
 
   @Override
@@ -57,14 +51,16 @@ public class VKeyTranslatorFactory extends AbstractSimpleTranslatorFactory<VKey,
       @Override
       public VKey loadValue(Key datastoreValue) {
         // TODO(mmuller): we need to call a method on refClass to also reconstitute the SQL key.
-        return VKey.createOfy(
-            classRegistry.get(datastoreValue.getKind()),
-            com.googlecode.objectify.Key.create(datastoreValue));
+        return datastoreValue == null
+            ? null
+            : VKey.createOfy(
+                CLASS_REGISTRY.get(datastoreValue.getKind()),
+                com.googlecode.objectify.Key.create(datastoreValue));
       }
 
       @Override
       public Key saveValue(VKey key) {
-        return key.getOfyKey().getRaw();
+        return key == null ? null : key.getOfyKey().getRaw();
       }
     };
   }