PLUGIN- 1883: Rework

Author: sgarg-CS

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

@@ -100,7 +100,7 @@ public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String
         return getParsedSchema(schema);
       }
       return getSchema(snowflakeAccessor, tableName, importQuery);
-    } catch (SchemaParseException | IllegalArgumentException e) {
+    } catch (SchemaParseException e) {
       collector.addFailure(String.format("Unable to retrieve output schema. Reason: '%s'", e.getMessage()),
                            null)
         .withStacktrace(e.getStackTrace())
@@ -115,7 +115,7 @@ private static Schema getParsedSchema(String schema) {
     }
     try {
       return Schema.parseJson(schema);
-    } catch (IOException | IllegalStateException e) {
+    } catch (IOException e) {
       throw new SchemaParseException(e);
     }
   }
@@ -126,8 +126,10 @@ private static Schema getSchema(SnowflakeAccessor snowflakeAccessor, @Nullable S
       List<SnowflakeFieldDescriptor> result;
       if (!Strings.isNullOrEmpty(tableName)) {
         result = snowflakeAccessor.describeTable(snowflakeAccessor.getSchema(), tableName);
-      } else {
+      } else if (!Strings.isNullOrEmpty(importQuery)) {
         result = snowflakeAccessor.describeQuery(importQuery);
+      } else {
+        throw new SchemaParseException("Either 'Import Query' or 'Table Name' must be provided");
       }
       List<Schema.Field> fields = result.stream()
         .map(fieldDescriptor -> Schema.Field.of(fieldDescriptor.getName(), getSchema(fieldDescriptor)))
@@ -136,7 +138,7 @@ private static Schema getSchema(SnowflakeAccessor snowflakeAccessor, @Nullable S
     } catch (SQLException e) {
       throw new SchemaParseException(e);
     } catch (IOException e) {
-      throw new RuntimeException(e);
+      throw new SchemaParseException(e);
     }
   }
 

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

