use of com.google.api.services.bigquery.model.JobStatus in project beam by apache.
the class BigQueryServicesImplTest method testPollJobFailed.
/**
* 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();
}
use of com.google.api.services.bigquery.model.JobStatus in project beam by apache.
the class BigQueryServicesImplTest method testPollJobSucceeds.
/**
* 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();
}
use of com.google.api.services.bigquery.model.JobStatus 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.JobStatus in project beam by apache.
the class FakeJobService method runCopyJob.
private JobStatus runCopyJob(JobConfigurationTableCopy copy) throws InterruptedException, IOException {
List<TableReference> sources = copy.getSourceTables();
TableReference destination = copy.getDestinationTable();
WriteDisposition writeDisposition = WriteDisposition.valueOf(copy.getWriteDisposition());
CreateDisposition createDisposition = CreateDisposition.valueOf(copy.getCreateDisposition());
Table existingTable = datasetService.getTable(destination);
if (!validateDispositions(existingTable, createDisposition, writeDisposition)) {
return new JobStatus().setState("FAILED").setErrorResult(new ErrorProto());
}
List<TableRow> allRows = Lists.newArrayList();
for (TableReference source : sources) {
allRows.addAll(datasetService.getAllRows(source.getProjectId(), source.getDatasetId(), source.getTableId()));
}
datasetService.createTable(new Table().setTableReference(destination));
datasetService.insertAll(destination, allRows, null);
return new JobStatus().setState("DONE");
}
use of com.google.api.services.bigquery.model.JobStatus 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));
}
}
Aggregations