Examples with Record com.amazonaws.services.kinesisfirehose.model.Record used on opensource projects

Example 1 with Record

use of com.amazonaws.services.kinesisfirehose.model.Record in project nifi by apache.

the class PutKinesisFirehose method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    final int batchSize = context.getProperty(BATCH_SIZE).asInteger();
    final long maxBufferSizeBytes = context.getProperty(MAX_MESSAGE_BUFFER_SIZE_MB).asDataSize(DataUnit.B).longValue();
    List<FlowFile> flowFiles = filterMessagesByMaxSize(session, batchSize, maxBufferSizeBytes, AWS_KINESIS_FIREHOSE_ERROR_MESSAGE);
    HashMap<String, List<FlowFile>> hashFlowFiles = new HashMap<>();
    HashMap<String, List<Record>> recordHash = new HashMap<String, List<Record>>();
    final AmazonKinesisFirehoseClient client = getClient();
    try {
        List<FlowFile> failedFlowFiles = new ArrayList<>();
        List<FlowFile> successfulFlowFiles = new ArrayList<>();
        // Prepare batch of records
        for (int i = 0; i < flowFiles.size(); i++) {
            FlowFile flowFile = flowFiles.get(i);
            final String firehoseStreamName = context.getProperty(KINESIS_FIREHOSE_DELIVERY_STREAM_NAME).evaluateAttributeExpressions(flowFile).getValue();
            final ByteArrayOutputStream baos = new ByteArrayOutputStream();
            session.exportTo(flowFile, baos);
            if (recordHash.containsKey(firehoseStreamName) == false) {
                recordHash.put(firehoseStreamName, new ArrayList<>());
            }
            if (hashFlowFiles.containsKey(firehoseStreamName) == false) {
                hashFlowFiles.put(firehoseStreamName, new ArrayList<>());
            }
            hashFlowFiles.get(firehoseStreamName).add(flowFile);
            recordHash.get(firehoseStreamName).add(new Record().withData(ByteBuffer.wrap(baos.toByteArray())));
        }
        for (Map.Entry<String, List<Record>> entryRecord : recordHash.entrySet()) {
            String streamName = entryRecord.getKey();
            List<Record> records = entryRecord.getValue();
            if (records.size() > 0) {
                // Send the batch
                PutRecordBatchRequest putRecordBatchRequest = new PutRecordBatchRequest();
                putRecordBatchRequest.setDeliveryStreamName(streamName);
                putRecordBatchRequest.setRecords(records);
                PutRecordBatchResult results = client.putRecordBatch(putRecordBatchRequest);
                // Separate out the successful and failed flow files
                List<PutRecordBatchResponseEntry> responseEntries = results.getRequestResponses();
                for (int i = 0; i < responseEntries.size(); i++) {
                    PutRecordBatchResponseEntry entry = responseEntries.get(i);
                    FlowFile flowFile = hashFlowFiles.get(streamName).get(i);
                    Map<String, String> attributes = new HashMap<>();
                    attributes.put(AWS_KINESIS_FIREHOSE_RECORD_ID, entry.getRecordId());
                    flowFile = session.putAttribute(flowFile, AWS_KINESIS_FIREHOSE_RECORD_ID, entry.getRecordId());
                    if (StringUtils.isBlank(entry.getErrorCode()) == false) {
                        attributes.put(AWS_KINESIS_FIREHOSE_ERROR_CODE, entry.getErrorCode());
                        attributes.put(AWS_KINESIS_FIREHOSE_ERROR_MESSAGE, entry.getErrorMessage());
                        flowFile = session.putAllAttributes(flowFile, attributes);
                        failedFlowFiles.add(flowFile);
                    } else {
                        flowFile = session.putAllAttributes(flowFile, attributes);
                        successfulFlowFiles.add(flowFile);
                    }
                }
                recordHash.get(streamName).clear();
                records.clear();
            }
        }
        if (failedFlowFiles.size() > 0) {
            session.transfer(failedFlowFiles, REL_FAILURE);
            getLogger().error("Failed to publish to kinesis firehose {}", new Object[] { failedFlowFiles });
        }
        if (successfulFlowFiles.size() > 0) {
            session.transfer(successfulFlowFiles, REL_SUCCESS);
            getLogger().info("Successfully published to kinesis firehose {}", new Object[] { successfulFlowFiles });
        }
    } catch (final Exception exception) {
        getLogger().error("Failed to publish to kinesis firehose {} with exception {}", new Object[] { flowFiles, exception });
        session.transfer(flowFiles, REL_FAILURE);
        context.yield();
    }
}

