Examples with RecordEnvelope org.apache.gobblin.stream.RecordEnvelope used on opensource projects

Example 16 with RecordEnvelope

use of org.apache.gobblin.stream.RecordEnvelope in project incubator-gobblin by apache.

the class PartitionedWriterTest method test.

@Test
public void test() throws IOException {
    State state = new State();
    state.setProp(ConfigurationKeys.WRITER_PARTITIONER_CLASS, TestPartitioner.class.getCanonicalName());
    TestPartitionAwareWriterBuilder builder = new TestPartitionAwareWriterBuilder();
    PartitionedDataWriter writer = new PartitionedDataWriter<String, String>(builder, state);
    Assert.assertEquals(builder.actions.size(), 0);
    String record1 = "abc";
    writer.writeEnvelope(new RecordEnvelope(record1));
    Assert.assertEquals(builder.actions.size(), 2);
    TestPartitionAwareWriterBuilder.Action action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "a");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.BUILD);
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "a");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.WRITE);
    Assert.assertEquals(action.getTarget(), record1);
    Assert.assertTrue(writer.isSpeculativeAttemptSafe());
    String record2 = "123";
    writer.writeEnvelope(new RecordEnvelope(record2));
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "1");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.BUILD);
    Assert.assertFalse(writer.isSpeculativeAttemptSafe());
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "1");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.WRITE);
    Assert.assertEquals(action.getTarget(), record2);
    writer.writeEnvelope(new RecordEnvelope(record1));
    Assert.assertEquals(builder.actions.size(), 1);
    action = builder.actions.poll();
    Assert.assertEquals(action.getPartition(), "a");
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.WRITE);
    Assert.assertEquals(action.getTarget(), record1);
    Assert.assertEquals(writer.recordsWritten(), 3);
    Assert.assertEquals(writer.bytesWritten(), 3);
    Assert.assertFalse(writer.isSpeculativeAttemptSafe());
    writer.cleanup();
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLEANUP);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLEANUP);
    writer.close();
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLOSE);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.CLOSE);
    writer.commit();
    Assert.assertEquals(builder.actions.size(), 2);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.COMMIT);
    action = builder.actions.poll();
    Assert.assertEquals(action.getType(), TestPartitionAwareWriterBuilder.Actions.COMMIT);
}

Example 17 with RecordEnvelope

use of org.apache.gobblin.stream.RecordEnvelope in project incubator-gobblin by apache.

the class PartitionedWriterTest method testControlMessageHandler.

@Test
public void testControlMessageHandler() throws IOException {
    State state = new State();
    state.setProp(ConfigurationKeys.WRITER_PARTITIONER_CLASS, TestPartitioner.class.getCanonicalName());
    TestPartitionAwareWriterBuilder builder = new TestPartitionAwareWriterBuilder();
    PartitionedDataWriter writer = new PartitionedDataWriter<String, String>(builder, state);
    Assert.assertEquals(builder.actions.size(), 0);
    String record1 = "abc";
    writer.writeEnvelope(new RecordEnvelope(record1));
    String record2 = "123";
    writer.writeEnvelope(new RecordEnvelope(record2));
    FlushControlMessage controlMessage = FlushControlMessage.builder().build();
    BasicAckableForTesting ackable = new BasicAckableForTesting();
    controlMessage.addCallBack(ackable);
    Assert.assertEquals(ackable.acked, 0);
    // when the control message is cloned properly then this does not raise an error
    writer.getMessageHandler().handleMessage(controlMessage);
    // message handler does not ack since consumeRecordStream does acking for control messages
    // this should be revisited when control message error handling is changed
    controlMessage.ack();
    Assert.assertEquals(ackable.acked, 1);
    writer.close();
}

Example 18 with RecordEnvelope

use of org.apache.gobblin.stream.RecordEnvelope in project incubator-gobblin by apache.

the class RetryWriterTest method retryTest.

public void retryTest() throws IOException {
    DataWriter<Void> writer = mock(DataWriter.class);
    doThrow(new RuntimeException()).when(writer).writeEnvelope(any(RecordEnvelope.class));
    DataWriterWrapperBuilder<Void> builder = new DataWriterWrapperBuilder<>(writer, new State());
    DataWriter<Void> retryWriter = builder.build();
    try {
        retryWriter.writeEnvelope(new RecordEnvelope<>(null));
        Assert.fail("Should have failed.");
    } catch (Exception e) {
    }
    verify(writer, times(5)).writeEnvelope(any(RecordEnvelope.class));
}

Example 19 with RecordEnvelope

use of org.apache.gobblin.stream.RecordEnvelope in project incubator-gobblin by apache.

the class RetryWriterTest method retryTestSuccess.

public void retryTestSuccess() throws IOException {
    DataWriter<Void> writer = mock(DataWriter.class);
    DataWriterWrapperBuilder<Void> builder = new DataWriterWrapperBuilder<>(writer, new State());
    DataWriter<Void> retryWriter = builder.build();
    retryWriter.writeEnvelope(new RecordEnvelope<>(null));
    verify(writer, times(1)).writeEnvelope(any(RecordEnvelope.class));
}

Example 20 with RecordEnvelope

use of org.apache.gobblin.stream.RecordEnvelope in project incubator-gobblin by apache.

the class AsyncConverter1to1Test method testFailedConversion.

@Test
public void testFailedConversion() throws Exception {
    MyAsyncConverter1to1 converter = new MyAsyncConverter1to1();
    List<Throwable> errors = Lists.newArrayList();
    AtomicBoolean done = new AtomicBoolean(false);
    WorkUnitState workUnitState = new WorkUnitState();
    workUnitState.setProp(AsyncConverter1to1.MAX_CONCURRENT_ASYNC_CONVERSIONS_KEY, 3);
    RecordStreamWithMetadata<String, String> stream = new RecordStreamWithMetadata<>(Flowable.just("0", MyAsyncConverter1to1.FAIL, "1").map(RecordEnvelope::new), GlobalMetadata.<String>builder().schema("schema").build());
    Set<String> outputRecords = Sets.newConcurrentHashSet();
    converter.processStream(stream, workUnitState).getRecordStream().subscribeOn(Schedulers.newThread()).subscribe(r -> outputRecords.add(((RecordEnvelope<String>) r).getRecord()), errors::add, () -> done.set(true));
    Assert.assertTrue(ExponentialBackoff.awaitCondition().maxWait(100L).callable(() -> errors.size() > 0).await());
    Assert.assertEquals(errors.size(), 1);
    Assert.assertEquals(errors.get(0).getCause().getMessage(), "injected failure");
}