@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 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 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 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 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);
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 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()); }
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"));
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());
@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()); }
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"));
@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()); } }));
@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()); }
@Override public String call() { return options .getBigQueryRpcV2() .open( new com.google.api.services.bigquery.model.Job() .setConfiguration(writeChannelConfiguration.toPb()) .setJobReference(jobId.toPb())); } },
@Override public void startQueryJob(JobReference jobRef, JobConfigurationQuery query) { 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)); } }