[vpj] Route repush ZSTD dictionary reads through configured adapter

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/VenicePushJob.java

@@ -6,6 +6,7 @@
 import static com.linkedin.venice.ConfigKeys.KAFKA_PRODUCER_REQUEST_TIMEOUT_MS;
 import static com.linkedin.venice.ConfigKeys.KAFKA_PRODUCER_RETRIES_CONFIG;
 import static com.linkedin.venice.ConfigKeys.MULTI_REGION;
+import static com.linkedin.venice.ConfigKeys.PUBSUB_BROKER_ADDRESS;
 import static com.linkedin.venice.ConfigKeys.VENICE_PARTITIONERS;
 import static com.linkedin.venice.VeniceConstants.DEFAULT_SSL_FACTORY_CLASS_NAME;
 import static com.linkedin.venice.status.BatchJobHeartbeatConfigs.HEARTBEAT_ENABLED_CONFIG;
@@ -866,13 +867,9 @@ public void run() {
       if (pushJobSetting.isSourceKafka) {
         if (pushJobSetting.sourceVersionCompressionStrategy == CompressionStrategy.ZSTD_WITH_DICT) {
           LOGGER.info("Source version uses ZSTD_WITH_DICT. Fetching source dictionary.");
-          Properties kafkaConsumerProperties = new Properties();
-          if (pushJobSetting.enableSSL) {
-            kafkaConsumerProperties.putAll(this.sslProperties.get());
-          }
-          kafkaConsumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, pushJobSetting.repushSourcePubsubBroker);
-          ByteBuffer sourceDict = DictionaryUtils
-              .readDictionaryFromKafka(pushJobSetting.kafkaInputTopic, new VeniceProperties(kafkaConsumerProperties));
+          ByteBuffer sourceDict = DictionaryUtils.readDictionaryFromKafka(
+              pushJobSetting.kafkaInputTopic,
+              new VeniceProperties(getRepushDictionaryConsumerProperties()));
           if (sourceDict != null) {
             pushJobSetting.sourceDictionary = ByteUtils.extractByteArray(sourceDict);
           }
@@ -1666,12 +1663,46 @@ private Optional<ByteBuffer> getCompressionDictionary() throws VeniceException {
     return Optional.of(emptyPushZstdDictionary.get());
   }
 
+  /**
+   * Build the pub-sub consumer properties used to read the repush source topic's compression dictionary.
+   * Seeds from the full job config so the configured pub-sub consumer adapter factory (e.g. xinfra) and its
+   * client configs are honored, then overlays SSL and the repush source broker. Seeding only SSL + broker would
+   * drop {@code pubsub.consumer.adapter.factory.class} and silently fall back to the default Apache Kafka
+   * consumer, which cannot read the source topic on a non-Kafka (e.g. xinfra) deployment.
+   */
+  private Properties getRepushDictionaryConsumerProperties() {
+    return buildRepushDictionaryConsumerProperties(
+        props,
+        pushJobSetting.enableSSL ? sslProperties.get() : new Properties(),
+        pushJobSetting.repushSourcePubsubBroker);
+  }
+
+  /**
+   * Assemble the pub-sub consumer properties used to read the repush source topic's compression dictionary. The
+   * source dictionary must be read with the same pub-sub client as the rest of the repush, so seed from the full
+   * job configuration (which carries {@code pubsub.consumer.adapter.factory.class} and any client-specific
+   * settings, such as xinfra's routing maps), overlay SSL, and point the broker properties at the repush source
+   * broker. Seeding only SSL and the broker address drops the adapter-factory class and silently falls back to
+   * the default Apache Kafka consumer, which cannot read the source topic on a non-Kafka (e.g. xinfra) deployment.
+   */
+  @VisibleForTesting
+  static Properties buildRepushDictionaryConsumerProperties(
+      VeniceProperties jobProperties,
+      Properties sslProperties,
+      String repushSourcePubsubBroker) {
+    Properties consumerProperties = jobProperties.toProperties();
+    consumerProperties.putAll(sslProperties);
+    consumerProperties.setProperty(PUBSUB_BROKER_ADDRESS, repushSourcePubsubBroker);
+    consumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, repushSourcePubsubBroker);
+    return consumerProperties;
+  }
+
   private ByteBuffer fetchOrBuildCompressionDictionary() throws VeniceException {
     // Prepare the param builder, which can be used by different scenarios.
     KafkaInputDictTrainer.ParamBuilder paramBuilder = new KafkaInputDictTrainer.ParamBuilder()
         .setKeySchema(AvroCompatibilityHelper.toParsingForm(pushJobSetting.storeKeySchema))
         .setNewKMESchemasFromController(pushJobSetting.newKmeSchemasFromController)
-        .setSslProperties(pushJobSetting.enableSSL ? sslProperties.get() : new Properties())
+        .setConsumerProperties(getRepushDictionaryConsumerProperties())
         .setCompressionDictSize(
             props.getInt(
                 COMPRESSION_DICTIONARY_SIZE_LIMIT,
@@ -1691,14 +1722,9 @@ private ByteBuffer fetchOrBuildCompressionDictionary() throws VeniceException {
           return ByteBuffer.wrap(dictTrainer.trainDict());
         } else {
           LOGGER.info("Reading Zstd dictionary from input topic: {}", pushJobSetting.kafkaInputTopic);
-          // set up ssl properties and kafka consumer properties
-          Properties kafkaConsumerProperties = new Properties();
-          if (pushJobSetting.enableSSL) {
-            kafkaConsumerProperties.putAll(this.sslProperties.get());
-          }
-          kafkaConsumerProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, pushJobSetting.repushSourcePubsubBroker);
-          return DictionaryUtils
-              .readDictionaryFromKafka(pushJobSetting.kafkaInputTopic, new VeniceProperties(kafkaConsumerProperties));
+          return DictionaryUtils.readDictionaryFromKafka(
+              pushJobSetting.kafkaInputTopic,
+              new VeniceProperties(getRepushDictionaryConsumerProperties()));
         }
       }
       LOGGER.info(

clients/venice-push-job/src/main/java/com/linkedin/venice/hadoop/input/kafka/KafkaInputDictTrainer.java

@@ -50,7 +50,7 @@ public static class Param {
     private final String kafkaInputBroker;
     private final String topicName;
     private final String keySchema;
-    private final Properties sslProperties;
+    private final Properties consumerProperties;
     private final int compressionDictSize;
     private final int dictSampleSize;
     private final CompressionStrategy sourceVersionCompressionStrategy;
@@ -62,7 +62,7 @@ public static class Param {
       this.kafkaInputBroker = builder.kafkaInputBroker;
       this.topicName = builder.topicName;
       this.keySchema = builder.keySchema;
-      this.sslProperties = builder.sslProperties;
+      this.consumerProperties = builder.consumerProperties;
       this.compressionDictSize = builder.compressionDictSize;
       this.dictSampleSize = builder.dictSampleSize;
       this.sourceVersionCompressionStrategy = builder.sourceVersionCompressionStrategy;
@@ -75,7 +75,7 @@ public static class ParamBuilder {
     private String kafkaInputBroker;
     private String topicName;
     private String keySchema;
-    private Properties sslProperties;
+    private Properties consumerProperties;
     private int compressionDictSize;
     private int dictSampleSize;
     private CompressionStrategy sourceVersionCompressionStrategy;
@@ -97,8 +97,14 @@ public ParamBuilder setKeySchema(String keySchema) {
       return this;
     }
 
-    public ParamBuilder setSslProperties(Properties sslProperties) {
-      this.sslProperties = sslProperties;
+    /**
+     * Properties the source-topic consumer is built from. These must include the job's pub-sub client
+     * configuration (e.g. {@code pubsub.consumer.adapter.factory.class} and any client-specific configs such as
+     * xinfra's) so the dictionary is read with the same pub-sub client as the rest of the repush; otherwise
+     * {@code PubSubClientsFactory.createConsumerFactory(...)} silently defaults to the Apache Kafka consumer.
+     */
+    public ParamBuilder setConsumerProperties(Properties consumerProperties) {
+      this.consumerProperties = consumerProperties;
       return this;
     }
 
@@ -170,7 +176,7 @@ protected KafkaInputDictTrainer(
     properties.setProperty(KAFKA_INPUT_TOPIC, param.topicName);
     properties.setProperty(KAFKA_SOURCE_KEY_SCHEMA_STRING_PROP, param.keySchema);
     this.sourceTopicName = param.topicName;
-    properties.putAll(param.sslProperties);
+    properties.putAll(param.consumerProperties);
     properties.setProperty(COMPRESSION_DICTIONARY_SIZE_LIMIT, Integer.toString(param.compressionDictSize));
     properties.setProperty(COMPRESSION_DICTIONARY_SAMPLE_SIZE, Integer.toString(param.dictSampleSize));
     properties

clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/VenicePushJobRepushTest.java

@@ -1,5 +1,9 @@
 package com.linkedin.venice.hadoop;
 
+import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS;
+import static com.linkedin.venice.ConfigKeys.PUBSUB_BROKER_ADDRESS;
+import static com.linkedin.venice.ConfigKeys.PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS;
+import static com.linkedin.venice.ConfigKeys.PUBSUB_SECURITY_PROTOCOL;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.ALLOW_REGULAR_PUSH_WITH_TTL_REPUSH;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.COMPLIANCE_PUSH;
 import static com.linkedin.venice.vpj.VenicePushJobConstants.KAFKA_INPUT_MAX_RECORDS_PER_MAPPER;
@@ -26,6 +30,7 @@
 import com.linkedin.venice.meta.StoreInfo;
 import com.linkedin.venice.meta.Version;
 import com.linkedin.venice.utils.Time;
+import com.linkedin.venice.utils.VeniceProperties;
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Properties;
@@ -42,6 +47,65 @@
  */
 
 public class VenicePushJobRepushTest extends VenicePushJobTestBase {
+  @Test
+  public void testBuildRepushDictionaryConsumerPropertiesRetainsAdapterConfigAndOverridesBrokers() {
+    Properties jobProperties = new Properties();
+    jobProperties.setProperty(
+        PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS,
+        "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory");
+    jobProperties.setProperty("xc.pubsub.broker.url.to.region.name.map", "northguard:ei4");
+    jobProperties.setProperty(PUBSUB_BROKER_ADDRESS, "destination-broker");
+    jobProperties.setProperty(KAFKA_BOOTSTRAP_SERVERS, "legacy-broker");
+
+    Properties consumerProperties = VenicePushJob.buildRepushDictionaryConsumerProperties(
+        new VeniceProperties(jobProperties),
+        new Properties(),
+        "source-broker");
+
+    // The configured pub-sub adapter factory and its client-specific configs must survive so the source
+    // dictionary is read with the same adapter as the rest of the repush instead of the implicit Kafka default.
+    assertEquals(
+        consumerProperties.getProperty(PUBSUB_CONSUMER_ADAPTER_FACTORY_CLASS),
+        "com.linkedin.venice.pubsub.adapter.xinfra.consumer.XcConsumerAdapterFactory");
+    assertEquals(consumerProperties.getProperty("xc.pubsub.broker.url.to.region.name.map"), "northguard:ei4");
+    // Both broker properties are pointed at the repush source broker.
+    assertEquals(consumerProperties.getProperty(PUBSUB_BROKER_ADDRESS), "source-broker");
+    assertEquals(consumerProperties.getProperty(KAFKA_BOOTSTRAP_SERVERS), "source-broker");
+  }
+
+  @Test
+  public void testBuildRepushDictionaryConsumerPropertiesAppliesSslOverrides() {
+    Properties jobProperties = new Properties();
+    jobProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "PLAINTEXT");
+    jobProperties.setProperty("ssl.keystore.location", "stale-keystore");
+
+    Properties sslProperties = new Properties();
+    sslProperties.setProperty(PUBSUB_SECURITY_PROTOCOL, "SSL");
+    sslProperties.setProperty("ssl.keystore.location", "credential-keystore");
+
+    Properties consumerProperties = VenicePushJob
+        .buildRepushDictionaryConsumerProperties(new VeniceProperties(jobProperties), sslProperties, "source-broker");
+
+    assertEquals(consumerProperties.getProperty(PUBSUB_SECURITY_PROTOCOL), "SSL");
+    assertEquals(consumerProperties.getProperty("ssl.keystore.location"), "credential-keystore");
+  }
+
+  @Test
+  public void testBuildRepushDictionaryConsumerPropertiesDoesNotMutateJobProperties() {
+    Properties jobProperties = new Properties();
+    jobProperties.setProperty(PUBSUB_BROKER_ADDRESS, "destination-broker");
+
+    VenicePushJob.buildRepushDictionaryConsumerProperties(
+        new VeniceProperties(jobProperties),
+        new Properties(),
+        "source-broker");
+
+    assertEquals(
+        jobProperties.getProperty(PUBSUB_BROKER_ADDRESS),
+        "destination-broker",
+        "Building consumer properties must not mutate the job properties");
+  }
+
   @Test(expectedExceptions = VeniceException.class, expectedExceptionsMessageRegExp = ".*Repush with TTL is only supported while using Kafka Input Format.*")
   public void testRepushTTLJobWithNonKafkaInput() {
     Properties repushProps = new Properties();

clients/venice-push-job/src/test/java/com/linkedin/venice/hadoop/input/kafka/TestKafkaInputDictTrainer.java

@@ -60,7 +60,7 @@ private KafkaInputDictTrainer.Param getParam(int sampleSize, CompressionStrategy
         .setKeySchema("\"string\"")
         .setCompressionDictSize(900 * 1024)
         .setDictSampleSize(sampleSize)
-        .setSslProperties(new Properties())
+        .setConsumerProperties(new Properties())
         .setSourceVersionCompressionStrategy(sourceVersionCompressionStrategy)
         .setNewKMESchemasFromController(allSchemaStr)
         .build();