/** Example of loading a parquet file from GCS to a table. */ public void loadTableGcsParquet(String datasetName) throws InterruptedException { // [START bigquery_load_table_gcs_parquet] String sourceUri = "gs://cloud-samples-data/bigquery/us-states/us-states.parquet"; TableId tableId = TableId.of(datasetName, "us_states"); LoadJobConfiguration configuration = LoadJobConfiguration.builder(tableId, sourceUri) .setFormatOptions(FormatOptions.parquet()) .build(); // Load the table Job loadJob = bigquery.create(JobInfo.of(configuration)); loadJob = loadJob.waitFor(); // Check the table StandardTableDefinition destinationTable = bigquery.getTable(tableId).getDefinition(); System.out.println("State: " + loadJob.getStatus().getState()); System.out.printf("Loaded %d rows.\n", destinationTable.getNumRows()); // [END bigquery_load_table_gcs_parquet] }
@Test public void testIsDone_NotExists() throws Exception { initializeExpectedJob(1); BigQuery.JobOption[] expectedOptions = {BigQuery.JobOption.fields(BigQuery.JobField.STATUS)}; expect(bigquery.getOptions()).andReturn(mockOptions); expect(bigquery.getJob(JOB_INFO.getJobId(), expectedOptions)).andReturn(null); replay(bigquery); initializeJob(); assertTrue(job.isDone()); }
@Override void run(BigQuery bigquery, JobInfo job) throws Exception { System.out.println("Creating job"); Job startedJob = bigquery.create(job); while (!startedJob.isDone()) { System.out.println("Waiting for job " + startedJob.getJobId().getJob() + " to complete"); Thread.sleep(1000L); } startedJob = startedJob.reload(); if (startedJob.getStatus().getError() == null) { System.out.println("Job " + startedJob.getJobId().getJob() + " succeeded"); } else { System.out.println("Job " + startedJob.getJobId().getJob() + " failed"); System.out.println("Error: " + startedJob.getStatus().getError()); } } }
/** * Sends a job cancel request. * * <p>Example of cancelling a job. * * <pre>{@code * if (job.cancel()) { * return true; // job successfully cancelled * } else { * // job not found * } * }</pre> * * @return {@code true} if cancel request was sent successfully, {@code false} if job was not * found * @throws BigQueryException upon failure */ public boolean cancel() { checkNotDryRun("cancel"); return bigquery.cancel(getJobId()); }
@Override public TableResult query(QueryJobConfiguration configuration, JobOption... options) throws InterruptedException, JobException { Job.checkNotDryRun(configuration, "query"); return create(JobInfo.of(configuration), options).getQueryResults(); }
/** * Fetches current job's latest information. Returns {@code null} if the job does not exist. * * <p>Example of reloading all fields until job status is DONE. * * <pre>{@code * while (!JobStatus.State.DONE.equals(job.getStatus().getState())) { * Thread.sleep(1000L); * job = job.reload(); * } * }</pre> * * <p>Example of reloading status field until job status is DONE. * * <pre>{@code * while (!JobStatus.State.DONE.equals(job.getStatus().getState())) { * Thread.sleep(1000L); * job = job.reload(BigQuery.JobOption.fields(BigQuery.JobField.STATUS)); * } * }</pre> * * @param options job options * @return a {@code Job} object with latest information or {@code null} if not found * @throws BigQueryException upon failure */ public Job reload(JobOption... options) { checkNotDryRun("reload"); return bigquery.getJob(getJobId(), options); }
@Test public void testWaitFor_Null() throws InterruptedException { initializeExpectedJob(1); BigQuery.JobOption[] expectedOptions = {BigQuery.JobOption.fields(BigQuery.JobField.STATUS)}; expect(bigquery.getOptions()).andReturn(mockOptions); expect(mockOptions.getClock()).andReturn(CurrentMillisClock.getDefaultClock()); expect(bigquery.getJob(JOB_INFO.getJobId(), expectedOptions)).andReturn(null); replay(bigquery, mockOptions); initializeJob(); assertNull(job.waitFor(TEST_RETRY_OPTIONS)); verify(mockOptions); }
@Test public void testGetJob() { EasyMock.expect(bigqueryRpcMock.getJob(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) .andReturn(COMPLETE_COPY_JOB.toPb()); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Job job = bigquery.getJob(JOB); assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_COPY_JOB)), job); }
.build(); JobInfo jobInfo = JobInfo.newBuilder(jobConfig) .setJobId(JOB_ID) .setStatistics(jobStatistics) JobStatus status = createStrictMock(JobStatus.class); expect(bigquery.getOptions()).andReturn(mockOptions); expect(mockOptions.getClock()).andReturn(CurrentMillisClock.getDefaultClock()).times(2); Job completedJob = expectedJob.toBuilder().setStatus(status).build(); .build(); expect(bigquery.getQueryResults(jobInfo.getJobId(), Job.DEFAULT_QUERY_WAIT_OPTIONS)) .andReturn(completedQuery); expect(bigquery.getJob(JOB_INFO.getJobId())).andReturn(completedJob); expect(bigquery.getQueryResults(jobInfo.getJobId(), Job.DEFAULT_QUERY_WAIT_OPTIONS)) .andReturn(completedQuery); replay(status, bigquery, mockOptions); initializeJob(jobInfo); assertThat(job.waitFor(TEST_RETRY_OPTIONS)).isSameAs(completedJob); assertThat(job.getQueryResults().iterateAll()).hasSize(0); verify(status, mockOptions);
@Test public void testWriteClosed() throws IOException { expect( bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andReturn(UPLOAD_ID); Capture<byte[]> capturedBuffer = Capture.newInstance(); expect( bigqueryRpcMock.write( eq(UPLOAD_ID), capture(capturedBuffer), eq(0), eq(0L), eq(0), eq(true))) .andReturn(job.toPb()); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); writer.close(); assertEquals(job, writer.getJob()); try { writer.write(ByteBuffer.allocate(MIN_CHUNK_SIZE)); fail("Expected TableDataWriteChannel write to throw IOException"); } catch (IOException ex) { // expected } }
@Test public void testExtractDataUris() throws Exception { initializeExpectedTable(2); expect(bigquery.getOptions()).andReturn(mockOptions); Job expectedJob = new Job(serviceMockReturnsOptions, new JobInfo.BuilderImpl(EXTRACT_JOB_INFO)); expect(bigquery.create(EXTRACT_JOB_INFO)).andReturn(expectedJob); replay(bigquery); initializeTable(); Job job = table.extract("CSV", ImmutableList.of("URI")); assertSame(expectedJob, job); }
@Test public void testIsDone_True() throws Exception { initializeExpectedJob(2); BigQuery.JobOption[] expectedOptions = {BigQuery.JobOption.fields(BigQuery.JobField.STATUS)}; JobStatus status = createStrictMock(JobStatus.class); expect(status.getState()).andReturn(JobStatus.State.DONE); expect(bigquery.getOptions()).andReturn(mockOptions); expect(bigquery.getJob(JOB_INFO.getJobId(), expectedOptions)) .andReturn(expectedJob.toBuilder().setStatus(status).build()); replay(status, bigquery); initializeJob(); assertTrue(job.isDone()); verify(status); }
@Before public void setUp() { rpcFactoryMock = createMock(BigQueryRpcFactory.class); bigqueryRpcMock = createMock(BigQueryRpc.class); bigqueryFactoryMock = createMock(BigQueryFactory.class); bigqueryMock = createMock(BigQuery.class); expect(bigqueryMock.getOptions()).andReturn(options).anyTimes(); replay(bigqueryMock); job = new Job(bigqueryMock, new JobInfo.BuilderImpl(JOB_INFO)); expect(rpcFactoryMock.create(anyObject(BigQueryOptions.class))).andReturn(bigqueryRpcMock); expect(bigqueryFactoryMock.create(anyObject(BigQueryOptions.class))) .andReturn(bigqueryMock) .anyTimes(); replay(rpcFactoryMock, bigqueryFactoryMock); options = BigQueryOptions.newBuilder() .setProjectId("projectid") .setServiceRpcFactory(rpcFactoryMock) .setServiceFactory(bigqueryFactoryMock) .build(); }
@Test public void testReload() throws Exception { initializeExpectedJob(4); JobInfo updatedInfo = JOB_INFO.toBuilder().setEtag("etag").build(); Job expectedJob = new Job(serviceMockReturnsOptions, new JobInfo.BuilderImpl(updatedInfo)); expect(bigquery.getOptions()).andReturn(mockOptions); expect(bigquery.getJob(JOB_INFO.getJobId())).andReturn(expectedJob); replay(bigquery); initializeJob(); Job updatedJob = job.reload(); compareJob(expectedJob, updatedJob); }
@Test public void testListJobs() { bigquery = options.getService(); ImmutableList<Job> jobList = ImmutableList.of( new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_QUERY_JOB)), new Job(bigquery, new JobInfo.BuilderImpl(COMPLETE_LOAD_JOB))); Tuple<String, Iterable<com.google.api.services.bigquery.model.Job>> result = Tuple.of( CURSOR, Iterables.transform( jobList, new Function<Job, com.google.api.services.bigquery.model.Job>() { @Override public com.google.api.services.bigquery.model.Job apply(Job job) { return job.toPb(); } })); EasyMock.expect(bigqueryRpcMock.listJobs(PROJECT, EMPTY_RPC_OPTIONS)).andReturn(result); EasyMock.replay(bigqueryRpcMock); Page<Job> page = bigquery.listJobs(); assertEquals(CURSOR, page.getNextPageToken()); assertArrayEquals(jobList.toArray(), Iterables.toArray(page.getValues(), Job.class)); }
private void initializeExpectedJob(int optionsCalls, JobInfo jobInfo) { expect(serviceMockReturnsOptions.getOptions()).andReturn(mockOptions).times(optionsCalls); replay(serviceMockReturnsOptions); bigquery = createStrictMock(BigQuery.class); expectedJob = new Job(serviceMockReturnsOptions, new JobInfo.BuilderImpl(jobInfo)); }
@Test public void testCreateJobWithProjectId() { JobInfo jobInfo = JobInfo.newBuilder(QUERY_JOB_CONFIGURATION.setProjectId(OTHER_PROJECT)) .setJobId(JobId.of(OTHER_PROJECT, JOB)) .build(); Capture<Map<BigQueryRpc.Option, Object>> capturedOptions = Capture.newInstance(); EasyMock.expect(bigqueryRpcMock.create(eq(jobInfo.toPb()), capture(capturedOptions))) .andReturn(jobInfo.toPb()); EasyMock.replay(bigqueryRpcMock); BigQueryOptions bigQueryOptions = createBigQueryOptionsForProject(OTHER_PROJECT, rpcFactoryMock); bigquery = bigQueryOptions.getService(); Job job = bigquery.create(jobInfo, JOB_OPTION_FIELDS); assertEquals(new Job(bigquery, new JobInfo.BuilderImpl(jobInfo)), job); String selector = (String) capturedOptions.getValue().get(JOB_OPTION_FIELDS.getRpcOption()); assertTrue(selector.contains("jobReference")); assertTrue(selector.contains("configuration")); assertTrue(selector.contains("user_email")); assertEquals(37, selector.length()); }
@Test public void testWaitForAndGetQueryResults_Unsupported() throws InterruptedException { initializeExpectedJob(1); expect(bigquery.getOptions()).andReturn(mockOptions); replay(bigquery, mockOptions); initializeJob(); thrown.expect(UnsupportedOperationException.class); job.getQueryResults(); }
.setSchema(TABLE_SCHEMA.toPb()); EasyMock.expect( bigqueryRpcMock.create( JOB_INFO.toPb(), Collections.<BigQueryRpc.Option, Object>emptyMap())) .andReturn(jobResponsePb); optionMap.put(pageSizeOption.getRpcOption(), pageSizeOption.getValue()); EasyMock.expect( bigqueryRpcMock.getQueryResults( PROJECT, JOB, null, BigQueryImpl.optionMap(Job.DEFAULT_QUERY_WAIT_OPTIONS))) .andReturn(responsePb); EasyMock.expect(bigqueryRpcMock.listTableData(PROJECT, DATASET, TABLE, optionMap)) .andReturn( new TableDataList() .setPageToken("") EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); Job job = bigquery.create(JobInfo.of(queryJob, QUERY_JOB_CONFIGURATION_FOR_QUERY)); TableResult result = job.getQueryResults(pageSizeOption); assertThat(result.getSchema()).isEqualTo(TABLE_SCHEMA); assertThat(result.getTotalRows()).isEqualTo(1);