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

// 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 java.io.BufferedInputStream;
import java.io.File;
import java.io.InputStream;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.nio.file.Files;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import org.eclipse.jgit.lib.ProgressMonitor;
import org.eclipse.jgit.lib.TextProgressMonitor;
import org.kohsuke.args4j.Option;

/**
 * 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);

    ProgressMonitor progress = new TextProgressMonitor();
    progress.beginTask("Importing entities", ProgressMonitor.UNKNOWN);
    try (ReviewDb db = schemaFactory.open()) {
      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(Files.newInputStream(file.toPath()))) {
        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();
    }

    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
  abstract static 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();
  }
}