refactor(ui): Adding caching for users, groups, and roles (#6673)

datahub-project · Dec 8, 2022 · bbcec3b · bbcec3b
1 parent 5b52534
commit bbcec3b
Show file tree

Hide file tree

Showing 22 changed files with 308 additions and 74 deletions.
diff --git a/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java b/datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java
@@ -30,6 +30,7 @@
 import com.linkedin.datahub.graphql.generated.Chart;
 import com.linkedin.datahub.graphql.generated.ChartInfo;
 import com.linkedin.datahub.graphql.generated.Container;
+import com.linkedin.datahub.graphql.generated.CorpGroup;
 import com.linkedin.datahub.graphql.generated.CorpGroupInfo;
 import com.linkedin.datahub.graphql.generated.CorpUser;
 import com.linkedin.datahub.graphql.generated.CorpUserInfo;
@@ -58,6 +59,7 @@
 import com.linkedin.datahub.graphql.generated.LineageRelationship;
 import com.linkedin.datahub.graphql.generated.ListAccessTokenResult;
 import com.linkedin.datahub.graphql.generated.ListDomainsResult;
+import com.linkedin.datahub.graphql.generated.ListGroupsResult;
 import com.linkedin.datahub.graphql.generated.ListTestsResult;
 import com.linkedin.datahub.graphql.generated.MLFeature;
 import com.linkedin.datahub.graphql.generated.MLFeatureProperties;
@@ -1069,6 +1071,12 @@ private void configureCorpGroupResolvers(final RuntimeWiring.Builder builder) {
                             (env) -> ((CorpGroupInfo) env.getSource()).getMembers().stream()
                                     .map(CorpUser::getUrn)
                                     .collect(Collectors.toList())))
+        )
+        .type("ListGroupsResult", typeWiring -> typeWiring
+            .dataFetcher("groups", new LoadableTypeBatchResolver<>(corpGroupType,
+                (env) -> ((ListGroupsResult) env.getSource()).getGroups().stream()
+                    .map(CorpGroup::getUrn)
+                    .collect(Collectors.toList())))
         );
     }
 

diff --git a/...l-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java b/...l-core/src/main/java/com/linkedin/datahub/graphql/resolvers/group/ListGroupsResolver.java
@@ -5,17 +5,18 @@
 import com.linkedin.datahub.graphql.authorization.AuthorizationUtils;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
 import com.linkedin.datahub.graphql.generated.CorpGroup;
+import com.linkedin.datahub.graphql.generated.EntityType;
 import com.linkedin.datahub.graphql.generated.ListGroupsInput;
 import com.linkedin.datahub.graphql.generated.ListGroupsResult;
-import com.linkedin.datahub.graphql.types.corpgroup.mappers.CorpGroupMapper;
 import com.linkedin.entity.EntityResponse;
 import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.query.filter.SortCriterion;
+import com.linkedin.metadata.query.filter.SortOrder;
 import com.linkedin.metadata.search.SearchEntity;
 import com.linkedin.metadata.search.SearchResult;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
