feat(policy): enable support for 10k+ policies

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/policy/ListPoliciesResolver.java

@@ -40,23 +40,15 @@ public CompletableFuture<ListPoliciesResult> get(final DataFetchingEnvironment e
  final Integer count = input.getCount() == null ? DEFAULT_COUNT : input.getCount();
  final String query = input.getQuery() == null ? DEFAULT_QUERY : input.getQuery();
 
- return CompletableFuture.supplyAsync(() -> {
- try {
- // First, get all policy Urns.
- final PolicyFetcher.PolicyFetchResult policyFetchResult =
- _policyFetcher.fetchPolicies(start, count, query, context.getAuthentication());
-
- // Now that we have entities we can bind this to a result.
- final ListPoliciesResult result = new ListPoliciesResult();
- result.setStart(start);
- result.setCount(count);
- result.setTotal(policyFetchResult.getTotal());
- result.setPolicies(mapEntities(policyFetchResult.getPolicies()));
- return result;
- } catch (Exception e) {
- throw new RuntimeException("Failed to list policies", e);
- }
- });
+ return _policyFetcher.fetchPolicies(start, query, count, context.getAuthentication())
+ .thenApply(policyFetchResult -> {
+ final ListPoliciesResult result = new ListPoliciesResult();
+ result.setStart(start);
+ result.setCount(count);
+ result.setTotal(policyFetchResult.getTotal());
+ result.setPolicies(mapEntities(policyFetchResult.getPolicies()));
+ return result;
+ });
  }
  throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
  }

metadata-io/src/main/java/com/linkedin/metadata/client/JavaEntityClient.java

