Merge pull request #53 from data-integrations/bugfix_release/CDAP-14531-bigquery-npe

CDAP-14531 Fix NPE when outputschema and input record schema do not m…

Author: CuriousVini

src/main/java/co/cask/gcp/bigquery/BigQuerySink.java

@@ -89,12 +89,13 @@ public BigQuerySink(BigQuerySinkConfig config) {
 
   @Override
   public void configurePipeline(PipelineConfigurer pipelineConfigurer) {
+    config.validate(pipelineConfigurer.getStageConfigurer().getInputSchema());
     super.configurePipeline(pipelineConfigurer);
-    config.validate();
   }
 
   @Override
   public void prepareRun(BatchSinkContext context) throws Exception {
+    config.validate(context.getInputSchema());
     BigQuery bigquery = BigQueryUtils.getBigQuery(config.getServiceAccountFilePath(), config.getProject());
     // create dataset if it does not exist
     if (bigquery.getDataset(config.getDataset()) == null) {
@@ -105,6 +106,7 @@ public void prepareRun(BatchSinkContext context) throws Exception {
       }
     }
 
+    // schema validation against bigquery table schema
     validateSchema();
 
     uuid = UUID.randomUUID();
@@ -155,10 +157,12 @@ public void initialize(BatchRuntimeContext context) throws Exception {
 
   @Override
   public void transform(StructuredRecord input, Emitter<KeyValue<JsonObject, NullWritable>> emitter) throws Exception {
-    List<Schema.Field> fields = config.getSchema().getFields();
     JsonObject object = new JsonObject();
-    for (Schema.Field field : fields) {
-      decodeSimpleTypes(object, field.getName(), input);
+    for (Schema.Field recordField : input.getSchema().getFields()) {
+      // From all the fields in input record, decode only those fields that are present in output schema
+      if (schema.getField(recordField.getName()) != null) {
+        decodeSimpleTypes(object, recordField.getName(), input);
+      }
     }
     emitter.emit(new KeyValue<>(object, NullWritable.get()));
   }
@@ -339,20 +343,8 @@ private void validateSchema() throws IOException {
 
     // Match output schema field type with bigquery column type
     for (Schema.Field field : config.getSchema().getFields()) {
-      validateSimpleTypes(field);
       BigQueryUtils.validateFieldSchemaMatches(bqFields.get(field.getName()),
                                                field, config.getDataset(), config.getTable());
     }
   }
-
-  private void validateSimpleTypes(Schema.Field field) {
-    String name = field.getName();
-    Schema fieldSchema = BigQueryUtils.getNonNullableSchema(field.getSchema());
-    Schema.Type type = fieldSchema.getType();
-
-    // Complex types like arrays, maps and unions are not supported in BigQuery plugins.
-    if (!type.isSimpleType()) {
-      throw new IllegalArgumentException(String.format("Field '%s' is of unsupported type '%s'.", name, type));
-    }
-  }
 }

src/main/java/co/cask/gcp/bigquery/BigQuerySinkConfig.java

@@ -88,4 +88,31 @@ public Schema getSchema() {
       throw new IllegalArgumentException("Invalid schema: " + e.getMessage());
     }
   }
+
+  /**
+   * Verifies if output schema only contains simple types. It also verifies if all the output schema fields are
+   * present in input schema.
+   *
+   * @param inputSchema input schema to bigquery sink
+   */
+  public void validate(@Nullable Schema inputSchema) {
+    super.validate();
+    if (!containsMacro("schema")) {
+      Schema outputSchema = getSchema();
+      for (Schema.Field field : outputSchema.getFields()) {
+        // check if the required fields are present in the input schema.
+        if (!field.getSchema().isNullable() && inputSchema != null && inputSchema.getField(field.getName()) == null) {
+          throw new IllegalArgumentException(String.format("Required output field '%s' is not present in input schema.",
+                                                           field.getName()));
+        }
+
+        Schema fieldSchema = BigQueryUtils.getNonNullableSchema(field.getSchema());
+
+        if (!fieldSchema.getType().isSimpleType()) {
+          throw new IllegalArgumentException(String.format("Field '%s' is of unsupported type '%s'.",
+                                                           field.getName(), fieldSchema.getType()));
+        }
+      }
+    }
+  }
 }

src/test/java/co/cask/gcp/bigquery/BigQuerySinkTest.java

@@ -0,0 +1,51 @@
+/*
+ * Copyright © 2018 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 co.cask.gcp.bigquery;
+
+import co.cask.cdap.api.data.schema.Schema;
+import org.junit.Test;
+
+/**
+ * Tests for {@link BigQuerySink}.
+ */
+public class BigQuerySinkTest {
+
+  @Test
+  public void testBigQuerySinkConfig() {
+    Schema schema = Schema.recordOf("record",
+                                    Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
+                                    Schema.Field.of("name", Schema.of(Schema.Type.STRING)),
+                                    Schema.Field.of("price", Schema.of(Schema.Type.DOUBLE)),
+                                    Schema.Field.of("dt", Schema.nullableOf(Schema.of(Schema.LogicalType.DATE))),
+                                    Schema.Field.of("bytedata", Schema.of(Schema.Type.BYTES)),
+                                    Schema.Field.of("timestamp",
+                                                    Schema.nullableOf(Schema.of(Schema.LogicalType.TIMESTAMP_MICROS))));
+
+    BigQuerySinkConfig config = new BigQuerySinkConfig("r", "ds", "tb", "bucket", schema.toString());
+    config.validate();
+  }
+
+  @Test(expected = IllegalArgumentException.class)
+  public void testBigQuerySinkInvalidConfig() {
+    Schema invalidSchema = Schema.recordOf("record",
+                                    Schema.Field.of("id", Schema.of(Schema.Type.LONG)),
+                                    Schema.Field.of("record", Schema.of(Schema.Type.RECORD)));
+
+    BigQuerySinkConfig config = new BigQuerySinkConfig("r", "ds", "tb", "bucket", invalidSchema.toString());
+    config.validate();
+  }
+}