com.cinchapi.concourse.importer.cli.ImportCli.java Source code

Java tutorial

Introduction

Here is the source code for com.cinchapi.concourse.importer.cli.ImportCli.java

Source

/*
 * Copyright (c) 2013-2016 Cinchapi 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.cinchapi.concourse.importer.cli;

import java.io.IOException;
import java.lang.reflect.Constructor;
import java.lang.reflect.Modifier;
import java.nio.file.DirectoryStream;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.text.MessageFormat;
import java.util.Collection;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import org.reflections.Reflections;

import jline.TerminalFactory;
import jline.console.ConsoleReader;

import com.beust.jcommander.Parameter;
import com.cinchapi.concourse.Concourse;
import com.cinchapi.concourse.cli.CommandLineInterface;
import com.cinchapi.concourse.cli.Options;
import com.cinchapi.concourse.importer.CsvImporter;
import com.cinchapi.concourse.importer.Headered;
import com.cinchapi.concourse.importer.Importer;
import com.cinchapi.concourse.importer.LegacyCsvImporter;
import com.cinchapi.concourse.util.FileOps;
import com.cinchapi.concourse.util.Reflection;
import com.cinchapi.concourse.util.Strings;
import com.google.common.base.CaseFormat;
import com.google.common.base.Stopwatch;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Queues;
import com.google.common.collect.Sets;

/**
 * A CLI that uses the import framework to import data into Concourse.
 * 
 * @author Jeff Nelson
 */
@SuppressWarnings("deprecation")
public class ImportCli extends CommandLineInterface {

    /**
     * Aliases for built-in importer types. These aliases should be specified
     * using the {@code -t} or {@code --type} flag when invoking the CLI.
     */
    private static Map<String, Class<? extends Importer>> importers = Maps.newHashMapWithExpectedSize(2);
    static {
        // NOTE: Be sure to increase the expectedSize parameter for the map when
        // adding additional aliases
        importers.put("csv", CsvImporter.class);
        importers.put(".csv", LegacyCsvImporter.class); // deprecated
    }

    /*
     * TODO
     * 2) add flags to whitelist or blacklist files in a directory
     */

    /**
     * Construct a new instance.
     * 
     * @param args
     */
    public ImportCli(String[] args) {
        super(args);
    }

