Skip to content

Commit

Permalink
fix(analytics): look at userEditableInfo to populate cells (datahub-p…
Browse files Browse the repository at this point in the history 
…roject#11909)
  • Loading branch information
@kevinkarchacryl
kevinkarchacryl authored Nov 22, 2024
1 parent 80ef541 commit dac80fb
Show file tree
Hide file tree
Showing 2 changed files with 169 additions and 21 deletions.
82 changes: 61 additions & 21 deletions ...phql-core/src/main/java/com/linkedin/datahub/graphql/analytics/service/AnalyticsUtil.java
View file
Original file line number Diff line number Diff line change
@@ -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;
Expand All @@ -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;
Expand All @@ -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;
Expand Down Expand Up @@ -169,36 +166,79 @@ public static void convertToUserInfoRows(
final Map<Urn, EntityResponse> gmsResponseByUser =
entityClient.batchGetV2(
opContext,
CORP_USER_INFO_ASPECT_NAME,
CORP_USER_ENTITY_NAME,
userUrns,
ImmutableSet.of(CORP_USER_INFO_ASPECT_NAME));
final Map<Urn, CorpUserInfo> urnToCorpUserInfo =
ImmutableSet.of(CORP_USER_INFO_ASPECT_NAME, CORP_USER_EDITABLE_INFO_ASPECT_NAME));
final Stream<Map.Entry<Urn, EntityResponse>> 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<Urn, Pair<CorpUserInfo, CorpUserEditableInfo>> 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<CorpUserInfo, CorpUserEditableInfo> 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);
}
Expand Down
108 changes: 108 additions & 0 deletions datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/utils/AnalyticsUtilTest.java
View file
Original file line number Diff line number Diff line change
@@ -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<Row> 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<Urn, EntityResponse> 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<Cell> 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
}
}

0 comments on commit dac80fb

Please sign in to comment.