Add a program to import from a protocol buffer dump file

This is the format used by database exports from googlesource.com, and
can also act as a general database-independent backup/transfer format.

Change-Id: I36e0390b839e3adb44bbe045ad8ca54d05be190c

gerrit-pgm/BUCK

@@ -106,7 +106,9 @@ java_library(
     '//gerrit-solr:solr',
     '//lib:args4j',
     '//lib:gwtorm',
+    '//lib:protobuf',
     '//lib:servlet-api-3_1',
+    '//lib/auto:auto-value',
     '//lib/prolog:cafeteria',
     '//lib/prolog:compiler',
     '//lib/prolog:runtime',

gerrit-pgm/src/main/java/com/google/gerrit/pgm/ProtobufImport.java

@@ -0,0 +1,157 @@
+// Copyright (C) 2015 The Android Open Source Project
+//
+// 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 com.google.gerrit.pgm;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.common.base.Preconditions.checkState;
+import static com.google.gerrit.server.schema.DataSourceProvider.Context.SINGLE_USER;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.Iterables;
+import com.google.gerrit.lifecycle.LifecycleManager;
+import com.google.gerrit.pgm.util.RuntimeShutdown;
+import com.google.gerrit.pgm.util.SiteProgram;
+import com.google.gerrit.reviewdb.server.ReviewDb;
+import com.google.gwtorm.protobuf.CodecFactory;
+import com.google.gwtorm.protobuf.ProtobufCodec;
+import com.google.gwtorm.schema.RelationModel;
+import com.google.gwtorm.schema.java.JavaSchemaModel;
+import com.google.gwtorm.server.Access;
+import com.google.gwtorm.server.OrmDuplicateKeyException;
+import com.google.gwtorm.server.OrmException;
+import com.google.gwtorm.server.SchemaFactory;
+import com.google.inject.Inject;
+import com.google.inject.Injector;
+import com.google.protobuf.ByteString;
+import com.google.protobuf.Parser;
+import com.google.protobuf.UnknownFieldSet;
+
+import org.eclipse.jgit.lib.ProgressMonitor;
+import org.eclipse.jgit.lib.TextProgressMonitor;
+import org.kohsuke.args4j.Option;
+
+import java.io.BufferedInputStream;
+import java.io.File;
+import java.io.FileInputStream;
+import java.io.InputStream;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Method;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Import data from a protocol buffer dump into the database.
+ * <p>
+ * Takes as input a file containing protocol buffers concatenated together with
+ * varint length encoding, as in {@link Parser#parseDelimitedFrom(InputStream)}.
+ * Each message contains a single field with a tag corresponding to the relation
+ * ID in the {@link com.google.gwtorm.server.Relation} annotation.
+ * <p>
+ * <strong>Warning</strong>: This method blindly upserts data into the database.
+ * It should only be used to restore a protobuf-formatted backup into a new,
+ * empty site.
+ */
+public class ProtobufImport extends SiteProgram {
+  @Option(name = "--file", aliases = {"-f"}, required = true, metaVar = "FILE",
+      usage = "File to import from")
+  private File file;
+
+  private final LifecycleManager manager = new LifecycleManager();
+  private final Map<Integer, Relation> relations = new HashMap<>();
+
+  @Inject
+  private SchemaFactory<ReviewDb> schemaFactory;
+
+  @Override
+  public int run() throws Exception {
+    mustHaveValidSite();
+
+    Injector dbInjector = createDbInjector(SINGLE_USER);
+    manager.add(dbInjector);
+    manager.start();
+    RuntimeShutdown.add(new Runnable() {
+      @Override
+      public void run() {
+        manager.stop();
+      }
+    });
+    dbInjector.injectMembers(this);
+
+    ReviewDb db = schemaFactory.open();
+
+    ProgressMonitor progress = new TextProgressMonitor();
+    progress.beginTask("Importing entities", ProgressMonitor.UNKNOWN);
+    try {
+      for (RelationModel model
+          : new JavaSchemaModel(ReviewDb.class).getRelations()) {
+        relations.put(model.getRelationID(), Relation.create(model, db));
+      }
+
+      Parser<UnknownFieldSet> parser =
+          UnknownFieldSet.getDefaultInstance().getParserForType();
+      try (InputStream in = new BufferedInputStream(new FileInputStream(file))) {
+        UnknownFieldSet msg;
+        while ((msg = parser.parseDelimitedFrom(in)) != null) {
+          Map.Entry<Integer, UnknownFieldSet.Field> e =
+              Iterables.getOnlyElement(msg.asMap().entrySet());
+          Relation rel = checkNotNull(relations.get(e.getKey()),
+              "unknown relation ID %s in message: %s", e.getKey(), msg);
+          List<ByteString> values = e.getValue().getLengthDelimitedList();
+          checkState(values.size() == 1,
+            "expected one string field in message: %s", msg);
+          upsert(rel, values.get(0));
+          progress.update(1);
+        }
+      }
+      progress.endTask();
+    } finally {
+      db.close();
+    }
+
+    return 0;
+  }
+
+  @SuppressWarnings({"rawtypes", "unchecked"})
+  private static void upsert(Relation rel, ByteString s)
+      throws OrmException {
+    Collection ents = Collections.singleton(rel.codec().decode(s));
+    try {
+      // Not all relations support update; fall back manually.
+      rel.access().insert(ents);
+    } catch (OrmDuplicateKeyException e) {
+      rel.access().delete(ents);
+      rel.access().insert(ents);
+    }
+  }
+
+  @AutoValue
+  static abstract class Relation {
+    private static Relation create(RelationModel model, ReviewDb db)
+        throws IllegalAccessException, InvocationTargetException,
+        NoSuchMethodException, ClassNotFoundException {
+      Method m = db.getClass().getMethod(model.getMethodName());
+      Class<?> clazz = Class.forName(model.getEntityTypeClassName());
+      return new AutoValue_ProtobufImport_Relation(
+          (Access<?, ?>) m.invoke(db),
+          CodecFactory.encoder(clazz));
+    }
+
+    abstract Access<?, ?> access();
+    abstract ProtobufCodec<?> codec();
+  }
+}