wonderdog 0.1.1 → 0.2.0

data/squirrel/squirrel.rb

@@ -0,0 +1,95 @@
+require "multi_json"
+require_relative "../squirrel/esbackup_stripped.rb"
+require_relative "../squirrel/replay.rb"
+require_relative "../squirrel/warmer_interface.rb"
+require_relative "../squirrel/clear_es_caches.rb"
+require_relative "../squirrel/change_es_index_settings.rb"
+
+class Squirrel
+
+  def initialize(command, options = {})
+    @command = command
+    @options = options
+  end
+
+  def determine_warmer_action(options = {})
+    options[:index] = options[:warmers_index]
+    unless options[:remove_warmer].nil?
+      puts "removing warmer"
+      options[:action] = "remove_warmer"
+      options[:warmer_name] = options[:remove_warmer]
+    else
+      if options[:warmers]
+        puts "enabling warmers"
+        options[:action] = "enable_warmer"
+      elsif options[:warmers] == false
+        puts "disabling warmers"
+        options[:action] = "disable_warmer"
+      end
+      unless options[:new_warmers_name].nil?
+        puts "adding warmer"
+        options[:action] = "add_warmer"
+        options[:warmer_name] = options[:new_warmers_name]
+        options[:query] = options[:create_warmer]
+      end
+    end
+    WarmerInterface.new(options).determine_interaction
+  end
+
+  def determine_cache_clear(options = {})
+    if options[:clear_all_cache]
+      options[:type] = "all"
+      ClearESCaches.new(options).run
+    end
+    if options[:clear_filter_cache]
+      options[:type] = "filter"
+      ClearESCaches.new(options).run
+    end
+    if options[:clear_fielddata]
+      options[:type] = "fielddata"
+      ClearESCaches.new(options).run
+    end
+  end
+
+  def cardinality(options)
+    options[:cardinality].each do |field|
+      output = `ruby getFields.rb --dump=#{options[:card_file]} --field=#{field} >> #{field}.txt ;
+        cat #{field}.txt |sort | uniq -c |sort -n | wc -l;`
+      puts "The number of values in #{field} form file #{ooptions[:card_file]} is #{output}"
+    end
+  end
+
+  def task_caller
+    puts "Running #{@command}"
+    case @command
+      when :restore
+        @options[:index] = @options[:restore_index]
+        @options[:mappings] = @options[:restore_mapping]
+        ESRestore.new(@options[:restore_file], @options).run
+      when :backup
+        @options[:index] = @options[:dump_index]
+        @options[:mappings] = @options[:dump_mapping]
+        ESBackup.new(@options[:output_dir], @options).run
+      when :duplicate
+        @options[:index] = @options[:duplicate_index]
+        @options[:mappings] = @options[:duplicate_mapping]
+        ESDup.new(@options[:duplicate_file], @options).run
+      when :cardinality
+        cardinality(@options)
+      when :warmer
+        determine_warmer_action(@options)
+      when :replay
+        Replay.new(@options[:execute_slow_queries], @options[:host], @options[:port], @options[:preference], @options[:routing]).run
+      when :cache
+        determine_cache_clear(@options)
+      when :index_settings
+        unless @options[:es_index_settings].nil? || @options[:es_index_settings_values].nil?
+          @options[:settings_and_values] = @options[:es_index_settings].zip(@options[:es_index_settings_values])
+          ChangeESIndexSettings.new(@options).run
+        else
+          puts "both --es_index_settings and --es_index_settings_values are required to change index settings"
+        end
+      else abort Settings.help("Must specify either backup, restore, duplicate, cardinality, warmer, replay, cache or index_settings.  Got <#{@command}> UPDATE THIS LINE!")
+    end
+  end
+end

data/squirrel/warmer_interface.rb

