@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()); }
/** Example of creating a channel with which to write to a table. */ // [TARGET writer(WriteChannelConfiguration)] // [VARIABLE "my_dataset_name"] // [VARIABLE "my_table_name"] // [VARIABLE "StringValue1\nStringValue2\n"] public long writeToTable(String datasetName, String tableName, String csvData) throws IOException, InterruptedException, TimeoutException { // [START ] TableId tableId = TableId.of(datasetName, tableName); WriteChannelConfiguration writeChannelConfiguration = WriteChannelConfiguration.newBuilder(tableId).setFormatOptions(FormatOptions.csv()).build(); TableDataWriteChannel writer = bigquery.writer(writeChannelConfiguration); // Write data to writer try { writer.write(ByteBuffer.wrap(csvData.getBytes(Charsets.UTF_8))); } finally { writer.close(); } // Get load job Job job = writer.getJob(); job = job.waitFor(); LoadStatistics stats = job.getStatistics(); return stats.getOutputRows(); // [END ] }
@Override protected StateImpl.Builder stateBuilder() { return StateImpl.builder(getOptions(), getEntity(), getUploadId(), job); }
@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 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()); }
ByteBuffer buffer1 = randomBuffer(DEFAULT_CHUNK_SIZE); ByteBuffer buffer2 = randomBuffer(DEFAULT_CHUNK_SIZE); writer = new TableDataWriteChannel(options, JOB_INFO.getJobId(), LOAD_CONFIGURATION); assertEquals(DEFAULT_CHUNK_SIZE, writer.write(buffer1)); assertArrayEquals(buffer1.array(), capturedBuffer.getValues().get(0)); assertEquals(new Long(0L), capturedPosition.getValues().get(0)); assertNull(writer.getJob()); RestorableState<WriteChannel> writerState = writer.capture(); WriteChannel restoredWriter = writerState.restore(); assertEquals(DEFAULT_CHUNK_SIZE, restoredWriter.write(buffer2));
@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()); }
@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()); }
@Override public WriteChannel restore() { TableDataWriteChannel channel = new TableDataWriteChannel(serviceOptions, entity, uploadId); channel.restore(this); channel.setJob(job); return channel; }
@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()); }
/** Example of writing a local file to a table. */ // [TARGET writer(WriteChannelConfiguration)] // [VARIABLE "my_dataset_name"] // [VARIABLE "my_table_name"] // [VARIABLE FileSystems.getDefault().getPath(".", "my-data.csv")] // [VARIABLE "us"] public long writeFileToTable(String datasetName, String tableName, Path csvPath, String location) throws IOException, InterruptedException, TimeoutException { // [START bigquery_load_from_file] TableId tableId = TableId.of(datasetName, tableName); WriteChannelConfiguration writeChannelConfiguration = WriteChannelConfiguration.newBuilder(tableId).setFormatOptions(FormatOptions.csv()).build(); // The location must be specified; other fields can be auto-detected. JobId jobId = JobId.newBuilder().setLocation(location).build(); TableDataWriteChannel writer = bigquery.writer(jobId, writeChannelConfiguration); // Write data to writer try (OutputStream stream = Channels.newOutputStream(writer)) { Files.copy(csvPath, stream); } // Get load job Job job = writer.getJob(); job = job.waitFor(); LoadStatistics stats = job.getStatistics(); return stats.getOutputRows(); // [END bigquery_load_from_file] }
bigquery.writer( JobId.newBuilder().setLocation(location).build(), writeChannelConfiguration)) { writer.write(ByteBuffer.wrap("foo".getBytes()));
@Override protected void flushBuffer(final int length, final boolean last) { try { com.google.api.services.bigquery.model.Job jobPb = runWithRetries( new Callable<com.google.api.services.bigquery.model.Job>() { @Override public com.google.api.services.bigquery.model.Job call() { return getOptions() .getBigQueryRpcV2() .write(getUploadId(), getBuffer(), 0, getPosition(), length, last); } }, getOptions().getRetrySettings(), BigQueryImpl.EXCEPTION_HANDLER, getOptions().getClock()); job = jobPb != null ? Job.fromPb(getOptions().getService(), jobPb) : null; } catch (RetryHelper.RetryHelperException e) { throw BigQueryException.translateAndThrow(e); } }
@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 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()); }
@Override public TableDataWriteChannel writer( JobId jobId, WriteChannelConfiguration writeChannelConfiguration) { return new TableDataWriteChannel( getOptions(), jobId.setProjectId(getOptions().getProjectId()), writeChannelConfiguration.setProjectId(getOptions().getProjectId())); }
@Override protected void flushBuffer(final int length, final boolean last) { try { com.google.api.services.bigquery.model.Job jobPb = runWithRetries( new Callable<com.google.api.services.bigquery.model.Job>() { @Override public com.google.api.services.bigquery.model.Job call() { return getOptions() .getBigQueryRpcV2() .write(getUploadId(), getBuffer(), 0, getPosition(), length, last); } }, getOptions().getRetrySettings(), BigQueryImpl.EXCEPTION_HANDLER, getOptions().getClock()); job = jobPb != null ? Job.fromPb(getOptions().getService(), jobPb) : null; } catch (RetryHelper.RetryHelperException e) { throw BigQueryException.translateAndThrow(e); } }
@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()); }
/** Example of creating a channel with which to write to a table. */ // [TARGET writer(JobId, WriteChannelConfiguration)] // [VARIABLE "my_dataset_name"] // [VARIABLE "my_table_name"] // [VARIABLE "StringValue1\nStringValue2\n"] // [VARIABLE "us"] public long writeToTableLocation( String datasetName, String tableName, String csvData, String location) throws IOException, InterruptedException, TimeoutException { // [START ] TableId tableId = TableId.of(datasetName, tableName); WriteChannelConfiguration writeChannelConfiguration = WriteChannelConfiguration.newBuilder(tableId).setFormatOptions(FormatOptions.csv()).build(); // The location must be specified; other fields can be auto-detected. JobId jobId = JobId.newBuilder().setLocation(location).build(); TableDataWriteChannel writer = bigquery.writer(jobId, writeChannelConfiguration); // Write data to writer try { writer.write(ByteBuffer.wrap(csvData.getBytes(Charsets.UTF_8))); } finally { writer.close(); } // Get load job Job job = writer.getJob(); job = job.waitFor(); LoadStatistics stats = job.getStatistics(); return stats.getOutputRows(); // [END ] }