PLUGIN-1883: Applied patch for Adding Named Table option to fetch schema

Author: sgarg-CS

docs/Snowflake-batchsource.md

@@ -25,7 +25,9 @@ log in to Snowflake, minus the "snowflakecomputing.com"). E.g. "myaccount.us-cen
 
 **Role:** Role to use (e.g. `ACCOUNTADMIN`).
 
-**Import Query:** Query for data import.
+**Import Query Type**  - Method used to retrieve schema from the source.
+* **Table Name**: The name of the table to retrieve the schema.
+* **Import Query**: Query for data import.
 
 ### Credentials
 

pom.xml

@@ -29,8 +29,8 @@
     <surefire.redirectTestOutputToFile>true</surefire.redirectTestOutputToFile>
     <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
     <!-- version properties -->
-    <cdap.version>6.11.0-SNAPSHOT</cdap.version>
-    <hydrator.version>2.13.0-SNAPSHOT</hydrator.version>
+    <cdap.version>6.11.0</cdap.version>
+    <hydrator.version>2.13.0</hydrator.version>
     <commons.csv.version>1.6</commons.csv.version>
     <hadoop.version>3.3.6</hadoop.version>
     <spark3.version>3.3.2</spark3.version>

src/main/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessor.java

@@ -37,6 +37,7 @@
 import java.io.IOException;
 import java.lang.reflect.Field;
 import java.sql.Connection;
+import java.sql.DatabaseMetaData;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
@@ -74,6 +75,38 @@ public void runSQL(String query) {
     }
   }
 
+  /**
+   * Returns field descriptors for specified table name
+   * @param schemaName The name of schema containing the table
+   * @param tableName The name of table whose metadata needs to be retrieved
+   * @return list of field descriptors
+   * @throws SQLException If an error occurs while retrieving metadata from the database
+   */
+  public List<SnowflakeFieldDescriptor> describeTable(String schemaName, String tableName) throws SQLException {
+    List<SnowflakeFieldDescriptor> fieldDescriptors = new ArrayList<>();
+    try (Connection connection = dataSource.getConnection()) {
+      DatabaseMetaData dbMetaData = connection.getMetaData();
+      try (ResultSet columns = dbMetaData.getColumns(null, schemaName, tableName, null)) {
+        while (columns.next()) {
+          String columnName = columns.getString("COLUMN_NAME");
+          int columnType = columns.getInt("DATA_TYPE");
+          boolean nullable = columns.getInt("NULLABLE") == DatabaseMetaData.columnNullable;
+          fieldDescriptors.add(new SnowflakeFieldDescriptor(columnName, columnType, nullable));
+        }
+      }
+    } catch (SQLException e) {
+      String errorMessage = String.format(
+        "Failed to retrieve table metadata with SQL State %s and error code %s with message: %s.",
+        e.getSQLState(), e.getErrorCode(), e.getMessage()
+      );
+      String errorReason = String.format("Failed to retrieve table metadata with SQL State %s and error " +
+                                           "code %s. For more details %s", e.getSQLState(), e.getErrorCode(),
+                                         DocumentUrlUtil.getSupportedDocumentUrl());
+      throw SnowflakeErrorType.fetchProgramFailureException(e, errorReason, errorMessage);
+    }
+    return fieldDescriptors;
+  }
+
   /**
    * Returns field descriptors for specified import query.
    *
@@ -193,4 +226,13 @@ private static String writeTextToTmpFile(String text) {
       throw new RuntimeException("Cannot write key to temporary file", e);
     }
   }
+
+  /**
+   * Retrieves schema name from the configuration
+   *
+   * @return The schema name
+   */
+  public String getSchema() {
+    return config.getSchemaName();
+  }
 }

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

@@ -16,6 +16,8 @@
 
 package io.cdap.plugin.snowflake.common.util;
 
+import com.google.common.base.Strings;
+
 /**
  * Transforms import query.
  */
