com.google.api.services.bigquery.model.Job

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

1. JobInfo#toPb()

Project: gcloud-java
Source File: JobInfo.java
View license
Job toPb() {
    Job jobPb = new Job();
    jobPb.setEtag(etag);
    jobPb.setId(generatedId);
    jobPb.setSelfLink(selfLink);
    jobPb.setUserEmail(userEmail);
    if (jobId != null) {
        jobPb.setJobReference(jobId.toPb());
    }
    if (status != null) {
        jobPb.setStatus(status.toPb());
    }
    if (statistics != null) {
        jobPb.setStatistics(statistics.toPb());
    }
    jobPb.setConfiguration(configuration.toPb());
    return jobPb;
}

2. 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();
}

3. 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();
}

4. 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();
}

5. 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();
}

6. 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();
}

7. 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();
}

8. 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();
}

9. 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();
}

10. 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);
    }
}

11. 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);
    }
}

12. 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();
}

13. BigQueryTableRowIteratorTest#testReadFromQuery()

View license
/**
   * Verifies that when the query runs, the correct data is returned and the temporary dataset and
   * table are both cleaned up.
   */
@Test
public void testReadFromQuery() throws IOException, InterruptedException {
    // Mock job inserting.
    Job insertedJob = new Job().setJobReference(new JobReference());
    when(mockJobsInsert.execute()).thenReturn(insertedJob);
    // Mock job polling.
    JobStatus status = new JobStatus().setState("DONE");
    TableReference tableRef = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setDestinationTable(tableRef);
    Job getJob = new Job().setJobReference(new JobReference()).setStatus(status).setConfiguration(new JobConfiguration().setQuery(queryConfig));
    when(mockJobsGet.execute()).thenReturn(getJob);
    // Mock table schema fetch.
    when(mockTablesGet.execute()).thenReturn(tableWithBasicSchema());
    // Mock table data fetch.
    when(mockTabledataList.execute()).thenReturn(rawDataList(rawRow("Arthur", 42)));
    // Run query and verify
    String query = "SELECT name, count from table";
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(query, "project", mockClient, null)) {
        iterator.open();
        assertTrue(iterator.advance());
        TableRow row = iterator.getCurrent();
        assertTrue(row.containsKey("name"));
        assertTrue(row.containsKey("answer"));
        assertEquals("Arthur", row.get("name"));
        assertEquals(42, row.get("answer"));
        assertFalse(iterator.advance());
    }
    // Temp dataset created and 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, 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();
    // Temp table get after query finish, deleted after reading.
    verify(mockClient, times(2)).tables();
    verify(mockTables).get("project", "dataset", "table");
    verify(mockTablesGet).execute();
    verify(mockTables).delete(anyString(), anyString(), anyString());
    verify(mockTablesDelete).execute();
    // Table data read.
    verify(mockClient).tabledata();
    verify(mockTabledata).list("project", "dataset", "table");
    verify(mockTabledataList).execute();
}

14. 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();
}

15. 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();
}

16. BigQueryIOTest#testBigQueryQuerySourceInitSplit()

