BigQueryHelper { public Job insertJobOrFetchDuplicate(String projectId, Job job) throws IOException { checkArgument( job.getJobReference() != null && job.getJobReference().getJobId() != null, "Require non-null JobReference and JobId inside; getJobReference() == '%s'", job.getJobReference()); Job response; try { response = service.jobs().insert(projectId, job).execute(); logger.atFine().log("Successfully inserted job '%s'. Response: '%s'", job, response); } catch (IOException insertJobException) { if (errorExtractor.itemAlreadyExists(insertJobException)) { logger.atInfo().withCause(insertJobException).log( "Fetching existing job after catching exception for duplicate jobId '%s'", job.getJobReference().getJobId()); try { response = service .jobs() .get(projectId, job.getJobReference().getJobId()) .setLocation(job.getJobReference().getLocation()) .execute(); } catch (IOException getJobException) { getJobException.addSuppressed(insertJobException); throw new IOException( String.format("Failed to get duplicate job '%s'", job), getJobException); } } else { throw new IOException( String.format("Unhandled exception trying to insert job '%s'", job), insertJobException); } } checkJobIdEquality(job, response); return response; } BigQueryHelper(Bigquery service); Bigquery getRawBigquery(); void importFederatedFromGcs(
      String projectId,
      TableReference tableRef,
      @Nullable TableSchema schema,
      BigQueryFileFormat sourceFormat,
      List<String> gcsPaths); void importFromGcs(
      String projectId,
      TableReference tableRef,
      @Nullable TableSchema schema,
      @Nullable TimePartitioning timePartitioning,
      @Nullable String kmsKeyName,
      BigQueryFileFormat sourceFormat,
      String createDisposition,
      String writeDisposition,
      List<String> gcsPaths,
      boolean awaitCompletion); void exportBigQueryToGcs(
      String projectId, TableReference tableRef, List<String> gcsPaths, boolean awaitCompletion); boolean tableExists(TableReference tableRef); Table getTable(TableReference tableRef); JobReference createJobReference(
      String projectId, String jobIdPrefix, @Nullable String location); void checkJobIdEquality(Job expected, Job actual); Job insertJobOrFetchDuplicate(String projectId, Job job); static final String BIGQUERY_JOB_ID_PATTERN; static final int BIGQUERY_JOB_ID_MAX_LENGTH; }
@Test public void testInsertJobOrFetchDuplicateBasicInsert() throws IOException { when(mockBigqueryJobsInsert.execute()).thenReturn(jobHandle); assertThat(helper.insertJobOrFetchDuplicate(jobProjectId, jobHandle)).isEqualTo(jobHandle); verify(mockBigquery, times(1)).jobs(); ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); verify(mockBigqueryJobs, times(1)).insert(eq(jobProjectId), jobCaptor.capture()); Job job = jobCaptor.getValue(); assertThat(job).isEqualTo(jobHandle); verify(mockBigqueryJobsInsert, times(1)).execute(); }
@Test public void testInsertJobOrFetchDuplicateAlreadyExistsException() throws IOException { IOException fakeConflictException = new IOException("fake 409 conflict"); when(mockBigqueryJobsInsert.execute()).thenThrow(fakeConflictException); when(mockErrorExtractor.itemAlreadyExists(any(IOException.class))).thenReturn(true); when(mockBigqueryJobsGet.execute()).thenReturn(jobHandle); assertThat(helper.insertJobOrFetchDuplicate(jobProjectId, jobHandle)).isEqualTo(jobHandle); verify(mockBigquery, times(2)).jobs(); ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); verify(mockBigqueryJobs, times(1)).insert(eq(jobProjectId), jobCaptor.capture()); Job job = jobCaptor.getValue(); assertThat(job).isEqualTo(jobHandle); verify(mockBigqueryJobsInsert, times(1)).execute(); verify(mockBigqueryJobs, times(1)).get(eq(jobProjectId), eq(jobId)); verify(mockBigqueryJobsGet, times(1)).setLocation(eq(null)); verify(mockBigqueryJobsGet, times(1)).execute(); verify(mockErrorExtractor).itemAlreadyExists(eq(fakeConflictException)); }
@Test public void testInsertJobOrFetchDuplicateUnhandledException() throws IOException { IOException unhandledException = new IOException("unhandled exception"); when(mockBigqueryJobsInsert.execute()).thenThrow(unhandledException); when(mockErrorExtractor.itemAlreadyExists(any(IOException.class))).thenReturn(false); IOException thrown = assertThrows( IOException.class, () -> helper.insertJobOrFetchDuplicate(jobProjectId, jobHandle)); assertThat(thrown).hasMessageThat().contains(jobHandle.toString()); assertThat(thrown).hasCauseThat().isEqualTo(unhandledException); verify(mockBigquery, times(1)).jobs(); ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); verify(mockBigqueryJobs, times(1)).insert(eq(jobProjectId), jobCaptor.capture()); Job job = jobCaptor.getValue(); assertThat(job).isEqualTo(jobHandle); verify(mockBigqueryJobsInsert, times(1)).execute(); verify(mockErrorExtractor).itemAlreadyExists(eq(unhandledException)); }