BigQueryHelper { public 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) throws IOException, InterruptedException { logger.atInfo().log( "Importing into table '%s' from %s paths; path[0] is '%s'; awaitCompletion: %s;" + " timePartitioning: %s", lazy(() -> BigQueryStrings.toString(tableRef)), gcsPaths.size(), gcsPaths.isEmpty() ? "(empty)" : gcsPaths.get(0), awaitCompletion, timePartitioning); JobConfigurationLoad loadConfig = new JobConfigurationLoad(); loadConfig.setSchema(schema); loadConfig.setSourceFormat(sourceFormat.getFormatIdentifier()); loadConfig.setSourceUris(gcsPaths); loadConfig.setDestinationTable(tableRef); loadConfig.setTimePartitioning(timePartitioning); loadConfig.setCreateDisposition(createDisposition); loadConfig.setWriteDisposition(writeDisposition); if (!Strings.isNullOrEmpty(kmsKeyName)) { loadConfig.setDestinationEncryptionConfiguration( new EncryptionConfiguration().setKmsKeyName(kmsKeyName)); } if (schema == null) { logger.atInfo().log("No import schema provided, auto detecting schema."); loadConfig.setAutodetect(true); } else { logger.atInfo().log("Using provided import schema '%s'.", schema); } JobConfiguration config = new JobConfiguration(); config.setLoad(loadConfig); Dataset dataset; try { dataset = service.datasets().get(tableRef.getProjectId(), tableRef.getDatasetId()).execute(); } catch (IOException ioe) { throw new IOException( String.format( "Failed to get dataset '%s' in project '%s' for table '%s'", tableRef.getDatasetId(), tableRef.getProjectId(), tableRef), ioe); } JobReference jobReference = createJobReference(projectId, "direct-bigqueryhelper-import", dataset.getLocation()); Job job = new Job(); job.setConfiguration(config); job.setJobReference(jobReference); insertJobOrFetchDuplicate(projectId, job); if (awaitCompletion) { BigQueryUtils.waitForJobCompletion(getRawBigquery(), projectId, jobReference, () -> {}); } } 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 testImportBigQueryFromGcs() throws Exception { when(mockBigqueryTablesGet.execute()).thenReturn(fakeTable); final ArgumentCaptor<Job> jobCaptor = ArgumentCaptor.forClass(Job.class); doAnswer( new Answer<Job>() { @Override public Job answer(InvocationOnMock invocationOnMock) throws Throwable { verify(mockBigqueryJobs, times(1)).insert(eq(jobProjectId), jobCaptor.capture()); return jobCaptor.getValue(); } }) .when(mockBigqueryJobsInsert) .execute(); when(mockBigqueryJobsGet.execute()).thenReturn(jobHandle); helper.importFromGcs( jobProjectId, tableRef, fakeTableSchema, null, kmsKeyName, BigQueryFileFormat.NEWLINE_DELIMITED_JSON, BigQueryConfiguration.OUTPUT_TABLE_CREATE_DISPOSITION.getDefault(), BigQueryConfiguration.OUTPUT_TABLE_WRITE_DISPOSITION.getDefault(), ImmutableList.of("test-import-path"), true); verify(mockBigquery, times(2)).jobs(); verify(mockBigqueryJobs, times(1)) .get(eq(jobProjectId), eq(jobCaptor.getValue().getJobReference().getJobId())); Job job = jobCaptor.getValue(); assertThat(job.getConfiguration().getLoad().getSourceUris()).contains("test-import-path"); assertThat(job.getConfiguration().getLoad().getDestinationTable()).isEqualTo(tableRef); assertThat(job.getJobReference().getLocation()).isEqualTo("test_location"); verify(mockBigqueryJobsGet).setLocation(eq("test_location")); verify(mockBigqueryJobsGet, times(1)).execute(); verify(mockBigqueryJobsInsert, times(1)).execute(); verify(mockBigquery).datasets(); }