Set context class loader inside Hive procedures

The procedures call Hadoop code which can load file systems or perform other actions that require the context class loader.
trinodb · Mar 9, 2019 · 77ee7bc · 77ee7bc
1 parent c493444
commit 77ee7bc
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/CreateEmptyPartitionProcedure.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/CreateEmptyPartitionProcedure.java
@@ -21,6 +21,7 @@
 import io.prestosql.plugin.hive.PartitionUpdate.UpdateMode;
 import io.prestosql.plugin.hive.metastore.ExtendedHiveMetastore;
 import io.prestosql.spi.PrestoException;
+import io.prestosql.spi.classloader.ThreadContextClassLoader;
 import io.prestosql.spi.connector.ConnectorSession;
 import io.prestosql.spi.procedure.Procedure;
 import io.prestosql.spi.procedure.Procedure.Argument;
@@ -83,6 +84,13 @@ public Procedure get()
     }
 
     public void createEmptyPartition(ConnectorSession session, String schema, String table, List<Object> partitionColumnNames, List<Object> partitionValues)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(getClass().getClassLoader())) {
+            doCreateEmptyPartition(session, schema, table, partitionColumnNames, partitionValues);
+        }
+    }
+
+    private void doCreateEmptyPartition(ConnectorSession session, String schema, String table, List<Object> partitionColumnNames, List<Object> partitionValues)
     {
         TransactionalMetadata hiveMetadata = hiveMetadataFactory.get();
 

diff --git a/presto-hive/src/main/java/io/prestosql/plugin/hive/SyncPartitionMetadataProcedure.java b/presto-hive/src/main/java/io/prestosql/plugin/hive/SyncPartitionMetadataProcedure.java
@@ -21,6 +21,7 @@
 import io.prestosql.plugin.hive.metastore.SemiTransactionalHiveMetastore;
 import io.prestosql.plugin.hive.metastore.Table;
 import io.prestosql.spi.PrestoException;
+import io.prestosql.spi.classloader.ThreadContextClassLoader;
 import io.prestosql.spi.connector.ConnectorSession;
 import io.prestosql.spi.connector.SchemaTableName;
 import io.prestosql.spi.connector.TableNotFoundException;
@@ -94,6 +95,13 @@ public Procedure get()
     }
 
     public void syncPartitionMetadata(ConnectorSession session, String schemaName, String tableName, String mode)
+    {
+        try (ThreadContextClassLoader ignored = new ThreadContextClassLoader(getClass().getClassLoader())) {
+            doSyncPartitionMetadata(session, schemaName, tableName, mode);
+        }
+    }
+
+    private void doSyncPartitionMetadata(ConnectorSession session, String schemaName, String tableName, String mode)
     {
         SyncMode syncMode = toSyncMode(mode);
         HdfsContext context = new HdfsContext(session, schemaName, tableName);