Examples with JobReference com.google.api.services.bigquery.model.JobReference used on opensource projects

Example 6 with JobReference

use of com.google.api.services.bigquery.model.JobReference in project beam by apache.

the class WriteRename method copy.

private void copy(JobService jobService, DatasetService datasetService, String jobIdPrefix, TableReference ref, List<TableReference> tempTables, WriteDisposition writeDisposition, CreateDisposition createDisposition, @Nullable String tableDescription) throws InterruptedException, IOException {
    JobConfigurationTableCopy copyConfig = new JobConfigurationTableCopy().setSourceTables(tempTables).setDestinationTable(ref).setWriteDisposition(writeDisposition.name()).setCreateDisposition(createDisposition.name());
    String projectId = ref.getProjectId();
    Job lastFailedCopyJob = null;
    for (int i = 0; i < BatchLoads.MAX_RETRY_JOBS; ++i) {
        String jobId = jobIdPrefix + "-" + i;
        JobReference jobRef = new JobReference().setProjectId(projectId).setJobId(jobId);
        jobService.startCopyJob(jobRef, copyConfig);
        Job copyJob = jobService.pollJob(jobRef, BatchLoads.LOAD_JOB_POLL_MAX_RETRIES);
        Status jobStatus = BigQueryHelpers.parseStatus(copyJob);
        switch(jobStatus) {
            case SUCCEEDED:
                if (tableDescription != null) {
                    datasetService.patchTableDescription(ref, tableDescription);
                }
                return;
            case UNKNOWN:
                throw new RuntimeException(String.format("UNKNOWN status of copy job [%s]: %s.", jobId, BigQueryHelpers.jobToPrettyString(copyJob)));
            case FAILED:
                lastFailedCopyJob = copyJob;
                continue;
            default:
                throw new IllegalStateException(String.format("Unexpected status [%s] of load job: %s.", jobStatus, BigQueryHelpers.jobToPrettyString(copyJob)));
        }
    }
    throw new RuntimeException(String.format("Failed to create copy job with id prefix %s, " + "reached max retries: %d, last failed copy job: %s.", jobIdPrefix, BatchLoads.MAX_RETRY_JOBS, BigQueryHelpers.jobToPrettyString(lastFailedCopyJob)));
}

Example 7 with JobReference

use of com.google.api.services.bigquery.model.JobReference in project beam by apache.

the class BigQueryTableRowIterator method executeQueryAndWaitForCompletion.

/**
   * Executes the specified query and returns a reference to the temporary BigQuery table created
   * to hold the results.
   *
   * @throws IOException if the query fails.
   */
private TableReference executeQueryAndWaitForCompletion() throws IOException, InterruptedException {
    checkState(projectId != null, "Unable to execute a query without a configured project id");
    checkState(queryConfig != null, "Unable to execute a query without a configured query");
    // Dry run query to get source table location
    Job dryRunJob = new Job().setConfiguration(new JobConfiguration().setQuery(queryConfig).setDryRun(true));
    JobStatistics jobStats = executeWithBackOff(client.jobs().insert(projectId, dryRunJob), String.format("Error when trying to dry run query %s.", queryConfig.toPrettyString())).getStatistics();
    // Let BigQuery to pick default location if the query does not read any tables.
    String location = null;
    @Nullable List<TableReference> tables = jobStats.getQuery().getReferencedTables();
    if (tables != null && !tables.isEmpty()) {
        Table table = getTable(tables.get(0));
        location = table.getLocation();
    }
    // Create a temporary dataset to store results.
    // Starting dataset name with an "_" so that it is hidden.
    Random rnd = new Random(System.currentTimeMillis());
    temporaryDatasetId = "_beam_temporary_dataset_" + rnd.nextInt(1000000);
    temporaryTableId = "beam_temporary_table_" + rnd.nextInt(1000000);
    createDataset(temporaryDatasetId, location);
    Job job = new Job();
    JobConfiguration config = new JobConfiguration();
    config.setQuery(queryConfig);
    job.setConfiguration(config);
    TableReference destinationTable = new TableReference();
    destinationTable.setProjectId(projectId);
    destinationTable.setDatasetId(temporaryDatasetId);
    destinationTable.setTableId(temporaryTableId);
    queryConfig.setDestinationTable(destinationTable);
    queryConfig.setAllowLargeResults(true);
    Job queryJob = executeWithBackOff(client.jobs().insert(projectId, job), String.format("Error when trying to execute the job for query %s.", queryConfig.toPrettyString()));
    JobReference jobId = queryJob.getJobReference();
    while (true) {
        Job pollJob = executeWithBackOff(client.jobs().get(projectId, jobId.getJobId()), String.format("Error when trying to get status of the job for query %s.", queryConfig.toPrettyString()));
        JobStatus status = pollJob.getStatus();
        if (status.getState().equals("DONE")) {
            // Job is DONE, but did not necessarily succeed.
            ErrorProto error = status.getErrorResult();
            if (error == null) {
                return pollJob.getConfiguration().getQuery().getDestinationTable();
            } else {
                // There will be no temporary table to delete, so null out the reference.
                temporaryTableId = null;
                throw new IOException(String.format("Executing query %s failed: %s", queryConfig.toPrettyString(), error.getMessage()));
            }
        }
        Uninterruptibles.sleepUninterruptibly(QUERY_COMPLETION_POLL_TIME.getMillis(), TimeUnit.MILLISECONDS);
    }
}

