/*
* Copyright 2013-2014 Amazon.com, Inc. or its affiliates. All Rights Reserved.
*
* Licensed under the Amazon Software License (the "License").
* You may not use this file except in compliance with the License.
* A copy of the License is located at
*
* http://aws.amazon.com/asl/
*
* or in the "license" file accompanying this file. This file 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 samples;
import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStream;
import java.io.InputStreamReader;
import java.nio.ByteBuffer;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import samples.utils.KinesisUtils;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.regions.RegionUtils;
import com.amazonaws.services.kinesis.AmazonKinesisClient;
import com.amazonaws.services.kinesis.connectors.KinesisConnectorConfiguration;
import com.amazonaws.services.kinesis.model.PutRecordRequest;
import com.fasterxml.jackson.databind.ObjectMapper;
/**
* This class is a data source for supplying input to the Amazon Kinesis stream. It reads lines from the
* input file specified in the constructor and emits them by calling String.getBytes() into the
* stream defined in the KinesisConnectorConfiguration.
*/
public class StreamSource implements Runnable {
private static Log LOG = LogFactory.getLog(StreamSource.class);
protected AmazonKinesisClient kinesisClient;
protected KinesisConnectorConfiguration config;
protected final String inputFile;
protected final boolean loopOverInputFile;
protected ObjectMapper objectMapper;
/**
* Creates a new StreamSource.
*
* @param config
* Configuration to determine which stream to put records to and get {@link AWSCredentialsProvider}
* @param inputFile
* File containing record data to emit on each line
*/
public StreamSource(KinesisConnectorConfiguration config, String inputFile) {
this(config, inputFile, false);
}
/**
* Creates a new StreamSource.
*
* @param config
* Configuration to determine which stream to put records to and get {@link AWSCredentialsProvider}
* @param inputFile
* File containing record data to emit on each line
* @param loopOverStreamSource
* Loop over the stream source to continually put records
*/
public StreamSource(KinesisConnectorConfiguration config, String inputFile, boolean loopOverStreamSource) {
this.config = config;
this.inputFile = inputFile;
this.loopOverInputFile = loopOverStreamSource;
this.objectMapper = new ObjectMapper();
kinesisClient = new AmazonKinesisClient(config.AWS_CREDENTIALS_PROVIDER);
kinesisClient.setRegion(RegionUtils.getRegion(config.REGION_NAME));
if (config.KINESIS_ENDPOINT != null) {
kinesisClient.setEndpoint(config.KINESIS_ENDPOINT);
}
KinesisUtils.createInputStream(config);
}
@Override
public void run() {
int iteration = 0;
do {
InputStream inputStream = Thread.currentThread().getContextClassLoader().getResourceAsStream(inputFile);
if (inputStream == null) {
throw new IllegalStateException("Could not find input file: " + inputFile);
}
if (loopOverInputFile) {
LOG.info("Starting iteration " + iteration + " over input file.");
}
try {
processInputStream(inputStream, iteration);
} catch (IOException e) {
LOG.error("Encountered exception while putting data in source stream.", e);
break;
}
iteration++;
} while (loopOverInputFile);
}
/**
* Process the input file and send PutRecordRequests to Amazon Kinesis.
*
* This function serves to Isolate StreamSource logic so subclasses
* can process input files differently.
*
* @param inputStream
* the input stream to process
* @param iteration
* the iteration if looping over file
* @throws IOException
* throw exception if error processing inputStream.
*/
protected void processInputStream(InputStream inputStream, int iteration) throws IOException {
try (BufferedReader br = new BufferedReader(new InputStreamReader(inputStream))) {
String line;
int lines = 0;
while ((line = br.readLine()) != null) {
KinesisMessageModel kinesisMessageModel = objectMapper.readValue(line, KinesisMessageModel.class);
PutRecordRequest putRecordRequest = new PutRecordRequest();
putRecordRequest.setStreamName(config.KINESIS_INPUT_STREAM);
putRecordRequest.setData(ByteBuffer.wrap(line.getBytes()));
putRecordRequest.setPartitionKey(Integer.toString(kinesisMessageModel.getUserid()));
kinesisClient.putRecord(putRecordRequest);
lines++;
}
LOG.info("Added " + lines + " records to stream source.");
}
}
@Override
protected void finalize() throws Throwable {
super.finalize();
}
}