Merge pull request #501 from data-integrations/cherrypick-plugin-474-015

[PLUGIN-474] Fixed BigQuery Plugins Not Deleting Temporary Files

Author: dli357

src/main/java/io/cdap/plugin/gcp/bigquery/sink/AbstractBigQuerySink.java

@@ -73,7 +73,7 @@ public abstract class AbstractBigQuerySink extends BatchSink<StructuredRecord, A
 
   private static final Logger LOG = LoggerFactory.getLogger(AbstractBigQuerySink.class);
 
-  private static final String gcsPathFormat = "gs://%s";
+  private static final String gcsPathFormat = "gs://%s/%s";
   private static final String temporaryBucketFormat = gcsPathFormat + "/input/%s-%s";
   public static final String RECORDS_UPDATED_METRIC = "records.updated";
 
@@ -115,17 +115,21 @@ public final void prepareRun(BatchSinkContext context) throws Exception {
 
   @Override
   public void onRunFinish(boolean succeeded, BatchSinkContext context) {
-    if (getConfig().getBucket() == null) {
-      Path gcsPath = new Path(String.format(gcsPathFormat, uuid.toString()));
-      try {
-        FileSystem fs = gcsPath.getFileSystem(baseConfiguration);
-        if (fs.exists(gcsPath)) {
-          fs.delete(gcsPath, true);
-          LOG.debug("Deleted temporary bucket '{}'", gcsPath);
-        }
-      } catch (IOException e) {
-        LOG.warn("Failed to delete bucket '{}': {}", gcsPath, e.getMessage());
+    Path gcsPath;
+    String bucket = getConfig().getBucket();
+    if (bucket == null) {
+      gcsPath = new Path(String.format(gcsPathFormat, uuid.toString(), uuid.toString()));
+    } else {
+      gcsPath = new Path(String.format(gcsPathFormat, bucket, uuid.toString()));
+    }
+    try {
+      FileSystem fs = gcsPath.getFileSystem(baseConfiguration);
+      if (fs.exists(gcsPath)) {
+        fs.delete(gcsPath, true);
+        LOG.debug("Deleted temporary directory '{}'", gcsPath);
       }
+    } catch (IOException e) {
+      LOG.warn("Failed to delete temporary directory '{}': {}", gcsPath, e.getMessage());
     }
   }
 
@@ -240,7 +244,7 @@ private Configuration getBaseConfiguration(@Nullable String cmekKey) throws IOEx
    * @return full path to temporary bucket
    */
   private String getTemporaryGcsPath(String bucket, String tableName) {
-    return String.format(temporaryBucketFormat, bucket, tableName, uuid);
+    return String.format(temporaryBucketFormat, bucket, uuid, tableName, uuid);
   }
 
   /**
@@ -257,7 +261,7 @@ private String configureBucket() {
       // So enable it only when bucket name is not provided.
       baseConfiguration.setBoolean("fs.gs.bucket.delete.enable", true);
     }
-    baseConfiguration.set("fs.gs.system.bucket", bucket);
+    baseConfiguration.set("fs.default.name", String.format(gcsPathFormat, bucket, uuid));
     baseConfiguration.setBoolean("fs.gs.impl.disable.cache", true);
     baseConfiguration.setBoolean("fs.gs.metadata.cache.enable", false);
     return bucket;

src/main/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySink.java

@@ -131,7 +131,7 @@ public void onRunFinish(boolean succeeded, BatchSinkContext context) {
     }
   }
 
-  private void recordMetric(boolean succeeded, BatchSinkContext context) {
+  void recordMetric(boolean succeeded, BatchSinkContext context) {
     if (!succeeded) {
       return;
     }

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

@@ -49,6 +49,7 @@
 import org.apache.avro.generic.GenericData;
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.io.LongWritable;
 import org.apache.hadoop.io.Text;
 import org.apache.hadoop.mapreduce.Job;
@@ -147,7 +148,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
                             cmekKey);
     }
 
-    configuration.set("fs.gs.system.bucket", bucket);
+    configuration.set("fs.default.name", String.format("gs://%s/%s/", bucket, uuid));
     configuration.setBoolean("fs.gs.impl.disable.cache", true);
     configuration.setBoolean("fs.gs.metadata.cache.enable", false);
 
@@ -171,7 +172,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
       configuration.set(BigQueryConstants.CONFIG_VIEW_MATERIALIZATION_DATASET, config.getViewMaterializationDataset());
     }
 
-    String temporaryGcsPath = String.format("gs://%s/hadoop/input/%s", bucket, uuid);
+    String temporaryGcsPath = String.format("gs://%s/%s/hadoop/input/%s", bucket, uuid, uuid);
     PartitionedBigQueryInputFormat.setTemporaryCloudStorageDirectory(configuration, temporaryGcsPath);
     BigQueryConfiguration.configureBigQueryInput(configuration, config.getDatasetProject(),
                                                  config.getDataset(), config.getTable());
@@ -210,16 +211,21 @@ public void transform(KeyValue<LongWritable, GenericData.Record> input, Emitter<
 
   @Override
   public void onRunFinish(boolean succeeded, BatchSourceContext context) {
-    org.apache.hadoop.fs.Path gcsPath = new org.apache.hadoop.fs.Path(String.format("gs://%s", uuid.toString()));
+    org.apache.hadoop.fs.Path gcsPath = null;
+    String bucket = config.getBucket();
+    if (bucket == null) {
+      gcsPath = new Path(String.format("gs://%s/%s", uuid, uuid));
+    } else {
+      gcsPath = new Path(String.format("gs://%s/%s", bucket, uuid));
+    }
     try {
-      if (config.getBucket() == null) {
-        FileSystem fs = gcsPath.getFileSystem(configuration);
-        if (fs.exists(gcsPath)) {
-          fs.delete(gcsPath, true);
-        }
+      FileSystem fs = gcsPath.getFileSystem(configuration);
+      if (fs.exists(gcsPath)) {
+        fs.delete(gcsPath, true);
+        LOG.debug("Deleted temporary directory '{}'", gcsPath);
       }
     } catch (IOException e) {
-      LOG.warn("Failed to delete bucket " + gcsPath.toUri().getPath() + ", " + e.getMessage());
+      LOG.warn("Failed to delete temporary directory '{}': {}", gcsPath, e.getMessage());
     }
   }
 

src/test/java/io/cdap/plugin/gcp/bigquery/sink/BigQuerySinkTest.java

@@ -28,7 +28,6 @@
 import com.google.cloud.bigquery.TableDefinition;
 import com.google.cloud.bigquery.TableId;
 import io.cdap.cdap.api.data.schema.Schema;
-import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.cdap.etl.api.batch.BatchSinkContext;
 import io.cdap.cdap.etl.api.validation.ValidationException;
 import io.cdap.cdap.etl.api.validation.ValidationFailure;
@@ -120,7 +119,7 @@ private void testMetric(Job mockJob, long expectedCount, int invocations)
     BigQuerySink sink = getSinkToTest(mockJob);
     MockStageMetrics mockStageMetrics = Mockito.spy(new MockStageMetrics("test"));
     BatchSinkContext context = getContextWithMetrics(mockStageMetrics);
-    sink.onRunFinish(true, context);
+    sink.recordMetric(true, context);
     if (expectedCount > -1) {
       Assert.assertEquals(expectedCount, mockStageMetrics.getCount(AbstractBigQuerySink.RECORDS_UPDATED_METRIC));
     }