Skip to content

Commit

Permalink
fix(openapi-v3): fix mcp alternative validator & test
Browse files Browse the repository at this point in the history
  • Loading branch information
@david-leifker
david-leifker committed Oct 29, 2024
1 parent 25a4eef commit 025aab4
Show file tree
Hide file tree
Showing 3 changed files with 77 additions and 4 deletions.
5 changes: 3 additions & 2 deletions ...enapi-servlet/src/main/java/io/datahubproject/openapi/v2/controller/EntityController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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 =
Expand All @@ -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(
Expand All @@ -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));
Expand Down
6 changes: 4 additions & 2 deletions ...enapi-servlet/src/main/java/io/datahubproject/openapi/v3/controller/EntityController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -463,6 +463,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()
Expand All @@ -472,7 +474,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)
Expand All @@ -495,7 +497,7 @@ protected AspectsBatch toMCPBatch(
.recordTemplate(
GenericRecordUtils.deserializeAspect(
ByteString.copyString(
objectMapper.writeValueAsString(aspect.getValue().get("value")),
objectMapper.writeValueAsString(jsonNodeAspect),
StandardCharsets.UTF_8),
GenericRecordUtils.JSON,
aspectSpec));
Expand Down
70 changes: 70 additions & 0 deletions ...i-servlet/src/test/java/io/datahubproject/openapi/v3/controller/EntityControllerTest.java
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -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() {
Expand Down Expand Up @@ -228,6 +241,63 @@ public void testDeleteEntity() throws Exception {
any(), eq(TEST_URN.toString()), eq(aspectName), anyMap(), eq(true)));
}

@Test
public void testAlternativeMCPValidation() throws InvalidUrnException, JsonProcessingException {
// 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());
AspectSpec aspectSpec =
entityRegistry
.getEntitySpec(STRUCTURED_PROPERTY_ENTITY_NAME)
.getAspectSpec(STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME);
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"));
}

@TestConfiguration
public static class EntityControllerTestConfig {
@MockBean public EntityServiceImpl entityService;
Expand Down

0 comments on commit 025aab4

Please sign in to comment.