NIFI-3574 - PutHiveStreaming UGI fixes

Signed-off-by: Matt Burgess <mattyb149@apache.org>

This closes #1578

nifi-commons/nifi-hadoop-utils/src/main/java/org/apache/nifi/hadoop/SecurityUtil.java

@@ -28,6 +28,8 @@ import java.io.IOException;
  * interfering with each other.
  */
 public class SecurityUtil {
+    public static final String HADOOP_SECURITY_AUTHENTICATION = "hadoop.security.authentication";
+    public static final String KERBEROS = "kerberos";
 
     /**
      * Initializes UserGroupInformation with the given Configuration and performs the login for the given principal
@@ -81,7 +83,7 @@ public class SecurityUtil {
      */
     public static boolean isSecurityEnabled(final Configuration config) {
         Validate.notNull(config);
-        return "kerberos".equalsIgnoreCase(config.get("hadoop.security.authentication"));
+        return KERBEROS.equalsIgnoreCase(config.get(HADOOP_SECURITY_AUTHENTICATION));
     }
 
     /**

nifi-nar-bundles/nifi-hive-bundle/nifi-hive-processors/src/main/java/org/apache/nifi/processors/hive/PutHiveStreaming.java

@@ -17,7 +17,6 @@
 package org.apache.nifi.processors.hive;
 
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
-import java.io.File;
 import org.apache.avro.Schema;
 import org.apache.avro.file.CodecFactory;
 import org.apache.avro.file.DataFileConstants;
@@ -32,6 +31,7 @@ import org.apache.hive.hcatalog.streaming.ConnectionError;
 import org.apache.hive.hcatalog.streaming.HiveEndPoint;
 import org.apache.hive.hcatalog.streaming.SerializationError;
 import org.apache.hive.hcatalog.streaming.StreamingException;
+import org.apache.nifi.annotation.behavior.RequiresInstanceClassLoading;
 import org.apache.nifi.annotation.behavior.TriggerSerially;
 import org.apache.nifi.annotation.behavior.WritesAttribute;
 import org.apache.nifi.annotation.behavior.WritesAttributes;
@@ -61,6 +61,7 @@ import org.apache.nifi.util.hive.HiveWriter;
 import org.json.JSONException;
 import org.json.JSONObject;
 
+import java.io.File;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
@@ -94,6 +95,7 @@ import java.util.regex.Pattern;
         @WritesAttribute(attribute = "hivestreaming.record.count", description = "This attribute is written on the flow files routed to the 'success' "
                 + "and 'failure' relationships, and contains the number of records from the incoming flow file written successfully and unsuccessfully, respectively.")
 })
+@RequiresInstanceClassLoading
 public class PutHiveStreaming extends AbstractProcessor {
 
     // Attributes
@@ -349,6 +351,8 @@ public class PutHiveStreaming extends AbstractProcessor {
             }
             log.info("Successfully logged in as principal {} with keytab {}", new Object[]{principal, keyTab});
             options = options.withKerberosPrincipal(principal).withKerberosKeytab(keyTab);
+        } else {
+            ugi = null;
         }
 
         allWriters = new ConcurrentHashMap<>();
@@ -662,14 +666,14 @@ public class PutHiveStreaming extends AbstractProcessor {
         callTimeoutPool.shutdown();
         try {
             while (!callTimeoutPool.isTerminated()) {
-                callTimeoutPool.awaitTermination(
+                callTimeoutPool.awaitTermination(options.getCallTimeOut(), TimeUnit.MILLISECONDS);
-                        options.getCallTimeOut(), TimeUnit.MILLISECONDS);
             }
         } catch (Throwable t) {
             log.warn("shutdown interrupted on " + callTimeoutPool, t);
         }
 
         callTimeoutPool = null;
+        ugi = null;
         hiveConfigurator.stopRenewer();
     }
 

nifi-nar-bundles/nifi-hive-bundle/nifi-hive-processors/src/main/java/org/apache/nifi/util/hive/HiveConfigurator.java

@@ -74,7 +74,8 @@ public class HiveConfigurator {
 
     public void preload(Configuration configuration) {
         try {
-            FileSystem.get(configuration);
+            FileSystem.get(configuration).close();
+            UserGroupInformation.setConfiguration(configuration);
         } catch (IOException ioe) {
             // Suppress exception as future uses of this configuration will fail
         }

nifi-nar-bundles/nifi-hive-bundle/nifi-hive-processors/src/main/java/org/apache/nifi/util/hive/HiveWriter.java

@@ -53,6 +53,7 @@ public class HiveWriter {
     private final ExecutorService callTimeoutPool;
     private final long callTimeout;
     private final Object txnBatchLock = new Object();
+    private final UserGroupInformation ugi;
     private TransactionBatch txnBatch;
     private long lastUsed; // time of last flush on this writer
     protected boolean closed; // flag indicating HiveWriter was closed
@@ -61,6 +62,7 @@ public class HiveWriter {
     public HiveWriter(HiveEndPoint endPoint, int txnsPerBatch, boolean autoCreatePartitions, long callTimeout, ExecutorService callTimeoutPool, UserGroupInformation ugi, HiveConf hiveConf)
             throws InterruptedException, ConnectFailure {
         try {
+            this.ugi = ugi;
             this.callTimeout = callTimeout;
             this.callTimeoutPool = callTimeoutPool;
             this.endPoint = endPoint;
@@ -348,7 +350,12 @@ public class HiveWriter {
      */
     private <T> T callWithTimeout(final CallRunner<T> callRunner)
             throws TimeoutException, StreamingException, InterruptedException {
-        Future<T> future = callTimeoutPool.submit(callRunner::call);
+        Future<T> future = callTimeoutPool.submit(() -> {
+            if (ugi == null) {
+                return callRunner.call();
+            }
+            return ugi.doAs((PrivilegedExceptionAction<T>) () -> callRunner.call());
+        });
         try {
             if (callTimeout > 0) {
                 return future.get(callTimeout, TimeUnit.MILLISECONDS);

nifi-nar-bundles/nifi-hive-bundle/nifi-hive-processors/src/test/java/org/apache/nifi/processors/hive/TestPutHiveStreaming.java

@@ -33,10 +33,12 @@ import org.apache.hive.hcatalog.streaming.StreamingConnection;
 import org.apache.hive.hcatalog.streaming.StreamingException;
 import org.apache.hive.hcatalog.streaming.TransactionBatch;
 import org.apache.nifi.hadoop.KerberosProperties;
+import org.apache.nifi.hadoop.SecurityUtil;
 import org.apache.nifi.stream.io.ByteArrayOutputStream;
 import org.apache.nifi.util.MockFlowFile;
 import org.apache.nifi.util.TestRunner;
 import org.apache.nifi.util.TestRunners;
+import org.apache.nifi.util.hive.AuthenticationFailedException;
 import org.apache.nifi.util.hive.HiveConfigurator;
 import org.apache.nifi.util.hive.HiveOptions;
 import org.apache.nifi.util.hive.HiveWriter;
@@ -60,7 +62,10 @@ import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertTrue;
+import static org.mockito.Matchers.any;
+import static org.mockito.Matchers.anyLong;
 import static org.mockito.Matchers.anyString;
+import static org.mockito.Matchers.eq;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -75,6 +80,7 @@ public class TestPutHiveStreaming {
     private KerberosProperties kerberosPropsWithFile;
     private HiveConfigurator hiveConfigurator;
     private HiveConf hiveConf;
+    private UserGroupInformation ugi;
 
     @Before
     public void setUp() throws Exception {
@@ -84,6 +90,7 @@ public class TestPutHiveStreaming {
         System.setProperty("java.security.krb5.realm", "nifi.com");
         System.setProperty("java.security.krb5.kdc", "nifi.kdc");
 
+        ugi = null;
         kerberosPropsWithFile = new KerberosProperties(new File("src/test/resources/krb5.conf"));
 
         processor = new MockPutHiveStreaming();
@@ -107,6 +114,36 @@ public class TestPutHiveStreaming {
         runner.run();
     }
 
+    @Test
+    public void testUgiGetsCleared() {
+        runner.setValidateExpressionUsage(false);
+        runner.setProperty(PutHiveStreaming.METASTORE_URI, "thrift://localhost:9083");
+        runner.setProperty(PutHiveStreaming.DB_NAME, "default");
+        runner.setProperty(PutHiveStreaming.TABLE_NAME, "users");
+        processor.ugi = mock(UserGroupInformation.class);
+        runner.run();
+        assertNull(processor.ugi);
+    }
+
+    @Test
+    public void testUgiGetsSetIfSecure() throws AuthenticationFailedException, IOException {
+        when(hiveConf.get(SecurityUtil.HADOOP_SECURITY_AUTHENTICATION)).thenReturn(SecurityUtil.KERBEROS);
+        ugi = mock(UserGroupInformation.class);
+        when(hiveConfigurator.authenticate(eq(hiveConf), anyString(), anyString(), anyLong(), any())).thenReturn(ugi);
+        runner.setValidateExpressionUsage(false);
+        runner.setProperty(PutHiveStreaming.METASTORE_URI, "thrift://localhost:9083");
+        runner.setProperty(PutHiveStreaming.DB_NAME, "default");
+        runner.setProperty(PutHiveStreaming.TABLE_NAME, "users");
+        Map<String, Object> user1 = new HashMap<String, Object>() {
+            {
+                put("name", "Joe");
+                put("favorite_number", 146);
+            }
+        };
+        runner.enqueue(createAvroRecord(Collections.singletonList(user1)));
+        runner.run();
+    }
+
     @Test
     public void testSetupBadPartitionColumns() throws Exception {
         runner.setValidateExpressionUsage(false);
@@ -608,6 +645,7 @@ public class TestPutHiveStreaming {
                 long callTimeout, ExecutorService callTimeoutPool, UserGroupInformation ugi, HiveConf hiveConf)
                 throws InterruptedException, ConnectFailure {
             super(endPoint, txnsPerBatch, autoCreatePartitions, callTimeout, callTimeoutPool, ugi, hiveConf);
+            assertEquals(TestPutHiveStreaming.this.ugi, ugi);
             this.endPoint = endPoint;
         }