Generate string to uniquely identify a SqlEntity (#1271)

* Generate string to uniquely identify a SqlEntity

Add a method to SqlEntity that returns a string built from the entity's
primary key(s). This string can be used in logging.

core/src/main/java/google/registry/beam/common/RegistryJpaIO.java

@@ -23,6 +23,7 @@ import com.google.common.collect.Streams;
 import google.registry.backup.AppEngineEnvironment;
 import google.registry.beam.common.RegistryQuery.CriteriaQuerySupplier;
 import google.registry.model.ofy.ObjectifyService;
+import google.registry.model.replay.SqlEntity;
 import google.registry.persistence.transaction.JpaTransactionManager;
 import google.registry.persistence.transaction.TransactionManagerFactory;
 import java.io.Serializable;
@@ -358,17 +359,17 @@ public final class RegistryJpaIO {
     @ProcessElement
     public void processElement(@Element KV<ShardedKey<Integer>, Iterable<T>> kv) {
       try (AppEngineEnvironment env = new AppEngineEnvironment()) {
-        ImmutableList<Object> ofyEntities =
+        ImmutableList<Object> entities =
             Streams.stream(kv.getValue())
                 .map(this.jpaConverter::apply)
                 // TODO(b/177340730): post migration delete the line below.
                 .filter(Objects::nonNull)
                 .collect(ImmutableList.toImmutableList());
         try {
-          jpaTm().transact(() -> jpaTm().putAll(ofyEntities));
+          jpaTm().transact(() -> jpaTm().putAll(entities));
-          counter.inc(ofyEntities.size());
+          counter.inc(entities.size());
         } catch (RuntimeException e) {
-          processSingly(ofyEntities);
+          processSingly(entities);
         }
       }
     }
@@ -377,19 +378,22 @@ public final class RegistryJpaIO {
      * Writes entities in a failed batch one by one to identify the first bad entity and throws a
      * {@link RuntimeException} on it.
      */
-    private void processSingly(ImmutableList<Object> ofyEntities) {
+    private void processSingly(ImmutableList<Object> entities) {
-      for (Object ofyEntity : ofyEntities) {
+      for (Object entity : entities) {
         try {
-          jpaTm().transact(() -> jpaTm().put(ofyEntity));
+          jpaTm().transact(() -> jpaTm().put(entity));
           counter.inc();
         } catch (RuntimeException e) {
-          throw new RuntimeException(toOfyKey(ofyEntity).toString(), e);
+          throw new RuntimeException(toEntityKeyString(entity), e);
         }
       }
     }
 
-    private com.googlecode.objectify.Key<?> toOfyKey(Object ofyEntity) {
+    private String toEntityKeyString(Object entity) {
-      return com.googlecode.objectify.Key.create(ofyEntity);
+      if (entity instanceof SqlEntity) {
+        return ((SqlEntity) entity).getPrimaryKeyString();
+      }
+      return "Non-SqlEntity: " + String.valueOf(entity);
     }
   }
 }

core/src/main/java/google/registry/model/registrar/RegistrarContact.java

@@ -30,6 +30,7 @@ import static google.registry.util.PasswordUtils.SALT_SUPPLIER;
 import static google.registry.util.PasswordUtils.hashPassword;
 import static java.util.stream.Collectors.joining;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.google.common.base.Enums;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
@@ -396,7 +397,8 @@ public class RegistrarContact extends ImmutableObject
   }
 
   /** Class to represent the composite primary key for {@link RegistrarContact} entity. */
-  static class RegistrarPocId extends ImmutableObject implements Serializable {
+  @VisibleForTesting
+  public static class RegistrarPocId extends ImmutableObject implements Serializable {
 
     String emailAddress;
 
@@ -405,7 +407,8 @@ public class RegistrarContact extends ImmutableObject
     // Hibernate requires this default constructor.
     private RegistrarPocId() {}
 
-    RegistrarPocId(String emailAddress, String registrarId) {
+    @VisibleForTesting
+    public RegistrarPocId(String emailAddress, String registrarId) {
       this.emailAddress = emailAddress;
       this.registrarId = registrarId;
     }

core/src/main/java/google/registry/model/replay/SqlEntity.java

@@ -14,6 +14,8 @@
 
 package google.registry.model.replay;
 
+import static google.registry.persistence.transaction.TransactionManagerFactory.jpaTm;
+
 import java.util.Optional;
 
 /**
@@ -29,4 +31,19 @@ public interface SqlEntity {
 
   /** A method that will ber called before the object is saved to SQL in asynchronous replay. */
   default void beforeSqlSaveOnReplay() {}
+
+  /* Returns this entity's primary key field(s) in a string. */
+  default String getPrimaryKeyString() {
+    return jpaTm()
+        .transact(
+            () ->
+                String.format(
+                    "%s_%s",
+                    this.getClass().getSimpleName(),
+                    jpaTm()
+                        .getEntityManager()
+                        .getEntityManagerFactory()
+                        .getPersistenceUnitUtil()
+                        .getIdentifier(this)));
+  }
 }

core/src/test/java/google/registry/schema/replay/SqlEntityTest.java

@@ -0,0 +1,71 @@
+// Copyright 2021 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.schema.replay;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.persistence.transaction.TransactionManagerFactory.tm;
+
+import google.registry.model.registrar.Registrar;
+import google.registry.model.registrar.RegistrarContact;
+import google.registry.model.registrar.RegistrarContact.RegistrarPocId;
+import google.registry.persistence.VKey;
+import google.registry.persistence.transaction.TransactionManagerFactory;
+import google.registry.testing.AppEngineExtension;
+import google.registry.testing.DatastoreEntityExtension;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Order;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.extension.RegisterExtension;
+
+/** Unit tests for {@link SqlEntity#getPrimaryKeyString}. */
+public class SqlEntityTest {
+
+  @RegisterExtension
+  @Order(1)
+  final DatastoreEntityExtension datastoreEntityExtension = new DatastoreEntityExtension();
+
+  @RegisterExtension
+  final AppEngineExtension database = new AppEngineExtension.Builder().withCloudSql().build();
+
+  @BeforeEach
+  void setup() throws Exception {
+    TransactionManagerFactory.setTmForTest(TransactionManagerFactory.jpaTm());
+    AppEngineExtension.loadInitialData();
+  }
+
+  @AfterEach
+  void teardown() {
+    TransactionManagerFactory.removeTmOverrideForTest();
+  }
+
+  @Test
+  void getPrimaryKeyString_oneIdColumn() {
+    // AppEngineExtension canned data: Registrar1
+    VKey<Registrar> key = Registrar.createVKey("NewRegistrar");
+    String expected = "NewRegistrar";
+    assertThat(tm().transact(() -> tm().loadByKey(key)).getPrimaryKeyString()).contains(expected);
+  }
+
+  @Test
+  void getPrimaryKeyString_multiId() {
+    // AppEngineExtension canned data: RegistrarContact1
+    VKey<RegistrarContact> key =
+        VKey.createSql(
+            RegistrarContact.class, new RegistrarPocId("janedoe@theregistrar.com", "NewRegistrar"));
+    String expected = "emailAddress=janedoe@theregistrar.com\n    registrarId=NewRegistrar";
+    assertThat(tm().transact(() -> tm().loadByKey(key)).getPrimaryKeyString()).contains(expected);
+  }
+}