com.linkedin.databus.bootstrap.utils.BootstrapAvroRecordDumper.java Source code

Introduction

Here is the source code for com.linkedin.databus.bootstrap.utils.BootstrapAvroRecordDumper.java

Source

package com.linkedin.databus.bootstrap.utils;
/*
 *
 * Copyright 2013 LinkedIn Corp. All rights reserved
 *
 * 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.
 *
*/

import java.io.File;
import java.io.IOException;

import org.apache.avro.Schema;
import org.apache.avro.file.DataFileWriter;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.commons.cli.CommandLine;
import org.apache.commons.cli.CommandLineParser;
import org.apache.commons.cli.GnuParser;
import org.apache.commons.cli.Option;
import org.apache.commons.cli.OptionBuilder;
import org.apache.commons.cli.Options;
import org.apache.commons.cli.ParseException;
import org.apache.log4j.Logger;

import com.linkedin.databus.core.DbusEvent;

public class BootstrapAvroRecordDumper {
    public static final String MODULE = BootstrapAvroRecordDumper.class.getName();
    public static final Logger LOG = Logger.getLogger(MODULE);

    public static final String OUTPUT_DIR_OPT_LONG_NAME = "output_dir";
    public static final char OUTPUT_DIR_OPT_CHAR = 'o';

    public static String outputDir = null;

    public static void main(String[] args) throws Exception {
        BootstrapTableReader.init(args);
        BootstrapTableReader reader = new BootstrapTableReader(
                new DumpEventHandler(outputDir, BootstrapTableReader.getSchema()));
        reader.execute();
    }

    @SuppressWarnings("static-access")
    public static void parseArgs(String[] args) throws IOException {
        CommandLineParser cliParser = new GnuParser();

        Option outputDirOption = OptionBuilder.withLongOpt(OUTPUT_DIR_OPT_LONG_NAME).withDescription("Help screen")
                .create(OUTPUT_DIR_OPT_CHAR);

        Options options = new Options();
        options.addOption(outputDirOption);

        CommandLine cmd = null;
        try {
            cmd = cliParser.parse(options, args);
        } catch (ParseException pe) {
            LOG.fatal("Bootstrap Avro Record Dumper: failed to parse command-line options.", pe);
            throw new RuntimeException("Bootstrap Avro Record Dumper: failed to parse command-line options.", pe);
        }

        if (cmd.hasOption(OUTPUT_DIR_OPT_CHAR)) {
            outputDir = cmd.getOptionValue(OUTPUT_DIR_OPT_CHAR);
        }
    }

    public static class DumpEventHandler implements BootstrapReaderEventHandler {
        private String _query = null;

        private DataFileWriter<GenericRecord> writer = null;
        private File directory = null;
        private File currentFile = null;

        public DumpEventHandler(String dir, Schema schema) throws IOException {
            DataFileWriter<GenericRecord> writeCreator = new DataFileWriter<GenericRecord>(
                    new GenericDatumWriter(schema));
            directory = new File(dir);

            if (!directory.isDirectory())
                throw new RuntimeException("The path (" + dir + ") either does not exist or is not a directory !!");

            currentFile = new File(directory.getAbsolutePath() + "/part-000.avro");
            writer = writeCreator.create(schema, currentFile);
        }

        @Override
        public void onRecord(DbusEvent event, GenericRecord record) {
            try {
                writer.append(record);
            } catch (IOException e) {
                e.printStackTrace();
                throw new RuntimeException(e);
            }
        }

        @Override
        public void onStart(String query) {
            _query = query;
        }

        @Override
        public void onEnd(int count) {
            System.out.println("Read " + count + " records by executing query :( " + _query + ")");
        }

    }
}