From b0091f727159ada94e466f83e6a003b05bea2451 Mon Sep 17 00:00:00 2001
From: Pere Miquel Brull <peremiquelbrull@gmail.com>
Date: Tue, 3 Jun 2025 12:04:24 +0200
Subject: [PATCH] FIX #1608 - Fix bot being used in pipelines and workflows
 (#21446)

Co-authored-by: Mohit Yadav <105265192+mohityadav766@users.noreply.github.com>
---
 .../automations/WorkflowResource.java         | 10 ++--
 .../util/OpenMetadataConnectionBuilder.java   |  3 --
 .../OpenMetadataConnectionBuilderTest.java    | 54 +++++++++++++++++++
 3 files changed, 60 insertions(+), 7 deletions(-)
 create mode 100644 openmetadata-service/src/test/java/org/openmetadata/service/util/OpenMetadataConnectionBuilderTest.java

diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/resources/automations/WorkflowResource.java b/openmetadata-service/src/main/java/org/openmetadata/service/resources/automations/WorkflowResource.java
index b8c86628f9d..e5029ca5553 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/resources/automations/WorkflowResource.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/resources/automations/WorkflowResource.java
@@ -85,7 +85,7 @@ public class WorkflowResource extends EntityResource<Workflow, WorkflowRepositor
   static final String FIELDS = "owners";
   private WorkflowMapper mapper;
   private PipelineServiceClientInterface pipelineServiceClient;
-  private OpenMetadataConnectionBuilder openMetadataConnectionBuilder;
+  private OpenMetadataApplicationConfig openMetadataApplicationConfig;
 
   public WorkflowResource(Authorizer authorizer, Limits limits) {
     super(Entity.WORKFLOW, authorizer, limits);
@@ -93,11 +93,11 @@ public class WorkflowResource extends EntityResource<Workflow, WorkflowRepositor
 
   @Override
   public void initialize(OpenMetadataApplicationConfig config) {
+    this.openMetadataApplicationConfig = config;
     this.mapper = new WorkflowMapper();
     this.pipelineServiceClient =
         PipelineServiceClientFactory.createPipelineServiceClient(
             config.getPipelineServiceClientConfiguration());
-    openMetadataConnectionBuilder = new OpenMetadataConnectionBuilder(config);
   }
 
   public static class WorkflowList extends ResultList<Workflow> {
@@ -359,7 +359,8 @@ public class WorkflowResource extends EntityResource<Workflow, WorkflowRepositor
       @Context SecurityContext securityContext) {
     EntityUtil.Fields fields = getFields(FIELD_OWNERS);
     Workflow workflow = repository.get(uriInfo, id, fields);
-    workflow.setOpenMetadataServerConnection(openMetadataConnectionBuilder.build());
+    workflow.setOpenMetadataServerConnection(
+        new OpenMetadataConnectionBuilder(openMetadataApplicationConfig).build());
     /*
      We will send the encrypted Workflow to the Pipeline Service Client
      It will be fetched from the API from there, since we are
@@ -594,7 +595,8 @@ public class WorkflowResource extends EntityResource<Workflow, WorkflowRepositor
       return workflowConverted;
     }
     Workflow workflowDecrypted = secretsManager.decryptWorkflow(workflow);
-    OpenMetadataConnection openMetadataServerConnection = openMetadataConnectionBuilder.build();
+    OpenMetadataConnection openMetadataServerConnection =
+        new OpenMetadataConnectionBuilder(openMetadataApplicationConfig).build();
     workflowDecrypted.setOpenMetadataServerConnection(
         secretsManager.encryptOpenMetadataConnection(openMetadataServerConnection, false));
     if (authorizer.shouldMaskPasswords(securityContext)) {
diff --git a/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataConnectionBuilder.java b/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataConnectionBuilder.java
index 5bf67f96e88..b78b0377f5f 100644
--- a/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataConnectionBuilder.java
+++ b/openmetadata-service/src/main/java/org/openmetadata/service/util/OpenMetadataConnectionBuilder.java
@@ -166,9 +166,6 @@ public class OpenMetadataConnectionBuilder {
   }
 
   public OpenMetadataConnection build() {
-    // Initialize the bot user while building to update any
-    // changes done on the bot like updating jwt token
-    initializeBotUser(Entity.INGESTION_BOT_NAME);
     return new OpenMetadataConnection()
         .withAuthProvider(authProvider)
         .withHostPort(openMetadataURL)
diff --git a/openmetadata-service/src/test/java/org/openmetadata/service/util/OpenMetadataConnectionBuilderTest.java b/openmetadata-service/src/test/java/org/openmetadata/service/util/OpenMetadataConnectionBuilderTest.java
new file mode 100644
index 00000000000..a100e043f13
--- /dev/null
+++ b/openmetadata-service/src/test/java/org/openmetadata/service/util/OpenMetadataConnectionBuilderTest.java
@@ -0,0 +1,54 @@
+package org.openmetadata.service.util;
+
+import com.auth0.jwt.JWT;
+import com.auth0.jwt.interfaces.DecodedJWT;
+import lombok.extern.slf4j.Slf4j;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.Test;
+import org.openmetadata.schema.api.configuration.pipelineServiceClient.PipelineServiceClientConfiguration;
+import org.openmetadata.schema.security.secrets.SecretsManagerClientLoader;
+import org.openmetadata.schema.security.secrets.SecretsManagerConfiguration;
+import org.openmetadata.schema.security.secrets.SecretsManagerProvider;
+import org.openmetadata.schema.security.ssl.VerifySSL;
+import org.openmetadata.schema.services.connections.metadata.OpenMetadataConnection;
+import org.openmetadata.service.OpenMetadataApplicationConfig;
+import org.openmetadata.service.OpenMetadataApplicationTest;
+import org.openmetadata.service.secrets.SecretsManagerFactory;
+
+@Slf4j
+public class OpenMetadataConnectionBuilderTest extends OpenMetadataApplicationTest {
+
+  private static SecretsManagerConfiguration config;
+  static final String CLUSTER_NAME = "test";
+
+  @BeforeAll
+  static void setUp() {
+    config = new SecretsManagerConfiguration();
+    config.setSecretsManager(SecretsManagerProvider.DB);
+    SecretsManagerFactory.createSecretsManager(config, CLUSTER_NAME);
+  }
+
+  @Test
+  void testOpenMetadataConnectionBuilder() {
+
+    OpenMetadataApplicationConfig openMetadataApplicationConfig =
+        new OpenMetadataApplicationConfig();
+    openMetadataApplicationConfig.setClusterName(CLUSTER_NAME);
+    openMetadataApplicationConfig.setPipelineServiceClientConfiguration(
+        new PipelineServiceClientConfiguration()
+            .withMetadataApiEndpoint("http://localhost:8585/api")
+            .withVerifySSL(VerifySSL.NO_SSL)
+            .withSecretsManagerLoader(SecretsManagerClientLoader.ENV));
+
+    String botName =
+        "autoClassification-bot"; // Whichever bot other than the ingestion-bot, which is the
+    // default
+    OpenMetadataConnection openMetadataServerConnection =
+        new OpenMetadataConnectionBuilder(openMetadataApplicationConfig, botName).build();
+
+    // The OM Connection passes the right JWT based on the incoming bot
+    DecodedJWT jwt = JWT.decode(openMetadataServerConnection.getSecurityConfig().getJwtToken());
+    Assertions.assertEquals("autoclassification-bot", jwt.getClaim("sub").asString());
+  }
+}