Add SQL Support for MERGE INTO In Presto prestodb#20578 (engine)

feat(connector): Add support for custom connector-provided serialization codecs (prestodb#26257)

Author: acarpente-denodo

presto-main-base/src/main/java/com/facebook/presto/metadata/MergeHandleJacksonModule.java

@@ -13,18 +13,46 @@
  */
 package com.facebook.presto.metadata;
 
+import com.facebook.presto.connector.ConnectorManager;
+import com.facebook.presto.spi.ConnectorCodec;
+import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorMergeTableHandle;
+import com.facebook.presto.spi.connector.ConnectorCodecProvider;
+import com.facebook.presto.sql.analyzer.FeaturesConfig;
+import jakarta.inject.Provider;
 
 import javax.inject.Inject;
 
+import java.util.Optional;
+import java.util.function.Function;
+
 public class MergeHandleJacksonModule
         extends AbstractTypedJacksonModule<ConnectorMergeTableHandle>
 {
     @Inject
-    public MergeHandleJacksonModule(HandleResolver handleResolver)
+    public MergeHandleJacksonModule(
+            HandleResolver handleResolver,
+            Provider<ConnectorManager> connectorManagerProvider,
+            FeaturesConfig featuresConfig)
+    {
+        super(ConnectorMergeTableHandle.class,
+                handleResolver::getId,
+                handleResolver::getMergeHandleClass,
+                featuresConfig.isUseConnectorProvidedSerializationCodecs(),
+                connectorId -> connectorManagerProvider.get()
+                        .getConnectorCodecProvider(connectorId)
+                        .flatMap(ConnectorCodecProvider::getConnectorMergeTableHandleCodec));
+    }
+
+    public MergeHandleJacksonModule(
+            HandleResolver handleResolver,
+            FeaturesConfig featuresConfig,
+            Function<ConnectorId, Optional<ConnectorCodec<ConnectorMergeTableHandle>>> codecExtractor)
     {
         super(ConnectorMergeTableHandle.class,
                 handleResolver::getId,
-                handleResolver::getMergeHandleClass);
+                handleResolver::getMergeHandleClass,
+                featuresConfig.isUseConnectorProvidedSerializationCodecs(),
+                codecExtractor);
     }
 }

presto-main/src/test/java/com/facebook/presto/server/remotetask/TestHttpRemoteTaskConnectorCodec.java

@@ -85,6 +85,7 @@
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorIndexHandle;
 import com.facebook.presto.spi.ConnectorInsertTableHandle;
+import com.facebook.presto.spi.ConnectorMergeTableHandle;
 import com.facebook.presto.spi.ConnectorOutputTableHandle;
 import com.facebook.presto.spi.ConnectorSplit;
 import com.facebook.presto.spi.ConnectorTableHandle;
@@ -984,6 +985,25 @@ public ConnectorDeleteTableHandle deserialize(byte[] data)
                 }
             });
         }
+
+        public Optional<ConnectorCodec<ConnectorMergeTableHandle>> getConnectorMergeTableHandleCodec()
+        {
+            return Optional.of(new ConnectorCodec<>()
+            {
+                @Override
+                public byte[] serialize(ConnectorMergeTableHandle handle)
+                {
+                    TestConnectorMergeTableHandle mergeTableHandle = (TestConnectorMergeTableHandle) handle;
+                    return mergeTableHandle.getTableName().getBytes(UTF_8);
+                }
+
+                @Override
+                public ConnectorMergeTableHandle deserialize(byte[] data)
+                {
+                    return new TestConnectorMergeTableHandle(new String(data, UTF_8));
+                }
+            });
+        }
     }
 
     /**
@@ -1287,6 +1307,53 @@ public int hashCode()
         }
     }
 
+    /**
+     * Test merge table handle with binary serialization support
+     */
+    public static class TestConnectorMergeTableHandle
+            implements ConnectorMergeTableHandle
+    {
+        private final String tableName;
+
+        @JsonCreator
+        public TestConnectorMergeTableHandle(
+                @JsonProperty("tableName") String tableName)
+        {
+            this.tableName = tableName;
+        }
+
+        @JsonProperty
+        public String getTableName()
+        {
+            return tableName;
+        }
+
+        @Override
+        public ConnectorTableHandle getTableHandle()
+        {
+            throw new UnsupportedOperationException("Merge table handles not supported");
+        }
+
+        @Override
+        public boolean equals(Object o)
+        {
+            if (this == o) {
+                return true;
+            }
+            if (o == null || getClass() != o.getClass()) {
+                return false;
+            }
+            TestConnectorMergeTableHandle that = (TestConnectorMergeTableHandle) o;
+            return tableName.equals(that.tableName);
+        }
+
+        @Override
+        public int hashCode()
+        {
+            return Objects.hash(tableName);
+        }
+    }
+
     public static class TestConnectorWithoutCodecSplit
             implements ConnectorSplit
     {
@@ -1396,5 +1463,11 @@ public Class<? extends ConnectorDeleteTableHandle> getDeleteTableHandleClass()
         {
             throw new UnsupportedOperationException("Delete table handles not supported");
         }
+
+        @Override
+        public Class<? extends ConnectorMergeTableHandle> getMergeTableHandleClass()
+        {
+            throw new UnsupportedOperationException("Merge table handles not supported");
+        }
     }
 }

presto-spi/src/main/java/com/facebook/presto/spi/connector/ConnectorCodecProvider.java

@@ -18,6 +18,7 @@
 import com.facebook.presto.spi.ConnectorDeleteTableHandle;
 import com.facebook.presto.spi.ConnectorIndexHandle;
 import com.facebook.presto.spi.ConnectorInsertTableHandle;
+import com.facebook.presto.spi.ConnectorMergeTableHandle;
 import com.facebook.presto.spi.ConnectorOutputTableHandle;
 import com.facebook.presto.spi.ConnectorSplit;
 import com.facebook.presto.spi.ConnectorTableHandle;
@@ -52,6 +53,11 @@ default Optional<ConnectorCodec<ConnectorDeleteTableHandle>> getConnectorDeleteT
         return Optional.empty();
     }
 
+    default Optional<ConnectorCodec<ConnectorMergeTableHandle>> getConnectorMergeTableHandleCodec()
+    {
+        return Optional.empty();
+    }
+
     default Optional<ConnectorCodec<ConnectorTableLayoutHandle>> getConnectorTableLayoutHandleCodec()
     {
         return Optional.empty();