Examples with FinalState org.apache.gobblin.util.FinalState used on opensource projects

Example 1 with FinalState

use of org.apache.gobblin.util.FinalState in project incubator-gobblin by apache.

the class RetryWriterTest method retryGetFinalState.

public void retryGetFinalState() throws IOException {
    PartitionedDataWriter writer = mock(PartitionedDataWriter.class);
    when(writer.getFinalState()).thenReturn(new State());
    DataWriterWrapperBuilder<Void> builder = new DataWriterWrapperBuilder<>(writer, new State());
    DataWriter<Void> retryWriter = builder.build();
    State state = ((FinalState) retryWriter).getFinalState();
    verify(writer, times(1)).getFinalState();
    Assert.assertTrue(state.contains(RetryWriter.FAILED_WRITES_KEY));
}

Example 2 with FinalState

use of org.apache.gobblin.util.FinalState in project incubator-gobblin by apache.

the class ThrottleWriterTest method testGetFinalState.

public void testGetFinalState() throws IOException {
    PartitionedDataWriter writer = mock(PartitionedDataWriter.class);
    when(writer.getFinalState()).thenReturn(new State());
    int parallelism = 2;
    int qps = 4;
    DataWriter<Void> throttleWriter = setup(writer, parallelism, qps, ThrottleType.QPS);
    State state = ((FinalState) throttleWriter).getFinalState();
    verify(writer, times(1)).getFinalState();
    Assert.assertTrue(state.contains(ThrottleWriter.THROTTLED_TIME_KEY));
}

Example 3 with FinalState

use of org.apache.gobblin.util.FinalState in project incubator-gobblin by apache.

the class RetryWriter method getFinalState.

@Override
public State getFinalState() {
    State state = new State();
    if (this.writer instanceof FinalState) {
        state.addAll(((FinalState) this.writer).getFinalState());
    } else {
        LOG.warn("Wrapped writer does not implement FinalState: " + this.writer.getClass());
    }
    state.setProp(FAILED_WRITES_KEY, this.failedWrites);
    return state;
}

Example 4 with FinalState

use of org.apache.gobblin.util.FinalState in project incubator-gobblin by apache.

the class PartitionedDataWriter method getFinalState.

@Override
public State getFinalState() {
    State state = new State();
    try {
        for (Map.Entry<GenericRecord, DataWriter<D>> entry : this.partitionWriters.asMap().entrySet()) {
            if (entry.getValue() instanceof FinalState) {
                State partitionFinalState = ((FinalState) entry.getValue()).getFinalState();
                if (this.shouldPartition) {
                    for (String key : partitionFinalState.getPropertyNames()) {
                        // Prevent overwriting final state across writers
                        partitionFinalState.setProp(key + "_" + AvroUtils.serializeAsPath(entry.getKey(), false, true), partitionFinalState.getProp(key));
                    }
                }
                state.addAll(partitionFinalState);
            }
        }
        state.setProp("RecordsWritten", recordsWritten());
        state.setProp("BytesWritten", bytesWritten());
    } catch (Exception exception) {
        log.warn("Failed to get final state." + exception.getMessage());
    // If Writer fails to return bytesWritten, it might not be implemented, or implemented incorrectly.
    // Omit property instead of failing.
    }
    return state;
}

Example 5 with FinalState

use of org.apache.gobblin.util.FinalState in project incubator-gobblin by apache.

the class ThrottleWriter method getFinalState.

@Override
public State getFinalState() {
    State state = new State();
    if (this.writer instanceof FinalState) {
        state.addAll(((FinalState) this.writer).getFinalState());
    } else {
        LOG.warn("Wrapped writer does not implement FinalState: " + this.writer.getClass());
    }
    state.setProp(THROTTLED_TIME_KEY, this.throttledTime);
    return state;
}