Merge pull request #72 from cloudsufi/feat/AbstractFileBatchSource

[PLUGIN-1906] GDrive abstract file batch source plugin extension

Author: psainics

src/main/java/io/cdap/plugin/google/common/GoogleDriveClient.java

@@ -68,7 +68,7 @@ public GoogleDriveClient(C config) throws IOException {
    * @return {@link Drive} client.
    * @throws IOException on issues with service account file reading.
    */
-  protected Drive getDriveClient() throws IOException {
+  public Drive getDriveClient() throws IOException {
     Drive drive;
     AuthType authType = config.getAuthType();
     switch (authType) {

src/main/java/io/cdap/plugin/google/drive/source/GoogleDriveFileSource.java

@@ -0,0 +1,50 @@
+/*
+ * Copyright © 2025 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap.plugin.google.drive.source;
+
+import io.cdap.cdap.etl.api.FailureCollector;
+import io.cdap.cdap.etl.api.batch.BatchSourceContext;
+import io.cdap.plugin.format.plugin.AbstractFileSource;
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * This class gives ability to read files from Google Drive with a structured schema.
+ */
+public class GoogleDriveFileSource extends AbstractFileSource<GoogleDriveSourceConfig> {
+  public static final String NAME = "GoogleDrive";
+  private final GoogleDriveSourceConfig config;
+
+  public GoogleDriveFileSource(GoogleDriveSourceConfig config) {
+    super(config);
+    this.config = config;
+  }
+
+  @Override
+  protected Map<String, String> getFileSystemProperties(BatchSourceContext context) {
+    FailureCollector collector = context == null ? null : context.getFailureCollector();
+    return new HashMap<>(config.getFileSystemProperties(collector));
+  }
+
+  @Override
+  protected boolean shouldGetSchema() {
+    return !config.containsMacro(GoogleDriveSourceConfig.NAME_FORMAT)
+        && !config.containsMacro(GoogleDriveSourceConfig.NAME_DELIMITER)
+        && !config.containsMacro(GoogleDriveSourceConfig.NAME_FILE_SYSTEM_PROPERTIES)
+        && !config.containsMacro(GoogleDriveSourceConfig.NAME_FILE_ENCODING);
+  }
+}

src/main/java/io/cdap/plugin/google/drive/source/GoogleDriveSource.java

@@ -23,14 +23,11 @@
 import io.cdap.cdap.api.data.batch.Input;
 import io.cdap.cdap.api.data.format.StructuredRecord;
 import io.cdap.cdap.api.data.schema.Schema;
-import io.cdap.cdap.api.dataset.lib.KeyValue;
-import io.cdap.cdap.etl.api.Emitter;
 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.plugin.common.LineageRecorder;
-import io.cdap.plugin.google.drive.common.FileFromFolder;
 import org.apache.hadoop.io.NullWritable;
 
 import java.util.stream.Collectors;
@@ -45,9 +42,11 @@ public class GoogleDriveSource extends BatchSource<NullWritable, StructuredRecor
   public static final String NAME = "GoogleDrive";
 
   private final GoogleDriveSourceConfig config;
+  private final GoogleDriveFileSource delegate;
 
   public GoogleDriveSource(GoogleDriveSourceConfig config) {
     this.config = config;
+    this.delegate = new GoogleDriveFileSource(config);
   }
 
   @Override
@@ -56,22 +55,29 @@ public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
     config.getValidationResult(failureCollector);
     failureCollector.getOrThrowException();
 
-    pipelineConfigurer.getStageConfigurer().setOutputSchema(config.getSchema());
+    if (!config.isStructuredSchemaRequired()) {
+      pipelineConfigurer.getStageConfigurer().setOutputSchema(config.getSchema());
+      return;
+    }
+    delegate.configurePipeline(pipelineConfigurer);
   }
 
   @Override
-  public void prepareRun(BatchSourceContext context) {
+  public void prepareRun(BatchSourceContext context) throws Exception {
     FailureCollector failureCollector = context.getFailureCollector();
     config.getValidationResult(failureCollector);
     failureCollector.getOrThrowException();
 
-    LineageRecorder lineageRecorder = new LineageRecorder(context, config.getReferenceName());
-    lineageRecorder.createExternalDataset(config.getSchema());
-    lineageRecorder.recordRead("Read", "Reading Google Drive files",
-                               Preconditions.checkNotNull(config.getSchema().getFields()).stream()
-                                 .map(Schema.Field::getName)
-                                 .collect(Collectors.toList()));
+    if (!config.isStructuredSchemaRequired()) {
+      LineageRecorder lineageRecorder = new LineageRecorder(context, config.getReferenceName());
+      lineageRecorder.createExternalDataset(config.getSchema());
+      lineageRecorder.recordRead("Read", "Reading Google Drive files",
+          Preconditions.checkNotNull(config.getSchema().getFields()).stream().map(Schema.Field::getName)
+              .collect(Collectors.toList()));
 
-    context.setInput(Input.of(config.getReferenceName(), new GoogleDriveInputFormatProvider(config)));
+      context.setInput(Input.of(config.getReferenceName(), new GoogleDriveInputFormatProvider(config)));
+      return;
+    }
+    delegate.prepareRun(context);
   }
 }

src/main/java/io/cdap/plugin/google/drive/source/GoogleDriveSourceConfig.java

@@ -17,7 +17,10 @@
 package io.cdap.plugin.google.drive.source;
 
 import com.google.common.base.Strings;
+import com.google.gson.Gson;
+import com.google.gson.GsonBuilder;
 import com.google.gson.JsonObject;
+import com.google.gson.reflect.TypeToken;
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Macro;
 import io.cdap.cdap.api.annotation.Name;
@@ -31,12 +34,16 @@
 import io.cdap.plugin.google.common.ValidationResult;
 import io.cdap.plugin.google.common.exceptions.InvalidPropertyTypeException;
 import io.cdap.plugin.google.common.utils.ExportedType;
+import io.cdap.plugin.google.drive.source.fs.GoogleDriveFileSystem;
 import io.cdap.plugin.google.drive.source.utils.BodyFormat;
 
 import java.io.IOException;
+import java.lang.reflect.Type;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -57,6 +64,9 @@ public class GoogleDriveSourceConfig extends GoogleFilteringSourceConfig impleme
   public static final String IS_STRUCTURED_SCHEMA_REQUIRED = "structuredSchemaRequired";
   public static final String NAME_SCHEMA = "schema";
   public static final String NAME_FORMAT = "format";
+  public static final String NAME_DELIMITER = "delimiter";
+  public static final String NAME_FILE_SYSTEM_PROPERTIES = "filesystemproperties";
+  public static final String NAME_FILE_ENCODING = "fileEncoding";
 
   public static final String DEFAULT_BODY_FORMAT = "bytes";
   public static final long DEFAULT_MAX_PARTITION_SIZE = 0;
@@ -75,6 +85,12 @@ public class GoogleDriveSourceConfig extends GoogleFilteringSourceConfig impleme
   public static final String GOOGLE_DRIVE_FOLDER_PATH_PREFIX = "/drive/folders";
   public static final String GOOGLE_DRIVE_DEFAULT_FILENAME = "default.txt";
 
+  public static final String GOOGLE_DRIVE_FILESYSTEM_IMPL = String.format("fs.%s.impl", GOOGLE_DRIVE_SCHEMA);
+  public static final String GOOGLE_DRIVE_FILESYSTEM_DISABLE_CACHE_KEY = String.format("fs.%s.impl.disable.cache",
+      GOOGLE_DRIVE_SCHEMA);
+  public static final Gson GSON = new GsonBuilder().create();
+  private static final Type MAP_STRING_STRING_TYPE = new TypeToken<Map<String, String>>() { }.getType();
+
   @Nullable
   @Name(FILE_METADATA_PROPERTIES)
   @Description("Properties that represent metadata of files. \n" +
@@ -168,6 +184,7 @@ public class GoogleDriveSourceConfig extends GoogleFilteringSourceConfig impleme
       "skipped.")
   private String override;
 
+  @Name(NAME_DELIMITER)
   @Macro
   @Nullable
   @Description("The delimiter to use if the format is 'delimited'. The delimiter will be ignored if the format "
@@ -186,12 +203,14 @@ public class GoogleDriveSourceConfig extends GoogleFilteringSourceConfig impleme
       "is 'csv', 'tsv' or 'delimited'. The default value is false.")
   protected Boolean enableQuotedValues;
 
+  @Name(NAME_FILE_SYSTEM_PROPERTIES)
   @Macro
   @Nullable
   @Description("Any additional properties to use when reading from the filesystem. "
       + "This is an advanced feature that requires knowledge of the properties supported by the underlying filesystem.")
   private String fileSystemProperties;
 
+  @Name(NAME_FILE_ENCODING)
   @Macro
   @Nullable
   @Description("File encoding for the source files. The default encoding is 'UTF-8'")
@@ -588,4 +607,31 @@ public static GoogleDriveSourceConfig of(JsonObject properties) throws IOExcepti
     }
     return googleDriveSourceConfig;
   }
+
+  public Map<String, String> getFileSystemProperties(@Nullable FailureCollector collector) {
+    Map<String, String> properties = new HashMap<>();
+    properties.put(GOOGLE_DRIVE_FILESYSTEM_IMPL, GoogleDriveFileSystem.class.getName());
+    properties.put(GoogleDriveInputFormatProvider.PROPERTY_CONFIG_JSON, GSON.toJson(getProperties()));
+    properties.put(GOOGLE_DRIVE_FILESYSTEM_DISABLE_CACHE_KEY, Boolean.TRUE.toString());
+
+    if (fileSystemProperties != null) {
+      try {
+        Map<String, String> userProperties = GSON.fromJson(fileSystemProperties, MAP_STRING_STRING_TYPE);
+        for (Map.Entry<String, String> entry : userProperties.entrySet()) {
+          if (entry.getKey() != null && !entry.getKey().isEmpty()) {
+            // let user override the default properties
+            properties.put(entry.getKey(), entry.getValue());
+          }
+        }
+      } catch (Exception e) {
+        String errorMessage = String.format("Unable to parse filesystem properties, %s: %s", e.getClass().getName(),
+            e.getMessage());
+        if (collector == null) {
+          throw new IllegalArgumentException(errorMessage, e);
+        }
+        collector.addFailure(errorMessage, "Ensure the properties are in a valid JSON format.");
+      }
+    }
+    return properties;
+  }
 }

