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

Example 1 with LeaseWorkItemResponse

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

the class DataflowWorkUnitClientTest method generateMockResponse.

private LowLevelHttpResponse generateMockResponse(WorkItem... workItems) throws Exception {
    MockLowLevelHttpResponse response = new MockLowLevelHttpResponse();
    response.setContentType(Json.MEDIA_TYPE);
    LeaseWorkItemResponse lease = new LeaseWorkItemResponse();
    lease.setWorkItems(Lists.newArrayList(workItems));
    // N.B. Setting the factory is necessary in order to get valid JSON.
    lease.setFactory(Transport.getJsonFactory());
    response.setContent(lease.toPrettyString());
    return response;
}

Example 2 with LeaseWorkItemResponse

use of com.google.api.services.dataflow.model.LeaseWorkItemResponse 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);
}