Maintain compatibility with custom connector-provided serialization

Author: hantangwangd

presto-main-base/src/main/java/com/facebook/presto/metadata/DistributedProcedureHandleJacksonModule.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.ConnectorDistributedProcedureHandle;
+import com.facebook.presto.spi.ConnectorId;
+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 DistributedProcedureHandleJacksonModule
         extends AbstractTypedJacksonModule<ConnectorDistributedProcedureHandle>
 {
     @Inject
-    public DistributedProcedureHandleJacksonModule(HandleResolver handleResolver)
+    public DistributedProcedureHandleJacksonModule(
+            HandleResolver handleResolver,
+            Provider<ConnectorManager> connectorManagerProvider,
+            FeaturesConfig featuresConfig)
+    {
+        super(ConnectorDistributedProcedureHandle.class,
+                handleResolver::getId,
+                handleResolver::getDistributedProcedureHandleClass,
+                featuresConfig.isUseConnectorProvidedSerializationCodecs(),
+                connectorId -> connectorManagerProvider.get()
+                        .getConnectorCodecProvider(connectorId)
+                        .flatMap(ConnectorCodecProvider::getConnectorDistributedProcedureHandleCodec));
+    }
+
+    public DistributedProcedureHandleJacksonModule(
+            HandleResolver handleResolver,
+            FeaturesConfig featuresConfig,
+            Function<ConnectorId, Optional<ConnectorCodec<ConnectorDistributedProcedureHandle>>> codecExtractor)
     {
         super(ConnectorDistributedProcedureHandle.class,
                 handleResolver::getId,
-                handleResolver::getDistributedProcedureHandleClass);
+                handleResolver::getDistributedProcedureHandleClass,
+                featuresConfig.isUseConnectorProvidedSerializationCodecs(),
+                codecExtractor);
     }
 }

presto-main-base/src/test/java/com/facebook/presto/metadata/TestAbstractTypedJacksonModule.java

@@ -372,7 +372,8 @@ public static class TestHandle
                        com.facebook.presto.spi.ConnectorDeleteTableHandle,
                        com.facebook.presto.spi.ConnectorIndexHandle,
                        com.facebook.presto.spi.connector.ConnectorPartitioningHandle,
-                       com.facebook.presto.spi.connector.ConnectorTransactionHandle
+                       com.facebook.presto.spi.connector.ConnectorTransactionHandle,
+                       com.facebook.presto.spi.ConnectorDistributedProcedureHandle
     {
         private final String connectorId;
         private final String id;
@@ -529,6 +530,12 @@ public Class<? extends com.facebook.presto.spi.connector.ConnectorTransactionHan
         {
             return TestHandle.class;
         }
+
+        @Override
+        public Class<? extends com.facebook.presto.spi.ConnectorDistributedProcedureHandle> getDistributedProcedureHandleClass()
+        {
+            return TestHandle.class;
+        }
     }
 
     // Test HandleResolver implementation

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

@@ -55,6 +55,7 @@
 import com.facebook.presto.metadata.ColumnHandleJacksonModule;
 import com.facebook.presto.metadata.DeleteTableHandle;
 import com.facebook.presto.metadata.DeleteTableHandleJacksonModule;
+import com.facebook.presto.metadata.DistributedProcedureHandleJacksonModule;
 import com.facebook.presto.metadata.FunctionAndTypeManager;
 import com.facebook.presto.metadata.FunctionHandleJacksonModule;
 import com.facebook.presto.metadata.HandleResolver;
@@ -81,6 +82,7 @@
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ConnectorCodec;
 import com.facebook.presto.spi.ConnectorDeleteTableHandle;
+import com.facebook.presto.spi.ConnectorDistributedProcedureHandle;
 import com.facebook.presto.spi.ConnectorHandleResolver;
 import com.facebook.presto.spi.ConnectorId;
 import com.facebook.presto.spi.ConnectorIndexHandle;
@@ -662,6 +664,8 @@ public void configure(Binder binder)
                                 connectorId -> Optional.empty();
                         Function<ConnectorId, Optional<ConnectorCodec<ConnectorPartitioningHandle>>> noOpPartitioningCodec =
                                 connectorId -> Optional.empty();
+                        Function<ConnectorId, Optional<ConnectorCodec<ConnectorDistributedProcedureHandle>>> noOpDistributedProcedureCodec =
+                                connectorId -> Optional.empty();
                         Function<ConnectorId, Optional<ConnectorCodec<ConnectorSplit>>> splitCodecExtractor =
                                 connectorId -> Optional.ofNullable(splitCodecMap.get(connectorId.getCatalogName()));
 
@@ -676,6 +680,7 @@ public void configure(Binder binder)
                         jsonBinder(binder).addModuleBinding().toInstance(new PartitioningHandleJacksonModule(handleResolver, featuresConfig, noOpPartitioningCodec));
                         jsonBinder(binder).addModuleBinding().toInstance(new FunctionHandleJacksonModule(handleResolver));
                         jsonBinder(binder).addModuleBinding().toInstance(new SplitJacksonModule(handleResolver, featuresConfig, splitCodecExtractor));
+                        jsonBinder(binder).addModuleBinding().toInstance(new DistributedProcedureHandleJacksonModule(handleResolver, featuresConfig, noOpDistributedProcedureCodec));
 
                         FunctionAndTypeManager functionAndTypeManager = createTestFunctionAndTypeManager();
                         binder.bind(TypeManager.class).toInstance(functionAndTypeManager);

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

@@ -16,6 +16,7 @@
 import com.facebook.presto.spi.ColumnHandle;
 import com.facebook.presto.spi.ConnectorCodec;
 import com.facebook.presto.spi.ConnectorDeleteTableHandle;
+import com.facebook.presto.spi.ConnectorDistributedProcedureHandle;
 import com.facebook.presto.spi.ConnectorIndexHandle;
 import com.facebook.presto.spi.ConnectorInsertTableHandle;
 import com.facebook.presto.spi.ConnectorOutputTableHandle;
@@ -76,4 +77,9 @@ default Optional<ConnectorCodec<ConnectorIndexHandle>> getConnectorIndexHandleCo
     {
         return Optional.empty();
     }
+
+    default Optional<ConnectorCodec<ConnectorDistributedProcedureHandle>> getConnectorDistributedProcedureHandleCodec()
+    {
+        return Optional.empty();
+    }
 }