Search in sources :

Example 6 with JobStatus

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

the class BigQueryServicesImplTest method testPollJobUnknown.

/**
   * Tests that {@link BigQueryServicesImpl.JobServiceImpl#pollJob} returns UNKNOWN.
   */
@Test
public void testPollJobUnknown() throws IOException, InterruptedException {
    Job testJob = new Job();
    testJob.setStatus(new JobStatus());
    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.STOP_BACKOFF);
    assertEquals(null, 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 7 with JobStatus

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

the class BigQueryIOTest method testBigQueryQuerySourceInitSplit.

@Test
public void testBigQueryQuerySourceInitSplit() throws Exception {
    TableReference dryRunTable = new TableReference();
    Job queryJob = new Job();
    JobStatistics queryJobStats = new JobStatistics();
    JobStatistics2 queryStats = new JobStatistics2();
    queryStats.setReferencedTables(ImmutableList.of(dryRunTable));
    queryJobStats.setQuery(queryStats);
    queryJob.setStatus(new JobStatus()).setStatistics(queryJobStats);
    Job extractJob = new Job();
    JobStatistics extractJobStats = new JobStatistics();
    JobStatistics4 extractStats = new JobStatistics4();
    extractStats.setDestinationUriFileCounts(ImmutableList.of(1L));
    extractJobStats.setExtract(extractStats);
    extractJob.setStatus(new JobStatus()).setStatistics(extractJobStats);
    FakeJobService fakeJobService = new FakeJobService();
    FakeDatasetService fakeDatasetService = new FakeDatasetService();
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(fakeJobService).withDatasetService(fakeDatasetService);
    List<TableRow> expected = ImmutableList.of(new TableRow().set("name", "a").set("number", 1L), new TableRow().set("name", "b").set("number", 2L), new TableRow().set("name", "c").set("number", 3L), new TableRow().set("name", "d").set("number", 4L), new TableRow().set("name", "e").set("number", 5L), new TableRow().set("name", "f").set("number", 6L));
    PipelineOptions options = PipelineOptionsFactory.create();
    BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
    bqOptions.setProject("project");
    String stepUuid = "testStepUuid";
    TableReference tempTableReference = createTempTableReference(bqOptions.getProject(), createJobIdToken(bqOptions.getJobName(), stepUuid));
    fakeDatasetService.createDataset(bqOptions.getProject(), tempTableReference.getDatasetId(), "", "");
    fakeDatasetService.createTable(new Table().setTableReference(tempTableReference).setSchema(new TableSchema().setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("number").setType("INTEGER")))));
    Path baseDir = Files.createTempDirectory(tempFolder, "testBigQueryQuerySourceInitSplit");
    String query = FakeBigQueryServices.encodeQuery(expected);
    BoundedSource<TableRow> bqSource = BigQueryQuerySource.create(stepUuid, StaticValueProvider.of(query), true, /* flattenResults */
    true, /* useLegacySql */
    fakeBqServices);
    options.setTempLocation(baseDir.toString());
    TableReference queryTable = new TableReference().setProjectId(bqOptions.getProject()).setDatasetId(tempTableReference.getDatasetId()).setTableId(tempTableReference.getTableId());
    fakeJobService.expectDryRunQuery(bqOptions.getProject(), query, new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(100L).setReferencedTables(ImmutableList.of(queryTable))));
    List<TableRow> read = SourceTestUtils.readFromSource(bqSource, options);
    assertThat(read, containsInAnyOrder(Iterables.toArray(expected, TableRow.class)));
    SourceTestUtils.assertSplitAtFractionBehavior(bqSource, 2, 0.3, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options);
    List<? extends BoundedSource<TableRow>> sources = bqSource.split(100, options);
    assertEquals(2, sources.size());
    BoundedSource<TableRow> actual = sources.get(0);
    assertThat(actual, CoreMatchers.instanceOf(TransformingSource.class));
}
Also used : Path(java.nio.file.Path) JobStatistics(com.google.api.services.bigquery.model.JobStatistics) JobStatistics2(com.google.api.services.bigquery.model.JobStatistics2) HashBasedTable(com.google.common.collect.HashBasedTable) Table(com.google.api.services.bigquery.model.Table) JobStatistics4(com.google.api.services.bigquery.model.JobStatistics4) TableSchema(com.google.api.services.bigquery.model.TableSchema) JsonSchemaToTableSchema(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.JsonSchemaToTableSchema) BigQueryHelpers.toJsonString(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.toJsonString) TableFieldSchema(com.google.api.services.bigquery.model.TableFieldSchema) JobStatus(com.google.api.services.bigquery.model.JobStatus) BigQueryHelpers.createTempTableReference(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createTempTableReference) TableReference(com.google.api.services.bigquery.model.TableReference) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) TableRow(com.google.api.services.bigquery.model.TableRow) Job(com.google.api.services.bigquery.model.Job) Test(org.junit.Test)

