HIVE-28965: tez.DagUtils: Failed to add credential supplier, ClassNotFoundException: org.apache.hadoop.hive.kafka.KafkaDagCredentialSupplier

Author: KiranVelumuri

ql/pom.xml

@@ -861,6 +861,21 @@
       <artifactId>jersey-multipart</artifactId>
       <version>${jersey.version}</version>
     </dependency>
+    <dependency>
+      <groupId>org.apache.kafka</groupId>
+      <artifactId>kafka-clients</artifactId>
+      <version>${kafka.version}</version>
+      <exclusions>
+        <exclusion>
+          <groupId>com.github.luben</groupId>
+          <artifactId>zstd-jni</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>org.apache.shiro</groupId>
+          <artifactId>shiro-crypto-cipher</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
     <dependency>
       <groupId>com.esri.geometry</groupId>
       <artifactId>esri-geometry-api</artifactId>

ql/src/java/org/apache/hadoop/hive/ql/exec/tez/DagUtils.java

@@ -41,6 +41,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Properties;
 import java.util.Set;
 import java.util.Stack;
 import java.util.concurrent.TimeUnit;
@@ -105,6 +106,7 @@
 import org.apache.hadoop.hive.ql.plan.MapWork;
 import org.apache.hadoop.hive.ql.plan.MergeJoinWork;
 import org.apache.hadoop.hive.ql.plan.OperatorDesc;
+import org.apache.hadoop.hive.ql.plan.PartitionDesc;
 import org.apache.hadoop.hive.ql.plan.ReduceWork;
 import org.apache.hadoop.hive.ql.plan.TableDesc;
 import org.apache.hadoop.hive.ql.plan.TezEdgeProperty;
@@ -133,6 +135,8 @@
 import org.apache.hadoop.yarn.api.records.URL;
 import org.apache.hadoop.yarn.util.ConverterUtils;
 import org.apache.hadoop.yarn.util.Records;
+import org.apache.kafka.clients.CommonClientConfigs;
+import org.apache.kafka.common.security.auth.SecurityProtocol;
 import org.apache.tez.common.TezUtils;
 import org.apache.tez.dag.api.DAG;
 import org.apache.tez.dag.api.DataSinkDescriptor;