@@ -21,20 +21,29 @@
  * Enum to specify the import query type used in Snowflake Batch Source.
  */
 public enum ImportQueryType {
-  IMPORT_QUERY,
-  TABLE_NAME;
+  IMPORT_QUERY ("importQuery"),
+  TABLE_NAME ("tableName");
+
+  private String value;
+
+  ImportQueryType(String value) {
+    this.value = value;
+  }
+
+  public String getValue() {
+    return value;
+  }
 
   public static ImportQueryType fromString(String value) {
     if (value == null) {
       return ImportQueryType.IMPORT_QUERY;
     }
 
     for (ImportQueryType type : ImportQueryType.values()) {
-      if (type.name().equalsIgnoreCase(value)) {
+      if (type.value.equalsIgnoreCase(value)) {
         return type;
       }
     }
-    
     return ImportQueryType.IMPORT_QUERY;
   }
 }

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

@@ -36,6 +36,7 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig {
   public static final String PROPERTY_MAX_SPLIT_SIZE = "maxSplitSize";
   public static final String PROPERTY_SCHEMA = "schema";
   public static final String PROPERTY_TABLE_NAME = "tableName";
+  public static final String PROPERTY_IMPORT_QUERY_TYPE = "importQueryType";
 
   @Name(PROPERTY_REFERENCE_NAME)
   @Description("This will be used to uniquely identify this source/sink for lineage, annotating metadata, etc.")
@@ -65,11 +66,11 @@ public class SnowflakeBatchSourceConfig extends BaseSnowflakeConfig {
   @Nullable
   private final String tableName;
 
-  @Name("importQueryType")
+  @Name(PROPERTY_IMPORT_QUERY_TYPE)
   @Description("Whether to select Table Name or Import Query to extract the data.")
   @Macro
   @Nullable
-  private final ImportQueryType importQueryType;
+  private final String importQueryType;
 
 
   public SnowflakeBatchSourceConfig(String referenceName, String accountName, String database,
@@ -81,7 +82,7 @@ public SnowflakeBatchSourceConfig(String referenceName, String accountName, Stri
                                     @Nullable String connectionArguments,
                                     @Nullable String schema,
                                     @Nullable String tableName,
-                                    @Nullable ImportQueryType importQueryType) {
+                                    @Nullable String importQueryType) {
     super(
       accountName, database, schemaName, username, password, keyPairEnabled, path, passphrase,
       oauth2Enabled, clientId, clientSecret, refreshToken, connectionArguments
@@ -91,7 +92,7 @@ public SnowflakeBatchSourceConfig(String referenceName, String accountName, Stri
     this.maxSplitSize = maxSplitSize;
     this.schema = schema;
     this.tableName = tableName;
-    this.importQueryType = importQueryType;
+    this.importQueryType = getImportQueryType();
   }
 
   public String getImportQuery() {
@@ -117,8 +118,8 @@ public String getTableName() {
   }
 
   @Nullable
-  public ImportQueryType getImportQueryType() {
-    return importQueryType == null ? ImportQueryType.IMPORT_QUERY : importQueryType;
+  public String getImportQueryType() {
+    return importQueryType == null ? ImportQueryType.IMPORT_QUERY.name() : importQueryType;
   }
 
   public void validate(FailureCollector collector) {
@@ -130,10 +131,27 @@ public void validate(FailureCollector collector) {
         .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);
+    if (!containsMacro(PROPERTY_IMPORT_QUERY_TYPE)) {
+      boolean isImportQuerySelected = ImportQueryType.IMPORT_QUERY.getValue().equals(importQueryType);
+      boolean isTableNameSelected = ImportQueryType.TABLE_NAME.getValue().equals(importQueryType);
+
+      if (isImportQuerySelected && !containsMacro(PROPERTY_IMPORT_QUERY) && Strings.isNullOrEmpty(importQuery)) {
+        collector.addFailure("Import Query cannot be empty", null)
+          .withConfigProperty(PROPERTY_IMPORT_QUERY);
+
+      } else if (isTableNameSelected && !containsMacro(PROPERTY_TABLE_NAME) && Strings.isNullOrEmpty(tableName)) {
+        collector.addFailure("Table Name cannot be empty", null)
+          .withConfigProperty(PROPERTY_TABLE_NAME);
+      }
+    } else {
+      boolean isImportQueryMissing = !containsMacro(PROPERTY_IMPORT_QUERY) && Strings.isNullOrEmpty(importQuery);
+      boolean isTableNameMissing = !containsMacro(PROPERTY_TABLE_NAME) && Strings.isNullOrEmpty(tableName);
+
+      if (isImportQueryMissing && isTableNameMissing) {
+        collector.addFailure("Either 'Import Query' or 'Table Name' must be provided.", null)
+          .withConfigProperty(PROPERTY_IMPORT_QUERY)
+          .withConfigProperty(PROPERTY_TABLE_NAME);
+      }
     }
   }
 }

src/test/java/io/cdap/plugin/snowflake/source/batch/SnowflakeBatchSourceConfigBuilder.java

@@ -40,7 +40,7 @@ public class SnowflakeBatchSourceConfigBuilder {
     "",
     "",
     "",
-    ImportQueryType.IMPORT_QUERY);
+    ImportQueryType.IMPORT_QUERY.name());
 
   private String referenceName;
   private String accountName;
@@ -60,7 +60,7 @@ public class SnowflakeBatchSourceConfigBuilder {
   private String connectionArguments;
   private String schema;
   private String tableName;
-  private ImportQueryType importQueryType;
+  private String importQueryType;
 
   public SnowflakeBatchSourceConfigBuilder() {
   }
@@ -177,7 +177,7 @@ public SnowflakeBatchSourceConfigBuilder setTableName(String tableName) {
     return this;
   }
 
-  public SnowflakeBatchSourceConfigBuilder setImportQueryType(ImportQueryType importQueryType) {
+  public SnowflakeBatchSourceConfigBuilder setImportQueryType(String importQueryType) {
     this.importQueryType = importQueryType;
     return this;
   }