Examples with ProducerRecord org.apache.kafka.clients.producer.ProducerRecord used on opensource projects

Example 91 with ProducerRecord

use of org.apache.kafka.clients.producer.ProducerRecord in project beam by apache.

the class KafkaIOTest method testKafkaWriteHeaders.

@Test
public void testKafkaWriteHeaders() throws Exception {
    // Set different output topic names
    int numElements = 1;
    SimpleEntry<String, String> header = new SimpleEntry<>("header_key", "header_value");
    try (MockProducerWrapper producerWrapper = new MockProducerWrapper()) {
        ProducerSendCompletionThread completionThread = new ProducerSendCompletionThread(producerWrapper.mockProducer).start();
        String defaultTopic = "test";
        p.apply(mkKafkaReadTransform(numElements, new ValueAsTimestampFn()).withoutMetadata()).apply(ParDo.of(new KV2ProducerRecord(defaultTopic, true, System.currentTimeMillis(), header))).setCoder(ProducerRecordCoder.of(VarIntCoder.of(), VarLongCoder.of())).apply(KafkaIO.<Integer, Long>writeRecords().withBootstrapServers("none").withKeySerializer(IntegerSerializer.class).withValueSerializer(LongSerializer.class).withInputTimestamp().withProducerFactoryFn(new ProducerFactoryFn(producerWrapper.producerKey)));
        p.run();
        completionThread.shutdown();
        // Verify that appropriate header is written with producer record
        List<ProducerRecord<Integer, Long>> sent = producerWrapper.mockProducer.history();
        for (int i = 0; i < numElements; i++) {
            ProducerRecord<Integer, Long> record = sent.get(i);
            Headers headers = record.headers();
            assertNotNull(headers);
            Header[] headersArray = headers.toArray();
            assertEquals(1, headersArray.length);
            assertEquals(header.getKey(), headersArray[0].key());
            assertEquals(header.getValue(), new String(headersArray[0].value(), StandardCharsets.UTF_8));
        }
    }
}

Example 92 with ProducerRecord

use of org.apache.kafka.clients.producer.ProducerRecord in project beam by apache.

the class ProducerRecordCoder method decode.

@Override
public ProducerRecord<K, V> decode(InputStream inStream) throws IOException {
    String topic = stringCoder.decode(inStream);
    Integer partition = intCoder.decode(inStream);
    if (partition == -1) {
        partition = null;
    }
    Long timestamp = longCoder.decode(inStream);
    if (timestamp == Long.MAX_VALUE) {
        timestamp = null;
    }
    Headers headers = (Headers) toHeaders(headerCoder.decode(inStream));
    KV<K, V> kv = kvCoder.decode(inStream);
    if (ConsumerSpEL.hasHeaders()) {
        return new ProducerRecord<>(topic, partition, timestamp, kv.getKey(), kv.getValue(), headers);
    }
    return new ProducerRecord<>(topic, partition, timestamp, kv.getKey(), kv.getValue());
}

Example 93 with ProducerRecord

use of org.apache.kafka.clients.producer.ProducerRecord in project beam by apache.

the class ProducerRecordCoderTest method testProducerRecordStructuralValueWithoutHeadersApi.

@Test
public void testProducerRecordStructuralValueWithoutHeadersApi() throws IOException {
    RecordHeaders headers = new RecordHeaders();
    headers.add("headerKey", "headerVal".getBytes(UTF_8));
    ProducerRecordCoder producerRecordCoder = ProducerRecordCoder.of(ByteArrayCoder.of(), ByteArrayCoder.of());
    ProducerRecord<byte[], byte[]> producerRecord = new ProducerRecord<>("topic", 1, null, "key".getBytes(UTF_8), "value".getBytes(UTF_8), headers);
    mockStatic(ConsumerSpEL.class);
    when(ConsumerSpEL.hasHeaders()).thenReturn(false);
    ProducerRecord testProducerRecord = (ProducerRecord) producerRecordCoder.structuralValue(producerRecord);
    assertEquals(testProducerRecord.headers(), new RecordHeaders());
}

Example 94 with ProducerRecord

use of org.apache.kafka.clients.producer.ProducerRecord in project beam by apache.

the class ProducerRecordCoderTest method testProducerRecordStructuralValueWithHeadersApi.

@Test
public void testProducerRecordStructuralValueWithHeadersApi() throws IOException {
    RecordHeaders headers = new RecordHeaders();
    headers.add("headerKey", "headerVal".getBytes(UTF_8));
    ProducerRecordCoder producerRecordCoder = ProducerRecordCoder.of(ByteArrayCoder.of(), ByteArrayCoder.of());
    ProducerRecord<byte[], byte[]> producerRecord = new ProducerRecord<>("topic", 1, null, "key".getBytes(UTF_8), "value".getBytes(UTF_8), headers);
    ProducerRecord testProducerRecord = (ProducerRecord) producerRecordCoder.structuralValue(producerRecord);
    assertEquals(testProducerRecord.headers(), headers);
}

Example 95 with ProducerRecord

use of org.apache.kafka.clients.producer.ProducerRecord in project beam by apache.

the class NestedPayloadKafkaTable method transformOutput.

// Suppress nullability warnings: ProducerRecord is supposed to accept null arguments.
@SuppressWarnings("argument.type.incompatible")
@VisibleForTesting
ProducerRecord<byte[], byte[]> transformOutput(Row row) {
    row = castRow(row, row.getSchema(), schema);
    String topic = Iterables.getOnlyElement(getTopics());
    byte[] key = null;
    byte[] payload;
    List<Header> headers = ImmutableList.of();
    Long timestampMillis = null;
    if (schema.hasField(Schemas.MESSAGE_KEY_FIELD)) {
        key = row.getBytes(Schemas.MESSAGE_KEY_FIELD);
    }
    if (schema.hasField(Schemas.EVENT_TIMESTAMP_FIELD)) {
        ReadableDateTime time = row.getDateTime(Schemas.EVENT_TIMESTAMP_FIELD);
        if (time != null) {
            timestampMillis = time.getMillis();
        }
    }
    if (schema.hasField(Schemas.HEADERS_FIELD)) {
        Collection<Row> headerRows = checkArgumentNotNull(row.getArray(Schemas.HEADERS_FIELD));
        ImmutableList.Builder<Header> headersBuilder = ImmutableList.builder();
        headerRows.forEach(entry -> {
            String headerKey = checkArgumentNotNull(entry.getString(Schemas.HEADERS_KEY_FIELD));
            Collection<byte[]> values = checkArgumentNotNull(entry.getArray(Schemas.HEADERS_VALUES_FIELD));
            values.forEach(value -> headersBuilder.add(new RecordHeader(headerKey, value)));
        });
        headers = headersBuilder.build();
    }
    if (payloadSerializer == null) {
        payload = row.getBytes(Schemas.PAYLOAD_FIELD);
    } else {
        payload = payloadSerializer.serialize(checkArgumentNotNull(row.getRow(Schemas.PAYLOAD_FIELD)));
    }
    return new ProducerRecord<>(topic, null, timestampMillis, key, payload, headers);
}