From 2b794347e63b8559a8d9ba0cc41f9c9d4a233bd3 Mon Sep 17 00:00:00 2001 From: Weimin Yu Date: Thu, 28 May 2020 13:38:00 -0400 Subject: [PATCH] Refactor LevelDbFileBuilder to accept DS Entity (#599) * Refactor LevelDbFileBuilder to accept DS Entity Builder now can directly work with Datastore Entity objects. No need to wrap data in ComparableEntity. --- .../registry/tools/CompareDbBackups.java | 18 +++--- ...mparableEntity.java => EntityWrapper.java} | 44 ++++++++++++-- .../registry/tools/RecordAccumulator.java | 8 +-- .../registry/tools/CompareDbBackupsTest.java | 50 +++++++++------- ...EntityTest.java => EntityWrapperTest.java} | 18 +++--- .../registry/tools/LevelDbFileBuilder.java | 33 ++--------- .../tools/LevelDbFileBuilderTest.java | 58 ++++++++++++++----- .../registry/tools/RecordAccumulatorTest.java | 33 ++++++----- 8 files changed, 155 insertions(+), 107 deletions(-) rename core/src/main/java/google/registry/tools/{ComparableEntity.java => EntityWrapper.java} (52%) rename core/src/test/java/google/registry/tools/{ComparableEntityTest.java => EntityWrapperTest.java} (89%) diff --git a/core/src/main/java/google/registry/tools/CompareDbBackups.java b/core/src/main/java/google/registry/tools/CompareDbBackups.java index 5f1766153..155f53602 100644 --- a/core/src/main/java/google/registry/tools/CompareDbBackups.java +++ b/core/src/main/java/google/registry/tools/CompareDbBackups.java @@ -39,16 +39,14 @@ class CompareDbBackups { return; } - ImmutableSet entities1 = - RecordAccumulator.readDirectory(new File(args[0]), DATA_FILE_MATCHER) - .getComparableEntitySet(); - ImmutableSet entities2 = - RecordAccumulator.readDirectory(new File(args[1]), DATA_FILE_MATCHER) - .getComparableEntitySet(); + ImmutableSet entities1 = + RecordAccumulator.readDirectory(new File(args[0]), DATA_FILE_MATCHER).getEntityWrapperSet(); + ImmutableSet entities2 = + RecordAccumulator.readDirectory(new File(args[1]), DATA_FILE_MATCHER).getEntityWrapperSet(); // Calculate the entities added and removed. - SetView added = Sets.difference(entities2, entities1); - SetView removed = Sets.difference(entities1, entities2); + SetView added = Sets.difference(entities2, entities1); + SetView removed = Sets.difference(entities1, entities2); printHeader( String.format("First backup: %d records", entities1.size()), @@ -56,14 +54,14 @@ class CompareDbBackups { if (!removed.isEmpty()) { printHeader(removed.size() + " records were removed:"); - for (ComparableEntity entity : removed) { + for (EntityWrapper entity : removed) { System.out.println(entity); } } if (!added.isEmpty()) { printHeader(added.size() + " records were added:"); - for (ComparableEntity entity : added) { + for (EntityWrapper entity : added) { System.out.println(entity); } } diff --git a/core/src/main/java/google/registry/tools/ComparableEntity.java b/core/src/main/java/google/registry/tools/EntityWrapper.java similarity index 52% rename from core/src/main/java/google/registry/tools/ComparableEntity.java rename to core/src/main/java/google/registry/tools/EntityWrapper.java index 533c77a95..c583f1b00 100644 --- a/core/src/main/java/google/registry/tools/ComparableEntity.java +++ b/core/src/main/java/google/registry/tools/EntityWrapper.java @@ -15,20 +15,32 @@ package google.registry.tools; import com.google.appengine.api.datastore.Entity; +import com.google.auto.value.AutoValue; import com.google.common.base.Objects; -/** Wraps {@link Entity} to do hashCode/equals based on both the entity's key and its properties. */ -final class ComparableEntity { +/** + * Wraps {@link Entity} for ease of processing in collections. + * + *

Note that the {@link #hashCode}/{@link #equals} methods are based on both the entity's key and + * its properties. + */ +final class EntityWrapper { + private static final String TEST_ENTITY_KIND = "TestEntity"; + private final Entity entity; - ComparableEntity(Entity entity) { + EntityWrapper(Entity entity) { this.entity = entity; } + public Entity getEntity() { + return entity; + } + @Override public boolean equals(Object that) { - if (that instanceof ComparableEntity) { - ComparableEntity thatEntity = (ComparableEntity) that; + if (that instanceof EntityWrapper) { + EntityWrapper thatEntity = (EntityWrapper) that; return entity.equals(thatEntity.entity) && entity.getProperties().equals(thatEntity.entity.getProperties()); } @@ -43,6 +55,26 @@ final class ComparableEntity { @Override public String toString() { - return "ComparableEntity(" + entity + ")"; + return "EntityWrapper(" + entity + ")"; + } + + public static EntityWrapper from(int id, Property... properties) { + Entity entity = new Entity(TEST_ENTITY_KIND, id); + for (Property prop : properties) { + entity.setProperty(prop.name(), prop.value()); + } + return new EntityWrapper(entity); + } + + @AutoValue + abstract static class Property { + + static Property create(String name, Object value) { + return new AutoValue_EntityWrapper_Property(name, value); + } + + abstract String name(); + + abstract Object value(); } } diff --git a/core/src/main/java/google/registry/tools/RecordAccumulator.java b/core/src/main/java/google/registry/tools/RecordAccumulator.java index 855656b96..4c2be48df 100644 --- a/core/src/main/java/google/registry/tools/RecordAccumulator.java +++ b/core/src/main/java/google/registry/tools/RecordAccumulator.java @@ -48,14 +48,14 @@ class RecordAccumulator { return new RecordAccumulator(builder.build()); } - /** Creates an entity set from the current set of raw records. */ - ImmutableSet getComparableEntitySet() { - ImmutableSet.Builder builder = new ImmutableSet.Builder<>(); + /** Creates an {@link EntityWrapper} set from the current set of raw records. */ + ImmutableSet getEntityWrapperSet() { + ImmutableSet.Builder builder = new ImmutableSet.Builder<>(); for (byte[] rawRecord : records) { // Parse the entity proto and create an Entity object from it. EntityProto proto = new EntityProto(); proto.parseFrom(rawRecord); - ComparableEntity entity = new ComparableEntity(EntityTranslator.createFromPb(proto)); + EntityWrapper entity = new EntityWrapper(EntityTranslator.createFromPb(proto)); builder.add(entity); } diff --git a/core/src/test/java/google/registry/tools/CompareDbBackupsTest.java b/core/src/test/java/google/registry/tools/CompareDbBackupsTest.java index db671bb5f..1a7a4aa83 100644 --- a/core/src/test/java/google/registry/tools/CompareDbBackupsTest.java +++ b/core/src/test/java/google/registry/tools/CompareDbBackupsTest.java @@ -19,7 +19,7 @@ import static java.nio.charset.StandardCharsets.UTF_8; import com.google.common.io.Resources; import google.registry.testing.DatastoreEntityExtension; -import google.registry.tools.LevelDbFileBuilder.Property; +import google.registry.tools.EntityWrapper.Property; import java.io.ByteArrayOutputStream; import java.io.File; import java.io.IOException; @@ -71,30 +71,38 @@ public class CompareDbBackupsTest { // Create two directories corresponding to data dumps. File dump1 = tempFs.newFolder("dump1"); LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(dump1, "output-data1")); - builder.addEntityProto( - BASE_ID, - Property.create("eeny", 100L), - Property.create("meeny", 200L), - Property.create("miney", 300L)); - builder.addEntityProto( - BASE_ID + 1, - Property.create("moxey", 100L), - Property.create("minney", 200L), - Property.create("motz", 300L)); + builder.addEntity( + EntityWrapper.from( + BASE_ID, + Property.create("eeny", 100L), + Property.create("meeny", 200L), + Property.create("miney", 300L)) + .getEntity()); + builder.addEntity( + EntityWrapper.from( + BASE_ID + 1, + Property.create("moxey", 100L), + Property.create("minney", 200L), + Property.create("motz", 300L)) + .getEntity()); builder.build(); File dump2 = tempFs.newFolder("dump2"); builder = new LevelDbFileBuilder(new File(dump2, "output-data2")); - builder.addEntityProto( - BASE_ID + 1, - Property.create("moxey", 100L), - Property.create("minney", 200L), - Property.create("motz", 300L)); - builder.addEntityProto( - BASE_ID + 2, - Property.create("blutzy", 100L), - Property.create("fishey", 200L), - Property.create("strutz", 300L)); + builder.addEntity( + EntityWrapper.from( + BASE_ID + 1, + Property.create("moxey", 100L), + Property.create("minney", 200L), + Property.create("motz", 300L)) + .getEntity()); + builder.addEntity( + EntityWrapper.from( + BASE_ID + 2, + Property.create("blutzy", 100L), + Property.create("fishey", 200L), + Property.create("strutz", 300L)) + .getEntity()); builder.build(); CompareDbBackups.main(new String[] {dump1.getCanonicalPath(), dump2.getCanonicalPath()}); diff --git a/core/src/test/java/google/registry/tools/ComparableEntityTest.java b/core/src/test/java/google/registry/tools/EntityWrapperTest.java similarity index 89% rename from core/src/test/java/google/registry/tools/ComparableEntityTest.java rename to core/src/test/java/google/registry/tools/EntityWrapperTest.java index 2369f5d22..16bfabb41 100644 --- a/core/src/test/java/google/registry/tools/ComparableEntityTest.java +++ b/core/src/test/java/google/registry/tools/EntityWrapperTest.java @@ -28,7 +28,7 @@ import org.junit.runner.RunWith; import org.junit.runners.JUnit4; @RunWith(JUnit4.class) -public final class ComparableEntityTest { +public final class EntityWrapperTest { private static final String TEST_ENTITY_KIND = "TestEntity"; private static final int ARBITRARY_KEY_ID = 1001; @@ -63,13 +63,13 @@ public final class ComparableEntityTest { Entity e2 = EntityTranslator.createFromPb(proto2); // Ensure that we have a normalized representation. - ComparableEntity ce1 = new ComparableEntity(e1); - ComparableEntity ce2 = new ComparableEntity(e2); + EntityWrapper ce1 = new EntityWrapper(e1); + EntityWrapper ce2 = new EntityWrapper(e2); assertThat(ce1).isEqualTo(ce2); assertThat(ce1.hashCode()).isEqualTo(ce2.hashCode()); // Ensure that the original entity is equal. - assertThat(new ComparableEntity(entity)).isEqualTo(ce1); + assertThat(new EntityWrapper(entity)).isEqualTo(ce1); } @Test @@ -90,8 +90,8 @@ public final class ComparableEntityTest { Entity e1 = EntityTranslator.createFromPb(proto1); Entity e2 = EntityTranslator.createFromPb(proto2); - ComparableEntity ce1 = new ComparableEntity(e1); - ComparableEntity ce2 = new ComparableEntity(e2); + EntityWrapper ce1 = new EntityWrapper(e1); + EntityWrapper ce2 = new EntityWrapper(e2); assertThat(e1).isEqualTo(e2); // The keys should still be the same. assertThat(ce1).isNotEqualTo(ce2); assertThat(ce1.hashCode()).isNotEqualTo(ce2.hashCode()); @@ -108,15 +108,15 @@ public final class ComparableEntityTest { Entity e1 = EntityTranslator.createFromPb(proto1); Entity e2 = EntityTranslator.createFromPb(proto2); - ComparableEntity ce1 = new ComparableEntity(e1); - ComparableEntity ce2 = new ComparableEntity(e2); + EntityWrapper ce1 = new EntityWrapper(e1); + EntityWrapper ce2 = new EntityWrapper(e2); assertThat(ce1).isNotEqualTo(ce2); assertThat(ce1.hashCode()).isNotEqualTo(ce2.hashCode()); } @Test public void testComparisonAgainstNonComparableEntities() { - ComparableEntity ce = new ComparableEntity(new Entity(TEST_ENTITY_KIND, ARBITRARY_KEY_ID)); + EntityWrapper ce = new EntityWrapper(new Entity(TEST_ENTITY_KIND, ARBITRARY_KEY_ID)); // Note: this has to be "isNotEqualTo()" and not isNotNull() because we want to test the // equals() method and isNotNull() just checks for "ce != null". assertThat(ce).isNotEqualTo(null); diff --git a/core/src/test/java/google/registry/tools/LevelDbFileBuilder.java b/core/src/test/java/google/registry/tools/LevelDbFileBuilder.java index e7837e9dc..a8a35c5c6 100644 --- a/core/src/test/java/google/registry/tools/LevelDbFileBuilder.java +++ b/core/src/test/java/google/registry/tools/LevelDbFileBuilder.java @@ -19,7 +19,6 @@ import static google.registry.tools.LevelDbLogReader.HEADER_SIZE; import com.google.appengine.api.datastore.Entity; import com.google.appengine.api.datastore.EntityTranslator; -import com.google.auto.value.AutoValue; import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; import google.registry.tools.LevelDbLogReader.ChunkType; import java.io.File; @@ -28,30 +27,19 @@ import java.io.FileOutputStream; import java.io.IOException; /** Utility class for building a leveldb logfile. */ -final class LevelDbFileBuilder { - private static final String TEST_ENTITY_KIND = "TestEntity"; - +public final class LevelDbFileBuilder { private final FileOutputStream out; private byte[] currentBlock = new byte[BLOCK_SIZE]; // Write position in the current block. private int currentPos = 0; - LevelDbFileBuilder(File file) throws FileNotFoundException { + public LevelDbFileBuilder(File file) throws FileNotFoundException { out = new FileOutputStream(file); } - /** - * Adds a record containing a new entity protobuf to the file. - * - *

Returns the ComparableEntity object rather than "this" so that we can check for the presence - * of the entity in the result set. - */ - ComparableEntity addEntityProto(int id, Property... properties) throws IOException { - Entity entity = new Entity(TEST_ENTITY_KIND, id); - for (Property prop : properties) { - entity.setProperty(prop.name(), prop.value()); - } + /** Adds an {@link Entity Datastore Entity object} to the leveldb log file. */ + LevelDbFileBuilder addEntity(Entity entity) throws IOException { EntityProto proto = EntityTranslator.convertToPb(entity); byte[] protoBytes = proto.toByteArray(); if (protoBytes.length > BLOCK_SIZE - (currentPos + HEADER_SIZE)) { @@ -61,7 +49,7 @@ final class LevelDbFileBuilder { } currentPos = LevelDbUtil.addRecord(currentBlock, currentPos, ChunkType.FULL, protoBytes); - return new ComparableEntity(entity); + return this; } /** Writes all remaining data and closes the block. */ @@ -69,15 +57,4 @@ final class LevelDbFileBuilder { out.write(currentBlock); out.close(); } - - @AutoValue - abstract static class Property { - static Property create(String name, Object value) { - return new AutoValue_LevelDbFileBuilder_Property(name, value); - } - - abstract String name(); - - abstract Object value(); - } } diff --git a/core/src/test/java/google/registry/tools/LevelDbFileBuilderTest.java b/core/src/test/java/google/registry/tools/LevelDbFileBuilderTest.java index 75c876efd..693c7ebce 100644 --- a/core/src/test/java/google/registry/tools/LevelDbFileBuilderTest.java +++ b/core/src/test/java/google/registry/tools/LevelDbFileBuilderTest.java @@ -15,13 +15,17 @@ package google.registry.tools; import static com.google.common.truth.Truth.assertThat; +import static google.registry.model.ofy.ObjectifyService.ofy; +import static google.registry.persistence.transaction.TransactionManagerFactory.tm; import com.google.appengine.api.datastore.Entity; import com.google.appengine.api.datastore.EntityTranslator; import com.google.common.collect.ImmutableList; import com.google.storage.onestore.v3.OnestoreEntity.EntityProto; +import google.registry.model.contact.ContactResource; import google.registry.testing.AppEngineRule; -import google.registry.tools.LevelDbFileBuilder.Property; +import google.registry.testing.DatastoreHelper; +import google.registry.tools.EntityWrapper.Property; import java.io.File; import java.io.IOException; import org.junit.Rule; @@ -45,19 +49,18 @@ public class LevelDbFileBuilderTest { File subdir = tempFs.newFolder("folder"); File logFile = new File(subdir, "testfile"); LevelDbFileBuilder builder = new LevelDbFileBuilder(logFile); - ComparableEntity entity = - builder.addEntityProto( + EntityWrapper entity = + EntityWrapper.from( BASE_ID, Property.create("first", 100L), Property.create("second", 200L)); + builder.addEntity(entity.getEntity()); builder.build(); ImmutableList records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath())); assertThat(records).hasSize(1); // Reconstitute an entity, make sure that what we've got is the same as what we started with. - EntityProto proto = new EntityProto(); - proto.parseFrom(records.get(0)); - Entity materializedEntity = EntityTranslator.createFromPb(proto); - assertThat(new ComparableEntity(materializedEntity)).isEqualTo(entity); + Entity materializedEntity = rawRecordToEntity(records.get(0)); + assertThat(new EntityWrapper(materializedEntity)).isEqualTo(entity); } @Test @@ -68,25 +71,50 @@ public class LevelDbFileBuilderTest { // Generate enough records to cross a block boundary. These records end up being around 80 // bytes, so 1000 works. - ImmutableList.Builder originalEntitiesBuilder = new ImmutableList.Builder<>(); + ImmutableList.Builder originalEntitiesBuilder = new ImmutableList.Builder<>(); for (int i = 0; i < 1000; ++i) { - ComparableEntity entity = - builder.addEntityProto( + EntityWrapper entity = + EntityWrapper.from( BASE_ID + i, Property.create("first", 100L), Property.create("second", 200L)); + builder.addEntity(entity.getEntity()); originalEntitiesBuilder.add(entity); } builder.build(); - ImmutableList originalEntities = originalEntitiesBuilder.build(); + ImmutableList originalEntities = originalEntitiesBuilder.build(); ImmutableList records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath())); assertThat(records).hasSize(1000); int index = 0; for (byte[] record : records) { - EntityProto proto = new EntityProto(); - proto.parseFrom(record); - Entity materializedEntity = EntityTranslator.createFromPb(proto); - assertThat(new ComparableEntity(materializedEntity)).isEqualTo(originalEntities.get(index)); + Entity materializedEntity = rawRecordToEntity(record); + assertThat(new EntityWrapper(materializedEntity)).isEqualTo(originalEntities.get(index)); ++index; } } + + @Test + public void testOfyEntityWrite() throws Exception { + File subdir = tempFs.newFolder("folder"); + File logFile = new File(subdir, "testfile"); + LevelDbFileBuilder builder = new LevelDbFileBuilder(logFile); + + ContactResource contact = DatastoreHelper.newContactResource("contact"); + builder.addEntity(tm().transact(() -> ofy().save().toEntity(contact))); + builder.build(); + + ImmutableList records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath())); + assertThat(records).hasSize(1); + ContactResource ofyEntity = rawRecordToOfyEntity(records.get(0), ContactResource.class); + assertThat(ofyEntity.getContactId()).isEqualTo(contact.getContactId()); + } + + private static Entity rawRecordToEntity(byte[] record) { + EntityProto proto = new EntityProto(); + proto.parseFrom(record); + return EntityTranslator.createFromPb(proto); + } + + private static T rawRecordToOfyEntity(byte[] record, Class expectedType) { + return expectedType.cast(ofy().load().fromEntity(rawRecordToEntity(record))); + } } diff --git a/core/src/test/java/google/registry/tools/RecordAccumulatorTest.java b/core/src/test/java/google/registry/tools/RecordAccumulatorTest.java index 853807029..631eff7a6 100644 --- a/core/src/test/java/google/registry/tools/RecordAccumulatorTest.java +++ b/core/src/test/java/google/registry/tools/RecordAccumulatorTest.java @@ -18,7 +18,7 @@ import static com.google.common.truth.Truth.assertThat; import com.google.common.collect.ImmutableSet; import google.registry.testing.AppEngineRule; -import google.registry.tools.LevelDbFileBuilder.Property; +import google.registry.tools.EntityWrapper.Property; import java.io.File; import java.io.IOException; import org.junit.Rule; @@ -44,39 +44,44 @@ public class RecordAccumulatorTest { // Note that we need to specify property values as "Long" for property comparisons to work // correctly because that's how they are deserialized from protos. - ComparableEntity e1 = - builder.addEntityProto( + EntityWrapper e1 = + EntityWrapper.from( BASE_ID, Property.create("eeny", 100L), Property.create("meeny", 200L), Property.create("miney", 300L)); - ComparableEntity e2 = - builder.addEntityProto( + builder.addEntity(e1.getEntity()); + EntityWrapper e2 = + EntityWrapper.from( BASE_ID + 1, Property.create("eeny", 100L), Property.create("meeny", 200L), Property.create("miney", 300L)); + builder.addEntity(e2.getEntity()); builder.build(); builder = new LevelDbFileBuilder(new File(subdir, "data2")); // Duplicate of the record in the other file. - builder.addEntityProto( - BASE_ID, - Property.create("eeny", 100L), - Property.create("meeny", 200L), - Property.create("miney", 300L)); + builder.addEntity( + EntityWrapper.from( + BASE_ID, + Property.create("eeny", 100L), + Property.create("meeny", 200L), + Property.create("miney", 300L)) + .getEntity()); - ComparableEntity e3 = - builder.addEntityProto( + EntityWrapper e3 = + EntityWrapper.from( BASE_ID + 2, Property.create("moxy", 100L), Property.create("fruvis", 200L), Property.create("cortex", 300L)); + builder.addEntity(e3.getEntity()); builder.build(); - ImmutableSet entities = - RecordAccumulator.readDirectory(subdir, any -> true).getComparableEntitySet(); + ImmutableSet entities = + RecordAccumulator.readDirectory(subdir, any -> true).getEntityWrapperSet(); assertThat(entities).containsExactly(e1, e2, e3); } }