PLUGIN-1883: Using schema to fetch the column names while constructing importQuery

Author: sgarg-CS

src/main/java/io/cdap/plugin/snowflake/common/util/SchemaHelper.java

@@ -109,7 +109,7 @@ public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String
     }
   }
 
-  private static Schema getParsedSchema(String schema) {
+  public static Schema getParsedSchema(String schema) {
     if (Strings.isNullOrEmpty(schema)) {
       return null;
     }

src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeSourceAccessor.java

@@ -18,10 +18,12 @@
 
 import au.com.bytecode.opencsv.CSVReader;
 import com.google.common.base.Strings;
+import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.plugin.snowflake.common.SnowflakeErrorType;
 import io.cdap.plugin.snowflake.common.client.SnowflakeAccessor;
 import io.cdap.plugin.snowflake.common.util.DocumentUrlUtil;
 import io.cdap.plugin.snowflake.common.util.QueryUtil;
+import io.cdap.plugin.snowflake.common.util.SchemaHelper;
 import net.snowflake.client.jdbc.SnowflakeConnection;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -35,6 +37,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.UUID;
+import java.util.stream.Collectors;
 
 /**
  * A class which accesses Snowflake API to do actions used by batch source.
@@ -81,7 +84,15 @@ public List<String> prepareStageSplits() {
     String importQuery = config.getImportQuery();
     if (Strings.isNullOrEmpty(importQuery)) {
       String tableName = config.getTableName();
-      importQuery = String.format("SELECT * FROM %s", tableName);
+      Schema schema = SchemaHelper.getParsedSchema(config.getSchema());
+      if (schema != null && schema.getFields() != null && !schema.getFields().isEmpty()) {
+        String columns = schema.getFields().stream()
+          .map(Schema.Field::getName)
+          .collect(Collectors.joining(","));
+        importQuery = String.format("SELECT %s FROM %s", columns, tableName);
+      } else {
+        importQuery = String.format("SELECT * FROM %s", tableName);
+      }
     }
     String copy = String.format(COMAND_COPY_INTO, QueryUtil.removeSemicolon(importQuery));
     if (config.getMaxSplitSize() > 0) {