Allow JDBC connector to define procedure

trinodb · Jan 29, 2019 · 120c28f · 120c28f
1 parent f958f81
commit 120c28f
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 1 deletion.
diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcConnector.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcConnector.java
@@ -13,6 +13,7 @@
  */
 package io.prestosql.plugin.jdbc;
 
+import com.google.common.collect.ImmutableSet;
 import io.airlift.bootstrap.LifeCycleManager;
 import io.airlift.log.Logger;
 import io.prestosql.spi.connector.Connector;
@@ -22,11 +23,13 @@
 import io.prestosql.spi.connector.ConnectorRecordSetProvider;
 import io.prestosql.spi.connector.ConnectorSplitManager;
 import io.prestosql.spi.connector.ConnectorTransactionHandle;
+import io.prestosql.spi.procedure.Procedure;
 import io.prestosql.spi.transaction.IsolationLevel;
 
 import javax.inject.Inject;
 
 import java.util.Optional;
+import java.util.Set;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 
@@ -46,6 +49,7 @@ public class JdbcConnector
     private final JdbcRecordSetProvider jdbcRecordSetProvider;
     private final JdbcPageSinkProvider jdbcPageSinkProvider;
     private final Optional<ConnectorAccessControl> accessControl;
+    private final Set<Procedure> procedures;
 
     private final ConcurrentMap<ConnectorTransactionHandle, JdbcMetadata> transactions = new ConcurrentHashMap<>();
 
@@ -56,14 +60,16 @@ public JdbcConnector(
             JdbcSplitManager jdbcSplitManager,
             JdbcRecordSetProvider jdbcRecordSetProvider,
             JdbcPageSinkProvider jdbcPageSinkProvider,
-            Optional<ConnectorAccessControl> accessControl)
+            Optional<ConnectorAccessControl> accessControl,
+            Set<Procedure> procedures)
     {
         this.lifeCycleManager = requireNonNull(lifeCycleManager, "lifeCycleManager is null");
         this.jdbcMetadataFactory = requireNonNull(jdbcMetadataFactory, "jdbcMetadataFactory is null");
         this.jdbcSplitManager = requireNonNull(jdbcSplitManager, "jdbcSplitManager is null");
         this.jdbcRecordSetProvider = requireNonNull(jdbcRecordSetProvider, "jdbcRecordSetProvider is null");
         this.jdbcPageSinkProvider = requireNonNull(jdbcPageSinkProvider, "jdbcPageSinkProvider is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
+        this.procedures = ImmutableSet.copyOf(requireNonNull(procedures, "procedures is null"));
     }
 
     @Override
@@ -127,6 +133,12 @@ public ConnectorAccessControl getAccessControl()
         return accessControl.orElseThrow(UnsupportedOperationException::new);
     }
 
+    @Override
+    public Set<Procedure> getProcedures()
+    {
+        return procedures;
+    }
+
     @Override
     public final void shutdown()
     {

diff --git a/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcModule.java b/presto-base-jdbc/src/main/java/io/prestosql/plugin/jdbc/JdbcModule.java
@@ -17,7 +17,9 @@
 import com.google.inject.Module;
 import com.google.inject.Scopes;
 import io.prestosql.spi.connector.ConnectorAccessControl;
+import io.prestosql.spi.procedure.Procedure;
 
+import static com.google.inject.multibindings.Multibinder.newSetBinder;
 import static com.google.inject.multibindings.OptionalBinder.newOptionalBinder;
 import static io.airlift.configuration.ConfigBinder.configBinder;
 import static java.util.Objects.requireNonNull;
@@ -36,6 +38,7 @@ public JdbcModule(String connectorId)
     public void configure(Binder binder)
     {
         newOptionalBinder(binder, ConnectorAccessControl.class);
+        newSetBinder(binder, Procedure.class);
         binder.bind(JdbcConnectorId.class).toInstance(new JdbcConnectorId(connectorId));
         binder.bind(JdbcMetadataFactory.class).in(Scopes.SINGLETON);
         binder.bind(JdbcSplitManager.class).in(Scopes.SINGLETON);