yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-12-12 18:47:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(structuredProps) Add validation for allowedTypes and harden API for invalid types (#12578)

Browse Source
This commit is contained in:
Chris Collins 2025-02-07 16:36:46 -05:00 committed by GitHub
parent 45c81233c9
commit 03f1f2d3e2
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 186 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/entitytype/EntityTypeUrnMapper.java
View File

@ -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);
}
}

18
datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/structuredproperty/StructuredPropertyMapper.java
View File

@ -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<EntityResponse, StructuredPropertyEntity> {
@ -141,8 +144,21 @@ public class StructuredPropertyMapper
final TypeQualifier typeQualifier = new TypeQualifier();
List<String> allowedTypes = gmsTypeQualifier.get(ALLOWED_TYPES);
if (allowedTypes != null) {
// filter out correct allowedTypes
List<String> 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;
}

55
metadata-io/src/main/java/com/linkedin/metadata/structuredproperties/validation/PropertyDefinitionValidator.java
View File

@ -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<AspectValidationException> allowedTypesCheck(
MCPItem item, @Nullable StringArrayMap typeQualifier, AspectRetriever aspectRetriever) {
if (typeQualifier == null || typeQualifier.get(ALLOWED_TYPES) == null) {
return Optional.empty();
}
List<String> allowedTypes = typeQualifier.get(ALLOWED_TYPES);
try {
List<Urn> 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<Urn, Boolean> 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();
}
}

112
metadata-io/src/test/java/com/linkedin/metadata/structuredproperties/validators/PropertyDefinitionValidatorTest.java
View File

@ -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<Urn, Boolean> 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<Urn, Boolean> 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;
}
}