@Override public String call() { return options .getBigQueryRpcV2() .open( new com.google.api.services.bigquery.model.Job() .setConfiguration(writeChannelConfiguration.toPb()) .setJobReference(jobId.toPb())); } },
@Test public void testCreateNonRetryableError() { expect( bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andThrow(new RuntimeException()); replay(bigqueryRpcMock); thrown.expect(RuntimeException.class); new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); }
@Test public void testWriteWithoutFlush() 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); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); assertEquals(MIN_CHUNK_SIZE, writer.write(ByteBuffer.allocate(MIN_CHUNK_SIZE))); assertNull(writer.getJob()); }
@Test public void testCreate() { expect( bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andReturn(UPLOAD_ID); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); assertTrue(writer.isOpen()); assertNull(writer.getJob()); }
@Test public void testStateEquals() { expect( bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andReturn(UPLOAD_ID) .times(2); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); // avoid closing when you don't want partial writes upon failure @SuppressWarnings("resource") WriteChannel writer2 = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); RestorableState<WriteChannel> state = writer.capture(); RestorableState<WriteChannel> state2 = writer2.capture(); assertEquals(state, state2); assertEquals(state.hashCode(), state2.hashCode()); assertEquals(state.toString(), state2.toString()); }
@Test public void testCreateRetryableError() { BigQueryException exception = new BigQueryException(new SocketException("Socket closed")); expect( bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andThrow(exception); expect( bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andReturn(UPLOAD_ID); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); assertTrue(writer.isOpen()); assertNull(writer.getJob()); }
Job toPb() { Job jobPb = new Job(); jobPb.setEtag(etag); jobPb.setId(generatedId); jobPb.setSelfLink(selfLink); jobPb.setUserEmail(userEmail); if (jobId != null) { jobPb.setJobReference(jobId.toPb()); } if (status != null) { jobPb.setStatus(status.toPb()); } if (statistics != null) { jobPb.setStatistics(statistics.toPb()); } if (configuration != null) { jobPb.setConfiguration(configuration.toPb()); } return jobPb; }
@Test public void testGetQueryResults() { JobId queryJob = JobId.of(JOB); GetQueryResultsResponse responsePb = new GetQueryResultsResponse() .setEtag("etag") .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true) .setCacheHit(false) .setPageToken(CURSOR) .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)); EasyMock.expect(bigqueryRpcMock.getQueryResults(PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) .andReturn(responsePb); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); QueryResponse response = bigquery.getQueryResults(queryJob); assertEquals(true, response.getCompleted()); assertEquals(null, response.getSchema()); }
@Test public void testGetQueryResultsWithProject() { JobId queryJob = JobId.of(OTHER_PROJECT, JOB); GetQueryResultsResponse responsePb = new GetQueryResultsResponse() .setEtag("etag") .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true) .setCacheHit(false) .setPageToken(CURSOR) .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)); EasyMock.expect(bigqueryRpcMock.getQueryResults(OTHER_PROJECT, JOB, null, EMPTY_RPC_OPTIONS)) .andReturn(responsePb); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); QueryResponse response = bigquery.getQueryResults(queryJob); assertTrue(response.getCompleted()); assertEquals(null, response.getSchema()); }
@Test public void testGetQueryResultsWithOptions() { JobId queryJob = JobId.of(PROJECT, JOB); GetQueryResultsResponse responsePb = new GetQueryResultsResponse() .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true) .setCacheHit(false) .setPageToken(CURSOR) .setTotalBytesProcessed(42L) .setTotalRows(BigInteger.valueOf(1L)); EasyMock.expect(bigqueryRpcMock.getQueryResults(PROJECT, JOB, null, QUERY_RESULTS_OPTIONS)) .andReturn(responsePb); EasyMock.replay(bigqueryRpcMock); bigquery = options.getService(); QueryResponse response = bigquery.getQueryResults( queryJob, QUERY_RESULTS_OPTION_TIME, QUERY_RESULTS_OPTION_INDEX, QUERY_RESULTS_OPTION_PAGE_SIZE, QUERY_RESULTS_OPTION_PAGE_TOKEN); assertEquals(true, response.getCompleted()); assertEquals(null, response.getSchema()); }
new com.google.api.services.bigquery.model.Job() .setConfiguration(QUERY_JOB_CONFIGURATION_FOR_QUERY.toPb()) .setJobReference(queryJob.toPb()) .setId(JOB); jobResponsePb1.getConfiguration().getQuery().setDestinationTable(TABLE_ID.toPb()); new GetQueryResultsResponse().setJobReference(queryJob.toPb()).setJobComplete(false); .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true)
bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andReturn(UPLOAD_ID);
@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 } }
bigqueryRpcMock.open( new com.google.api.services.bigquery.model.Job() .setJobReference(JOB_INFO.getJobId().toPb()) .setConfiguration(LOAD_CONFIGURATION.toPb()))) .andReturn(UPLOAD_ID);
new com.google.api.services.bigquery.model.Job() .setConfiguration(QUERY_JOB_CONFIGURATION_FOR_QUERY.toPb()) .setJobReference(queryJob.toPb()) .setId(JOB) .setStatus(new com.google.api.services.bigquery.model.JobStatus().setState("DONE")); GetQueryResultsResponse responsePb = new GetQueryResultsResponse() .setJobReference(queryJob.toPb()) .setRows(ImmutableList.of(TABLE_ROW)) .setJobComplete(true)
@Test public void testWriteWithFlush() 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(CUSTOM_CHUNK_SIZE), eq(false))) .andReturn(null); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); writer.setChunkSize(CUSTOM_CHUNK_SIZE); ByteBuffer buffer = randomBuffer(CUSTOM_CHUNK_SIZE); assertEquals(CUSTOM_CHUNK_SIZE, writer.write(buffer)); assertArrayEquals(buffer.array(), capturedBuffer.getValue()); assertNull(writer.getJob()); }
@Test public void testCloseWithoutFlush() 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); assertTrue(writer.isOpen()); writer.close(); assertArrayEquals(new byte[0], capturedBuffer.getValue()); assertTrue(!writer.isOpen()); assertEquals(job, writer.getJob()); }
@Test public void testCloseWithFlush() 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(); ByteBuffer buffer = randomBuffer(MIN_CHUNK_SIZE); expect( bigqueryRpcMock.write( eq(UPLOAD_ID), capture(capturedBuffer), eq(0), eq(0L), eq(MIN_CHUNK_SIZE), eq(true))) .andReturn(job.toPb()); replay(bigqueryRpcMock); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); assertTrue(writer.isOpen()); writer.write(buffer); writer.close(); assertEquals(DEFAULT_CHUNK_SIZE, capturedBuffer.getValue().length); assertArrayEquals(buffer.array(), Arrays.copyOf(capturedBuffer.getValue(), MIN_CHUNK_SIZE)); assertTrue(!writer.isOpen()); assertEquals(job, writer.getJob()); }
@Test public void testSaveAndRestoreClosed() 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()); RestorableState<WriteChannel> writerState = writer.capture(); RestorableState<WriteChannel> expectedWriterState = TableDataWriteChannel.StateImpl.builder(options, LOAD_CONFIGURATION, UPLOAD_ID, job) .setBuffer(null) .setChunkSize(DEFAULT_CHUNK_SIZE) .setIsOpen(false) .setPosition(0) .build(); WriteChannel restoredWriter = writerState.restore(); assertArrayEquals(new byte[0], capturedBuffer.getValue()); assertEquals(expectedWriterState, restoredWriter.capture()); }