Search in sources :

Example 6 with ByteArrayBasedKafkaRecord

use of org.apache.gobblin.kafka.client.ByteArrayBasedKafkaRecord in project incubator-gobblin by apache.

the class KafkaDeserializerExtractorTest method testConfluentAvroDeserializer.

@Test
public void testConfluentAvroDeserializer() throws IOException, RestClientException {
    WorkUnitState mockWorkUnitState = getMockWorkUnitState();
    mockWorkUnitState.setProp("schema.registry.url", TEST_URL);
    Schema schema = SchemaBuilder.record(TEST_RECORD_NAME).namespace(TEST_NAMESPACE).fields().name(TEST_FIELD_NAME).type().stringType().noDefault().endRecord();
    GenericRecord testGenericRecord = new GenericRecordBuilder(schema).set(TEST_FIELD_NAME, "testValue").build();
    SchemaRegistryClient mockSchemaRegistryClient = mock(SchemaRegistryClient.class);
    when(mockSchemaRegistryClient.getByID(any(Integer.class))).thenReturn(schema);
    Serializer<Object> kafkaEncoder = new KafkaAvroSerializer(mockSchemaRegistryClient);
    Deserializer<Object> kafkaDecoder = new KafkaAvroDeserializer(mockSchemaRegistryClient);
    ByteBuffer testGenericRecordByteBuffer = ByteBuffer.wrap(kafkaEncoder.serialize(TEST_TOPIC_NAME, testGenericRecord));
    KafkaSchemaRegistry<Integer, Schema> mockKafkaSchemaRegistry = mock(KafkaSchemaRegistry.class);
    KafkaDeserializerExtractor kafkaDecoderExtractor = new KafkaDeserializerExtractor(mockWorkUnitState, Optional.fromNullable(Deserializers.CONFLUENT_AVRO), kafkaDecoder, mockKafkaSchemaRegistry);
    ByteArrayBasedKafkaRecord mockMessageAndOffset = getMockMessageAndOffset(testGenericRecordByteBuffer);
    Assert.assertEquals(kafkaDecoderExtractor.decodeRecord(mockMessageAndOffset), testGenericRecord);
}
Also used : WorkUnitState(org.apache.gobblin.configuration.WorkUnitState) Schema(org.apache.avro.Schema) KafkaAvroSerializer(io.confluent.kafka.serializers.KafkaAvroSerializer) KafkaAvroDeserializer(io.confluent.kafka.serializers.KafkaAvroDeserializer) ByteBuffer(java.nio.ByteBuffer) ByteArrayBasedKafkaRecord(org.apache.gobblin.kafka.client.ByteArrayBasedKafkaRecord) GenericRecordBuilder(org.apache.avro.generic.GenericRecordBuilder) GenericRecord(org.apache.avro.generic.GenericRecord) SchemaRegistryClient(io.confluent.kafka.schemaregistry.client.SchemaRegistryClient) Test(org.testng.annotations.Test)

Example 7 with ByteArrayBasedKafkaRecord

use of org.apache.gobblin.kafka.client.ByteArrayBasedKafkaRecord in project incubator-gobblin by apache.

the class SimpleKafkaSpecConsumer method changedSpecs.

