diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java index 334faf753c..205b6b7d11 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java @@ -111,4 +111,8 @@ public class EntityTypeUrnMapper { } return ENTITY_NAME_TO_ENTITY_TYPE_URN.get(name); } + + public static boolean isValidEntityType(String entityTypeUrn) { + return ENTITY_TYPE_URN_TO_NAME.containsKey(entityTypeUrn); + } } diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/structuredproperty/StructuredPropertyMapper.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/structuredproperty/StructuredPropertyMapper.java index 5dc73d9ad0..98d48d7aa1 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/structuredproperty/StructuredPropertyMapper.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/structuredproperty/StructuredPropertyMapper.java @@ -20,6 +20,7 @@ import com.linkedin.datahub.graphql.generated.StructuredPropertyEntity; import com.linkedin.datahub.graphql.generated.StructuredPropertySettings; import com.linkedin.datahub.graphql.generated.TypeQualifier; import com.linkedin.datahub.graphql.types.common.mappers.util.MappingHelper; +import com.linkedin.datahub.graphql.types.entitytype.EntityTypeUrnMapper; import com.linkedin.datahub.graphql.types.mappers.MapperUtils; import com.linkedin.datahub.graphql.types.mappers.ModelMapper; import com.linkedin.entity.EntityResponse; @@ -30,7 +31,9 @@ import java.util.List; import java.util.stream.Collectors; import javax.annotation.Nonnull; import javax.annotation.Nullable; +import lombok.extern.slf4j.Slf4j; +@Slf4j public class StructuredPropertyMapper implements ModelMapper { @@ -141,8 +144,21 @@ public class StructuredPropertyMapper final TypeQualifier typeQualifier = new TypeQualifier(); List allowedTypes = gmsTypeQualifier.get(ALLOWED_TYPES); if (allowedTypes != null) { + // filter out correct allowedTypes + List validAllowedTypes = + allowedTypes.stream() + .filter(EntityTypeUrnMapper::isValidEntityType) + .collect(Collectors.toList()); + if (validAllowedTypes.size() != allowedTypes.size()) { + log.error( + String.format( + "Property has invalid allowed types set. Current list of allowed types: %s", + allowedTypes)); + } typeQualifier.setAllowedTypes( - allowedTypes.stream().map(this::createEntityTypeEntity).collect(Collectors.toList())); + validAllowedTypes.stream() + .map(this::createEntityTypeEntity) + .collect(Collectors.toList())); } return typeQualifier; } diff --git a/metadata-io/src/main/java/com/linkedin/metadata/structuredproperties/validation/PropertyDefinitionValidator.java b/metadata-io/src/main/java/com/linkedin/metadata/structuredproperties/validation/PropertyDefinitionValidator.java index 6e047c12da..18553d6930 100644 --- a/metadata-io/src/main/java/com/linkedin/metadata/structuredproperties/validation/PropertyDefinitionValidator.java +++ b/metadata-io/src/main/java/com/linkedin/metadata/structuredproperties/validation/PropertyDefinitionValidator.java @@ -1,13 +1,14 @@ package com.linkedin.metadata.structuredproperties.validation; -import static com.linkedin.metadata.Constants.STATUS_ASPECT_NAME; -import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME; +import static com.linkedin.metadata.Constants.*; import static com.linkedin.structured.PropertyCardinality.*; import com.google.common.collect.ImmutableSet; import com.linkedin.common.Status; import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.template.GetMode; +import com.linkedin.data.template.StringArrayMap; import com.linkedin.entity.Aspect; import com.linkedin.metadata.Constants; import com.linkedin.metadata.aspect.AspectRetriever; @@ -24,6 +25,8 @@ import com.linkedin.structured.PropertyValue; import com.linkedin.structured.StructuredPropertyDefinition; import java.util.Collection; import java.util.Collections; +import java.util.HashSet; +import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Set; @@ -42,6 +45,8 @@ import lombok.experimental.Accessors; public class PropertyDefinitionValidator extends AspectPayloadValidator { private AspectPluginConfig config; + private static String ALLOWED_TYPES = "allowedTypes"; + /** * Prevent deletion of the definition or key aspect (only soft delete) * @@ -92,6 +97,9 @@ public class PropertyDefinitionValidator extends AspectPayloadValidator { urnIdCheck(item).ifPresent(exceptions::addException); qualifiedNameCheck(item, newDefinition.getQualifiedName()) .ifPresent(exceptions::addException); + allowedTypesCheck( + item, newDefinition.getTypeQualifier(), retrieverContext.getAspectRetriever()) + .ifPresent(exceptions::addException); if (item.getPreviousSystemAspect() != null) { @@ -211,4 +219,47 @@ public class PropertyDefinitionValidator extends AspectPayloadValidator { } return Optional.empty(); } + + private static Optional allowedTypesCheck( + MCPItem item, @Nullable StringArrayMap typeQualifier, AspectRetriever aspectRetriever) { + if (typeQualifier == null || typeQualifier.get(ALLOWED_TYPES) == null) { + return Optional.empty(); + } + List allowedTypes = typeQualifier.get(ALLOWED_TYPES); + try { + List allowedTypesUrns = + allowedTypes.stream().map(UrnUtils::getUrn).collect(Collectors.toList()); + + // ensure all types are entityTypes + if (allowedTypesUrns.stream() + .anyMatch(t -> !t.getEntityType().equals(ENTITY_TYPE_ENTITY_NAME))) { + return Optional.of( + AspectValidationException.forItem( + item, + String.format( + "Provided allowedType that is not an entityType entity. List of allowedTypes: %s", + allowedTypes))); + } + + // ensure all types exist as entities + Map existsMap = aspectRetriever.entityExists(new HashSet<>(allowedTypesUrns)); + if (existsMap.containsValue(false)) { + return Optional.of( + AspectValidationException.forItem( + item, + String.format( + "Provided allowedType that does not exist. List of allowedTypes: %s", + allowedTypes))); + } + } catch (Exception e) { + return Optional.of( + AspectValidationException.forItem( + item, + String.format( + "Issue resolving allowedTypes inside of typeQualifier. These must be entity type urns. List of allowedTypes: %s", + allowedTypes))); + } + + return Optional.empty(); + } } diff --git a/metadata-io/src/test/java/com/linkedin/metadata/structuredproperties/validators/PropertyDefinitionValidatorTest.java b/metadata-io/src/test/java/com/linkedin/metadata/structuredproperties/validators/PropertyDefinitionValidatorTest.java index 18949f0566..681b31a2db 100644 --- a/metadata-io/src/test/java/com/linkedin/metadata/structuredproperties/validators/PropertyDefinitionValidatorTest.java +++ b/metadata-io/src/test/java/com/linkedin/metadata/structuredproperties/validators/PropertyDefinitionValidatorTest.java @@ -1,5 +1,6 @@ package com.linkedin.metadata.structuredproperties.validators; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; @@ -8,6 +9,8 @@ import com.linkedin.common.UrnArray; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.data.template.SetMode; +import com.linkedin.data.template.StringArray; +import com.linkedin.data.template.StringArrayMap; import com.linkedin.metadata.aspect.AspectRetriever; import com.linkedin.metadata.aspect.GraphRetriever; import com.linkedin.metadata.aspect.RetrieverContext; @@ -22,6 +25,7 @@ import com.linkedin.structured.StructuredPropertyDefinition; import com.linkedin.test.metadata.aspect.TestEntityRegistry; import com.linkedin.test.metadata.aspect.batch.TestMCP; import java.net.URISyntaxException; +import java.util.HashMap; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -37,6 +41,8 @@ public class PropertyDefinitionValidatorTest { testPropertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); AspectRetriever mockAspectRetriever = mock(AspectRetriever.class); when(mockAspectRetriever.getEntityRegistry()).thenReturn(entityRegistry); + HashMap map = new HashMap<>(); + when(mockAspectRetriever.entityExists(any())).thenReturn(map); GraphRetriever mockGraphRetriever = mock(GraphRetriever.class); mockRetrieverContext = mock(RetrieverContext.class); when(mockRetrieverContext.getAspectRetriever()).thenReturn(mockAspectRetriever); @@ -433,4 +439,110 @@ public class PropertyDefinitionValidatorTest { .count(), 1); } + + @Test + public void testValidAllowedTypes() + throws URISyntaxException, CloneNotSupportedException, AspectValidationException { + Urn propertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); + StructuredPropertyDefinition newProperty = createValidPropertyDefinition(); + StringArrayMap typeQualifier = new StringArrayMap(); + typeQualifier.put("allowedTypes", new StringArray("urn:li:entityType:datahub.dataset")); + newProperty.setTypeQualifier(typeQualifier); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(propertyUrn, null, newProperty, entityRegistry), + mockRetrieverContext) + .count(), + 0); + } + + @Test + public void testInvalidUrnsInAllowedTypes() + throws URISyntaxException, CloneNotSupportedException, AspectValidationException { + Urn propertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); + StructuredPropertyDefinition newProperty = createValidPropertyDefinition(); + StringArrayMap typeQualifier = new StringArrayMap(); + // invalid urn here + typeQualifier.put("allowedTypes", new StringArray("invalidUrn")); + newProperty.setTypeQualifier(typeQualifier); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(propertyUrn, null, newProperty, entityRegistry), + mockRetrieverContext) + .count(), + 1); + } + + @Test + public void testNotEntityTypeInAllowedTypes() + throws URISyntaxException, CloneNotSupportedException, AspectValidationException { + Urn propertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); + StructuredPropertyDefinition newProperty = createValidPropertyDefinition(); + StringArrayMap typeQualifier = new StringArrayMap(); + // urn that is not an entityType + typeQualifier.put("allowedTypes", new StringArray("urn:li:dataPlatform:snowflake")); + newProperty.setTypeQualifier(typeQualifier); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(propertyUrn, null, newProperty, entityRegistry), + mockRetrieverContext) + .count(), + 1); + } + + @Test + public void testEntityTypeDoesNotExistInAllowedTypes() + throws URISyntaxException, CloneNotSupportedException, AspectValidationException { + AspectRetriever mockAspectRetriever = mock(AspectRetriever.class); + when(mockAspectRetriever.getEntityRegistry()).thenReturn(entityRegistry); + HashMap map = new HashMap<>(); + map.put(UrnUtils.getUrn("urn:li:entityType:datahub.fakeEntity"), false); + when(mockAspectRetriever.entityExists(any())).thenReturn(map); + GraphRetriever mockGraphRetriever = mock(GraphRetriever.class); + RetrieverContext retrieverContext = mock(RetrieverContext.class); + when(retrieverContext.getAspectRetriever()).thenReturn(mockAspectRetriever); + when(retrieverContext.getGraphRetriever()).thenReturn(mockGraphRetriever); + + Urn propertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); + StructuredPropertyDefinition newProperty = createValidPropertyDefinition(); + StringArrayMap typeQualifier = new StringArrayMap(); + // urn that doesn't exist + typeQualifier.put("allowedTypes", new StringArray("urn:li:entityType:datahub.fakeEntity")); + newProperty.setTypeQualifier(typeQualifier); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(propertyUrn, null, newProperty, entityRegistry), retrieverContext) + .count(), + 1); + } + + @Test + public void testAllowedTypesMixOfValidAndInvalid() + throws URISyntaxException, CloneNotSupportedException, AspectValidationException { + Urn propertyUrn = UrnUtils.getUrn("urn:li:structuredProperty:foo.bar"); + StructuredPropertyDefinition newProperty = createValidPropertyDefinition(); + StringArrayMap typeQualifier = new StringArrayMap(); + // urn that is not an entityType + typeQualifier.put( + "allowedTypes", + new StringArray("urn:li:entityType:datahub.dataset", "urn:li:dataPlatform:snowflake")); + newProperty.setTypeQualifier(typeQualifier); + assertEquals( + PropertyDefinitionValidator.validateDefinitionUpserts( + TestMCP.ofOneMCP(propertyUrn, null, newProperty, entityRegistry), + mockRetrieverContext) + .count(), + 1); + } + + private StructuredPropertyDefinition createValidPropertyDefinition() throws URISyntaxException { + StructuredPropertyDefinition newProperty = new StructuredPropertyDefinition(); + newProperty.setEntityTypes( + new UrnArray(Urn.createFromString("urn:li:entityType:datahub.dataset"))); + newProperty.setDisplayName("oldProp"); + newProperty.setQualifiedName("foo.bar"); + newProperty.setCardinality(PropertyCardinality.MULTIPLE); + newProperty.setValueType(Urn.createFromString("urn:li:dataType:datahub.urn")); + return newProperty; + } }