IndirectBigQueryOutputCommitter extends ForwardingBigQueryFileOutputCommitter { @Override public void commitJob(JobContext context) throws IOException { super.commitJob(context); Configuration conf = context.getConfiguration(); TableReference destTable = BigQueryOutputConfiguration.getTableReference(conf); String jobProjectId = BigQueryOutputConfiguration.getJobProjectId(conf); String writeDisposition = BigQueryOutputConfiguration.getWriteDisposition(conf); String createDisposition = BigQueryOutputConfiguration.getCreateDisposition(conf); Optional<BigQueryTableSchema> destSchema = BigQueryOutputConfiguration.getTableSchema(conf); Optional<BigQueryTimePartitioning> timePartitioning = BigQueryOutputConfiguration.getTablePartitioning(conf); String kmsKeyName = BigQueryOutputConfiguration.getKmsKeyName(conf); BigQueryFileFormat outputFileFormat = BigQueryOutputConfiguration.getFileFormat(conf); List<String> sourceUris = getOutputFileURIs(); try { getBigQueryHelper() .importFromGcs( jobProjectId, destTable, destSchema.isPresent() ? destSchema.get().get() : null, timePartitioning.isPresent() ? timePartitioning.get().get() : null, kmsKeyName, outputFileFormat, createDisposition, writeDisposition, sourceUris, true); } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new IOException("Failed to import GCS into BigQuery", e); } cleanup(context); } IndirectBigQueryOutputCommitter(TaskAttemptContext context, OutputCommitter delegate); @Override void commitJob(JobContext context); @Override void abortJob(JobContext context, State state);  }
@Test public void testCommitJob() throws Exception { generateSampleFiles(); committer.commitJob(job); @SuppressWarnings({"rawtypes", "unchecked", "cast"}) Class<List<String>> listClass = (Class<List<String>>) (Class) List.class; ArgumentCaptor<List<String>> gcsOutputFileCaptor = ArgumentCaptor.forClass(listClass); verify(mockBigQueryHelper) .importFromGcs( eq(TEST_PROJECT_ID), eq(outputTableRef), eq(TEST_TABLE_SCHEMA.get()), eq(TEST_TIME_PARTITIONING.get()), eq(TEST_KMS_KEY_NAME), eq(TEST_FILE_FORMAT), eq(TEST_CREATE_DISPOSITION), eq(TEST_WRITE_DISPOSITION), gcsOutputFileCaptor.capture(), eq(true)); verify(mockCommitter).commitJob(eq(job)); assertThat(gcsOutputFileCaptor.getValue()).contains(TEST_OUTPUT_FILE_STRING); }
@SuppressWarnings("unchecked") @Test public void testCommitJobInterrupt() throws Exception { generateSampleFiles(); InterruptedException helperInterruptedException = new InterruptedException("Test exception"); doThrow(helperInterruptedException) .when(mockBigQueryHelper) .importFromGcs( any(String.class), any(TableReference.class), any(TableSchema.class), any(TimePartitioning.class), anyString(), any(BigQueryFileFormat.class), any(String.class), any(String.class), any(List.class), eq(true)); IOException thrown = assertThrows(IOException.class, () -> committer.commitJob(job)); assertThat(thrown).hasCauseThat().isEqualTo(helperInterruptedException); verify(mockBigQueryHelper) .importFromGcs( eq(TEST_PROJECT_ID), eq(outputTableRef), eq(TEST_TABLE_SCHEMA.get()), eq(TEST_TIME_PARTITIONING.get()), eq(TEST_KMS_KEY_NAME), eq(TEST_FILE_FORMAT), eq(TEST_CREATE_DISPOSITION), eq(TEST_WRITE_DISPOSITION), any(List.class), eq(true)); verify(mockCommitter).commitJob(eq(job)); }