Add plumbing for python use case.

Author: fozzie15

sdks/java/io/expansion-service/build.gradle

@@ -76,6 +76,8 @@ dependencies {
   permitUnusedDeclared project(":sdks:java:io:kafka") // BEAM-11761
   implementation project(":sdks:java:io:kafka:upgrade")
   permitUnusedDeclared project(":sdks:java:io:kafka:upgrade") // BEAM-11761
+  implementation project(":sdks:java:extensions:kafka-factories")
+  permitUnusedDeclared project(":sdks:java:extensions:kafka-factories")
 
   if (JavaVersion.current().compareTo(JavaVersion.VERSION_11) >= 0 && project.findProperty('testJavaVersion') != '8') {
     // iceberg ended support for Java 8 in 1.7.0

sdks/java/io/kafka/src/main/java/org/apache/beam/sdk/io/kafka/KafkaIO.java

@@ -27,6 +27,7 @@
 import io.confluent.kafka.serializers.KafkaAvroDeserializer;
 import java.io.InputStream;
 import java.io.OutputStream;
+import java.lang.reflect.Constructor;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -94,6 +95,8 @@
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.Manual;
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.MonotonicallyIncreasing;
 import org.apache.beam.sdk.transforms.splittabledofn.WatermarkEstimators.WallTime;
+import org.apache.beam.sdk.util.InstanceBuilder;
+import org.apache.beam.sdk.transforms.windowing.GlobalWindow;
 import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.sdk.util.construction.PTransformMatchers;
 import org.apache.beam.sdk.util.construction.ReplacementOutputs;
@@ -930,6 +933,25 @@ static <K, V> void setupExternalBuilder(
           builder.setOffsetDeduplication(false);
           builder.setRedistributeByRecordKey(false);
         }
+
+        if (config.consumerFactoryFnClass != null) {
+          if (config.consumerFactoryFnClass.contains("KerberosConsumerFactoryFn")) {
+            try {
+              if (!config.consumerFactoryFnClass.contains("krb5Location")) {
+                throw new IllegalArgumentException("The KerberosConsumerFactoryFn requires a location for the krb5.conf file. " +
+                    "Please provide either a GCS location or Google Secret Manager location for this file.");
+              }
+                builder.setConsumerFactoryFn(InstanceBuilder.ofType(new TypeDescriptor<SerializableFunction<Map<String, Object>, Consumer<byte[], byte[]>>>() {})
+                    .fromClassName(config.consumerFactoryFnClass)
+                    .withArg(
+                            String.class,
+                            config.consumerFactoryFnParams.get("krb5Location"))
+                    .build());
+            } catch (Exception e) {
+              throw new RuntimeException("Unable to construct FactoryFn " + config.consumerFactoryFnClass + ": " + e.getMessage(), e);
+            }
+          }
+        }
       }
 
       private static <T> Coder<T> resolveCoder(Class<Deserializer<T>> deserializer) {
@@ -1000,6 +1022,8 @@ public static class Configuration {
         private Boolean offsetDeduplication;
         private Boolean redistributeByRecordKey;
         private Long dynamicReadPollIntervalSeconds;
+        private String consumerFactoryFnClass;
+        private Map<String, String> consumerFactoryFnParams;
 
         public void setConsumerConfig(Map<String, String> consumerConfig) {
           this.consumerConfig = consumerConfig;
@@ -1068,6 +1092,14 @@ public void setRedistributeByRecordKey(Boolean redistributeByRecordKey) {
         public void setDynamicReadPollIntervalSeconds(Long dynamicReadPollIntervalSeconds) {
           this.dynamicReadPollIntervalSeconds = dynamicReadPollIntervalSeconds;
         }
+
+        public void setConsumerFactoryFnClass(String consumerFactoryFnClass) {
+            this.consumerFactoryFnClass = consumerFactoryFnClass;
+        }
+
+        public void setConsumerFactoryFnParams(Map<String, String> consumerFactoryFnParams) {
+            this.consumerFactoryFnParams = consumerFactoryFnParams;
+          }
       }
     }
 

sdks/python/apache_beam/io/external/xlang_kafkaio_it_test.py

@@ -41,6 +41,10 @@
 from apache_beam.transforms.userstate import BagStateSpec
 from apache_beam.transforms.userstate import CombiningValueStateSpec
 from apache_beam.utils import subprocess_server
+from apache_beam.options.pipeline_options import (
+    PipelineOptions,
+    GoogleCloudOptions
+)
 
 NUM_RECORDS = 1000
 
@@ -68,7 +72,7 @@ def process(
 
 class CrossLanguageKafkaIO(object):
   def __init__(
-      self, bootstrap_servers, topic, null_key, expansion_service=None):
+      self, bootstrap_servers=None, topic=None, null_key=None, expansion_service=None):
     self.bootstrap_servers = bootstrap_servers
     self.topic = topic
     self.null_key = null_key
@@ -107,6 +111,34 @@ def build_read_pipeline(self, pipeline, max_num_records=None):
         | 'CalculateSum' >> beam.ParDo(CollectingFn())
         | 'SetSumCounter' >> beam.Map(self.sum_counter.inc))
 
+  def build_read_pipeline_with_kerberos(self, p, max_num_records=None):
+
+      jaas_config = (
+          f'com.sun.security.auth.module.Krb5LoginModule required '
+          f'useKeyTab=true storeKey=true '
+          f'keyTab="secretValue:projects/dataflow-testing-311516/secrets/kafka-client-keytab/versions/latest" '
+          f'principal="[email protected]";'
+      )
+
+      kafka_records = (
+          p
+          | 'ReadFromKafka' >> ReadFromKafka(
+              consumer_config={
+                  'bootstrap.servers': 'fozzie-test-kafka-broker.us-central1-c.c.dataflow-testing-311516.internal:9092',
+                  'auto.offset.reset': 'earliest',
+                  'max_num_records': max_num_records,
+                  'security.protocol': 'SASL_PLAINTEXT',
+                  'sasl.mechanism': 'GSSAPI',
+                  'sasl.kerberos.service.name': 'kafka',
+                  'sasl.jaas.config': jaas_config
+              },
+              topics=['fozzie_test_kerberos_topic'],
+              key_deserializer='org.apache.kafka.common.serialization.StringDeserializer',
+              value_deserializer='org.apache.kafka.common.serialization.StringDeserializer',
+              consumer_factory_fn_class='org.apache.beam.sdk.extensions.kafka.factories.KerberosConsumerFactoryFn',
+              consumer_factory_fn_params={'krb5Location': 'gs://fozzie_testing_bucket/kerberos/krb5.conf'}))
+      return kafka_records
+
   def run_xlang_kafkaio(self, pipeline):
     self.build_write_pipeline(pipeline)
     self.build_read_pipeline(pipeline)
@@ -182,6 +214,17 @@ def test_hosted_kafkaio_null_key(self):
     self.run_kafka_write(pipeline_creator)
     self.run_kafka_read(pipeline_creator, None)
 
+  def test_hosted_kafkaio_null_key_kerberos(self):
+      kafka_topic = 'xlang_kafkaio_test_null_key_{}'.format(uuid.uuid4())
+      bootstrap_servers = 'fozzie-test-kafka-broker.us-central1-c.c.dataflow-testing-311516.internal:9092'
+      pipeline_creator = CrossLanguageKafkaIO(
+          bootstrap_servers,
+          kafka_topic,
+          True,
+          'localhost:%s' % os.environ.get('EXPANSION_PORT'))
+
+      self.run_kafka_read_with_kerberos(pipeline_creator)
+
   def run_kafka_write(self, pipeline_creator):
     with TestPipeline() as pipeline:
       pipeline.not_use_test_runner_api = True
@@ -196,6 +239,23 @@ def run_kafka_read(self, pipeline_creator, expected_key):
           equal_to([(expected_key, str(i).encode())
                     for i in range(NUM_RECORDS)]))
 
+  def run_kafka_read_with_kerberos(self, pipeline_creator):
+      options_dict = {
+          'runner': 'DataflowRunner',
+          'project': 'dataflow-testing-311516',
+          'region': 'us-central1',
+          'streaming': False
+      }
+      options = PipelineOptions.from_dictionary(options_dict)
+      expected_records = [f'test{i}' for i in range(1, 12)]
+      with beam.Pipeline(options=options) as p:
+          pipeline.not_use_test_runner_api = True
+          result = pipeline_creator.build_read_pipeline_with_kerberos(p, max_num_records=11)
+          assert_that(
+              result,
+              equal_to(expected_records)
+          )
+
   def get_platform_localhost(self):
     if sys.platform == 'darwin':
       return 'host.docker.internal'

sdks/python/apache_beam/io/kafka.py

@@ -124,7 +124,9 @@
         ('redistribute', typing.Optional[bool]),
         ('redistribute_num_keys', typing.Optional[np.int32]),
         ('allow_duplicates', typing.Optional[bool]),
-        ('dynamic_read_poll_interval_seconds', typing.Optional[int]),
+        ('dynamic_read_poll_interval_seconds', typing.Optional[int])
+        ('consumer_factory_fn', typing.Optional[str]),
+        ('consumer_factory_fn_params', typing.Optional[typing.Mapping[str, str]])
     ])
 
 
