Replace Confluent ProtobufSchema parser with protoc compiler (awslabs#2796)

Jithendar12 · ritiktrianz · commit 39ecbdacd8ec · 2025-08-20T23:31:29.000+05:30
diff --git a/athena-msk/pom.xml b/athena-msk/pom.xml
@@ -43,28 +43,6 @@
                 <version>2.1.21</version>
                 <scope>runtime</scope>
             </dependency>
-            <dependency>
-                <groupId>com.squareup.wire</groupId>
-                <artifactId>wire-schema</artifactId>
-                <version>5.3.3</version>
-            </dependency>
-            <dependency>
-                <groupId>com.squareup.wire</groupId>
-                <artifactId>wire-schema-jvm</artifactId>
-                <version>4.9.0</version>
-            </dependency>
-            <dependency>
-                <groupId>com.squareup.wire</groupId>
-                <artifactId>wire-runtime-jvm</artifactId>
-                <version>5.3.3</version>
-                <scope>runtime</scope>
-            </dependency>
-            <dependency>
-                <groupId>com.squareup.wire</groupId>
-                <artifactId>wire-compiler</artifactId>
-                <version>5.3.3</version>
-                <scope>runtime</scope>
-            </dependency>
             <dependency>
                 <groupId>software.amazon.msk</groupId>
                 <artifactId>aws-msk-iam-auth</artifactId>
@@ -83,7 +61,7 @@
         <dependency>
             <groupId>org.apache.kafka</groupId>
             <artifactId>kafka-clients</artifactId>
-            <version>7.9.1-ce</version>
+            <version>4.0.0</version>
         </dependency>
         <dependency>
             <groupId>org.apache.avro</groupId>
@@ -95,16 +73,17 @@
             <artifactId>protobuf-java</artifactId>
             <version>3.25.5</version>
         </dependency>
+        <dependency>
+            <groupId>com.github.os72</groupId>
+            <artifactId>protoc-jar</artifactId>
+            <version>3.11.4</version>
+        </dependency>
+
         <dependency>
             <groupId>software.amazon.glue</groupId>
             <artifactId>schema-registry-serde</artifactId>
             <version>1.1.23</version>
         </dependency>
-        <dependency>
-            <groupId>io.confluent</groupId>
-            <artifactId>kafka-protobuf-provider</artifactId>
-            <version>7.7.2</version>
-        </dependency>
         <dependency>
             <groupId>com.fasterxml.jackson.core</groupId>
             <artifactId>jackson-annotations</artifactId>
@@ -274,10 +253,4 @@
             </plugin>
         </plugins>
     </build>
-    <repositories>
-        <repository>
-            <id>confluent</id>
-            <url>https://packages.confluent.io/maven/</url>
-        </repository>
-    </repositories>
 </project>
diff --git a/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandler.java b/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandler.java
@@ -38,12 +38,11 @@
 import com.amazonaws.athena.connector.lambda.metadata.ListTablesResponse;
 import com.amazonaws.athena.connector.util.PaginatedRequestIterator;
 import com.amazonaws.athena.connectors.msk.dto.AvroTopicSchema;
+import com.amazonaws.athena.connectors.msk.dto.MSKField;
 import com.amazonaws.athena.connectors.msk.dto.SplitParameters;
 import com.amazonaws.athena.connectors.msk.dto.TopicPartitionPiece;
 import com.amazonaws.athena.connectors.msk.dto.TopicSchema;
 import com.google.common.annotations.VisibleForTesting;
-import com.google.protobuf.Descriptors;
-import io.confluent.kafka.schemaregistry.protobuf.ProtobufSchema;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.FieldType;
 import org.apache.arrow.vector.types.pojo.Schema;
@@ -460,19 +459,21 @@ private Schema getSchema(String glueRegistryName, String glueSchemaName) throws
             schemaBuilder.addMetadata("dataFormat", AVRO_DATA_FORMAT);
         }
         else if (dataFormat.equalsIgnoreCase(PROTOBUF_DATA_FORMAT)) {
-            // Get protobuf topic schema from Glue registry
-            String glueSchema = registryReader.getSchemaDef(glueRegistryName, glueSchemaName);
-            ProtobufSchema protobufSchema = new ProtobufSchema(glueSchema);
-            Descriptors.Descriptor descriptor = protobufSchema.toDescriptor();
-            // Creating ArrowType for each field in the protobuf topic schema.
-            for (Descriptors.FieldDescriptor fieldDescriptor : descriptor.getFields()) {
+            // Get protobuf topic schema from Glue registry using protoc compiler
+            List<MSKField> fields = registryReader.getProtobufFields(glueRegistryName, glueSchemaName);
+            // Creating ArrowType for each field in the protobuf topic schema
+            for (MSKField field : fields) {
                 FieldType fieldType = new FieldType(
                         true,
-                        AmazonMskUtils.toArrowType(fieldDescriptor.getType().toString()),
-                        null
+                        AmazonMskUtils.toArrowType(field.getType()),
+                        null,
+                        com.google.common.collect.ImmutableMap.of(
+                                "name", field.getName(),
+                                "type", field.getType()
+                        )
                 );
-                Field field = new Field(fieldDescriptor.getName(), fieldType, null);
-                schemaBuilder.addField(field);
+                Field arrowField = new Field(field.getName(), fieldType, null);
+                schemaBuilder.addField(arrowField);
             }
             schemaBuilder.addMetadata("dataFormat", PROTOBUF_DATA_FORMAT);
         }
diff --git a/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/AmazonMskUtils.java b/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/AmazonMskUtils.java
@@ -100,7 +100,6 @@ enum AuthType {
     private static final String KAFKA_MAX_PARTITION_FETCH_BYTES_CONFIG = "max.partition.fetch.bytes";
     private static final String KAFKA_KEY_DESERIALIZER_CLASS_CONFIG = "key.deserializer";
     private static final String KAFKA_VALUE_DESERIALIZER_CLASS_CONFIG = "value.deserializer";
-
     private static final ObjectMapper objectMapper = new ObjectMapper();
 
     private AmazonMskUtils() {}
diff --git a/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/GlueRegistryReader.java b/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/GlueRegistryReader.java
@@ -19,26 +19,139 @@
  */
 package com.amazonaws.athena.connectors.msk;
 
+import com.amazonaws.athena.connector.lambda.exceptions.AthenaConnectorException;
+import com.amazonaws.athena.connectors.msk.dto.MSKField;
 import com.fasterxml.jackson.databind.DeserializationFeature;
 import com.fasterxml.jackson.databind.ObjectMapper;
+import com.github.os72.protocjar.Protoc;
+import com.google.protobuf.DescriptorProtos.DescriptorProto;
+import com.google.protobuf.DescriptorProtos.FieldDescriptorProto;
+import com.google.protobuf.DescriptorProtos.FileDescriptorSet;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import software.amazon.awssdk.services.glue.GlueClient;
+import software.amazon.awssdk.services.glue.model.ErrorDetails;
+import software.amazon.awssdk.services.glue.model.FederationSourceErrorCode;
 import software.amazon.awssdk.services.glue.model.GetSchemaRequest;
 import software.amazon.awssdk.services.glue.model.GetSchemaResponse;
 import software.amazon.awssdk.services.glue.model.GetSchemaVersionRequest;
 import software.amazon.awssdk.services.glue.model.GetSchemaVersionResponse;
 import software.amazon.awssdk.services.glue.model.SchemaId;
 import software.amazon.awssdk.services.glue.model.SchemaVersionNumber;
 
+import java.io.FileInputStream;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.UUID;
+
 public class GlueRegistryReader
 {
+    private static final Logger logger = LoggerFactory.getLogger(GlueRegistryReader.class);
     private static final ObjectMapper objectMapper;
+    private static final String PROTO_FILE = "schema.proto";
+    private static final String DESC_FILE = "schema.desc";
 
     static {
         objectMapper = new ObjectMapper();
         objectMapper.enable(DeserializationFeature.ACCEPT_EMPTY_STRING_AS_NULL_OBJECT);
         objectMapper.disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES);
     }
 
+    /**
+     * Parse protobuf schema definition from Glue Schema Registry using protoc compiler
+     * @param glueRegistryName Registry name
+     * @param glueSchemaName Schema name
+     * @return List of MSKField objects containing field information
+     * @throws AthenaConnectorException if schema parsing fails
+     */
+    public List<MSKField> getProtobufFields(String glueRegistryName, String glueSchemaName)
+    {
+        // Get schema from Glue
+        GetSchemaVersionResponse schemaVersionResponse = getSchemaVersionResult(glueRegistryName, glueSchemaName);
+        String schemaDef = schemaVersionResponse.schemaDefinition();
+
+        // Create a unique temp directory using UUID
+        Path protoDir = Paths.get("/tmp", "proto_" + UUID.randomUUID());
+        Path protoFile = protoDir.resolve(PROTO_FILE);
+        Path descFile = protoDir.resolve(DESC_FILE);
+
+        try {
+            Files.createDirectories(protoDir);
+            Files.writeString(protoFile, schemaDef);
+            // Compile using protoc-jar
+            int exitCode = Protoc.runProtoc(new String[]{
+                    "--descriptor_set_out=" + descFile.toAbsolutePath(),
+                    "--proto_path=" + protoDir.toAbsolutePath(),
+                    protoFile.getFileName().toString()
+            });
+
+            if (exitCode != 0 || !Files.exists(descFile)) {
+                throw new AthenaConnectorException(
+                        "Failed to generate descriptor set with protoc",
+                        ErrorDetails.builder()
+                                .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
+                                .build()
+                );
+            }
+
+            try (FileInputStream fis = new FileInputStream(descFile.toFile())) {
+                FileDescriptorSet descriptorSet = FileDescriptorSet.parseFrom(fis);
+
+                if (descriptorSet.getFileList().isEmpty() ||
+                        descriptorSet.getFile(0).getMessageTypeList().isEmpty()) {
+                    throw new AthenaConnectorException(
+                            "No message types found in compiled schema",
+                            ErrorDetails.builder()
+                                    .errorCode(FederationSourceErrorCode.INVALID_RESPONSE_EXCEPTION.toString())
+                                    .build()
+                    );
+                }
+
+                List<MSKField> fields = new ArrayList<>();
+                DescriptorProto messageType = descriptorSet.getFile(0).getMessageType(0);
+                for (FieldDescriptorProto field : messageType.getFieldList()) {
+                    String fieldType = getFieldTypeString(field);
+                    fields.add(new MSKField(field.getName(), fieldType));
+                }
+
+                return fields;
+            }
+        }
+        catch (IOException | InterruptedException e) {
+            throw new AthenaConnectorException(
+                    "Error while handling schema files or protoc execution",
+                    ErrorDetails.builder()
+                            .errorCode(FederationSourceErrorCode.INTERNAL_SERVICE_EXCEPTION.toString())
+                            .build()
+            );
+        }
+        finally {
+            // Clean up temporary files
+            try {
+                Files.deleteIfExists(protoFile);
+                Files.deleteIfExists(descFile);
+                Files.deleteIfExists(protoDir);
+            }
+            catch (IOException e) {
+                logger.warn("Failed to clean up temporary proto directory: {}", protoDir.toAbsolutePath(), e);
+            }
+        }
+    }
+
+    /**
+     * Convert protobuf field type to string representation
+     */
+    private String getFieldTypeString(FieldDescriptorProto field)
+    {
+        String baseType = field.getType().toString().toLowerCase().replace("type_", "");
+        return field.getLabel() == FieldDescriptorProto.Label.LABEL_REPEATED ? 
+               "repeated " + baseType : baseType;
+    }
+
     /**
      * Fetch glue schema content for latest version
      * @param glueRegistryName
@@ -62,6 +175,7 @@ public GetSchemaVersionResponse getSchemaVersionResult(String glueRegistryName,
                 .build()
         );
     }
+
     /**
      * fetch schema file content from glue schema.
      *
@@ -77,14 +191,10 @@ public <T> T getGlueSchema(String glueRegistryName, String glueSchemaName, Class
         GetSchemaVersionResponse result = getSchemaVersionResult(glueRegistryName, glueSchemaName);
         return objectMapper.readValue(result.schemaDefinition(), clazz);
     }
+
     public String getGlueSchemaType(String glueRegistryName, String glueSchemaName)
     {
         GetSchemaVersionResponse result = getSchemaVersionResult(glueRegistryName, glueSchemaName);
         return result.dataFormatAsString();
     }
-    public String getSchemaDef(String glueRegistryName, String glueSchemaName)
-    {
-        GetSchemaVersionResponse result = getSchemaVersionResult(glueRegistryName, glueSchemaName);
-        return result.schemaDefinition();
-    }
 }
diff --git a/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/dto/MSKField.java b/athena-msk/src/main/java/com/amazonaws/athena/connectors/msk/dto/MSKField.java
@@ -40,6 +40,12 @@ public MSKField(String name, String mapping, String type, String formatHint, Obj
         this.value = value;
     }
 
+    public MSKField(String name, String type)
+    {
+        this.name = name;
+        this.type = type;
+    }
+
     public String getName()
     {
         return name;
diff --git a/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandlerTest.java b/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskMetadataHandlerTest.java
@@ -27,6 +27,7 @@
 import com.amazonaws.athena.connector.lambda.domain.predicate.Constraints;
 import com.amazonaws.athena.connector.lambda.metadata.*;
 import com.amazonaws.athena.connector.lambda.security.FederatedIdentity;
+import org.apache.arrow.vector.types.Types;
 import org.apache.kafka.clients.consumer.MockConsumer;
 import org.apache.kafka.clients.consumer.OffsetResetStrategy;
 import org.apache.kafka.common.PartitionInfo;
@@ -57,8 +58,10 @@
 import java.util.Map;
 import java.util.stream.Collectors;
 
+import static com.amazonaws.athena.connectors.msk.AmazonMskConstants.PROTOBUF_DATA_FORMAT;
 import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.assertNull;
+import static org.junit.Assert.fail;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
@@ -181,6 +184,65 @@ public void testDoGetTable() throws Exception {
         assertEquals(1, getTableResponse.getSchema().getFields().size());
     }
 
+    @Test
+    public void testDoGetTableWithProtobufSchema()
+    {
+        String arn = "defaultarn";
+        String schemaName = "defaultschemaname";
+        String schemaVersionId = "defaultversionid";
+        Long latestSchemaVersion = 123L;
+        GetSchemaResponse getSchemaResponse = GetSchemaResponse.builder()
+                .schemaArn(arn)
+                .schemaName(schemaName)
+                .latestSchemaVersion(latestSchemaVersion)
+                .build();
+        GetSchemaVersionResponse getSchemaVersionResponse = GetSchemaVersionResponse.builder()
+                .schemaArn(arn)
+                .schemaVersionId(schemaVersionId)
+                .schemaDefinition("syntax = \"proto3\";\n" +
+                        "package test;\n" +
+                        "message TestMessage {\n" +
+                        "    int32 id = 1;\n" +
+                        "    string name = 2;\n" +
+                        "    double value = 3;\n" +
+                        "}")
+                .dataFormat(DataFormat.PROTOBUF)
+                .build();
+        Mockito.when(awsGlue.getSchema(any(GetSchemaRequest.class))).thenReturn(getSchemaResponse);
+        Mockito.when(awsGlue.getSchemaVersion(any(GetSchemaVersionRequest.class))).thenReturn(getSchemaVersionResponse);
+
+        GetTableRequest getTableRequest = new GetTableRequest(
+            federatedIdentity, 
+            QUERY_ID, 
+            "kafka",
+            new TableName("default", "testmessage"), 
+            Collections.emptyMap()
+        );
+
+        GetTableResponse getTableResponse;
+        try {
+            getTableResponse = amazonMskMetadataHandler.doGetTable(blockAllocator, getTableRequest);
+        }
+        catch (Exception e) {
+            fail("Unexpected exception in doGetTable():" + e.getMessage());
+            return;
+        }
+        
+        // Verify schema field names
+        assertEquals(3, getTableResponse.getSchema().getFields().size());
+        assertEquals("id", getTableResponse.getSchema().getFields().get(0).getName());
+        assertEquals("name", getTableResponse.getSchema().getFields().get(1).getName());
+        assertEquals("value", getTableResponse.getSchema().getFields().get(2).getName());
+
+        // verify schema field types
+        assertEquals(Types.MinorType.INT.getType(), getTableResponse.getSchema().getFields().get(0).getType());
+        assertEquals(Types.MinorType.VARCHAR.getType(), getTableResponse.getSchema().getFields().get(1).getType());
+        assertEquals(Types.MinorType.FLOAT8.getType(), getTableResponse.getSchema().getFields().get(2).getType());
+
+        // Verify data format metadata
+        assertEquals(PROTOBUF_DATA_FORMAT, getTableResponse.getSchema().getCustomMetadata().get("dataFormat"));
+    }
+
     @Test
     public void testDoGetSplits() throws Exception
     {
diff --git a/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java b/athena-msk/src/test/java/com/amazonaws/athena/connectors/msk/AmazonMskRecordHandlerTest.java

Original file line number	Diff line number	Diff line change
`@@ -40,6 +40,12 @@ public MSKField(String name, String mapping, String type, String formatHint, Obj`
`40`	`40`	`this.value = value;`
`41`	`41`	`}`
`42`	`42`
	`43`	`+ public MSKField(String name, String type)`
	`44`	`+ {`
	`45`	`+ this.name = name;`
	`46`	`+ this.type = type;`
	`47`	`+ }`
	`48`	`+`
`43`	`49`	`public String getName()`
`44`	`50`	`{`
`45`	`51`	`return name;`