@@ -29,6 +31,9 @@ private QueryUtil() {
   }
 
   public static String removeSemicolon(String importQuery) {
+    if (Strings.isNullOrEmpty(importQuery)) {
+      return null;
+    }
     if (importQuery.endsWith(";")) {
       importQuery = importQuery.substring(0, importQuery.length() - 1);
     }

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

@@ -27,11 +27,13 @@
 import io.cdap.plugin.snowflake.source.batch.SnowflakeInputFormatProvider;
 import io.cdap.plugin.snowflake.source.batch.SnowflakeSourceAccessor;
 import java.io.IOException;
+import java.sql.SQLException;
 import java.sql.Types;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
 import java.util.stream.Collectors;
+import javax.annotation.Nullable;
 
 /**
  * Resolves schema.
@@ -58,24 +60,47 @@ public class SchemaHelper {
   private SchemaHelper() {
   }
 
+  /**
+   * Retrieves schema for the Snowflake batch source based on the given configuration.
+   *
+   * @param config    The configuration for Snowflake batch source
+   * @param collector The failure collector to capture any schema retrieval errors.
+   * @return The resolved schema for Snowflake source
+   */
   public static Schema getSchema(SnowflakeBatchSourceConfig config, FailureCollector collector) {
     if (!config.canConnect()) {
       return getParsedSchema(config.getSchema());
     }
 
     SnowflakeSourceAccessor snowflakeSourceAccessor =
       new SnowflakeSourceAccessor(config, SnowflakeInputFormatProvider.PROPERTY_DEFAULT_ESCAPE_CHAR);
-    return getSchema(snowflakeSourceAccessor, config.getSchema(), collector, config.getImportQuery());
+    return getSchema(
+      snowflakeSourceAccessor,
+      config.getSchema(),
+      collector,
+      config.getTableName(),
+      config.getImportQuery()
+    );
   }
 
+  /**
+   * Retrieves schema for a Snowflake source based on the provided parameters.
+   *
+   * @param snowflakeAccessor The {@link SnowflakeSourceAccessor} used to connect to Snowflake.
+   * @param schema            A JSON-format schema string
+   * @param collector         The {@link FailureCollector} to collect errors if schema retrieval fails.
+   * @param tableName         The name of the table in Snowflake.
+   * @param importQuery       The query to fetch data from Snowflake, used when `tableName` is not provided.
+   * @return The parsed {@link Schema} if successful, or {@code null} if an error occurs.
+   */
   public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String schema,
-                                 FailureCollector collector, String importQuery) {
+                                 FailureCollector collector, String tableName, String importQuery) {
     try {
       if (!Strings.isNullOrEmpty(schema)) {
         return getParsedSchema(schema);
       }
-      return Strings.isNullOrEmpty(importQuery) ? null : getSchema(snowflakeAccessor, importQuery);
-    } catch (SchemaParseException e) {
+      return getSchema(snowflakeAccessor, tableName, importQuery);
+    } catch (SchemaParseException | IllegalArgumentException e) {
       collector.addFailure(String.format("Unable to retrieve output schema. Reason: '%s'", e.getMessage()),
                            null)
         .withStacktrace(e.getStackTrace())
@@ -95,6 +120,26 @@ private static Schema getParsedSchema(String schema) {
     }
   }
 
+  private static Schema getSchema(SnowflakeAccessor snowflakeAccessor, @Nullable String tableName,
+                                  @Nullable String importQuery) {
+    try {
+      List<SnowflakeFieldDescriptor> result;
+      if (!Strings.isNullOrEmpty(tableName)) {
+        result = snowflakeAccessor.describeTable(snowflakeAccessor.getSchema(), tableName);
+      } else {
+        result = snowflakeAccessor.describeQuery(importQuery);
+      }
+      List<Schema.Field> fields = result.stream()
+        .map(fieldDescriptor -> Schema.Field.of(fieldDescriptor.getName(), getSchema(fieldDescriptor)))
+        .collect(Collectors.toList());
+      return Schema.recordOf("data", fields);
+    } catch (SQLException e) {
+      throw new SchemaParseException(e);
+    } catch (IOException e) {
+      throw new RuntimeException(e);
+    }
+  }
+
   public static Schema getSchema(SnowflakeAccessor snowflakeAccessor, String importQuery) {
     try {
       List<SnowflakeFieldDescriptor> result = snowflakeAccessor.describeQuery(importQuery);

src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java

@@ -23,7 +23,6 @@
 import io.cdap.plugin.snowflake.common.BaseSnowflakeConfig;
 import io.cdap.plugin.snowflake.common.client.SnowflakeAccessor;
 import io.cdap.plugin.snowflake.common.util.SchemaHelper;
-
 import javax.annotation.Nullable;
 
 /**

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

@@ -0,0 +1,36 @@
+/*
+ * Copyright © 2020 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.snowflake.source.batch;
+
+
+/**
+ * Enum to specify the import query type used in Snowflake Batch Source.
+ */
+public enum ImportQueryType {
+  IMPORT_QUERY,
+  TABLE_NAME;
+
+  public static ImportQueryType fromString(String value) {
+    for (ImportQueryType type : ImportQueryType.values()) {
+      if (type.name().equalsIgnoreCase(value)) {
+        return type;
+      }
+    }
+    throw new IllegalArgumentException(
+      String.format("Unsupported importQueryType '%s'. Supported types are: importQuery, tableName.", value));
+  }
+}

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

@@ -16,12 +16,12 @@
 
 package io.cdap.plugin.snowflake.source.batch;
 
+import com.google.common.base.Strings;
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Macro;
 import io.cdap.cdap.api.annotation.Name;
 import io.cdap.cdap.etl.api.FailureCollector;
 import io.cdap.plugin.snowflake.common.BaseSnowflakeConfig;
-
 import java.util.Objects;
 import javax.annotation.Nullable;
 
@@ -34,6 +34,7 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig {
   public static final String PROPERTY_IMPORT_QUERY = "importQuery";
   public static final String PROPERTY_MAX_SPLIT_SIZE = "maxSplitSize";
   public static final String PROPERTY_SCHEMA = "schema";
+  public static final String PROPERTY_TABLE_NAME = "tableName";
 
   @Name(PROPERTY_REFERENCE_NAME)
   @Description("This will be used to uniquely identify this source/sink for lineage, annotating metadata, etc.")
@@ -42,6 +43,7 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig {
   @Name(PROPERTY_IMPORT_QUERY)
   @Description("Query for import data.")
   @Macro
+  @Nullable
   private String importQuery;
 
   @Name(PROPERTY_MAX_SPLIT_SIZE)
@@ -55,19 +57,39 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig {
   @Macro
   private String schema;
 
+
+  @Name(PROPERTY_TABLE_NAME)
+  @Nullable
+  @Description("The name of the table used to retrieve the schema.")
+  private final String tableName;
+
+  @Name("importQueryType")
+  @Description("Whether to select Table Name or Import Query to extract the data.")
+  @Macro
+  @Nullable
+  private final ImportQueryType importQueryType;
+
+
   public SnowflakeBatchSourceConfig(String referenceName, String accountName, String database,
-                                    String schemaName, String importQuery, String username, String password,
+                                    String schemaName, @Nullable String importQuery, String username, String password,
                                     @Nullable Boolean keyPairEnabled, @Nullable String path,
                                     @Nullable String passphrase, @Nullable Boolean oauth2Enabled,
                                     @Nullable String clientId, @Nullable String clientSecret,
                                     @Nullable String refreshToken, Long maxSplitSize,
-                                    @Nullable String connectionArguments, @Nullable String schema) {
-    super(accountName, database, schemaName, username, password,
-          keyPairEnabled, path, passphrase, oauth2Enabled, clientId, clientSecret, refreshToken, connectionArguments);
+                                    @Nullable String connectionArguments,
+                                    @Nullable String schema,
+                                    @Nullable String tableName,
+                                    @Nullable ImportQueryType importQueryType) {
+    super(
+      accountName, database, schemaName, username, password, keyPairEnabled, path, passphrase,
+      oauth2Enabled, clientId, clientSecret, refreshToken, connectionArguments
+    );
     this.referenceName = referenceName;
     this.importQuery = importQuery;
     this.maxSplitSize = maxSplitSize;
     this.schema = schema;
+    this.tableName = tableName;
+    this.importQueryType = importQueryType;
   }
 
   public String getImportQuery() {
@@ -87,6 +109,15 @@ public String getSchema() {
     return schema;
   }
 
+  @Nullable
+  public String getTableName() {
+    return tableName;
+  }
+
+  public ImportQueryType getImportQueryType() {
+    return importQueryType;
+  }
+
   public void validate(FailureCollector collector) {
     super.validate(collector);
 
@@ -95,5 +126,11 @@ public void validate(FailureCollector collector) {
       collector.addFailure("Maximum Slit Size cannot be a negative number.", null)
         .withConfigProperty(PROPERTY_MAX_SPLIT_SIZE);
     }
+
+    if (Strings.isNullOrEmpty(importQuery) && Strings.isNullOrEmpty(tableName)) {
+      collector.addFailure("Either 'Schema' or 'Table Name' must be provided.", null)
+        .withConfigProperty(PROPERTY_IMPORT_QUERY)
+        .withConfigProperty(PROPERTY_TABLE_NAME);
+    }
   }
 }

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

@@ -17,6 +17,7 @@
 package io.cdap.plugin.snowflake.source.batch;
 
 import au.com.bytecode.opencsv.CSVReader;
+import com.google.common.base.Strings;
 import io.cdap.plugin.snowflake.common.SnowflakeErrorType;
 import io.cdap.plugin.snowflake.common.client.SnowflakeAccessor;
 import io.cdap.plugin.snowflake.common.util.DocumentUrlUtil;
@@ -77,7 +78,12 @@ public SnowflakeSourceAccessor(SnowflakeBatchSourceConfig config, String escapeC
    */
   public List<String> prepareStageSplits() {
     LOG.info("Loading data into stage: '{}'", STAGE_PATH);
-    String copy = String.format(COMAND_COPY_INTO, QueryUtil.removeSemicolon(config.getImportQuery()));
+    String importQuery = config.getImportQuery();
+    if (Strings.isNullOrEmpty(importQuery)) {
+      String tableName = config.getTableName();
+      importQuery = String.format("SELECT * FROM %s", tableName);
+    }
+    String copy = String.format(COMAND_COPY_INTO, QueryUtil.removeSemicolon(importQuery));
     if (config.getMaxSplitSize() > 0) {
       copy = copy + String.format(COMMAND_MAX_FILE_SIZE, config.getMaxSplitSize());
     }
@@ -94,10 +100,13 @@ public List<String> prepareStageSplits() {
       }
     } catch (SQLException e) {
       String errorReason = String.format("Failed to load data into stage '%s' with sqlState %s and errorCode %s. " +
-          "For more details, see %s.", STAGE_PATH, e.getErrorCode(), e.getSQLState(),
-        DocumentUrlUtil.getSupportedDocumentUrl());
-      String errorMessage = String.format("Failed to load data into stage '%s' with sqlState %s and errorCode %s. " +
-        "Failed to execute query with message: %s.", STAGE_PATH, e.getSQLState(), e.getErrorCode(), e.getMessage());
+                                           "For more details, see %s.", STAGE_PATH, e.getErrorCode(), e.getSQLState(),
+                                         DocumentUrlUtil.getSupportedDocumentUrl());
+      String errorMessage = String.format(
+        "Failed to load data into stage '%s' with sqlState %s and errorCode %s. "
+          + "Failed to execute query with message: %s.",
+        STAGE_PATH, e.getSQLState(), e.getErrorCode(), e.getMessage()
+      );
       throw SnowflakeErrorType.fetchProgramFailureException(e, errorReason, errorMessage);
     }
     return stageSplits;