com.google.api.services.bigquery.model.JobReference

Here are the examples of the java api class com.google.api.services.bigquery.model.JobReference taken from open source projects.

1. BigQueryServicesImplTest#testStartLoadJobSucceedsAlreadyExists()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds
   * with an already exist job.
   */
@Test
public void testStartLoadJobSucceedsAlreadyExists() throws IOException, InterruptedException {
    Job testJob = new Job();
    JobReference jobRef = new JobReference();
    jobRef.setJobId("jobId");
    jobRef.setProjectId("projectId");
    testJob.setJobReference(jobRef);
    // 409 means already exists
    when(response.getStatusCode()).thenReturn(409);
    TableReference ref = new TableReference();
    ref.setProjectId("projectId");
    JobConfigurationLoad loadConfig = new JobConfigurationLoad();
    loadConfig.setDestinationTable(ref);
    Sleeper sleeper = new FastNanoClockAndSleeper();
    BackOff backoff = new AttemptBoundedExponentialBackOff(5, /* attempts */
    1000);
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, backoff);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

2. BigQueryServicesImplTest#testStartLoadJobSucceedsAlreadyExists()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds
   * with an already exist job.
   */
@Test
public void testStartLoadJobSucceedsAlreadyExists() throws IOException, InterruptedException {
    Job testJob = new Job();
    JobReference jobRef = new JobReference();
    jobRef.setJobId("jobId");
    jobRef.setProjectId("projectId");
    testJob.setJobReference(jobRef);
    // 409 means already exists
    when(response.getStatusCode()).thenReturn(409);
    TableReference ref = new TableReference();
    ref.setProjectId("projectId");
    JobConfigurationLoad loadConfig = new JobConfigurationLoad();
    loadConfig.setDestinationTable(ref);
    Sleeper sleeper = new FastNanoClockAndSleeper();
    BackOff backoff = new AttemptBoundedExponentialBackOff(5, /* attempts */
    1000);
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, backoff);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

3. BigQueryServicesImplTest#testPollJobSucceeds()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} succeeds.
   */
@Test
public void testPollJobSucceeds() throws IOException, InterruptedException {
    Job testJob = new Job();
    testJob.setStatus(new JobStatus().setState("DONE"));
    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();
}

4. BigQueryServicesImplTest#testPollJobFailed()

View license
/**
   * 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();
}

5. BigQueryServicesImplTest#testPollJobUnknown()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} returns UNKNOWN.
   */
@Test
public void testPollJobUnknown() throws IOException, InterruptedException {
    Job testJob = new Job();
    testJob.setStatus(new JobStatus());
    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.STOP_BACKOFF);
    assertEquals(null, job);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

6. BigQueryServicesImplTest#testPollJobSucceeds()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} succeeds.
   */
@Test
public void testPollJobSucceeds() throws IOException, InterruptedException {
    Job testJob = new Job();
    testJob.setStatus(new JobStatus().setState("DONE"));
    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();
}

7. BigQueryServicesImplTest#testPollJobFailed()

View license
/**
   * 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();
}

8. BigQueryServicesImplTest#testPollJobUnknown()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} returns UNKNOWN.
   */
@Test
public void testPollJobUnknown() throws IOException, InterruptedException {
    Job testJob = new Job();
    testJob.setStatus(new JobStatus());
    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.STOP_BACKOFF);
    assertEquals(null, job);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

9. BigQueryServicesImplTest#testStartLoadJobSucceeds()

View license
/**
   * 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));
    TableReference ref = new TableReference();
    ref.setProjectId("projectId");
    JobConfigurationLoad loadConfig = new JobConfigurationLoad();
    loadConfig.setDestinationTable(ref);
    Sleeper sleeper = new FastNanoClockAndSleeper();
    BackOff backoff = new AttemptBoundedExponentialBackOff(5, /* attempts */
    1000);
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, backoff);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

10. BigQueryServicesImplTest#testStartLoadJobRetry()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds with a retry.
   */
@Test
public void testStartLoadJobRetry() throws IOException, InterruptedException {
    Job testJob = new Job();
    JobReference jobRef = new JobReference();
    jobRef.setJobId("jobId");
    jobRef.setProjectId("projectId");
    testJob.setJobReference(jobRef);
    // First response is 403 rate limited, second response has valid payload.
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403))).thenReturn(toStream(testJob));
    TableReference ref = new TableReference();
    ref.setProjectId("projectId");
    JobConfigurationLoad loadConfig = new JobConfigurationLoad();
    loadConfig.setDestinationTable(ref);
    Sleeper sleeper = new FastNanoClockAndSleeper();
    BackOff backoff = new AttemptBoundedExponentialBackOff(5, /* attempts */
    1000);
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, backoff);
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
}

