Search in sources :

Example 1 with Job

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

the class BigQueryInterpreter method pollJob.

//Function to poll a job for completion. Future use
public static Job pollJob(final Bigquery.Jobs.Get request, final long interval) throws IOException, InterruptedException {
    Job job = request.execute();
    while (!job.getStatus().getState().equals("DONE")) {
        System.out.println("Job is " + job.getStatus().getState() + " waiting " + interval + " milliseconds...");
        Thread.sleep(interval);
        job = request.execute();
    }
    return job;
}
Also used : Job(com.google.api.services.bigquery.model.Job)

Example 2 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 3 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 4 with Job

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

the class FakeJobService method startCopyJob.

@Override
public void startCopyJob(JobReference jobRef, JobConfigurationTableCopy copyConfig) throws IOException, InterruptedException {
    synchronized (allJobs) {
        Job job = new Job();
        job.setJobReference(jobRef);
        job.setConfiguration(new JobConfiguration().setCopy(copyConfig));
        job.setKind(" bigquery#job");
        job.setStatus(new JobStatus().setState("PENDING"));
        allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
    }
}
Also used : JobStatus(com.google.api.services.bigquery.model.JobStatus) Job(com.google.api.services.bigquery.model.Job) JobConfiguration(com.google.api.services.bigquery.model.JobConfiguration)

Example 5 with Job

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

the class BigQueryTableRowIteratorTest method testQueryFailed.

/**
   * Verifies that when the query fails, the user gets a useful exception and the temporary dataset
   * is cleaned up. Also verifies that the temporary table (which is never created) is not
   * erroneously attempted to be deleted.
   */
@Test
public void testQueryFailed() throws IOException {
    // Job state polled with an error.
    String errorReason = "bad query";
    Exception exception = new IOException(errorReason);
    when(mockJobsInsert.execute()).thenThrow(exception, exception, exception, exception);
    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setQuery("NOT A QUERY");
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(queryConfig, "project", mockClient)) {
        iterator.open();
        fail();
    } catch (Exception expected) {
        // Verify message explains cause and reports the query.
        assertThat(expected.getMessage(), containsString("Error"));
        assertThat(expected.getMessage(), containsString("NOT A QUERY"));
        assertThat(expected.getCause().getMessage(), containsString(errorReason));
    }
    // Job inserted to run the query, then polled once.
    verify(mockClient, times(1)).jobs();
    verify(mockJobs).insert(anyString(), any(Job.class));
    verify(mockJobsInsert, times(4)).execute();
}
Also used : JobConfigurationQuery(com.google.api.services.bigquery.model.JobConfigurationQuery) Matchers.anyString(org.mockito.Matchers.anyString) Matchers.containsString(org.hamcrest.Matchers.containsString) IOException(java.io.IOException) Job(com.google.api.services.bigquery.model.Job) ExpectedException(org.junit.rules.ExpectedException) IOException(java.io.IOException) Test(org.junit.Test)

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