runWithContextClassLoader in getNextPage to solve classloader issue

JingsongLi · Mar 10, 2023 · 9c493a8 · 9c493a8
1 parent c1451f1
commit 9c493a8
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/src/main/java/org/apache/flink/table/store/trino/TrinoPageSourceBase.java b/src/main/java/org/apache/flink/table/store/trino/TrinoPageSourceBase.java
@@ -52,6 +52,7 @@
 import java.math.BigDecimal;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.function.Supplier;
 
 import static io.airlift.slice.Slices.wrappedBuffer;
 import static io.trino.spi.StandardErrorCode.GENERIC_INTERNAL_ERROR;
@@ -108,11 +109,13 @@ public boolean isFinished() {
 
     @Override
     public Page getNextPage() {
-        try {
-            return nextPage();
-        } catch (IOException e) {
-            throw new UncheckedIOException(e);
-        }
+        return ClassLoaderUtils.runWithContextClassLoader(() -> {
+            try {
+                return nextPage();
+            } catch (IOException e) {
+                throw new UncheckedIOException(e);
+            }
+        }, TrinoPageSourceBase.class.getClassLoader());
     }
 
     private Page nextPage() throws IOException {