Examples with PendingJobManager org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.PendingJobManager used on opensource projects

Example 1 with PendingJobManager

use of org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.PendingJobManager in project beam by apache.

the class WriteRename method finishBundle.

@FinishBundle
public void finishBundle(FinishBundleContext c) throws Exception {
    DatasetService datasetService = getDatasetService(c.getPipelineOptions().as(BigQueryOptions.class));
    PendingJobManager jobManager = new PendingJobManager();
    for (PendingJobData pendingJob : pendingJobs) {
        jobManager.addPendingJob(pendingJob.retryJob, j -> {
            try {
                if (pendingJob.tableDestination.getTableDescription() != null) {
                    TableReference ref = pendingJob.tableDestination.getTableReference();
                    datasetService.patchTableDescription(ref.clone().setTableId(BigQueryHelpers.stripPartitionDecorator(ref.getTableId())), pendingJob.tableDestination.getTableDescription());
                }
                c.output(pendingJob.tableDestination, pendingJob.window.maxTimestamp(), pendingJob.window);
                removeTemporaryTables(datasetService, pendingJob.tempTables);
                return null;
            } catch (IOException | InterruptedException e) {
                return e;
            }
        });
    }
    jobManager.waitForDone();
}

Example 2 with PendingJobManager

use of org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.PendingJobManager in project beam by apache.

the class BigQueryHelpersTest method testPendingJobManager.

@Test
public void testPendingJobManager() throws Exception {
    PendingJobManager jobManager = new PendingJobManager(BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.withMaxRetries(Integer.MAX_VALUE).withInitialBackoff(Duration.millis(10)).withMaxBackoff(Duration.millis(10)).backoff()));
    Set<String> succeeded = Sets.newHashSet();
    for (int i = 0; i < 5; i++) {
        Job currentJob = new Job();
        currentJob.setKind(" bigquery#job");
        PendingJob pendingJob = new PendingJob(retryId -> {
            if (new Random().nextInt(2) == 0) {
                throw new RuntimeException("Failing to start.");
            }
            currentJob.setJobReference(new JobReference().setProjectId("").setLocation("").setJobId(retryId.getJobId()));
            return null;
        }, retryId -> {
            if (retryId.getRetryIndex() < 5) {
                currentJob.setStatus(new JobStatus().setErrorResult(new ErrorProto()));
            } else {
                currentJob.setStatus(new JobStatus().setErrorResult(null));
            }
            return currentJob;
        }, retryId -> {
            if (retryId.getJobId().equals(currentJob.getJobReference().getJobId())) {
                return currentJob;
            } else {
                return null;
            }
        }, 100, "JOB_" + i);
        jobManager.addPendingJob(pendingJob, j -> {
            succeeded.add(j.currentJobId.getJobId());
            return null;
        });
    }
    jobManager.waitForDone();
    Set<String> expectedJobs = ImmutableSet.of("JOB_0-5", "JOB_1-5", "JOB_2-5", "JOB_3-5", "JOB_4-5");
    assertEquals(expectedJobs, succeeded);
}