Refactor the connector SPI to expose procedure registry to connectors

Author: hantangwangd

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

@@ -23,6 +23,7 @@
 import com.facebook.presto.spi.function.FunctionMetadataManager;
 import com.facebook.presto.spi.function.StandardFunctionResolution;
 import com.facebook.presto.spi.plan.FilterStatsCalculatorService;
+import com.facebook.presto.spi.procedure.IProcedureRegistry;
 import com.facebook.presto.spi.relation.RowExpressionService;
 
 import static java.util.Objects.requireNonNull;
@@ -32,6 +33,7 @@ public class ConnectorContextInstance
 {
     private final NodeManager nodeManager;
     private final TypeManager typeManager;
+    private final IProcedureRegistry procedureRegistry;
     private final FunctionMetadataManager functionMetadataManager;
     private final StandardFunctionResolution functionResolution;
     private final PageSorter pageSorter;
@@ -44,6 +46,7 @@ public class ConnectorContextInstance
     public ConnectorContextInstance(
             NodeManager nodeManager,
             TypeManager typeManager,
+            IProcedureRegistry procedureRegistry,
             FunctionMetadataManager functionMetadataManager,
             StandardFunctionResolution functionResolution,
             PageSorter pageSorter,
@@ -55,6 +58,7 @@ public ConnectorContextInstance(
     {
         this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
+        this.procedureRegistry = requireNonNull(procedureRegistry, "procedureRegistry is null");
         this.functionMetadataManager = requireNonNull(functionMetadataManager, "functionMetadataManager is null");
         this.functionResolution = requireNonNull(functionResolution, "functionResolution is null");
         this.pageSorter = requireNonNull(pageSorter, "pageSorter is null");
@@ -77,6 +81,12 @@ public TypeManager getTypeManager()
         return typeManager;
     }
 
+    @Override
+    public IProcedureRegistry getProcedureRegistry()
+    {
+        return procedureRegistry;
+    }
+
     @Override
     public FunctionMetadataManager getFunctionMetadataManager()
     {

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

@@ -52,6 +52,7 @@
 import com.facebook.presto.spi.connector.ConnectorRecordSetProvider;
 import com.facebook.presto.spi.connector.ConnectorSplitManager;
 import com.facebook.presto.spi.function.table.ConnectorTableFunction;
+import com.facebook.presto.spi.procedure.IProcedureRegistry;
 import com.facebook.presto.spi.procedure.Procedure;
 import com.facebook.presto.spi.relation.DeterminismEvaluator;
 import com.facebook.presto.spi.relation.DomainTranslator;
@@ -110,6 +111,7 @@ public class ConnectorManager
     private final HandleResolver handleResolver;
     private final InternalNodeManager nodeManager;
     private final TypeManager typeManager;
+    private final IProcedureRegistry procedureRegistry;
     private final PageSorter pageSorter;
     private final PageIndexerFactory pageIndexerFactory;
     private final NodeInfo nodeInfo;
@@ -146,6 +148,7 @@ public ConnectorManager(
             InternalNodeManager nodeManager,
             NodeInfo nodeInfo,
             TypeManager typeManager,
+            IProcedureRegistry procedureRegistry,
             PageSorter pageSorter,
             PageIndexerFactory pageIndexerFactory,
             TransactionManager transactionManager,
@@ -170,6 +173,7 @@ public ConnectorManager(
         this.handleResolver = requireNonNull(handleResolver, "handleResolver is null");
         this.nodeManager = requireNonNull(nodeManager, "nodeManager is null");
         this.typeManager = requireNonNull(typeManager, "typeManager is null");
+        this.procedureRegistry = requireNonNull(procedureRegistry, "procedureRegistry is null");
         this.pageSorter = requireNonNull(pageSorter, "pageSorter is null");
         this.pageIndexerFactory = requireNonNull(pageIndexerFactory, "pageIndexerFactory is null");
         this.nodeInfo = requireNonNull(nodeInfo, "nodeInfo is null");
@@ -374,6 +378,7 @@ private Connector createConnector(ConnectorId connectorId, ConnectorFactory fact
         ConnectorContext context = new ConnectorContextInstance(
                 new ConnectorAwareNodeManager(nodeManager, nodeInfo.getEnvironment(), connectorId),
                 typeManager,
+                procedureRegistry,
                 metadataManager.getFunctionAndTypeManager(),
                 new FunctionResolution(metadataManager.getFunctionAndTypeManager().getFunctionAndTypeResolver()),
                 pageSorter,

presto-main-base/src/main/java/com/facebook/presto/testing/LocalQueryRunner.java

@@ -506,6 +506,7 @@ private LocalQueryRunner(Session defaultSession, FeaturesConfig featuresConfig,
                 nodeManager,
                 nodeInfo,
                 metadata.getFunctionAndTypeManager(),
+                procedureRegistry,
                 pageSorter,
                 pageIndexerFactory,
                 transactionManager,

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

@@ -22,6 +22,7 @@
 import com.facebook.presto.spi.function.FunctionMetadataManager;
 import com.facebook.presto.spi.function.StandardFunctionResolution;
 import com.facebook.presto.spi.plan.FilterStatsCalculatorService;
+import com.facebook.presto.spi.procedure.IProcedureRegistry;
 import com.facebook.presto.spi.relation.RowExpressionService;
 
 public interface ConnectorContext
@@ -36,6 +37,11 @@ default TypeManager getTypeManager()
         throw new UnsupportedOperationException();
     }
 
+    default IProcedureRegistry getProcedureRegistry()
+    {
+        throw new UnsupportedOperationException();
+    }
+
     default FunctionMetadataManager getFunctionMetadataManager()
     {
         throw new UnsupportedOperationException();