org.springframework.data.hadoop.impala.r.RCommands.java Source code

Introduction

Here is the source code for org.springframework.data.hadoop.impala.r.RCommands.java

Source

/*
 * Copyright 2011-2012 the original author or authors.
 * 
 * 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.hadoop.impala.r;

import java.io.BufferedReader;
import java.io.File;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

import org.apache.commons.io.IOUtils;
import org.springframework.core.io.FileSystemResourceLoader;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.PathMatchingResourcePatternResolver;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.shell.core.CommandMarker;
import org.springframework.shell.core.annotation.CliCommand;
import org.springframework.shell.core.annotation.CliOption;
import org.springframework.shell.support.util.StringUtils;
import org.springframework.stereotype.Component;

/**
 * Provider of R commands.
 * 
 * @author Costin Leau
 */
@Component
public class RCommands implements CommandMarker {

    private static final String PREFIX = "r ";
    private static final String BIN = "bin" + File.separator + "Rscript";
    private String home = "";
    private String cmd = home + BIN;
    private File wkdir = new File(".");

    private ResourcePatternResolver resourceResolver = new PathMatchingResourcePatternResolver(
            new FileSystemResourceLoader());

    public RCommands() {
        String os = System.getProperty("os.name").toLowerCase();

        // set default R paths - best effort
        if (os.contains("win")) {
            home = "C:\\Program Files\\R\\";
            File f = new File(home);
            if (f.exists() && f.isDirectory()) {
                File[] listFiles = f.listFiles();
                if (listFiles != null && listFiles.length > 0) {
                    home = listFiles[0].getAbsolutePath() + "\\";
                }
            }
        } else if (os.contains("mac")) {
            home = "/Library/Frameworks/R.framework/Resources/";
        } else if (os.contains("sunos")) {
            home = "/usr/lib/";
        } else {
            // else linux/unix
            home = "/usr/";
        }

        cmd = home + BIN;
    }

    private String info() {
        StringBuilder sb = new StringBuilder();
        sb.append("R [");
        //String version = "unknown";
        //sb.append(version);
        //sb.append("][");
        sb.append("home=");
        sb.append(home);
        sb.append("][workDir=");
        sb.append(wkdir.getAbsolutePath());
        sb.append("]");

        return sb.toString();
    }

    @CliCommand(value = { PREFIX + "cfg" }, help = "Configures R")
    public String config(@CliOption(key = { "home" }, mandatory = false, help = "R home directory") String home,
            @CliOption(key = { "workDir" }, mandatory = false, help = "working directory") String wkdir)
            throws Exception {

        if (StringUtils.hasText(home)) {
            if (new File(home).exists()) {
                if (!home.endsWith(File.separator)) {
                    home += File.separator;
                }
                this.home = home;
                cmd = home + BIN;
            } else {
                return "Cannot find path [" + home + "]";
            }
        }

        if (StringUtils.hasText(wkdir)) {
            File f = new File(wkdir);
            if (f.exists()) {
                this.wkdir = f;
            } else {
                return "Cannot find path [" + wkdir + "]";
            }
        }

        return info();
    }

    @CliCommand(value = { PREFIX + "script" }, help = "Executes a R script")
    public String script(
            @CliOption(key = { "", "location" }, mandatory = true, help = "Script location") String location,
            @CliOption(key = { "args" }, mandatory = false, help = "Script arguments") String args) {

        if (!new File(cmd).exists()) {
            return "Cannot find R command [" + cmd + "]";
        }

        Resource resource = resourceResolver.getResource(fixLocation(location));
        if (!resource.exists()) {
            return "No resource found at " + location;
        }

        String uri = location;

        try {
            uri = resource.getFile().getAbsolutePath();
        } catch (IOException ex) {
            // ignore - we'll use location
        }

        Process process;
        List<String> cmds = new ArrayList<String>();

        cmds.add(cmd);
        cmds.add(uri);

        if (StringUtils.hasText(args)) {
            cmds.addAll(Arrays.asList(StringUtils.tokenizeToStringArray(args, " ")));
        }

        BufferedReader output = null;
        try {
            process = new ProcessBuilder(cmds).directory(wkdir).redirectErrorStream(true).start();
            output = new BufferedReader(new InputStreamReader(process.getInputStream()));
            String line;
            while ((line = output.readLine()) != null) {
                System.out.println(line);
            }
            int code = process.waitFor();
            return "R script executed with exit code - " + code;
        } catch (Exception ex) {
            return "R script execution failed - " + ex;
        } finally {
            IOUtils.closeQuietly(output);
        }
    }

    private static String fixLocation(String location) {
        if (StringUtils.hasText(location) && !location.contains(":")) {
            return "file:" + location;
        }
        return location;
    }
}