Extend DistributedProcedure hierarchy for more extension types

Refactor `Procedure` and `DistributedProcedure` into abstract classes.
Use a subclass `TableDataRewriteDistributedProcedure` for table rewrite
tasks, for example, merge small data files, sort table data,
repartition table data etc. And introduce a new class `LocalProcedure`
to represent the former coordinator-only procedures.

Rename `IProcedureRegistry` to `ProcedureRegistry`, and accordingly
rename previous `ProcedureRegistry` to `BuiltInProcedureRegistry`.

Author: hantangwangd

presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/Analysis.java

@@ -32,6 +32,7 @@
 import com.facebook.presto.spi.function.FunctionKind;
 import com.facebook.presto.spi.function.table.Argument;
 import com.facebook.presto.spi.function.table.ConnectorTableFunctionHandle;
+import com.facebook.presto.spi.procedure.DistributedProcedure;
 import com.facebook.presto.spi.security.AccessControl;
 import com.facebook.presto.spi.security.AccessControlContext;
 import com.facebook.presto.spi.security.AllowAllAccessControl;
@@ -175,8 +176,9 @@ public class Analysis
     private final Map<NodeRef<Table>, Map<String, Expression>> columnMasks = new LinkedHashMap<>();
 
     // for call distributed procedure
-    private Optional<QualifiedObjectName> procedureName;
-    private Optional<Object[]> procedureArguments;
+    private Optional<DistributedProcedure.DistributedProcedureType> distributedProcedureType = Optional.empty();
+    private Optional<QualifiedObjectName> procedureName = Optional.empty();
+    private Optional<Object[]> procedureArguments = Optional.empty();
     private Optional<TableHandle> callTarget = Optional.empty();
     private Optional<QuerySpecification> targetQuery = Optional.empty();
 
@@ -682,6 +684,16 @@ public void setProcedureName(Optional<QualifiedObjectName> procedureName)
         this.procedureName = procedureName;
     }
 
