diff --git a/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java b/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java index c0144ebbe..2744195e5 100644 --- a/core/src/main/java/google/registry/beam/common/RegistryJpaIO.java +++ b/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, Iterable> kv) { try (AppEngineEnvironment env = new AppEngineEnvironment()) { - ImmutableList ofyEntities = + ImmutableList 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)); - counter.inc(ofyEntities.size()); + jpaTm().transact(() -> jpaTm().putAll(entities)); + 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 ofyEntities) { - for (Object ofyEntity : ofyEntities) { + private void processSingly(ImmutableList entities) { + 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) { - return com.googlecode.objectify.Key.create(ofyEntity); + private String toEntityKeyString(Object entity) { + if (entity instanceof SqlEntity) { + return ((SqlEntity) entity).getPrimaryKeyString(); + } + return "Non-SqlEntity: " + String.valueOf(entity); } } } diff --git a/core/src/main/java/google/registry/model/registrar/RegistrarContact.java b/core/src/main/java/google/registry/model/registrar/RegistrarContact.java index 810d67c52..bcaf96aff 100644 --- a/core/src/main/java/google/registry/model/registrar/RegistrarContact.java +++ b/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; } diff --git a/core/src/main/java/google/registry/model/replay/SqlEntity.java b/core/src/main/java/google/registry/model/replay/SqlEntity.java index 69144c6cb..da9313795 100644 --- a/core/src/main/java/google/registry/model/replay/SqlEntity.java +++ b/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))); + } } diff --git a/core/src/test/java/google/registry/schema/replay/SqlEntityTest.java b/core/src/test/java/google/registry/schema/replay/SqlEntityTest.java new file mode 100644 index 000000000..03e155055 --- /dev/null +++ b/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 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 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); + } +}