Project: incubator-beam
Source File: BigQueryIOTest.java
View license
@Test
public void testBigQueryQuerySourceInitSplit() throws Exception {
    TableReference dryRunTable = new TableReference();
    Job queryJob = new Job();
    JobStatistics queryJobStats = new JobStatistics();
    JobStatistics2 queryStats = new JobStatistics2();
    queryStats.setReferencedTables(ImmutableList.of(dryRunTable));
    queryJobStats.setQuery(queryStats);
    queryJob.setStatus(new JobStatus()).setStatistics(queryJobStats);
    Job extractJob = new Job();
    JobStatistics extractJobStats = new JobStatistics();
    JobStatistics4 extractStats = new JobStatistics4();
    extractStats.setDestinationUriFileCounts(ImmutableList.of(1L));
    extractJobStats.setExtract(extractStats);
    extractJob.setStatus(new JobStatus()).setStatistics(extractJobStats);
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(mockJobService).withDatasetService(mockDatasetService).readerReturns(toJsonString(new TableRow().set("name", "a").set("number", "1")), toJsonString(new TableRow().set("name", "b").set("number", "2")), toJsonString(new TableRow().set("name", "c").set("number", "3")));
    String jobIdToken = "testJobIdToken";
    String extractDestinationDir = "mock://tempLocation";
    TableReference destinationTable = BigQueryIO.parseTableSpec("project:data_set.table_name");
    BoundedSource<TableRow> bqSource = BigQueryQuerySource.create(jobIdToken, "query", destinationTable, true, /* flattenResults */
    extractDestinationDir, fakeBqServices);
    List<TableRow> expected = ImmutableList.of(new TableRow().set("name", "a").set("number", "1"), new TableRow().set("name", "b").set("number", "2"), new TableRow().set("name", "c").set("number", "3"));
    PipelineOptions options = PipelineOptionsFactory.create();
    options.setTempLocation(extractDestinationDir);
    TableReference queryTable = new TableReference().setProjectId("testProejct").setDatasetId("testDataset").setTableId("testTable");
    when(mockJobService.dryRunQuery(anyString(), anyString())).thenReturn(new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(100L).setReferencedTables(ImmutableList.of(queryTable))));
    when(mockDatasetService.getTable(eq(queryTable.getProjectId()), eq(queryTable.getDatasetId()), eq(queryTable.getTableId()))).thenReturn(new Table().setSchema(new TableSchema()));
    when(mockDatasetService.getTable(eq(destinationTable.getProjectId()), eq(destinationTable.getDatasetId()), eq(destinationTable.getTableId()))).thenReturn(new Table().setSchema(new TableSchema()));
    IOChannelUtils.setIOFactory("mock", mockIOChannelFactory);
    when(mockIOChannelFactory.resolve(anyString(), anyString())).thenReturn("mock://tempLocation/output");
    when(mockJobService.pollJob(Mockito.<JobReference>any(), Mockito.anyInt())).thenReturn(extractJob);
    Assert.assertThat(SourceTestUtils.readFromSource(bqSource, options), CoreMatchers.is(expected));
    SourceTestUtils.assertSplitAtFractionBehavior(bqSource, 2, 0.3, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options);
    List<? extends BoundedSource<TableRow>> sources = bqSource.splitIntoBundles(100, options);
    assertEquals(1, sources.size());
    BoundedSource<TableRow> actual = sources.get(0);
    assertThat(actual, CoreMatchers.instanceOf(TransformingSource.class));
    Mockito.verify(mockJobService).startQueryJob(Mockito.<JobReference>any(), Mockito.<JobConfigurationQuery>any());
    Mockito.verify(mockJobService).startExtractJob(Mockito.<JobReference>any(), Mockito.<JobConfigurationExtract>any());
    Mockito.verify(mockDatasetService).createDataset(anyString(), anyString(), anyString(), anyString());
}

17. 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();
}

18. BigQueryTableRowIteratorTest#testReadFromQuery()

View license
/**
   * Verifies that when the query runs, the correct data is returned and the temporary dataset and
   * table are both cleaned up.
   */
@Test
public void testReadFromQuery() throws IOException, InterruptedException {
    // Mock job inserting.
    Job insertedJob = new Job().setJobReference(new JobReference());
    when(mockJobsInsert.execute()).thenReturn(insertedJob);
    // Mock job polling.
    JobStatus status = new JobStatus().setState("DONE");
    TableReference tableRef = new TableReference().setProjectId("project").setDatasetId("dataset").setTableId("table");
    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setDestinationTable(tableRef);
    Job getJob = new Job().setJobReference(new JobReference()).setStatus(status).setConfiguration(new JobConfiguration().setQuery(queryConfig));
    when(mockJobsGet.execute()).thenReturn(getJob);
    // Mock table schema fetch.
    when(mockTablesGet.execute()).thenReturn(tableWithBasicSchema());
    // Mock table data fetch.
    when(mockTabledataList.execute()).thenReturn(rawDataList(rawRow("Arthur", 42)));
    // Run query and verify
    String query = "SELECT name, count from table";
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(query, "project", mockClient, null)) {
        iterator.open();
        assertTrue(iterator.advance());
        TableRow row = iterator.getCurrent();
        assertTrue(row.containsKey("name"));
        assertTrue(row.containsKey("answer"));
        assertEquals("Arthur", row.get("name"));
        assertEquals(42, row.get("answer"));
        assertFalse(iterator.advance());
    }
    // Temp dataset created and 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, 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();
    // Temp table get after query finish, deleted after reading.
    verify(mockClient, times(2)).tables();
    verify(mockTables).get("project", "dataset", "table");
    verify(mockTablesGet).execute();
    verify(mockTables).delete(anyString(), anyString(), anyString());
    verify(mockTablesDelete).execute();
    // Table data read.
    verify(mockClient).tabledata();
    verify(mockTabledata).list("project", "dataset", "table");
    verify(mockTabledataList).execute();
}

19. 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();
}

20. 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();
}

21. BigQueryIOTest#testBigQueryQuerySourceInitSplit()