-import java.util.Collection;
-import java.util.Collections;
+import java.util.ArrayList;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
@@ -53,7 +54,11 @@ public CompletableFuture<ListGroupsResult> get(final DataFetchingEnvironment env
         try {
           // First, get all group Urns.
           final SearchResult gmsResult =
-              _entityClient.search(CORP_GROUP_ENTITY_NAME, query, Collections.emptyMap(), start, count, context.getAuthentication());
+              _entityClient.search(CORP_GROUP_ENTITY_NAME,
+                  query,
+                  null,
+                  new SortCriterion().setField(CORP_GROUP_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
+                  start, count, context.getAuthentication());
 
           // Then, get hydrate all groups.
           final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(CORP_GROUP_ENTITY_NAME,
@@ -66,7 +71,9 @@ public CompletableFuture<ListGroupsResult> get(final DataFetchingEnvironment env
           result.setStart(gmsResult.getFrom());
           result.setCount(gmsResult.getPageSize());
           result.setTotal(gmsResult.getNumEntities());
-          result.setGroups(mapEntities(entities.values()));
+          result.setGroups(mapUnresolvedGroups(gmsResult.getEntities().stream()
+              .map(SearchEntity::getEntity)
+              .collect(Collectors.toList())));
           return result;
         } catch (Exception e) {
           throw new RuntimeException("Failed to list groups", e);
@@ -76,9 +83,15 @@ public CompletableFuture<ListGroupsResult> get(final DataFetchingEnvironment env
     throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
   }
 
-  private List<CorpGroup> mapEntities(final Collection<EntityResponse> entities) {
-    return entities.stream()
-        .map(CorpGroupMapper::map)
-        .collect(Collectors.toList());
+  // This method maps urns returned from the list endpoint into Partial Group objects which will be resolved be a separate Batch resolver.
+  private List<CorpGroup> mapUnresolvedGroups(final List<Urn> entityUrns) {
+    final List<CorpGroup> results = new ArrayList<>();
+    for (final Urn urn : entityUrns) {
+      final CorpGroup unresolvedGroup = new CorpGroup();
+      unresolvedGroup.setUrn(urn.toString());
+      unresolvedGroup.setType(EntityType.CORP_GROUP);
+      results.add(unresolvedGroup);
+    }
+    return results;
   }
 }
diff --git a/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx b/datahub-web-react/src/app/entity/shared/EntityDropdown/EntityDropdown.tsx
@@ -56,12 +56,18 @@ const StyledMenuItem = styled(Menu.Item)<{ disabled: boolean }>`
             : ''}
 `;
 
+interface Options {
+    hideDeleteMessage?: boolean;
+    skipDeleteWait?: boolean;
+}
+
 interface Props {
     urn: string;
     entityType: EntityType;
     entityData?: any;
     menuItems: Set<EntityMenuItems>;
     size?: number;
+    options?: Options;
     refetchForEntity?: () => void;
     refetchForTerms?: () => void;
     refetchForNodes?: () => void;
@@ -81,11 +87,19 @@ function EntityDropdown(props: Props) {
         refreshBrowser,
         onDeleteEntity: onDelete,
         size,
+        options,
     } = props;
 
     const entityRegistry = useEntityRegistry();
     const [updateDeprecation] = useUpdateDeprecationMutation();
-    const { onDeleteEntity, hasBeenDeleted } = useDeleteEntity(urn, entityType, entityData, onDelete);
+    const { onDeleteEntity, hasBeenDeleted } = useDeleteEntity(
+        urn,
+        entityType,
+        entityData,
+        onDelete,
+        options?.hideDeleteMessage,
+        options?.skipDeleteWait,
+    );
 
     const [isCreateTermModalVisible, setIsCreateTermModalVisible] = useState(false);
     const [isCreateNodeModalVisible, setIsCreateNodeModalVisible] = useState(false);

diff --git a/datahub-web-react/src/app/entity/shared/EntityDropdown/useDeleteEntity.tsx b/datahub-web-react/src/app/entity/shared/EntityDropdown/useDeleteEntity.tsx
@@ -12,7 +12,14 @@ import analytics, { EventType } from '../../../analytics';
  * @param type the type of the entity to delete
  * @param name the name of the entity to delete
  */
-function useDeleteEntity(urn: string, type: EntityType, entityData: any, onDelete?: () => void) {
+function useDeleteEntity(
+    urn: string,
+    type: EntityType,
+    entityData: any,
+    onDelete?: () => void,
+    hideMessage?: boolean,
+    skipWait?: boolean,
+) {
     const [hasBeenDeleted, setHasBeenDeleted] = useState(false);
     const entityRegistry = useEntityRegistry();
 
@@ -31,18 +38,25 @@ function useDeleteEntity(urn: string, type: EntityType, entityData: any, onDelet
                     entityUrn: urn,
                     entityType: type,
                 });
-                message.loading({
-                    content: 'Deleting...',
-                    duration: 2,
-                });
-                setTimeout(() => {
-                    setHasBeenDeleted(true);
-                    onDelete?.();
-                    message.success({
-                        content: `Deleted ${entityRegistry.getEntityName(type)}!`,
+                if (!hideMessage && !skipWait) {
+                    message.loading({
+                        content: 'Deleting...',
                         duration: 2,
                     });
-                }, 2000);
+                }
+                setTimeout(
+                    () => {
+                        setHasBeenDeleted(true);
+                        onDelete?.();
+                        if (!hideMessage) {
+                            message.success({
+                                content: `Deleted ${entityRegistry.getEntityName(type)}!`,
+                                duration: 2,
+                            });
+                        }
+                    },
+                    skipWait ? 0 : 2000,
+                );
             })
             .catch((e) => {
                 message.destroy();

diff --git a/datahub-web-react/src/app/identity/group/CreateGroupModal.tsx b/datahub-web-react/src/app/identity/group/CreateGroupModal.tsx
@@ -4,10 +4,11 @@ import { useCreateGroupMutation } from '../../../graphql/group.generated';
 import { useEnterKeyListener } from '../../shared/useEnterKeyListener';
 import { groupIdTextValidation } from '../../shared/textUtil';
 import analytics, { EventType } from '../../analytics';
+import { CorpGroup, EntityType } from '../../../types.generated';
 
 type Props = {
     onClose: () => void;
-    onCreate: (name: string, description: string) => void;
+    onCreate: (group: CorpGroup) => void;
 };
 
 export default function CreateGroupModal({ onClose, onCreate }: Props) {
@@ -28,23 +29,31 @@ export default function CreateGroupModal({ onClose, onCreate }: Props) {
                 },
             },
         })
-            .then(({ errors }) => {
+            .then(({ data, errors }) => {
                 if (!errors) {
                     analytics.event({
                         type: EventType.CreateGroupEvent,
                     });
+                    message.success({
+                        content: `Created group!`,
+                        duration: 3,
+                    });
+                    // TODO: Get a full corp group back from create endpoint.
+                    onCreate({
+                        urn: data?.createGroup || '',
+                        type: EntityType.CorpGroup,
+                        name: stagedName,
+                        info: {
+                            description: stagedDescription,
+                        },
+                    });
                 }
             })
             .catch((e) => {
                 message.destroy();
                 message.error({ content: `Failed to create group!: \n ${e.message || ''}`, duration: 3 });
             })
             .finally(() => {
-                message.success({
-                    content: `Created group!`,
-                    duration: 3,
-                });
-                onCreate(stagedName, stagedDescription);
                 setStagedName('');
                 setStagedDescription('');
             });

diff --git a/datahub-web-react/src/app/identity/group/GroupList.tsx b/datahub-web-react/src/app/identity/group/GroupList.tsx
@@ -15,6 +15,7 @@ import { useEntityRegistry } from '../../useEntityRegistry';
 import { scrollToTop } from '../../shared/searchUtils';
 import { GROUPS_CREATE_GROUP_ID, GROUPS_INTRO_ID } from '../../onboarding/config/GroupsOnboardingConfig';
 import { OnboardingTour } from '../../onboarding/OnboardingTour';
+import { addGroupToListGroupsCache, DEFAULT_GROUP_LIST_PAGE_SIZE, removeGroupFromListGroupsCache } from './cacheUtils';
 
 const GroupContainer = styled.div``;
 
@@ -30,8 +31,6 @@ const GroupPaginationContainer = styled.div`
     justify-content: center;
 `;
 
-const DEFAULT_PAGE_SIZE = 25;
-
 export const GroupList = () => {
     const entityRegistry = useEntityRegistry();
     const location = useLocation();
@@ -42,39 +41,32 @@ export const GroupList = () => {
 
     const [page, setPage] = useState(1);
     const [isCreatingGroup, setIsCreatingGroup] = useState(false);
-    const [removedUrns, setRemovedUrns] = useState<string[]>([]);
 
     // Policy list paging.
-    const pageSize = DEFAULT_PAGE_SIZE;
+    const pageSize = DEFAULT_GROUP_LIST_PAGE_SIZE;
     const start = (page - 1) * pageSize;
 
-    const { loading, error, data, refetch } = useListGroupsQuery({
+    const { loading, error, data, refetch, client } = useListGroupsQuery({
         variables: {
             input: {
                 start,
                 count: pageSize,
-                query,
+                query: (query?.length && query) || undefined,
             },
         },
-        fetchPolicy: 'no-cache',
+        fetchPolicy: (query?.length || 0) > 0 ? 'no-cache' : 'cache-first',
     });
 
     const totalGroups = data?.listGroups?.total || 0;
     const groups = data?.listGroups?.groups || [];
-    const filteredGroups = groups.filter((group) => !removedUrns.includes(group.urn));
 
     const onChangePage = (newPage: number) => {
         scrollToTop();
         setPage(newPage);
     };
 
     const handleDelete = (urn: string) => {
-        // Hack to deal with eventual consistency.
-        const newRemovedUrns = [...removedUrns, urn];
-        setRemovedUrns(newRemovedUrns);
-        setTimeout(function () {
-            refetch?.();
-        }, 3000);
+        removeGroupFromListGroupsCache(urn, client);
     };
 
     return (
@@ -84,11 +76,9 @@ export const GroupList = () => {
             {error && <Message type="error" content="Failed to load groups! An unexpected error occurred." />}
             <GroupContainer>
                 <TabToolbar>
-                    <div>
-                        <Button id={GROUPS_CREATE_GROUP_ID} type="text" onClick={() => setIsCreatingGroup(true)}>
-                            <UsergroupAddOutlined /> Create group
-                        </Button>
-                    </div>
+                    <Button id={GROUPS_CREATE_GROUP_ID} type="text" onClick={() => setIsCreatingGroup(true)}>
+                        <UsergroupAddOutlined /> Create group
+                    </Button>
                     <SearchBar
                         initialQuery={query || ''}
                         placeholderText="Search groups..."
@@ -112,7 +102,7 @@ export const GroupList = () => {
                     locale={{
                         emptyText: <Empty description="No Groups!" image={Empty.PRESENTED_IMAGE_SIMPLE} />,
                     }}
-                    dataSource={filteredGroups}
+                    dataSource={groups}
                     renderItem={(item: any) => (
                         <GroupListItem onDelete={() => handleDelete(item.urn)} group={item as CorpGroup} />
                     )}
@@ -131,11 +121,9 @@ export const GroupList = () => {
                 {isCreatingGroup && (
                     <CreateGroupModal
                         onClose={() => setIsCreatingGroup(false)}
-                        onCreate={() => {
-                            // Hack to deal with eventual consistency.
-                            setTimeout(function () {
-                                refetch?.();
-                            }, 2000);
+                        onCreate={(group) => {
+                            addGroupToListGroupsCache(group, client);
+                            setTimeout(() => refetch(), 3000);
                         }}
                     />
                 )}

diff --git a/datahub-web-react/src/app/identity/group/GroupListItem.tsx b/datahub-web-react/src/app/identity/group/GroupListItem.tsx
@@ -73,6 +73,7 @@ export default function GroupListItem({ group, onDelete }: Props) {
                         menuItems={new Set([EntityMenuItems.DELETE])}
                         size={20}
                         onDeleteEntity={onDelete}
+                        options={{ hideDeleteMessage: false, skipDeleteWait: true }}
                     />
                 </GroupItemButtonGroup>
             </GroupItemContainer>