restrict policy fetch

datahub-project · Apr 9, 2024 · ffd8b82 · ffd8b82
1 parent 68ba65e
commit ffd8b82
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java b/metadata-service/auth-impl/src/main/java/com/datahub/authorization/PolicyFetcher.java
@@ -21,6 +21,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -107,7 +108,11 @@ public PolicyFetchResult fetchPolicies(
 
     // Fetch DataHubPolicyInfo aspects for each urn
     final Map<Urn, EntityResponse> policyEntities =
-        entityClient.batchGetV2(opContext, POLICY_ENTITY_NAME, new HashSet<>(policyUrns), null);
+        entityClient.batchGetV2(
+            opContext,
+            POLICY_ENTITY_NAME,
+            new HashSet<>(policyUrns),
+            Set.of(DATAHUB_POLICY_INFO_ASPECT_NAME));
     return new PolicyFetchResult(
         policyUrns.stream()
             .map(policyEntities::get)

diff --git a/...data-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java b/...data-service/auth-impl/src/test/java/com/datahub/authorization/DataHubAuthorizerTest.java
@@ -4,6 +4,7 @@
 import static com.linkedin.metadata.authorization.PoliciesConfig.ACTIVE_POLICY_STATE;
 import static com.linkedin.metadata.authorization.PoliciesConfig.INACTIVE_POLICY_STATE;
 import static com.linkedin.metadata.authorization.PoliciesConfig.METADATA_POLICY_TYPE;
+import static org.mockito.ArgumentMatchers.anySet;
 import static org.mockito.ArgumentMatchers.isNull;
 import static org.mockito.ArgumentMatchers.nullable;
 import static org.mockito.Mockito.any;
@@ -211,7 +212,7 @@ public void setupTest() throws Exception {
         .thenReturn(policySearchResult5);
 
     when(_entityClient.batchGetV2(
-            any(OperationContext.class), eq(POLICY_ENTITY_NAME), any(), eq(null)))
+            any(OperationContext.class), eq(POLICY_ENTITY_NAME), any(), anySet()))
         .thenAnswer(
             args -> {
               Set<Urn> inputUrns = args.getArgument(2);