From beb92ca532a333e79eb18fb9f67eec3378f691ed Mon Sep 17 00:00:00 2001
From: Yuya Ebihara <ebyhry@gmail.com>
Date: Fri, 18 Aug 2023 09:46:42 +0900
Subject: [PATCH] Add support for guessing Decimal128 in MongoDB

---
 .../io/trino/plugin/mongodb/MongoSession.java | 12 +++++++++++
 .../mongodb/TestMongoConnectorTest.java       | 21 +++++++++++++++++++
 .../plugin/mongodb/TestMongoTypeMapping.java  |  8 +++++++
 3 files changed, 41 insertions(+)

diff --git a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java
index 366b4e478942..94fb348f0172 100644
--- a/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java
+++ b/plugin/trino-mongodb/src/main/java/io/trino/plugin/mongodb/MongoSession.java
@@ -68,6 +68,7 @@
 import org.bson.types.Decimal128;
 import org.bson.types.ObjectId;
 
+import java.math.BigDecimal;
 import java.time.Instant;
 import java.time.LocalDate;
 import java.time.LocalDateTime;
@@ -98,6 +99,7 @@
 import static io.trino.spi.type.Chars.padSpaces;
 import static io.trino.spi.type.DateTimeEncoding.unpackMillisUtc;
 import static io.trino.spi.type.DateType.DATE;
+import static io.trino.spi.type.DecimalType.createDecimalType;
 import static io.trino.spi.type.DoubleType.DOUBLE;
 import static io.trino.spi.type.SmallintType.SMALLINT;
 import static io.trino.spi.type.TimeType.TIME_MILLIS;
@@ -900,6 +902,16 @@ else if (value instanceof Boolean) {
         else if (value instanceof Float || value instanceof Double) {
             typeSignature = DOUBLE.getTypeSignature();
         }
+        else if (value instanceof Decimal128 decimal128) {
+            BigDecimal decimal;
+            try {
+                decimal = decimal128.bigDecimalValue();
+            }
+            catch (ArithmeticException e) {
+                return Optional.empty();
+            }
+            typeSignature = createDecimalType(decimal.precision(), decimal.scale()).getTypeSignature();
+        }
         else if (value instanceof Date) {
             typeSignature = TIMESTAMP_MILLIS.getTypeSignature();
         }
diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java
index ab98e1551a02..6de385a809c3 100644
--- a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java
+++ b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoConnectorTest.java
@@ -33,6 +33,7 @@
 import io.trino.testing.TestingConnectorBehavior;
 import io.trino.testing.sql.TestTable;
 import org.bson.Document;
+import org.bson.types.Decimal128;
 import org.bson.types.ObjectId;
 import org.testng.SkipException;
 import org.testng.annotations.AfterClass;
@@ -524,6 +525,26 @@ public void testSkipUnknownTypes()
         assertUpdate("DROP TABLE test." + allUnknownFieldTable);
     }
 
+    @Test
+    public void testSkipUnsupportedDecimal128()
+    {
+        String tableName = "test_unsupported_decimal128" + randomNameSuffix();
+
+        Document document = new Document(ImmutableMap.<String, Object>builder()
+                .put("col", 1)
+                .put("nan", Decimal128.NaN)
+                .put("negative_nan", Decimal128.NEGATIVE_NaN)
+                .put("positive_infinity", Decimal128.POSITIVE_INFINITY)
+                .put("negative_infinity", Decimal128.NEGATIVE_INFINITY)
+                .put("negative_zero", Decimal128.NEGATIVE_ZERO)
+                .buildOrThrow());
+        client.getDatabase("test").getCollection(tableName).insertOne(document);
+        assertQuery("SHOW COLUMNS FROM test." + tableName, "SELECT 'col', 'bigint', '', ''");
+        assertQuery("SELECT col FROM test." + tableName, "SELECT 1");
+
+        assertUpdate("DROP TABLE test." + tableName);
+    }
+
     @Test(dataProvider = "dbRefProvider")
     public void testDBRef(Object objectId, String expectedValue, String expectedType)
     {
diff --git a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTypeMapping.java b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTypeMapping.java
index c557dc3dcee0..afd6b5032b8a 100644
--- a/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTypeMapping.java
+++ b/plugin/trino-mongodb/src/test/java/io/trino/plugin/mongodb/TestMongoTypeMapping.java
@@ -185,6 +185,14 @@ public void testDecimal()
                 .addRoundTrip("decimal(38, 0)", "CAST(NULL AS decimal(38, 0))", createDecimalType(38, 0), "CAST(NULL AS decimal(38, 0))")
                 .execute(getQueryRunner(), trinoCreateAsSelect("test_decimal"))
                 .execute(getQueryRunner(), trinoCreateAndInsert("test_decimal"));
+
+        SqlDataTypeTest.create()
+                .addRoundTrip("NumberDecimal(\"2\")", "CAST('2' AS decimal(1, 0))")
+                .addRoundTrip("NumberDecimal(\"2.3\")", "CAST('2.3' AS decimal(2, 1))")
+                .addRoundTrip("NumberDecimal(\"-2.3\")", "CAST('-2.3' AS decimal(2, 1))")
+                .addRoundTrip("NumberDecimal(\"1234567890123456789012345678901234\")", "CAST('1234567890123456789012345678901234' AS decimal(34, 0))") // 34 is the max precision in Decimal128
+                .addRoundTrip("NumberDecimal(\"1234567890123456.789012345678901234\")", "CAST('1234567890123456.789012345678901234' AS decimal(34, 18))")
+                .execute(getQueryRunner(), mongoCreateAndInsert(getSession(), "tpch", "test_decimal"));
     }
 
     @Test