@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 testToPbAndFromPb() { assertNull(LOAD_CONFIGURATION_CSV.toPb().getLoad().getSourceUris()); compareLoadConfiguration( LOAD_CONFIGURATION_CSV, WriteChannelConfiguration.fromPb(LOAD_CONFIGURATION_CSV.toPb())); WriteChannelConfiguration configuration = WriteChannelConfiguration.of(TABLE_ID); compareLoadConfiguration(configuration, WriteChannelConfiguration.fromPb(configuration.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()); }
@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 } }
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(CaptureType.ALL);
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();
@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()); }
TableDataWriteChannel(BigQueryOptions options, WriteChannelConfiguration writeChannelConfiguration) { this(options, writeChannelConfiguration, options.rpc().open(writeChannelConfiguration.toPb())); }
@Override public String call() { return options .getBigQueryRpcV2() .open( new com.google.api.services.bigquery.model.Job() .setConfiguration(writeChannelConfiguration.toPb()) .setJobReference(jobId.toPb())); } },