@@ -179,13 +183,29 @@ public class DagUtils {
   public static final String TEZ_TMP_DIR_KEY = "_hive_tez_tmp_dir";
   private static final Logger LOG = LoggerFactory.getLogger(DagUtils.class.getName());
   private static final String TEZ_DIR = "_tez_scratch_dir";
-  private static final DagUtils instance = new DagUtils(defaultCredentialSuppliers());
+  private static final DagUtils instance = new DagUtils();
   // The merge file being currently processed.
   public static final String TEZ_MERGE_CURRENT_MERGE_FILE_PREFIX =
       "hive.tez.current.merge.file.prefix";
   // A comma separated list of work names used as prefix.
   public static final String TEZ_MERGE_WORK_FILE_PREFIXES = "hive.tez.merge.file.prefixes";
-  private final List<DagCredentialSupplier> credentialSuppliers;
+  private static final List<DagCredentialSupplier> credentialSuppliers = new ArrayList<>();
+
+  /**
+   * MANDATORY Table property indicating kafka broker(s) connection string.
+   */
+  private static final String HIVE_KAFKA_BOOTSTRAP_SERVERS = "kafka.bootstrap.servers";
+
+  /**
+   * Table property prefix used to inject kafka consumer properties, e.g "kafka.consumer.max.poll.records" = "5000"
+   * this will lead to inject max.poll.records=5000 to the Kafka Consumer. NOT MANDATORY defaults to nothing
+   */
+  private static final String CONSUMER_CONFIGURATION_PREFIX = "kafka.consumer";
+
+  /**
+   * Table property prefix used to inject kafka producer properties, e.g "kafka.producer.lingers.ms" = "100".
+   */
+  private static final String PRODUCER_CONFIGURATION_PREFIX = "kafka.producer";
   /**
    * Notifiers to synchronize resource localization across threads. If one thread is localizing
    * a file, other threads can wait on the corresponding notifier object instead of just sleeping
@@ -286,6 +306,23 @@ private void getCredentialsFromSuppliers(BaseWork work, Set<TableDesc> tables, D
     if (!UserGroupInformation.isSecurityEnabled()){
       return;
     }
+    Map<String, PartitionDesc> partitions = ((MapWork) work).getAliasToPartnInfo();
+
+    // We don't need to iterate on all partitions, and check the same TableDesc.
+    PartitionDesc partition = partitions.values().stream().findFirst().orElse(null);
+    if (partition != null) {
+      TableDesc tableDesc = partition.getTableDesc();
+      if (isTokenRequired(tableDesc)) {
+        addCredentialSuppliers(Collections.singletonList("org.apache.hadoop.hive.kafka.KafkaDagCredentialSupplier"));
+      }
+    }
+
+    for (TableDesc tableDesc : tables) {
+      if (isTokenRequired(tableDesc)) {
+        addCredentialSuppliers(Collections.singletonList("org.apache.hadoop.hive.kafka.KafkaDagCredentialSupplier"));
+        break;
+      }
+    }
     for (DagCredentialSupplier supplier : credentialSuppliers) {
       Text alias = supplier.getTokenAlias();
       Token<?> t = dag.getCredentials().getToken(alias);
@@ -300,11 +337,10 @@ private void getCredentialsFromSuppliers(BaseWork work, Set<TableDesc> tables, D
     }
   }
 
-  private static List<DagCredentialSupplier> defaultCredentialSuppliers() {
+  @VisibleForTesting
+  static void addCredentialSuppliers(List<String> supplierClassNames) {
     // Class names of credential providers that should be used when adding credentials to the dag.
     // Use plain strings instead of {@link Class#getName()} to avoid compile scope dependencies to other modules.
-    List<String> supplierClassNames =
-        Collections.singletonList("org.apache.hadoop.hive.kafka.KafkaDagCredentialSupplier");
     List<DagCredentialSupplier> dagSuppliers = new ArrayList<>();
     for (String s : supplierClassNames) {
       try {
@@ -314,7 +350,46 @@ private static List<DagCredentialSupplier> defaultCredentialSuppliers() {
         LOG.error("Failed to add credential supplier", e);
       }
     }
-    return dagSuppliers;
+    credentialSuppliers.addAll(dagSuppliers);
+  }
+
+  /**
+   * Returns the security protocol if one is defined in the properties and null otherwise.
+   * <p>The following properties are examined to determine the protocol:</p>
+   * <ol>
+   *   <li>security.protocol</li>
+   *   <li>kafka.consumer.security.protocol</li>
+   *   <li>kafka.producer.security.protocol</li>
+   * </ol>
+   * <p>and the first non null/not empty is returned.</p>
+   * <p>Defining multiple security protocols at the same time is invalid but this method is lenient and tries to pick
+   * the most reasonable option.</p>
+   * @param props the properties from which to obtain the protocol.
+   * @return the security protocol if one is defined in the properties and null otherwise.
+   */
+  static SecurityProtocol getSecurityProtocol(Properties props) {
+    String[] securityProtocolConfigs = new String[] { CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+            CONSUMER_CONFIGURATION_PREFIX + "." + CommonClientConfigs.SECURITY_PROTOCOL_CONFIG,
+            PRODUCER_CONFIGURATION_PREFIX + "." + CommonClientConfigs.SECURITY_PROTOCOL_CONFIG };
+    for (String c : securityProtocolConfigs) {
+      String v = props.getProperty(c);
+      if (v != null && !v.isEmpty()) {
+        return SecurityProtocol.forName(v);
+      }
+    }
+    return null;
+  }
+
+  /**
+   * Returns whether a Kafka token is required for performing operations on the specified table.
+   * If "security.protocol" is set to "PLAINTEXT", we don't need to collect delegation token at all.
+   *
+   * @return true if a Kafka token is required for performing operations on the specified table and false otherwise.
+   */
+  private boolean isTokenRequired(TableDesc tableDesc) {
+    String kafkaBrokers = tableDesc.getProperties().getProperty(HIVE_KAFKA_BOOTSTRAP_SERVERS);
+    SecurityProtocol protocol = getSecurityProtocol(tableDesc.getProperties());
+    return !StringUtils.isEmpty(kafkaBrokers) && SecurityProtocol.PLAINTEXT != protocol;
   }
 
   private void collectNeededFileSinkData(BaseWork work, Set<URI> fileSinkUris, Set<TableDesc> fileSinkTableDescs) {
@@ -1697,8 +1772,8 @@ public static String getUserSpecifiedDagName(Configuration conf) {
   }
 
   @VisibleForTesting
-  DagUtils(List<DagCredentialSupplier> suppliers) {
-    this.credentialSuppliers = suppliers;
+  DagUtils() {
+    // don't instantiate
   }
 
   /**

ql/src/test/org/apache/hadoop/hive/ql/exec/tez/TestDagUtilsSecurityEnabled.java

@@ -59,7 +59,8 @@ public static void clear() {
   @Test
   public void testAddCredentialsWithCredentialSupplierNewTokenAdded() {
     IncrementalIntDagCredentialSupplier supplier = new IncrementalIntDagCredentialSupplier();
-    DagUtils dagUtils = new DagUtils(Collections.singletonList(supplier));
+    DagUtils dagUtils = new DagUtils();
+    DagUtils.addCredentialSuppliers(Collections.singletonList("org.apache.hadoop.hive.ql.exec.tez.IncrementalIntDagCredentialSupplier"));
     DAG dag = DAG.create("test_credentials_dag");
 
     dagUtils.addCredentials(mock(MapWork.class), dag, null);
@@ -70,7 +71,8 @@ public void testAddCredentialsWithCredentialSupplierNewTokenAdded() {
   @Test
   public void testAddCredentialsWithCredentialSupplierTokenExistsNothingAdded() {
     IncrementalIntDagCredentialSupplier supplier = new IncrementalIntDagCredentialSupplier();
-    DagUtils dagUtils = new DagUtils(Collections.singletonList(supplier));
+    DagUtils dagUtils = new DagUtils();
+    DagUtils.addCredentialSuppliers(Collections.singletonList("org.apache.hadoop.hive.ql.exec.tez.IncrementalIntDagCredentialSupplier"));
     DAG dag = DAG.create("test_credentials_dag");
     Token<TokenIdentifier> oldToken = new Token<>();
     // Add explicitly the token in the DAG before calling addCredentials simulating the use-case where the DAG has already the token