@@ -0,0 +1,59 @@
+require 'multi_json'
+require 'httparty'
+
+class WarmerInterface
+  def initialize(options = {})
+    puts options.inspect
+    @host = options[:host]
+    @port = options[:port]
+    @query = MultiJson.dump(options[:query])
+    @warmer_name = options[:warmer_name]
+    @index = options[:index]
+    @action = options[:action]
+    @warmer_state = nil
+  end
+
+  def add_warmer
+    url = "http://#{@host}:#{@port}/#{@index}/_warmer/#{@warmer_name}"
+    puts url
+    puts @query
+    HTTParty.put(url, {:body => @query})
+  end
+
+  def remove_warmer
+    puts "removing warmer #{@warmer_name}"
+    `curl -s -XDELETE #{@host}:#{@port}/#{@index}/_warmer/#{@warmer_name}`
+  end
+
+  def enable_warmer
+    puts "closing #{@index}"
+    `curl -s -XPOST '#{@host}:#{@port}/#{@index}/_close'`
+    puts "enabling warmer"
+    `curl -s -XPUT '#{@host}:#{@port}/#{@index}/_settings?pretty=true' -d '{"index.warmer.enabled":"true"}'`
+    puts "opening #{@index}"
+    `curl -s -XPOST '#{@host}:#{@port}/#{@index}/_open'`
+  end
+
+  def disable_warmer
+    puts "closing #{@index}"
+    `curl -s -XPOST '#{@host}:#{@port}/#{@index}/_close'`
+    puts "disabling warmer"
+    `curl -s -XPUT '#{@host}:#{@port}/#{@index}/_settings?pretty=true' -d '{"index.warmer.enabled":"false"}'`
+    puts "opening #{@index}"
+    `curl -s -XPOST '#{@host}:#{@port}/#{@index}/_open'`
+  end
+
+  def determine_interaction
+    unless @index.nil? || @host.nil? || @port.nil?
+      case command = @action.to_sym
+        when :add_warmer then add_warmer
+        when :remove_warmer then remove_warmer
+        when :enable_warmer then enable_warmer
+        when :disable_warmer then disable_warmer
+        else abort "#{command} is not a recognized action for determine_interaction from warmers_interface"
+      end
+    else
+      puts "index, host and port are required to interact with the warmers"
+    end
+  end
+end

data/src/main/java/com/infochimps/elasticsearch/ElasticSearchInputFormat.java

@@ -133,7 +133,7 @@ public class ElasticSearchInputFormat extends InputFormat<Text, Text> implements
             .setSize(requestSize)
             .execute()
             .actionGet();
-        this.numHits = response.hits().totalHits();
+        this.numHits = response.getHits().totalHits();
         if(numSplits > numHits) numSplits = numHits; // This could be bad
         this.numSplitRecords = (numHits/numSplits);
     }
@@ -206,7 +206,7 @@ public class ElasticSearchInputFormat extends InputFormat<Text, Text> implements
                 .setQuery(QueryBuilders.queryString(queryString))                
                 .execute()
                 .actionGet();
-            return response.hits().iterator();
+            return response.getHits().iterator();
         }
         
         @Override

data/src/main/java/com/infochimps/elasticsearch/ElasticSearchStreamingInputFormat.java

@@ -63,12 +63,24 @@ public class ElasticSearchStreamingInputFormat<K, V> implements InputFormat<K, V
     private static final String ES_PLUGINS     = "/usr/local/share/elasticsearch/plugins";
 
     private static final String ES_UNICAST_HOSTS_NAME  = "discovery.zen.ping.unicast.hosts";
+
+    private static final String ES_TRANSPORT_OPT     = "elasticsearch.transport";
+    private static final String ES_TRANSPORT         = "false";
+
+    private static final String ES_TRANSPORT_HOST_OPT     = "elasticsearch.transport.host";
+    private static final String ES_TRANSPORT_HOST         = "localhost";
+
+    private static final String ES_TRANSPORT_PORT_OPT     = "elasticsearch.transport.port";
+    private static final String ES_TRANSPORT_PORT         = "9300";
     
     private TransportClient client;
 
     public RecordReader<K, V> getRecordReader(InputSplit split, JobConf conf, Reporter reporter) {
 	setLocalElasticSearchInstallation(conf);
-        return (RecordReader) new ElasticSearchStreamingRecordReader(split, conf);
+	boolean esTransport        = new Boolean(conf.get(ES_TRANSPORT_OPT, ES_TRANSPORT));
+	String  esTransportHost    = conf.get(ES_TRANSPORT_HOST_OPT, ES_TRANSPORT_HOST);
+	Integer esTransportPort    = Integer.parseInt(conf.get(ES_TRANSPORT_PORT_OPT, ES_TRANSPORT_PORT));
+        return (RecordReader) new ElasticSearchStreamingRecordReader(split, conf, esTransport, esTransportHost, esTransportPort);
     }
 
     public InputSplit[] getSplits(JobConf conf, int requestedNumSplits) {
@@ -185,7 +197,7 @@ public class ElasticSearchStreamingInputFormat<K, V> implements InputFormat<K, V
 	    request.setQuery(queryJSON);
 	}
         SearchResponse response = request.execute().actionGet();
