Merge remote-tracking branch 'acryl/jj--add-models-for-persona-captur…

…e' into jj--add-models-for-persona-capture
datahub-project · Jan 16, 2024 · d8a9dd1 · d8a9dd1
2 parents 63126bd + ed91fce
commit d8a9dd1
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 2 deletions.
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/Constants.java
@@ -1,5 +1,8 @@
 package com.linkedin.datahub.graphql;
 
+import com.google.common.collect.ImmutableSet;
+import java.util.Set;
+
 /** Constants relating to GraphQL type system & execution. */
 public class Constants {
 
@@ -24,4 +27,11 @@ private Constants() {}
   public static final String VERSION_STAMP_FIELD_NAME = "versionStamp";
 
   public static final String ENTITY_FILTER_NAME = "_entityType";
+
+  public static final Set<String> DEFAULT_PERSONA_URNS =
+      ImmutableSet.of(
+          "urn:li:dataHubPersona:technicalUser",
+          "urn:li:dataHubPersona:businessUser",
+          "urn:li:dataHubPersona:dataLeader",
+          "urn:li:dataHubPersona:dataSteward");
 }
diff --git a/...-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java b/...-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/corpuser/CorpUserType.java
@@ -1,5 +1,6 @@
 package com.linkedin.datahub.graphql.types.corpuser;
 
+import static com.linkedin.datahub.graphql.Constants.DEFAULT_PERSONA_URNS;
 import static com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils.*;
 import static com.linkedin.metadata.Constants.*;
 
@@ -15,6 +16,8 @@
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
+import com.linkedin.datahub.graphql.exception.DataHubGraphQLErrorCode;
+import com.linkedin.datahub.graphql.exception.DataHubGraphQLException;
 import com.linkedin.datahub.graphql.featureflags.FeatureFlags;
 import com.linkedin.datahub.graphql.generated.AutoCompleteResults;
 import com.linkedin.datahub.graphql.generated.CorpUser;
@@ -254,8 +257,13 @@ private RecordTemplate mapCorpUserEditableInfo(
               input.getPlatformUrns().stream().map(UrnUtils::getUrn).collect(Collectors.toList())));
     }
     if (input.getPersonaUrn() != null) {
-      // TODO: Verify that the persona exists before accepting.
-      result.setPersona(UrnUtils.getUrn(input.getPersonaUrn()));
+      if (DEFAULT_PERSONA_URNS.contains(input.getPersonaUrn())) {
+        result.setPersona(UrnUtils.getUrn(input.getPersonaUrn()));
+      } else {
+        throw new DataHubGraphQLException(
+            String.format("Provided persona urn %s does not exist", input.getPersonaUrn()),
+            DataHubGraphQLErrorCode.NOT_FOUND);
+      }
     }
     return result;
   }