    @Override
    protected void doTask() {
        final ImportOptions opts = (ImportOptions) options;
        final Set<Long> records;
        final Constructor<? extends Importer> constructor = getConstructor(opts.type);
        if (opts.data == null) { // Import data from stdin
            Importer importer = Reflection.newInstance(constructor, concourse);
            if (!opts.dynamic.isEmpty()) {
                importer.setParams(options.dynamic);
            }
            if (importer instanceof Headered && !opts.header.isEmpty()) {
                ((Headered) importer).parseHeader(opts.header);
            }
            try {
                ConsoleReader reader = new ConsoleReader();
                String line;
                records = Sets.newLinkedHashSet();
                Runtime.getRuntime().addShutdownHook(new Thread(new Runnable() {

                    // Interactive import is ended when user presses CTRL + C,
                    // so we need this shutdown hook to ensure that they get
                    // feedback about the import before the JVM dies.

                    @Override
                    public void run() {
                        if (options.verbose) {
                            System.out.println(records);
                        }
                        System.out.println(Strings.format("Imported data into {} records", records.size()));
                    }

                }));
                try {
                    final AtomicBoolean lock = new AtomicBoolean(false);
                    new Thread(new Runnable() { // If there is no input in
                                                // 100ms, assume that the
                                                // session is interactive (i.e.
                                                // not piped) and display a
                                                // prompt

                        @Override
                        public void run() {
                            try {
                                Thread.sleep(100);
                                if (lock.compareAndSet(false, true)) {
                                    System.out.println("Importing from stdin. Press " + "CTRL + C when finished");
                                }
                            } catch (InterruptedException e) {
                            }
                        }

                    }).start();
                    while ((line = reader.readLine()) != null) {
                        try {
                            lock.set(true);
                            records.addAll(importer.importString(line));
                        } catch (Exception e) {
                            System.err.println(e);
                        }
                    }
                } catch (IOException e) {
                    throw Throwables.propagate(e);
                }
            } catch (IOException e) {
                throw Throwables.propagate(e);
            } finally {
                try {
                    TerminalFactory.get().restore();
                } catch (Exception e) {
                    throw Throwables.propagate(e);
                }
            }
        } else {
            String path = FileOps.expandPath(opts.data, getLaunchDirectory());
            Collection<String> files = FileOps.isDirectory(path) ? scan(Paths.get(path)) : ImmutableList.of(path);
            Stopwatch watch = Stopwatch.createUnstarted();
            if (files.size() > 1) {
                records = Sets.newConcurrentHashSet();
                final Queue<String> filesQueue = (Queue<String>) files;
                List<Runnable> runnables = Lists.newArrayListWithCapacity(opts.numThreads);
                // Create just enough Runnables with instantiated Importers in
                // advance. Each of those Runnables will work until #filesQueue
                // is exhausted.
                opts.numThreads = Math.min(opts.numThreads, files.size());
                for (int i = 0; i < opts.numThreads; ++i) {
                    final Importer importer0 = Reflection.newInstance(constructor,
                            i == 0 ? concourse
                                    : Concourse.connect(opts.host, opts.port, opts.username, opts.password,
                                            opts.environment));
                    if (!opts.dynamic.isEmpty()) {
                        importer0.setParams(opts.dynamic);
                    }
                    if (importer0 instanceof Headered && !opts.header.isEmpty()) {
                        ((Headered) importer0).parseHeader(opts.header);
                    }
                    runnables.add(new Runnable() {

                        private final Importer importer = importer0;

                        @Override
                        public void run() {
                            String file;
                            while ((file = filesQueue.poll()) != null) {
                                records.addAll(importer.importFile(file));
                            }
                        }

                    });
                }
                ExecutorService executor = Executors.newFixedThreadPool(runnables.size());
                System.out.println("Starting import...");
                watch.start();
                for (Runnable runnable : runnables) {
                    executor.execute(runnable);
                }
                executor.shutdown();
                try {
                    if (!executor.awaitTermination(1, TimeUnit.MINUTES)) {
                        while (!executor.isTerminated()) {
                            System.out.print('.'); // block until all tasks are
                                                   // completed and provide some
                                                   // feedback to the user
                        }
                    }
                } catch (InterruptedException e) {
                    throw Throwables.propagate(e);
                }
            } else {
                Importer importer = Reflection.newInstance(constructor, concourse);
                if (!opts.dynamic.isEmpty()) {
                    importer.setParams(opts.dynamic);
                }
                if (importer instanceof Headered && !opts.header.isEmpty()) {
                    ((Headered) importer).parseHeader(opts.header);
                }
                System.out.println("Starting import...");
                watch.start();
                records = importer.importFile(files.iterator().next());
            }
            watch.stop();
            long elapsed = watch.elapsed(TimeUnit.MILLISECONDS);
            double seconds = elapsed / 1000.0;
            if (options.verbose) {
                System.out.println(records);
            }
            System.out.println(MessageFormat.format("Imported data " + "into {0} records in {1} seconds",
                    records.size(), seconds));
        }
    }

    @Override
    protected Options getOptions() {
        return new ImportOptions();
    }

