CDAP-15546 Oracle db plugin enhancements: all data types support and proper test coverage

Author: vhlinskyi

aurora-postgresql-plugin/src/main/java/io/cdap/plugin/auroradb/postgres/AuroraPostgresDBRecord.java

@@ -18,18 +18,20 @@
 
 import io.cdap.cdap.api.data.format.StructuredRecord;
 import io.cdap.cdap.api.data.schema.Schema;
+import io.cdap.plugin.db.ColumnType;
 import io.cdap.plugin.db.DBRecord;
 import io.cdap.plugin.db.SchemaReader;
 
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.List;
 
 /**
  * Writable class for Aurora DB PostgreSQL Source/Sink
  */
 public class AuroraPostgresDBRecord extends DBRecord {
 
-  public AuroraPostgresDBRecord(StructuredRecord record, int[] columnTypes) {
+  public AuroraPostgresDBRecord(StructuredRecord record, List<ColumnType> columnTypes) {
     super(record, columnTypes);
   }
 

aurora-postgresql-plugin/src/main/java/io/cdap/plugin/auroradb/postgres/AuroraPostgresSink.java

@@ -53,7 +53,6 @@ public AuroraPostgresSink(AuroraPostgresSinkConfig auroraPostgresSinkConfig) {
   protected void setColumnsInfo(List<Schema.Field> fields) {
     List<String> columnsList = new ArrayList<>();
     StringJoiner columnsJoiner = new StringJoiner(",");
-
     for (Schema.Field field : fields) {
       columnsList.add(field.getName());
       columnsJoiner.add(ESCAPE_CHAR + field.getName() + ESCAPE_CHAR);

database-commons/src/main/java/io/cdap/plugin/db/ColumnType.java

@@ -0,0 +1,39 @@
+/*
+ * Copyright © 2019 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.db;
+
+/**
+ * Stores SQL column name and type.
+ */
+public class ColumnType {
+
+  private final String name;
+  private final int type;
+
+  public ColumnType(String name, int type) {
+    this.name = name;
+    this.type = type;
+  }
+
+  public String getName() {
+    return name;
+  }
+
+  public int getType() {
+    return type;
+  }
+}

database-commons/src/main/java/io/cdap/plugin/db/DBRecord.java

@@ -44,6 +44,7 @@
 import java.time.ZoneId;
 import java.time.ZoneOffset;
 import java.util.List;
+import javax.annotation.Nullable;
 import javax.sql.rowset.serial.SerialBlob;
 
 /**
@@ -59,18 +60,16 @@ public class DBRecord implements Writable, DBWritable, Configurable {
   private final Lazy<Schema> schema = new Lazy<>(this::computeSchema);
 
   /**
-   * Need to cache {@link ResultSetMetaData} of the record for use during writing to a table.
-   * This is because we cannot rely on JDBC drivers to properly set metadata in the {@link PreparedStatement}
-   * passed to the #write method in this class.
+   * Need to cache column types to set fields of the input record on {@link PreparedStatement} in the right order.
    */
-  protected int[] columnTypes;
+  protected List<ColumnType> columnTypes;
 
   /**
    * Used to construct a DBRecord from a StructuredRecord in the ETL Pipeline
    *
    * @param record the {@link StructuredRecord} to construct the {@link DBRecord} from
    */
-  public DBRecord(StructuredRecord record, int[] columnTypes) {
+  public DBRecord(StructuredRecord record, List<ColumnType> columnTypes) {
     this.record = record;
     this.columnTypes = columnTypes;
   }
@@ -114,7 +113,7 @@ public void readFields(ResultSet resultSet) throws SQLException {
     record = recordBuilder.build();
   }
 
-  private Schema getSchema() {
+  protected Schema getSchema() {
     return schema.getOrCompute();
   }
 
@@ -202,10 +201,10 @@ public void write(DataOutput out) throws IOException {
    * @param stmt the {@link PreparedStatement} to write the {@link StructuredRecord} to
    */
   public void write(PreparedStatement stmt) throws SQLException {
-    Schema recordSchema = record.getSchema();
-    List<Schema.Field> schemaFields = recordSchema.getFields();
-    for (int i = 0; i < schemaFields.size(); i++) {
-      writeToDB(stmt, schemaFields.get(i), i);
+    for (int i = 0; i < columnTypes.size(); i++) {
+      ColumnType columnType = columnTypes.get(i);
+      Schema.Field field = record.getSchema().getField(columnType.getName());
+      writeToDB(stmt, field, i);
     }
   }
 
@@ -264,16 +263,24 @@ private void writeToDataOut(DataOutput out, Schema.Field field) throws IOExcepti
     }
   }
 
-  protected void writeToDB(PreparedStatement stmt, Schema.Field field, int fieldIndex) throws SQLException {
+  protected void writeToDB(PreparedStatement stmt, @Nullable Schema.Field field, int fieldIndex) throws SQLException {
+
+    int sqlIndex = fieldIndex + 1;
+    int sqlType = columnTypes.get(fieldIndex).getType();
+    if (field == null) {
+      // Some of the fields can be absent in the record
+      stmt.setNull(sqlIndex, sqlType);
+      return;
+    }
+
     String fieldName = field.getName();
     Schema fieldSchema = getNonNullableSchema(field);
     Schema.Type fieldType = fieldSchema.getType();
     Schema.LogicalType fieldLogicalType = fieldSchema.getLogicalType();
     Object fieldValue = record.get(fieldName);
-    int sqlIndex = fieldIndex + 1;
 
     if (fieldValue == null) {
-      stmt.setNull(sqlIndex, columnTypes[fieldIndex]);
+      stmt.setNull(sqlIndex, columnTypes.get(fieldIndex).getType());
       return;
     }
 
@@ -299,7 +306,7 @@ protected void writeToDB(PreparedStatement stmt, Schema.Field field, int fieldIn
 
     switch (fieldType) {
       case NULL:
-        stmt.setNull(sqlIndex, columnTypes[fieldIndex]);
+        stmt.setNull(sqlIndex, columnTypes.get(fieldIndex).getType());
         break;
       case STRING:
         // clob can also be written to as setString
@@ -333,7 +340,7 @@ protected void writeToDB(PreparedStatement stmt, Schema.Field field, int fieldIn
   protected void writeBytes(PreparedStatement stmt, int fieldIndex, int sqlIndex, Object fieldValue)
     throws SQLException {
     byte[] byteValue = fieldValue instanceof ByteBuffer ? Bytes.toBytes((ByteBuffer) fieldValue) : (byte[]) fieldValue;
-    int parameterType = columnTypes[fieldIndex];
+    int parameterType = columnTypes.get(fieldIndex).getType();
     if (Types.BLOB == parameterType) {
       stmt.setBlob(sqlIndex, new SerialBlob(byteValue));
       return;
@@ -342,9 +349,9 @@ protected void writeBytes(PreparedStatement stmt, int fieldIndex, int sqlIndex,
     stmt.setBytes(sqlIndex, byteValue);
   }
 
-  private void writeInt(PreparedStatement stmt, int fieldIndex, int sqlIndex, Object fieldValue) throws SQLException {
+  protected void writeInt(PreparedStatement stmt, int fieldIndex, int sqlIndex, Object fieldValue) throws SQLException {
     Integer intValue = (Integer) fieldValue;
-    int parameterType = columnTypes[fieldIndex];
+    int parameterType = columnTypes.get(fieldIndex).getType();
     if (Types.TINYINT == parameterType || Types.SMALLINT == parameterType) {
       stmt.setShort(sqlIndex, intValue.shortValue());
       return;

database-commons/src/main/java/io/cdap/plugin/db/batch/sink/AbstractDBSink.java

@@ -38,6 +38,7 @@
 import io.cdap.plugin.common.ReferenceBatchSink;
 import io.cdap.plugin.common.ReferencePluginConfig;
 import io.cdap.plugin.common.batch.sink.SinkOutputFormatProvider;
+import io.cdap.plugin.db.ColumnType;
 import io.cdap.plugin.db.CommonSchemaReader;
 import io.cdap.plugin.db.ConnectionConfig;
 import io.cdap.plugin.db.ConnectionConfigAccessor;
@@ -59,6 +60,7 @@
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
 import java.sql.Statement;
+import java.sql.Types;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashSet;
@@ -80,10 +82,9 @@ public abstract class AbstractDBSink extends ReferenceBatchSink<StructuredRecord
   private final DBSinkConfig dbSinkConfig;
   private Class<? extends Driver> driverClass;
   private DriverCleanup driverCleanup;
-  protected int[] columnTypes;
   protected List<String> columns;
+  protected List<ColumnType> columnTypes;
   protected String dbColumns;
-  private Schema outputSchema;
 
   public AbstractDBSink(DBSinkConfig dbSinkConfig) {
     super(new ReferencePluginConfig(dbSinkConfig.referenceName));
@@ -116,7 +117,7 @@ public void prepareRun(BatchSinkContext context) {
               dbSinkConfig.jdbcPluginName,
               connectionString);
 
-    outputSchema = context.getInputSchema();
+    Schema outputSchema = context.getInputSchema();
 
     // Load the plugin class to make sure it is available.
     Class<? extends Driver> driverClass = context.loadPluginClass(getJDBCPluginId());
@@ -176,7 +177,7 @@ protected void setColumnsInfo(List<Schema.Field> fields) {
   public void initialize(BatchRuntimeContext context) throws Exception {
     super.initialize(context);
     driverClass = context.loadPluginClass(getJDBCPluginId());
-    outputSchema = Optional.ofNullable(context.getInputSchema()).orElse(inferSchema(driverClass));
+    Schema outputSchema = Optional.ofNullable(context.getInputSchema()).orElse(inferSchema(driverClass));
 
     setColumnsInfo(outputSchema.getFields());
     setResultSetMetadata();
@@ -209,23 +210,11 @@ private Schema inferSchema(Class<? extends Driver> driverClass) {
 
   @Override
   public void transform(StructuredRecord input, Emitter<KeyValue<DBRecord, NullWritable>> emitter) {
-    // Create StructuredRecord that only has the columns in this.columns
-    List<Schema.Field> outputFields = new ArrayList<>();
-    for (Schema.Field field : input.getSchema().getFields()) {
-      Preconditions.checkArgument(columns.contains(field.getName()), "Input field '%s' is not found in columns",
-                                  field.getName());
-      outputFields.add(field);
-    }
-    StructuredRecord.Builder output = StructuredRecord.builder(outputSchema);
-    for (String column : columns) {
-      output.set(column, input.get(column));
-    }
-
-    emitter.emit(new KeyValue<>(getDBRecord(output), null));
+    emitter.emit(new KeyValue<>(getDBRecord(input), null));
   }
 
-  protected DBRecord getDBRecord(StructuredRecord.Builder output) {
-    return new DBRecord(output.build(), columnTypes);
+  protected DBRecord getDBRecord(StructuredRecord output) {
+    return new DBRecord(output, columnTypes);
   }
 
   protected SchemaReader getSchemaReader() {
@@ -272,12 +261,12 @@ private void setResultSetMetadata() throws Exception {
       }
     }
 
-    columnTypes = new int[columns.size()];
-    for (int i = 0; i < columnTypes.length; i++) {
-      String name = columns.get(i);
-      Preconditions.checkArgument(columnToType.containsKey(name), "Missing column '%s' in SQL table", name);
-      columnTypes[i] = columnToType.get(name);
-    }
+    this.columnTypes = columns.stream()
+      .map(name -> {
+        Preconditions.checkArgument(columnToType.containsKey(name), "Missing column '%s' in SQL table", name);
+        return new ColumnType(name, columnToType.get(name));
+      })
+      .collect(Collectors.collectingAndThen(Collectors.toList(), Collections::unmodifiableList));
   }
 
   private void validateSchema(Class<? extends Driver> jdbcDriverClass, String tableName, Schema inputSchema) {
@@ -324,7 +313,23 @@ private void validateFields(Schema inputSchema, ResultSet rs) throws SQLExceptio
     Set<String> invalidFields = new HashSet<>();
     for (Schema.Field field : inputSchema.getFields()) {
       int columnIndex = rs.findColumn(field.getName());
+      boolean isColumnNullable = (ResultSetMetaData.columnNullable == rsMetaData.isNullable(columnIndex));
+      boolean isNotNullAssignable = !isColumnNullable && field.getSchema().isNullable();
+      if (isNotNullAssignable) {
+        LOG.error("Field '{}' was given as nullable but the database column is not nullable", field.getName());
+        invalidFields.add(field.getName());
+      }
+
       if (!isFieldCompatible(field, rsMetaData, columnIndex)) {
+        String sqlTypeName = rsMetaData.getColumnTypeName(columnIndex);
+        Schema fieldSchema = field.getSchema().isNullable() ? field.getSchema().getNonNullable() : field.getSchema();
+        Schema.Type fieldType = fieldSchema.getType();
+        Schema.LogicalType fieldLogicalType = fieldSchema.getLogicalType();
+        LOG.error("Field '{}' was given as type '{}' but the database column is actually of type '{}'.",
+                  field.getName(),
+                  fieldLogicalType != null ? fieldLogicalType.getToken() : fieldType,
+                  sqlTypeName
+        );
         invalidFields.add(field.getName());
       }
     }
@@ -335,34 +340,70 @@ private void validateFields(Schema inputSchema, ResultSet rs) throws SQLExceptio
   }
 
   /**
-   * Checks if field of the input schema is compatible with corresponding database column.
+   * Checks if field is compatible to be written into database column of the given sql index.
    * @param field field of the explicit input schema.
    * @param metadata resultSet metadata.
    * @param index sql column index.
-   * @return 'true' if field is compatible, 'false' otherwise.
+   * @return 'true' if field is compatible to be written, 'false' otherwise.
    */
   protected boolean isFieldCompatible(Schema.Field field, ResultSetMetaData metadata, int index) throws SQLException {
-    boolean isColumnNullable = (ResultSetMetaData.columnNullable == metadata.isNullable(index));
-    boolean isNotNullAssignable = !isColumnNullable && field.getSchema().isNullable();
-    if (isNotNullAssignable) {
-      LOG.error("Field '{}' was given as nullable but the database column is not nullable", field.getName());
-      return false;
+    Schema fieldSchema = field.getSchema().isNullable() ? field.getSchema().getNonNullable() : field.getSchema();
+    Schema.Type fieldType = fieldSchema.getType();
+    Schema.LogicalType fieldLogicalType = fieldSchema.getLogicalType();
+
+    int sqlType = metadata.getColumnType(index);
+
+    // Handle logical types first
+    if (fieldLogicalType != null) {
+      switch (fieldLogicalType) {
+        case DATE:
+          return sqlType == Types.DATE;
+        case TIME_MILLIS:
+        case TIME_MICROS:
+          return sqlType == Types.TIME;
+        case TIMESTAMP_MILLIS:
+        case TIMESTAMP_MICROS:
+          return sqlType == Types.TIMESTAMP;
+        case DECIMAL:
+          return sqlType == Types.NUMERIC
+            || sqlType == Types.DECIMAL;
+      }
     }
 
-    int type = metadata.getColumnType(index);
-    int precision = metadata.getPrecision(index);
-    int scale = metadata.getScale(index);
-
-    Schema inputFieldSchema = field.getSchema().isNullable() ? field.getSchema().getNonNullable() : field.getSchema();
-    Schema outputFieldSchema = DBUtils.getSchema(type, precision, scale);
-    if (!Objects.equals(inputFieldSchema.getType(), outputFieldSchema.getType()) ||
-      !Objects.equals(inputFieldSchema.getLogicalType(), outputFieldSchema.getLogicalType())) {
-      LOG.error("Field '{}' was given as type '{}' but the database column is actually of type '{}'.",
-                field.getName(), inputFieldSchema.getType(), outputFieldSchema.getType());
-      return false;
+    switch (fieldType) {
+      case NULL:
+        return true;
+      case BOOLEAN:
+        return sqlType == Types.BOOLEAN
+          || sqlType == Types.BIT;
+      case INT:
+        return sqlType == Types.INTEGER
+          ||  sqlType == Types.SMALLINT
+          || sqlType == Types.TINYINT;
+      case LONG:
+        return sqlType == Types.BIGINT;
+      case FLOAT:
+        return sqlType == Types.REAL
+          || sqlType == Types.FLOAT;
+      case DOUBLE:
+        return sqlType == Types.DOUBLE;
+      case BYTES:
+        return sqlType == Types.BINARY
+          || sqlType == Types.VARBINARY
+          || sqlType == Types.LONGVARBINARY
+          || sqlType == Types.BLOB;
+      case STRING:
+        return sqlType == Types.VARCHAR
+          || sqlType == Types.CHAR
+          || sqlType == Types.CLOB
+          || sqlType == Types.LONGNVARCHAR
+          || sqlType == Types.LONGVARCHAR
+          || sqlType == Types.NCHAR
+          || sqlType == Types.NCLOB
+          || sqlType == Types.NVARCHAR;
+      default:
+        return false;
     }
-
-    return true;
   }
 
   private void emitLineage(BatchSinkContext context, List<Schema.Field> fields) {