Improve robustness of JdbcExecutor by safely handling multiple consumers

wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java

@@ -47,7 +47,6 @@
 import org.apache.wayang.core.platform.PartialExecution;
 import org.apache.wayang.core.platform.Platform;
 import org.apache.wayang.core.platform.lineage.ExecutionLineageNode;
-import org.apache.wayang.core.util.WayangCollections;
 import org.apache.wayang.jdbc.channels.SqlQueryChannel;
 import org.apache.wayang.jdbc.compiler.FunctionCompiler;
 import org.apache.wayang.jdbc.operators.JdbcExecutionOperator;
@@ -338,16 +337,20 @@ private Collection<ExecutionLineageNode> createExecutionLineageNodes(
      * unless it is not comprising a {@link JdbcExecutionOperator} and/or not in the
      * given {@link ExecutionStage}.
      *
-     * @param task  whose follow-up {@link ExecutionTask} is requested; should have
-     *              a single follower
+     * @param task  whose follow-up {@link ExecutionTask} is requested
      * @param stage in which the follow-up {@link ExecutionTask} should be
      * @return the said follow-up {@link ExecutionTask} or {@code null} if none
      */
     private static ExecutionTask findJdbcExecutionOperatorTaskInStage(final ExecutionTask task,
             final ExecutionStage stage) {
-        assert task.getNumOuputChannels() == 1;
+        if (task.getNumOuputChannels() != 1) {
+            return null;
+        }
         final Channel outputChannel = task.getOutputChannel(0);
-        final ExecutionTask consumer = WayangCollections.getSingle(outputChannel.getConsumers());
+        if (outputChannel == null || outputChannel.getConsumers().size() != 1) {
+            return null;
+        }
+        final ExecutionTask consumer = outputChannel.getConsumers().iterator().next();
         return consumer.getStage() == stage && consumer.getOperator() instanceof JdbcExecutionOperator 
                 ? consumer
                 : null;

wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java

@@ -40,6 +40,7 @@
 
 import java.sql.SQLException;
 import java.util.Collections;
+import java.util.Set;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.mockito.Mockito.mock;
@@ -244,4 +245,62 @@ void testExecuteWithProjectionAndFilters() throws SQLException {
                 sqlQueryChannelInstance.getSqlQuery()
         );
     }
+
+    @Test
+    void testExecuteWithNoTerminalConsumers() {
+        assertQueryWithTerminalConsumers(0);
+    }
+
+    @Test
+    void testExecuteWithMultipleTerminalConsumers() {
+        assertQueryWithTerminalConsumers(2);
+    }
+
+    private void assertQueryWithTerminalConsumers(int consumerCount) {
+        Configuration configuration = new Configuration();
+        Job job = mock(Job.class);
+        when(job.getConfiguration()).thenReturn(configuration);
+        CrossPlatformExecutor executionState = new CrossPlatformExecutor(job, new NoInstrumentationStrategy());
+        when(job.getCrossPlatformExecutor()).thenReturn(executionState);
+        SqlQueryChannel.Descriptor descriptor = HsqldbPlatform.getInstance().getSqlQueryChannelDescriptor();
+        ExecutionStage stage = mock(ExecutionStage.class);
+
+        JdbcTableSource source = new HsqldbTableSource("customer");
+        ExecutionTask sourceTask = new ExecutionTask(source);
+        sourceTask.setStage(stage);
+        sourceTask.setOutputChannel(0, new SqlQueryChannel(descriptor, source.getOutput(0)));
+
+        JdbcFilterOperator filter = new HsqldbFilterOperator(new PredicateDescriptor<Record>(
+                record -> true, Record.class).withSqlImplementation("age >= 18"));
+        ExecutionTask filterTask = new ExecutionTask(filter);
+        filterTask.setStage(stage);
+        sourceTask.getOutputChannel(0).addConsumer(filterTask, 0);
+        // A reusable channel lets the stage boundary have multiple downstream consumers.
+        SqlQueryChannel output = new SqlQueryChannel(descriptor, filter.getOutput(0)) {
+            @Override
+            public boolean isReusable() {
+                return true;
+            }
+        };
+        filterTask.setOutputChannel(0, output);
+        when(stage.getStartTasks()).thenReturn(Collections.singleton(sourceTask));
+        when(stage.getTerminalTasks()).thenReturn(Collections.singleton(filterTask));
+        when(stage.getAllTasks()).thenReturn(Set.of(sourceTask, filterTask));
+
+        ExecutionStage nextStage = mock(ExecutionStage.class);
+        for (int i = 0; i < consumerCount; i++) {
+            ExecutionTask consumer = new ExecutionTask(new SqlToStreamOperator(HsqldbPlatform.getInstance()));
+            consumer.setStage(nextStage);
+            output.addConsumer(consumer, 0);
+        }
+
+        JdbcExecutor executor = new JdbcExecutor(HsqldbPlatform.getInstance(), job);
+        try {
+            executor.execute(stage, new DefaultOptimizationContext(job), executionState);
+            SqlQueryChannel.Instance result = (SqlQueryChannel.Instance) executionState.getChannelInstance(output);
+            assertEquals("SELECT * FROM customer WHERE age >= 18", result.getSqlQuery());
+        } finally {
+            executor.dispose();
+        }
+    }
 }