Example 2 with Record

use of com.amazonaws.services.kinesisfirehose.model.Record in project bender by Nextdoor.

the class FirehoseTransportBufferBatch method add.

@Override
public boolean add(InternalEvent ievent) throws IllegalStateException, IOException {
    byte[] record = serializer.serialize(ievent);
    /*
     * Restrict size of individual record
     */
    if (record.length > MAX_RECORD_SIZE) {
        throw new IOException("serialized event is " + record.length + " larger than max of " + MAX_RECORD_SIZE);
    }
    /*
     * Write record if there's room in buffer
     */
    if (dataRecords.size() >= MAX_RECORDS) {
        logger.trace("hit record index max");
        throw new IllegalStateException("reached max payload size");
    } else {
        if (cos.getByteCount() + record.length < MAX_RECORD_SIZE) {
            cos.write(record);
            return true;
        }
        /*
       * If current record is full then flush buffer to a Firehose Record and create a new buffer
       */
        logger.trace("creating new datarecord");
        ByteBuffer data = ByteBuffer.wrap(baos.toByteArray());
        this.dataRecords.add(new Record().withData(data));
        baos.reset();
        cos.resetByteCount();
        cos.resetCount();
        /*
       * If we hit the max number of Firehose Records (4) then notify IPC service that this buffer
       * needs to be sent.
       */
        if (dataRecords.size() >= MAX_RECORDS) {
            logger.trace("hit record index max");
            throw new IllegalStateException("reached max payload size");
        }
        /*
       * Otherwise write the record to the empty internal buffer
       */
        cos.write(record);
    }
    return true;
}

Example 3 with Record

use of com.amazonaws.services.kinesisfirehose.model.Record in project bender by Nextdoor.

the class FirehoseTransportBufferBatch method close.

@Override
public void close() {
    if (this.cos.getByteCount() != 0 && this.dataRecords.size() < MAX_RECORDS) {
        logger.trace("flushing remainder of buffer");
        ByteBuffer data = ByteBuffer.wrap(baos.toByteArray());
        this.dataRecords.add(new Record().withData(data));
    }
    try {
        this.baos.close();
    } catch (IOException e) {
    }
}

Example 4 with Record

use of com.amazonaws.services.kinesisfirehose.model.Record in project camel by apache.

the class KinesisFirehoseProducer method createRequest.

private PutRecordRequest createRequest(Exchange exchange) {
    ByteBuffer body = exchange.getIn().getBody(ByteBuffer.class);
    Record record = new Record();
    record.setData(body);
    PutRecordRequest putRecordRequest = new PutRecordRequest();
    putRecordRequest.setDeliveryStreamName(getEndpoint().getStreamName());
    putRecordRequest.setRecord(record);
    return putRecordRequest;
}

Example 5 with Record

use of com.amazonaws.services.kinesisfirehose.model.Record in project bender by Nextdoor.

the class FirehoseTransportBufferSimple method add.

@Override
public boolean add(InternalEvent ievent) throws IllegalStateException, IOException {
    if (dataRecords.size() >= MAX_RECORDS) {
        logger.trace("hit record index max");
        throw new IllegalStateException("reached max payload size");
    }
    byte[] record = this.serializer.serialize(ievent);
    /*
     * Restrict size of individual record
     */
    if (record.length > MAX_RECORD_SIZE) {
        throw new IOException("serialized event is " + record.length + " larger than max of " + MAX_RECORD_SIZE);
    }
    ByteBuffer data = ByteBuffer.wrap(record);
    dataRecords.add(new Record().withData(data));
    return true;
}