Fix lack of valid class loader for TableChangesSplitProcessor in Delta/Iceberg

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeTableFunctionProcessorProvider.java

@@ -0,0 +1,53 @@
+/*
+ * 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 io.trino.plugin.base.classloader;
+
+import io.trino.spi.classloader.ThreadContextClassLoader;
+import io.trino.spi.connector.ConnectorSession;
+import io.trino.spi.connector.ConnectorSplit;
+import io.trino.spi.function.table.ConnectorTableFunctionHandle;
+import io.trino.spi.function.table.TableFunctionDataProcessor;
+import io.trino.spi.function.table.TableFunctionProcessorProvider;
+import io.trino.spi.function.table.TableFunctionSplitProcessor;
+
+import static java.util.Objects.requireNonNull;
+
+public final class ClassLoaderSafeTableFunctionProcessorProvider
+        implements TableFunctionProcessorProvider
+{
+    private final TableFunctionProcessorProvider delegate;
+    private final ClassLoader classLoader;
+
+    public ClassLoaderSafeTableFunctionProcessorProvider(TableFunctionProcessorProvider delegate, ClassLoader classLoader)
+    {
+        this.delegate = requireNonNull(delegate, "delegate is null");
+        this.classLoader = requireNonNull(classLoader, "classLoader is null");
+    }
+
+    @Override
+    public TableFunctionDataProcessor getDataProcessor(ConnectorTableFunctionHandle handle)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getDataProcessor(handle);
+        }
+    }
+
+    @Override
+    public TableFunctionSplitProcessor getSplitProcessor(ConnectorSession session, ConnectorTableFunctionHandle handle, ConnectorSplit split)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getSplitProcessor(session, handle, split);
+        }
+    }
+}

lib/trino-plugin-toolkit/src/main/java/io/trino/plugin/base/classloader/ClassLoaderSafeTableFunctionSplitProcessor.java

@@ -0,0 +1,41 @@
+/*
+ * 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 io.trino.plugin.base.classloader;
+
+import io.trino.spi.classloader.ThreadContextClassLoader;
+import io.trino.spi.function.table.TableFunctionProcessorState;
+import io.trino.spi.function.table.TableFunctionSplitProcessor;
+
+import static java.util.Objects.requireNonNull;
+
+public final class ClassLoaderSafeTableFunctionSplitProcessor
+        implements TableFunctionSplitProcessor
+{
+    private final TableFunctionSplitProcessor delegate;
+    private final ClassLoader classLoader;
+
+    public ClassLoaderSafeTableFunctionSplitProcessor(TableFunctionSplitProcessor delegate, ClassLoader classLoader)
+    {
+        this.delegate = requireNonNull(delegate, "delegate is null");
+        this.classLoader = requireNonNull(classLoader, "classLoader is null");
+    }
+
+    @Override
+    public TableFunctionProcessorState process()
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.process();
+        }
+    }
+}

lib/trino-plugin-toolkit/src/test/java/io/trino/plugin/base/classloader/TestClassLoaderSafeWrappers.java

@@ -27,6 +27,8 @@
 import io.trino.spi.connector.SystemTable;
 import io.trino.spi.eventlistener.EventListener;
 import io.trino.spi.function.table.ConnectorTableFunction;
+import io.trino.spi.function.table.TableFunctionProcessorProvider;
+import io.trino.spi.function.table.TableFunctionSplitProcessor;
 import org.junit.jupiter.api.Test;
 
 import java.lang.reflect.Method;
@@ -58,6 +60,8 @@ public void test()
         testClassLoaderSafe(RecordSet.class, ClassLoaderSafeRecordSet.class);
         testClassLoaderSafe(EventListener.class, ClassLoaderSafeEventListener.class);
         testClassLoaderSafe(ConnectorTableFunction.class, ClassLoaderSafeConnectorTableFunction.class);
+        testClassLoaderSafe(TableFunctionSplitProcessor.class, ClassLoaderSafeTableFunctionSplitProcessor.class);
+        testClassLoaderSafe(TableFunctionProcessorProvider.class, ClassLoaderSafeTableFunctionProcessorProvider.class);
     }
 
     private static <I, C extends I> void testClassLoaderSafe(Class<I> iface, Class<C> clazz)

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/DeltaLakeFunctionProvider.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.deltalake;
 
 import com.google.inject.Inject;
+import io.trino.plugin.base.classloader.ClassLoaderSafeTableFunctionProcessorProvider;
 import io.trino.plugin.deltalake.functions.tablechanges.TableChangesProcessorProvider;
 import io.trino.plugin.deltalake.functions.tablechanges.TableChangesTableFunctionHandle;
 import io.trino.spi.function.FunctionProvider;
@@ -37,7 +38,7 @@ public DeltaLakeFunctionProvider(TableChangesProcessorProvider tableChangesProce
     public TableFunctionProcessorProvider getTableFunctionProcessorProvider(ConnectorTableFunctionHandle functionHandle)
     {
         if (functionHandle instanceof TableChangesTableFunctionHandle) {
-            return tableChangesProcessorProvider;
+            return new ClassLoaderSafeTableFunctionProcessorProvider(tableChangesProcessorProvider, getClass().getClassLoader());
         }
         throw new UnsupportedOperationException("Unsupported function: " + functionHandle);
     }

plugin/trino-delta-lake/src/main/java/io/trino/plugin/deltalake/functions/tablechanges/TableChangesProcessorProvider.java

@@ -16,6 +16,7 @@
 import com.google.inject.Inject;
 import io.trino.filesystem.TrinoFileSystemFactory;
 import io.trino.parquet.ParquetReaderOptions;
+import io.trino.plugin.base.classloader.ClassLoaderSafeTableFunctionSplitProcessor;
 import io.trino.plugin.deltalake.DeltaLakeConfig;
 import io.trino.plugin.hive.FileFormatDataSourceStats;
 import io.trino.plugin.hive.parquet.ParquetReaderConfig;
@@ -54,14 +55,15 @@ public TableChangesProcessorProvider(
     @Override
     public TableFunctionSplitProcessor getSplitProcessor(ConnectorSession session, ConnectorTableFunctionHandle handle, ConnectorSplit split)
     {
-        return new TableChangesFunctionProcessor(
+        return new ClassLoaderSafeTableFunctionSplitProcessor(new TableChangesFunctionProcessor(
                 session,
                 fileSystemFactory,
                 parquetDateTimeZone,
                 domainCompactionThreshold,
                 fileFormatDataSourceStats,
                 parquetReaderOptions,
                 (TableChangesTableFunctionHandle) handle,
-                (TableChangesSplit) split);
+                (TableChangesSplit) split),
+                getClass().getClassLoader());
     }
 }

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveConnector.java

@@ -70,7 +70,7 @@ public class HiveConnector
 
     private final HiveTransactionManager transactionManager;
     private final Set<ConnectorTableFunction> connectorTableFunctions;
-    private final Optional<FunctionProvider> functionProvider;
+    private final FunctionProvider functionProvider;
     private final boolean singleStatementWritesOnly;
 
     public HiveConnector(
@@ -92,7 +92,7 @@ public HiveConnector(
             List<PropertyMetadata<?>> materializedViewProperties,
             Optional<ConnectorAccessControl> accessControl,
             Set<ConnectorTableFunction> connectorTableFunctions,
-            Optional<FunctionProvider> functionProvider,
+            FunctionProvider functionProvider,
             boolean singleStatementWritesOnly,
             ClassLoader classLoader)
     {
@@ -243,7 +243,7 @@ public final void shutdown()
     @Override
     public Optional<FunctionProvider> getFunctionProvider()
     {
-        return functionProvider;
+        return Optional.of(functionProvider);
     }
 
     @Override

plugin/trino-hive/src/main/java/io/trino/plugin/hive/HiveModule.java

@@ -19,7 +19,6 @@
 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;
@@ -57,7 +56,6 @@
 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;
 
@@ -149,7 +147,7 @@ public void configure(Binder binder)
         configBinder(binder).bindConfig(ParquetWriterConfig.class);
         fileWriterFactoryBinder.addBinding().to(ParquetFileWriterFactory.class).in(Scopes.SINGLETON);
 
-        newOptionalBinder(binder, new TypeLiteral<Optional<FunctionProvider>>(){}).setDefault().toInstance(Optional.empty());
+        newOptionalBinder(binder, FunctionProvider.class).setDefault().toInstance(new NoopFunctionProvider());
         newSetBinder(binder, ConnectorTableFunction.class);
     }
 

plugin/trino-hive/src/main/java/io/trino/plugin/hive/InternalHiveConnectorFactory.java

@@ -174,7 +174,7 @@ public static Connector createConnector(
                     hiveMaterializedViewPropertiesProvider.getMaterializedViewProperties(),
                     hiveAccessControl,
                     injector.getInstance(Key.get(new TypeLiteral<Set<ConnectorTableFunction>>() {})),
-                    injector.getInstance(Key.get(new TypeLiteral<Optional<FunctionProvider>>() {})),
+                    injector.getInstance(FunctionProvider.class),
                     injector.getInstance(HiveConfig.class).isSingleStatementWritesOnly(),
                     classLoader);
         }

plugin/trino-hive/src/main/java/io/trino/plugin/hive/NoopFunctionProvider.java

@@ -0,0 +1,20 @@
+/*
+ * 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 io.trino.plugin.hive;
+
+import io.trino.spi.function.FunctionProvider;
+
+public class NoopFunctionProvider
+        implements FunctionProvider
+{}

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/IcebergFunctionProvider.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.iceberg.functions;
 
 import com.google.inject.Inject;
+import io.trino.plugin.base.classloader.ClassLoaderSafeTableFunctionProcessorProvider;
 import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionHandle;
 import io.trino.plugin.iceberg.functions.tablechanges.TableChangesFunctionProcessorProvider;
 import io.trino.spi.function.FunctionProvider;
@@ -37,7 +38,7 @@ public IcebergFunctionProvider(TableChangesFunctionProcessorProvider tableChange
     public TableFunctionProcessorProvider getTableFunctionProcessorProvider(ConnectorTableFunctionHandle functionHandle)
     {
         if (functionHandle instanceof TableChangesFunctionHandle) {
-            return tableChangesFunctionProcessorProvider;
+            return new ClassLoaderSafeTableFunctionProcessorProvider(tableChangesFunctionProcessorProvider, getClass().getClassLoader());
         }
 
         throw new UnsupportedOperationException("Unsupported function: " + functionHandle);

plugin/trino-iceberg/src/main/java/io/trino/plugin/iceberg/functions/tablechanges/TableChangesFunctionProcessorProvider.java

@@ -14,6 +14,7 @@
 package io.trino.plugin.iceberg.functions.tablechanges;
 
 import com.google.inject.Inject;
+import io.trino.plugin.base.classloader.ClassLoaderSafeTableFunctionSplitProcessor;
 import io.trino.plugin.iceberg.IcebergPageSourceProvider;
 import io.trino.spi.connector.ConnectorSession;
 import io.trino.spi.connector.ConnectorSplit;
@@ -40,10 +41,11 @@ public TableFunctionSplitProcessor getSplitProcessor(
             ConnectorTableFunctionHandle handle,
             ConnectorSplit split)
     {
-        return new TableChangesFunctionProcessor(
+        return new ClassLoaderSafeTableFunctionSplitProcessor(new TableChangesFunctionProcessor(
                 session,
                 (TableChangesFunctionHandle) handle,
                 (TableChangesSplit) split,
-                icebergPageSourceProvider);
+                icebergPageSourceProvider),
+                getClass().getClassLoader());
     }
 }