From 5488e1b3234aac72bb89b6bb5b027847588f7185 Mon Sep 17 00:00:00 2001
From: Michael Muller <mmuller@google.com>
Date: Fri, 11 Sep 2020 13:45:51 -0400
Subject: [PATCH] Fix accessing superclass fields in checkExists() (#799)

* Fix accessing superclass fields in checkExists()

JpaTransactionManagerImpl doesn't respect @Id fields in mapped superclasses.
Replace calls to getDeclaredId() and getDeclaredField() with superclass
friendly counterparts.
---
 .../JpaTransactionManagerImpl.java            | 21 +++++++++++++---
 .../transaction/TransactionManagerTest.java   | 24 +++++++++++++++----
 2 files changed, 38 insertions(+), 7 deletions(-)

diff --git a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java
index ff17843bf..912448766 100644
--- a/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java
+++ b/core/src/main/java/google/registry/persistence/transaction/JpaTransactionManagerImpl.java
@@ -391,7 +391,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   private static ImmutableSet<EntityId> getEntityIdsFromEntity(
       EntityType<?> entityType, Object entity) {
     if (entityType.hasSingleIdAttribute()) {
-      String idName = entityType.getDeclaredId(entityType.getIdType().getJavaType()).getName();
+      String idName = entityType.getId(entityType.getIdType().getJavaType()).getName();
       Object idValue = getFieldValue(entity, idName);
       return ImmutableSet.of(new EntityId(idName, idValue));
     } else {
@@ -402,7 +402,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
   private static ImmutableSet<EntityId> getEntityIdsFromSqlKey(
       EntityType<?> entityType, Object sqlKey) {
     if (entityType.hasSingleIdAttribute()) {
-      String idName = entityType.getDeclaredId(entityType.getIdType().getJavaType()).getName();
+      String idName = entityType.getId(entityType.getIdType().getJavaType()).getName();
       return ImmutableSet.of(new EntityId(idName, sqlKey));
     } else {
       return getEntityIdsFromIdContainer(entityType, sqlKey);
@@ -429,7 +429,7 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
 
   private static Object getFieldValue(Object object, String fieldName) {
     try {
-      Field field = object.getClass().getDeclaredField(fieldName);
+      Field field = getField(object.getClass(), fieldName);
       field.setAccessible(true);
       return field.get(object);
     } catch (NoSuchFieldException | IllegalAccessException e) {
@@ -437,6 +437,21 @@ public class JpaTransactionManagerImpl implements JpaTransactionManager {
     }
   }
 
+  /** Gets the field definition from clazz or any superclass. */
+  private static Field getField(Class clazz, String fieldName) throws NoSuchFieldException {
+    try {
+      // Note that we have to use getDeclaredField() for this, getField() just finds public fields.
+      return clazz.getDeclaredField(fieldName);
+    } catch (NoSuchFieldException e) {
+      Class base = clazz.getSuperclass();
+      if (base != null) {
+        return getField(base, fieldName);
+      } else {
+        throw e;
+      }
+    }
+  }
+
   private static class TransactionInfo {
     EntityManager entityManager;
     boolean inTransaction = false;
diff --git a/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java b/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java
index 9290e4ce7..50d6044da 100644
--- a/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java
+++ b/core/src/test/java/google/registry/persistence/transaction/TransactionManagerTest.java
@@ -37,6 +37,8 @@ import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Set;
 import java.util.stream.Stream;
+import javax.persistence.Embeddable;
+import javax.persistence.MappedSuperclass;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.TestTemplate;
 import org.junit.jupiter.api.extension.RegisterExtension;
@@ -65,7 +67,7 @@ public class TransactionManagerTest {
           .withClock(fakeClock)
           .withDatastoreAndCloudSql()
           .withOfyTestEntities(TestEntity.class)
-          .withJpaUnitTestEntities(TestEntity.class)
+          .withJpaUnitTestEntities(TestEntity.class, TestEntityBase.class)
           .build();
 
   TransactionManagerTest() {}
@@ -324,17 +326,31 @@ public class TransactionManagerTest {
     entities.forEach(TransactionManagerTest::assertEntityNotExist);
   }
 
+  /**
+   * We put the id field into a base class to test that id fields can be discovered in a base class.
+   */
+  @MappedSuperclass
+  @Embeddable
+  private static class TestEntityBase extends ImmutableObject {
+    @Id @javax.persistence.Id protected String name;
+
+    TestEntityBase(String name) {
+      this.name = name;
+    }
+
+    TestEntityBase() {}
+  }
+
   @Entity(name = "TxnMgrTestEntity")
   @javax.persistence.Entity(name = "TestEntity")
-  private static class TestEntity extends ImmutableObject {
-    @Id @javax.persistence.Id private String name;
+  private static class TestEntity extends TestEntityBase {
 
     private String data;
 
     private TestEntity() {}
 
     private TestEntity(String name, String data) {
-      this.name = name;
+      super(name);
       this.data = data;
     }