ElasticSearchSink.java

package org.elasticsearch.flume;

import static org.elasticsearch.common.xcontent.XContentFactory.jsonBuilder;
import static org.elasticsearch.node.NodeBuilder.nodeBuilder;

import java.io.IOException;
import java.nio.charset.Charset;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Date;
import java.util.List;
import java.util.Map;
import java.util.concurrent.atomic.AtomicLong;

import com.cloudera.flume.conf.SinkFactory.SinkBuilder;
import com.cloudera.flume.core.Event;
import com.cloudera.flume.core.EventSink;
import com.cloudera.flume.reporter.ReportEvent;
import com.cloudera.util.Pair;
import org.elasticsearch.client.Client;
import org.elasticsearch.client.transport.TransportClient;
import org.elasticsearch.cluster.ClusterName;
import org.elasticsearch.common.settings.ImmutableSettings;
import org.elasticsearch.common.settings.Settings;
import org.elasticsearch.common.transport.InetSocketTransportAddress;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.common.xcontent.XContentFactory;
import org.elasticsearch.common.xcontent.XContentParser;
import org.elasticsearch.common.xcontent.XContentType;
import org.elasticsearch.node.Node;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class ElasticSearchSink extends EventSink.Base {

    private static final Logger LOG = LoggerFactory.getLogger(ElasticSearchSink.class);
    private static final String DEFAULT_INDEX_NAME = "flume";
    private static final String DEFAULT_LOG_TYPE = "log";
    private static final int DEFAULT_ELASTICSEARCH_PORT = 9300;

    private Node node;
    private Client client;
    private String indexName = DEFAULT_INDEX_NAME;
    private String indexPattern = null;
    private String indexType = DEFAULT_LOG_TYPE;

    private Charset charset = Charset.defaultCharset();

    private String[] hostNames = new String[0];
    private String clusterName = ClusterName.DEFAULT.value();

    // Enabled only for testing
    private boolean localOnly = false;

    private AtomicLong eventErrorCount = new AtomicLong(0L);
    private static final String NO_OF_FAILED_EVENTS = "NO_OF_FAILED_EVENTS";

    @Override
    public void append(Event e) throws IOException {
        try {
            XContentBuilder builder = jsonBuilder()
                    .startObject()
                    .field("timestamp", new Date(e.getTimestamp()))
                    .field("host", e.getHost())
                    .field("priority", e.getPriority().name());

            addBody(builder, e.getBody());

            addAttrs(builder, e.getAttrs());

            index(e, builder);
        } catch (Exception ex) {
            LOG.error("Error Processing event: {}", e.toString(), ex);
            eventErrorCount.incrementAndGet();
        }
    }

    @Override
    public synchronized ReportEvent getMetrics() {
        ReportEvent event = new ReportEvent("ElasticSearchSink");
        event.setLongMetric(NO_OF_FAILED_EVENTS, eventErrorCount.longValue());
        return event;
    }

    private void addBody(XContentBuilder builder, byte[] data) throws IOException {
        XContentType contentType = XContentFactory.xContentType(data);

        if (contentType == null) {
            builder.startObject("message");
            addSimpleField(builder, "text", data);
            builder.endObject();
        } else {
            addComplexField(builder, "message", contentType, data);
        }
    }

    private void addAttrs(XContentBuilder builder, Map<String, byte[]> attrs) throws IOException {
        builder.startObject("fields");
        for (Map.Entry<String, byte[]> entry : attrs.entrySet()) {
            if (LOG.isDebugEnabled()) {
                LOG.debug("field: {}, data: {}", entry.getKey(), new String(entry.getValue()));
            }
            addField(builder, entry.getKey(), entry.getValue());
        }
        builder.endObject();
    }

    private void addField(XContentBuilder builder, String fieldName, byte[] data) throws IOException {
        XContentType contentType = XContentFactory.xContentType(data);
        if (contentType == null) {
            addSimpleField(builder, fieldName, data);
        } else {
            addComplexField(builder, fieldName, contentType, data);
        }
    }

    private void addSimpleField(XContentBuilder builder, String fieldName, byte[] data) throws IOException {
        builder.field(fieldName, new String(data, charset));
    }

    private void addComplexField(XContentBuilder builder, String fieldName, XContentType contentType, byte[] data)
            throws IOException {
        XContentParser parser = null;
        try {
            parser = XContentFactory.xContent(contentType).createParser(data);
            parser.nextToken();
            builder.field(fieldName).copyCurrentStructure(parser);
        } finally {
            if (parser != null) {
                parser.close();
            }
        }
    }

    private void index(Event e, XContentBuilder builder) {
        String iName = indexName;
        if (indexPattern != null) {
            iName = e.escapeString(indexPattern);
        }
        client.prepareIndex(iName, indexType, null)
                .setSource(builder)
                .execute()
                .actionGet();

        if (!iName.equals(indexName)) {
            client.admin().indices().prepareAliases().addAlias(iName, indexName).execute().actionGet();
        }
    }

    @Override
    public void close() throws IOException, InterruptedException {
        super.close();

        if (client != null) {
            client.close();
        }
        if (node != null) {
            node.close();
        }
    }

    @Override
    public void open() throws IOException, InterruptedException {
        super.open();

        if (hostNames.length == 0) {
            LOG.info("Using ES AutoDiscovery mode");
            node = nodeBuilder().client(true).clusterName(clusterName).local(localOnly).node();
            client = node.client();
        } else {
            LOG.info("Using provided ES hostnames: {} ", Arrays.toString(hostNames));

            Settings settings = ImmutableSettings.settingsBuilder()
                    .put("cluster.name", clusterName)
                    .build();

            TransportClient transportClient = new TransportClient(settings);
            for (String esHostName : hostNames) {
                LOG.info("Adding TransportClient: {}", esHostName);
                transportClient = transportClient.addTransportAddress(new InetSocketTransportAddress(esHostName,
                        DEFAULT_ELASTICSEARCH_PORT));
            }
            client = transportClient;
        }

    }

    /**
     * This is a special function used by the SourceFactory to pull in this class as a plugin sink.
     */
    public static List<Pair<String, SinkBuilder>> getSinkBuilders() {
        List<Pair<String, SinkBuilder>> builders =
                new ArrayList<Pair<String, SinkBuilder>>();

        builders.add(new Pair<String, SinkBuilder>("elasticSearchSink", new ElasticSearchSinkBuilder()));
        return builders;
    }

    public String getClusterName() {
        return clusterName;
    }

    public void setClusterName(String clusterName) {
        this.clusterName = clusterName;
    }

    public String getIndexName() {
        return indexName;
    }

    public void setIndexName(String indexName) {
        this.indexName = indexName;
    }

    public String getIndexPattern() {
        return indexPattern;
    }

    public void setIndexPattern(String indexPattern) {
        this.indexPattern = indexPattern;
    }

    public String getIndexType() {
        return indexType;
    }

    public void setIndexType(String indexType) {
        this.indexType = indexType;
    }

    public void setHostNames(String[] hostNames) {
        this.hostNames = hostNames;
    }

    public String[] getHostNames() {
        return hostNames;
    }

    void setLocalOnly(boolean localOnly) {
        this.localOnly = localOnly;
    }

    boolean isLocalOnly() {
        return localOnly;
    }
}