Search in sources :

Example 11 with Job

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

the class BigQueryServicesImplTest method testStartLoadJobSucceeds.

/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds.
   */
@Test
public void testStartLoadJobSucceeds() throws IOException, InterruptedException {
    Job testJob = new Job();
    JobReference jobRef = new JobReference();
    jobRef.setJobId("jobId");
    jobRef.setProjectId("projectId");
    testJob.setJobReference(jobRef);
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(testJob));
    Sleeper sleeper = new FastNanoClockAndSleeper();
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, BackOffAdapter.toGcpBackOff(FluentBackoff.DEFAULT.backoff()));
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
    expectedLogs.verifyInfo(String.format("Started BigQuery job: %s", jobRef));
}
Also used : JobReference(com.google.api.services.bigquery.model.JobReference) FastNanoClockAndSleeper(org.apache.beam.sdk.util.FastNanoClockAndSleeper) MockSleeper(com.google.api.client.testing.util.MockSleeper) FastNanoClockAndSleeper(org.apache.beam.sdk.util.FastNanoClockAndSleeper) Sleeper(com.google.api.client.util.Sleeper) Job(com.google.api.services.bigquery.model.Job) ApiErrorExtractor(com.google.cloud.hadoop.util.ApiErrorExtractor) Test(org.junit.Test)

Example 12 with Job

use of com.google.api.services.bigquery.model.Job 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)));
}
Also used : Status(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.Status) JobReference(com.google.api.services.bigquery.model.JobReference) JobConfigurationTableCopy(com.google.api.services.bigquery.model.JobConfigurationTableCopy) Job(com.google.api.services.bigquery.model.Job)

Example 13 with Job

use of com.google.api.services.bigquery.model.Job 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);
    }
}
Also used : JobStatistics(com.google.api.services.bigquery.model.JobStatistics) Table(com.google.api.services.bigquery.model.Table) JobReference(com.google.api.services.bigquery.model.JobReference) ErrorProto(com.google.api.services.bigquery.model.ErrorProto) IOException(java.io.IOException) JobStatus(com.google.api.services.bigquery.model.JobStatus) TableReference(com.google.api.services.bigquery.model.TableReference) Random(java.util.Random) Job(com.google.api.services.bigquery.model.Job) JobConfiguration(com.google.api.services.bigquery.model.JobConfiguration) Nullable(javax.annotation.Nullable)

Example 14 with Job

use of com.google.api.services.bigquery.model.Job 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);
}
Also used : JobReference(com.google.api.services.bigquery.model.JobReference) IOException(java.io.IOException) JobConfigurationExtract(com.google.api.services.bigquery.model.JobConfigurationExtract) Job(com.google.api.services.bigquery.model.Job)

Example 15 with Job

use of com.google.api.services.bigquery.model.Job in project google-cloud-java by GoogleCloudPlatform.

the class HttpBigQueryRpc method write.

@Override
public Job write(String uploadId, byte[] toWrite, int toWriteOffset, long destOffset, int length, boolean last) {
    try {
        GenericUrl url = new GenericUrl(uploadId);
        HttpRequest httpRequest = bigquery.getRequestFactory().buildPutRequest(url, new ByteArrayContent(null, toWrite, toWriteOffset, length));
        httpRequest.setParser(bigquery.getObjectParser());
        long limit = destOffset + length;
        StringBuilder range = new StringBuilder("bytes ");
        range.append(destOffset).append('-').append(limit - 1).append('/');
        if (last) {
            range.append(limit);
        } else {
            range.append('*');
        }
        httpRequest.getHeaders().setContentRange(range.toString());
        int code;
        String message;
        IOException exception = null;
        HttpResponse response = null;
        try {
            response = httpRequest.execute();
            code = response.getStatusCode();
            message = response.getStatusMessage();
        } catch (HttpResponseException ex) {
            exception = ex;
            code = ex.getStatusCode();
            message = ex.getStatusMessage();
        }
        if (!last && code != HTTP_RESUME_INCOMPLETE || last && !(code == HTTP_OK || code == HTTP_CREATED)) {
            if (exception != null) {
                throw exception;
            }
            throw new BigQueryException(code, message);
        }
        return last && response != null ? response.parseAs(Job.class) : null;
    } catch (IOException ex) {
        throw translate(ex);
    }
}
Also used : HttpRequest(com.google.api.client.http.HttpRequest) HttpResponse(com.google.api.client.http.HttpResponse) HttpResponseException(com.google.api.client.http.HttpResponseException) GenericUrl(com.google.api.client.http.GenericUrl) IOException(java.io.IOException) BigQueryException(com.google.cloud.bigquery.BigQueryException) ByteArrayContent(com.google.api.client.http.ByteArrayContent) Job(com.google.api.services.bigquery.model.Job)

Aggregations

Job (com.google.api.services.bigquery.model.Job)29 JobStatus (com.google.api.services.bigquery.model.JobStatus)16 JobReference (com.google.api.services.bigquery.model.JobReference)15 Test (org.junit.Test)14 IOException (java.io.IOException)8 JobConfiguration (com.google.api.services.bigquery.model.JobConfiguration)7 JobStatistics (com.google.api.services.bigquery.model.JobStatistics)6 TableReference (com.google.api.services.bigquery.model.TableReference)6 TableRow (com.google.api.services.bigquery.model.TableRow)5 JobServiceImpl (org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.JobServiceImpl)5 Sleeper (com.google.api.client.util.Sleeper)4 JobConfigurationQuery (com.google.api.services.bigquery.model.JobConfigurationQuery)4 JobStatistics2 (com.google.api.services.bigquery.model.JobStatistics2)4 Table (com.google.api.services.bigquery.model.Table)4 MockSleeper (com.google.api.client.testing.util.MockSleeper)3 JobStatistics4 (com.google.api.services.bigquery.model.JobStatistics4)3 TableFieldSchema (com.google.api.services.bigquery.model.TableFieldSchema)3 TableSchema (com.google.api.services.bigquery.model.TableSchema)3 ApiErrorExtractor (com.google.cloud.hadoop.util.ApiErrorExtractor)3 HashBasedTable (com.google.common.collect.HashBasedTable)3