Search in sources :

Example 11 with JobStatus

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();
}
Also used : JobStatus(com.google.api.services.bigquery.model.JobStatus) ErrorProto(com.google.api.services.bigquery.model.ErrorProto) JobReference(com.google.api.services.bigquery.model.JobReference) JobServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.JobServiceImpl) JobServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.JobServiceImpl) Job(com.google.api.services.bigquery.model.Job) Test(org.junit.Test)

Example 12 with JobStatus

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();
}
Also used : JobStatus(com.google.api.services.bigquery.model.JobStatus) JobReference(com.google.api.services.bigquery.model.JobReference) JobServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.JobServiceImpl) JobServiceImpl(org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.JobServiceImpl) Job(com.google.api.services.bigquery.model.Job) Test(org.junit.Test)

Example 13 with JobStatus

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));
    }
}
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 14 with JobStatus

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");
}
Also used : JobStatus(com.google.api.services.bigquery.model.JobStatus) TableReference(com.google.api.services.bigquery.model.TableReference) CreateDisposition(org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.CreateDisposition) HashBasedTable(com.google.common.collect.HashBasedTable) Table(com.google.api.services.bigquery.model.Table) ErrorProto(com.google.api.services.bigquery.model.ErrorProto) TableRow(com.google.api.services.bigquery.model.TableRow) WriteDisposition(org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO.Write.WriteDisposition)

Example 15 with JobStatus

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

Aggregations

JobStatus (com.google.api.services.bigquery.model.JobStatus)20 Job (com.google.api.services.bigquery.model.Job)16 TableReference (com.google.api.services.bigquery.model.TableReference)9 TableRow (com.google.api.services.bigquery.model.TableRow)9 Test (org.junit.Test)9 JobConfiguration (com.google.api.services.bigquery.model.JobConfiguration)7 JobReference (com.google.api.services.bigquery.model.JobReference)7 JobStatistics (com.google.api.services.bigquery.model.JobStatistics)7 Table (com.google.api.services.bigquery.model.Table)7 HashBasedTable (com.google.common.collect.HashBasedTable)6 TableSchema (com.google.api.services.bigquery.model.TableSchema)5 ErrorProto (com.google.api.services.bigquery.model.ErrorProto)4 JobStatistics2 (com.google.api.services.bigquery.model.JobStatistics2)4 JobStatistics4 (com.google.api.services.bigquery.model.JobStatistics4)4 JobServiceImpl (org.apache.beam.sdk.io.gcp.bigquery.BigQueryServicesImpl.JobServiceImpl)4 TableFieldSchema (com.google.api.services.bigquery.model.TableFieldSchema)3 IOException (java.io.IOException)3 JsonSchemaToTableSchema (org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.JsonSchemaToTableSchema)3 BigQueryHelpers.createTempTableReference (org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createTempTableReference)3 BigQueryHelpers.toJsonString (org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.toJsonString)3