com.shapestone.event.stream.producer.KinesisProducer Maven / Gradle / Ivy

Go to download
package com.shapestone.event.stream.producer;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
import com.amazonaws.handlers.AsyncHandler;
import com.amazonaws.services.kinesis.AmazonKinesisAsyncClient;
import com.amazonaws.services.kinesis.model.PutRecordRequest;
import com.amazonaws.services.kinesis.model.PutRecordResult;
import rx.Observable;

import java.io.UnsupportedEncodingException;
import java.nio.ByteBuffer;
import java.util.UUID;

import static rx.Observable.create;
import static rx.schedulers.Schedulers.io;

/**
 * Name: Michael Williams
 * Date: 3/4/17.
 */
public class KinesisProducer {

  private AmazonKinesisAsyncClient amazonKinesisAsyncClient;
  private String streamName;

  public KinesisProducer(String streamName) {
    final ProfileCredentialsProvider provider = new ProfileCredentialsProvider();
    this.amazonKinesisAsyncClient = new AmazonKinesisAsyncClient(provider);
    this.streamName = streamName;

  }

  public KinesisProducer(String streamName, AWSCredentials awsCredentials) {
    this.amazonKinesisAsyncClient = new AmazonKinesisAsyncClient(awsCredentials);
    this.streamName = streamName;
  }

  public Observable writeStream(String blobMessage) {
    final byte[] utf8s;
    try {
      utf8s = blobMessage.getBytes("UTF8");
    } catch (UnsupportedEncodingException e) {
      throw new IllegalArgumentException(e);
    }
    return writeStream(ByteBuffer.wrap(utf8s));
  }

  public Observable writeStream(byte[] bytes) {
    return writeStream(ByteBuffer.wrap(bytes));
  }

  public Observable writeStream(ByteBuffer byteBuffer) {
    final Observable observable = create(subscriber -> {

      final PutRecordRequest putRecordRequest = new PutRecordRequest()
        .withStreamName(streamName)
        .withPartitionKey(UUID.randomUUID().toString())
        .withData(byteBuffer);

      this.amazonKinesisAsyncClient.putRecordAsync(putRecordRequest,
        new AsyncHandler() {
        @Override
        public void onError(Exception exception) {
          subscriber.onError(exception);
          subscriber.onCompleted();
        }

        @Override
        public void onSuccess(PutRecordRequest request, PutRecordResult putRecordResult) {
          subscriber.onNext(Boolean.TRUE);
          subscriber.onCompleted();
        }
      });
    });

    return observable.subscribeOn(io());
  }

}