@@ -381,7 +381,7 @@ public SearchResult searchAcrossEntities(
  @Nonnull
  @Override
  public ScrollResult scrollAcrossEntities(@Nonnull List<String> entities, @Nonnull String input,
- @Nullable Filter filter, @Nullable String scrollId, @Nonnull String keepAlive, int count,
+ @Nullable Filter filter, @Nullable String scrollId, @Nullable String keepAlive, int count,
  @Nullable SearchFlags searchFlags, @Nonnull Authentication authentication)
  throws RemoteInvocationException {
  final SearchFlags finalFlags = searchFlags != null ? searchFlags : new SearchFlags().setFulltext(true);

metadata-io/src/main/java/com/linkedin/metadata/search/SearchService.java

@@ -147,15 +147,23 @@ public SearchResult searchAcrossEntities(@Nonnull List<String> entities, @Nonnul
  return result;
  }
 
+ /**
+ * If no entities are provided, fallback to the list of non-empty entities
+ * @param inputEntities the requested entities
+ * @return some entities to search
+ */
  private List<String> getEntitiesToSearch(@Nonnull List<String> inputEntities) {
  List<String> nonEmptyEntities;
  List<String> lowercaseEntities = inputEntities.stream().map(String::toLowerCase).collect(Collectors.toList());
- try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getNonEmptyEntities").time()) {
- nonEmptyEntities = _entityDocCountCache.getNonEmptyEntities();
- }
- if (!inputEntities.isEmpty()) {
- nonEmptyEntities = nonEmptyEntities.stream().filter(lowercaseEntities::contains).collect(Collectors.toList());
+
+ if (lowercaseEntities.isEmpty()) {
+ try (Timer.Context ignored = MetricUtils.timer(this.getClass(), "getNonEmptyEntities").time()) {
+ nonEmptyEntities = _entityDocCountCache.getNonEmptyEntities();
+ }
+ } else {
+ nonEmptyEntities = lowercaseEntities;
  }
+
  return nonEmptyEntities;
  }
 

metadata-service/auth-impl/src/main/java/com/datahub/authorization/DataHubAuthorizer.java

@@ -72,11 +72,13 @@ public DataHubAuthorizer(
  final EntityClient entityClient,
  final int delayIntervalSeconds,
  final int refreshIntervalSeconds,
- final AuthorizationMode mode) {
+ final AuthorizationMode mode,
+ final int policyFetchSize) {
  _systemAuthentication = Objects.requireNonNull(systemAuthentication);
  _mode = Objects.requireNonNull(mode);
  _policyEngine = new PolicyEngine(systemAuthentication, Objects.requireNonNull(entityClient));
- _policyRefreshRunnable = new PolicyRefreshRunnable(systemAuthentication, new PolicyFetcher(entityClient), _policyCache, readWriteLock.writeLock());
+ _policyRefreshRunnable = new PolicyRefreshRunnable(systemAuthentication, new PolicyFetcher(entityClient), _policyCache,
+ readWriteLock.writeLock(), policyFetchSize);
  _refreshExecutorService.scheduleAtFixedRate(_policyRefreshRunnable, delayIntervalSeconds, refreshIntervalSeconds, TimeUnit.SECONDS);
  }
 
@@ -244,29 +246,28 @@ static class PolicyRefreshRunnable implements Runnable {
  private final PolicyFetcher _policyFetcher;
  private final Map<String, List<DataHubPolicyInfo>> _policyCache;
  private final Lock writeLock;
+ private final int count;
 
  @Override
  public void run() {
  try {
  // Populate new cache and swap.
  Map<String, List<DataHubPolicyInfo>> newCache = new HashMap<>();
+ Integer total = null;
+ String scrollId = null;
 
- int start = 0;
- int count = 30;
- int total = 30;
-
- while (start < total) {
+ while (total == null || scrollId != null) {
  try {
  final PolicyFetcher.PolicyFetchResult
- policyFetchResult = _policyFetcher.fetchPolicies(start, count, _systemAuthentication);
+  policyFetchResult = _policyFetcher.fetchPolicies(count, scrollId, _systemAuthentication);
 
  addPoliciesToCache(newCache, policyFetchResult.getPolicies());
 
  total = policyFetchResult.getTotal();
- start = start + count;
+ scrollId = policyFetchResult.getScrollId();
  } catch (Exception e) {
  log.error(
- "Failed to retrieve policy urns! Skipping updating policy cache until next refresh. start: {}, count: {}", start, count, e);
+  "Failed to retrieve policy urns! Skipping updating policy cache until next refresh. count: {}, scrollId: {}", count, scrollId, e);
  return;
  }
  }

metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java

@@ -8,8 +8,8 @@
 import com.linkedin.metadata.query.SearchFlags;
 import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.query.filter.SortOrder;
+import com.linkedin.metadata.search.ScrollResult;
 import com.linkedin.metadata.search.SearchEntity;
-import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.policy.DataHubPolicyInfo;
 import com.linkedin.r2.RemoteInvocationException;
 import java.net.URISyntaxException;
@@ -18,11 +18,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import lombok.RequiredArgsConstructor;
 import lombok.Value;
 import lombok.extern.slf4j.Slf4j;
 
+import javax.annotation.Nullable;
+
 import static com.linkedin.metadata.Constants.DATAHUB_POLICY_INFO_ASPECT_NAME;
 import static com.linkedin.metadata.Constants.POLICY_ENTITY_NAME;
 
@@ -38,22 +41,53 @@ public class PolicyFetcher {
  private static final SortCriterion POLICY_SORT_CRITERION =
  new SortCriterion().setField("lastUpdatedTimestamp").setOrder(SortOrder.DESCENDING);
 
- public PolicyFetchResult fetchPolicies(int start, int count, Authentication authentication)
- throws RemoteInvocationException, URISyntaxException {
- return fetchPolicies(start, count, "", authentication);
+ /**
+ * This is to provide a scroll implementation using the start/count api. It is not efficient
+ * and the scroll native functions should be used instead. This does fix a failure to fetch
+ * policies when deep pagination happens where there are >10k policies.
+ * Exists primarily to prevent breaking change to the graphql api.
+ */
+ @Deprecated
+ public CompletableFuture<PolicyFetchResult> fetchPolicies(int start, String query, int count, Authentication authentication) {
+ return CompletableFuture.supplyAsync(() -> {
+ try {
+ PolicyFetchResult result = PolicyFetchResult.EMPTY;
+ String scrollId = "";
+ int fetchedResults = 0;
+
+ while (PolicyFetchResult.EMPTY.equals(result) && scrollId != null) {
+ PolicyFetchResult tmpResult = fetchPolicies(query, count, scrollId.isEmpty() ? null : scrollId, authentication);
+ fetchedResults += tmpResult.getPolicies().size();
+ scrollId = tmpResult.getScrollId();
+ if (fetchedResults > start) {
+ result = tmpResult;
+ }
+ }
+
+ return result;
+ } catch (Exception e) {
+ throw new RuntimeException("Failed to list policies", e);
+ }
+ });
  }
 
- public PolicyFetchResult fetchPolicies(int start, int count, String query, Authentication authentication)
+ public PolicyFetchResult fetchPolicies(int count, @Nullable String scrollId, Authentication authentication)
+ throws RemoteInvocationException, URISyntaxException {
+ return fetchPolicies("", count, scrollId, authentication);
+ }
+
+ public PolicyFetchResult fetchPolicies(String query, int count, @Nullable String scrollId, Authentication authentication)
  throws RemoteInvocationException, URISyntaxException {
- log.debug(String.format("Batch fetching policies. start: %s, count: %s ", start, count));
- // First fetch all policy urns from start - start + count
- SearchResult result =
- _entityClient.search(POLICY_ENTITY_NAME, query, null, POLICY_SORT_CRITERION, start, count, authentication,
- new SearchFlags().setFulltext(true));
+ log.debug(String.format("Batch fetching policies. count: %s, scroll: %s", count, scrollId));
+
+ // First fetch all policy urns
+ ScrollResult result = _entityClient.scrollAcrossEntities(List.of(POLICY_ENTITY_NAME), query, null, scrollId,
+ null, count, new SearchFlags().setSkipCache(true).setSkipAggregates(true)
+ .setSkipHighlighting(true).setFulltext(true), authentication);
  List<Urn> policyUrns = result.getEntities().stream().map(SearchEntity::getEntity).collect(Collectors.toList());
 
  if (policyUrns.isEmpty()) {
- return new PolicyFetchResult(Collections.emptyList(), 0);
+ return PolicyFetchResult.EMPTY;
  }
 
  // Fetch DataHubPolicyInfo aspects for each urn
@@ -64,7 +98,7 @@ public PolicyFetchResult fetchPolicies(int start, int count, String query, Authe
  .filter(Objects::nonNull)
  .map(this::extractPolicy)
  .filter(Objects::nonNull)
- .collect(Collectors.toList()), result.getNumEntities());
+ .collect(Collectors.toList()), result.getNumEntities(), result.getScrollId());
  }
 
  private Policy extractPolicy(EntityResponse entityResponse) {
@@ -82,6 +116,10 @@ private Policy extractPolicy(EntityResponse entityResponse) {
  public static class PolicyFetchResult {
  List<Policy> policies;
  int total;
+ @Nullable
+ String scrollId;
+
+ public static final PolicyFetchResult EMPTY = new PolicyFetchResult(Collections.emptyList(), 0, null);
  }
 
  @Value

metadata-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java

@@ -22,6 +22,7 @@
 import com.linkedin.entity.EnvelopedAspectMap;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.query.SearchFlags;
+import com.linkedin.metadata.search.ScrollResult;
 import com.linkedin.metadata.search.SearchEntity;
 import com.linkedin.metadata.search.SearchEntityArray;
 import com.linkedin.metadata.search.SearchResult;
@@ -35,6 +36,8 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.Set;
+
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
@@ -89,30 +92,58 @@ public void setupTest() throws Exception {
  final EnvelopedAspectMap childDomainPolicyAspectMap = new EnvelopedAspectMap();
  childDomainPolicyAspectMap.put(DATAHUB_POLICY_INFO_ASPECT_NAME, new EnvelopedAspect().setValue(new Aspect(childDomainPolicy.data())));
 
- final SearchResult policySearchResult = new SearchResult();
- policySearchResult.setNumEntities(3);
- policySearchResult.setEntities(
- new SearchEntityArray(
- ImmutableList.of(
- new SearchEntity().setEntity(activePolicyUrn),
- new SearchEntity().setEntity(inactivePolicyUrn),
- new SearchEntity().setEntity(parentDomainPolicyUrn),
- new SearchEntity().setEntity(childDomainPolicyUrn)
- )
- )
- );
-
- when(_entityClient.search(eq("dataHubPolicy"), eq(""), isNull(), any(), anyInt(), anyInt(), any(),
- eq(new SearchFlags().setFulltext(true)))).thenReturn(policySearchResult);
- when(_entityClient.batchGetV2(eq(POLICY_ENTITY_NAME),
- eq(ImmutableSet.of(activePolicyUrn, inactivePolicyUrn, parentDomainPolicyUrn, childDomainPolicyUrn)), eq(null), any())).thenReturn(
- ImmutableMap.of(
- activePolicyUrn, new EntityResponse().setUrn(activePolicyUrn).setAspects(activeAspectMap),
- inactivePolicyUrn, new EntityResponse().setUrn(inactivePolicyUrn).setAspects(inactiveAspectMap),
- parentDomainPolicyUrn, new EntityResponse().setUrn(parentDomainPolicyUrn).setAspects(parentDomainPolicyAspectMap),
- childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspects(childDomainPolicyAspectMap)
- )
- );
+ final ScrollResult policySearchResult1 = new ScrollResult()
+ .setScrollId("1")
+ .setNumEntities(4)
+ .setEntities(
+ new SearchEntityArray(
+ ImmutableList.of(new SearchEntity().setEntity(activePolicyUrn))));
+
+ final ScrollResult policySearchResult2 = new ScrollResult()
+ .setScrollId("2")
+ .setNumEntities(4)
+ .setEntities(
+ new SearchEntityArray(
+ ImmutableList.of(new SearchEntity().setEntity(inactivePolicyUrn))));
+
+ final ScrollResult policySearchResult3 = new ScrollResult()
+ .setScrollId("3")
+ .setNumEntities(4)
+ .setEntities(
+ new SearchEntityArray(
+ ImmutableList.of(new SearchEntity().setEntity(parentDomainPolicyUrn))));
+
+ final ScrollResult policySearchResult4 = new ScrollResult()
+ .setNumEntities(4)
+ .setEntities(
+ new SearchEntityArray(
+ ImmutableList.of(
+ new SearchEntity().setEntity(childDomainPolicyUrn))));
+
+ when(_entityClient.scrollAcrossEntities(eq(List.of("dataHubPolicy")), eq(""), isNull(), any(), isNull(),
+ anyInt(), eq(new SearchFlags().setFulltext(true).setSkipAggregates(true).setSkipHighlighting(true).setSkipCache(true)), any()))
+ .thenReturn(policySearchResult1)
+ .thenReturn(policySearchResult2)
+ .thenReturn(policySearchResult3)
+ .thenReturn(policySearchResult4);
+
+ when(_entityClient.batchGetV2(eq(POLICY_ENTITY_NAME), any(), eq(null), any())).thenAnswer(args -> {
+ Set<Urn> inputUrns = args.getArgument(1);
+ Urn urn = inputUrns.stream().findFirst().get();
+
+ switch (urn.toString()) {
+ case "urn:li:dataHubPolicy:0":
+ return Map.of(activePolicyUrn, new EntityResponse().setUrn(activePolicyUrn).setAspects(activeAspectMap));
+ case "urn:li:dataHubPolicy:1":
+ return Map.of(inactivePolicyUrn, new EntityResponse().setUrn(inactivePolicyUrn).setAspects(inactiveAspectMap));
+ case "urn:li:dataHubPolicy:2":
+ return Map.of(parentDomainPolicyUrn, new EntityResponse().setUrn(parentDomainPolicyUrn).setAspects(parentDomainPolicyAspectMap));
+ case "urn:li:dataHubPolicy:3":
+ return Map.of(childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspects(childDomainPolicyAspectMap));
+ default:
+ throw new IllegalStateException();
+ }
+ });
 
  final List<Urn> userUrns = ImmutableList.of(Urn.createFromString("urn:li:corpuser:user3"), Urn.createFromString("urn:li:corpuser:user4"));
  final List<Urn> groupUrns = ImmutableList.of(Urn.createFromString("urn:li:corpGroup:group3"), Urn.createFromString("urn:li:corpGroup:group4"));
@@ -146,7 +177,8 @@ childDomainPolicyUrn, new EntityResponse().setUrn(childDomainPolicyUrn).setAspec
  _entityClient,
  10,
  10,
- DataHubAuthorizer.AuthorizationMode.DEFAULT
+ DataHubAuthorizer.AuthorizationMode.DEFAULT,
+ 1 // force pagination logic
  );
  _dataHubAuthorizer.init(Collections.emptyMap(), createAuthorizerContext(systemAuthentication, _entityClient));
  _dataHubAuthorizer.invalidateCache();

metadata-service/configuration/src/main/resources/application.yml

@@ -39,6 +39,7 @@ authorization:
  defaultAuthorizer:
  enabled: ${AUTH_POLICIES_ENABLED:true}
  cacheRefreshIntervalSecs: ${POLICY_CACHE_REFRESH_INTERVAL_SECONDS:120}
+ cachePolicyFetchSize: ${POLICY_CACHE_FETCH_SIZE:1000}
  # Enables authorization of reads, writes, and deletes on REST APIs. Defaults to false for backwards compatibility, but should become true down the road
  restApiAuthorization: ${REST_API_AUTHORIZATION_ENABLED:false}
 

metadata-service/factories/src/main/java/com/linkedin/gms/factory/auth/DataHubAuthorizerFactory.java

@@ -32,6 +32,9 @@ public class DataHubAuthorizerFactory {
  @Value("${authorization.defaultAuthorizer.cacheRefreshIntervalSecs}")
  private Integer policyCacheRefreshIntervalSeconds;
 
+ @Value("${authorization.defaultAuthorizer.cachePolicyFetchSize}")
+ private Integer policyCacheFetchSize;
+
  @Value("${authorization.defaultAuthorizer.enabled:true}")
  private Boolean policiesEnabled;
 
@@ -44,6 +47,6 @@ protected DataHubAuthorizer getInstance() {
  : DataHubAuthorizer.AuthorizationMode.ALLOW_ALL;
 
  return new DataHubAuthorizer(systemAuthentication, entityClient, 10,
- policyCacheRefreshIntervalSeconds, mode);
+ policyCacheRefreshIntervalSeconds, mode, policyCacheFetchSize);
  }
 }

metadata-service/restli-client/src/main/java/com/linkedin/entity/client/EntityClient.java

@@ -241,7 +241,7 @@ public SearchResult searchAcrossEntities(@Nonnull List<String> entities, @Nonnul
  */
  @Nonnull
  ScrollResult scrollAcrossEntities(@Nonnull List<String> entities, @Nonnull String input,
- @Nullable Filter filter, @Nullable String scrollId, @Nonnull String keepAlive, int count, @Nullable SearchFlags searchFlags,
+ @Nullable Filter filter, @Nullable String scrollId, @Nullable String keepAlive, int count, @Nullable SearchFlags searchFlags,
  @Nonnull Authentication authentication)
  throws RemoteInvocationException;