refactor(ui): Caching Ingestion Secrets (#6772)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/CreateSecretResolver.java

@@ -1,5 +1,7 @@
 package com.linkedin.datahub.graphql.resolvers.ingest.secret;
 
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.data.template.SetMode;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.exception.AuthorizationException;
@@ -63,6 +65,7 @@ public class CreateSecretResolver implements DataFetcher<CompletableFuture<Strin
           value.setName(input.getName());
           value.setValue(_secretService.encrypt(input.getValue()));
           value.setDescription(input.getDescription(), SetMode.IGNORE_NULL);
+          value.setCreated(new AuditStamp().setActor(UrnUtils.getUrn(context.getActorUrn())).setTime(System.currentTimeMillis()));
 
           proposal.setEntityType(Constants.SECRETS_ENTITY_NAME);
           proposal.setAspectName(Constants.SECRET_VALUE_ASPECT_NAME);

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolver.java

@@ -14,22 +14,24 @@ import com.linkedin.entity.EnvelopedAspect;
 import com.linkedin.entity.EnvelopedAspectMap;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
+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 com.linkedin.secret.DataHubSecretValue;
 import graphql.schema.DataFetcher;
 import graphql.schema.DataFetchingEnvironment;
 import java.util.ArrayList;
-import java.util.Collection;
-import java.util.Collections;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
 import lombok.extern.slf4j.Slf4j;
 
 import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.*;
+import static com.linkedin.metadata.Constants.*;
 
 
 /**
@@ -62,8 +64,14 @@ public class ListSecretsResolver implements DataFetcher<CompletableFuture<ListSe
       return CompletableFuture.supplyAsync(() -> {
         try {
           // First, get all secrets
-          final SearchResult
+          final SearchResult gmsResult = _entityClient.search(
-              gmsResult = _entityClient.search(Constants.SECRETS_ENTITY_NAME, query, Collections.emptyMap(), start, count, context.getAuthentication());
+              Constants.SECRETS_ENTITY_NAME,
+              query,
+              null,
+              new SortCriterion().setField(DOMAIN_CREATED_TIME_INDEX_FIELD_NAME).setOrder(SortOrder.DESCENDING),
+              start,
+              count,
+              context.getAuthentication());
 
           // Then, resolve all secrets
           final Map<Urn, EntityResponse> entities = _entityClient.batchGetV2(
@@ -79,7 +87,10 @@ public class ListSecretsResolver implements DataFetcher<CompletableFuture<ListSe
           result.setStart(gmsResult.getFrom());
           result.setCount(gmsResult.getPageSize());
           result.setTotal(gmsResult.getNumEntities());
-          result.setSecrets(mapEntities(entities.values()));
+          result.setSecrets(mapEntities(gmsResult.getEntities().stream()
+              .map(entity -> entities.get(entity.getEntity()))
+              .filter(Objects::nonNull)
+              .collect(Collectors.toList())));
           return result;
 
         } catch (Exception e) {
@@ -90,7 +101,7 @@ public class ListSecretsResolver implements DataFetcher<CompletableFuture<ListSe
     throw new AuthorizationException("Unauthorized to perform this action. Please contact your DataHub administrator.");
   }
 
-  private List<Secret> mapEntities(final Collection<EntityResponse> entities) {
+  private List<Secret> mapEntities(final List<EntityResponse> entities) {
     final List<Secret> results = new ArrayList<>();
     for (EntityResponse response : entities) {
       final Urn entityUrn = response.getUrn();

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/CreateSecretResolverMatcherTest.java

@@ -0,0 +1,45 @@
+package com.linkedin.datahub.graphql.resolvers.ingest.secret;
+
+import com.linkedin.metadata.utils.GenericRecordUtils;
+import com.linkedin.mxe.GenericAspect;
+import com.linkedin.mxe.MetadataChangeProposal;
+import com.linkedin.secret.DataHubSecretValue;
+import org.mockito.ArgumentMatcher;
+
+
+public class CreateSecretResolverMatcherTest implements ArgumentMatcher<MetadataChangeProposal> {
+
+  private MetadataChangeProposal left;
+
+  public CreateSecretResolverMatcherTest(MetadataChangeProposal left) {
+    this.left = left;
+  }
+
+  @Override
+  public boolean matches(MetadataChangeProposal right) {
+    return left.getEntityType().equals(right.getEntityType())
+        && left.getAspectName().equals(right.getAspectName())
+        && left.getChangeType().equals(right.getChangeType())
+        && secretPropertiesMatch(left.getAspect(), right.getAspect());
+  }
+
+  private boolean secretPropertiesMatch(GenericAspect left, GenericAspect right) {
+    DataHubSecretValue leftProps = GenericRecordUtils.deserializeAspect(
+        left.getValue(),
+        "application/json",
+        DataHubSecretValue.class
+    );
+
+    DataHubSecretValue rightProps = GenericRecordUtils.deserializeAspect(
+        right.getValue(),
+        "application/json",
+        DataHubSecretValue.class
+    );
+
+    // Omit timestamp comparison.
+    return leftProps.getName().equals(rightProps.getName())
+        && leftProps.getValue().equals(rightProps.getValue())
+        && leftProps.getDescription().equals(rightProps.getDescription())
+        && leftProps.getCreated().getActor().equals(rightProps.getCreated().getActor());
+  }
+}

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/CreateSecretResolverTest.java

@@ -2,6 +2,8 @@ package com.linkedin.datahub.graphql.resolvers.ingest.secret;
 
 
 import com.datahub.authentication.Authentication;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.UrnUtils;
 import com.linkedin.datahub.graphql.QueryContext;
 import com.linkedin.datahub.graphql.generated.CreateSecretInput;
 import com.linkedin.datahub.graphql.resolvers.ingest.source.UpsertIngestionSourceResolver;
@@ -55,16 +57,15 @@ public class CreateSecretResolverTest {
     value.setValue("encryptedvalue");
     value.setName(TEST_INPUT.getName());
     value.setDescription(TEST_INPUT.getDescription());
+    value.setCreated(new AuditStamp().setActor(UrnUtils.getUrn("urn:li:corpuser:test")).setTime(0L));
 
     Mockito.verify(mockClient, Mockito.times(1)).ingestProposal(
-        Mockito.eq(
+        Mockito.argThat(new CreateSecretResolverMatcherTest(new MetadataChangeProposal()
-            new MetadataChangeProposal()
             .setChangeType(ChangeType.UPSERT)
             .setEntityType(Constants.SECRETS_ENTITY_NAME)
             .setAspectName(Constants.SECRET_VALUE_ASPECT_NAME)
             .setAspect(GenericRecordUtils.serializeAspect(value))
-                .setEntityKeyAspect(GenericRecordUtils.serializeAspect(key))
+            .setEntityKeyAspect(GenericRecordUtils.serializeAspect(key)))),
-        ),
         Mockito.any(Authentication.class)
     );
   }

datahub-graphql-core/src/test/java/com/linkedin/datahub/graphql/resolvers/ingest/secret/ListSecretsResolverTest.java

@@ -11,13 +11,13 @@ import com.linkedin.entity.EnvelopedAspect;
 import com.linkedin.entity.EnvelopedAspectMap;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.Constants;
+import com.linkedin.metadata.query.filter.SortCriterion;
 import com.linkedin.metadata.search.SearchEntity;
 import com.linkedin.metadata.search.SearchEntityArray;
 import com.linkedin.metadata.search.SearchResult;
 import com.linkedin.r2.RemoteInvocationException;
 import com.linkedin.secret.DataHubSecretValue;
 import graphql.schema.DataFetchingEnvironment;
-import java.util.Collections;
 import java.util.HashSet;
 import org.mockito.Mockito;
 import org.testng.annotations.Test;
@@ -42,7 +42,8 @@ public class ListSecretsResolverTest {
     Mockito.when(mockClient.search(
         Mockito.eq(Constants.SECRETS_ENTITY_NAME),
         Mockito.eq(""),
-        Mockito.eq(Collections.emptyMap()),
+        Mockito.eq(null),
+        Mockito.any(SortCriterion.class),
         Mockito.eq(0),
         Mockito.eq(20),
         Mockito.any(Authentication.class)
@@ -109,7 +110,8 @@ public class ListSecretsResolverTest {
     Mockito.verify(mockClient, Mockito.times(0)).search(
         Mockito.any(),
         Mockito.eq(""),
-        Mockito.anyMap(),
+        Mockito.eq(null),
+        Mockito.any(SortCriterion.class),
         Mockito.anyInt(),
         Mockito.anyInt(),
         Mockito.any(Authentication.class));

datahub-web-react/src/app/ingest/secret/SecretsList.tsx

@@ -17,6 +17,7 @@ import { StyledTable } from '../../entity/shared/components/styled/StyledTable';
 import { SearchBar } from '../../search/SearchBar';
 import { useEntityRegistry } from '../../useEntityRegistry';
 import { scrollToTop } from '../../shared/searchUtils';
+import { addSecretToListSecretsCache, removeSecretFromListSecretsCache } from './cacheUtils';
 
 const DeleteButtonContainer = styled.div`
     display: flex;
@@ -45,24 +46,22 @@ export const SecretsList = () => {
 
     // Whether or not there is an urn to show in the modal
     const [isCreatingSecret, setIsCreatingSecret] = useState<boolean>(false);
-    const [removedUrns, setRemovedUrns] = useState<string[]>([]);
 
     const [deleteSecretMutation] = useDeleteSecretMutation();
     const [createSecretMutation] = useCreateSecretMutation();
-    const { loading, error, data, refetch } = useListSecretsQuery({
+    const { loading, error, data, client } = useListSecretsQuery({
         variables: {
             input: {
                 start,
                 count: pageSize,
-                query,
+                query: query && query.length > 0 ? query : undefined,
             },
         },
-        fetchPolicy: 'no-cache',
+        fetchPolicy: query && query.length > 0 ? 'no-cache' : 'cache-first',
     });
 
     const totalSecrets = data?.listSecrets?.total || 0;
     const secrets = data?.listSecrets?.secrets || [];
-    const filteredSecrets = secrets.filter((user) => !removedUrns.includes(user.urn));
 
     const deleteSecret = async (urn: string) => {
         deleteSecretMutation({
@@ -70,11 +69,7 @@ export const SecretsList = () => {
         })
             .then(() => {
                 message.success({ content: 'Removed secret.', duration: 2 });
-                const newRemovedUrns = [...removedUrns, urn];
+                removeSecretFromListSecretsCache(urn, client, page, pageSize);
-                setRemovedUrns(newRemovedUrns);
-                setTimeout(function () {
-                    refetch?.();
-                }, 3000);
             })
             .catch((e: unknown) => {
                 message.destroy();
@@ -99,14 +94,22 @@ export const SecretsList = () => {
                 },
             },
         })
-            .then(() => {
+            .then((res) => {
                 message.success({
                     content: `Successfully created Secret!`,
                     duration: 3,
                 });
                 resetBuilderState();
                 setIsCreatingSecret(false);
-                setTimeout(() => refetch(), 3000);
+                addSecretToListSecretsCache(
+                    {
+                        urn: res.data?.createSecret || '',
+                        name: state.name,
+                        description: state.description,
+                    },
+                    client,
+                    pageSize,
+                );
             })
             .catch((e) => {
                 message.destroy();
@@ -160,7 +163,7 @@ export const SecretsList = () => {
         },
     ];
 
-    const tableData = filteredSecrets?.map((secret) => ({
+    const tableData = secrets?.map((secret) => ({
         urn: secret.urn,
         name: secret.name,
         description: secret.description,

datahub-web-react/src/app/ingest/secret/cacheUtils.ts

@@ -0,0 +1,70 @@
+import { ListSecretsDocument, ListSecretsQuery } from '../../../graphql/ingestion.generated';
+
+export const removeSecretFromListSecretsCache = (urn, client, page, pageSize) => {
+    const currData: ListSecretsQuery | null = client.readQuery({
+        query: ListSecretsDocument,
+        variables: {
+            input: {
+                start: (page - 1) * pageSize,
+                count: pageSize,
+            },
+        },
+    });
+
+    const newSecrets = [...(currData?.listSecrets?.secrets || []).filter((secret) => secret.urn !== urn)];
+
+    client.writeQuery({
+        query: ListSecretsDocument,
+        variables: {
+            input: {
+                start: (page - 1) * pageSize,
+                count: pageSize,
+            },
+        },
+        data: {
+            listSecrets: {
+                start: currData?.listSecrets?.start || 0,
+                count: (currData?.listSecrets?.count || 1) - 1,
+                total: (currData?.listSecrets?.total || 1) - 1,
+                secrets: newSecrets,
+            },
+        },
+    });
+};
+
+export const addSecretToListSecretsCache = (secret, client, pageSize) => {
+    const currData: ListSecretsQuery | null = client.readQuery({
+        query: ListSecretsDocument,
+        variables: {
+            input: {
+                start: 0,
+                count: pageSize,
+            },
+        },
+    });
+
+    const newSecrets = [secret, ...(currData?.listSecrets?.secrets || [])];
+
+    client.writeQuery({
+        query: ListSecretsDocument,
+        variables: {
+            input: {
+                start: 0,
+                count: pageSize,
+            },
+        },
+        data: {
+            listSecrets: {
+                start: currData?.listSecrets?.start || 0,
+                count: (currData?.listSecrets?.count || 1) + 1,
+                total: (currData?.listSecrets?.total || 1) + 1,
+                secrets: newSecrets,
+            },
+        },
+    });
+};
+
+export const clearSecretListCache = (client) => {
+    // Remove any caching of 'listSecrets'
+    client.cache.evict({ id: 'ROOT_QUERY', fieldName: 'listSecrets' });
+};

datahub-web-react/src/app/ingest/source/builder/RecipeForm/SecretField/SecretField.tsx

@@ -1,10 +1,12 @@
 import React, { ReactNode } from 'react';
 import { AutoComplete, Divider, Form } from 'antd';
+import { useApolloClient } from '@apollo/client';
 import styled from 'styled-components/macro';
 import { Secret } from '../../../../../../types.generated';
 import CreateSecretButton from './CreateSecretButton';
 import { RecipeField } from '../common';
 import { ANTD_GRAY } from '../../../../../entity/shared/constants';
+import { clearSecretListCache } from '../../../../secret/cacheUtils';
 
 const StyledDivider = styled(Divider)`
     margin: 0;
@@ -86,6 +88,7 @@ const encodeSecret = (secretName: string) => {
 
 function SecretField({ field, secrets, removeMargin, updateFormValue, refetchSecrets }: SecretFieldProps) {
     const options = secrets.map((secret) => ({ value: encodeSecret(secret.name), label: secret.name }));
+    const apolloClient = useApolloClient();
 
     return (
         <StyledFormItem
@@ -108,7 +111,10 @@ function SecretField({ field, secrets, removeMargin, updateFormValue, refetchSec
                             {menu}
                             <StyledDivider />
                             <CreateSecretButton
-                                onSubmit={(state) => updateFormValue(field.name, encodeSecret(state.name as string))}
+                                onSubmit={(state) => {
+                                    updateFormValue(field.name, encodeSecret(state.name as string));
+                                    setTimeout(() => clearSecretListCache(apolloClient), 3000);
+                                }}
                                 refetchSecrets={refetchSecrets}
                             />
                         </>

metadata-models/src/main/pegasus/com/linkedin/secret/DataHubSecretValue.pdl

@@ -1,5 +1,7 @@
 namespace com.linkedin.secret
 
+import com.linkedin.common.AuditStamp
+
 /**
  * The value of a DataHub Secret
  */
@@ -24,4 +26,15 @@ record DataHubSecretValue {
    * Description of the secret
    */
   description: optional string
+
+  /**
+   * Created Audit stamp
+   */
+   @Searchable = {
+     "/time": {
+      "fieldName": "createdTime",
+      "fieldType": "DATETIME"
+     }
+   }
+   created: optional AuditStamp
 }