private List<ResourceId> executeExtract( String jobId, TableReference table, JobService jobService, String executingProject, String extractDestinationDir, String bqLocation) throws InterruptedException, IOException { JobReference jobRef = new JobReference().setProjectId(executingProject).setLocation(bqLocation).setJobId(jobId); String destinationUri = BigQueryIO.getExtractDestinationUri(extractDestinationDir); JobConfigurationExtract extract = new JobConfigurationExtract() .setSourceTable(table) .setDestinationFormat("AVRO") .setDestinationUris(ImmutableList.of(destinationUri)); LOG.info("Starting BigQuery extract job: {}", jobId); jobService.startExtractJob(jobRef, extract); Job extractJob = jobService.pollJob(jobRef, JOB_POLL_MAX_RETRIES); if (BigQueryHelpers.parseStatus(extractJob) != Status.SUCCEEDED) { throw new IOException( String.format( "Extract job %s failed, status: %s.", extractJob.getJobReference().getJobId(), BigQueryHelpers.statusToPrettyString(extractJob.getStatus()))); } LOG.info("BigQuery extract job completed: {}", jobId); return BigQueryIO.getExtractFilePaths(extractDestinationDir, extractJob); }
@Override void cleanup(PassThroughThenCleanup.ContextContainer c) throws Exception { PipelineOptions options = c.getPipelineOptions(); BigQueryOptions bqOptions = options.as(BigQueryOptions.class); String jobUuid = c.getJobId(); final String extractDestinationDir = resolveTempLocation(bqOptions.getTempLocation(), "BigQueryExtractTemp", jobUuid); final String executingProject = bqOptions.getProject(); JobReference jobRef = new JobReference() .setProjectId(executingProject) .setJobId(getExtractJobId(createJobIdToken(bqOptions.getJobName(), jobUuid))); Job extractJob = getBigQueryServices().getJobService(bqOptions).getJob(jobRef); if (extractJob != null) { List<ResourceId> extractFiles = getExtractFilePaths(extractDestinationDir, extractJob); if (extractFiles != null && !extractFiles.isEmpty()) { FileSystems.delete( extractFiles, MoveOptions.StandardMoveOptions.IGNORE_MISSING_FILES); } } } };