Examples with SourceSplitRequest com.google.api.services.dataflow.model.SourceSplitRequest used on opensource projects

Example 1 with SourceSplitRequest

use of com.google.api.services.dataflow.model.SourceSplitRequest in project beam by apache.

the class SourceOperationExecutorFactoryTest method testCreateDefault.

@Test
public void testCreateDefault() throws Exception {
    SourceOperationRequest request = new SourceOperationRequest().setName("name").setOriginalName("original").setSystemName("system").setStageName("stage").setSplit(new SourceSplitRequest());
    DataflowOperationContext mockOperation = Mockito.mock(DataflowOperationContext.class);
    Mockito.when(executionContext.createOperationContext(Mockito.isA(NameContext.class))).thenReturn(mockOperation);
    SourceOperationExecutor sourceOperationExecutor = SourceOperationExecutorFactory.create(PipelineOptionsFactory.create(), request, null, executionContext, "STAGE");
    assertThat(sourceOperationExecutor, instanceOf(WorkerCustomSourceOperationExecutor.class));
}

Example 2 with SourceSplitRequest

use of com.google.api.services.dataflow.model.SourceSplitRequest in project beam by apache.

the class NoOpSourceOperationExecutorTest method setUp.

@Before
public void setUp() {
    SourceSplitRequest splitRequest = new SourceSplitRequest();
    SourceOperationRequest operationRequest = new SourceOperationRequest().setSplit(splitRequest);
    executor = new NoOpSourceOperationExecutor(operationRequest);
}

Example 3 with SourceSplitRequest

use of com.google.api.services.dataflow.model.SourceSplitRequest in project beam by apache.

the class WorkerCustomSourceOperationExecutor method execute.

@Override
public void execute() throws Exception {
    LOG.debug("Executing source operation");
    try (Closeable stateTrackerCloser = executionStateTracker.activate()) {
        try (Closeable stateCloser = executionStateTracker.enterState(executionState)) {
            SourceSplitRequest split = request.getSplit();
            if (split != null) {
                this.response = WorkerCustomSources.performSplit(split, options);
            } else {
                throw new UnsupportedOperationException("Unsupported source operation request: " + request);
            }
        }
    }
    LOG.debug("Source operation execution complete");
}

Example 4 with SourceSplitRequest

use of com.google.api.services.dataflow.model.SourceSplitRequest in project beam by apache.

the class WorkerCustomSourcesTest method performSplit.

static SourceSplitResponse performSplit(com.google.api.services.dataflow.model.Source source, PipelineOptions options, @Nullable Long desiredBundleSizeBytes, @Nullable Integer numBundlesLimitForTest, @Nullable Long apiByteLimitForTest) throws Exception {
    SourceSplitRequest splitRequest = new SourceSplitRequest();
    splitRequest.setSource(source);
    if (desiredBundleSizeBytes != null) {
        splitRequest.setOptions(new SourceSplitOptions().setDesiredBundleSizeBytes(desiredBundleSizeBytes));
    }
    SourceOperationResponse response = WorkerCustomSources.performSplitWithApiLimit(splitRequest, options, MoreObjects.firstNonNull(numBundlesLimitForTest, WorkerCustomSources.DEFAULT_NUM_BUNDLES_LIMIT), MoreObjects.firstNonNull(apiByteLimitForTest, WorkerCustomSources.DATAFLOW_SPLIT_RESPONSE_API_SIZE_LIMIT));
    return response.getSplit();
}