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