addressed PR comments

Author: Aditya Pratap Singh

gobblin-service/src/main/java/org/apache/gobblin/service/modules/flowgraph/Dag.java

@@ -261,17 +261,11 @@ public static class DagNode<T> {
     private T value;
     //List of parent Nodes that are dependencies of this Node.
     private List<DagNode<T>> parentNodes;
-    @Setter
-    private boolean isFailedDag;
 
     //Constructor
     public DagNode(T value) {
       this.value = value;
     }
-    public DagNode(T value,boolean isFailedDag) {
-      this.value = value;
-      this.isFailedDag = isFailedDag;
-    }
 
 
     public void addParentNode(DagNode<T> node) {

gobblin-service/src/main/java/org/apache/gobblin/service/modules/orchestration/MysqlDagStateStoreWithDagNodes.java

@@ -77,18 +77,20 @@ public class MysqlDagStateStoreWithDagNodes implements DagStateStoreWithDagNodes
   protected final GsonSerDe<List<JobExecutionPlan>> serDe;
   private final JobExecutionPlanDagFactory jobExecPlanDagFactory;
 
-  protected static final String CREATE_TABLE_STATEMENT = "CREATE TABLE IF NOT EXISTS %s ("
-      + "dag_node_id VARCHAR(" + ServiceConfigKeys.MAX_DAG_NODE_ID_LENGTH
-      + ") CHARACTER SET latin1 COLLATE latin1_bin NOT NULL, " + "parent_dag_id VARCHAR("
-      + ServiceConfigKeys.MAX_DAG_ID_LENGTH + ") NOT NULL, " + "dag_node JSON NOT NULL, "
-      + "modified_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, "
-      + "is_failed_dag INT NOT NULL DEFAULT 0, " + "PRIMARY KEY (dag_node_id), "
-      + "UNIQUE INDEX dag_node_index (dag_node_id), " + "INDEX dag_index (parent_dag_id))";
-
-  protected static final String INSERT_STATEMENT = "INSERT INTO %s (dag_node_id, parent_dag_id, dag_node, is_failed_dag) "
-      + "VALUES (?, ?, ?, ?) AS new ON DUPLICATE KEY UPDATE dag_node = new.dag_node, is_failed_dag = new.is_failed_dag";
-  protected static final String GET_DAG_NODES_STATEMENT = "SELECT dag_node,is_failed_dag FROM %s WHERE parent_dag_id = ?";
-  protected static final String GET_DAG_NODE_STATEMENT = "SELECT dag_node,is_failed_dag FROM %s WHERE dag_node_id = ?";
+  protected static final String CREATE_TABLE_STATEMENT =
+      "CREATE TABLE IF NOT EXISTS %s (" + "dag_node_id VARCHAR(" + ServiceConfigKeys.MAX_DAG_NODE_ID_LENGTH
+          + ") CHARACTER SET latin1 COLLATE latin1_bin NOT NULL, " + "parent_dag_id VARCHAR("
+          + ServiceConfigKeys.MAX_DAG_ID_LENGTH + ") NOT NULL, " + "dag_node JSON NOT NULL, "
+          + "modified_time timestamp NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, "
+          + "is_failed_dag TINYINT(1) DEFAULT 0, " + "PRIMARY KEY (dag_node_id), "
+          + "UNIQUE INDEX dag_node_index (dag_node_id), " + "INDEX dag_index (parent_dag_id))";
+
+  protected static final String INSERT_STATEMENT =
+      "INSERT INTO %s (dag_node_id, parent_dag_id, dag_node, is_failed_dag) "
+          + "VALUES (?, ?, ?, ?) AS new ON DUPLICATE KEY UPDATE dag_node = new.dag_node, is_failed_dag = new.is_failed_dag";
+  protected static final String GET_DAG_NODES_STATEMENT =
+      "SELECT dag_node FROM %s WHERE parent_dag_id = ?";
+  protected static final String GET_DAG_NODE_STATEMENT = "SELECT dag_node FROM %s WHERE dag_node_id = ?";
   protected static final String DELETE_DAG_STATEMENT = "DELETE FROM %s WHERE parent_dag_id = ?";
   private final ContextAwareCounter totalDagCount;
 
@@ -103,7 +105,8 @@ public MysqlDagStateStoreWithDagNodes(Config config, Map<URI, TopologySpec> topo
     DataSource dataSource = MysqlDataSourceFactory.get(config, SharedResourcesBrokerFactory.getImplicitBroker());
 
     try (Connection connection = dataSource.getConnection();
-        PreparedStatement createStatement = connection.prepareStatement(String.format(CREATE_TABLE_STATEMENT, tableName))) {
+        PreparedStatement createStatement = connection.prepareStatement(
+            String.format(CREATE_TABLE_STATEMENT, tableName))) {
       createStatement.executeUpdate();
       connection.commit();
     } catch (SQLException e) {
@@ -165,7 +168,8 @@ public void cleanUp(String dagId) throws IOException {
   @Override
   public List<Dag<JobExecutionPlan>> getDags() throws IOException {
     throw new NotSupportedException(getClass().getSimpleName() + " does not need this legacy API that originated with "
-        + "the DagManager that is replaced by DagProcessingEngine"); }
+        + "the DagManager that is replaced by DagProcessingEngine");
+  }
 
   @Override
   public Dag<JobExecutionPlan> getDag(DagManager.DagId dagId) throws IOException {
@@ -189,14 +193,7 @@ private Dag<JobExecutionPlan> convertDagNodesIntoDag(Set<Dag.DagNode<JobExecutio
     if (dagNodes.isEmpty()) {
       return null;
     }
-    Dag<JobExecutionPlan> dag = jobExecPlanDagFactory.createDag(dagNodes.stream().map(Dag.DagNode::getValue).collect(Collectors.toList()));
-
-    // if any node of the dag is failed it means that the dag has been marked as failed, update the is_failed_dag field of the dag and it's nodes as true
-    if (dag.getNodes().stream().anyMatch(Dag.DagNode::isFailedDag)) {
-      dag.setFailedDag(true);
-      dag.getNodes().forEach(node -> node.setFailedDag(true));
-    }
-    return dag;
+    return jobExecPlanDagFactory.createDag(dagNodes.stream().map(Dag.DagNode::getValue).collect(Collectors.toList()));
   }
 
   @Override
@@ -224,7 +221,7 @@ public Set<Dag.DagNode<JobExecutionPlan>> getDagNodes(DagManager.DagId parentDag
           HashSet<Dag.DagNode<JobExecutionPlan>> dagNodes = new HashSet<>();
           try (ResultSet rs = getStatement.executeQuery()) {
             while (rs.next()) {
-              dagNodes.add(new Dag.DagNode<>(this.serDe.deserialize(rs.getString(1)).get(0), rs.getBoolean(2)));
+              dagNodes.add(new Dag.DagNode<>(this.serDe.deserialize(rs.getString(1)).get(0)));
             }
             return dagNodes;
           } catch (SQLException e) {
@@ -239,7 +236,7 @@ public Optional<Dag.DagNode<JobExecutionPlan>> getDagNode(DagNodeId dagNodeId) t
       getStatement.setString(1, dagNodeId.toString());
       try (ResultSet rs = getStatement.executeQuery()) {
         if (rs.next()) {
-          return Optional.of(new Dag.DagNode<>(this.serDe.deserialize(rs.getString(1)).get(0), rs.getBoolean(2)));
+          return Optional.of(new Dag.DagNode<>(this.serDe.deserialize(rs.getString(1)).get(0)));
         }
         return Optional.empty();
       } catch (SQLException e) {

gobblin-service/src/test/java/org/apache/gobblin/service/modules/orchestration/MysqlDagStateStoreWithDagNodesTest.java

@@ -17,10 +17,18 @@
 
 package org.apache.gobblin.service.modules.orchestration;
 
+import java.io.IOException;
 import java.net.URI;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
 
+import lombok.extern.slf4j.Slf4j;
+
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -36,17 +44,22 @@
 import org.apache.gobblin.service.ExecutionStatus;
 import org.apache.gobblin.service.modules.flowgraph.Dag;
 import org.apache.gobblin.service.modules.spec.JobExecutionPlan;
-
+import org.apache.gobblin.broker.SharedResourcesBrokerFactory;
+import org.apache.gobblin.metastore.MysqlDataSourceFactory;
+import org.apache.gobblin.util.DBStatementExecutor;
 
 /**
  * Mainly testing functionalities related to DagStateStore but not Mysql-related components.
  */
+@Slf4j
 public class MysqlDagStateStoreWithDagNodesTest {
 
   private DagStateStore dagStateStore;
-
   private static final String TEST_USER = "testUser";
   private static ITestMetastoreDatabase testDb;
+  private DBStatementExecutor dbStatementExecutor;
+  private static final String GET_DAG_NODES_STATEMENT = "SELECT dag_node, is_failed_dag FROM %s WHERE parent_dag_id = ?";
+  private static final String tableName = "dag_node_state_store";
 
   @BeforeClass
   public void setUp() throws Exception {
@@ -63,6 +76,8 @@ public void setUp() throws Exception {
     URI specExecURI = new URI(specExecInstance);
     topologySpecMap.put(specExecURI, topologySpec);
     this.dagStateStore = new MysqlDagStateStoreWithDagNodes(configBuilder.build(), topologySpecMap);
+    dbStatementExecutor = new DBStatementExecutor(
+        MysqlDataSourceFactory.get(configBuilder.build(), SharedResourcesBrokerFactory.getImplicitBroker()), log);
   }
 
   @AfterClass(alwaysRun = true)
@@ -74,7 +89,7 @@ public void tearDown() throws Exception {
   }
 
   @Test
-  public void testAddGetAndDeleteDag() throws Exception{
+  public void testAddGetAndDeleteDag() throws Exception {
     Dag<JobExecutionPlan> originalDag1 = DagTestUtils.buildDag("random_1", 123L);
     Dag<JobExecutionPlan> originalDag2 = DagTestUtils.buildDag("random_2", 456L);
     DagManager.DagId dagId1 = DagManagerUtils.generateDagId(originalDag1);
@@ -140,26 +155,48 @@ public void testAddGetAndDeleteDag() throws Exception{
 
   @Test
   public void testMarkDagAsFailed() throws Exception {
-    //Set up initial conditions
+    // Set up initial conditions
     Dag<JobExecutionPlan> dag = DagTestUtils.buildDag("test_dag", 789L);
     DagManager.DagId dagId = DagManagerUtils.generateDagId(dag);
 
     this.dagStateStore.writeCheckpoint(dag);
-    //Check Initial State
-    for (Dag.DagNode<JobExecutionPlan> node : dag.getNodes()) {
-      Assert.assertFalse(node.isFailedDag());
+
+    // Fetch all initial states into a list
+    List<Boolean> initialStates = fetchDagNodeStates(dagId.toString());
+
+    // Check Initial State
+    for (Boolean state : initialStates) {
+      Assert.assertFalse(state);
     }
+    // Set the DAG as failed
     dag.setFailedDag(true);
     this.dagStateStore.writeCheckpoint(dag);
 
-    Dag<JobExecutionPlan> updatedDag = this.dagStateStore.getDag(dagId);
-    for (Dag.DagNode<JobExecutionPlan> node : updatedDag.getNodes()) {
-      Assert.assertTrue(node.isFailedDag());
-    }
+    // Fetch all states after marking the DAG as failed
+    List<Boolean> failedStates = fetchDagNodeStates(dagId.toString());
 
-    // Cleanup
+    // Check if all states are now true (indicating failure)
+    for (Boolean state : failedStates) {
+      Assert.assertTrue(state);
+    }
     dagStateStore.cleanUp(dagId);
     Assert.assertNull(this.dagStateStore.getDag(dagId));
   }
 
+  private List<Boolean> fetchDagNodeStates(String dagId) throws IOException {
+    List<Boolean> states = new ArrayList<>();
+    dbStatementExecutor.withPreparedStatement(String.format(GET_DAG_NODES_STATEMENT, tableName), getStatement -> {
+      getStatement.setString(1, dagId.toString());
+      HashSet<Dag.DagNode<JobExecutionPlan>> dagNodes = new HashSet<>();
+      try (ResultSet rs = getStatement.executeQuery()) {
+        while (rs.next()) {
+          states.add(rs.getBoolean(2));
+        }
+        return dagNodes;
+      } catch (SQLException e) {
+        throw new IOException(String.format("Failure get dag nodes for dag %s", dagId), e);
+      }
+    }, true);
+    return states;
+  }
 }