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

Example 26 with Job

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

the class BigQueryTableRowIteratorTest method testReadFromQueryNoTables.

/**
   * Verifies that queries that reference no data can be read.
   */
@Test
public void testReadFromQueryNoTables() throws IOException, InterruptedException {
    // Mock job inserting.
    Job dryRunJob = new Job().setStatistics(new JobStatistics().setQuery(new JobStatistics2()));
    Job insertedJob = new Job().setJobReference(new JobReference());
    when(mockJobsInsert.execute()).thenReturn(dryRunJob, insertedJob);
    // Mock job polling.
    JobStatus status = new JobStatus().setState("DONE");
    JobConfigurationQuery resultQueryConfig = new JobConfigurationQuery().setDestinationTable(new TableReference().setProjectId("project").setDatasetId("tempdataset").setTableId("temptable"));
    Job getJob = new Job().setJobReference(new JobReference()).setStatus(status).setConfiguration(new JobConfiguration().setQuery(resultQueryConfig));
    when(mockJobsGet.execute()).thenReturn(getJob);
    // Mock table schema fetch.
    when(mockTablesGet.execute()).thenReturn(noTableQuerySchema());
    byte[] photoBytes = "photograph".getBytes();
    String photoBytesEncoded = BaseEncoding.base64().encode(photoBytes);
    // Mock table data fetch.
    when(mockTabledataList.execute()).thenReturn(rawDataList(rawRow("Arthur", 42, photoBytesEncoded)));
    // Run query and verify
    String query = String.format("SELECT \"Arthur\" as name, 42 as count, \"%s\" as photo", photoBytesEncoded);
    JobConfigurationQuery queryConfig = new JobConfigurationQuery().setQuery(query);
    try (BigQueryTableRowIterator iterator = BigQueryTableRowIterator.fromQuery(queryConfig, "project", mockClient)) {
        iterator.open();
        assertTrue(iterator.advance());
        TableRow row = iterator.getCurrent();
        assertTrue(row.containsKey("name"));
        assertTrue(row.containsKey("count"));
        assertTrue(row.containsKey("photo"));
        assertEquals("Arthur", row.get("name"));
        assertEquals(42, row.get("count"));
        assertEquals(photoBytesEncoded, row.get("photo"));
        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(3)).jobs();
    verify(mockJobs, times(2)).insert(anyString(), any(Job.class));
    verify(mockJobsInsert, times(2)).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, times(1)).get(anyString(), anyString(), anyString());
    verify(mockTablesGet, times(1)).execute();
    verify(mockTables).delete(anyString(), anyString(), anyString());
    verify(mockTablesDelete).execute();
    // Table data read.
    verify(mockClient).tabledata();
    verify(mockTabledata).list("project", "tempdataset", "temptable");
    verify(mockTabledataList).execute();
}

Example 27 with Job

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

the class FakeJobService method startExtractJob.

@Override
public void startExtractJob(JobReference jobRef, JobConfigurationExtract extractConfig) throws InterruptedException, IOException {
    checkArgument(extractConfig.getDestinationFormat().equals("AVRO"), "Only extract to AVRO is supported");
    synchronized (allJobs) {
        ++numExtractJobCalls;
        Job job = new Job();
        job.setJobReference(jobRef);
        job.setConfiguration(new JobConfiguration().setExtract(extractConfig));
        job.setKind(" bigquery#job");
        job.setStatus(new JobStatus().setState("PENDING"));
        allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
    }
}

Example 28 with Job

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

the class BigQueryQuerySource method executeQuery.

private void executeQuery(String executingProject, String jobId, TableReference destinationTable, JobService jobService) throws IOException, InterruptedException {
    JobReference jobRef = new JobReference().setProjectId(executingProject).setJobId(jobId);
    JobConfigurationQuery queryConfig = createBasicQueryConfig().setAllowLargeResults(true).setCreateDisposition("CREATE_IF_NEEDED").setDestinationTable(destinationTable).setPriority("BATCH").setWriteDisposition("WRITE_EMPTY");
    jobService.startQueryJob(jobRef, queryConfig);
    Job job = jobService.pollJob(jobRef, JOB_POLL_MAX_RETRIES);
    if (BigQueryHelpers.parseStatus(job) != Status.SUCCEEDED) {
        throw new IOException(String.format("Query job %s failed, status: %s.", jobId, BigQueryHelpers.statusToPrettyString(job.getStatus())));
    }
}

Example 29 with Job

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

the class WriteTables method load.

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