Merge pull request #1146 from data-integrations/feat/add-fqn-to-gcs

CDAP-19623: add fqn support to gcs plugins

Author: mrahanjam

src/main/java/io/cdap/plugin/gcp/gcs/sink/GCSBatchSink.java

@@ -40,10 +40,12 @@
 import io.cdap.cdap.etl.api.batch.BatchSinkContext;
 import io.cdap.cdap.etl.api.connector.Connector;
 import io.cdap.cdap.etl.api.validation.ValidatingOutputFormat;
+import io.cdap.plugin.common.Asset;
 import io.cdap.plugin.common.ConfigUtil;
 import io.cdap.plugin.common.Constants;
 import io.cdap.plugin.common.IdUtils;
 import io.cdap.plugin.common.LineageRecorder;
+import io.cdap.plugin.common.ReferenceNames;
 import io.cdap.plugin.format.FileFormat;
 import io.cdap.plugin.format.plugin.AbstractFileSink;
 import io.cdap.plugin.format.plugin.FileSinkProperties;
@@ -84,6 +86,7 @@ public class GCSBatchSink extends AbstractFileSink<GCSBatchSink.GCSBatchSinkConf
 
   private final GCSBatchSinkConfig config;
   private String outputPath;
+  private Asset asset;
 
   public GCSBatchSink(GCSBatchSinkConfig config) {
     super(config);
@@ -109,7 +112,6 @@ public ValidatingOutputFormat getOutputFormatForRun(BatchSinkContext context) th
 
   @Override
   public void prepareRun(BatchSinkContext context) throws Exception {
-    super.prepareRun(context);
     FailureCollector collector = context.getFailureCollector();
     CryptoKeyName cmekKeyName = CmekUtils.getCmekKey(config.cmekKey, context.getArguments().asMap(), collector);
     collector.getOrThrowException();
@@ -125,17 +127,30 @@ public void prepareRun(BatchSinkContext context) throws Exception {
       null : GCPUtils.loadServiceAccountCredentials(config.connection.getServiceAccount(), isServiceAccountFilePath);
     Storage storage = GCPUtils.getStorage(config.connection.getProject(), credentials);
     Bucket bucket;
+    String location;
     try {
       bucket = storage.get(config.getBucket());
     } catch (StorageException e) {
       throw new RuntimeException(
         String.format("Unable to access or create bucket %s. ", config.getBucket())
           + "Ensure you entered the correct bucket path and have permissions for it.", e);
     }
-    if (bucket == null) {
+    if (bucket != null) {
+      location = bucket.getLocation();
+    } else {
       GCPUtils.createBucket(storage, config.getBucket(), config.getLocation(), cmekKeyName);
+      location = config.getLocation();
     }
     this.outputPath = getOutputDir(context);
+    // create asset for lineage
+    String referenceName = Strings.isNullOrEmpty(config.getReferenceName())
+      ? ReferenceNames.normalizeFqn(config.getPath())
+      : config.getReferenceName();
+    asset = Asset.builder(referenceName)
+      .setFqn(config.getPath()).setLocation(location).build();
+    
+    // super is called down here to avoid instantiating the lineage recorder with a null asset
+    super.prepareRun(context);
   }
 
   @Override
@@ -154,6 +169,11 @@ protected Map<String, String> getFileSystemProperties(BatchSinkContext context)
     return properties;
   }
 
+  @Override
+  protected LineageRecorder getLineageRecorder(BatchSinkContext context) {
+    return new LineageRecorder(context, asset);
+  }
+
   @Override
   protected void recordLineage(LineageRecorder lineageRecorder, List<String> outputFields) {
     lineageRecorder.recordWrite("Write", "Wrote to Google Cloud Storage.", outputFields);

src/main/java/io/cdap/plugin/gcp/gcs/source/GCSSource.java

@@ -16,6 +16,9 @@
 
 package io.cdap.plugin.gcp.gcs.source;
 
+import com.google.auth.Credentials;
+import com.google.cloud.storage.Storage;
+import com.google.cloud.storage.StorageException;
 import com.google.common.base.Strings;
 import com.google.gson.Gson;
 import com.google.gson.reflect.TypeToken;
@@ -25,27 +28,22 @@
 import io.cdap.cdap.api.annotation.MetadataProperty;
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.api.annotation.Plugin;
-import io.cdap.cdap.api.data.schema.Schema;
-import io.cdap.cdap.api.plugin.PluginConfig;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.cdap.etl.api.PipelineConfigurer;
 import io.cdap.cdap.etl.api.batch.BatchSource;
 import io.cdap.cdap.etl.api.batch.BatchSourceContext;
 import io.cdap.cdap.etl.api.connector.Connector;
+import io.cdap.plugin.common.Asset;
 import io.cdap.plugin.common.ConfigUtil;
-import io.cdap.plugin.common.Constants;
-import io.cdap.plugin.common.IdUtils;
 import io.cdap.plugin.common.LineageRecorder;
-import io.cdap.plugin.format.FileFormat;
-import io.cdap.plugin.format.charset.fixedlength.FixedLengthCharset;
+import io.cdap.plugin.common.ReferenceNames;
 import io.cdap.plugin.format.input.PathTrackingInputFormat;
 import io.cdap.plugin.format.plugin.AbstractFileSource;
 import io.cdap.plugin.format.plugin.AbstractFileSourceConfig;
 import io.cdap.plugin.format.plugin.FileSourceProperties;
 import io.cdap.plugin.gcp.common.GCPConnectorConfig;
 import io.cdap.plugin.gcp.common.GCPUtils;
 import io.cdap.plugin.gcp.crypto.EncryptedFileSystem;
-import io.cdap.plugin.gcp.gcs.Formats;
 import io.cdap.plugin.gcp.gcs.GCSPath;
 import io.cdap.plugin.gcp.gcs.connector.GCSConnector;
 
@@ -67,6 +65,7 @@
 public class GCSSource extends AbstractFileSource<GCSSource.GCSSourceConfig> {
   public static final String NAME = "GCSFile";
   private final GCSSourceConfig config;
+  private Asset asset;
 
   public GCSSource(GCSSourceConfig config) {
     super(config);
@@ -78,6 +77,34 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
     super.configurePipeline(pipelineConfigurer);
   }
 
+  @Override
+  public void prepareRun(BatchSourceContext context) throws Exception {
+    // Get location of the source for lineage
+    String location;
+    String bucketName = GCSPath.from(config.getPath()).getBucket();
+    Credentials credentials = config.connection.getServiceAccount() == null ?
+      null : GCPUtils.loadServiceAccountCredentials(config.connection.getServiceAccount(),
+                                                    config.connection.isServiceAccountFilePath());
+    Storage storage = GCPUtils.getStorage(config.connection.getProject(), credentials);
+    try {
+      location = storage.get(bucketName).getLocation();
+    } catch (StorageException e) {
+      throw new RuntimeException(
+        String.format("Unable to access bucket %s. ", bucketName)
+          + "Ensure you entered the correct bucket path and have permissions for it.", e);
+    }
+
+    // create asset for lineage
+    String referenceName = Strings.isNullOrEmpty(config.getReferenceName())
+      ? ReferenceNames.normalizeFqn(config.getPath())
+      : config.getReferenceName();
+    asset = Asset.builder(referenceName)
+      .setFqn(config.getPath()).setLocation(location).build();
+
+    // super is called down here to avoid instantiating the lineage recorder with a null asset
+    super.prepareRun(context);
+  }
+
   @Override
   protected Map<String, String> getFileSystemProperties(BatchSourceContext context) {
     Map<String, String> properties = GCPUtils.getFileSystemProperties(config.connection, config.getPath(),
@@ -101,6 +128,11 @@ protected Map<String, String> getFileSystemProperties(BatchSourceContext context
     return properties;
   }
 
+  @Override
+  protected LineageRecorder getLineageRecorder(BatchSourceContext context) {
+    return new LineageRecorder(context, asset);
+  }
+
   @Override
   protected void recordLineage(LineageRecorder lineageRecorder, List<String> outputFields) {
     lineageRecorder.recordRead("Read", String.format("Read%sfrom Google Cloud Storage.",