com.stratio.ingestion.sink.decision.StratioDecisionSink.java Source code

Java tutorial

Introduction

Here is the source code for com.stratio.ingestion.sink.decision.StratioDecisionSink.java

Source

/**
 * Copyright (C) 2014 Stratio (http://stratio.com)
 *
 * 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 com.stratio.ingestion.sink.decision;

import java.io.File;
import java.io.FileInputStream;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Locale;
import java.util.Map;

import org.apache.commons.io.IOUtils;
import org.apache.flume.Channel;
import org.apache.flume.ChannelException;
import org.apache.flume.Context;
import org.apache.flume.Event;
import org.apache.flume.EventDeliveryException;
import org.apache.flume.Transaction;
import org.apache.flume.conf.Configurable;
import org.apache.flume.instrumentation.SinkCounter;
import org.apache.flume.sink.AbstractSink;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.stratio.decision.api.IStratioStreamingAPI;
import com.stratio.decision.api.StratioStreamingAPIFactory;
import com.stratio.decision.api.messaging.ColumnNameType;
import com.stratio.decision.api.messaging.ColumnNameValue;
import com.stratio.decision.commons.constants.ColumnType;
import com.stratio.decision.commons.exceptions.StratioEngineConnectionException;
import com.stratio.decision.commons.exceptions.StratioStreamingException;

public class StratioDecisionSink extends AbstractSink implements Configurable {

    private static final Logger log = LoggerFactory.getLogger(StratioDecisionSink.class);
    private static final int DEFAULT_BATCH_SIZE = 20;
    private static final String DEFAULT_ZOOKEEPER = "localhost:2181";
    private static final String DEFAULT_KAFKA = "localhost:9092";
    private static final String DEFAULT_TOPIC = "stratio_decision_data";
    private static final String CONF_BATCH_SIZE = "batchSize";
    private static final String ZOOKEEPER = "zookeeper";
    private static final String KAFKA = "kafka";
    private static final String STREAM_DEFINITION_FILE = "streamDefinitionFile";
    private static final String TOPIC = "topic";
    private static final String TOPIC_EVENT_HEADER = "_topic";

    private SinkCounter sinkCounter;
    private int batchsize;
    private IStratioStreamingAPI stratioStreamingAPI;
    private String zookeeper;
    private String kafka;
    private String streamName;
    private String topic = "";
    private List<StreamField> streamFields;

    public StratioDecisionSink() {
        super();
    }

    public void configure(Context context) {
        this.batchsize = context.getInteger(CONF_BATCH_SIZE, DEFAULT_BATCH_SIZE);
        this.sinkCounter = new SinkCounter(this.getName());
        this.zookeeper = context.getString(ZOOKEEPER, DEFAULT_ZOOKEEPER);
        this.kafka = context.getString(KAFKA, DEFAULT_KAFKA);

        //if (context.getString(TOPIC) != null)
        //    this.topic= context.getString(TOPIC);
        this.topic = context.getString(TOPIC, "");

        log.info("Configuring Stratio Decision Sink: {zookeeper= " + this.zookeeper + ", kafka= " + this.kafka
                + ", topic= " + this.topic + ", batchSize= " + this.batchsize + ", sinkCounter= " + this.sinkCounter
                + "}");
        //else
        //this.topic= "";

        String columnDefinitionFile = context.getString(STREAM_DEFINITION_FILE);
        com.stratio.ingestion.sink.decision.StreamDefinitionParser parser = new StreamDefinitionParser(
                readJsonFromFile(new File(columnDefinitionFile)));
        StreamDefinition theStreamDefinition = parser.parse();
        this.streamName = theStreamDefinition.getStreamName();
        this.streamFields = theStreamDefinition.getFields();

        try {
            this.stratioStreamingAPI = StratioStreamingAPIFactory.create().withServerConfig(kafka, zookeeper)
                    .init();
        } catch (StratioEngineConnectionException e) {
            throw new StratioDecisionSinkException(e);
        }
    }

    @Override
    public synchronized void start() {
        super.start();
        createStream();
        this.sinkCounter.start();
    }

    private void createStream() {
        try {
            String streamName = this.streamName;
            List<ColumnNameType> columnList = new ArrayList<ColumnNameType>();
            for (StreamField streamField : this.streamFields) {
                ColumnNameType streamColumn = new ColumnNameType(streamField.getName(),
                        parseStreamField(streamField.getType()));
                columnList.add(streamColumn);
            }
            stratioStreamingAPI.createStream(streamName, columnList);
        } catch (StratioStreamingException e) {
            e.printStackTrace();
        }
    }

    public Status process() throws EventDeliveryException {
        Status status = Status.BACKOFF;
        Transaction transaction = this.getChannel().getTransaction();
        try {
            transaction.begin();
            List<Event> eventList = this.takeEventsFromChannel(this.getChannel());
            status = Status.READY;
            if (!eventList.isEmpty()) {
                if (eventList.size() == this.batchsize) {
                    this.sinkCounter.incrementBatchCompleteCount();
                } else {
                    this.sinkCounter.incrementBatchUnderflowCount();
                }
                for (Event event : eventList) {
                    List<ColumnNameValue> columnNameValueList = getColumnNameValueListFromEvent(event);

                    if (event.getHeaders().containsKey(TOPIC_EVENT_HEADER)) {
                        // If we've defined the _topic header using in the event, we send the data to this topic
                        stratioStreamingAPI.insertData(this.streamName, columnNameValueList,
                                getEventDefinedTopicName(event.getHeaders().get(TOPIC_EVENT_HEADER)), false);

                    } else if (!this.topic.isEmpty()) {
                        // If we've specified a topic in the properties file we send the data to that topic
                        stratioStreamingAPI.insertData(this.streamName, columnNameValueList, this.topic, false);
                    } else {
                        // In other case we send the data to default topic, don't specifying any topic name
                        stratioStreamingAPI.insertData(this.streamName, columnNameValueList);
                    }
                }
                this.sinkCounter.addToEventDrainSuccessCount(eventList.size());
            } else {
                this.sinkCounter.incrementBatchEmptyCount();
            }
            transaction.commit();
            status = Status.READY;
        } catch (ChannelException e) {
            e.printStackTrace();
            transaction.rollback();
            status = Status.BACKOFF;
            this.sinkCounter.incrementConnectionFailedCount();
        } catch (Throwable t) {
            t.printStackTrace();
            transaction.rollback();
            status = Status.BACKOFF;
            if (t instanceof Error) {
                throw new StratioDecisionSinkException(t);
            }
        } finally {
            transaction.close();
        }
        return status;
    }

    private String getEventDefinedTopicName(String eventTopicName) {
        //        return DEFAULT_TOPIC + "_" + eventTopicName;
        return eventTopicName;
    }

    private List<ColumnNameValue> getColumnNameValueListFromEvent(Event event) {
        List<ColumnNameValue> columnNameValues = new ArrayList<ColumnNameValue>();
        Map<String, String> headers = event.getHeaders();
        for (StreamField field : streamFields) {
            String fieldContent = headers.get(field.getName());
            columnNameValues.add(new ColumnNameValue(field.getName(), fieldContent));
        }
        return columnNameValues;
    }

    private List<Event> takeEventsFromChannel(Channel channel) {
        List<Event> events = new ArrayList<Event>();
        for (int i = 0; i < this.batchsize; i++) {
            this.sinkCounter.incrementEventDrainAttemptCount();
            events.add(channel.take());
        }
        events.removeAll(Collections.singleton(null));
        return events;
    }

    private static final String readJsonFromFile(File file) {
        FileInputStream inputStream;
        try {
            inputStream = new FileInputStream(file);
            return IOUtils.toString(inputStream, "UTF-8");
        } catch (Exception e) {
            throw new StratioDecisionSinkException(e);
        }
    }

    private ColumnType parseStreamField(String field) {
        try {
            return ColumnType.valueOf(field.toUpperCase(Locale.ENGLISH));
        } catch (IllegalArgumentException ex) {
            //TODO: log something
        }
        return ColumnType.STRING;
    }
}