jjeong

package org.apache.solr.util;

/**

 * Licensed to the Apache Software Foundation (ASF) under one or more

 * contributor license agreements.  See the NOTICE file distributed with

 * this work for additional information regarding copyright ownership.

 * The ASF licenses this file to You 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.FileInputStream;

import java.io.IOException;

import java.io.InputStream;

import java.io.ByteArrayInputStream;

import java.io.OutputStream;

import java.io.UnsupportedEncodingException;

import java.util.Set;

import java.util.HashSet;

import java.net.HttpURLConnection;

import java.net.MalformedURLException;

import java.net.ProtocolException;

import java.net.URL;

/**

 * A simple utility class for posting raw updates to a Solr server, 

 * has a main method so it can be run on the command line.

 * 

 */

public class SimplePostTool {

  public static final String DEFAULT_POST_URL = "http://localhost:8983/solr/update";

  public static final String VERSION_OF_THIS_TOOL = "1.4";

  private static final String DEFAULT_COMMIT = "yes";

  private static final String DEFAULT_OPTIMIZE = "no";

  private static final String DEFAULT_OUT = "no";

  public static final String DEFAULT_DATA_TYPE = "application/xml";

  private static final String DATA_MODE_FILES = "files";

  private static final String DATA_MODE_ARGS = "args";

  private static final String DATA_MODE_STDIN = "stdin";

  private static final String DEFAULT_DATA_MODE = DATA_MODE_FILES;

  private static final Set<String> DATA_MODES = new HashSet<String>();

  static {

    DATA_MODES.add(DATA_MODE_FILES);

    DATA_MODES.add(DATA_MODE_ARGS);

    DATA_MODES.add(DATA_MODE_STDIN);

  }

  protected URL solrUrl;

  public static void main(String[] args) {

    info("version " + VERSION_OF_THIS_TOOL);

    if (0 < args.length && ("-help".equals(args[0]) || "--help".equals(args[0]) || "-h".equals(args[0]))) {

      System.out.println

        ("This is a simple command line tool for POSTing raw data to a Solr\n"+

         "port.  Data can be read from files specified as commandline args,\n"+

         "as raw commandline arg strings, or via STDIN.\n"+

         "Examples:\n"+

         "  java -jar post.jar *.xml\n"+

         "  java -Ddata=args  -jar post.jar '<delete><id>42</id></delete>'\n"+

         "  java -Ddata=stdin -jar post.jar < hd.xml\n"+

         "  java -Durl=http://localhost:8983/solr/update/csv -Dtype=text/csv -jar post.jar *.csv\n"+

         "  java -Durl=http://localhost:8983/solr/update/json -Dtype=application/json -jar post.jar *.json\n"+

         "  java -Durl=http://localhost:8983/solr/update/extract?literal.id=a -Dtype=application/pdf -jar post.jar a.pdf\n"+

         "Other options controlled by System Properties include the Solr\n"+

         "URL to POST to, the Content-Type of the data, whether a commit\n"+

         "or optimize should be executed, and whether the response should\n"+

         "be written to STDOUT. These are the defaults for all System Properties:\n"+

         "  -Ddata=" + DEFAULT_DATA_MODE + "\n"+

         "  -Dtype=" + DEFAULT_DATA_TYPE + "\n"+

         "  -Durl=" + DEFAULT_POST_URL + "\n"+

         "  -Dcommit=" + DEFAULT_COMMIT + "\n"+

         "  -Doptimize=" + DEFAULT_OPTIMIZE + "\n"+

         "  -Dout=" + DEFAULT_OUT + "\n");

      return;

    }

    OutputStream out = null;

    final String type = System.getProperty("type", DEFAULT_DATA_TYPE);

    URL u = null;

    try {

      u = new URL(System.getProperty("url", DEFAULT_POST_URL));

    } catch (MalformedURLException e) {

      fatal("System Property 'url' is not a valid URL: " + u);

    }

    final SimplePostTool t = new SimplePostTool(u);

    final String mode = System.getProperty("data", DEFAULT_DATA_MODE);

    if (! DATA_MODES.contains(mode)) {

      fatal("System Property 'data' is not valid for this tool: " + mode);

    }

    if ("yes".equals(System.getProperty("out", DEFAULT_OUT))) {

      out = System.out;

    }

    try {

      if (DATA_MODE_FILES.equals(mode)) {

        if (0 < args.length) {

          info("POSTing files to " + u + "..");

          t.postFiles(args, 0, out, type);

        } else {

          info("No files specified. (Use -h for help)");

        }

      } else if (DATA_MODE_ARGS.equals(mode)) {

        if (0 < args.length) {

          info("POSTing args to " + u + "..");

          for (String a : args) {

            t.postData(SimplePostTool.stringToStream(a), null, out, type);

          }

        }

      } else if (DATA_MODE_STDIN.equals(mode)) {

        info("POSTing stdin to " + u + "..");

        t.postData(System.in, null, out, type);

      }

      if ("yes".equals(System.getProperty("commit",DEFAULT_COMMIT))) {

        info("COMMITting Solr index changes..");

        t.commit();

      }

      if ("yes".equals(System.getProperty("optimize",DEFAULT_OPTIMIZE))) {

        info("Performing an OPTIMIZE..");

        t.optimize();

      }

    } catch(RuntimeException e) {

      e.printStackTrace();

      fatal("RuntimeException " + e);

    }

  }