    /**
     * Return the appropriate constructor for creating instances of the
     * {@code type} {@link Importer}.
     * 
     * <p>
     * To make the {@link Importer}, call
     * {@link Constructor#newInstance(Object...)} with a non-shared Concourse
     * connection.
     * </p>
     * 
     * @param type the {@link #importers alias} or fully qualified name for the
     *            desired {@link Importer} class
     * @return the constructor
     */
    private static Constructor<? extends Importer> getConstructor(String type) {
        Class<? extends Importer> clz = importers.get(type);
        if (clz == null) {
            try {
                clz = getCustomImporterClass(type);
            } catch (ClassNotFoundException e) {
                throw new RuntimeException(Strings.format("{} is not a valid importer type.", type));
            }
        }
        try {
            return clz.getDeclaredConstructor(Concourse.class);
        } catch (NoSuchMethodException e) {
            // This should never happen because Importer base class mandates the
            // existence of the constructor in the subclass by not defining a
            // no-arg alternative
            throw Throwables.propagate(e);
        }
    }

    /**
     * Given an alias (or fully qualified class name) attempt to load a "custom"
     * importer that is not already defined in the {@link #importers built-in}
     * collection.
     * 
     * @param alias a conventional alias (FileTypeImporter --> file-type) or a
     *            fully qualified class name
     * @return the {@link Class} that corresponds to the custom importer
     * @throws ClassNotFoundException
     */
    @SuppressWarnings("unchecked")
    private static Class<? extends Importer> getCustomImporterClass(String alias) throws ClassNotFoundException {
        try {
            return (Class<? extends Importer>) Class.forName(alias);
        } catch (ClassNotFoundException e) {
            // Attempt to determine the correct class name from the alias by
            // loading the server's classpath. For the record, this is hella
            // slow.
            Reflections.log = null; // turn off reflection logging
            Reflections reflections = new Reflections();
            char firstChar = alias.charAt(0);
            for (Class<? extends Importer> clazz : reflections.getSubTypesOf(Importer.class)) {
                String name = clazz.getSimpleName();
                if (name.length() == 0) { // Skip anonymous subclasses
                    continue;
                }
                char nameFirstChar = name.charAt(0);
                if (!Modifier.isAbstract(clazz.getModifiers()) && (nameFirstChar == Character.toUpperCase(firstChar)
                        || nameFirstChar == Character.toLowerCase(firstChar))) {
                    String expected = CaseFormat.UPPER_CAMEL.to(CaseFormat.LOWER_HYPHEN, clazz.getSimpleName())
                            .replaceAll("-importer", "");
                    if (alias.equals(expected)) {
                        return clazz;
                    }
                }
            }
            throw e;
        }
    }

    /**
     * Recursively scan and collect all the files in the directory defined by
     * {@code path}.
     * 
     * @param path a {@link Path} that is already verified to
     *            {@link java.nio.file.Files#isDirectory(Path, java.nio.file.LinkOption...)
     *            to be a directory}.
     * @return the collection of files in the directory
     */
    private static Queue<String> scan(Path path) {
        try (DirectoryStream<Path> stream = java.nio.file.Files.newDirectoryStream(path)) {
            Iterator<Path> it = stream.iterator();
            Queue<String> files = Queues.newConcurrentLinkedQueue();
            while (it.hasNext()) {
                Path thePath = it.next();
                if (java.nio.file.Files.isDirectory(thePath)) {
                    files.addAll(scan(thePath));
                } else {
                    files.add(thePath.toString());
                }
            }
            return files;
        } catch (IOException e) {
            throw Throwables.propagate(e);
        }
    }

    /**
     * Import specific {@link Options}.
     * 
     * @author jnelson
     */
    protected static class ImportOptions extends Options {

        @Parameter(names = { "-d",
                "--data" }, description = "The path to the file or directory to import; if no source is provided read from stdin")
        public String data;

        @Parameter(names = "--numThreads", description = "The number of worker threads to use for a multithreaded import")
        public int numThreads = Runtime.getRuntime().availableProcessors();

        @Parameter(names = { "-r",
                "--resolveKey" }, description = "The key to use when resolving data into existing records")
        public String resolveKey = null;

        @Parameter(names = { "-t", "--type" }, description = "The name/type of the importer to use")
        public String type = "csv";

        @Parameter(names = "--header", description = "A custom header to assign for supporting importers")
        public String header = "";

    }

}