Example 8 with JobStatus

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

the class BigQueryIOTest method testBigQueryNoTableQuerySourceInitSplit.

@Test
public void testBigQueryNoTableQuerySourceInitSplit() throws Exception {
    TableReference dryRunTable = new TableReference();
    Job queryJob = new Job();
    JobStatistics queryJobStats = new JobStatistics();
    JobStatistics2 queryStats = new JobStatistics2();
    queryStats.setReferencedTables(ImmutableList.of(dryRunTable));
    queryJobStats.setQuery(queryStats);
    queryJob.setStatus(new JobStatus()).setStatistics(queryJobStats);
    Job extractJob = new Job();
    JobStatistics extractJobStats = new JobStatistics();
    JobStatistics4 extractStats = new JobStatistics4();
    extractStats.setDestinationUriFileCounts(ImmutableList.of(1L));
    extractJobStats.setExtract(extractStats);
    extractJob.setStatus(new JobStatus()).setStatistics(extractJobStats);
    FakeDatasetService datasetService = new FakeDatasetService();
    FakeJobService jobService = new FakeJobService();
    FakeBigQueryServices fakeBqServices = new FakeBigQueryServices().withJobService(jobService).withDatasetService(datasetService);
    PipelineOptions options = PipelineOptionsFactory.create();
    BigQueryOptions bqOptions = options.as(BigQueryOptions.class);
    bqOptions.setProject("project");
    String stepUuid = "testStepUuid";
    TableReference tempTableReference = createTempTableReference(bqOptions.getProject(), createJobIdToken(bqOptions.getJobName(), stepUuid));
    List<TableRow> expected = ImmutableList.of(new TableRow().set("name", "a").set("number", 1L), new TableRow().set("name", "b").set("number", 2L), new TableRow().set("name", "c").set("number", 3L), new TableRow().set("name", "d").set("number", 4L), new TableRow().set("name", "e").set("number", 5L), new TableRow().set("name", "f").set("number", 6L));
    datasetService.createDataset(tempTableReference.getProjectId(), tempTableReference.getDatasetId(), "", "");
    Table table = new Table().setTableReference(tempTableReference).setSchema(new TableSchema().setFields(ImmutableList.of(new TableFieldSchema().setName("name").setType("STRING"), new TableFieldSchema().setName("number").setType("INTEGER"))));
    datasetService.createTable(table);
    String query = FakeBigQueryServices.encodeQuery(expected);
    jobService.expectDryRunQuery("project", query, new JobStatistics().setQuery(new JobStatistics2().setTotalBytesProcessed(100L).setReferencedTables(ImmutableList.of(table.getTableReference()))));
    Path baseDir = Files.createTempDirectory(tempFolder, "testBigQueryNoTableQuerySourceInitSplit");
    BoundedSource<TableRow> bqSource = BigQueryQuerySource.create(stepUuid, StaticValueProvider.of(query), true, /* flattenResults */
    true, /* useLegacySql */
    fakeBqServices);
    options.setTempLocation(baseDir.toString());
    List<TableRow> read = convertBigDecimaslToLong(SourceTestUtils.readFromSource(bqSource, options));
    assertThat(read, containsInAnyOrder(Iterables.toArray(expected, TableRow.class)));
    SourceTestUtils.assertSplitAtFractionBehavior(bqSource, 2, 0.3, ExpectedSplitOutcome.MUST_BE_CONSISTENT_IF_SUCCEEDS, options);
    List<? extends BoundedSource<TableRow>> sources = bqSource.split(100, options);
    assertEquals(2, sources.size());
    BoundedSource<TableRow> actual = sources.get(0);
    assertThat(actual, CoreMatchers.instanceOf(TransformingSource.class));
}
Also used : Path(java.nio.file.Path) JobStatistics(com.google.api.services.bigquery.model.JobStatistics) JobStatistics2(com.google.api.services.bigquery.model.JobStatistics2) HashBasedTable(com.google.common.collect.HashBasedTable) Table(com.google.api.services.bigquery.model.Table) JobStatistics4(com.google.api.services.bigquery.model.JobStatistics4) TableSchema(com.google.api.services.bigquery.model.TableSchema) JsonSchemaToTableSchema(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.JsonSchemaToTableSchema) BigQueryHelpers.toJsonString(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.toJsonString) TableFieldSchema(com.google.api.services.bigquery.model.TableFieldSchema) JobStatus(com.google.api.services.bigquery.model.JobStatus) BigQueryHelpers.createTempTableReference(org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers.createTempTableReference) TableReference(com.google.api.services.bigquery.model.TableReference) PipelineOptions(org.apache.beam.sdk.options.PipelineOptions) TableRow(com.google.api.services.bigquery.model.TableRow) Job(com.google.api.services.bigquery.model.Job) Test(org.junit.Test)

