Get schema without import query

AnkitCLI · AnkitCLI · commit 89bb5dfe2eba · 2025-04-14T20:38:24.000+05:30
diff --git a/docs/Snowflake-batchsource.md b/docs/Snowflake-batchsource.md
@@ -25,6 +25,8 @@ log in to Snowflake, minus the "snowflakecomputing.com"). E.g. "myaccount.us-cen
 
 **Role:** Role to use (e.g. `ACCOUNTADMIN`).
 
+**Table Name:** The name of the table to retrieve the schema.
+
 **Import Query:** Query for data import.
 
 ### Credentials
diff --git a/src/main/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessor.java b/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,39 @@ 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.
    *
@@ -163,6 +197,7 @@ public void checkConnection() {
         errorReason, errorMessage, ErrorType.USER, true, e);
     }
   }
+
   // SnowflakeBasicDataSource doesn't provide access for additional properties.
   private void addConnectionArguments(SnowflakeBasicDataSource dataSource, String connectionArguments) {
     try {
@@ -193,4 +228,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();
+  }
 }
diff --git a/src/main/java/io/cdap/plugin/snowflake/common/util/QueryUtil.java b/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);
     }
diff --git a/src/main/java/io/cdap/plugin/snowflake/common/util/SchemaHelper.java b/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,25 @@ 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);
diff --git a/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java b/src/main/java/io/cdap/plugin/snowflake/sink/batch/SnowflakeSinkConfig.java
@@ -44,6 +44,7 @@ public class SnowflakeSinkConfig extends BaseSnowflakeConfig {
   @Name(PROPERTY_TABLE_NAME)
   @Description("Name of the table to insert records into.")
   @Macro
+  @Nullable
   private String tableName;
 
   @Name(PROPERTY_MAX_FILE_SIZE)
diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfig.java b/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,29 @@ 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;
+
   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) {
+    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;
   }
 
   public String getImportQuery() {
@@ -87,6 +99,11 @@ public String getSchema() {
     return schema;
   }
 
+  @Nullable
+  public String getTableName() {
+    return tableName;
+  }
+
   public void validate(FailureCollector collector) {
     super.validate(collector);
 
@@ -95,5 +112,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);
+    }
   }
 }
diff --git a/src/main/java/io/cdap/plugin/snowflake/source/batch/SnowflakeSourceAccessor.java b/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;
diff --git a/src/test/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessorTest.java b/src/test/java/io/cdap/plugin/snowflake/common/client/SnowflakeAccessorTest.java
@@ -90,6 +90,41 @@ public void testDescribeQuery() throws Exception {
     Assert.assertEquals(expected, actual);
   }
 
+  @Test
+  public void testDescribeTable() throws Exception {
+    String schemaName = "TEST_SCHEMA";
+    String tableName = "TEST_TABLE";
+    List<SnowflakeFieldDescriptor> expected = Arrays.asList(
+      new SnowflakeFieldDescriptor("COLUMN_NUMBER", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_DECIMAL", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_NUMERIC", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_INT", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_INTEGER", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_BIGINT", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_SMALLINT", -5, true),
+      new SnowflakeFieldDescriptor("COLUMN_FLOAT", 8, true),
+      new SnowflakeFieldDescriptor("COLUMN_DOUBLE", 8, true),
+      new SnowflakeFieldDescriptor("COLUMN_REAL", 8, true),
+      new SnowflakeFieldDescriptor("COLUMN_VARCHAR", 12, true),
+      new SnowflakeFieldDescriptor("COLUMN_CHAR", 12, true),
+      new SnowflakeFieldDescriptor("COLUMN_TEXT", 12, true),
+      new SnowflakeFieldDescriptor("COLUMN_BINARY", -2, true),
+      new SnowflakeFieldDescriptor("COLUMN_BOOLEAN", 16, true),
+      new SnowflakeFieldDescriptor("COLUMN_DATE", 91, true),
+      new SnowflakeFieldDescriptor("COLUMN_TIMESTAMP", 93, true),
+      new SnowflakeFieldDescriptor("COLUMN_VARIANT", 12, true),
+      new SnowflakeFieldDescriptor("COLUMN_OBJECT", 12, true),
+      new SnowflakeFieldDescriptor("COLUMN_ARRAY", 12, true)
+    );
+
+    List<SnowflakeFieldDescriptor> actual = snowflakeAccessor.describeTable(String.valueOf(Constants.TEST_TABLE_SCHEMA),
+                                                                            Constants.TEST_TABLE);
+
+    Assert.assertNotNull(actual);
+    Assert.assertFalse(actual.isEmpty());
+    Assert.assertEquals(expected, actual);
+  }
+
   @Test
   public void testPrepareStageSplits() throws Exception {
     Pattern expected = Pattern.compile("cdap_stage/result.*data__0_0_0\\.csv\\.gz");
diff --git a/src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java b/src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java
diff --git a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java
diff --git a/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigTest.java b/src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigTest.java
diff --git a/widgets/Snowflake-batchsource.json b/widgets/Snowflake-batchsource.json
