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.

core/src/main/java/google/registry/tools/CompareDbBackups.java

@@ -39,16 +39,14 @@ class CompareDbBackups {
       return;
     }
 
-    ImmutableSet<ComparableEntity> entities1 =
+    ImmutableSet<EntityWrapper> entities1 =
-        RecordAccumulator.readDirectory(new File(args[0]), DATA_FILE_MATCHER)
+        RecordAccumulator.readDirectory(new File(args[0]), DATA_FILE_MATCHER).getEntityWrapperSet();
-            .getComparableEntitySet();
+    ImmutableSet<EntityWrapper> entities2 =
-    ImmutableSet<ComparableEntity> entities2 =
+        RecordAccumulator.readDirectory(new File(args[1]), DATA_FILE_MATCHER).getEntityWrapperSet();
-        RecordAccumulator.readDirectory(new File(args[1]), DATA_FILE_MATCHER)
-            .getComparableEntitySet();
 
     // Calculate the entities added and removed.
-    SetView<ComparableEntity> added = Sets.difference(entities2, entities1);
+    SetView<EntityWrapper> added = Sets.difference(entities2, entities1);
-    SetView<ComparableEntity> removed = Sets.difference(entities1, entities2);
+    SetView<EntityWrapper> 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);
       }
     }

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.
+ *
+ * <p>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) {
+    if (that instanceof EntityWrapper) {
-      ComparableEntity thatEntity = (ComparableEntity) that;
+      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();
   }
 }

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. */
+  /** Creates an {@link EntityWrapper} set from the current set of raw records. */
-  ImmutableSet<ComparableEntity> getComparableEntitySet() {
+  ImmutableSet<EntityWrapper> getEntityWrapperSet() {
-    ImmutableSet.Builder<ComparableEntity> builder = new ImmutableSet.Builder<>();
+    ImmutableSet.Builder<EntityWrapper> 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);
     }

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(
+    builder.addEntity(
-        BASE_ID,
+        EntityWrapper.from(
-        Property.create("eeny", 100L),
+                BASE_ID,
-        Property.create("meeny", 200L),
+                Property.create("eeny", 100L),
-        Property.create("miney", 300L));
+                Property.create("meeny", 200L),
-    builder.addEntityProto(
+                Property.create("miney", 300L))
-        BASE_ID + 1,
+            .getEntity());
-        Property.create("moxey", 100L),
+    builder.addEntity(
-        Property.create("minney", 200L),
+        EntityWrapper.from(
-        Property.create("motz", 300L));
+                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(
+    builder.addEntity(
-        BASE_ID + 1,
+        EntityWrapper.from(
-        Property.create("moxey", 100L),
+                BASE_ID + 1,
-        Property.create("minney", 200L),
+                Property.create("moxey", 100L),
-        Property.create("motz", 300L));
+                Property.create("minney", 200L),
-    builder.addEntityProto(
+                Property.create("motz", 300L))
-        BASE_ID + 2,
+            .getEntity());
-        Property.create("blutzy", 100L),
+    builder.addEntity(
-        Property.create("fishey", 200L),
+        EntityWrapper.from(
-        Property.create("strutz", 300L));
+                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()});

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);
+    EntityWrapper ce1 = new EntityWrapper(e1);
-    ComparableEntity ce2 = new ComparableEntity(e2);
+    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);
+    EntityWrapper ce1 = new EntityWrapper(e1);
-    ComparableEntity ce2 = new ComparableEntity(e2);
+    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);
+    EntityWrapper ce1 = new EntityWrapper(e1);
-    ComparableEntity ce2 = new ComparableEntity(e2);
+    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);

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 {
+public final class LevelDbFileBuilder {
-  private static final String TEST_ENTITY_KIND = "TestEntity";
-
   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 an {@link Entity Datastore Entity object} to the leveldb log file. */
-   * Adds a record containing a new entity protobuf to the file.
+  LevelDbFileBuilder addEntity(Entity entity) throws IOException {
-   *
-   * <p>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());
-    }
     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();
-  }
 }

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 =
+    EntityWrapper entity =
-        builder.addEntityProto(
+        EntityWrapper.from(
             BASE_ID, Property.create("first", 100L), Property.create("second", 200L));
+    builder.addEntity(entity.getEntity());
     builder.build();
 
     ImmutableList<byte[]> 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();
+    Entity materializedEntity = rawRecordToEntity(records.get(0));
-    proto.parseFrom(records.get(0));
+    assertThat(new EntityWrapper(materializedEntity)).isEqualTo(entity);
-    Entity materializedEntity = EntityTranslator.createFromPb(proto);
-    assertThat(new ComparableEntity(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<ComparableEntity> originalEntitiesBuilder = new ImmutableList.Builder<>();
+    ImmutableList.Builder<EntityWrapper> originalEntitiesBuilder = new ImmutableList.Builder<>();
     for (int i = 0; i < 1000; ++i) {
-      ComparableEntity entity =
+      EntityWrapper entity =
-          builder.addEntityProto(
+          EntityWrapper.from(
               BASE_ID + i, Property.create("first", 100L), Property.create("second", 200L));
+      builder.addEntity(entity.getEntity());
       originalEntitiesBuilder.add(entity);
     }
     builder.build();
-    ImmutableList<ComparableEntity> originalEntities = originalEntitiesBuilder.build();
+    ImmutableList<EntityWrapper> originalEntities = originalEntitiesBuilder.build();
 
     ImmutableList<byte[]> records = ImmutableList.copyOf(LevelDbLogReader.from(logFile.getPath()));
     assertThat(records).hasSize(1000);
     int index = 0;
     for (byte[] record : records) {
-      EntityProto proto = new EntityProto();
+      Entity materializedEntity = rawRecordToEntity(record);
-      proto.parseFrom(record);
+      assertThat(new EntityWrapper(materializedEntity)).isEqualTo(originalEntities.get(index));
-      Entity materializedEntity = EntityTranslator.createFromPb(proto);
-      assertThat(new ComparableEntity(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<byte[]> 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> T rawRecordToOfyEntity(byte[] record, Class<T> expectedType) {
+    return expectedType.cast(ofy().load().fromEntity(rawRecordToEntity(record)));
+  }
 }

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 =
+    EntityWrapper e1 =
-        builder.addEntityProto(
+        EntityWrapper.from(
             BASE_ID,
             Property.create("eeny", 100L),
             Property.create("meeny", 200L),
             Property.create("miney", 300L));
-    ComparableEntity e2 =
+    builder.addEntity(e1.getEntity());
-        builder.addEntityProto(
+    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(
+    builder.addEntity(
-        BASE_ID,
+        EntityWrapper.from(
-        Property.create("eeny", 100L),
+                BASE_ID,
-        Property.create("meeny", 200L),
+                Property.create("eeny", 100L),
-        Property.create("miney", 300L));
+                Property.create("meeny", 200L),
+                Property.create("miney", 300L))
+            .getEntity());
 
-    ComparableEntity e3 =
+    EntityWrapper e3 =
-        builder.addEntityProto(
+        EntityWrapper.from(
             BASE_ID + 2,
             Property.create("moxy", 100L),
             Property.create("fruvis", 200L),
             Property.create("cortex", 300L));
+    builder.addEntity(e3.getEntity());
     builder.build();
 
-    ImmutableSet<ComparableEntity> entities =
+    ImmutableSet<EntityWrapper> entities =
-        RecordAccumulator.readDirectory(subdir, any -> true).getComparableEntitySet();
+        RecordAccumulator.readDirectory(subdir, any -> true).getEntityWrapperSet();
     assertThat(entities).containsExactly(e1, e2, e3);
   }
 }