Project: DataflowJavaSDK
Source File: BigQueryIOTest.java
View license
@Test
public void testBigQueryQuerySourceInitSplit() throws Exception {
    TableReference dryRunTable = new TableReference();
    Job queryJob = new Job();
    JobStatistics queryJobStats = new JobStatistics();
    JobStatistics2 queryStats = new JobStatistics2();
    queryStats.setReferencedTables(ImmutableList.of(dryRunTable));
    queryJobStats.setQuery(queryStats);
    queryJob.setStatus(new JobStatus()).setStatistics(queryJobStats);
    Job extractJob = new Job();
    JobStatistics extractJobStats = new JobStatistics();
    JobStatistics4 extractStats = new JobStatistics4();
    extractStats.setDestinationUriFileCounts(ImmutableList.of(1L));
    extractJobStats.setExtract(extractStats);
    extractJob.setStatus(new JobStatus()).setStatistics(extractJobStats);
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(mockJobService).withDatasetService(mockDatasetService).readerReturns(toJsonString(new TableRow().set("name", "a").set("number", "1")), toJsonString(new TableRow().set("name", "b").set("number", "2")), toJsonString(new TableRow().set("name", "c").set("number", "3")));
    String jobIdToken = "testJobIdToken";
    String extractDestinationDir = "mock://tempLocation";
    TableReference destinationTable = BigQueryIO.parseTableSpec("project:data_set.table_name");
    BoundedSource<TableRow> bqSource = BigQueryQuerySource.create(jobIdToken, "query", destinationTable, true, /* flattenResults */
    extractDestinationDir, fakeBqServices);
    List<TableRow> expected = ImmutableList.of(new TableRow().set("name", "a").set("number", "1"), new TableRow().set("name", "b").set("number", "2"), new TableRow().set("name", "c").set("number", "3"));
    PipelineOptions options = PipelineOptionsFactory.create();
    options.setTempLocation(extractDestinationDir);
    TableReference queryTable = new TableReference().setProjectId("testProejct").setDatasetId("testDataset").setTableId("testTable");
    when(mockJobService.dryRunQuery(anyString(), anyString())).thenReturn(new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(100L).setReferencedTables(ImmutableList.of(queryTable))));
    when(mockDatasetService.getTable(eq(queryTable.getProjectId()), eq(queryTable.getDatasetId()), eq(queryTable.getTableId()))).thenReturn(new Table().setSchema(new TableSchema()));
    when(mockDatasetService.getTable(eq(destinationTable.getProjectId()), eq(destinationTable.getDatasetId()), eq(destinationTable.getTableId()))).thenReturn(new Table().setSchema(new TableSchema()));
    IOChannelUtils.setIOFactory("mock", mockIOChannelFactory);
    when(mockIOChannelFactory.resolve(anyString(), anyString())).thenReturn("mock://tempLocation/output");
    when(mockJobService.pollJob(Mockito.<JobReference>any(), Mockito.anyInt())).thenReturn(extractJob);
    Assert.assertThat(SourceTestUtils.readFromSource(bqSource, options), CoreMatchers.is(expected));
    SourceTestUtils.assertSplitAtFractionBehavior(bqSource, 2, 0.3, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options);
    List<? extends BoundedSource<TableRow>> sources = bqSource.splitIntoBundles(100, options);
    assertEquals(1, sources.size());
    BoundedSource<TableRow> actual = sources.get(0);
    assertThat(actual, CoreMatchers.instanceOf(TransformingSource.class));
    Mockito.verify(mockJobService).startQueryJob(Mockito.<JobReference>any(), Mockito.<JobConfigurationQuery>any());
    Mockito.verify(mockJobService).startExtractJob(Mockito.<JobReference>any(), Mockito.<JobConfigurationExtract>any());
    Mockito.verify(mockDatasetService).createDataset(anyString(), anyString(), anyString(), anyString());
}

22. BigQueryIOTest#testBigQueryTableSourceInitSplit()

