fix: support Apollo caching for settings / Policies

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

@@ -4,6 +4,7 @@ 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 { useApolloClient } from '@apollo/client';
 import PolicyBuilderModal from './PolicyBuilderModal';
 import {
     Policy,
@@ -26,7 +27,7 @@ import {
     useUpdatePolicyMutation,
 } from '../../../graphql/policy.generated';
 import { Message } from '../../shared/Message';
-import { EMPTY_POLICY } from './policyUtils';
+import { EMPTY_POLICY, removeFromListPoliciesCache, updateListPoliciesCache } from './policyUtils';
 import TabToolbar from '../../entity/shared/components/styled/TabToolbar';
 import { StyledTable } from '../../entity/shared/components/styled/StyledTable';
 import AvatarsGroup from '../AvatarsGroup';
@@ -141,6 +142,7 @@ const toPolicyInput = (policy: Omit<Policy, 'urn'>): PolicyUpdateInput => {
 export const ManagePolicies = () => {
     const entityRegistry = useEntityRegistry();
     const location = useLocation();
+    const client = useApolloClient();
     const params = QueryString.parse(location.search, { arrayFormat: 'comma' });
     const paramsQuery = (params?.query as string) || undefined;
     const [query, setQuery] = useState<undefined | string>(undefined);
@@ -258,17 +260,21 @@ export const ManagePolicies = () => {
             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();
+                deletePolicy({ variables: { urn: policy?.urn as string } })
+                .then(()=>{
+                    // There must be a focus policy urn.
+                    analytics.event({
+                        type: EventType.DeleteEntityEvent,
+                        entityUrn: policy?.urn,
+                        entityType: EntityType.DatahubPolicy,
+                    });
+                    message.success('Successfully removed policy.');
+                    removeFromListPoliciesCache(client,policy?.urn, DEFAULT_PAGE_SIZE);
+                    setTimeout(() => {
+                        policiesRefetch();
+                    }, 3000);
+                    onCancelViewPolicy();
+                })
             },
             onCancel() {},
             okText: 'Yes',
@@ -289,35 +295,66 @@ export const ManagePolicies = () => {
                 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);
+        }).then(()=>{
+            const updatePolicies= {
+                ...newPolicy,
+                __typename: 'ListPoliciesResult',
+            }
+            updateListPoliciesCache(client,updatePolicies,DEFAULT_PAGE_SIZE);
+            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,
-            });
+            updatePolicy({ variables: { urn: focusPolicyUrn, input: toPolicyInput(savePolicy) } })
+            .then(()=>{
+                const newPolicy = {
+                    __typename: 'ListPoliciesResult',
+                    urn: focusPolicyUrn,
+                    ...savePolicy,
+                };
+                analytics.event({
+                    type: EventType.UpdatePolicyEvent,
+                    policyUrn: focusPolicyUrn,
+                });
+                message.success('Successfully saved policy.');
+                updateListPoliciesCache(client,newPolicy,DEFAULT_PAGE_SIZE);
+                setTimeout(() => {
+                    policiesRefetch();
+                }, 1000);
+                onClosePolicyBuilder();
+            })
         } 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,
-            });
+            createPolicy({ variables: { input: toPolicyInput(savePolicy) } })
+            .then((result)=>{
+                const newPolicy = {
+                    __typename: 'ListPoliciesResult',
+                    urn: result?.data?.createPolicy,
+                    ...savePolicy,
+                    type: null,
+                    actors: null,
+                    resources: null,
+                };
+                analytics.event({
+                    type: EventType.CreatePolicyEvent,
+                });
+                message.success('Successfully saved policy.');
+                setTimeout(() => {
+                    policiesRefetch();
+                }, 1000);
+                updateListPoliciesCache(client,newPolicy,DEFAULT_PAGE_SIZE);
+                onClosePolicyBuilder();
+            })
         }
-        message.success('Successfully saved policy.');
-        setTimeout(() => {
-            policiesRefetch();
-        }, 3000);
-        onClosePolicyBuilder();
     };
 
     const tableColumns = [

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),
+      });
+    });
+  });
+

datahub-web-react/src/app/permissions/policy/policyUtils.ts

@@ -10,6 +10,7 @@ import {
     ResourceFilter,
     ResourcePrivileges,
 } from '../../../types.generated';
+import { ListPoliciesDocument, ListPoliciesQuery } from '../../../graphql/policy.generated';
 
 export const EMPTY_POLICY = {
     type: PolicyType.Metadata,
@@ -126,3 +127,98 @@ export const setFieldValues = (
     }
     return { ...filter, criteria: [...restCriteria, createCriterion(resourceFieldType, fieldValues)] };
 };
+
+export const addOrUpdatePoliciesInList = (existingPolicies, newPolicies) => {
+    const policies = [...existingPolicies];
+    let didUpdate = false;
+    const updatedPolicies = policies.map((policy) => {
+        if (policy.urn === newPolicies.urn) {
+            didUpdate = true;
+            return newPolicies;
+        }
+        return policy;
+    });
+    return didUpdate ? updatedPolicies : [newPolicies, ...existingPolicies];
+};
+
+/**
+ * Add an entry to the ListPolicies cache.
+ */
+export const updateListPoliciesCache = (client, policies, pageSize) => {
+    // Read the data from our cache for this query.
+    const currData: ListPoliciesQuery | null = client.readQuery({
+        query: ListPoliciesDocument,
+        variables: {
+            input: {
+                start: 0,
+                count: pageSize,
+                query: undefined,
+            },
+        },
+    });
+
+    // Add our new policy into the existing list.
+    const existingPolicies = [...(currData?.listPolicies?.policies || [])];
+    const newPolicies = addOrUpdatePoliciesInList(existingPolicies, policies);
+    const didAddTest = newPolicies.length > existingPolicies.length;
+
+    // Write our data back to the cache.
+    client.writeQuery({
+        query: ListPoliciesDocument,
+        variables: {
+            input: {
+                start: 0,
+                count: pageSize,
+                query: undefined,
+            },
+        },
+        data: {
+
+            listPolicies: {
+                __typename: 'ListPoliciesResult',
+                start: 0,
+                count: didAddTest ? (currData?.listPolicies?.count || 0) + 1 : currData?.listPolicies?.count,
+                total: didAddTest ? (currData?.listPolicies?.total || 0) + 1 : currData?.listPolicies?.total,
+                policies: newPolicies,
+            },
+        },
+    });
+};
+
+/**
+ * Remove an entry from the ListTests cache.
+ */
+export const removeFromListPoliciesCache = (client, urn, pageSize) => {
+    // Read the data from our cache for this query.
+    const currData: ListPoliciesQuery | null = client.readQuery({
+        query: ListPoliciesDocument,
+        variables: {
+            input: {
+                start: 0,
+                count: pageSize,
+            },
+        },
+    });
+
+    // Remove the policy from the existing tests set.
+    const newPolicies = [...(currData?.listPolicies?.policies || []).filter((policy) => policy.urn !== urn)];
+
+    // Write our data back to the cache.
+    client.writeQuery({
+        query: ListPoliciesDocument,
+        variables: {
+            input: {
+                start: 0,
+                count: pageSize,
+            },
+        },
+        data: {
+            listPolicies: {
+                start: currData?.listPolicies?.start || 0,
+                count: (currData?.listPolicies?.count || 1) - 1,
+                total: (currData?.listPolicies?.total || 1) - 1,
+                policies: newPolicies,
+            },
+        },
+    });
+};