  /**

   * @deprecated use {@link #postData(InputStream, Integer, OutputStream, String)} instead

   */

  @Deprecated

  int postFiles(String [] args,int startIndexInArgs, OutputStream out) {

    final String type = System.getProperty("type", DEFAULT_DATA_TYPE);

    return postFiles(args, startIndexInArgs, out, type);

  }

  /** Post all filenames provided in args, return the number of files posted*/

  int postFiles(String [] args,int startIndexInArgs, OutputStream out, String type) {

    int filesPosted = 0;

    for (int j = startIndexInArgs; j < args.length; j++) {

      File srcFile = new File(args[j]);

      if (srcFile.canRead()) {

        info("POSTing file " + srcFile.getName());

        postFile(srcFile, out, type);

        filesPosted++;

      } else {

        warn("Cannot read input file: " + srcFile);

      }

    }

    return filesPosted;

  }

  static void warn(String msg) {

    System.err.println("SimplePostTool: WARNING: " + msg);

  }

  static void info(String msg) {

    System.out.println("SimplePostTool: " + msg);

  }

  static void fatal(String msg) {

    System.err.println("SimplePostTool: FATAL: " + msg);

    System.exit(1);

  }

  /**

   * Constructs an instance for posting data to the specified Solr URL 

   * (ie: "http://localhost:8983/solr/update")

   */

  public SimplePostTool(URL solrUrl) {

    this.solrUrl = solrUrl;

  }

  /**

   * Does a simple commit operation 

   */

  public void commit() {

    doGet(appendParam(solrUrl.toString(), "commit=true"));

  }

  /**

   * Does a simple optimize operation 

   */

  public void optimize() {

    doGet(appendParam(solrUrl.toString(), "optimize=true"));

  }

  private String appendParam(String url, String param) {

    return url + (url.indexOf('?')>0 ? "&" : "?") + param;

  }

  /**

   * @deprecated use {@link #postFile(File, OutputStream, String)} instead

   */

  public void postFile(File file, OutputStream output) {

    final String type = System.getProperty("type", DEFAULT_DATA_TYPE);

    postFile(file, output, type);

  }

  /**

   * Opens the file and posts it's contents to the solrUrl,

   * writes to response to output.

   * @throws UnsupportedEncodingException 

   */