Project: incubator-beam
Source File: BigQueryIOTest.java
View license
@Test
public void testBigQueryTableSourceInitSplit() throws Exception {
    Job extractJob = new Job();
    JobStatistics jobStats = new JobStatistics();
    JobStatistics4 extractStats = new JobStatistics4();
    extractStats.setDestinationUriFileCounts(ImmutableList.of(1L));
    jobStats.setExtract(extractStats);
    extractJob.setStatus(new JobStatus()).setStatistics(jobStats);
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(mockJobService).withDatasetService(mockDatasetService).readerReturns(toJsonString(new TableRow().set("name", "a").set("number", "1")), toJsonString(new TableRow().set("name", "b").set("number", "2")), toJsonString(new TableRow().set("name", "c").set("number", "3")));
    String jobIdToken = "testJobIdToken";
    TableReference table = BigQueryIO.parseTableSpec("project:data_set.table_name");
    String extractDestinationDir = "mock://tempLocation";
    BoundedSource<TableRow> bqSource = BigQueryTableSource.create(jobIdToken, table, extractDestinationDir, fakeBqServices, "project");
    List<TableRow> expected = ImmutableList.of(new TableRow().set("name", "a").set("number", "1"), new TableRow().set("name", "b").set("number", "2"), new TableRow().set("name", "c").set("number", "3"));
    when(mockJobService.pollJob(Mockito.<JobReference>any(), Mockito.anyInt())).thenReturn(extractJob);
    PipelineOptions options = PipelineOptionsFactory.create();
    options.setTempLocation("mock://tempLocation");
    IOChannelUtils.setIOFactory("mock", mockIOChannelFactory);
    when(mockIOChannelFactory.resolve(anyString(), anyString())).thenReturn("mock://tempLocation/output");
    when(mockDatasetService.getTable(anyString(), anyString(), anyString())).thenReturn(new Table().setSchema(new TableSchema()));
    Assert.assertThat(SourceTestUtils.readFromSource(bqSource, options), CoreMatchers.is(expected));
    SourceTestUtils.assertSplitAtFractionBehavior(bqSource, 2, 0.3, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options);
    List<? extends BoundedSource<TableRow>> sources = bqSource.splitIntoBundles(100, options);
    assertEquals(1, sources.size());
    BoundedSource<TableRow> actual = sources.get(0);
    assertThat(actual, CoreMatchers.instanceOf(TransformingSource.class));
    Mockito.verify(mockJobService).startExtractJob(Mockito.<JobReference>any(), Mockito.<JobConfigurationExtract>any());
}

23. BigQueryIOTest#testBigQueryTableSourceInitSplit()

Project: DataflowJavaSDK
Source File: BigQueryIOTest.java
View license
@Test
public void testBigQueryTableSourceInitSplit() throws Exception {
    Job extractJob = new Job();
    JobStatistics jobStats = new JobStatistics();
    JobStatistics4 extractStats = new JobStatistics4();
    extractStats.setDestinationUriFileCounts(ImmutableList.of(1L));
    jobStats.setExtract(extractStats);
    extractJob.setStatus(new JobStatus()).setStatistics(jobStats);
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(mockJobService).withDatasetService(mockDatasetService).readerReturns(toJsonString(new TableRow().set("name", "a").set("number", "1")), toJsonString(new TableRow().set("name", "b").set("number", "2")), toJsonString(new TableRow().set("name", "c").set("number", "3")));
    String jobIdToken = "testJobIdToken";
    TableReference table = BigQueryIO.parseTableSpec("project:data_set.table_name");
    String extractDestinationDir = "mock://tempLocation";
    BoundedSource<TableRow> bqSource = BigQueryTableSource.create(jobIdToken, table, extractDestinationDir, fakeBqServices, "project");
    List<TableRow> expected = ImmutableList.of(new TableRow().set("name", "a").set("number", "1"), new TableRow().set("name", "b").set("number", "2"), new TableRow().set("name", "c").set("number", "3"));
    when(mockJobService.pollJob(Mockito.<JobReference>any(), Mockito.anyInt())).thenReturn(extractJob);
    PipelineOptions options = PipelineOptionsFactory.create();
    options.setTempLocation("mock://tempLocation");
    IOChannelUtils.setIOFactory("mock", mockIOChannelFactory);
    when(mockIOChannelFactory.resolve(anyString(), anyString())).thenReturn("mock://tempLocation/output");
    when(mockDatasetService.getTable(anyString(), anyString(), anyString())).thenReturn(new Table().setSchema(new TableSchema()));
    Assert.assertThat(SourceTestUtils.readFromSource(bqSource, options), CoreMatchers.is(expected));
    SourceTestUtils.assertSplitAtFractionBehavior(bqSource, 2, 0.3, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options);
    List<? extends BoundedSource<TableRow>> sources = bqSource.splitIntoBundles(100, options);
    assertEquals(1, sources.size());
    BoundedSource<TableRow> actual = sources.get(0);
    assertThat(actual, CoreMatchers.instanceOf(TransformingSource.class));
    Mockito.verify(mockJobService).startExtractJob(Mockito.<JobReference>any(), Mockito.<JobConfigurationExtract>any());
}