feat (resolver): Add new endpoint for editing secrets (#9665)

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/GmsGraphQLEngine.java

@@ -170,6 +170,7 @@
 import com.linkedin.datahub.graphql.resolvers.ingest.secret.DeleteSecretResolver;
 import com.linkedin.datahub.graphql.resolvers.ingest.secret.GetSecretValuesResolver;
 import com.linkedin.datahub.graphql.resolvers.ingest.secret.ListSecretsResolver;
+import com.linkedin.datahub.graphql.resolvers.ingest.secret.UpdateSecretResolver;
 import com.linkedin.datahub.graphql.resolvers.ingest.source.DeleteIngestionSourceResolver;
 import com.linkedin.datahub.graphql.resolvers.ingest.source.GetIngestionSourceResolver;
 import com.linkedin.datahub.graphql.resolvers.ingest.source.ListIngestionSourcesResolver;
@@ -1086,6 +1087,8 @@ private void configureMutationResolvers(final RuntimeWiring.Builder builder) {
                 .dataFetcher(
                     "createSecret", new CreateSecretResolver(this.entityClient, this.secretService))
                 .dataFetcher("deleteSecret", new DeleteSecretResolver(this.entityClient))
+                .dataFetcher(
+                    "updateSecret", new UpdateSecretResolver(this.entityClient, this.secretService))
                 .dataFetcher(
                     "createAccessToken", new CreateAccessTokenResolver(this.statefulTokenService))
                 .dataFetcher(

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

@@ -6,11 +6,11 @@
 
 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;
 import com.linkedin.datahub.graphql.generated.CreateSecretInput;
 import com.linkedin.datahub.graphql.resolvers.ingest.IngestionAuthUtils;
+import com.linkedin.datahub.graphql.types.ingest.secret.mapper.DataHubSecretValueMapper;
 import com.linkedin.entity.client.EntityClient;
 import com.linkedin.metadata.key.DataHubSecretKey;
 import com.linkedin.metadata.secret.SecretService;
@@ -58,14 +58,15 @@ public CompletableFuture<String> get(final DataFetchingEnvironment environment)
               }
 
               // Create the secret value.
-              final DataHubSecretValue value = new DataHubSecretValue();
-              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()));
+              final DataHubSecretValue value =
+                  DataHubSecretValueMapper.map(
+                      null,
+                      input.getName(),
+                      _secretService.encrypt(input.getValue()),
+                      input.getDescription(),
+                      new AuditStamp()
+                          .setActor(UrnUtils.getUrn(context.getActorUrn()))
+                          .setTime(System.currentTimeMillis()));
 
               final MetadataChangeProposal proposal =
                   buildMetadataChangeProposalWithKey(

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

@@ -0,0 +1,82 @@
+package com.linkedin.datahub.graphql.resolvers.ingest.secret;
+
+import static com.linkedin.datahub.graphql.resolvers.ResolverUtils.bindArgument;
+import static com.linkedin.datahub.graphql.resolvers.mutate.MutationUtils.buildMetadataChangeProposalWithUrn;
+import static com.linkedin.metadata.Constants.SECRET_VALUE_ASPECT_NAME;
+
+import com.linkedin.common.urn.Urn;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.exception.AuthorizationException;
+import com.linkedin.datahub.graphql.generated.UpdateSecretInput;
+import com.linkedin.datahub.graphql.resolvers.ingest.IngestionAuthUtils;
+import com.linkedin.datahub.graphql.types.ingest.secret.mapper.DataHubSecretValueMapper;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.secret.SecretService;
+import com.linkedin.mxe.MetadataChangeProposal;
+import com.linkedin.secret.DataHubSecretValue;
+import graphql.schema.DataFetcher;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.Set;
+import java.util.concurrent.CompletableFuture;
+import lombok.RequiredArgsConstructor;
+import lombok.extern.slf4j.Slf4j;
+
+/**
+ * Creates an encrypted DataHub secret. Uses AES symmetric encryption / decryption. Requires the
+ * MANAGE_SECRETS privilege.
+ */
+@Slf4j
+@RequiredArgsConstructor
+public class UpdateSecretResolver implements DataFetcher<CompletableFuture<String>> {
+  private final EntityClient entityClient;
+  private final SecretService secretService;
+
+  @Override
+  public CompletableFuture<String> get(final DataFetchingEnvironment environment) throws Exception {
+    final QueryContext context = environment.getContext();
+    final UpdateSecretInput input =
+        bindArgument(environment.getArgument("input"), UpdateSecretInput.class);
+    final Urn secretUrn = Urn.createFromString(input.getUrn());
+    return CompletableFuture.supplyAsync(
+        () -> {
+          if (IngestionAuthUtils.canManageSecrets(context)) {
+
+            try {
+              EntityResponse response =
+                  entityClient.getV2(
+                      secretUrn.getEntityType(),
+                      secretUrn,
+                      Set.of(SECRET_VALUE_ASPECT_NAME),
+                      context.getAuthentication());
+              if (!entityClient.exists(secretUrn, context.getAuthentication())
+                  || response == null) {
+                throw new IllegalArgumentException(
+                    String.format("Secret for urn %s doesn't exists!", secretUrn));
+              }
+
+              DataHubSecretValue updatedVal =
+                  DataHubSecretValueMapper.map(
+                      response,
+                      input.getName(),
+                      secretService.encrypt(input.getValue()),
+                      input.getDescription(),
+                      null);
+
+              final MetadataChangeProposal proposal =
+                  buildMetadataChangeProposalWithUrn(
+                      secretUrn, SECRET_VALUE_ASPECT_NAME, updatedVal);
+              return entityClient.ingestProposal(proposal, context.getAuthentication(), false);
+            } catch (Exception e) {
+              throw new RuntimeException(
+                  String.format(
+                      "Failed to update a secret with urn %s and name %s",
+                      secretUrn, input.getName()),
+                  e);
+            }
+          }
+          throw new AuthorizationException(
+              "Unauthorized to perform this action. Please contact your DataHub administrator.");
+        });
+  }
+}

datahub-graphql-core/src/main/java/com/linkedin/datahub/graphql/types/ingest/secret/mapper/DataHubSecretValueMapper.java

@@ -0,0 +1,55 @@
+package com.linkedin.datahub.graphql.types.ingest.secret.mapper;
+
+import static com.linkedin.metadata.Constants.SECRET_VALUE_ASPECT_NAME;
+
+import com.linkedin.common.AuditStamp;
+import com.linkedin.data.template.RecordTemplate;
+import com.linkedin.data.template.SetMode;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.secret.DataHubSecretValue;
+import java.util.Objects;
+import javax.annotation.Nonnull;
+
+/**
+ * Maps Pegasus {@link RecordTemplate} objects to objects conforming to the GQL schema.
+ *
+ * <p>To be replaced by auto-generated mappers implementations
+ */
+public class DataHubSecretValueMapper {
+
+  public static final DataHubSecretValueMapper INSTANCE = new DataHubSecretValueMapper();
+
+  public static DataHubSecretValue map(
+      EntityResponse fromSecret,
+      @Nonnull final String name,
+      @Nonnull final String value,
+      String description,
+      AuditStamp auditStamp) {
+    return INSTANCE.apply(fromSecret, name, value, description, auditStamp);
+  }
+
+  public DataHubSecretValue apply(
+      EntityResponse existingSecret,
+      @Nonnull final String name,
+      @Nonnull final String value,
+      String description,
+      AuditStamp auditStamp) {
+    final DataHubSecretValue result;
+    if (Objects.nonNull(existingSecret)) {
+      result =
+          new DataHubSecretValue(
+              existingSecret.getAspects().get(SECRET_VALUE_ASPECT_NAME).getValue().data());
+    } else {
+      result = new DataHubSecretValue();
+    }
+
+    result.setName(name);
+    result.setValue(value);
+    result.setDescription(description, SetMode.IGNORE_NULL);
+    if (Objects.nonNull(auditStamp)) {
+      result.setCreated(auditStamp);
+    }
+
+    return result;
+  }
+}

datahub-graphql-core/src/main/resources/ingestion.graphql

@@ -36,6 +36,11 @@ extend type Mutation {
   """
   createSecret(input: CreateSecretInput!): String
 
+  """
+  Update a Secret
+  """
+  updateSecret(input: UpdateSecretInput!): String
+
   """
   Delete a Secret
   """
@@ -560,6 +565,31 @@ input CreateSecretInput {
   description: String
 }
 
+"""
+Input arguments for updating a Secret
+"""
+input UpdateSecretInput {
+  """
+  The primary key of the Secret to update
+  """
+  urn: String!
+
+  """
+  The name of the secret for reference in ingestion recipes
+  """
+  name: String!
+
+  """
+  The value of the secret, to be encrypted and stored
+  """
+  value: String!
+
+  """
+  An optional description for the secret
+  """
+  description: String
+}
+
 """
 Input arguments for retrieving the plaintext values of a set of secrets
 """

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

@@ -0,0 +1,98 @@
+package com.linkedin.datahub.graphql.resolvers.ingest.secret;
+
+import static com.linkedin.datahub.graphql.resolvers.ingest.IngestTestUtils.getMockAllowContext;
+import static com.linkedin.datahub.graphql.resolvers.ingest.IngestTestUtils.getMockDenyContext;
+import static com.linkedin.metadata.Constants.SECRET_VALUE_ASPECT_NAME;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyBoolean;
+import static org.mockito.Mockito.when;
+
+import com.datahub.authentication.Authentication;
+import com.linkedin.common.AuditStamp;
+import com.linkedin.common.urn.Urn;
+import com.linkedin.common.urn.UrnUtils;
+import com.linkedin.datahub.graphql.QueryContext;
+import com.linkedin.datahub.graphql.exception.AuthorizationException;
+import com.linkedin.datahub.graphql.generated.UpdateSecretInput;
+import com.linkedin.entity.Aspect;
+import com.linkedin.entity.EntityResponse;
+import com.linkedin.entity.EnvelopedAspect;
+import com.linkedin.entity.EnvelopedAspectMap;
+import com.linkedin.entity.client.EntityClient;
+import com.linkedin.metadata.secret.SecretService;
+import com.linkedin.secret.DataHubSecretValue;
+import graphql.schema.DataFetchingEnvironment;
+import java.util.concurrent.CompletionException;
+import org.mockito.Mockito;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class UpdateSecretResolverTest {
+
+  private static final Urn TEST_URN = UrnUtils.getUrn("urn:li:secret:secret-id");
+
+  private static final UpdateSecretInput TEST_INPUT =
+      new UpdateSecretInput(TEST_URN.toString(), "MY_SECRET", "mysecretvalue", "dummy");
+
+  private DataFetchingEnvironment mockEnv;
+  private EntityClient mockClient;
+  private SecretService mockSecretService;
+  private UpdateSecretResolver resolver;
+
+  @BeforeMethod
+  public void before() {
+    mockClient = Mockito.mock(EntityClient.class);
+    mockSecretService = Mockito.mock(SecretService.class);
+
+    resolver = new UpdateSecretResolver(mockClient, mockSecretService);
+  }
+
+  private DataHubSecretValue createSecretAspect() {
+    DataHubSecretValue secretAspect = new DataHubSecretValue();
+    secretAspect.setValue("encryptedvalue.updated");
+    secretAspect.setName(TEST_INPUT.getName() + ".updated");
+    secretAspect.setDescription(TEST_INPUT.getDescription() + ".updated");
+    secretAspect.setCreated(
+        new AuditStamp().setActor(UrnUtils.getUrn("urn:li:corpuser:test")).setTime(0L));
+    return secretAspect;
+  }
+
+  @Test
+  public void testGetSuccess() throws Exception {
+    // with valid context
+    QueryContext mockContext = getMockAllowContext();
+    mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+    Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT);
+    Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+
+    Mockito.when(mockClient.exists(any(), any())).thenReturn(true);
+    Mockito.when(mockSecretService.encrypt(any())).thenReturn("encrypted_value");
+    final EntityResponse entityResponse = new EntityResponse();
+    final EnvelopedAspectMap aspectMap = new EnvelopedAspectMap();
+    aspectMap.put(
+        SECRET_VALUE_ASPECT_NAME,
+        new EnvelopedAspect().setValue(new Aspect(createSecretAspect().data())));
+    entityResponse.setAspects(aspectMap);
+
+    when(mockClient.getV2(any(), any(), any(), any())).thenReturn(entityResponse);
+
+    // Invoke the resolver
+    resolver.get(mockEnv).join();
+    Mockito.verify(mockClient, Mockito.times(1)).ingestProposal(any(), any(), anyBoolean());
+  }
+
+  @Test(
+      description = "validate if nothing provided throws Exception",
+      expectedExceptions = {AuthorizationException.class, CompletionException.class})
+  public void testGetUnauthorized() throws Exception {
+    // Execute resolver
+    QueryContext mockContext = getMockDenyContext();
+    mockEnv = Mockito.mock(DataFetchingEnvironment.class);
+    Mockito.when(mockEnv.getArgument(Mockito.eq("input"))).thenReturn(TEST_INPUT);
+    Mockito.when(mockEnv.getContext()).thenReturn(mockContext);
+
+    resolver.get(mockEnv).join();
+    Mockito.verify(mockClient, Mockito.times(0))
+        .ingestProposal(any(), any(Authentication.class), anyBoolean());
+  }
+}

datahub-web-react/src/graphql/ingestion.graphql

@@ -145,6 +145,10 @@ mutation createSecret($input: CreateSecretInput!) {
     createSecret(input: $input)
 }
 
+mutation updateSecret($input: UpdateSecretInput!) {
+    updateSecret(input: $input)
+}
+
 mutation deleteSecret($urn: String!) {
     deleteSecret(urn: $urn)
 }

smoke-test/tests/managed-ingestion/managed_ingestion_test.py

@@ -260,6 +260,27 @@ def test_create_list_get_remove_secret(frontend_session):
     # Get new count of secrets
     _ensure_secret_increased(frontend_session, before_count)
 
+    # Update existing secret
+    json_q = {
+        "query": """mutation updateSecret($input: UpdateSecretInput!) {\n
+            updateSecret(input: $input)
+        }""",
+        "variables": {"input": {"urn": secret_urn, "name": "SMOKE_TEST", "value": "mytestvalue.updated"}},
+    }
+
+    response = frontend_session.post(
+        f"{get_frontend_url()}/api/v2/graphql", json=json_q
+    )
+    response.raise_for_status()
+    res_data = response.json()
+
+    assert res_data
+    assert res_data["data"]
+    assert res_data["data"]["updateSecret"] is not None
+    assert "errors" not in res_data
+
+    secret_urn = res_data["data"]["updateSecret"]
+
     # Get the secret value back
     json_q = {
         "query": """query getSecretValues($input: GetSecretValuesInput!) {\n
@@ -285,7 +306,7 @@ def test_create_list_get_remove_secret(frontend_session):
 
     secret_values = res_data["data"]["getSecretValues"]
     secret_value = [x for x in secret_values if x["name"] == "SMOKE_TEST"][0]
-    assert secret_value["value"] == "mytestvalue"
+    assert secret_value["value"] == "mytestvalue.updated"
 
     # Now cleanup and remove the secret
     json_q = {