refactor(ui): support Apollo caching for settings / Policies (#9442)

datahub-web-react/src/app/permissions/policy/ManagePolicies.tsx

@@ -1,32 +1,21 @@
 import React, { useEffect, useMemo, useState } from 'react';
-import { Button, Empty, message, Modal, Pagination, Tag } from 'antd';
+import { Button, Empty, message, Pagination, Tag } from 'antd';
 import styled from 'styled-components/macro';
 import * as QueryString from 'query-string';
 import { DeleteOutlined, PlusOutlined } from '@ant-design/icons';
 import { useLocation } from 'react-router';
 import PolicyBuilderModal from './PolicyBuilderModal';
 import {
     Policy,
-    PolicyUpdateInput,
     PolicyState,
-    PolicyType,
-    Maybe,
-    ResourceFilterInput,
-    PolicyMatchFilter,
-    PolicyMatchFilterInput,
-    PolicyMatchCriterionInput,
-    EntityType,
 } from '../../../types.generated';
 import { useAppConfig } from '../../useAppConfig';
 import PolicyDetailsModal from './PolicyDetailsModal';
 import {
-    useCreatePolicyMutation,
-    useDeletePolicyMutation,
     useListPoliciesQuery,
-    useUpdatePolicyMutation,
 } from '../../../graphql/policy.generated';
 import { Message } from '../../shared/Message';
-import { EMPTY_POLICY } from './policyUtils';
+import { DEFAULT_PAGE_SIZE, EMPTY_POLICY } from './policyUtils';
 import TabToolbar from '../../entity/shared/components/styled/TabToolbar';
 import { StyledTable } from '../../entity/shared/components/styled/StyledTable';
 import AvatarsGroup from '../AvatarsGroup';
@@ -37,6 +26,7 @@ import { scrollToTop } from '../../shared/searchUtils';
 import analytics, { EventType } from '../../analytics';
 import { POLICIES_CREATE_POLICY_ID, POLICIES_INTRO_ID } from '../../onboarding/config/PoliciesOnboardingConfig';
 import { OnboardingTour } from '../../onboarding/OnboardingTour';
+import { usePolicy } from './usePolicy';
 
 const SourceContainer = styled.div`
     overflow: auto;
@@ -84,58 +74,6 @@ const PageContainer = styled.span`
     overflow: auto;
 `;
 
-const DEFAULT_PAGE_SIZE = 10;
-
-type PrivilegeOptionType = {
-    type?: string;
-    name?: Maybe<string>;
-};
-
-const toFilterInput = (filter: PolicyMatchFilter): PolicyMatchFilterInput => {
-    return {
-        criteria: filter.criteria?.map((criterion): PolicyMatchCriterionInput => {
-            return {
-                field: criterion.field,
-                values: criterion.values.map((criterionValue) => criterionValue.value),
-                condition: criterion.condition,
-            };
-        }),
-    };
-};
-
-const toPolicyInput = (policy: Omit<Policy, 'urn'>): PolicyUpdateInput => {
-    let policyInput: PolicyUpdateInput = {
-        type: policy.type,
-        name: policy.name,
-        state: policy.state,
-        description: policy.description,
-        privileges: policy.privileges,
-        actors: {
-            users: policy.actors.users,
-            groups: policy.actors.groups,
-            allUsers: policy.actors.allUsers,
-            allGroups: policy.actors.allGroups,
-            resourceOwners: policy.actors.resourceOwners,
-            resourceOwnersTypes: policy.actors.resourceOwnersTypes,
-        },
-    };
-    if (policy.resources !== null && policy.resources !== undefined) {
-        let resourceFilter: ResourceFilterInput = {
-            type: policy.resources.type,
-            resources: policy.resources.resources,
-            allResources: policy.resources.allResources,
-        };
-        if (policy.resources.filter) {
-            resourceFilter = { ...resourceFilter, filter: toFilterInput(policy.resources.filter) };
-        }
-        // Add the resource filters.
-        policyInput = {
-            ...policyInput,
-            resources: resourceFilter,
-        };
-    }
-    return policyInput;
-};
 
 // TODO: Cleanup the styling.
 export const ManagePolicies = () => {
@@ -163,9 +101,7 @@ export const ManagePolicies = () => {
     const [focusPolicyUrn, setFocusPolicyUrn] = useState<undefined | string>(undefined);
     const [focusPolicy, setFocusPolicy] = useState<Omit<Policy, 'urn'>>(EMPTY_POLICY);
 
-    // Construct privileges
-    const platformPrivileges = policiesConfig?.platformPrivileges || [];
-    const resourcePrivileges = policiesConfig?.resourcePrivileges || [];
+
 
     const {
         loading: policiesLoading,
@@ -183,15 +119,6 @@ export const ManagePolicies = () => {
         fetchPolicy: (query?.length || 0) > 0 ? 'no-cache' : 'cache-first',
     });
 
-    // Any time a policy is removed, edited, or created, refetch the list.
-    const [createPolicy, { error: createPolicyError }] = useCreatePolicyMutation();
-
-    const [updatePolicy, { error: updatePolicyError }] = useUpdatePolicyMutation();
-
-    const [deletePolicy, { error: deletePolicyError }] = useDeletePolicyMutation();
-
-    const updateError = createPolicyError || updatePolicyError || deletePolicyError;
-
     const totalPolicies = policiesData?.listPolicies?.total || 0;
     const policies = useMemo(() => policiesData?.listPolicies?.policies || [], [policiesData]);
 
@@ -212,28 +139,6 @@ export const ManagePolicies = () => {
         setShowPolicyBuilderModal(false);
     };
 
-    const getPrivilegeNames = (policy: Omit<Policy, 'urn'>) => {
-        let privileges: PrivilegeOptionType[] = [];
-        if (policy?.type === PolicyType.Platform) {
-            privileges = platformPrivileges
-                .filter((platformPrivilege) => policy.privileges.includes(platformPrivilege.type))
-                .map((platformPrivilege) => {
-                    return { type: platformPrivilege.type, name: platformPrivilege.displayName };
-                });
-        } else {
-            const allResourcePriviliges = resourcePrivileges.find(
-                (resourcePrivilege) => resourcePrivilege.resourceType === 'all',
-            );
-            privileges =
-                allResourcePriviliges?.privileges
-                    .filter((resourcePrivilege) => policy.privileges.includes(resourcePrivilege.type))
-                    .map((b) => {
-                        return { type: b.type, name: b.displayName };
-                    }) || [];
-        }
-        return privileges;
-    };
-
     const onViewPolicy = (policy: Policy) => {
         setShowViewPolicyModal(true);
         setFocusPolicyUrn(policy?.urn);
@@ -247,79 +152,30 @@ export const ManagePolicies = () => {
     };
 
     const onEditPolicy = (policy: Policy) => {
-        setShowPolicyBuilderModal(true);
-        setFocusPolicyUrn(policy?.urn);
-        setFocusPolicy({ ...policy });
-    };
-
-    // On Delete Policy handler
-    const onRemovePolicy = (policy: Policy) => {
-        Modal.confirm({
-            title: `Delete ${policy?.name}`,
-            content: `Are you sure you want to remove policy?`,
-            onOk() {
-                deletePolicy({ variables: { urn: policy?.urn as string } }); // There must be a focus policy urn.
-                analytics.event({
-                    type: EventType.DeleteEntityEvent,
-                    entityUrn: policy?.urn,
-                    entityType: EntityType.DatahubPolicy,
-                });
-                message.success('Successfully removed policy.');
-                setTimeout(() => {
-                    policiesRefetch();
-                }, 3000);
-                onCancelViewPolicy();
-            },
-            onCancel() {},
-            okText: 'Yes',
-            maskClosable: true,
-            closable: true,
-        });
+    setShowPolicyBuilderModal(true);
+    setFocusPolicyUrn(policy?.urn);
+    setFocusPolicy({ ...policy });
     };
 
-    // On Activate and deactivate Policy handler
-    const onToggleActiveDuplicate = (policy: Policy) => {
-        const newState = policy?.state === PolicyState.Active ? PolicyState.Inactive : PolicyState.Active;
-        const newPolicy = {
-            ...policy,
-            state: newState,
-        };
-        updatePolicy({
-            variables: {
-                urn: policy?.urn as string, // There must be a focus policy urn.
-                input: toPolicyInput(newPolicy),
-            },
-        });
-        message.success(`Successfully ${newState === PolicyState.Active ? 'activated' : 'deactivated'} policy.`);
-        setTimeout(() => {
-            policiesRefetch();
-        }, 3000);
-        setShowViewPolicyModal(false);
-    };
-
-    // On Add/Update Policy handler
-    const onSavePolicy = (savePolicy: Omit<Policy, 'urn'>) => {
-        if (focusPolicyUrn) {
-            // If there's an URN associated with the focused policy, then we are editing an existing policy.
-            updatePolicy({ variables: { urn: focusPolicyUrn, input: toPolicyInput(savePolicy) } });
-            analytics.event({
-                type: EventType.UpdatePolicyEvent,
-                policyUrn: focusPolicyUrn,
-            });
-        } else {
-            // If there's no URN associated with the focused policy, then we are creating.
-            createPolicy({ variables: { input: toPolicyInput(savePolicy) } });
-            analytics.event({
-                type: EventType.CreatePolicyEvent,
-            });
-        }
-        message.success('Successfully saved policy.');
-        setTimeout(() => {
-            policiesRefetch();
-        }, 3000);
-        onClosePolicyBuilder();
-    };
+    const {
+        createPolicyError,
+        updatePolicyError,
+        deletePolicyError,
+        onSavePolicy,
+        onToggleActiveDuplicate,
+        onRemovePolicy,
+        getPrivilegeNames
+    } = usePolicy(
+        policiesConfig,
+        focusPolicyUrn,
+        policiesRefetch,
+        setShowViewPolicyModal,
+        onCancelViewPolicy,
+        onClosePolicyBuilder
+    );
 
+    const updateError = createPolicyError || updatePolicyError || deletePolicyError;
+
     const tableColumns = [
         {
             title: 'Name',

datahub-web-react/src/app/permissions/policy/_tests_/policyUtils.test.tsx

@@ -0,0 +1,110 @@
+import {
+    addOrUpdatePoliciesInList,
+    updateListPoliciesCache,
+    removeFromListPoliciesCache,
+  } from '../policyUtils';
+
+  // Mock the Apollo Client readQuery and writeQuery methods
+  const mockReadQuery = jest.fn();
+  const mockWriteQuery = jest.fn();
+
+  jest.mock('@apollo/client', () => ({
+    ...jest.requireActual('@apollo/client'),
+    useApolloClient: () => ({
+      readQuery: mockReadQuery,
+      writeQuery: mockWriteQuery,
+    }),
+  }));
+
+  describe('addOrUpdatePoliciesInList', () => {
+    it('should add a new policy to the list', () => {
+      const existingPolicies = [{ urn: 'existing-urn' }];
+      const newPolicies = { urn: 'new-urn' };
+
+      const result = addOrUpdatePoliciesInList(existingPolicies, newPolicies);
+
+      expect(result.length).toBe(existingPolicies.length + 1);
+      expect(result).toContain(newPolicies);
+    });
+
+    it('should update an existing policy in the list', () => {
+      const existingPolicies = [{ urn: 'existing-urn' }];
+      const newPolicies = { urn: 'existing-urn', updatedField: 'new-value' };
+
+      const result = addOrUpdatePoliciesInList(existingPolicies, newPolicies);
+
+      expect(result.length).toBe(existingPolicies.length);
+      expect(result).toContainEqual(newPolicies);
+    });
+  });
+
+  describe('updateListPoliciesCache', () => {
+    // Mock client.readQuery response
+    const mockReadQueryResponse = {
+      listPolicies: {
+        start: 0,
+        count: 1,
+        total: 1,
+        policies: [{ urn: 'existing-urn' }],
+      },
+    };
+
+    beforeEach(() => {
+      mockReadQuery.mockReturnValueOnce(mockReadQueryResponse);
+    });
+
+    it('should update the list policies cache with a new policy', () => {
+      const mockClient = {
+        readQuery: mockReadQuery,
+        writeQuery: mockWriteQuery,
+      };
+
+      const policiesToAdd = [{ urn: 'new-urn' }];
+      const pageSize = 10;
+
+      updateListPoliciesCache(mockClient, policiesToAdd, pageSize);
+
+      // Ensure writeQuery is called with the expected data
+      expect(mockWriteQuery).toHaveBeenCalledWith({
+        query: expect.any(Object),
+        variables: { input: { start: 0, count: pageSize, query: undefined } },
+        data: expect.any(Object),
+      });
+    });
+  });
+
+  describe('removeFromListPoliciesCache', () => {
+    // Mock client.readQuery response
+    const mockReadQueryResponse = {
+      listPolicies: {
+        start: 0,
+        count: 1,
+        total: 1,
+        policies: [{ urn: 'existing-urn' }],
+      },
+    };
+
+    beforeEach(() => {
+      mockReadQuery.mockReturnValueOnce(mockReadQueryResponse);
+    });
+
+    it('should remove a policy from the list policies cache', () => {
+      const mockClient = {
+        readQuery: mockReadQuery,
+        writeQuery: mockWriteQuery,
+      };
+
+      const urnToRemove = 'existing-urn';
+      const pageSize = 10;
+
+      removeFromListPoliciesCache(mockClient, urnToRemove, pageSize);
+
+      // Ensure writeQuery is called with the expected data
+      expect(mockWriteQuery).toHaveBeenCalledWith({
+        query: expect.any(Object),
+        variables: { input: { start: 0, count: pageSize } },
+        data: expect.any(Object),
+      });
+    });
+  });
+