diff --git a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java index 4b816b2eb..344f89ea8 100644 --- a/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java +++ b/wayang-platforms/wayang-jdbc-template/src/main/java/org/apache/wayang/jdbc/execution/JdbcExecutor.java @@ -371,14 +371,29 @@ private static ExecutionTask findJdbcExecutionOperatorTaskInStage(final Executio final Channel outputChannel = task.getOutputChannel(0); - if (outputChannel.getConsumers().size() != 1) { + final Collection consumers = outputChannel.getConsumers(); + + // No consumer → end of pipeline + if (consumers.isEmpty()) { return null; } - final ExecutionTask consumer = outputChannel.getConsumers().iterator().next(); + // Multiple consumers are currently unsupported + if (consumers.size() != 1) { + throw new WayangException( + String.format("Expected a single consumer for task %s but found %d.", + task, + consumers.size() + ) + ); + } + + final ExecutionTask consumer = consumers.iterator().next(); - return consumer.getStage() == stage && consumer.getOperator() instanceof JdbcExecutionOperator ? consumer - : null; + return consumer.getStage() == stage + && consumer.getOperator() instanceof JdbcExecutionOperator + ? consumer + : null; } private static SqlQueryChannel.Instance instantiateOutboundChannel(final ExecutionTask task, diff --git a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java index 0dfd8b698..0a4f1da95 100644 --- a/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java +++ b/wayang-platforms/wayang-jdbc-template/src/test/java/org/apache/wayang/jdbc/execution/JdbcExecutorTest.java @@ -21,6 +21,7 @@ import org.apache.wayang.basic.data.Record; import org.apache.wayang.core.api.Configuration; import org.apache.wayang.core.api.Job; +import org.apache.wayang.core.api.exception.WayangException; import org.apache.wayang.core.function.PredicateDescriptor; import org.apache.wayang.core.optimizer.DefaultOptimizationContext; import org.apache.wayang.core.plan.executionplan.ExecutionStage; @@ -42,6 +43,7 @@ import java.util.Collections; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.when; @@ -86,6 +88,48 @@ void testExecuteWithPlainTableSource() throws SQLException { ); } + @Test + void testExecuteWithMultipleConsumers() throws SQLException { + Configuration configuration = new Configuration(); + Job job = mock(Job.class); + when(job.getConfiguration()).thenReturn(configuration); + when(job.getCrossPlatformExecutor()) + .thenReturn(new CrossPlatformExecutor(job, new NoInstrumentationStrategy())); + + SqlQueryChannel.Descriptor sqlChannelDescriptor = + HsqldbPlatform.getInstance().getSqlQueryChannelDescriptor(); + + ExecutionStage sqlStage = mock(ExecutionStage.class); + + JdbcTableSource tableSource = new HsqldbTableSource("customer"); + ExecutionTask tableSourceTask = new ExecutionTask(tableSource); + + SqlQueryChannel outputChannel = + new SqlQueryChannel(sqlChannelDescriptor, tableSource.getOutput(0)); + tableSourceTask.setOutputChannel(0, outputChannel); + tableSourceTask.setStage(sqlStage); + + when(sqlStage.getStartTasks()).thenReturn(Collections.singleton(tableSourceTask)); + when(sqlStage.getTerminalTasks()).thenReturn(Collections.singleton(tableSourceTask)); + + ExecutionTask firstConsumer = mock(ExecutionTask.class); + ExecutionTask secondConsumer = mock(ExecutionTask.class); + + outputChannel.getConsumers().add(firstConsumer); + outputChannel.getConsumers().add(secondConsumer); + + JdbcExecutor executor = new JdbcExecutor(HsqldbPlatform.getInstance(), job); + + assertThrows( + WayangException.class, + () -> executor.execute( + sqlStage, + new DefaultOptimizationContext(job), + job.getCrossPlatformExecutor() + ) + ); + } + @Test void testExecuteWithFilter() throws SQLException { Configuration configuration = new Configuration();