Enabling support for the precision less numbers in PostgreSQL plugin. (#404)

Author: MrRahulSharma

postgresql-plugin/pom.xml

@@ -53,10 +53,6 @@
       <version>42.2.20</version>
       <scope>test</scope>
     </dependency>
-    <dependency>
-      <groupId>org.mockito</groupId>
-      <artifactId>mockito-core</artifactId>
-    </dependency>
     <dependency>
       <groupId>io.cdap.plugin</groupId>
       <artifactId>database-commons</artifactId>
@@ -72,10 +68,6 @@
       <groupId>io.cdap.cdap</groupId>
       <artifactId>cdap-data-pipeline3_2.12</artifactId>
     </dependency>
-    <dependency>
-      <groupId>junit</groupId>
-      <artifactId>junit</artifactId>
-    </dependency>
     <dependency>
       <groupId>io.cdap.cdap</groupId>
       <artifactId>cdap-api</artifactId>
@@ -87,6 +79,16 @@
       <version>RELEASE</version>
       <scope>compile</scope>
     </dependency>
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-core</artifactId>
+      <scope>test</scope>
+    </dependency>
+    <dependency>
+      <groupId>junit</groupId>
+      <artifactId>junit</artifactId>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
   <build>
     <plugins>

postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresDBRecord.java

@@ -26,6 +26,8 @@
 
 import java.lang.reflect.InvocationTargetException;
 import java.lang.reflect.Method;
+import java.math.BigDecimal;
+import java.math.RoundingMode;
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
@@ -57,8 +59,9 @@ public PostgresDBRecord() {
   @Override
   protected void handleField(ResultSet resultSet, StructuredRecord.Builder recordBuilder, Schema.Field field,
                              int columnIndex, int sqlType, int sqlPrecision, int sqlScale) throws SQLException {
-    String columnTypeName = resultSet.getMetaData().getColumnTypeName(columnIndex);
-    if (isUseSchema(resultSet.getMetaData(), columnIndex)) {
+    ResultSetMetaData metadata = resultSet.getMetaData();
+    String columnTypeName = metadata.getColumnTypeName(columnIndex);
+    if (isUseSchema(metadata, columnIndex)) {
       setFieldAccordingToSchema(resultSet, recordBuilder, field, columnIndex);
     } else if (sqlType == Types.TIMESTAMP && columnTypeName.equalsIgnoreCase("timestamp")) {
       Timestamp timestamp = resultSet.getTimestamp(columnIndex, DBUtils.PURE_GREGORIAN_CALENDAR);
@@ -80,6 +83,24 @@ protected void handleField(ResultSet resultSet, StructuredRecord.Builder recordB
         recordBuilder.set(field.getName(), null);
       }
     } else {
+      int columnType = metadata.getColumnType(columnIndex);
+      if (columnType == Types.NUMERIC) {
+        Schema nonNullableSchema = field.getSchema().isNullable() ?
+                field.getSchema().getNonNullable() : field.getSchema();
+        int precision = metadata.getPrecision(columnIndex);
+        if (precision == 0 && Schema.Type.STRING.equals(nonNullableSchema.getType())) {
+          // When output schema is set to String for precision less numbers
+          recordBuilder.set(field.getName(), resultSet.getString(columnIndex));
+        } else if (Schema.LogicalType.DECIMAL.equals(nonNullableSchema.getLogicalType())) {
+          BigDecimal orgValue = resultSet.getBigDecimal(columnIndex);
+          if (orgValue != null) {
+            BigDecimal decimalValue = new BigDecimal(orgValue.toPlainString())
+                    .setScale(nonNullableSchema.getScale(), RoundingMode.HALF_EVEN);
+            recordBuilder.setDecimal(field.getName(), decimalValue);
+          }
+        }
+        return;
+      }
       setField(resultSet, recordBuilder, field, columnIndex, sqlType, sqlPrecision, sqlScale);
     }
   }
@@ -98,14 +119,14 @@ private void setZonedDateTimeBasedOnOuputSchema(StructuredRecord.Builder recordB
   }
 
   private static boolean isUseSchema(ResultSetMetaData metadata, int columnIndex) throws SQLException {
-    switch (metadata.getColumnTypeName(columnIndex)) {
-      case "bit":
-      case "timetz":
-      case "money":
-        return true;
-      default:
-        return PostgresSchemaReader.STRING_MAPPED_POSTGRES_TYPES.contains(metadata.getColumnType(columnIndex));
+    String columnTypeName = metadata.getColumnTypeName(columnIndex);
+    // If the column Type Name is present in the String mapped PostgreSQL types then return true.
+    if (PostgresSchemaReader.STRING_MAPPED_POSTGRES_TYPES_NAMES.contains(columnTypeName)
+        || PostgresSchemaReader.STRING_MAPPED_POSTGRES_TYPES.contains(metadata.getColumnType(columnIndex))) {
+      return true;
     }
+
+    return false;
   }
 
   private Object createPGobject(String type, String value, ClassLoader classLoader) throws SQLException {
@@ -133,9 +154,17 @@ protected void writeNonNullToDB(PreparedStatement stmt, Schema fieldSchema,
       stmt.setObject(sqlIndex, createPGobject(columnType.getTypeName(),
                                               record.get(fieldName),
                                               stmt.getClass().getClassLoader()));
-    } else {
-      super.writeNonNullToDB(stmt, fieldSchema, fieldName, fieldIndex);
+      return;
+    } else if (columnType.getType() == Types.NUMERIC) {
+      if (record.get(fieldName) != null) {
+        if (fieldSchema.getType() == Schema.Type.STRING) {
+          stmt.setBigDecimal(sqlIndex, new BigDecimal((String) record.get(fieldName)));
+          return;
+        }
+      }
     }
+
+    super.writeNonNullToDB(stmt, fieldSchema, fieldName, fieldIndex);
   }
 
   @Override

postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresFieldsValidator.java

@@ -21,6 +21,7 @@
 
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
+import java.sql.Types;
 import java.util.Objects;
 
 /**
@@ -46,6 +47,14 @@ public boolean isFieldCompatible(Schema.Field field, ResultSetMetaData metadata,
       }
     }
 
+    // Since Numeric types without precision and scale are getting converted into CDAP String type at the Source
+    // plugin, hence making the String type compatible with the Numeric type at the Sink as well.
+    if (fieldType.equals(Schema.Type.STRING)) {
+      if (Types.NUMERIC == columnType) {
+        return true;
+      }
+    }
+    
     if (colTypeName.equalsIgnoreCase("timestamp")
         && schema.getLogicalType().equals(Schema.LogicalType.DATETIME)) {
       return true;

postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresSchemaReader.java

@@ -19,6 +19,8 @@
 import com.google.common.collect.ImmutableSet;
 import io.cdap.cdap.api.data.schema.Schema;
 import io.cdap.plugin.db.CommonSchemaReader;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
@@ -30,6 +32,8 @@
  */
 public class PostgresSchemaReader extends CommonSchemaReader {
 
+  private static final Logger LOG = LoggerFactory.getLogger(PostgresSchemaReader.class);
+
   public static final Set<Integer> STRING_MAPPED_POSTGRES_TYPES = ImmutableSet.of(
     Types.OTHER, Types.ARRAY, Types.SQLXML
   );
@@ -58,6 +62,18 @@ public Schema getSchema(ResultSetMetaData metadata, int index) throws SQLExcepti
       return Schema.of(Schema.Type.STRING);
     }
 
+    // If it is a numeric type without precision then use the Schema of String to avoid any precision loss
+    if (Types.NUMERIC == columnType) {
+      int precision = metadata.getPrecision(index);
+      if (precision == 0) {
+        LOG.warn(String.format("Field '%s' is a %s type without precision and scale, "
+                        + "converting into STRING type to avoid any precision loss.",
+                metadata.getColumnName(index),
+                metadata.getColumnTypeName(index)));
+        return Schema.of(Schema.Type.STRING);
+      }
+    }
+    
     if (typeName.equalsIgnoreCase("timestamp")) {
       return Schema.of(Schema.LogicalType.DATETIME);
     }

postgresql-plugin/src/main/java/io/cdap/plugin/postgres/PostgresSource.java

@@ -142,6 +142,17 @@ public void validate(FailureCollector collector) {
     @Override
     protected void validateField(FailureCollector collector, Schema.Field field,
                                  Schema actualFieldSchema, Schema expectedFieldSchema) {
+
+      // This change is needed to make sure that the pipeline upgrade continues to work post upgrade.
+      // Since the older handling of the precision less used to convert to the decimal type,
+      // and the new version would try to convert to the String type. In that case the output schema would
+      // contain Decimal(38, 0) (or something similar), and the code internally would try to identify
+      // the schema of the field(without precision and scale) as String.
+      if (Schema.LogicalType.DECIMAL.equals(expectedFieldSchema.getLogicalType())
+              && actualFieldSchema.getType().equals(Schema.Type.STRING)) {
+        return;
+      }
+
       // This change is needed to make sure that the pipeline upgrade continues to work post upgrade.
       // Since the older handling of the Timestamp used to convert to CDAP TIMESTAMP type,
       // but since PostgreSQL Timestamp does not have a timezone information, hence it should ideally map to

postgresql-plugin/src/test/java/io/cdap/plugin/postgres/PostgresDBRecordUnitTest.java

@@ -25,6 +25,7 @@
 import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
 
+import java.math.BigDecimal;
 import java.sql.ResultSet;
 import java.sql.ResultSetMetaData;
 import java.sql.SQLException;
@@ -42,6 +43,50 @@
  */
 @RunWith(MockitoJUnitRunner.class)
 public class PostgresDBRecordUnitTest {
+
+    private static final int DEFAULT_PRECISION = 38;
+
+    /**
+     * Validate the precision less Numbers handling against following use cases.
+     * 1. Ensure that the numeric type with [p,s] set as [38,4] detect as BigDecimal(38,4) in cdap.
+     * 2. Ensure that the numeric type without [p,s] detect as String type in cdap.
+     * @throws Exception
+     */
+    @Test
+    public void validatePrecisionLessDecimalParsing() throws Exception {
+        Schema.Field field1 = Schema.Field.of("ID1", Schema.decimalOf(DEFAULT_PRECISION, 4));
+        Schema.Field field2 = Schema.Field.of("ID2", Schema.of(Schema.Type.STRING));
+
+        Schema schema = Schema.recordOf(
+            "dbRecord",
+            field1,
+            field2
+        );
+
+        ResultSetMetaData resultSetMetaData = Mockito.mock(ResultSetMetaData.class);
+        when(resultSetMetaData.getColumnType(eq(1))).thenReturn(Types.NUMERIC);
+        when(resultSetMetaData.getPrecision(eq(1))).thenReturn(DEFAULT_PRECISION);
+        when(resultSetMetaData.getColumnType(eq(2))).thenReturn(Types.NUMERIC);
+        when(resultSetMetaData.getPrecision(eq(2))).thenReturn(0);
+
+        ResultSet resultSet = Mockito.mock(ResultSet.class);
+
+        when(resultSet.getMetaData()).thenReturn(resultSetMetaData);
+        when(resultSet.getBigDecimal(eq(1))).thenReturn(BigDecimal.valueOf(123.4568));
+        when(resultSet.getString(eq(2))).thenReturn("123.4568");
+
+        StructuredRecord.Builder builder = StructuredRecord.builder(schema);
+        PostgresDBRecord dbRecord = new PostgresDBRecord(null, null, null, null);
+        dbRecord.handleField(resultSet, builder, field1, 1, Types.NUMERIC, DEFAULT_PRECISION, 4);
+        dbRecord.handleField(resultSet, builder, field2, 2, Types.NUMERIC, 0, -127);
+
+        StructuredRecord record = builder.build();
+        Assert.assertTrue(record.getDecimal("ID1") instanceof BigDecimal);
+        Assert.assertEquals(record.getDecimal("ID1"), BigDecimal.valueOf(123.4568));
+        Assert.assertTrue(record.get("ID2") instanceof String);
+        Assert.assertEquals(record.get("ID2"), "123.4568");
+    }
+
     @Test
     public void validateTimestampType() throws SQLException {
         OffsetDateTime offsetDateTime = OffsetDateTime.of(2023, 1, 1, 1, 0, 0, 0, ZoneOffset.UTC);