11. BigQueryServicesImplTest#testStartLoadJobSucceeds()

View license
/**
   * 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));
    TableReference ref = new TableReference();
    ref.setProjectId("projectId");
    JobConfigurationLoad loadConfig = new JobConfigurationLoad();
    loadConfig.setDestinationTable(ref);
    Sleeper sleeper = new FastNanoClockAndSleeper();
    BackOff backoff = new AttemptBoundedExponentialBackOff(5, /* attempts */
    1000);
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, backoff);
    verify(response, times(1)).getStatusCode();
    verify(response, times(1)).getContent();
    verify(response, times(1)).getContentType();
}

12. BigQueryServicesImplTest#testStartLoadJobRetry()

View license
/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#startLoadJob} succeeds with a retry.
   */
@Test
public void testStartLoadJobRetry() throws IOException, InterruptedException {
    Job testJob = new Job();
    JobReference jobRef = new JobReference();
    jobRef.setJobId("jobId");
    jobRef.setProjectId("projectId");
    testJob.setJobReference(jobRef);
    // First response is 403 rate limited, second response has valid payload.
    when(response.getContentType()).thenReturn(Json.MEDIA_TYPE);
    when(response.getStatusCode()).thenReturn(403).thenReturn(200);
    when(response.getContent()).thenReturn(toStream(errorWithReasonAndStatus("rateLimitExceeded", 403))).thenReturn(toStream(testJob));
    TableReference ref = new TableReference();
    ref.setProjectId("projectId");
    JobConfigurationLoad loadConfig = new JobConfigurationLoad();
    loadConfig.setDestinationTable(ref);
    Sleeper sleeper = new FastNanoClockAndSleeper();
    BackOff backoff = new AttemptBoundedExponentialBackOff(5, /* attempts */
    1000);
    JobServiceImpl.startJob(testJob, new ApiErrorExtractor(), bigquery, sleeper, backoff);
    verify(response, times(2)).getStatusCode();
    verify(response, times(2)).getContent();
    verify(response, times(2)).getContentType();
}

13. BigQueryTableRowIterator#executeQueryAndWaitForCompletion()