-        this.numHits = response.hits().totalHits();
+        this.numHits = response.getHits().totalHits();
 
 	LOG.info("Ran query: "+String.valueOf(numHits)+" hits");
     }

data/src/main/java/com/infochimps/elasticsearch/ElasticSearchStreamingOutputFormat.java

@@ -37,21 +37,28 @@ public class ElasticSearchStreamingOutputFormat<K, V> implements OutputFormat<K,
 
     private static final String ES_INDEX_FIELD_OPT = "elasticsearch.output.index.field";
     private static final String ES_INDEX_FIELD     = "_index";
-    private              String indexFieldName;
 
     private static final String ES_MAPPING_FIELD_OPT = "elasticsearch.output.mapping.field";
     private static final String ES_MAPPING_FIELD     = "_mapping";
-    private              String mappingFieldName;
 	
     private static final String ES_ID_FIELD_OPT = "elasticsearch.output.id.field";
     private static final String ES_ID_FIELD     = "_id";
-    private              String idFieldName;
+
+    private static final String ES_ROUTING_FIELD_OPT = "elasticsearch.output.routing.field";
+    private static final String ES_ROUTING_FIELD     = "_routing";
 	
     private static final String ES_BULK_SIZE_OPT     = "elasticsearch.output.bulk_size";
     private static final String ES_BULK_SIZE         = "1000";
-    private              int    bulkSize;
 
+    private static final String ES_TRANSPORT_OPT     = "elasticsearch.transport";
+    private static final String ES_TRANSPORT         = "false";
+
+    private static final String ES_TRANSPORT_HOST_OPT     = "elasticsearch.transport.host";
+    private static final String ES_TRANSPORT_HOST         = "localhost";
 
+    private static final String ES_TRANSPORT_PORT_OPT     = "elasticsearch.transport.port";
+    private static final String ES_TRANSPORT_PORT         = "9300";
+    
     // Elasticsearch internal settings required to make a client
     // connection.
     private static final String ES_CONFIG_OPT        = "es.config";
@@ -67,8 +74,16 @@ public class ElasticSearchStreamingOutputFormat<K, V> implements OutputFormat<K,
 	String  indexFieldName     = conf.get(ES_INDEX_FIELD_OPT,    ES_INDEX_FIELD);
 	String  mappingFieldName   = conf.get(ES_MAPPING_FIELD_OPT,  ES_MAPPING_FIELD);
 	String  idFieldName        = conf.get(ES_ID_FIELD_OPT,       ES_ID_FIELD);
+	String  routingFieldName   = conf.get(ES_ROUTING_FIELD_OPT,  ES_ROUTING_FIELD);
 	Integer bulkSize           = Integer.parseInt(conf.get(ES_BULK_SIZE_OPT, ES_BULK_SIZE));