  public void postFile(File file, OutputStream output, String type) {

    InputStream is = null;

    try {

      is = new FileInputStream(file);

      postData(is, (int)file.length(), output, type);

    } catch (IOException e) {

      fatal("Can't open/read file: " + file);

    } finally {

      try {

        if(is!=null) is.close();

      } catch (IOException e) {

        fatal("IOException while closing file: "+ e);

      }

    }

  }

  /**

   * Performs a simple get on the given URL

   * @param url

   */

  public static void doGet(String url) {

    try {

      doGet(new URL(url));

    } catch (MalformedURLException e) {

      fatal("The specified URL "+url+" is not a valid URL. Please check");

    }

  }

  /**

   * Performs a simple get on the given URL

   * @param url

   */

  public static void doGet(URL url) {

    try {

      HttpURLConnection urlc = (HttpURLConnection) url.openConnection();

      if (HttpURLConnection.HTTP_OK != urlc.getResponseCode()) {

        fatal("Solr returned an error #" + urlc.getResponseCode() + 

            " " + urlc.getResponseMessage());

      }

    } catch (IOException e) {

      fatal("An error occured posting data to "+url+". Please check that Solr is running.");

    }

  }

  /**

   * @deprecated use {@link #postData(InputStream, Integer, OutputStream, String)} instead

   */

  @Deprecated

  public void postData(InputStream data, Integer length, OutputStream output) {

    final String type = System.getProperty("type", DEFAULT_DATA_TYPE);

    postData(data, length, output, type);

  }

  /**

   * Reads data from the data stream and posts it to solr,

   * writes to the response to output

   */

  public void postData(InputStream data, Integer length, OutputStream output, String type) {

    HttpURLConnection urlc = null;

    try {

      try {

        urlc = (HttpURLConnection) solrUrl.openConnection();

        try {

          urlc.setRequestMethod("POST");

        } catch (ProtocolException e) {

          fatal("Shouldn't happen: HttpURLConnection doesn't support POST??"+e);

        }

        urlc.setDoOutput(true);

        urlc.setDoInput(true);

        urlc.setUseCaches(false);

        urlc.setAllowUserInteraction(false);

        urlc.setRequestProperty("Content-type", type);

        if (null != length) urlc.setFixedLengthStreamingMode(length);

      } catch (IOException e) {

        fatal("Connection error (is Solr running at " + solrUrl + " ?): " + e);

      }

      OutputStream out = null;

      try {

        out = urlc.getOutputStream();

        pipe(data, out);

      } catch (IOException e) {

        fatal("IOException while posting data: " + e);

      } finally {

        try { if(out!=null) out.close(); } catch (IOException x) { /*NOOP*/ }

      }

      InputStream in = null;

      try {

        if (HttpURLConnection.HTTP_OK != urlc.getResponseCode()) {

          fatal("Solr returned an error #" + urlc.getResponseCode() + 

                " " + urlc.getResponseMessage());

        }

        in = urlc.getInputStream();

        pipe(in, output);

      } catch (IOException e) {

        fatal("IOException while reading response: " + e);

      } finally {

        try { if(in!=null) in.close(); } catch (IOException x) { /*NOOP*/ }

      }

    } finally {

      if(urlc!=null) urlc.disconnect();

    }

  }

  public static InputStream stringToStream(String s) {

    InputStream is = null;

    try {

      is = new ByteArrayInputStream(s.getBytes("UTF-8"));

    } catch (UnsupportedEncodingException e) {

      fatal("Shouldn't happen: UTF-8 not supported?!?!?!");

    }

    return is;

  }

  /**

   * Pipes everything from the source to the dest.  If dest is null, 

   * then everything is read from source and thrown away.

   */

  private static void pipe(InputStream source, OutputStream dest) throws IOException {

    byte[] buf = new byte[1024];

    int read = 0;

    while ( (read = source.read(buf) ) >= 0) {

      if (null != dest) dest.write(buf, 0, read);

    }

    if (null != dest) dest.flush();

  }

}