fix(openapi-v3): fix mcp alternative validator & test (#11744)

metadata-service/openapi-servlet/models/src/main/java/io/datahubproject/openapi/v3/models/GenericEntityV3.java

@@ -44,14 +44,6 @@ public static class GenericEntityV3Builder {
 
     public GenericEntityV3 build(
         ObjectMapper objectMapper, @Nonnull Urn urn, Map<String, AspectItem> aspects) {
-      return build(objectMapper, urn, aspects, false);
-    }
-
-    public GenericEntityV3 build(
-        ObjectMapper objectMapper,
-        @Nonnull Urn urn,
-        Map<String, AspectItem> aspects,
-        boolean isAsyncAlternateValidation) {
       Map<String, GenericAspectV3> jsonObjectMap =
           aspects.entrySet().stream()
               .map(
@@ -64,11 +56,6 @@ public GenericEntityV3 build(
                                   .getBytes(StandardCharsets.UTF_8),
                               new TypeReference<>() {});
 
-                      Map<String, Object> aspectValue =
-                          isAsyncAlternateValidation
-                              ? (Map<String, Object>) aspectValueMap.get("value")
-                              : aspectValueMap;
-
                       Map<String, Object> systemMetadata =
                           entry.getValue().getSystemMetadata() != null
                               ? objectMapper.convertValue(
@@ -83,7 +70,7 @@ public GenericEntityV3 build(
                       return Map.entry(
                           aspectName,
                           GenericAspectV3.builder()
-                              .value(aspectValue)
+                              .value(aspectValueMap)
                               .systemMetadata(systemMetadata)
                               .auditStamp(auditStamp)
                               .build());

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/controller/GenericEntitiesController.java

@@ -163,9 +163,9 @@ protected abstract List<E> buildEntityVersionedAspectList(
       throws URISyntaxException;
 
   protected abstract List<E> buildEntityList(
+      OperationContext opContext,
       Collection<IngestResult> ingestResults,
-      boolean withSystemMetadata,
-      boolean isAsyncAlternateValidation);
+      boolean withSystemMetadata);
 
   protected abstract E buildGenericEntity(
       @Nonnull String aspectName,
@@ -515,11 +515,10 @@ public ResponseEntity<List<E>> createEntity(
     List<IngestResult> results = entityService.ingestProposal(opContext, batch, async);
 
     if (!async) {
-      return ResponseEntity.ok(buildEntityList(results, withSystemMetadata, false));
+      return ResponseEntity.ok(buildEntityList(opContext, results, withSystemMetadata));
     } else {
-      boolean isAsyncAlternateValidation = opContext.getValidationContext().isAlternateValidation();
       return ResponseEntity.accepted()
-          .body(buildEntityList(results, withSystemMetadata, isAsyncAlternateValidation));
+          .body(buildEntityList(opContext, results, withSystemMetadata));
     }
   }
 

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java

@@ -163,6 +163,7 @@ protected AspectsBatch toMCPBatch(
           Map.Entry<String, JsonNode> aspect = aspectItr.next();
 
           AspectSpec aspectSpec = lookupAspectSpec(entityUrn, aspect.getKey()).get();
+          JsonNode jsonNodeAspect = aspect.getValue().get("value");
 
           if (opContext.getValidationContext().isAlternateValidation()) {
             ProposedItem.ProposedItemBuilder builder =
@@ -173,7 +174,7 @@ protected AspectsBatch toMCPBatch(
                             .setAspectName(aspect.getKey())
                             .setEntityType(entityUrn.getEntityType())
                             .setChangeType(ChangeType.UPSERT)
-                            .setAspect(GenericRecordUtils.serializeAspect(aspect.getValue()))
+                            .setAspect(GenericRecordUtils.serializeAspect(jsonNodeAspect))
                             .setSystemMetadata(SystemMetadataUtils.createDefaultSystemMetadata()))
                     .auditStamp(AuditStampUtils.createAuditStamp(actor.toUrnStr()))
                     .entitySpec(
@@ -191,7 +192,7 @@ protected AspectsBatch toMCPBatch(
                     .recordTemplate(
                         GenericRecordUtils.deserializeAspect(
                             ByteString.copyString(
-                                objectMapper.writeValueAsString(aspect.getValue().get("value")),
+                                objectMapper.writeValueAsString(jsonNodeAspect),
                                 StandardCharsets.UTF_8),
                             GenericRecordUtils.JSON,
                             aspectSpec));
@@ -282,11 +283,10 @@ private List<GenericEntityV2> toRecordTemplates(
         true);
   }
 
-  @Override
   protected List<GenericEntityV2> buildEntityList(
+      OperationContext opContext,
       Collection<IngestResult> ingestResults,
-      boolean withSystemMetadata,
-      boolean isAsyncAlternateValidation) {
+      boolean withSystemMetadata) {
     List<GenericEntityV2> responseList = new LinkedList<>();
 
     Map<Urn, List<IngestResult>> entityMap =
@@ -305,7 +305,10 @@ protected List<GenericEntityV2> buildEntityList(
       responseList.add(
           GenericEntityV2.builder()
               .urn(urnAspects.getKey().toString())
-              .build(objectMapper, aspectsMap, isAsyncAlternateValidation));
+              .build(
+                  objectMapper,
+                  aspectsMap,
+                  opContext.getValidationContext().isAlternateValidation()));
     }
     return responseList;
   }

metadata-service/openapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java

@@ -285,9 +285,9 @@ private Map<String, AspectItem> toAspectItemMap(
 
   @Override
   protected List<GenericEntityV3> buildEntityList(
+      OperationContext opContext,
       Collection<IngestResult> ingestResults,
-      boolean withSystemMetadata,
-      boolean isAsyncAlternateValidation) {
+      boolean withSystemMetadata) {
     List<GenericEntityV3> responseList = new LinkedList<>();
 
     Map<Urn, List<IngestResult>> entityMap =
@@ -310,8 +310,7 @@ protected List<GenericEntityV3> buildEntityList(
                               .build()))
               .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
       responseList.add(
-          GenericEntityV3.builder()
-              .build(objectMapper, urnAspects.getKey(), aspectsMap, isAsyncAlternateValidation));
+          GenericEntityV3.builder().build(objectMapper, urnAspects.getKey(), aspectsMap));
     }
     return responseList;
   }
@@ -463,6 +462,8 @@ protected AspectsBatch toMCPBatch(
                     aspect.getValue().get("headers"), new TypeReference<>() {});
           }
 
+          JsonNode jsonNodeAspect = aspect.getValue().get("value");
+
           if (opContext.getValidationContext().isAlternateValidation()) {
             ProposedItem.ProposedItemBuilder builder =
                 ProposedItem.builder()
@@ -472,7 +473,7 @@ protected AspectsBatch toMCPBatch(
                             .setAspectName(aspect.getKey())
                             .setEntityType(entityUrn.getEntityType())
                             .setChangeType(ChangeType.UPSERT)
-                            .setAspect(GenericRecordUtils.serializeAspect(aspect.getValue()))
+                            .setAspect(GenericRecordUtils.serializeAspect(jsonNodeAspect))
                             .setHeaders(
                                 headers != null ? new StringMap(headers) : null,
                                 SetMode.IGNORE_NULL)
@@ -495,7 +496,7 @@ protected AspectsBatch toMCPBatch(
                     .recordTemplate(
                         GenericRecordUtils.deserializeAspect(
                             ByteString.copyString(
-                                objectMapper.writeValueAsString(aspect.getValue().get("value")),
+                                objectMapper.writeValueAsString(jsonNodeAspect),
                                 StandardCharsets.UTF_8),
                             GenericRecordUtils.JSON,
                             aspectSpec));

metadata-service/openapi-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java

@@ -1,6 +1,9 @@
 package io.datahubproject.openapi.v3.controller;
 
 import static com.linkedin.metadata.Constants.DATASET_ENTITY_NAME;
+import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME;
+import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME;
+import static com.linkedin.metadata.utils.GenericRecordUtils.JSON;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyInt;
 import static org.mockito.ArgumentMatchers.anyMap;
@@ -9,24 +12,29 @@
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status;
 import static org.testng.Assert.assertNotNull;
+import static org.testng.AssertJUnit.assertEquals;
 
 import com.datahub.authentication.Actor;
 import com.datahub.authentication.ActorType;
 import com.datahub.authentication.Authentication;
 import com.datahub.authentication.AuthenticationContext;
 import com.datahub.authorization.AuthorizationResult;
 import com.datahub.authorization.AuthorizerChain;
+import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.linkedin.common.Status;
 import com.linkedin.common.urn.Urn;
 import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.data.template.RecordTemplate;
 import com.linkedin.entity.Aspect;
 import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.metadata.aspect.batch.AspectsBatch;
 import com.linkedin.metadata.entity.EntityService;
 import com.linkedin.metadata.entity.EntityServiceImpl;
 import com.linkedin.metadata.graph.elastic.ElasticSearchGraphService;
@@ -38,9 +46,13 @@
 import com.linkedin.metadata.search.SearchEntity;
 import com.linkedin.metadata.search.SearchEntityArray;
 import com.linkedin.metadata.search.SearchService;
+import com.linkedin.metadata.utils.GenericRecordUtils;
 import com.linkedin.metadata.utils.SearchUtil;
+import com.linkedin.mxe.GenericAspect;
 import io.datahubproject.metadata.context.OperationContext;
+import io.datahubproject.metadata.context.ValidationContext;
 import io.datahubproject.openapi.config.SpringWebConfig;
+import io.datahubproject.openapi.exception.InvalidUrnException;
 import io.datahubproject.test.metadata.context.TestOperationContexts;
 import java.util.Collections;
 import java.util.List;
@@ -74,6 +86,7 @@ public class EntityControllerTest extends AbstractTestNGSpringContextTests {
   @Autowired private SearchService mockSearchService;
   @Autowired private EntityService<?> mockEntityService;
   @Autowired private EntityRegistry entityRegistry;
+  @Autowired private OperationContext opContext;
 
   @Test
   public void initTest() {
@@ -228,6 +241,79 @@ public void testDeleteEntity() throws Exception {
                         any(), eq(TEST_URN.toString()), eq(aspectName), anyMap(), eq(true)));
   }
 
+  @Test
+  public void testAlternativeMCPValidation() throws InvalidUrnException, JsonProcessingException {
+    final AspectSpec aspectSpec =
+        entityRegistry
+            .getEntitySpec(STRUCTURED_PROPERTY_ENTITY_NAME)
+            .getAspectSpec(STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME);
+
+    // Enable Alternative MCP Validation via mock
+    OperationContext opContextSpy = spy(opContext);
+    ValidationContext mockValidationContext = mock(ValidationContext.class);
+    when(mockValidationContext.isAlternateValidation()).thenReturn(true);
+    when(opContextSpy.getValidationContext()).thenReturn(mockValidationContext);
+
+    final String testBody =
+        "[\n"
+            + "    {\n"
+            + "      \"urn\": \"urn:li:structuredProperty:io.acryl.privacy.retentionTime05\",\n"
+            + "      \"propertyDefinition\": {\n"
+            + "        \"value\": {\n"
+            + "          \"allowedValues\": [\n"
+            + "            {\n"
+            + "              \"value\": {\n"
+            + "                \"string\": \"foo2\"\n"
+            + "              },\n"
+            + "              \"description\": \"test foo2 value\"\n"
+            + "            },\n"
+            + "            {\n"
+            + "              \"value\": {\n"
+            + "                \"string\": \"bar2\"\n"
+            + "              },\n"
+            + "              \"description\": \"test bar2 value\"\n"
+            + "            }\n"
+            + "          ],\n"
+            + "          \"entityTypes\": [\n"
+            + "            \"urn:li:entityType:datahub.dataset\"\n"
+            + "          ],\n"
+            + "          \"qualifiedName\": \"io.acryl.privacy.retentionTime05\",\n"
+            + "          \"displayName\": \"Retention Time 03\",\n"
+            + "          \"cardinality\": \"SINGLE\",\n"
+            + "          \"valueType\": \"urn:li:dataType:datahub.string\"\n"
+            + "        }\n"
+            + "      }\n"
+            + "    }\n"
+            + "]";
+
+    AspectsBatch testAspectsBatch =
+        entityController.toMCPBatch(
+            opContextSpy,
+            testBody,
+            opContext.getSessionActorContext().getAuthentication().getActor());
+
+    GenericAspect aspect =
+        testAspectsBatch.getMCPItems().get(0).getMetadataChangeProposal().getAspect();
+    RecordTemplate propertyDefinition =
+        GenericRecordUtils.deserializeAspect(aspect.getValue(), JSON, aspectSpec);
+    assertEquals(
+        propertyDefinition.data().get("entityTypes"), List.of("urn:li:entityType:datahub.dataset"));
+
+    // test alternative
+    reset(mockValidationContext);
+    when(mockValidationContext.isAlternateValidation()).thenReturn(false);
+    testAspectsBatch =
+        entityController.toMCPBatch(
+            opContextSpy,
+            testBody,
+            opContext.getSessionActorContext().getAuthentication().getActor());
+
+    aspect = testAspectsBatch.getMCPItems().get(0).getMetadataChangeProposal().getAspect();
+    propertyDefinition = GenericRecordUtils.deserializeAspect(aspect.getValue(), JSON, aspectSpec);
+    assertEquals(
+        propertyDefinition.data().get("entityTypes"), List.of("urn:li:entityType:datahub.dataset"));
+  }
+
   @TestConfiguration
   public static class EntityControllerTestConfig {
     @MockBean public EntityServiceImpl entityService;