View license
/**
   * 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 {
    // Create a temporary dataset to store results.
    // Starting dataset name with an "_" so that it is hidden.
    Random rnd = new Random(System.currentTimeMillis());
    temporaryDatasetId = "_dataflow_temporary_dataset_" + rnd.nextInt(1000000);
    temporaryTableId = "dataflow_temporary_table_" + rnd.nextInt(1000000);
    createDataset(temporaryDatasetId);
    Job job = new Job();
    JobConfiguration config = new JobConfiguration();
    JobConfigurationQuery queryConfig = new JobConfigurationQuery();
    config.setQuery(queryConfig);
    job.setConfiguration(config);
    queryConfig.setQuery(query);
    queryConfig.setAllowLargeResults(true);
    queryConfig.setFlattenResults(flattenResults);
    TableReference destinationTable = new TableReference();
    destinationTable.setProjectId(projectId);
    destinationTable.setDatasetId(temporaryDatasetId);
    destinationTable.setTableId(temporaryTableId);
    queryConfig.setDestinationTable(destinationTable);
    Insert insert = client.jobs().insert(projectId, job);
    Job queryJob = executeWithBackOff(insert, "Error when trying to execute the job for query " + query + " :{}");
    JobReference jobId = queryJob.getJobReference();
    while (true) {
        Job pollJob = executeWithBackOff(client.jobs().get(projectId, jobId.getJobId()), "Error when trying to get status of the job for query " + query + " :{}");
        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("Executing query " + query + " failed: " + error.getMessage());
            }
        }
        Uninterruptibles.sleepUninterruptibly(QUERY_COMPLETION_POLL_TIME.getMillis(), TimeUnit.MILLISECONDS);
    }
}

14. BigQueryTableRowIteratorTest#testQueryFailed()

View license
/**
   * 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 can be created.
    JobReference ref = new JobReference();
    Job insertedJob = new Job().setJobReference(ref);
    when(mockJobsInsert.execute()).thenReturn(insertedJob);
    // Job state polled with an error.
    String errorReason = "bad query";
    JobStatus status = new JobStatus().setState("DONE").setErrorResult(new ErrorProto().setMessage(errorReason));
    Job getJob = new Job().setJobReference(ref).setStatus(status);
    when(mockJobsGet.execute()).thenReturn(getJob);
    String query = "NOT A QUERY";
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(query, "project", mockClient, null)) {
        try {
            iterator.open();
            fail();
        } catch (Exception expected) {
            assertThat(expected.getMessage(), containsString("failed"));
            assertThat(expected.getMessage(), containsString(errorReason));
            assertThat(expected.getMessage(), containsString(query));
        }
    }
    // Temp dataset created and then later deleted.
    verify(mockClient, times(2)).datasets();
    verify(mockDatasets).insert(anyString(), any(Dataset.class));
    verify(mockDatasetsInsert).execute();
    verify(mockDatasets).delete(anyString(), anyString());
    verify(mockDatasetsDelete).execute();
    // Job inserted to run the query, then polled once.
    verify(mockClient, times(2)).jobs();
    verify(mockJobs).insert(anyString(), any(Job.class));
    verify(mockJobsInsert).execute();
    verify(mockJobs).get(anyString(), anyString());
    verify(mockJobsGet).execute();
}

15. BigQueryTableRowIterator#executeQueryAndWaitForCompletion()

View license
/**
   * 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 {
    // Create a temporary dataset to store results.
    // Starting dataset name with an "_" so that it is hidden.
    Random rnd = new Random(System.currentTimeMillis());
    temporaryDatasetId = "_dataflow_temporary_dataset_" + rnd.nextInt(1000000);
    temporaryTableId = "dataflow_temporary_table_" + rnd.nextInt(1000000);
    createDataset(temporaryDatasetId);
    Job job = new Job();
    JobConfiguration config = new JobConfiguration();
    JobConfigurationQuery queryConfig = new JobConfigurationQuery();
    config.setQuery(queryConfig);
    job.setConfiguration(config);
    queryConfig.setQuery(query);
    queryConfig.setAllowLargeResults(true);
    queryConfig.setFlattenResults(flattenResults);
    TableReference destinationTable = new TableReference();
    destinationTable.setProjectId(projectId);
    destinationTable.setDatasetId(temporaryDatasetId);
    destinationTable.setTableId(temporaryTableId);
    queryConfig.setDestinationTable(destinationTable);
    Insert insert = client.jobs().insert(projectId, job);
    Job queryJob = executeWithBackOff(insert, "Error when trying to execute the job for query " + query + " :{}");
    JobReference jobId = queryJob.getJobReference();
    while (true) {
        Job pollJob = executeWithBackOff(client.jobs().get(projectId, jobId.getJobId()), "Error when trying to get status of the job for query " + query + " :{}");
        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("Executing query " + query + " failed: " + error.getMessage());
            }
        }
        Uninterruptibles.sleepUninterruptibly(QUERY_COMPLETION_POLL_TIME.getMillis(), TimeUnit.MILLISECONDS);
    }
}

16. BigQueryTableRowIteratorTest#testQueryFailed()

View license
/**
   * 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 can be created.
    JobReference ref = new JobReference();
    Job insertedJob = new Job().setJobReference(ref);
    when(mockJobsInsert.execute()).thenReturn(insertedJob);
    // Job state polled with an error.
    String errorReason = "bad query";
    JobStatus status = new JobStatus().setState("DONE").setErrorResult(new ErrorProto().setMessage(errorReason));
    Job getJob = new Job().setJobReference(ref).setStatus(status);
    when(mockJobsGet.execute()).thenReturn(getJob);
    String query = "NOT A QUERY";
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(query, "project", mockClient, null)) {
        try {
            iterator.open();
            fail();
        } catch (Exception expected) {
            assertThat(expected.getMessage(), containsString("failed"));
            assertThat(expected.getMessage(), containsString(errorReason));
            assertThat(expected.getMessage(), containsString(query));
        }
    }
    // Temp dataset created and then later deleted.
    verify(mockClient, times(2)).datasets();
    verify(mockDatasets).insert(anyString(), any(Dataset.class));
    verify(mockDatasetsInsert).execute();
    verify(mockDatasets).delete(anyString(), anyString());
    verify(mockDatasetsDelete).execute();
    // Job inserted to run the query, then polled once.
    verify(mockClient, times(2)).jobs();
    verify(mockJobs).insert(anyString(), any(Job.class));
    verify(mockJobsInsert).execute();
    verify(mockJobs).get(anyString(), anyString());
    verify(mockJobsGet).execute();
}