Example 8 with JobReference

use of com.google.api.services.bigquery.model.JobReference in project beam by apache.

the class BigQuerySourceBase method executeExtract.

private List<ResourceId> executeExtract(String jobId, TableReference table, JobService jobService, String executingProject, String extractDestinationDir) throws InterruptedException, IOException {
    JobReference jobRef = new JobReference().setProjectId(executingProject).setJobId(jobId);
    String destinationUri = BigQueryIO.getExtractDestinationUri(extractDestinationDir);
    JobConfigurationExtract extract = new JobConfigurationExtract().setSourceTable(table).setDestinationFormat("AVRO").setDestinationUris(ImmutableList.of(destinationUri));
    LOG.info("Starting BigQuery extract job: {}", jobId);
    jobService.startExtractJob(jobRef, extract);
    Job extractJob = jobService.pollJob(jobRef, JOB_POLL_MAX_RETRIES);
    if (BigQueryHelpers.parseStatus(extractJob) != Status.SUCCEEDED) {
        throw new IOException(String.format("Extract job %s failed, status: %s.", extractJob.getJobReference().getJobId(), BigQueryHelpers.statusToPrettyString(extractJob.getStatus())));
    }
    LOG.info("BigQuery extract job completed: {}", jobId);
    return BigQueryIO.getExtractFilePaths(extractDestinationDir, extractJob);
}

Example 9 with JobReference

use of com.google.api.services.bigquery.model.JobReference in project beam by apache.

the class BigQueryServicesImplTest method testGetJobThrows.

@Test
public void testGetJobThrows() throws Exception {
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(401);
    BigQueryServicesImpl.JobServiceImpl jobService = new BigQueryServicesImpl.JobServiceImpl(bigquery);
    JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
    thrown.expect(IOException.class);
    thrown.expectMessage(String.format("Unable to find BigQuery job: %s", jobRef));
    jobService.getJob(jobRef, Sleeper.DEFAULT, BackOff.STOP_BACKOFF);
}

Example 10 with JobReference

use of com.google.api.services.bigquery.model.JobReference in project beam by apache.

the class BigQueryServicesImplTest method testPollJobFailed.

/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} fails.
   */
@Test
public void testPollJobFailed() throws IOException, InterruptedException {
    Job testJob = new Job();
    testJob.setStatus(new JobStatus().setState("DONE").setErrorResult(new ErrorProto()));
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(testJob));
    BigQueryServicesImpl.JobServiceImpl jobService = new BigQueryServicesImpl.JobServiceImpl(bigquery);
    JobReference jobRef = new JobReference().setProjectId("projectId").setJobId("jobId");
    Job job = jobService.pollJob(jobRef, Sleeper.DEFAULT, BackOff.ZERO_BACKOFF);
    assertEquals(testJob, job);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}