fix(structuredProps) Fix adding new allowed types in updateStructured…

…Prop endpoint (datahub-project#11424)
hsheth2 · Sep 19, 2024 · c353332 · c353332
1 parent 06e0deb
commit c353332
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/...edin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolver.java b/...edin/datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolver.java
@@ -1,6 +1,7 @@
 package com.linkedin.datahub.graphql.resolvers.structuredproperties;
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME;
 import static com.linkedin.metadata.Constants.STRUCTURED_PROPERTY_ENTITY_NAME;
 
 import com.linkedin.common.urn.Urn;
@@ -21,17 +22,21 @@
 import com.linkedin.structured.PrimitivePropertyValue;
 import com.linkedin.structured.PropertyCardinality;
 import com.linkedin.structured.PropertyValue;
+import com.linkedin.structured.StructuredPropertyDefinition;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 
 public class UpdateStructuredPropertyResolver
     implements DataFetcher<CompletableFuture<StructuredPropertyEntity>> {
 
   private final EntityClient _entityClient;
 
+  private static final String ALLOWED_TYPES = "allowedTypes";
+
   public UpdateStructuredPropertyResolver(@Nonnull final EntityClient entityClient) {
     _entityClient = Objects.requireNonNull(entityClient, "entityClient must not be null");
   }
@@ -52,6 +57,8 @@ public CompletableFuture<StructuredPropertyEntity> get(final DataFetchingEnviron
                   "Unable to update structured property. Please contact your admin.");
             }
             final Urn propertyUrn = UrnUtils.getUrn(input.getUrn());
+            StructuredPropertyDefinition existingDefinition =
+                getExistingStructuredProperty(context, propertyUrn);
             StructuredPropertyDefinitionPatchBuilder builder =
                 new StructuredPropertyDefinitionPatchBuilder().urn(propertyUrn);
 
@@ -65,7 +72,7 @@ public CompletableFuture<StructuredPropertyEntity> get(final DataFetchingEnviron
               builder.setImmutable(input.getImmutable());
             }
             if (input.getTypeQualifier() != null) {
-              buildTypeQualifier(input, builder);
+              buildTypeQualifier(input, builder, existingDefinition);
             }
             if (input.getNewAllowedValues() != null) {
               buildAllowedValues(input, builder);
@@ -97,10 +104,16 @@ public CompletableFuture<StructuredPropertyEntity> get(final DataFetchingEnviron
 
   private void buildTypeQualifier(
       @Nonnull final UpdateStructuredPropertyInput input,
-      @Nonnull final StructuredPropertyDefinitionPatchBuilder builder) {
+      @Nonnull final StructuredPropertyDefinitionPatchBuilder builder,
+      @Nullable final StructuredPropertyDefinition existingDefinition) {
     if (input.getTypeQualifier().getNewAllowedTypes() != null) {
       final StringArrayMap typeQualifier = new StringArrayMap();
       StringArray allowedTypes = new StringArray();
+      if (existingDefinition != null
+          && existingDefinition.getTypeQualifier() != null
+          && existingDefinition.getTypeQualifier().get(ALLOWED_TYPES) != null) {
+        allowedTypes.addAll(existingDefinition.getTypeQualifier().get(ALLOWED_TYPES));
+      }
       allowedTypes.addAll(input.getTypeQualifier().getNewAllowedTypes());
       typeQualifier.put("allowedTypes", allowedTypes);
       builder.setTypeQualifier(typeQualifier);
@@ -127,4 +140,18 @@ private void buildAllowedValues(
               builder.addAllowedValue(value);
             });
   }
+
+  private StructuredPropertyDefinition getExistingStructuredProperty(
+      @Nonnull final QueryContext context, @Nonnull final Urn propertyUrn) throws Exception {
+    EntityResponse response =
+        _entityClient.getV2(
+            context.getOperationContext(), STRUCTURED_PROPERTY_ENTITY_NAME, propertyUrn, null);
+
+    if (response != null
+        && response.getAspects().containsKey(STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME)) {
+      return new StructuredPropertyDefinition(
+          response.getAspects().get(STRUCTURED_PROPERTY_DEFINITION_ASPECT_NAME).getValue().data());
+    }
+    return null;
+  }
 }
diff --git a/.../datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolverTest.java b/.../datahub/graphql/resolvers/structuredproperties/UpdateStructuredPropertyResolverTest.java
@@ -89,8 +89,8 @@ public void testGetFailure() throws Exception {
 
     assertThrows(CompletionException.class, () -> resolver.get(mockEnv).join());
 
-    // Validate that ingest was called, but that caused a failure
-    Mockito.verify(mockEntityClient, Mockito.times(1))
+    // Validate that ingest was not called since there was a get failure before ingesting
+    Mockito.verify(mockEntityClient, Mockito.times(0))
         .ingestProposal(any(), any(MetadataChangeProposal.class), Mockito.eq(false));
   }