Create an entity record accumulator

RecordAccumulator builds a set of datastore Entity records from a set of
leveldb logfiles in a directory tree (which is how we receive them for
database backup/restore testing).

This CL also refactors some of the logfile test code out of
LevelDbLogReaderTest so that we can reuse it for building test logs.

-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=167313553

java/google/registry/tools/RecordAccumulator.java

@@ -0,0 +1,59 @@
+// Copyright 2017 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.tools;
+
+import com.google.appengine.api.datastore.EntityTranslator;
+import com.google.common.collect.ImmutableSet;
+import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.IOException;
+
+/** Utility class that accumulates Entity records from level db files. */
+class RecordAccumulator {
+  private final LevelDbLogReader reader = new LevelDbLogReader();
+
+  /** Recursively reads all records in the directory. */
+  public final RecordAccumulator readDirectory(File dir) {
+    for (File child : dir.listFiles()) {
+      if (child.isDirectory()) {
+        readDirectory(child);
+      } else if (child.isFile()) {
+        try {
+          reader.readFrom(new FileInputStream(child));
+        } catch (IOException e) {
+          throw new RuntimeException("IOException reading from file: " + child, e);
+        }
+      }
+    }
+
+    return this;
+  }
+
+  /** Creates an entity set from the current set of raw records. */
+  ImmutableSet<ComparableEntity> getComparableEntitySet() {
+    ImmutableSet.Builder<ComparableEntity> builder = new ImmutableSet.Builder<>();
+    for (byte[] rawRecord : reader.getRecords()) {
+      // 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));
+
+      builder.add(entity);
+    }
+
+    return builder.build();
+  }
+}

javatests/google/registry/tools/BUILD

