From 581eb87017d5766279184817bdb55a7479121a7f Mon Sep 17 00:00:00 2001
From: exxiang <543985040@qq.com>
Date: Tue, 6 Aug 2024 11:24:10 +0800
Subject: [PATCH] Extend MongoDB connector support binData

---
 .../facebook/presto/mongodb/MongoSession.java    |  5 +++++
 .../mongodb/TestMongoIntegrationSmokeTest.java   | 13 +++++++++++++
 .../presto/mongodb/TestMongoSession.java         | 16 ++++++++++++++++
 3 files changed, 34 insertions(+)

diff --git a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
index 36b67f8b8cb2e..779ea39f38a7a 100644
--- a/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
+++ b/presto-mongodb/src/main/java/com/facebook/presto/mongodb/MongoSession.java
@@ -24,6 +24,7 @@
 import com.facebook.presto.common.type.TypeManager;
 import com.facebook.presto.common.type.TypeSignature;
 import com.facebook.presto.common.type.TypeSignatureParameter;
+import com.facebook.presto.common.type.VarbinaryType;
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.SchemaNotFoundException;
@@ -46,6 +47,7 @@
 import com.mongodb.client.result.DeleteResult;
 import io.airlift.slice.Slice;
 import org.bson.Document;
+import org.bson.types.Binary;
 import org.bson.types.ObjectId;
 
 import java.util.ArrayList;
@@ -519,6 +521,9 @@ else if (value instanceof Date) {
         else if (value instanceof ObjectId) {
             typeSignature = OBJECT_ID.getTypeSignature();
         }
+        else if (value instanceof Binary) {
+            typeSignature = VarbinaryType.VARBINARY.getTypeSignature();
+        }
         else if (value instanceof List) {
             List<Optional<TypeSignature>> subTypes = ((List<?>) value).stream()
                     .map(this::guessFieldType)
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
index 38474722104b4..7819e15017b65 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoIntegrationSmokeTest.java
@@ -247,6 +247,19 @@ public void testObjectIds()
         assertOneNotNullResult("SELECT id FROM tmp_objectid WHERE id = ObjectId('ffffffffffffffffffffffff')");
     }
 
+    @Test
+    public void testBinarys()
+    {
+        assertUpdate("CREATE TABLE tmp_binary AS SELECT cast('value' as varbinary) AS _varbinary", 1);
+        assertOneNotNullResult("SELECT _varbinary FROM tmp_binary");
+
+        MaterializedResult results = getQueryRunner().execute(getSession(), "SELECT _varbinary FROM tmp_binary").toTestTypes();
+        assertEquals(results.getRowCount(), 1);
+
+        MaterializedRow row = results.getMaterializedRows().get(0);
+        assertEquals(row.getField(0), "value".getBytes(UTF_8));
+    }
+
     private void assertOneNotNullResult(String query)
     {
         MaterializedResult results = getQueryRunner().execute(getSession(), query).toTestTypes();
diff --git a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoSession.java b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoSession.java
index 5cd4595e1e4d2..2755bfdc5a189 100644
--- a/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoSession.java
+++ b/presto-mongodb/src/test/java/com/facebook/presto/mongodb/TestMongoSession.java
@@ -28,8 +28,11 @@
 import static com.facebook.presto.common.predicate.Range.lessThan;
 import static com.facebook.presto.common.predicate.Range.range;
 import static com.facebook.presto.common.type.BigintType.BIGINT;
+import static com.facebook.presto.common.type.VarbinaryType.VARBINARY;
 import static com.facebook.presto.common.type.VarcharType.createUnboundedVarcharType;
 import static io.airlift.slice.Slices.utf8Slice;
+import static io.airlift.slice.Slices.wrappedBuffer;
+import static java.nio.charset.StandardCharsets.UTF_8;
 import static java.util.Arrays.asList;
 import static org.testng.Assert.assertEquals;
 
@@ -37,6 +40,7 @@ public class TestMongoSession
 {
     private static final MongoColumnHandle COL1 = new MongoColumnHandle("col1", BIGINT, false);
     private static final MongoColumnHandle COL2 = new MongoColumnHandle("col2", createUnboundedVarcharType(), false);
+    private static final MongoColumnHandle COL3 = new MongoColumnHandle("col3", VARBINARY, false);
 
     @Test
     public void testBuildQuery()
@@ -52,6 +56,18 @@ public void testBuildQuery()
         assertEquals(query, expected);
     }
 
+    @Test
+    public void testBuildQueryBinaryType()
+    {
+        TupleDomain<ColumnHandle> tupleDomain = TupleDomain.withColumnDomains(ImmutableMap.of(
+                COL3, Domain.singleValue(VARBINARY, wrappedBuffer("VarBinary Value".getBytes(UTF_8)))));
+
+        Document query = MongoSession.buildQuery(tupleDomain);
+        Document expected = new Document()
+                .append(COL3.getName(), new Document().append("$eq", "VarBinary Value"));
+        assertEquals(query, expected);
+    }
+
     @Test
     public void testBuildQueryStringType()
     {