Search in sources :

Example 6 with JobConfiguration

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

the class FakeJobService method startCopyJob.

@Override
public void startCopyJob(JobReference jobRef, JobConfigurationTableCopy copyConfig) throws IOException {
    synchronized (allJobs) {
        verifyUniqueJobId(jobRef.getJobId());
        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 7 with JobConfiguration

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

the class FakeJobService method startExtractJob.

@Override
public void startExtractJob(JobReference jobRef, JobConfigurationExtract extractConfig) throws IOException {
    checkArgument("AVRO".equals(extractConfig.getDestinationFormat()), "Only extract to AVRO is supported");
    synchronized (allJobs) {
        verifyUniqueJobId(jobRef.getJobId());
        ++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));
    }
}
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 8 with JobConfiguration

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

the class FakeJobService method startQueryJob.

@Override
public void startQueryJob(JobReference jobRef, JobConfigurationQuery query) throws IOException, InterruptedException {
    synchronized (allJobs) {
        Job job = new Job();
        job.setJobReference(jobRef);
        job.setConfiguration(new JobConfiguration().setQuery(query));
        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 9 with JobConfiguration

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

the class FakeJobService method startLoadJob.

@Override
public void startLoadJob(JobReference jobRef, JobConfigurationLoad loadConfig) throws InterruptedException, IOException {
    synchronized (allJobs) {
        Job job = new Job();
        job.setJobReference(jobRef);
        job.setConfiguration(new JobConfiguration().setLoad(loadConfig));
        job.setKind(" bigquery#job");
        job.setStatus(new JobStatus().setState("PENDING"));
        // the caller.
        if (loadConfig.getSourceUris().size() > 0) {
            ImmutableList.Builder<ResourceId> sourceFiles = ImmutableList.builder();
            ImmutableList.Builder<ResourceId> loadFiles = ImmutableList.builder();
            for (String filename : loadConfig.getSourceUris()) {
                sourceFiles.add(FileSystems.matchNewResource(filename, false));
                loadFiles.add(FileSystems.matchNewResource(filename + ThreadLocalRandom.current().nextInt(), false));
            }
            FileSystems.copy(sourceFiles.build(), loadFiles.build(), MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES);
            filesForLoadJobs.put(jobRef.getProjectId(), jobRef.getJobId(), loadFiles.build());
        }
        allJobs.put(jobRef.getProjectId(), jobRef.getJobId(), new JobInfo(job));
    }
}
Also used : JobStatus(com.google.api.services.bigquery.model.JobStatus) ResourceId(org.apache.beam.sdk.io.fs.ResourceId) ImmutableList(com.google.common.collect.ImmutableList) Job(com.google.api.services.bigquery.model.Job) JobConfiguration(com.google.api.services.bigquery.model.JobConfiguration)

Example 10 with JobConfiguration

use of com.google.api.services.bigquery.model.JobConfiguration 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();
}
Also used : JobStatistics(com.google.api.services.bigquery.model.JobStatistics) JobStatistics2(com.google.api.services.bigquery.model.JobStatistics2) JobReference(com.google.api.services.bigquery.model.JobReference) JobConfigurationQuery(com.google.api.services.bigquery.model.JobConfigurationQuery) Dataset(com.google.api.services.bigquery.model.Dataset) Matchers.anyString(org.mockito.Matchers.anyString) Matchers.containsString(org.hamcrest.Matchers.containsString) JobStatus(com.google.api.services.bigquery.model.JobStatus) TableReference(com.google.api.services.bigquery.model.TableReference) TableRow(com.google.api.services.bigquery.model.TableRow) Job(com.google.api.services.bigquery.model.Job) JobConfiguration(com.google.api.services.bigquery.model.JobConfiguration) Test(org.junit.Test)

Aggregations

Job (com.google.api.services.bigquery.model.Job)12 JobConfiguration (com.google.api.services.bigquery.model.JobConfiguration)12 JobStatus (com.google.api.services.bigquery.model.JobStatus)11 TableReference (com.google.api.services.bigquery.model.TableReference)4 JobConfigurationQuery (com.google.api.services.bigquery.model.JobConfigurationQuery)3 JobReference (com.google.api.services.bigquery.model.JobReference)3 JobStatistics (com.google.api.services.bigquery.model.JobStatistics)3 TableRow (com.google.api.services.bigquery.model.TableRow)3 Dataset (com.google.api.services.bigquery.model.Dataset)2 JobStatistics2 (com.google.api.services.bigquery.model.JobStatistics2)2 Table (com.google.api.services.bigquery.model.Table)2 Random (java.util.Random)2 ResourceId (org.apache.beam.sdk.io.fs.ResourceId)2 Matchers.containsString (org.hamcrest.Matchers.containsString)2 Test (org.junit.Test)2 Matchers.anyString (org.mockito.Matchers.anyString)2 ErrorProto (com.google.api.services.bigquery.model.ErrorProto)1 GetQueryResultsResponse (com.google.api.services.bigquery.model.GetQueryResultsResponse)1 TableSchema (com.google.api.services.bigquery.model.TableSchema)1 ImmutableList (com.google.common.collect.ImmutableList)1