From 4361e4ed11a63cafaef7d7a943fdf55813a6fd09 Mon Sep 17 00:00:00 2001
From: Konrad Dziedzic <konraddziedzic@gmail.com>
Date: Wed, 15 Nov 2023 16:26:59 +0100
Subject: [PATCH] Add table functions to hive connector

---
 .../io/trino/plugin/hive/HiveConnector.java   | 20 +++++++++++++++++++
 .../java/io/trino/plugin/hive/HiveModule.java |  7 +++++++
 .../hive/InternalHiveConnectorFactory.java    |  4 ++++
 3 files changed, 31 insertions(+)

diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java
index 862c380f0a40..6771a8216ade 100644
--- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java
+++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java
@@ -30,6 +30,8 @@
 import io.trino.spi.connector.ConnectorTransactionHandle;
 import io.trino.spi.connector.TableProcedureMetadata;
 import io.trino.spi.eventlistener.EventListener;
+import io.trino.spi.function.FunctionProvider;
+import io.trino.spi.function.table.ConnectorTableFunction;
 import io.trino.spi.procedure.Procedure;
 import io.trino.spi.session.PropertyMetadata;
 import io.trino.spi.transaction.IsolationLevel;
@@ -67,6 +69,8 @@ public class HiveConnector
     private final ClassLoader classLoader;
 
     private final HiveTransactionManager transactionManager;
+    private final Set<ConnectorTableFunction> connectorTableFunctions;
+    private final Optional<FunctionProvider> functionProvider;
     private final boolean singleStatementWritesOnly;
 
     public HiveConnector(
@@ -87,6 +91,8 @@ public HiveConnector(
             List<PropertyMetadata<?>> analyzeProperties,
             List<PropertyMetadata<?>> materializedViewProperties,
             Optional<ConnectorAccessControl> accessControl,
+            Set<ConnectorTableFunction> connectorTableFunctions,
+            Optional<FunctionProvider> functionProvider,
             boolean singleStatementWritesOnly,
             ClassLoader classLoader)
     {
@@ -109,6 +115,8 @@ public HiveConnector(
         this.analyzeProperties = ImmutableList.copyOf(requireNonNull(analyzeProperties, "analyzeProperties is null"));
         this.materializedViewProperties = requireNonNull(materializedViewProperties, "materializedViewProperties is null");
         this.accessControl = requireNonNull(accessControl, "accessControl is null");
+        this.connectorTableFunctions = ImmutableSet.copyOf(requireNonNull(connectorTableFunctions, "connectorTableFunctions is null"));
+        this.functionProvider = requireNonNull(functionProvider, "functionProvider is null");
         this.singleStatementWritesOnly = singleStatementWritesOnly;
         this.classLoader = requireNonNull(classLoader, "classLoader is null");
     }
@@ -232,6 +240,18 @@ public final void shutdown()
         lifeCycleManager.stop();
     }
 
+    @Override
+    public Optional<FunctionProvider> getFunctionProvider()
+    {
+        return functionProvider;
+    }
+
+    @Override
+    public Set<ConnectorTableFunction> getTableFunctions()
+    {
+        return connectorTableFunctions;
+    }
+
     @Override
     public Set<TableProcedureMetadata> getTableProcedures()
     {
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java
index e44202c0acf1..2f0e6a6b97ea 100644
--- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java
+++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java
@@ -19,6 +19,7 @@
 import com.google.inject.Provides;
 import com.google.inject.Scopes;
 import com.google.inject.Singleton;
+import com.google.inject.TypeLiteral;
 import com.google.inject.multibindings.Multibinder;
 import io.airlift.event.client.EventClient;
 import io.trino.plugin.base.CatalogName;
@@ -53,7 +54,10 @@
 import io.trino.spi.connector.ConnectorPageSinkProvider;
 import io.trino.spi.connector.ConnectorPageSourceProvider;
 import io.trino.spi.connector.ConnectorSplitManager;
+import io.trino.spi.function.FunctionProvider;
+import io.trino.spi.function.table.ConnectorTableFunction;
 
+import java.util.Optional;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 
@@ -144,6 +148,9 @@ public void configure(Binder binder)
         configBinder(binder).bindConfig(ParquetReaderConfig.class);
         configBinder(binder).bindConfig(ParquetWriterConfig.class);
         fileWriterFactoryBinder.addBinding().to(ParquetFileWriterFactory.class).in(Scopes.SINGLETON);
+
+        newOptionalBinder(binder, new TypeLiteral<Optional<FunctionProvider>>(){}).setDefault().toInstance(Optional.empty());
+        newSetBinder(binder, ConnectorTableFunction.class);
     }
 
     @Singleton
diff --git a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java
index 69a5fe1089d1..b740d5f42fb0 100644
--- a/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java
+++ b/plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java
@@ -62,6 +62,8 @@
 import io.trino.spi.connector.MetadataProvider;
 import io.trino.spi.connector.TableProcedureMetadata;
 import io.trino.spi.eventlistener.EventListener;
+import io.trino.spi.function.FunctionProvider;
+import io.trino.spi.function.table.ConnectorTableFunction;
 import io.trino.spi.procedure.Procedure;
 import org.weakref.jmx.guice.MBeanModule;
 
@@ -173,6 +175,8 @@ public static Connector createConnector(
                     hiveAnalyzeProperties.getAnalyzeProperties(),
                     hiveMaterializedViewPropertiesProvider.getMaterializedViewProperties(),
                     hiveAccessControl,
+                    injector.getInstance(Key.get(new TypeLiteral<Set<ConnectorTableFunction>>() {})),
+                    injector.getInstance(Key.get(new TypeLiteral<Optional<FunctionProvider>>() {})),
                     injector.getInstance(HiveConfig.class).isSingleStatementWritesOnly(),
                     classLoader);
         }