diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsUtil.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsUtil.java index a17745948eb823..88ac29b72dee83 100644 --- a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsUtil.java +++ b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsUtil.java @@ -1,19 +1,14 @@ package com.linkedin.datahub.graphql.analytics.service; +import static com.linkedin.metadata.Constants.CORP_USER_EDITABLE_INFO_ASPECT_NAME; +import static com.linkedin.metadata.Constants.CORP_USER_ENTITY_NAME; import static com.linkedin.metadata.Constants.CORP_USER_INFO_ASPECT_NAME; import com.google.common.collect.ImmutableSet; import com.linkedin.common.urn.Urn; import com.linkedin.common.urn.UrnUtils; import com.linkedin.dashboard.DashboardInfo; -import com.linkedin.datahub.graphql.generated.BarSegment; -import com.linkedin.datahub.graphql.generated.Cell; -import com.linkedin.datahub.graphql.generated.Entity; -import com.linkedin.datahub.graphql.generated.EntityProfileParams; -import com.linkedin.datahub.graphql.generated.LinkParams; -import com.linkedin.datahub.graphql.generated.NamedBar; -import com.linkedin.datahub.graphql.generated.Row; -import com.linkedin.datahub.graphql.generated.SearchParams; +import com.linkedin.datahub.graphql.generated.*; import com.linkedin.datahub.graphql.types.common.mappers.UrnToEntityMapper; import com.linkedin.dataplatform.DataPlatformInfo; import com.linkedin.dataset.DatasetProperties; @@ -22,6 +17,7 @@ import com.linkedin.entity.EnvelopedAspect; import com.linkedin.entity.client.EntityClient; import com.linkedin.glossary.GlossaryTermInfo; +import com.linkedin.identity.CorpUserEditableInfo; import com.linkedin.identity.CorpUserInfo; import com.linkedin.metadata.Constants; import com.linkedin.metadata.key.GlossaryTermKey; @@ -35,6 +31,7 @@ import java.util.Set; import java.util.function.Function; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.Nonnull; import javax.annotation.Nullable; import lombok.extern.slf4j.Slf4j; @@ -169,36 +166,79 @@ public static void convertToUserInfoRows( final Map gmsResponseByUser = entityClient.batchGetV2( opContext, - CORP_USER_INFO_ASPECT_NAME, + CORP_USER_ENTITY_NAME, userUrns, - ImmutableSet.of(CORP_USER_INFO_ASPECT_NAME)); - final Map urnToCorpUserInfo = + ImmutableSet.of(CORP_USER_INFO_ASPECT_NAME, CORP_USER_EDITABLE_INFO_ASPECT_NAME)); + final Stream> entityStream = gmsResponseByUser.entrySet().stream() .filter( entry -> entry.getValue() != null - && entry.getValue().getAspects().containsKey(CORP_USER_INFO_ASPECT_NAME)) - .collect( - Collectors.toMap( - Map.Entry::getKey, - entry -> + && (entry.getValue().getAspects().containsKey(CORP_USER_INFO_ASPECT_NAME) + || entry + .getValue() + .getAspects() + .containsKey(CORP_USER_EDITABLE_INFO_ASPECT_NAME))); + final Map> urnToCorpUserInfo = + entityStream.collect( + Collectors.toMap( + Map.Entry::getKey, + entry -> { + CorpUserInfo userInfo = null; + CorpUserEditableInfo editableInfo = null; + try { + userInfo = new CorpUserInfo( entry .getValue() .getAspects() .get(CORP_USER_INFO_ASPECT_NAME) .getValue() - .data()))); + .data()); + } catch (Exception e) { + // nothing to do + } + try { + + editableInfo = + new CorpUserEditableInfo( + entry + .getValue() + .getAspects() + .get(CORP_USER_EDITABLE_INFO_ASPECT_NAME) + .getValue() + .data()); + } catch (Exception e) { + // nothing to do + } + + return Pair.of(userInfo, editableInfo); + })); // Populate a row with the user link, title, and email. rows.forEach( row -> { Urn urn = UrnUtils.getUrn(row.getCells().get(0).getValue()); EntityResponse response = gmsResponseByUser.get(urn); String maybeDisplayName = response != null ? getUserName(response).orElse(null) : null; - String maybeEmail = - urnToCorpUserInfo.containsKey(urn) ? urnToCorpUserInfo.get(urn).getEmail() : null; - String maybeTitle = - urnToCorpUserInfo.containsKey(urn) ? urnToCorpUserInfo.get(urn).getTitle() : null; + String maybeEmail = null; + String maybeTitle = null; + if (urnToCorpUserInfo.containsKey(urn)) { + Pair pair = urnToCorpUserInfo.get(urn); + if (pair.getLeft() != null) { + CorpUserInfo userInfo = pair.getLeft(); + maybeEmail = userInfo.getEmail(); + maybeTitle = userInfo.getTitle(); + } + if (pair.getRight() != null) { + CorpUserEditableInfo userInfo = pair.getRight(); + if (maybeEmail == null) { + maybeEmail = userInfo.getEmail(); + } + if (maybeTitle == null) { + maybeTitle = userInfo.getTitle(); + } + } + } if (maybeDisplayName != null) { row.getCells().get(0).setValue(maybeDisplayName); } diff --git a/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/utils/AnalyticsUtilTest.java b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/utils/AnalyticsUtilTest.java new file mode 100644 index 00000000000000..ab1140d2380315 --- /dev/null +++ b/datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/utils/AnalyticsUtilTest.java @@ -0,0 +1,108 @@ +package com.linkedin.datahub.graphql.utils; + +import static org.mockito.ArgumentMatchers.anySet; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; + +import com.linkedin.common.urn.Urn; +import com.linkedin.common.urn.UrnUtils; +import com.linkedin.data.DataMap; +import com.linkedin.datahub.graphql.analytics.service.AnalyticsUtil; +import com.linkedin.datahub.graphql.generated.Cell; +import com.linkedin.datahub.graphql.generated.Row; +import com.linkedin.entity.EntityResponse; +import com.linkedin.entity.EnvelopedAspect; +import com.linkedin.entity.EnvelopedAspectMap; +import com.linkedin.entity.client.EntityClient; +import com.linkedin.identity.CorpUserEditableInfo; +import com.linkedin.identity.CorpUserInfo; +import com.linkedin.metadata.Constants; +import io.datahubproject.metadata.context.OperationContext; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.mockito.Mock; +import org.mockito.MockitoAnnotations; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + +public class AnalyticsUtilTest { + + @Mock private OperationContext mockOpContext; + + @Mock private EntityClient mockEntityClient; + + final String TEST_CORP_USER_INFO_TEST_USER = "Corp User"; + final String TEST_CORP_USER_EDITABLE_INFO_TEST_TITLE = "Editable Info Title"; + final String TEST_CORP_USER_EDITABLE_INFO_TEST_EMAIL = "Editable Info Email"; + + @BeforeMethod + public void setUp() { + MockitoAnnotations.openMocks(this); + } + + @Test + public void testConvertToUserInfoRows() throws Exception { + List rows = new ArrayList<>(); + rows.add(new Row(null, Arrays.asList(new Cell("urn:li:corpuser:testuser", null, null)))); + + // create a CorpUserInfo with only display name set + CorpUserInfo corpUserInfo = new CorpUserInfo(); + corpUserInfo.setActive(true); + corpUserInfo.setDisplayName(TEST_CORP_USER_INFO_TEST_USER); + + // create an editableInfo with the email and title set + CorpUserEditableInfo corpUserEditableInfo = new CorpUserEditableInfo(); + corpUserEditableInfo.setEmail(TEST_CORP_USER_EDITABLE_INFO_TEST_EMAIL); // Overriding email + corpUserEditableInfo.setTitle(TEST_CORP_USER_EDITABLE_INFO_TEST_TITLE); // Overriding title + + DataMap corpUserInfoDataMap = new DataMap(); + corpUserInfoDataMap.put("name", Constants.CORP_USER_INFO_ASPECT_NAME); + corpUserInfoDataMap.put("type", "VERSIONED"); + corpUserInfoDataMap.put("value", corpUserInfo.data()); + + DataMap corpUserEditableInfoDataMap = new DataMap(); + corpUserEditableInfoDataMap.put("name", Constants.CORP_USER_EDITABLE_INFO_ASPECT_NAME); + corpUserEditableInfoDataMap.put("type", "VERSIONED"); + corpUserEditableInfoDataMap.put("value", corpUserEditableInfo.data()); + + EnvelopedAspect corpUserInfoEnvelopedAspect = new EnvelopedAspect(corpUserInfoDataMap); + EnvelopedAspect corpUserEditableInfoEnvelopedAspect = + new EnvelopedAspect(corpUserEditableInfoDataMap); + + EnvelopedAspectMap aspectMap = new EnvelopedAspectMap(); + aspectMap.put(Constants.CORP_USER_INFO_ASPECT_NAME, corpUserInfoEnvelopedAspect); + aspectMap.put( + Constants.CORP_USER_EDITABLE_INFO_ASPECT_NAME, corpUserEditableInfoEnvelopedAspect); + + EntityResponse entityResponse = new EntityResponse(); + entityResponse.setAspects(aspectMap); + + Map entityResponseMap = new HashMap<>(); + Urn userUrn = UrnUtils.getUrn("urn:li:corpuser:testuser"); + entityResponseMap.put(userUrn, entityResponse); + + // method of the entity client we need to mock to retrieve the response map + when(mockEntityClient.batchGetV2( + eq(mockOpContext), eq(Constants.CORP_USER_ENTITY_NAME), anySet(), anySet())) + .thenReturn(entityResponseMap); + + // function we are testing + AnalyticsUtil.convertToUserInfoRows(mockOpContext, mockEntityClient, rows); + + Row updatedRow = rows.get(0); + List updatedCells = updatedRow.getCells(); + + // asserting that the display user is from CorpUserInfo and email, title are from EditableInfo + assertEquals(updatedCells.get(0).getValue(), TEST_CORP_USER_INFO_TEST_USER); + assertEquals( + updatedCells.get(1).getValue(), + TEST_CORP_USER_EDITABLE_INFO_TEST_TITLE); // Overriding title + assertEquals( + updatedCells.get(2).getValue(), + TEST_CORP_USER_EDITABLE_INFO_TEST_EMAIL); // Overriding email + } +}