Add Support to Managed KafkaIO for GCP specific Schema Registries that require extra parameters. (#35085)

Author: fozzie15

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

@@ -21,6 +21,7 @@
 import static org.apache.beam.sdk.util.construction.BeamUrns.getUrn;
 
 import com.google.auto.service.AutoService;
+import io.confluent.kafka.serializers.KafkaAvroDeserializerConfig;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
@@ -132,6 +133,13 @@ public List<String> outputCollectionNames() {
 
   static class KafkaReadSchemaTransform extends SchemaTransform {
     private final KafkaReadSchemaTransformConfiguration configuration;
+    private static final String googleManagedSchemaRegistryPrefix =
+        "https://managedkafka.googleapis.com/";
+
+    enum SchemaRegistryProvider {
+      UNSPECIFIED,
+      GOOGLE_MANAGED
+    }
 
     KafkaReadSchemaTransform(KafkaReadSchemaTransformConfiguration configuration) {
       this.configuration = configuration;
@@ -151,6 +159,13 @@ Row getConfigurationRow() {
       }
     }
 
+    private SchemaRegistryProvider getSchemaRegistryProvider(String confluentSchemaRegUrl) {
+      if (confluentSchemaRegUrl.contains(googleManagedSchemaRegistryPrefix)) {
+        return SchemaRegistryProvider.GOOGLE_MANAGED;
+      }
+      return SchemaRegistryProvider.UNSPECIFIED;
+    }
+
     @Override
     public PCollectionRowTuple expand(PCollectionRowTuple input) {
       configuration.validate();
@@ -178,16 +193,41 @@ public PCollectionRowTuple expand(PCollectionRowTuple input) {
       if (confluentSchemaRegUrl != null) {
         final String confluentSchemaRegSubject =
             checkArgumentNotNull(configuration.getConfluentSchemaRegistrySubject());
-        KafkaIO.Read<byte[], GenericRecord> kafkaRead =
+        KafkaIO.Read<byte[], GenericRecord> kafkaRead;
+
+        kafkaRead =
             KafkaIO.<byte[], GenericRecord>read()
                 .withTopic(configuration.getTopic())
                 .withConsumerFactoryFn(new ConsumerFactoryWithGcsTrustStores())
                 .withBootstrapServers(configuration.getBootstrapServers())
                 .withConsumerConfigUpdates(consumerConfigs)
-                .withKeyDeserializer(ByteArrayDeserializer.class)
-                .withValueDeserializer(
+                .withKeyDeserializer(ByteArrayDeserializer.class);
+
+        SchemaRegistryProvider provider = getSchemaRegistryProvider(confluentSchemaRegUrl);
+        switch (provider) {
+          case GOOGLE_MANAGED:
+            // Custom configs to authenticate with Google's Managed Schema Registry
+            Map<String, Object> configs = new HashMap<>();
+            configs.put(
+                KafkaAvroDeserializerConfig.SCHEMA_REGISTRY_URL_CONFIG, confluentSchemaRegUrl);
+            configs.put(KafkaAvroDeserializerConfig.BEARER_AUTH_CREDENTIALS_SOURCE, "CUSTOM");
+            configs.put(
+                "bearer.auth.custom.provider.class",
+                "com.google.cloud.hosted.kafka.auth.GcpBearerAuthCredentialProvider");
+
+            LOG.info("Constructing read transform with Google Managed Schema Registry URL.");
+            kafkaRead =
+                kafkaRead.withValueDeserializer(
+                    ConfluentSchemaRegistryDeserializerProvider.of(
+                        confluentSchemaRegUrl, confluentSchemaRegSubject, null, configs));
+            break;
+          case UNSPECIFIED:
+            kafkaRead =
+                kafkaRead.withValueDeserializer(
                     ConfluentSchemaRegistryDeserializerProvider.of(
                         confluentSchemaRegUrl, confluentSchemaRegSubject));
+        }
+
         Integer maxReadTimeSeconds = configuration.getMaxReadTimeSeconds();
         if (maxReadTimeSeconds != null) {
           kafkaRead = kafkaRead.withMaxReadTime(Duration.standardSeconds(maxReadTimeSeconds));

sdks/java/io/kafka/src/test/java/org/apache/beam/sdk/io/kafka/KafkaReadSchemaTransformProviderTest.java

@@ -155,6 +155,26 @@ public void testBuildTransformWithAvroSchema() {
             .build());
   }
 
+  @Test
+  public void testBuildTransformWithAvroSchemaRegistry() {
+    ServiceLoader<SchemaTransformProvider> serviceLoader =
+        ServiceLoader.load(SchemaTransformProvider.class);
+    List<SchemaTransformProvider> providers =
+        StreamSupport.stream(serviceLoader.spliterator(), false)
+            .filter(provider -> provider.getClass() == KafkaReadSchemaTransformProvider.class)
+            .collect(Collectors.toList());
+    KafkaReadSchemaTransformProvider kafkaProvider =
+        (KafkaReadSchemaTransformProvider) providers.get(0);
+    kafkaProvider.from(
+        KafkaReadSchemaTransformConfiguration.builder()
+            .setFormat("AVRO")
+            .setTopic("anytopic")
+            .setBootstrapServers("anybootstrap")
+            .setConfluentSchemaRegistryUrl("anyschemaregistryurl")
+            .setConfluentSchemaRegistrySubject("anysubject")
+            .build());
+  }
+
   @Test
   public void testBuildTransformWithJsonSchema() throws IOException {
     ServiceLoader<SchemaTransformProvider> serviceLoader =