From 120c28fb643d812608be606a48e1316d654bdce4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Grzegorz=20Kokosi=C5=84ski?= <grzegorz@starburstdata.com>
Date: Wed, 16 Jan 2019 11:17:18 +0100
Subject: [PATCH] Allow JDBC connector to define procedure

---
 .../io/prestosql/plugin/jdbc/JdbcConnector.java    | 14 +++++++++++++-
 .../java/io/prestosql/plugin/jdbc/JdbcModule.java  |  3 +++
 2 files changed, 16 insertions(+), 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
index 9c36ec1a1b02..84803a0b3914 100644
--- 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,7 +60,8 @@ 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");
@@ -64,6 +69,7 @@ public JdbcConnector(
         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
index 352be87d83a7..e2275cc464d5 100644
--- 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);