Fixed null output schema

Author: psainics

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

@@ -58,7 +58,7 @@ private SchemaHelper() {
 
   public static Schema getSchema(SnowflakeBatchSourceConfig config, FailureCollector collector) {
     if (!config.canConnect()) {
-      return null;
+      return getParsedSchema(config.getSchema());
     }
 
     SnowflakeSourceAccessor snowflakeSourceAccessor = new SnowflakeSourceAccessor(config);
@@ -69,11 +69,7 @@ public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String
                                  FailureCollector collector, String importQuery) {
     try {
       if (!Strings.isNullOrEmpty(schema)) {
-        try {
-          return Schema.parseJson(schema);
-        } catch (IOException | IllegalStateException e) {
-          throw new SchemaParseException(e);
-        }
+        return getParsedSchema(schema);
       }
       return Strings.isNullOrEmpty(importQuery) ? null : getSchema(snowflakeAccessor, importQuery);
     } catch (SchemaParseException e) {
@@ -85,6 +81,17 @@ public static Schema getSchema(SnowflakeSourceAccessor snowflakeAccessor, String
     }
   }
 
+  private static Schema getParsedSchema(String schema) {
+    if (Strings.isNullOrEmpty(schema)) {
+      return null;
+    }
+    try {
+      return Schema.parseJson(schema);
+    } catch (IOException | IllegalStateException e) {
+      throw new SchemaParseException(e);
+    }
+  }
+
   public static Schema getSchema(SnowflakeAccessor snowflakeAccessor, String importQuery) {
     try {
       List<SnowflakeFieldDescriptor> result = snowflakeAccessor.describeQuery(importQuery);

src/test/java/io/cdap/plugin/snowflake/common/util/SchemaHelperTest.java

@@ -25,6 +25,7 @@
 import org.junit.Assert;
 import org.junit.Test;
 import org.mockito.Mockito;
+
 import java.io.IOException;
 import java.sql.Types;
 import java.util.ArrayList;
@@ -144,4 +145,37 @@ public void testGetSchemaFromSnowflake() throws IOException {
     Assert.assertTrue(collector.getValidationFailures().isEmpty());
     Assert.assertEquals(expected, actual);
   }
+
+  @Test
+  public void testGetSchemaWhenMacroIsEnabled() {
+    Schema expected = Schema.recordOf("test",
+            Schema.Field.of("test_field", Schema.nullableOf(Schema.of(Schema.Type.LONG)))
+    );
+
+    SnowflakeBatchSourceConfig mockConfig = Mockito.mock(SnowflakeBatchSourceConfig.class);
+    Mockito.when(mockConfig.canConnect()).thenReturn(false);
+    Mockito.when(mockConfig.getSchema()).thenReturn(expected.toString());
+
+    MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE);
+    Schema actual = SchemaHelper.getSchema(mockConfig, collector);
+
+    Assert.assertTrue(collector.getValidationFailures().isEmpty());
+    Assert.assertEquals(expected, actual);
+
+  }
+
+  @Test
+  public void testGetSchemaWhenMacroIsEnabledSchemaIsNull() {
+
+    SnowflakeBatchSourceConfig mockConfig = Mockito.mock(SnowflakeBatchSourceConfig.class);
+    Mockito.when(mockConfig.canConnect()).thenReturn(false);
+    Mockito.when(mockConfig.getSchema()).thenReturn(null);
+
+    MockFailureCollector collector = new MockFailureCollector(MOCK_STAGE);
+    Schema actual = SchemaHelper.getSchema(mockConfig, collector);
+
+    Assert.assertTrue(collector.getValidationFailures().isEmpty());
+    Assert.assertNull(actual);
+
+  }
 }