fix(protobuf): use DynamicMessage for MESSAGE-type extension defaults

metadata-integration/java/datahub-protobuf/src/main/java/datahub/protobuf/ProtobufUtils.java

@@ -169,12 +169,11 @@ private static void addToRegistry(
     if (fieldDescriptor.getJavaType() != Descriptors.FieldDescriptor.JavaType.MESSAGE) {
       registry.add(fieldDescriptor);
     } else {
-      fileDescriptorProto.getMessageTypeList().stream()
-          .filter(typ -> typ.getName().equals(fieldDescriptor.getMessageType().getName()))
-          .findFirst()
-          .ifPresent(
-              messageType ->
-                  registry.add(fieldDescriptor, messageType.getDefaultInstanceForType()));
+      // For MESSAGE-type extensions, we need to use DynamicMessage to create the default instance
+      // from the Descriptor, as there's no compiled Java class for custom proto extensions
+      registry.add(
+          fieldDescriptor,
+          com.google.protobuf.DynamicMessage.getDefaultInstance(fieldDescriptor.getMessageType()));
       fieldDescriptor.getMessageType().getFields().stream()
           .filter(Descriptors.FieldDescriptor::isExtension)
           .forEach(f -> addToRegistry(fileDescriptorProto, f, registry));

metadata-integration/java/datahub-protobuf/src/test/java/datahub/protobuf/ProtobufUtilsTest.java

@@ -15,6 +15,7 @@ public class ProtobufUtilsTest {
 
   @Test
   public void registryTest() throws IOException, IllegalArgumentException {
+    // Test that primitive and enum-type extensions are properly handled
     byte[] protocBytes = getTestProtoc("extended_protobuf", "messageA").readAllBytes();
     DescriptorProtos.FileDescriptorSet fileSet =
         getTestProtobufFileSet("extended_protobuf", "messageA");
@@ -24,10 +25,7 @@ public void registryTest() throws IOException, IllegalArgumentException {
 
     assertNotEquals(fileSet, fileSetWithRegistry);
 
-    /*
-     *
-     * Without the ExtensionRegistry we get field numbers instead of the names.
-     */
+    // Without the ExtensionRegistry we get field numbers instead of the names.
     ProtobufGraph graph = new ProtobufGraph(fileSet, null);
     assertEquals(
         "[meta.msg.classification_enum]: HighlyConfidential\n"
@@ -46,6 +44,32 @@ public void registryTest() throws IOException, IllegalArgumentException {
         graph.root().messageProto().getOptions().toString());
   }
 
+  @Test
+  public void messageTypeExtensionRegistryTest() throws IOException {
+    // Test that MESSAGE-type extensions (not just primitive/enum) are properly handled
+    byte[] protocBytes = getTestProtoc("extended_protobuf", "messageF").readAllBytes();
+    DescriptorProtos.FileDescriptorSet fileSet =
+        getTestProtobufFileSet("extended_protobuf", "messageF");
+    ExtensionRegistry registry = ProtobufUtils.buildRegistry(fileSet);
+    DescriptorProtos.FileDescriptorSet fileSetWithRegistry =
+        DescriptorProtos.FileDescriptorSet.parseFrom(protocBytes, registry);
+
+    assertNotEquals(fileSet, fileSetWithRegistry);
+
+    // Verify that message-type extension options are properly parsed
+    ProtobufGraph graph = new ProtobufGraph(fileSetWithRegistry, "extended_protobuf.Product");
+    String options = graph.root().messageProto().getOptions().toString();
+
+    // The extension should be properly parsed with field names (not field numbers)
+    assertTrue(
+        options.contains("[meta.nested_msg.metadata]"),
+        "Expected MESSAGE-type extension to be properly registered");
+    assertTrue(
+        options.contains("description: \"Product catalog message\""),
+        "Expected extension field values to be parsed");
+    assertTrue(options.contains("priority: 10"), "Expected extension field values to be parsed");
+  }
+
   @Test
   public void testCollapseLocationCommentsWithUTF8() {
     DescriptorProtos.SourceCodeInfo.Location location =

metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/messageF.proto

@@ -0,0 +1,37 @@
+syntax = "proto3";
+package extended_protobuf;
+
+import "extended_protobuf/meta/meta_nested.proto";
+
+/*
+  Test message-type extensions (not just primitive/enum extensions)
+*/
+message Product {
+  option(meta.nested_msg.metadata) = {
+    description: "Product catalog message"
+    priority: 10
+  };
+
+  string product_id = 1;
+  string product_name = 2;
+}
+
+enum ProductStatus {
+  option(meta.nested_enum.metadata) = {
+    description: "Product status enum"
+    priority: 5
+  };
+
+  PRODUCT_STATUS_UNKNOWN = 0 [(meta.nested_enum_value.metadata) = {
+    description: "Unknown product status"
+    priority: 1
+  }];
+  PRODUCT_STATUS_ACTIVE = 1 [(meta.nested_enum_value.metadata) = {
+    description: "Active product"
+    priority: 10
+  }];
+  PRODUCT_STATUS_INACTIVE = 2 [(meta.nested_enum_value.metadata) = {
+    description: "Inactive product"
+    priority: 0
+  }];
+}

metadata-integration/java/datahub-protobuf/src/test/resources/extended_protobuf/meta/meta_nested.proto

@@ -0,0 +1,41 @@
+syntax = "proto3";
+
+package meta;
+
+import "google/protobuf/descriptor.proto";
+
+/*
+  Message-type metadata for testing nested message extensions.
+  This tests the case where extension values are MESSAGE types, not primitives/enums.
+*/
+message NestedMetadata {
+  string description = 1;
+  int32 priority = 2;
+}
+
+/*
+  Extension for MessageOptions using a MESSAGE type (not primitive/enum)
+*/
+message nested_msg {
+  extend google.protobuf.MessageOptions {
+    NestedMetadata metadata = 7000;
+  }
+}
+
+/*
+  Extension for EnumOptions using a MESSAGE type
+*/
+message nested_enum {
+  extend google.protobuf.EnumOptions {
+    NestedMetadata metadata = 7100;
+  }
+}
+
+/*
+  Extension for EnumValueOptions using a MESSAGE type
+*/
+message nested_enum_value {
+  extend google.protobuf.EnumValueOptions {
+    NestedMetadata metadata = 7200;
+  }
+}