Example 9 with JobStatus

use of com.google.api.services.bigquery.model.JobStatus in project google-cloud-java by GoogleCloudPlatform.

the class HttpBigQueryRpc method listJobs.

@Override
public Tuple<String, Iterable<Job>> listJobs(String projectId, Map<Option, ?> options) {
    try {
        JobList jobsList = bigquery.jobs().list(projectId).setAllUsers(Option.ALL_USERS.getBoolean(options)).setFields(Option.FIELDS.getString(options)).setStateFilter(Option.STATE_FILTER.<List<String>>get(options)).setMaxResults(Option.MAX_RESULTS.getLong(options)).setPageToken(Option.PAGE_TOKEN.getString(options)).setProjection(DEFAULT_PROJECTION).execute();
        Iterable<JobList.Jobs> jobs = jobsList.getJobs();
        return Tuple.of(jobsList.getNextPageToken(), Iterables.transform(jobs != null ? jobs : ImmutableList.<JobList.Jobs>of(), new Function<JobList.Jobs, Job>() {

            @Override
            public Job apply(JobList.Jobs jobPb) {
                JobStatus statusPb = jobPb.getStatus() != null ? jobPb.getStatus() : new JobStatus();
                if (statusPb.getState() == null) {
                    statusPb.setState(jobPb.getState());
                }
                if (statusPb.getErrorResult() == null) {
                    statusPb.setErrorResult(jobPb.getErrorResult());
                }
                return new Job().setConfiguration(jobPb.getConfiguration()).setId(jobPb.getId()).setJobReference(jobPb.getJobReference()).setKind(jobPb.getKind()).setStatistics(jobPb.getStatistics()).setStatus(statusPb).setUserEmail(jobPb.getUserEmail());
            }
        }));
    } catch (IOException ex) {
        throw translate(ex);
    }
}
Also used : JobStatus(com.google.api.services.bigquery.model.JobStatus) Function(com.google.common.base.Function) JobList(com.google.api.services.bigquery.model.JobList) IOException(java.io.IOException) Job(com.google.api.services.bigquery.model.Job)

Example 10 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)

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