+    public Optional<DistributedProcedure.DistributedProcedureType> getDistributedProcedureType()
+    {
+        return distributedProcedureType;
+    }
+
+    public void setDistributedProcedureType(Optional<DistributedProcedure.DistributedProcedureType> distributedProcedureType)
+    {
+        this.distributedProcedureType = distributedProcedureType;
+    }
+
     public Optional<Object[]> getProcedureArguments()
     {
         return procedureArguments;

presto-analyzer/src/main/java/com/facebook/presto/sql/analyzer/BuiltInQueryPreparer.java

@@ -23,7 +23,7 @@
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.analyzer.AnalyzerOptions;
 import com.facebook.presto.spi.analyzer.QueryPreparer;
-import com.facebook.presto.spi.procedure.IProcedureRegistry;
+import com.facebook.presto.spi.procedure.ProcedureRegistry;
 import com.facebook.presto.sql.analyzer.utils.StatementUtils;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.tree.Call;
@@ -64,12 +64,12 @@ public class BuiltInQueryPreparer
         implements QueryPreparer
 {
     private final SqlParser sqlParser;
-    private final IProcedureRegistry procedureRegistry;
+    private final ProcedureRegistry procedureRegistry;
 
     @Inject
     public BuiltInQueryPreparer(
             SqlParser sqlParser,
-            IProcedureRegistry procedureRegistry)
+            ProcedureRegistry procedureRegistry)
     {
         this.sqlParser = requireNonNull(sqlParser, "sqlParser is null");
         this.procedureRegistry = requireNonNull(procedureRegistry, "procedureRegistry is null");

presto-analyzer/src/test/java/com/facebook/presto/sql/analyzer/TestBuiltInQueryPreparer.java

@@ -18,10 +18,11 @@
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.WarningCollector;
 import com.facebook.presto.spi.analyzer.AnalyzerOptions;
-import com.facebook.presto.spi.procedure.DistributedProcedure;
-import com.facebook.presto.spi.procedure.IProcedureRegistry;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
-import com.facebook.presto.spi.procedure.TestProcedureRegistry;
+import com.facebook.presto.spi.procedure.Procedure.Argument;
+import com.facebook.presto.spi.procedure.ProcedureRegistry;
+import com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure;
 import com.facebook.presto.sql.analyzer.BuiltInQueryPreparer.BuiltInPreparedQuery;
 import com.facebook.presto.sql.parser.SqlParser;
 import com.facebook.presto.sql.tree.AllColumns;
@@ -40,8 +41,8 @@
 
 import static com.facebook.presto.common.type.StandardTypes.VARCHAR;
 import static com.facebook.presto.spi.StandardErrorCode.NOT_FOUND;
-import static com.facebook.presto.spi.procedure.DistributedProcedure.SCHEMA;
-import static com.facebook.presto.spi.procedure.DistributedProcedure.TABLE_NAME;
+import static com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure.SCHEMA;
+import static com.facebook.presto.spi.procedure.TableDataRewriteDistributedProcedure.TABLE_NAME;
 import static com.facebook.presto.sql.QueryUtil.selectList;
 import static com.facebook.presto.sql.QueryUtil.simpleQuery;
 import static com.facebook.presto.sql.QueryUtil.table;
@@ -55,23 +56,24 @@ public class TestBuiltInQueryPreparer
     private static final SqlParser SQL_PARSER = new SqlParser();
     private static final Map<String, String> emptyPreparedStatements = ImmutableMap.of();
     private static final AnalyzerOptions testAnalyzerOptions = AnalyzerOptions.builder().build();
-    private static IProcedureRegistry procedureRegistry;
+    private static ProcedureRegistry procedureRegistry;
     private static BuiltInQueryPreparer queryPreparer;
 
     @BeforeClass
     public void setup()
     {
         procedureRegistry = new TestProcedureRegistry();
-        List<Procedure.Argument> arguments = new ArrayList<>();
-        arguments.add(new Procedure.Argument(SCHEMA, VARCHAR));
-        arguments.add(new Procedure.Argument(TABLE_NAME, VARCHAR));
+        List<Argument> arguments = new ArrayList<>();
+        arguments.add(new Argument(SCHEMA, VARCHAR));
+        arguments.add(new Argument(TABLE_NAME, VARCHAR));
 
         List<Procedure> procedures = new ArrayList<>();
-        procedures.add(new Procedure("system", "fun", arguments));
-        procedures.add(new DistributedProcedure("system", "distributed_fun",
+        procedures.add(new LocalProcedure("system", "fun", arguments));
+        procedures.add(new TableDataRewriteDistributedProcedure("system", "distributed_fun",
                 arguments,
                 (session, transactionContext, procedureHandle, fragments) -> null,
-                (transactionContext, procedureHandle, fragments) -> {}));
+                (transactionContext, procedureHandle, fragments) -> {},
+                TestProcedureRegistry.TestProcedureContext::new));
         procedureRegistry.addProcedures(new ConnectorId("test"), procedures);
         queryPreparer = new BuiltInQueryPreparer(SQL_PARSER, procedureRegistry);
     }

presto-analyzer/src/test/java/com/facebook/presto/sql/analyzer/TestProcedureRegistry.java

@@ -0,0 +1,96 @@
+/*
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.facebook.presto.sql.analyzer;
+
+import com.facebook.presto.spi.ConnectorId;
+import com.facebook.presto.spi.PrestoException;
+import com.facebook.presto.spi.SchemaTableName;
+import com.facebook.presto.spi.connector.ConnectorProcedureContext;
+import com.facebook.presto.spi.procedure.DistributedProcedure;
+import com.facebook.presto.spi.procedure.Procedure;
+import com.facebook.presto.spi.procedure.ProcedureRegistry;
+
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+import java.util.stream.Collectors;
+
+import static com.facebook.presto.spi.StandardErrorCode.PROCEDURE_NOT_FOUND;
+import static java.util.Objects.requireNonNull;
+
+public class TestProcedureRegistry
+        implements ProcedureRegistry
+{
+    private final Map<ConnectorId, Map<SchemaTableName, Procedure>> connectorProcedures = new ConcurrentHashMap<>();
+
+    @Override
+    public void addProcedures(ConnectorId connectorId, Collection<Procedure> procedures)
+    {
+        requireNonNull(connectorId, "connectorId is null");
+        requireNonNull(procedures, "procedures is null");
+
+        Map<SchemaTableName, Procedure> proceduresByName = procedures.stream().collect(Collectors.toMap(
+                procedure -> new SchemaTableName(procedure.getSchema(), procedure.getName()),
+                Function.identity()));
+        if (connectorProcedures.putIfAbsent(connectorId, proceduresByName) != null) {
+            throw new IllegalStateException("Procedures already registered for connector: " + connectorId);
+        }
+    }
+
+    @Override
+    public void removeProcedures(ConnectorId connectorId)
+    {
+        connectorProcedures.remove(connectorId);
+    }
+
+    @Override
+    public Procedure resolve(ConnectorId connectorId, SchemaTableName name)
+    {
+        Map<SchemaTableName, Procedure> procedures = connectorProcedures.get(connectorId);
+        if (procedures != null) {
+            Procedure procedure = procedures.get(name);
+            if (procedure != null) {
+                return procedure;
+            }
+        }
+        throw new PrestoException(PROCEDURE_NOT_FOUND, "Procedure not registered: " + name);
+    }
+
+    @Override
+    public DistributedProcedure resolveDistributed(ConnectorId connectorId, SchemaTableName name)
+    {
+        Map<SchemaTableName, Procedure> procedures = connectorProcedures.get(connectorId);
+        if (procedures != null) {
+            Procedure procedure = procedures.get(name);
+            if (procedure != null && procedure instanceof DistributedProcedure) {
+                return (DistributedProcedure) procedure;
+            }
+        }
+        throw new PrestoException(PROCEDURE_NOT_FOUND, "Distributed procedure not registered: " + name);
+    }
+
+    @Override
+    public boolean isDistributedProcedure(ConnectorId connectorId, SchemaTableName name)
+    {
+        Map<SchemaTableName, Procedure> procedures = connectorProcedures.get(connectorId);
+        return procedures != null &&
+                procedures.containsKey(name) &&
+                procedures.get(name) instanceof DistributedProcedure;
+    }
+
+    public static class TestProcedureContext
+            implements ConnectorProcedureContext
+    {}
+}

presto-hive-metastore/src/main/java/com/facebook/presto/hive/metastore/InvalidateMetastoreCacheProcedure.java

@@ -17,6 +17,7 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.procedure.Procedure.Argument;
 import com.google.common.collect.ImmutableList;
@@ -65,7 +66,7 @@ public InvalidateMetastoreCacheProcedure(ExtendedHiveMetastore extendedHiveMetas
     @Override
     public Procedure get()
     {
-        return new Procedure(
+        return new LocalProcedure(
                 "system",
                 "invalidate_metastore_cache",
                 ImmutableList.of(

presto-hive/src/main/java/com/facebook/presto/hive/CreateEmptyPartitionProcedure.java

@@ -22,6 +22,7 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.procedure.Procedure.Argument;
 import com.google.common.collect.ImmutableList;
@@ -83,7 +84,7 @@ public CreateEmptyPartitionProcedure(
     @Override
     public Procedure get()
     {
-        return new Procedure(
+        return new LocalProcedure(
                 "system",
                 "create_empty_partition",
                 ImmutableList.of(

presto-hive/src/main/java/com/facebook/presto/hive/DirectoryListCacheInvalidationProcedure.java

@@ -16,6 +16,7 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.PrestoException;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.procedure.Procedure.Argument;
 import com.google.common.collect.ImmutableList;
@@ -52,7 +53,7 @@ public DirectoryListCacheInvalidationProcedure(DirectoryLister directoryLister)
     @Override
     public Procedure get()
     {
-        return new Procedure(
+        return new LocalProcedure(
                 "system",
                 "invalidate_directory_list_cache",
                 ImmutableList.of(

presto-hive/src/main/java/com/facebook/presto/hive/SyncPartitionMetadataProcedure.java

@@ -24,6 +24,7 @@
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.TableNotFoundException;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.procedure.Procedure.Argument;
 import com.google.common.collect.ImmutableList;
@@ -95,7 +96,7 @@ public SyncPartitionMetadataProcedure(
     @Override
     public Procedure get()
     {
-        return new Procedure(
+        return new LocalProcedure(
                 "system",
                 "sync_partition_metadata",
                 ImmutableList.of(

presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/ExpireSnapshotsProcedure.java

@@ -20,7 +20,9 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.classloader.ThreadContextClassLoader;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
+import com.facebook.presto.spi.procedure.Procedure.Argument;
 import com.google.common.collect.ImmutableList;
 import jakarta.inject.Inject;
 import org.apache.iceberg.ExpireSnapshots;
@@ -61,15 +63,15 @@ public ExpireSnapshotsProcedure(IcebergMetadataFactory metadataFactory)
     @Override
     public Procedure get()
     {
-        return new Procedure(
+        return new LocalProcedure(
                 "system",
                 "expire_snapshots",
                 ImmutableList.of(
-                        new Procedure.Argument("schema", VARCHAR),
-                        new Procedure.Argument("table_name", VARCHAR),
-                        new Procedure.Argument("older_than", TIMESTAMP, false, null),
-                        new Procedure.Argument("retain_last", INTEGER, false, null),
-                        new Procedure.Argument("snapshot_ids", "array(bigint)", false, null)),
+                        new Argument("schema", VARCHAR),
+                        new Argument("table_name", VARCHAR),
+                        new Argument("older_than", TIMESTAMP, false, null),
+                        new Argument("retain_last", INTEGER, false, null),
+                        new Argument("snapshot_ids", "array(bigint)", false, null)),
                 EXPIRE_SNAPSHOTS.bindTo(this));
     }
 

presto-iceberg/src/main/java/com/facebook/presto/iceberg/procedure/FastForwardBranchProcedure.java

@@ -17,7 +17,9 @@
 import com.facebook.presto.spi.ConnectorSession;
 import com.facebook.presto.spi.SchemaTableName;
 import com.facebook.presto.spi.connector.ConnectorMetadata;
+import com.facebook.presto.spi.procedure.LocalProcedure;
 import com.facebook.presto.spi.procedure.Procedure;
+import com.facebook.presto.spi.procedure.Procedure.Argument;
 import com.google.common.collect.ImmutableList;
 import jakarta.inject.Inject;
 import org.apache.iceberg.Table;
@@ -54,14 +56,14 @@ public FastForwardBranchProcedure(IcebergMetadataFactory metadataFactory)
     @Override
     public Procedure get()
     {
-        return new Procedure(
+        return new LocalProcedure(
                 "system",
                 "fast_forward",
                 ImmutableList.of(
-                        new Procedure.Argument("schema", VARCHAR),
-                        new Procedure.Argument("table_name", VARCHAR),
-                        new Procedure.Argument("branch", VARCHAR),
-                        new Procedure.Argument("to", VARCHAR)),
+                        new Argument("schema", VARCHAR),
+                        new Argument("table_name", VARCHAR),
+                        new Argument("branch", VARCHAR),
+                        new Argument("to", VARCHAR)),
                 FAST_FORWARD.bindTo(this));
     }