Upgrade CompareDbBackup for Datastore V3 (#543)

* Upgrade CompareDbBackup for Datastore V3

Upgrade the CompareDbBackup class to work with latest
Datastore backup directory structure.

Also fixed a few unrelated minor issues:
- Remaining cases of improper use of System.setOut
- Wrong import order in one class

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

@@ -18,9 +18,20 @@ import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.Sets;
 import com.google.common.collect.Sets.SetView;
 import java.io.File;
+import java.util.function.Predicate;
 
-/** Compare two database backups. */
+/**
+ * Compares two Datastore backups in V3 format on local file system. This is for use in tests and
+ * experiments with small data sizes.
+ *
+ * <p>This utility only supports the current Datastore backup format (version 3). A backup is a
+ * two-level directory hierarchy with data files in level-db format (output-*) and Datastore
+ * metadata files (*.export_metadata).
+ */
 class CompareDbBackups {
+  private static final String DS_V3_BACKUP_FILE_PREFIX = "output-";
+  private static final Predicate<File> DATA_FILE_MATCHER =
+      file -> file.isFile() && file.getName().startsWith(DS_V3_BACKUP_FILE_PREFIX);
 
   public static void main(String[] args) {
     if (args.length != 2) {
@@ -29,9 +40,13 @@ class CompareDbBackups {
     }
 
     ImmutableSet<ComparableEntity> entities1 =
-        new RecordAccumulator().readDirectory(new File(args[0])).getComparableEntitySet();
+        new RecordAccumulator()
+            .readDirectory(new File(args[0]), DATA_FILE_MATCHER)
+            .getComparableEntitySet();
     ImmutableSet<ComparableEntity> entities2 =
-        new RecordAccumulator().readDirectory(new File(args[1])).getComparableEntitySet();
+        new RecordAccumulator()
+            .readDirectory(new File(args[1]), DATA_FILE_MATCHER)
+            .getComparableEntitySet();
 
     // Calculate the entities added and removed.
     SetView<ComparableEntity> added = Sets.difference(entities2, entities1);
@@ -54,6 +69,10 @@ class CompareDbBackups {
         System.out.println(entity);
       }
     }
+
+    if (added.isEmpty() && removed.isEmpty()) {
+      System.out.printf("\nBoth sets have the same %d entities.\n", entities1.size());
+    }
   }
 
   /** Print out multi-line text in a pretty ASCII header frame. */

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

@@ -20,17 +20,18 @@ import com.google.storage.onestore.v3.OnestoreEntity.EntityProto;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
+import java.util.function.Predicate;
 
-/** Utility class that accumulates Entity records from level db files. */
+/** Accumulates Entity records from level db files under a directory hierarchy. */
 class RecordAccumulator {
   private final LevelDbLogReader reader = new LevelDbLogReader();
 
   /** Recursively reads all records in the directory. */
-  public final RecordAccumulator readDirectory(File dir) {
+  public final RecordAccumulator readDirectory(File dir, Predicate<File> fileMatcher) {
     for (File child : dir.listFiles()) {
       if (child.isDirectory()) {
-        readDirectory(child);
+        readDirectory(child, fileMatcher);
-      } else if (child.isFile()) {
+      } else if (fileMatcher.test(child)) {
         try {
           reader.readFrom(new FileInputStream(child));
         } catch (IOException e) {

core/src/test/java/google/registry/persistence/VKeyTest.java

@@ -17,11 +17,11 @@ import static com.google.common.truth.Truth.assertThat;
 
 import com.googlecode.objectify.Key;
 import google.registry.testing.AppEngineRule;
+import google.registry.testing.TestObject;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.runner.RunWith;
 import org.junit.runners.JUnit4;
-import google.registry.testing.TestObject;
 
 @RunWith(JUnit4.class)
 public class VKeyTest {

core/src/test/java/google/registry/tools/CompareDbBackupsTest.java

@@ -17,11 +17,15 @@ package google.registry.tools;
 import static com.google.common.truth.Truth.assertThat;
 import static java.nio.charset.StandardCharsets.UTF_8;
 
+import com.google.common.io.Resources;
 import google.registry.testing.AppEngineRule;
 import google.registry.tools.LevelDbFileBuilder.Property;
 import java.io.ByteArrayOutputStream;
 import java.io.File;
 import java.io.PrintStream;
+import java.net.URL;
+import org.junit.After;
+import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.TemporaryFolder;
@@ -35,18 +39,38 @@ public class CompareDbBackupsTest {
 
   // Capture standard output.
   private final ByteArrayOutputStream stdout = new ByteArrayOutputStream();
+  private PrintStream orgStdout;
 
   @Rule public final TemporaryFolder tempFs = new TemporaryFolder();
 
   @Rule
   public final AppEngineRule appEngine = AppEngineRule.builder().withDatastoreAndCloudSql().build();
 
+  @Before
+  public void before() {
+    orgStdout = System.out;
+    System.setOut(new PrintStream(stdout));
+  }
+
+  @After
+  public void after() {
+    System.setOut(orgStdout);
+  }
+
   @Test
-  public void testCommand() throws Exception {
+  public void testLoadBackup() {
+    URL backupRootFolder = Resources.getResource("google/registry/tools/datastore-export");
+    CompareDbBackups.main(new String[] {backupRootFolder.getPath(), backupRootFolder.getPath()});
+    String output = new String(stdout.toByteArray(), UTF_8);
+    assertThat(output).containsMatch("Both sets have the same 41 entities");
+  }
+
+  @Test
+  public void testCompareBackups() throws Exception {
 
     // Create two directories corresponding to data dumps.
     File dump1 = tempFs.newFolder("dump1");
-    LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(dump1, "data1"));
+    LevelDbFileBuilder builder = new LevelDbFileBuilder(new File(dump1, "output-data1"));
     builder.addEntityProto(
         BASE_ID,
         Property.create("eeny", 100L),
@@ -60,7 +84,7 @@ public class CompareDbBackupsTest {
     builder.build();
 
     File dump2 = tempFs.newFolder("dump2");
-    builder = new LevelDbFileBuilder(new File(dump2, "data2"));
+    builder = new LevelDbFileBuilder(new File(dump2, "output-data2"));
     builder.addEntityProto(
         BASE_ID + 1,
         Property.create("moxey", 100L),
@@ -73,7 +97,6 @@ public class CompareDbBackupsTest {
         Property.create("strutz", 300L));
     builder.build();
 
-    System.setOut(new PrintStream(stdout));
     CompareDbBackups.main(new String[] {dump1.getCanonicalPath(), dump2.getCanonicalPath()});
     String output = new String(stdout.toByteArray(), UTF_8);
     assertThat(output)

core/src/test/java/google/registry/tools/GhostrydeCommandTest.java

@@ -23,11 +23,8 @@ import google.registry.rde.Ghostryde;
 import google.registry.testing.BouncyCastleProviderRule;
 import google.registry.testing.FakeKeyringModule;
 import google.registry.testing.InjectRule;
-import java.io.ByteArrayOutputStream;
-import java.io.PrintStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
-import org.junit.After;
 import org.junit.Before;
 import org.junit.Rule;
 import org.junit.Test;
@@ -61,19 +58,12 @@ public class GhostrydeCommandTest extends CommandTestCase<GhostrydeCommand> {
   public final BouncyCastleProviderRule bouncy = new BouncyCastleProviderRule();
 
   private Keyring keyring;
-  private PrintStream orgStdout;
 
   @Before
   public void before() {
     keyring = new FakeKeyringModule().get();
     command.rdeStagingDecryptionKey = keyring::getRdeStagingDecryptionKey;
     command.rdeStagingEncryptionKey = keyring::getRdeStagingEncryptionKey;
-    orgStdout = System.out;
-  }
-
-  @After
-  public void after() {
-    System.setOut(orgStdout);
   }
 
   @Test
@@ -153,9 +143,7 @@ public class GhostrydeCommandTest extends CommandTestCase<GhostrydeCommand> {
     Path inFile = tmpDir.newFolder().toPath().resolve("atrain.ghostryde");
     Files.write(
         inFile, Ghostryde.encode(SONG_BY_CHRISTINA_ROSSETTI, keyring.getRdeStagingEncryptionKey()));
-    ByteArrayOutputStream out = new ByteArrayOutputStream();
-    System.setOut(new PrintStream(out));
     runCommand("--decrypt", "--input=" + inFile);
-    assertThat(out.toByteArray()).isEqualTo(SONG_BY_CHRISTINA_ROSSETTI);
+    assertThat(getStdoutAsString().getBytes(UTF_8)).isEqualTo(SONG_BY_CHRISTINA_ROSSETTI);
   }
 }

core/src/test/java/google/registry/tools/RecordAccumulatorTest.java

@@ -76,7 +76,7 @@ public class RecordAccumulatorTest {
     builder.build();
 
     ImmutableSet<ComparableEntity> entities =
-        new RecordAccumulator().readDirectory(subdir).getComparableEntitySet();
+        new RecordAccumulator().readDirectory(subdir, any -> true).getComparableEntitySet();
     assertThat(entities).containsExactly(e1, e2, e3);
   }
 }