ZNodeBulkLoader.java

/*
 * Copyright (C) 2012 Facebook, Inc.
 *
 * 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.facebook.zookeeper.cmd;

import com.facebook.util.StreamImporter;
import com.facebook.zookeeper.path.ZkGenericPath;
import com.facebook.zookeeper.convenience.ZkQuickConnectionManager;
import com.facebook.zookeeper.convenience.ZkScript;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.Options;

import java.io.File;
import java.io.FileInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Iterator;
import java.util.List;

public class ZNodeBulkLoader extends ZkScript {
  private volatile boolean verbose = false;

  public ZNodeBulkLoader(ZkQuickConnectionManager zkQuickConnectionManager) {
    super(zkQuickConnectionManager);
  }

  public ZNodeBulkLoader() {
    this(new ZkQuickConnectionManager());
  }

  public void setVerbose(boolean verbose) {
    this.verbose = verbose;
  }

  public void load(File template)
    throws InterruptedException, IOException, KeeperException {
    InputStream in = new FileInputStream(template);
    try {
      load(in);
    } finally {
      in.close();
    }
  }

  public void load(InputStream in)
    throws InterruptedException, IOException, KeeperException {
    load(StreamImporter.importLines(in));
  }

  public void load(List<String> pathStrs)
    throws InterruptedException, KeeperException {
    for (String pathStr : pathStrs) {
      createEntirePath(pathStr);
    }
  }

  private void createEntirePath(String pathStr)
    throws InterruptedException, KeeperException {
    ZkGenericPath path = ZkGenericPath.parse("/", pathStr);
    Iterator<ZkGenericPath> lineageIter = path.lineageIterator();
    while (lineageIter.hasNext()) {
      String currentPathStr = lineageIter.next().toString();
      try {
        getZk().create(
          currentPathStr,
          null,
          ZooDefs.Ids.OPEN_ACL_UNSAFE,
          CreateMode.PERSISTENT
        );
      } catch (KeeperException.NodeExistsException e) {
        // Ignore if it already exist
      }
    }
    if (verbose) {
      System.out.println("Created ZNode: " + pathStr);
    }
  }

  @Override
  protected String getName() {
    return ZNodeBulkLoader.class.getName();
  }

  @Override
  protected Options getSpecificOptions() {
    Options options = new Options();
    options.addOption(
      "f",
      "file-template",
      true,
      "Path to file containing new-line delimited list of ZNodes to create " +
        "(recursive). If this parameter is not specified, expects the values " +
        "to be provided via standard input."
    );
    options.addOption(
      "v",
      "verbose",
      false,
      "Print verbose messages [Default: off]"
    );
    return options;
  }

  @Override
  protected boolean verifySpecificOptions(CommandLine cmd) {
    if (cmd.hasOption("file-template")) {
      File template = new File(cmd.getOptionValue("file-template"));
      if (!template.exists()) {
        System.err.println("Error: invalid file-template path.\n");
        return false;
      }
    }
    return true;
  }

  @Override
  protected void runScript(CommandLine cmd) throws Exception {
    setVerbose(cmd.hasOption("verbose"));
    if (cmd.hasOption("file-template")) {
      File template = new File(cmd.getOptionValue("file-template"));
      load(template);
    } else {
      load(System.in);
    }
  }

  public static void main(String[] args) throws Exception {
    ZkScript script = new ZNodeBulkLoader();
    script.runMain(args);
    System.out.println("DONE");
  }
}