@Override
public Future<? extends List<Pair<SpecExecutor.Verb, Spec>>> changedSpecs() {
    List<Pair<SpecExecutor.Verb, Spec>> changesSpecs = new ArrayList<>();
    initializeWatermarks();
    this.currentPartitionIdx = -1;
    while (!allPartitionsFinished()) {
        if (currentPartitionFinished()) {
            moveToNextPartition();
            continue;
        }
        if (this.messageIterator == null || !this.messageIterator.hasNext()) {
            try {
                this.messageIterator = fetchNextMessageBuffer();
            } catch (Exception e) {
                log.error(String.format("Failed to fetch next message buffer for partition %s. Will skip this partition.", getCurrentPartition()), e);
                moveToNextPartition();
                continue;
            }
            if (this.messageIterator == null || !this.messageIterator.hasNext()) {
                moveToNextPartition();
                continue;
            }
        }
        while (!currentPartitionFinished()) {
            if (!this.messageIterator.hasNext()) {
                break;
            }
            KafkaConsumerRecord nextValidMessage = this.messageIterator.next();
            // until we get to x.
            if (nextValidMessage.getOffset() < _nextWatermark.get(this.currentPartitionIdx)) {
                continue;
            }
            _nextWatermark.set(this.currentPartitionIdx, nextValidMessage.getNextOffset());
            try {
                final AvroJobSpec record;
                if (nextValidMessage instanceof ByteArrayBasedKafkaRecord) {
                    record = decodeRecord((ByteArrayBasedKafkaRecord) nextValidMessage);
                } else if (nextValidMessage instanceof DecodeableKafkaRecord) {
                    record = ((DecodeableKafkaRecord<?, AvroJobSpec>) nextValidMessage).getValue();
                } else {
                    throw new IllegalStateException("Unsupported KafkaConsumerRecord type. The returned record can either be ByteArrayBasedKafkaRecord" + " or DecodeableKafkaRecord");
                }
                JobSpec.Builder jobSpecBuilder = JobSpec.builder(record.getUri());
                Properties props = new Properties();
                props.putAll(record.getProperties());
                jobSpecBuilder.withJobCatalogURI(record.getUri()).withVersion(record.getVersion()).withDescription(record.getDescription()).withConfigAsProperties(props);
                if (!record.getTemplateUri().isEmpty()) {
                    jobSpecBuilder.withTemplate(new URI(record.getTemplateUri()));
                }
                String verbName = record.getMetadata().get(VERB_KEY);
                SpecExecutor.Verb verb = SpecExecutor.Verb.valueOf(verbName);
                changesSpecs.add(new ImmutablePair<SpecExecutor.Verb, Spec>(verb, jobSpecBuilder.build()));
            } catch (Throwable t) {
                log.error("Could not decode record at partition " + this.currentPartitionIdx + " offset " + nextValidMessage.getOffset());
            }
        }
    }
    return new CompletedFuture(changesSpecs, null);
}
Also used : DecodeableKafkaRecord(org.apache.gobblin.kafka.client.DecodeableKafkaRecord) ArrayList(java.util.ArrayList) KafkaConsumerRecord(org.apache.gobblin.kafka.client.KafkaConsumerRecord) AvroJobSpec(org.apache.gobblin.runtime.job_spec.AvroJobSpec) Properties(java.util.Properties) ByteArrayBasedKafkaRecord(org.apache.gobblin.kafka.client.ByteArrayBasedKafkaRecord) URI(java.net.URI) KafkaOffsetRetrievalFailureException(org.apache.gobblin.source.extractor.extract.kafka.KafkaOffsetRetrievalFailureException) IOException(java.io.IOException) InvocationTargetException(java.lang.reflect.InvocationTargetException) SpecExecutor(org.apache.gobblin.runtime.api.SpecExecutor) JobSpec(org.apache.gobblin.runtime.api.JobSpec) AvroJobSpec(org.apache.gobblin.runtime.job_spec.AvroJobSpec) JobSpec(org.apache.gobblin.runtime.api.JobSpec) Spec(org.apache.gobblin.runtime.api.Spec) AvroJobSpec(org.apache.gobblin.runtime.job_spec.AvroJobSpec) Pair(org.apache.commons.lang3.tuple.Pair) ImmutablePair(org.apache.commons.lang3.tuple.ImmutablePair) CompletedFuture(org.apache.gobblin.util.CompletedFuture)

Aggregations

ByteArrayBasedKafkaRecord (org.apache.gobblin.kafka.client.ByteArrayBasedKafkaRecord)7 ByteBuffer (java.nio.ByteBuffer)4 WorkUnitState (org.apache.gobblin.configuration.WorkUnitState)4 Test (org.testng.annotations.Test)4 SchemaRegistryClient (io.confluent.kafka.schemaregistry.client.SchemaRegistryClient)2 KafkaAvroDeserializer (io.confluent.kafka.serializers.KafkaAvroDeserializer)2 KafkaAvroSerializer (io.confluent.kafka.serializers.KafkaAvroSerializer)2 IOException (java.io.IOException)2 Schema (org.apache.avro.Schema)2 GenericRecord (org.apache.avro.generic.GenericRecord)2 GenericRecordBuilder (org.apache.avro.generic.GenericRecordBuilder)2 DecodeableKafkaRecord (org.apache.gobblin.kafka.client.DecodeableKafkaRecord)2 KafkaConsumerRecord (org.apache.gobblin.kafka.client.KafkaConsumerRecord)2 KafkaJsonDeserializer (io.confluent.kafka.serializers.KafkaJsonDeserializer)1 KafkaJsonSerializer (io.confluent.kafka.serializers.KafkaJsonSerializer)1 InvocationTargetException (java.lang.reflect.InvocationTargetException)1 URI (java.net.URI)1 ArrayList (java.util.ArrayList)1 Properties (java.util.Properties)1 ImmutablePair (org.apache.commons.lang3.tuple.ImmutablePair)1