@@ -173,6 +175,8 @@ def __init__(
       redistribute_num_keys=np.int32(0),
       allow_duplicates=False,
       dynamic_read_poll_interval_seconds: typing.Optional[int] = None,
+      consumer_factory_fn_class=None,
+      consumer_factory_fn_params=None
   ):
     """
     Initializes a read operation from Kafka.
@@ -216,6 +220,13 @@ def __init__(
     :param dynamic_read_poll_interval_seconds: The interval in seconds at which
         to check for new partitions. If not None, dynamic partition discovery
         is enabled.
+    :param consumer_factory_fn_class: A fully qualified classpath to an existing provided
+        consumerFactoryFn. If not None, this will construct Kafka consumers with a custom
+        configuration.
+    :param consumer_factory_fn_params: A map which specifies the parameters for the provided
+        consumer_factory_fn_class. IF not None, the values in this map will be used when
+        constructing the consumer_factory_fn_class object. This cannot be null
+        if the consumer_factory_fn_class is not null.
     """
     if timestamp_policy not in [ReadFromKafka.processing_time_policy,
                                 ReadFromKafka.create_time_policy,
@@ -242,7 +253,9 @@ def __init__(
                 redistribute_num_keys=redistribute_num_keys,
                 allow_duplicates=allow_duplicates,
                 dynamic_read_poll_interval_seconds=
-                dynamic_read_poll_interval_seconds)),
+                dynamic_read_poll_interval_seconds,
+                consumer_factory_fn_class=consumer_factory_fn_class,
+                consumer_factory_fn_params=consumer_factory_fn_params)),
         expansion_service or default_io_expansion_service())