-        return (RecordWriter) new ElasticSearchStreamingRecordWriter(defaultIndexName, defaultMappingName, indexFieldName, mappingFieldName, idFieldName, bulkSize);
+	boolean esTransport        = new Boolean(conf.get(ES_TRANSPORT_OPT, ES_TRANSPORT));
+	String  esTransportHost    = conf.get(ES_TRANSPORT_HOST_OPT, ES_TRANSPORT_HOST);
+	Integer esTransportPort    = Integer.parseInt(conf.get(ES_TRANSPORT_PORT_OPT, ES_TRANSPORT_PORT));
+
+        return (RecordWriter) new ElasticSearchStreamingRecordWriter(defaultIndexName, defaultMappingName,
+								     indexFieldName, mappingFieldName, idFieldName, routingFieldName,
+								     bulkSize,
+								     esTransport, esTransportHost, esTransportPort);
     }
     
     public void setLocalElasticSearchInstallation(JobConf conf) {

data/src/main/java/com/infochimps/elasticsearch/ElasticSearchStreamingRecordReader.java

@@ -17,6 +17,10 @@ import org.elasticsearch.common.unit.TimeValue;
 import org.elasticsearch.node.Node;
 import org.elasticsearch.node.NodeBuilder;
 import org.elasticsearch.client.Client;
+import org.elasticsearch.common.transport.InetSocketTransportAddress;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.settings.ImmutableSettings;
+import org.elasticsearch.client.transport.TransportClient;
 import org.elasticsearch.action.search.SearchRequestBuilder;
 import org.elasticsearch.action.search.SearchScrollRequestBuilder;
 
@@ -42,29 +46,66 @@ class ElasticSearchStreamingRecordReader<K, V> implements RecordReader<K, V> {
     private Node                        node;
     private Client                      client;
     private ElasticSearchStreamingSplit split;
-
+    private              boolean transport;
+    private              String  transportHost;
+    private              Integer transportPort;
+    
     private String  scrollId;
     private Integer recordsRead;
     private Iterator<SearchHit> hitsItr = null;
         
-    public ElasticSearchStreamingRecordReader(InputSplit split, JobConf conf) {
+    public ElasticSearchStreamingRecordReader(InputSplit split, JobConf conf, boolean transport, String transportHost, Integer transportPort) {
 	this.split         = (ElasticSearchStreamingSplit) split;
 	this.recordsRead   = 0;
 	this.requestSize   = Integer.parseInt(conf.get(ES_REQUEST_SIZE_OPT, ES_REQUEST_SIZE));
 	this.scrollTimeout = conf.get(ES_SCROLL_TIMEOUT_OPT, ES_SCROLL_TIMEOUT);
 	this.scroll        = new Scroll(TimeValue.parseTimeValue(this.scrollTimeout, defaultScrollTimeout));
+
+	this.transport     = transport;
+	this.transportHost = transportHost;
+	this.transportPort = transportPort;
 	
 	LOG.info("Initializing "+this.split.getSummary());
-	startEmbeddedClient();
+	if (transport) {
+	    this.client = buildTransportClient();
+	} else {
+	    startNode();
+	    this.client = node.client();
+	}
 	fetchNextHits();
     }
 
+    /**
+       Build a transport client that will connect to some
+       Elasticsearch node.
+       
+     */
+    private Client buildTransportClient() {
+	LOG.info("Connecting transport client to "+transportHost+":"+Integer.toString(transportPort));
+	Settings settings = ImmutableSettings.settingsBuilder().put("client.transport.ignore_cluster_name", "true").build();
+	return new TransportClient(settings).addTransportAddress(new InetSocketTransportAddress(transportHost, transportPort));
+    }
+
+    /**
+       Start an embedded Elasticsearch node.
+
+       The node will not store any data locally (non-datanode) but
+       will connect to a cluster using the default Elasticsearch
+       settings (those available in
+       /etc/elasticsearch/elasticsearch.yml).
+    */
+    private void startNode() {
+	LOG.info("Starting embedded Elasticsearch client (non-datanode)...");
+	this.node   = NodeBuilder.nodeBuilder().client(true).node();
+	LOG.info("Successfully joined Elasticsearch cluster '"+ClusterName.clusterNameFromSettings(node.settings())+'"');
+    }
+    
     private void fetchNextHits() {
 	if (scrollId == null) {
 	    LOG.info("Running initial scroll with timeout "+scrollTimeout);
 	    SearchRequestBuilder request  = split.initialScrollRequest(client, scroll, requestSize);
 	    SearchResponse       response = request.execute().actionGet();
-	    this.scrollId = response.scrollId();
+	    this.scrollId = response.getScrollId();
 	    LOG.info("Got scroll ID "+scrollId);
 	    // Do we need to call fetchNextHits() again here?  Or does
 	    // the initial request also itself contain the first set
@@ -75,9 +116,9 @@ class ElasticSearchStreamingRecordReader<K, V> implements RecordReader<K, V> {
 	    // LOG.info("Running query for scroll ID "+scrollId+" with timeout "+scrollTimeout);
 	    SearchScrollRequestBuilder request  = split.scrollRequest(client, scroll, scrollId);
 	    SearchResponse             response = request.execute().actionGet();
-	    this.scrollId = response.scrollId();
+	    this.scrollId = response.getScrollId();
 	    // LOG.info("Got scroll ID "+scrollId);
-	    this.hitsItr = response.hits().iterator();
+	    this.hitsItr = response.getHits().iterator();
 	}
     }
 
@@ -151,26 +192,14 @@ class ElasticSearchStreamingRecordReader<K, V> implements RecordReader<K, V> {
 
     @Override
 	public void close() throws IOException {
-	stopEmbeddedClient();
-    }
-
-    //
-    // == Connecting to Elasticsearch ==
-    // 
-
-    private void startEmbeddedClient() {
-	LOG.info("Starting embedded Elasticsearch client (non-datanode)...");
-	this.node   = NodeBuilder.nodeBuilder().client(true).node();
-	this.client = node.client();
-	LOG.info("Successfully joined Elasticsearch cluster '"+ClusterName.clusterNameFromSettings(node.settings())+'"');
-    }
-
-    private void stopEmbeddedClient() {
-	LOG.info("Stopping embedded Elasticsearch client...");
-	if (client != null) client.close();
-	if (node   != null) node.close();
-	LOG.info("Left Elasticsearch cluster");
+	if (client != null) {
+	    LOG.info("Shutting down Elasticsearch client...");
+	    client.close();
+	}
+	if (node != null) {
+	    LOG.info("Shutting down Elasticsearch node...");
+	    node.close();
+	}
     }
     
-
 }

data/src/main/java/com/infochimps/elasticsearch/ElasticSearchStreamingRecordWriter.java

@@ -20,8 +20,13 @@ import org.elasticsearch.node.Node;
 import org.elasticsearch.node.NodeBuilder;
 import org.elasticsearch.client.Client;
 import org.elasticsearch.client.Requests;
+import org.elasticsearch.common.transport.InetSocketTransportAddress;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.common.settings.ImmutableSettings;
+import org.elasticsearch.client.transport.TransportClient;
 import org.elasticsearch.action.bulk.BulkRequestBuilder;
 import org.elasticsearch.action.bulk.BulkResponse;
+import org.elasticsearch.action.index.IndexRequest;
 import org.elasticsearch.ExceptionsHelper;
 
 import org.codehaus.jackson.map.ObjectMapper;
@@ -36,6 +41,7 @@ class ElasticSearchStreamingRecordWriter<K, V> implements RecordWriter<K, V> {
     private String  indexFieldName;
     private String  mappingFieldName;
     private String  idFieldName;
+    private String  routingFieldName;
     private Integer bulkSize;
 
     // Bookkeeping
@@ -48,6 +54,9 @@ class ElasticSearchStreamingRecordWriter<K, V> implements RecordWriter<K, V> {
     private              Node               node;
     private              Client             client;
     private volatile     BulkRequestBuilder currentRequest;
+    private              boolean transport;
+    private              String  transportHost;
+    private              Integer transportPort;
 
     // JSON parsing
     private              ObjectMapper       mapper;
@@ -56,35 +65,53 @@ class ElasticSearchStreamingRecordWriter<K, V> implements RecordWriter<K, V> {
     // == Lifecycle ==
     // 
 
-    public ElasticSearchStreamingRecordWriter(String defaultIndexName, String defaultMappingName, String indexFieldName, String mappingFieldName, String idFieldName, Integer bulkSize) {
-	this.defaultIndexName = defaultIndexName;
-	this.defaultMappingName  = defaultMappingName;
-	this.indexFieldName   = indexFieldName;
-	this.mappingFieldName    = mappingFieldName;
-	this.idFieldName      = idFieldName;
-	this.bulkSize         = bulkSize;
+    public ElasticSearchStreamingRecordWriter(String defaultIndexName, String defaultMappingName, String indexFieldName, String mappingFieldName, String idFieldName, String routingFieldName, Integer bulkSize, boolean transport, String transportHost, Integer transportPort) {
+	this.defaultIndexName   = defaultIndexName;
+	this.defaultMappingName = defaultMappingName;
+	this.indexFieldName     = indexFieldName;
+	this.mappingFieldName   = mappingFieldName;
+	this.idFieldName        = idFieldName;
+	this.routingFieldName   = routingFieldName;
+	this.bulkSize           = bulkSize;
+	this.transport          = transport;
+	this.transportHost      = transportHost;
+	this.transportPort      = transportPort;
 	
 	LOG.info("Writing "+Integer.toString(bulkSize)+" records per batch");
 	LOG.info("Using default target /"+defaultIndexName+"/"+defaultMappingName);
 	LOG.info("Records override default target with index field '"+indexFieldName+"', mapping field '"+mappingFieldName+"', and ID field '"+idFieldName);
-            
-	startEmbeddedClient();
+	if (transport) {
+	    this.client = buildTransportClient();
+	} else {
+	    startNode();
+	    this.client = node.client();
+	}
 	this.currentRequest = client.prepareBulk();
 	this.mapper = new ObjectMapper();
     }
 
     /**
-       Start an embedded Elasticsearch client.  The client will not be
-       a data node and will not store data locally.
+       Build a transport client that will connect to some
+       Elasticsearch node.
+       
+     */
+    private Client buildTransportClient() {
+	LOG.info("Connecting transport client to "+transportHost+":"+Integer.toString(transportPort));
+	Settings settings = ImmutableSettings.settingsBuilder().put("client.transport.ignore_cluster_name", "true").build();
+	return new TransportClient(settings).addTransportAddress(new InetSocketTransportAddress(transportHost, transportPort));
+    }
+
+    /**
+       Start an embedded Elasticsearch node.
 
-       The client will connect to the target Elasticsearch cluster as
-       a client node, enabling one-hop writes for all data.  See
-       http://www.elasticsearch.org/guide/reference/java-api/client.html
+       The node will not store any data locally (non-datanode) but
+       will connect to a cluster using the default Elasticsearch
+       settings (those available in
+       /etc/elasticsearch/elasticsearch.yml).
     */
-    private void startEmbeddedClient() {
+    private void startNode() {
 	LOG.info("Starting embedded Elasticsearch client (non-datanode)...");
 	this.node   = NodeBuilder.nodeBuilder().client(true).node();
-	this.client = node.client();
 	LOG.info("Successfully joined Elasticsearch cluster '"+ClusterName.clusterNameFromSettings(node.settings())+'"');
     }
 
@@ -95,10 +122,14 @@ class ElasticSearchStreamingRecordWriter<K, V> implements RecordWriter<K, V> {
     */
     public void close(Reporter reporter) throws IOException {
 	sendBulkRequestIfMoreThan(0);
-	LOG.info("Shutting down Elasticsearch client...");
-	if (client != null) client.close();
-	if (node   != null) node.close();
-	LOG.info("Successfully shut down Elasticsearch client");
+	if (client != null) {
+	    LOG.info("Shutting down Elasticsearch client...");
+	    client.close();
+	}
+	if (node != null) {
+	    LOG.info("Shutting down Elasticsearch node...");
+	    node.close();
+	}
     }
 
     //
@@ -122,12 +153,18 @@ class ElasticSearchStreamingRecordWriter<K, V> implements RecordWriter<K, V> {
 
     private void index(String json) throws IOException {
 	Map<String, Object> record = mapper.readValue(json, Map.class);
+	IndexRequest request = null;
 	if (record.containsKey(idFieldName)) {
 	    Object idValue = record.get(idFieldName);
-	    currentRequest.add(Requests.indexRequest(indexNameForRecord(record)).id(String.valueOf(idValue)).type(mappingNameForRecord(record)).create(false).source(json));
+	     request = Requests.indexRequest(indexNameForRecord(record)).id(String.valueOf(idValue)).type(mappingNameForRecord(record)).create(false).source(json);
 	} else {
-	    currentRequest.add(Requests.indexRequest(indexNameForRecord(record)).type(mappingNameForRecord(record)).source(json));
+	    request = Requests.indexRequest(indexNameForRecord(record)).type(mappingNameForRecord(record)).source(json);
+	}
+	if (record.containsKey(routingFieldName)) {
+	    Object routingValue = record.get(routingFieldName);
+		request.routing(String.valueOf(routingValue));
 	}
+	currentRequest.add(request);
     }
 
     private String indexNameForRecord(Map<String, Object> record) {