Add ClassLoaderSafeSystemTable

presto-iceberg/src/main/java/io/prestosql/plugin/iceberg/IcebergMetadata.java

@@ -43,6 +43,7 @@
 import io.prestosql.spi.connector.SchemaTablePrefix;
 import io.prestosql.spi.connector.SystemTable;
 import io.prestosql.spi.connector.TableNotFoundException;
+import io.prestosql.spi.connector.classloader.ClassLoaderSafeSystemTable;
 import io.prestosql.spi.statistics.ComputedStatistics;
 import io.prestosql.spi.type.DecimalType;
 import io.prestosql.spi.type.TimestampType;
@@ -153,7 +154,8 @@ public Optional<SystemTable> getSystemTable(ConnectorSession session, SchemaTabl
         IcebergTableHandle table = IcebergTableHandle.from(tableName);
         if (table.getTableType() == TableType.PARTITIONS) {
             org.apache.iceberg.Table icebergTable = getIcebergTable(metastore, hdfsEnvironment, session, table.getSchemaTableName());
-            return Optional.of(new PartitionTable(table, session, typeManager, icebergTable));
+            SystemTable systemTable = new PartitionTable(table, session, typeManager, icebergTable);
+            return Optional.of(new ClassLoaderSafeSystemTable(systemTable, getClass().getClassLoader()));
         }
         return Optional.empty();
     }

presto-spi/src/main/java/io/prestosql/spi/connector/classloader/ClassLoaderSafeSystemTable.java

@@ -0,0 +1,70 @@
+/*
+ * 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.prestosql.spi.connector.classloader;
+
+import io.prestosql.spi.classloader.ThreadContextClassLoader;
+import io.prestosql.spi.connector.ConnectorPageSource;
+import io.prestosql.spi.connector.ConnectorSession;
+import io.prestosql.spi.connector.ConnectorTableMetadata;
+import io.prestosql.spi.connector.ConnectorTransactionHandle;
+import io.prestosql.spi.connector.RecordCursor;
+import io.prestosql.spi.connector.SystemTable;
+import io.prestosql.spi.predicate.TupleDomain;
+
+import static java.util.Objects.requireNonNull;
+
+public class ClassLoaderSafeSystemTable
+        implements SystemTable
+{
+    private final SystemTable delegate;
+    private final ClassLoader classLoader;
+
+    public ClassLoaderSafeSystemTable(SystemTable delegate, ClassLoader classLoader)
+    {
+        this.delegate = requireNonNull(delegate, "delegate is null");
+        this.classLoader = requireNonNull(classLoader, "classLoader is null");
+    }
+
+    @Override
+    public Distribution getDistribution()
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getDistribution();
+        }
+    }
+
+    @Override
+    public ConnectorTableMetadata getTableMetadata()
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.getTableMetadata();
+        }
+    }
+
+    @Override
+    public RecordCursor cursor(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.cursor(transactionHandle, session, constraint);
+        }
+    }
+
+    @Override
+    public ConnectorPageSource pageSource(ConnectorTransactionHandle transactionHandle, ConnectorSession session, TupleDomain<Integer> constraint)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(classLoader)) {
+            return delegate.pageSource(transactionHandle, session, constraint);
+        }
+    }
+}

presto-spi/src/test/java/io/prestosql/spi/connector/classloader/TestClassLoaderSafeWrappers.java

@@ -20,6 +20,7 @@
 import io.prestosql.spi.connector.ConnectorPageSourceProvider;
 import io.prestosql.spi.connector.ConnectorSplitManager;
 import io.prestosql.spi.connector.ConnectorSplitSource;
+import io.prestosql.spi.connector.SystemTable;
 import org.testng.annotations.Test;
 
 import static io.prestosql.spi.testing.InterfaceTestUtils.assertAllMethodsOverridden;
@@ -36,5 +37,6 @@ public void testAllMethodsOverridden()
         assertAllMethodsOverridden(ConnectorSplitManager.class, ClassLoaderSafeConnectorSplitManager.class);
         assertAllMethodsOverridden(ConnectorNodePartitioningProvider.class, ClassLoaderSafeNodePartitioningProvider.class);
         assertAllMethodsOverridden(ConnectorSplitSource.class, ClassLoaderSafeConnectorSplitSource.class);
+        assertAllMethodsOverridden(SystemTable.class, ClassLoaderSafeSystemTable.class);
     }
 }