src/main/java/io/cdap/plugin/google/drive/source/fs/GoogleDriveFileSystem.java

@@ -0,0 +1,143 @@
+/*
+ * Copyright © 2025 Cask Data, Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not
+ * use this file except in compliance with the License. You may obtain a copy of
+ * the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+ * License for the specific language governing permissions and limitations under
+ * the License.
+ */
+
+package io.cdap.plugin.google.drive.source.fs;
+
+import com.google.api.services.drive.Drive;
+import com.google.common.base.Strings;
+import io.cdap.plugin.google.common.GoogleDriveClient;
+import io.cdap.plugin.google.drive.source.GoogleDriveInputFormatProvider;
+import io.cdap.plugin.google.drive.source.GoogleDriveSourceConfig;
+import org.apache.hadoop.conf.Configuration;
+import org.apache.hadoop.fs.FSDataInputStream;
+import org.apache.hadoop.fs.FSDataOutputStream;
+import org.apache.hadoop.fs.FileStatus;
+import org.apache.hadoop.fs.FileSystem;
+import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.permission.FsPermission;
+import org.apache.hadoop.util.Progressable;
+
+import java.io.FileNotFoundException;
+import java.io.IOException;
+import java.net.URI;
+
+/**
+ * A custom Hadoop FileSystem implementation for Google Drive.
+ * This class provides methods to interact with Google Drive files and directories.
+ */
+public class GoogleDriveFileSystem extends FileSystem {
+  private URI uri;
+  private Path workingDir;
+  private Drive driveService;
+  private String filter;
+
+  @Override
+  public void initialize(URI name, Configuration conf) throws IOException {
+    super.initialize(name, conf);
+    setConf(conf);
+    this.uri = URI.create(name.getScheme() + "://" + name.getAuthority());
+    this.workingDir = new Path("/");
+
+    GoogleDriveSourceConfig googleDriveSourceConfig = GoogleDriveInputFormatProvider.extractPropertiesFromConfig(conf);
+    GoogleDriveClient<GoogleDriveSourceConfig> client = new GoogleDriveClient<>(googleDriveSourceConfig);
+
+    // Initialize Google Drive service (using OAuth2 or service account)
+    this.driveService = client.getDriveClient();
+
+    // Initialize filter to be passed down to listStatus
+    if (!Strings.isNullOrEmpty(googleDriveSourceConfig.getFilter())) {
+      this.filter = googleDriveSourceConfig.getFilter();
+    }
+  }
+  @Override
+  public URI getUri() {
+    return this.uri;
+  }
+
+  @Override
+  public FSDataInputStream open(Path f, int bufferSize) throws IOException {
+    return new GoogleDriveInputStream(new GoogleDriveInputStreamWrapper(driveService, f));
+  }
+
+  @Override
+  public FSDataOutputStream create(Path f,
+                                   FsPermission permission,
+                                   boolean overwrite,
+                                   int bufferSize,
+                                   short replication,
+                                   long blockSize,
+                                   Progressable progress) throws IOException {
+    throw new UnsupportedOperationException(
+      "GDrive does not support: FSDataOutputStream create(Path, FsPermission, boolean, int, " +
+        "short, long, Progressable)");
+  }
+
+  @Override
+  public FSDataOutputStream append(Path f, int bufferSize, Progressable progress) throws IOException {
+    throw new UnsupportedOperationException(
+      "GDrive does not support: FSDataOutputStream append(Path, int, Progressable)");
+  }
+
+  @Override
+  public boolean rename(Path src, Path dst) throws IOException {
+    throw  new UnsupportedOperationException("GDrive does not support: boolean rename(Path, Path)");
+  }
+
+  @Override
+  public boolean delete(Path f, boolean recursive) throws IOException {
+    throw  new UnsupportedOperationException("GDrive does not support: boolean delete(Path, boolean)");
+  }
+
+  @Override
+  public FileStatus[] listStatus(Path f) throws FileNotFoundException, IOException {
+    if (isPathDirectory(f)) {
+      return GoogleDriveUtils.listStatus(driveService, f, filter);
+    }
+    return new FileStatus[]{GoogleDriveUtils.listObjectStatus(driveService, f, false, GoogleDriveUtils.getFileId(f))};
+  }
+
+  @Override
+  public void setWorkingDirectory(Path newDir) {
+    throw new UnsupportedOperationException("GDrive does not support: void setWorkingDirectory(Path)");
+  }
+
+  @Override
+  public Path getWorkingDirectory() {
+    return workingDir;
+  }
+
+  @Override
+  public boolean mkdirs(Path f, FsPermission permission) throws IOException {
+    throw new UnsupportedOperationException(
+      "GDrive does not support: boolean mkdirs(Path, FsPermission)");
+  }
+
+  @Override
+  public FileStatus getFileStatus(Path f) throws IOException {
+    if (isPathDirectory(f)) {
+      return GoogleDriveUtils.listObjectStatus(driveService, f, true, f.getName());
+    }
+    return GoogleDriveUtils.listObjectStatus(driveService, f, false, GoogleDriveUtils.getFileId(f));
+  }
+
+  private boolean isPathDirectory(Path path) {
+    return path != null && path.toString().startsWith(String.format("%s://%s%s/",
+        GoogleDriveSourceConfig.GOOGLE_DRIVE_SCHEMA,
+        GoogleDriveSourceConfig.GOOGLE_DRIVE_AUTHORITY,
+        GoogleDriveSourceConfig.GOOGLE_DRIVE_FOLDER_PATH_PREFIX));
+  }
+
+}