Use a dedicated spi method in Connector for distributed procedure

Author: hantangwangd

presto-iceberg/src/main/java/com/facebook/presto/iceberg/IcebergConnector.java

@@ -32,6 +32,8 @@
 import com.facebook.presto.spi.connector.ConnectorTransactionHandle;
 import com.facebook.presto.spi.connector.classloader.ClassLoaderSafeConnectorMetadata;
 import com.facebook.presto.spi.procedure.BaseProcedure;
+import com.facebook.presto.spi.procedure.DistributedProcedure;
+import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.session.PropertyMetadata;
 import com.facebook.presto.spi.transaction.IsolationLevel;
 import com.google.common.collect.ImmutableList;
@@ -151,11 +153,15 @@ public Set<SystemTable> getSystemTables()
     }
 
     @Override
-    public <T extends BaseProcedure<?>> Set<T> getProcedures(Class<T> clazz)
+    public Set<Procedure> getProcedures()
     {
-        return procedures.stream().filter(clazz::isInstance)
-                .map(clazz::cast)
-                .collect(Collectors.toSet());
+        return getProcedures(Procedure.class);
+    }
+
+    @Override
+    public Set<DistributedProcedure> getDistributedProcedures()
+    {
+        return getProcedures(DistributedProcedure.class);
     }
 
     @Override
@@ -233,4 +239,11 @@ public Set<Class<?>> getSystemFunctions()
                 .add(IcebergBucketFunction.Bucket.class)
                 .build();
     }
+
+    private <T extends BaseProcedure<?>> Set<T> getProcedures(Class<T> clazz)
+    {
+        return procedures.stream().filter(clazz::isInstance)
+                .map(clazz::cast)
+                .collect(Collectors.toSet());
+    }
 }

presto-main-base/src/main/java/com/facebook/presto/connector/ConnectorManager.java

@@ -451,7 +451,7 @@ public MaterializedConnector(ConnectorId connectorId, Connector connector)
             Set<Procedure> procedures = connector.getProcedures();
             requireNonNull(procedures, "Connector %s returned a null procedures set");
             proceduresBuilder.addAll(procedures);
-            Set<DistributedProcedure> distributedProcedures = connector.getProcedures(DistributedProcedure.class);
+            Set<DistributedProcedure> distributedProcedures = connector.getDistributedProcedures();
             requireNonNull(distributedProcedures, "Connector %s returned a null distributedProcedures set");
             proceduresBuilder.addAll(distributedProcedures);
             this.procedures = ImmutableSet.copyOf(proceduresBuilder.build());

presto-main-base/src/test/java/com/facebook/presto/sql/planner/TestLogicalPlanner.java

@@ -70,7 +70,6 @@
 import com.facebook.presto.util.MorePredicates;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
-import com.google.common.collect.ImmutableSet;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
@@ -82,6 +81,7 @@
 import java.util.Set;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
+import java.util.stream.Collectors;
 
 import static com.facebook.presto.SystemSessionProperties.DISTRIBUTED_SORT;
 import static com.facebook.presto.SystemSessionProperties.ENFORCE_FIXED_DISTRIBUTION_FOR_OUTPUT_OPERATOR;
@@ -243,14 +243,11 @@ public ConnectorSplitManager getSplitManager()
                             }
 
                             @Override
-                            public Set<BaseProcedure<?>> getProcedures(Class clazz)
+                            public Set<DistributedProcedure> getDistributedProcedures()
                             {
-                                if (DistributedProcedure.class.isAssignableFrom(clazz)) {
-                                    return procedures;
-                                }
-                                else {
-                                    return ImmutableSet.of();
-                                }
+                                return procedures.stream().filter(DistributedProcedure.class::isInstance)
+                                        .map(DistributedProcedure.class::cast)
+                                        .collect(Collectors.toSet());
                             }
                         };
                     }

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

@@ -15,7 +15,7 @@
 
 import com.facebook.presto.spi.SystemTable;
 import com.facebook.presto.spi.function.table.ConnectorTableFunction;
-import com.facebook.presto.spi.procedure.BaseProcedure;
+import com.facebook.presto.spi.procedure.DistributedProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.session.PropertyMetadata;
 import com.facebook.presto.spi.transaction.IsolationLevel;
@@ -108,13 +108,13 @@ default Set<SystemTable> getSystemTables()
      */
     default Set<Procedure> getProcedures()
     {
-        return getProcedures(Procedure.class);
+        return emptySet();
     }
 
     /**
-     * @return the set of procedures with target class provided by this connector
+     * @return the set of distributed procedures provided by this connector
      */
-    default <T extends BaseProcedure<?>> Set<T> getProcedures(Class<T> targetClz)
+    default Set<DistributedProcedure> getDistributedProcedures()
     {
         return emptySet();
     }