org.springframework.data.keyvalue.riak.util.RiakClassFileLoader.java Source code

Introduction

Here is the source code for org.springframework.data.keyvalue.riak.util.RiakClassFileLoader.java

Source

/*
 * Copyright (c) 2011 by J. Brisbin <jon@jbrisbin.com>
 * Portions (c) 2011 by NPC International, Inc. or the
 * original author(s).
 *
 * 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 org.springframework.data.keyvalue.riak.util;

import org.apache.commons.cli.*;
import org.springframework.data.keyvalue.riak.core.RiakTemplate;

import java.io.*;
import java.net.URLEncoder;
import java.util.zip.ZipEntry;
import java.util.zip.ZipInputStream;

/**
 * @author J. Brisbin <jon@jbrisbin.com>
 */
public class RiakClassFileLoader {

    static Options opts = new Options();

    static {
        opts.addOption("v", false, "Verbose output");
        opts.addOption("u", true, "URL to Riak (defaults to: 'http://localhost:8098/riak/{bucket}/{key}')");
        opts.addOption("b", true, "Bucket to load class files into");
        opts.addOption("k", true, "Key under which to store an individual class file");
        opts.addOption("j", true, "JAR file to load into Riak");
        opts.addOption("c", true, "Class file to load into Riak");
        opts.addOption("d", true, "Directory from which to load all JAR files into Riak");
    }

    public static void main(String[] args) {
        Parser p = new BasicParser();
        CommandLine cl = null;
        try {
            cl = p.parse(opts, args);
        } catch (ParseException e) {
            System.err.println("Error parsing command line: " + e.getMessage());
        }

        if (null != cl) {
            boolean verbose = cl.hasOption('v');
            RiakTemplate riak = new RiakTemplate();
            riak.getRestTemplate().setErrorHandler(new Ignore404sErrorHandler());
            if (cl.hasOption('u')) {
                riak.setDefaultUri(cl.getOptionValue('u'));
            }
            try {
                riak.afterPropertiesSet();
            } catch (Exception e) {
                System.err.println("Error creating RiakTemplate: " + e.getMessage());
            }
            String[] files = cl.getOptionValues('j');
            if (null != files) {
                for (String file : files) {
                    if (verbose) {
                        System.out.println(String.format("Loading JAR file %s into Riak...", file));
                    }
                    try {
                        File zfile = new File(file);
                        ZipInputStream zin = new ZipInputStream(new FileInputStream(zfile));
                        ZipEntry entry;
                        while (null != (entry = zin.getNextEntry())) {
                            ByteArrayOutputStream bout = new ByteArrayOutputStream();
                            byte[] buff = new byte[16384];
                            for (int bytesRead = zin.read(buff); bytesRead > 0; bytesRead = zin.read(buff)) {
                                bout.write(buff, 0, bytesRead);
                            }

                            if (entry.getName().endsWith(".class")) {
                                String name = entry.getName().replaceAll("/", ".");
                                name = URLEncoder.encode(name.substring(0, name.length() - 6), "UTF-8");
                                String bucket;
                                if (cl.hasOption('b')) {
                                    bucket = cl.getOptionValue('b');
                                } else {
                                    bucket = URLEncoder.encode(zfile.getCanonicalFile().getName(), "UTF-8");
                                }
                                if (verbose) {
                                    System.out.println(String.format("Uploading to %s/%s", bucket, name));
                                }

                                // Load these bytes into Riak
                                riak.setAsBytes(bucket, name, bout.toByteArray());
                            }
                        }
                    } catch (FileNotFoundException e) {
                        System.err.println("Error reading JAR file: " + e.getMessage());
                    } catch (IOException e) {
                        System.err.println("Error reading JAR file: " + e.getMessage());
                    }
                }
            }

            String[] classFiles = cl.getOptionValues('c');
            if (null != classFiles) {
                for (String classFile : classFiles) {
                    try {
                        FileInputStream fin = new FileInputStream(classFile);
                        ByteArrayOutputStream bout = new ByteArrayOutputStream();
                        byte[] buff = new byte[16384];
                        for (int bytesRead = fin.read(buff); bytesRead > 0; bytesRead = fin.read(buff)) {
                            bout.write(buff, 0, bytesRead);
                        }

                        String name;
                        if (cl.hasOption('k')) {
                            name = cl.getOptionValue('k');
                        } else {
                            throw new IllegalStateException(
                                    "Must specify a Riak key in which to store the data if loading individual class files.");
                        }
                        String bucket;
                        if (cl.hasOption('b')) {
                            bucket = cl.getOptionValue('b');
                        } else {
                            throw new IllegalStateException(
                                    "Must specify a Riak bucket in which to store the data if loading individual class files.");
                        }
                        if (verbose) {
                            System.out.println(String.format("Uploading to %s/%s", bucket, name));
                        }

                        // Load these bytes into Riak
                        riak.setAsBytes(bucket, name, bout.toByteArray());

                    } catch (FileNotFoundException e) {
                        System.err.println("Error reading class file: " + e.getMessage());
                    } catch (IOException e) {
                        System.err.println("Error reading class file: " + e.getMessage());
                    }
                }
            }
        }
    }

}