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

Example 1 with LeaseWorkItemRequest

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

the class DataflowWorkUnitClientTest method testCloudServiceCallNoWorkId.

@Test
public void testCloudServiceCallNoWorkId() throws Exception {
    // If there's no work the service should return an empty work item.
    WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
    workItem.setId(null);
    when(request.execute()).thenReturn(generateMockResponse(workItem));
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    assertEquals(Optional.absent(), client.getWorkItem());
    LeaseWorkItemRequest actualRequest = Transport.getJsonFactory().fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
    assertEquals(WORKER_ID, actualRequest.getWorkerId());
    assertEquals(ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities());
    assertEquals(ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes());
}

Example 2 with LeaseWorkItemRequest

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

the class DataflowWorkUnitClient method getWorkItemInternal.

private Optional<WorkItem> getWorkItemInternal(List<String> workItemTypes, List<String> capabilities) throws IOException {
    LeaseWorkItemRequest request = new LeaseWorkItemRequest();
    request.setFactory(Transport.getJsonFactory());
    request.setWorkItemTypes(workItemTypes);
    request.setWorkerCapabilities(capabilities);
    request.setWorkerId(options.getWorkerId());
    request.setCurrentWorkerTime(toCloudTime(DateTime.now()));
    // This shouldn't be necessary, but a valid cloud duration string is
    // required by the Google API parsing framework.  TODO: Fix the framework
    // so that an empty or not-present string can be used as a default value.
    request.setRequestedLeaseDuration(toCloudDuration(Duration.millis(WorkProgressUpdater.DEFAULT_LEASE_DURATION_MILLIS)));
    logger.debug("Leasing work: {}", request);
    LeaseWorkItemResponse response = dataflow.projects().locations().jobs().workItems().lease(options.getProject(), options.getRegion(), options.getJobId(), request).execute();
    logger.debug("Lease work response: {}", response);
    List<WorkItem> workItems = response.getWorkItems();
    if (workItems == null || workItems.isEmpty()) {
        // We didn't lease any work.
        return Optional.absent();
    } else if (workItems.size() > 1) {
        throw new IOException("This version of the SDK expects no more than one work item from the service: " + response);
    }
    WorkItem work = response.getWorkItems().get(0);
    // Looks like the work's a'ight.
    return Optional.of(work);
}

Example 3 with LeaseWorkItemRequest

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

the class DataflowWorkUnitClientTest method testCloudServiceCallNoWorkPresent.

@Test
public void testCloudServiceCallNoWorkPresent() throws Exception {
    // If there's no work the service should return an empty work item.
    WorkItem workItem = new WorkItem();
    when(request.execute()).thenReturn(generateMockResponse(workItem));
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    assertEquals(Optional.absent(), client.getWorkItem());
    LeaseWorkItemRequest actualRequest = Transport.getJsonFactory().fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
    assertEquals(WORKER_ID, actualRequest.getWorkerId());
    assertEquals(ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities());
    assertEquals(ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes());
}

Example 4 with LeaseWorkItemRequest

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

the class DataflowWorkUnitClientTest method testCloudServiceCall.

@Test
public void testCloudServiceCall() throws Exception {
    WorkItem workItem = createWorkItem(PROJECT_ID, JOB_ID);
    when(request.execute()).thenReturn(generateMockResponse(workItem));
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    assertEquals(Optional.of(workItem), client.getWorkItem());
    LeaseWorkItemRequest actualRequest = Transport.getJsonFactory().fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
    assertEquals(WORKER_ID, actualRequest.getWorkerId());
    assertEquals(ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities());
    assertEquals(ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes());
    assertEquals("1234", DataflowWorkerLoggingMDC.getWorkId());
}

Example 5 with LeaseWorkItemRequest

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

the class DataflowWorkUnitClientTest method testCloudServiceCallNoWorkItem.

@Test
public void testCloudServiceCallNoWorkItem() throws Exception {
    when(request.execute()).thenReturn(generateMockResponse());
    WorkUnitClient client = new DataflowWorkUnitClient(pipelineOptions, LOG);
    assertEquals(Optional.absent(), client.getWorkItem());
    LeaseWorkItemRequest actualRequest = Transport.getJsonFactory().fromString(request.getContentAsString(), LeaseWorkItemRequest.class);
    assertEquals(WORKER_ID, actualRequest.getWorkerId());
    assertEquals(ImmutableList.<String>of(WORKER_ID, "remote_source", "custom_source"), actualRequest.getWorkerCapabilities());
    assertEquals(ImmutableList.<String>of("map_task", "seq_map_task", "remote_source_task"), actualRequest.getWorkItemTypes());
}