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

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

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

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.*;
-import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME;
 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();
+  }
 }

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<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;
+  }
 }