Merge pull request #515 from AdaptiveScale/cherrypick-release/0.14/PLUGIN-499

CherryPick PLUGIN-499 for release/0.14

Author: CuriousVini

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySource.java

@@ -24,6 +24,7 @@
 import com.google.cloud.bigquery.StandardTableDefinition;
 import com.google.cloud.bigquery.Table;
 import com.google.cloud.bigquery.TableDefinition.Type;
+import com.google.cloud.bigquery.TableId;
 import com.google.cloud.bigquery.TimePartitioning;
 import com.google.cloud.hadoop.io.bigquery.BigQueryConfiguration;
 import io.cdap.cdap.api.annotation.Description;
@@ -118,8 +119,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
     Schema configuredSchema = getOutputSchema(collector);
 
     String serviceAccount = config.getServiceAccount();
-    Credentials credentials = serviceAccount == null ?
-      null : GCPUtils.loadServiceAccountCredentials(serviceAccount, config.isServiceAccountFilePath());
+    Credentials credentials = getCredentials(serviceAccount);
     BigQuery bigQuery = GCPUtils.getBigQuery(config.getDatasetProject(), credentials);
 
     uuid = UUID.randomUUID();
@@ -163,6 +163,9 @@ public void prepareRun(BatchSourceContext context) throws Exception {
     if (config.getViewMaterializationDataset() != null) {
       configuration.set(BigQueryConstants.CONFIG_VIEW_MATERIALIZATION_DATASET, config.getViewMaterializationDataset());
     }
+    String temporaryTableName = String.format("_%s_%s", config.getTable(),
+                                              UUID.randomUUID().toString().replaceAll("-", "_"));
+    configuration.set(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME, temporaryTableName);
 
     String temporaryGcsPath = String.format("gs://%s/%s/hadoop/input/%s", bucket, uuid, uuid);
     PartitionedBigQueryInputFormat.setTemporaryCloudStorageDirectory(configuration, temporaryGcsPath);
@@ -203,6 +206,24 @@ public void transform(KeyValue<LongWritable, GenericData.Record> input, Emitter<
 
   @Override
   public void onRunFinish(boolean succeeded, BatchSourceContext context) {
+    deleteGcsTemporaryDirectory();
+    deleteBigQueryTemporaryTable();
+  }
+
+  private void deleteBigQueryTemporaryTable() {
+    String temporaryTable = configuration.get(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME);
+    try {
+      String serviceAccount = config.getServiceAccount();
+      Credentials credentials = getCredentials(serviceAccount);
+      BigQuery bigQuery = GCPUtils.getBigQuery(config.getDatasetProject(), credentials);
+      bigQuery.delete(TableId.of(config.getProject(), config.getDataset(), temporaryTable));
+      LOG.debug("Deleted temporary table '{}'", temporaryTable);
+    } catch (IOException e) {
+      LOG.error("Failed to load service account credentials: {}", e.getMessage(), e);
+    }
+  }
+
+  private void deleteGcsTemporaryDirectory() {
     org.apache.hadoop.fs.Path gcsPath = null;
     String bucket = config.getBucket();
     if (bucket == null) {
@@ -221,6 +242,12 @@ public void onRunFinish(boolean succeeded, BatchSourceContext context) {
     }
   }
 
+  @Nullable
+  private Credentials getCredentials(@Nullable String serviceAccount) throws IOException {
+    return serviceAccount == null ?
+            null : GCPUtils.loadServiceAccountCredentials(serviceAccount, config.isServiceAccountFilePath());
+  }
+
   public Schema getSchema(FailureCollector collector) {
     com.google.cloud.bigquery.Schema bqSchema = getBQSchema(collector);
     return BigQueryUtil.getTableSchema(bqSchema, collector);

src/main/java/io/cdap/plugin/gcp/bigquery/source/PartitionedBigQueryInputFormat.java

@@ -110,7 +110,7 @@ private void processQuery(JobContext context) throws IOException, InterruptedExc
       TableReference sourceTable = new TableReference().setDatasetId(datasetId).setProjectId(inputProjectId)
         .setTableId(tableName);
       String location = bigQueryHelper.getTable(sourceTable).getLocation();
-      String temporaryTableName = String.format("%s_%s", tableName, UUID.randomUUID().toString().replaceAll("-", "_"));
+      String temporaryTableName = configuration.get(BigQueryConstants.CONFIG_TEMPORARY_TABLE_NAME);
       TableReference exportTableReference = createExportTableReference(type, inputProjectId, datasetId,
                                                                        temporaryTableName, configuration);
       runQuery(bigQueryHelper, inputProjectId, exportTableReference, query, location);

src/main/java/io/cdap/plugin/gcp/bigquery/util/BigQueryConstants.java

@@ -42,4 +42,5 @@ public interface BigQueryConstants {
   String CONFIG_PARTITION_INTEGER_RANGE_START = "cdap.bq.sink.partition.integer.range.start";
   String CONFIG_PARTITION_INTEGER_RANGE_END = "cdap.bq.sink.partition.integer.range.end";
   String CONFIG_PARTITION_INTEGER_RANGE_INTERVAL = "cdap.bq.sink.partition.integer.range.interval";
+  String CONFIG_TEMPORARY_TABLE_NAME = "cdap.bq.source.temporary.table.name";
 }