@@ -41,6 +41,7 @@ java_library(
         "@com_google_api_client",
         "@com_google_appengine_api_1_0_sdk//:testonly",
         "@com_google_appengine_remote_api//:link",
+        "@com_google_auto_value",
         "@com_google_code_findbugs_jsr305",
         "@com_google_guava",
         "@com_google_http_client",

javatests/google/registry/tools/LevelDbLogReaderTest.java

@@ -15,10 +15,12 @@
 package google.registry.tools;
 
 import static com.google.common.truth.Truth.assertThat;
-import static google.registry.tools.LevelDbLogReader.ChunkType;
+import static google.registry.tools.LevelDbUtil.MAX_RECORD;
+import static google.registry.tools.LevelDbUtil.addRecord;
 
 import com.google.common.collect.ImmutableList;
 import com.google.common.primitives.Bytes;
+import google.registry.tools.LevelDbLogReader.ChunkType;
 import java.io.ByteArrayInputStream;
 import java.io.IOException;
 import java.util.List;
@@ -30,8 +32,6 @@ import org.junit.runners.JUnit4;
 @RunWith(JUnit4.class)
 public final class LevelDbLogReaderTest {
 
-  private static final int MAX_RECORD = LevelDbLogReader.BLOCK_SIZE - LevelDbLogReader.HEADER_SIZE;
-
   // Size of the test record.  Any value < 256 will do.
   private static final int TEST_RECORD_SIZE = 231;
 
@@ -40,51 +40,12 @@ public final class LevelDbLogReaderTest {
   private static final int MAX_TEST_RECORD_OFFSET =
       LevelDbLogReader.BLOCK_SIZE - (LevelDbLogReader.HEADER_SIZE + TEST_RECORD_SIZE);
 
-  /**
-   * Adds a record of bytes of 'val' of the given size to bytes.
-   *
-   * <p>This currently doesn't write a real checksum since we're not doing anything with that in the
-   * leveldb reader.
-   *
-   * <p>Returns the new offset for the next block.
-   */
-  private static int addRecord(
-      byte[] bytes, int pos, ChunkType type, int size, int val) {
-
-    // Write a bogus checksum.
-    for (int i = 0; i < 4; ++i) {
-      bytes[pos++] = -1;
-    }
-
-    // Write size and type.
-    bytes[pos++] = (byte) size;
-    bytes[pos++] = (byte) (size >> 8);
-    bytes[pos++] = (byte) type.getCode();
-
-    // Write "size" bytes of data.
-    for (int i = 0; i < size; ++i) {
-      bytes[pos + i] = (byte) val;
-
-      // Swap the least significant bytes in val so we can have more than 256 different same-sized
-      // records.
-      val = (val >> 8) | ((val & 0xff) << 8);
-    }
-
-    return pos + size;
-  }
-
   private TestBlock makeBlockOfRepeatingBytes(int startVal) {
     byte[] block = new byte[LevelDbLogReader.BLOCK_SIZE];
     int pos = 0;
     int recordCount = 0;
     while (pos < MAX_TEST_RECORD_OFFSET) {
-      pos =
-          addRecord(
-              block,
-              pos,
-              ChunkType.FULL,
-              TEST_RECORD_SIZE,
-              0xffff & (pos + startVal));
+      pos = addRecord(block, pos, ChunkType.FULL, TEST_RECORD_SIZE, 0xffff & (pos + startVal));
       ++recordCount;
     }
     return new TestBlock(block, recordCount);

javatests/google/registry/tools/LevelDbUtil.java

@@ -0,0 +1,87 @@
+// Copyright 2017 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.tools;
+
+import static com.google.common.base.Preconditions.checkArgument;
+import google.registry.tools.LevelDbLogReader.ChunkType;
+
+class LevelDbUtil {
+
+  public static final int MAX_RECORD = LevelDbLogReader.BLOCK_SIZE - LevelDbLogReader.HEADER_SIZE;
+
+  /** Adds a new record header to "bytes" at "pos", returns the new position. */
+  private static int addRecordHeader(byte[] bytes, int pos, ChunkType type, int size) {
+    // Write a bogus checksum.
+    for (int i = 0; i < 4; ++i) {
+      bytes[pos++] = -1;
+    }
+
+    // Write size and type.
+    bytes[pos++] = (byte) size;
+    bytes[pos++] = (byte) (size >> 8);
+    bytes[pos++] = (byte) type.getCode();
+
+    return pos;
+  }
+
+  /**
+   * Adds a record of repeating bytes of 'val' of the given size to bytes at pos.
+   *
+   * <p>Writes the two least-significant bytes of 'val', alternating the order of them. So if the
+   * value of 'val' is 0x1234, writes 0x12 0x34 0x34 0x12 0x12 ... If the value is greater than
+   * 0xffff, it will be truncated to 16 bits.
+   *
+   * <p>This currently doesn't write a real checksum since we're not doing anything with that in the
+   * leveldb reader.
+   *
+   * <p>Returns the new offset for the next block.
+   */
+  static int addRecord(byte[] bytes, int pos, ChunkType type, int size, int val) {
+    pos = addRecordHeader(bytes, pos, type, size);
+
+    // Write "size" bytes of data.
+    for (int i = 0; i < size; ++i) {
+      bytes[pos + i] = (byte) val;
+
+      // Swap the least significant bytes in val so we can have more than 256 different same-sized
+      // records.
+      val = ((val >> 8) & 0xff) | ((val & 0xff) << 8);
+    }
+
+    return pos + size;
+  }
+
+  /**
+   * Adds a record containing "data" to "bytes".
+   *
+   * <p>This currently doesn't write a real checksum since we're not doing anything with that in the
+   * leveldb reader.
+   *
+   * <p>Returns the new offset for the next block.
+   */
+  static int addRecord(byte[] bytes, int pos, ChunkType type, byte[] data) {
+    checkArgument(
+        data.length < MAX_RECORD, "Record length (%d) > max record size (%d)", data.length,
+        MAX_RECORD);
+    pos = addRecordHeader(bytes, pos, type, data.length);
+
+    // Write the contents of "data".
+    for (int i = 0; i < data.length; ++i) {
+      bytes[pos + i] = data[i];
+    }
+
+    return pos + data.length;
+  }
+}

javatests/google/registry/tools/RecordAccumulatorTest.java

@@ -0,0 +1,140 @@
+// Copyright 2017 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.tools;
+
+import static com.google.common.truth.Truth.assertThat;
+import static google.registry.tools.LevelDbLogReader.BLOCK_SIZE;
+
+import com.google.appengine.api.datastore.Entity;
+import com.google.appengine.api.datastore.EntityTranslator;
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableSet;
+import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
+import google.registry.testing.AppEngineRule;
+import google.registry.tools.LevelDbLogReader.ChunkType;
+import java.io.File;
+import java.io.FileNotFoundException;
+import java.io.FileOutputStream;
+import java.io.IOException;
+import org.junit.Rule;
+import org.junit.Test;
+import org.junit.rules.TemporaryFolder;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+@RunWith(JUnit4.class)
+public class RecordAccumulatorTest {
+
+  private static final int BASE_ID = 1001;
+  private static final String TEST_ENTITY_KIND = "TestEntity";
+
+  @Rule public final TemporaryFolder tempFs = new TemporaryFolder();
+  @Rule public final AppEngineRule appEngine = AppEngineRule.builder().withDatastore().build();
+
+  @Test
+  public void testReadDirectory() throws IOException {
+    File subdir = tempFs.newFolder("folder");
+    LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(subdir, "data1"));
+
+    // 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(
+            BASE_ID,
+            Property.create("eeny", 100L),
+            Property.create("meeny", 200L),
+            Property.create("miney", 300L));
+    ComparableEntity e2 =
+        builder.addEntityProto(
+            BASE_ID + 1,
+            Property.create("eeny", 100L),
+            Property.create("meeny", 200L),
+            Property.create("miney", 300L));
+    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));
+
+    ComparableEntity e3 =
+        builder.addEntityProto(
+            BASE_ID + 2,
+            Property.create("moxy", 100L),
+            Property.create("fruvis", 200L),
+            Property.create("cortex", 300L));
+    builder.build();
+
+    ImmutableSet<ComparableEntity> entities =
+        new RecordAccumulator().readDirectory(subdir).getComparableEntitySet();
+    assertThat(entities).containsExactly(e1, e2, e3);
+  }
+
+  /** Utility class for building a leveldb logfile. */
+  private static 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 {
+      out = new FileOutputStream(file);
+    }
+
+    /**
+     * Adds a record containing a new entity protobuf to the file.
+     *
+     * <p>Returns the ComparableEntity object rather than "this" so that we can check for the
+     * presence of the entity in the result set.
+     */
+    private 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) {
+        out.write(currentBlock);
+        currentBlock = new byte[BLOCK_SIZE];
+      }
+
+      currentPos = LevelDbUtil.addRecord(currentBlock, currentPos, ChunkType.FULL, protoBytes);
+      return new ComparableEntity(entity);
+    }
+
+    /** Writes all remaining data and closes the block. */
+    void build() throws IOException {
+      out.write(currentBlock);
+      out.close();
+    }
+  }
+
+  @AutoValue
+  abstract static class Property {
+    static Property create(String name, Object value) {
+      return new AutoValue_RecordAccumulatorTest_Property(name, value);
+    }
+
+    abstract String name();
+
+    abstract Object value();
+  }
+}