Yaml test option to force continuations (#3075)

* Carry over from previous PR: Use NoOpConfig, protect from inadvertant execution

* Add additionalOptions to the yaml execution context, add force-continuations option

* Initial implementation of force continuation execution and aggregated result set

* Bug fixes, create test configuration

* Add the ability to set maxRows:0 which would differentiate "not set" from "set to 0", so we can skip the forceContinuations option if maxRows explicitly set to 0

* Analyze test failures, add @disabled comments

* Style

* PMD

* Revert use of NoOpConfig since retaining the original SupportedVersionConfig is required for assertions

* PR Comments

* PR Comments

* PR Comments

* Style

* PR comments

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/AbstractAggregateResultSet.java

@@ -0,0 +1,202 @@
+/*
+ * AbstractAggregateResultSet.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.apple.foundationdb.relational.yamltests;
+
+import com.apple.foundationdb.relational.api.RelationalArray;
+import com.apple.foundationdb.relational.api.RelationalResultSet;
+import com.apple.foundationdb.relational.api.RelationalResultSetMetaData;
+import com.apple.foundationdb.relational.api.RelationalStruct;
+import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
+
+import java.sql.SQLException;
+
+/**
+ * A result set implementation that aggregates many result sets into one, where each internal result set acts as a row
+ * in the aggregate.
+ * TODO: This is a copy of the AbstractMockResultSet. We should clear that tech debt by integrating the types
+ * TODO: of testing results sets into more reusable code.
+ */
+public abstract class AbstractAggregateResultSet implements RelationalResultSet {
+    private boolean isClosed = false;
+    private final RelationalResultSetMetaData metadata;
+    protected RelationalResultSet currentRow;
+
+    protected AbstractAggregateResultSet(RelationalResultSetMetaData metadata) {
+        this.metadata = metadata;
+    }
+
+    protected abstract boolean hasNext();
+
+    protected abstract RelationalResultSet advanceRow() throws SQLException;
+
+    @Override
+    public boolean next() throws SQLException {
+        currentRow = advanceRow();
+        return currentRow != null;
+    }
+
+    @Override
+    public void close() throws SQLException {
+        isClosed = true;
+    }
+
+    @Override
+    public boolean wasNull() throws SQLException {
+        checkCurrentRow();
+        return currentRow.wasNull();
+    }
+
+    @Override
+    public String getString(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getString(columnIndex);
+    }
+
+    @Override
+    public boolean getBoolean(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getBoolean(columnIndex);
+    }
+
+    @Override
+    public int getInt(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getInt(columnIndex);
+    }
+
+    @Override
+    public long getLong(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getLong(columnIndex);
+    }
+
+    @Override
+    public float getFloat(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getFloat(columnIndex);
+    }
+
+    @Override
+    public double getDouble(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getDouble(columnIndex);
+    }
+
+    @Override
+    public byte[] getBytes(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getBytes(columnIndex);
+    }
+
+    @Override
+    public Object getObject(int columnIndex) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getObject(columnIndex);
+    }
+
+    @Override
+    public RelationalStruct getStruct(int oneBasedPosition) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getStruct(oneBasedPosition);
+    }
+
+    @Override
+    public RelationalArray getArray(int oneBasedPosition) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getArray(oneBasedPosition);
+    }
+
+    @Override
+    public String getString(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getString(columnLabel);
+    }
+
+    @Override
+    public boolean getBoolean(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getBoolean(columnLabel);
+    }
+
+    @Override
+    public int getInt(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getInt(columnLabel);
+    }
+
+    @Override
+    public long getLong(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getLong(columnLabel);
+    }
+
+    @Override
+    public float getFloat(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getFloat(columnLabel);
+    }
+
+    @Override
+    public double getDouble(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getDouble(columnLabel);
+    }
+
+    @Override
+    public byte[] getBytes(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getBytes(columnLabel);
+    }
+
+    @Override
+    public Object getObject(String columnLabel) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getObject(columnLabel);
+    }
+
+    @Override
+    public RelationalStruct getStruct(String fieldName) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getStruct(fieldName);
+    }
+
+    @Override
+    public RelationalArray getArray(String fieldName) throws SQLException {
+        checkCurrentRow();
+        return currentRow.getArray(fieldName);
+    }
+
+    @Override
+    public RelationalResultSetMetaData getMetaData() throws SQLException {
+        return metadata;
+    }
+
+    @Override
+    public boolean isClosed() throws SQLException {
+        return isClosed;
+    }
+
+    private void checkCurrentRow() throws SQLException {
+        if ((currentRow == null) || isClosed()) {
+            throw new SQLException("ResultSet exhausted", ErrorCode.INVALID_CURSOR_STATE.getErrorCode());
+        }
+    }
+}

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/AggregateResultSet.java

@@ -0,0 +1,76 @@
+/*
+ * AggregateResultSet.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2025 Apple Inc. and the FoundationDB project authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.apple.foundationdb.relational.yamltests;
+
+import com.apple.foundationdb.relational.api.Continuation;
+import com.apple.foundationdb.relational.api.RelationalResultSet;
+import com.apple.foundationdb.relational.api.RelationalResultSetMetaData;
+import com.apple.foundationdb.relational.api.exceptions.ErrorCode;
+import org.junit.jupiter.api.Assertions;
+
+import javax.annotation.Nonnull;
+import java.sql.SQLException;
+import java.util.Iterator;
+
+/**
+ * An implementation of a result set that is made up of a set of inner result sets.
+ * This result set delegates all data get* calls to the inner delegates.
+ * At this time, each inner result set if assumed to represent a single row.
+ * TODO: It may be beneficial to create a more sophisticated aggregate that can take inner result sets that have more than
+ * TODO: one row each.
+ */
+public class AggregateResultSet extends AbstractAggregateResultSet {
+    private final Continuation continuation;
+    private final Iterator<? extends RelationalResultSet> rowIterator;
+
+    public AggregateResultSet(RelationalResultSetMetaData metadata, Continuation continuation, Iterator<RelationalResultSet> rowIterator) {
+        super(metadata);
+        this.continuation = continuation;
+        this.rowIterator = rowIterator;
+    }
+
+    @Override
+    protected boolean hasNext() {
+        return rowIterator.hasNext();
+    }
+
+    @Override
+    protected RelationalResultSet advanceRow() throws SQLException {
+        if (!hasNext()) {
+            return null;
+        }
+        if (currentRow != null) {
+            // Ensure that result sets don't have additional rows. Since we assume maxRows:1 here and advance the "outer" iterator,
+            // we should assert that the actual result doesn't contain extra "stuff"
+            Assertions.assertFalse(currentRow.next());
+        }
+        return rowIterator.next();
+    }
+
+    @Nonnull
+    @Override
+    public Continuation getContinuation() throws SQLException {
+        if (hasNext()) {
+            throw new SQLException("Continuation can only be returned once the result set has been exhausted", ErrorCode.UNSUPPORTED_OPERATION.getErrorCode());
+        }
+        return continuation;
+    }
+}

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/YamlExecutionContext.java

@@ -36,11 +36,13 @@
 import java.nio.charset.StandardCharsets;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.function.Supplier;
 
 @SuppressWarnings({"PMD.GuardLogStatement"}) // It already is, but PMD is confused and reporting error in unrelated locations.
 public final class YamlExecutionContext {
+    public static final String OPTION_FORCE_CONTINUATIONS = "optionForceContinuations";
 
     private static final Logger logger = LogManager.getLogger(YamlRunner.class);
 
@@ -54,6 +56,8 @@ public final class YamlExecutionContext {
     private final List<Block> finalizeBlocks = new ArrayList<>();
     @SuppressWarnings("AbbreviationAsWordInName")
     private final List<String> connectionURIs = new ArrayList<>();
+    // Additional options that can be set by the runners to impact test execution
+    private final Map<String, Object> additionalOptions;
 
     public static class YamlExecutionError extends RuntimeException {
 
@@ -64,10 +68,11 @@ public static class YamlExecutionError extends RuntimeException {
         }
     }
 
-    YamlExecutionContext(@Nonnull String resourcePath, @Nonnull YamlRunner.YamlConnectionFactory factory, boolean correctExplain) throws RelationalException {
+    YamlExecutionContext(@Nonnull String resourcePath, @Nonnull YamlRunner.YamlConnectionFactory factory, boolean correctExplain, @Nonnull final Map<String, Object> additionalOptions) throws RelationalException {
         this.connectionFactory = factory;
         this.resourcePath = resourcePath;
         this.editedFileStream = correctExplain ? loadFileToMemory(resourcePath) : null;
+        this.additionalOptions = Map.copyOf(additionalOptions);
         if (isNightly()) {
             logger.info("ℹ️ Running in the NIGHTLY context.");
             logger.info("ℹ️ Number of threads to be used for parallel execution " + getNumThreads());
@@ -225,6 +230,18 @@ public YamlExecutionError wrapContext(@Nullable Throwable e, @Nonnull Supplier<S
         }
     }
 
+    /**
+     * Return the value of an additional option, or a default value.
+     * Additional options are options set by the test execution environment that can control the test execution, in additional
+     * to the "core" set of options defined in this class.
+     * @param name the name of the option
+     * @param defaultValue the default value (if option is undefined)
+     * @return the defined value of the option, or the default value, if undefined
+     */
+    public Object getOption(String name, Object defaultValue) {
+        return additionalOptions.getOrDefault(name, defaultValue);
+    }
+
     @Nonnull
     private static List<String> loadFileToMemory(@Nonnull final String resourcePath) throws RelationalException {
         final ClassLoader classLoader = Thread.currentThread().getContextClassLoader();

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/YamlRunner.java

@@ -46,6 +46,7 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
@@ -97,9 +98,9 @@ default boolean isMultiServer() {
         }
     }
 
-    public YamlRunner(@Nonnull String resourcePath, @Nonnull YamlConnectionFactory factory, boolean correctExplain) throws RelationalException {
+    public YamlRunner(@Nonnull String resourcePath, @Nonnull YamlConnectionFactory factory, boolean correctExplain, @Nonnull final Map<String, Object> additionalOptions) throws RelationalException {
         this.resourcePath = resourcePath;
-        this.executionContext = new YamlExecutionContext(resourcePath, factory, correctExplain);
+        this.executionContext = new YamlExecutionContext(resourcePath, factory, correctExplain, additionalOptions);
     }
 
     public void run() throws Exception {

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/command/QueryCommand.java

@@ -155,7 +155,7 @@ private void executeInternal(@Nonnull final RelationalConnection connection, boo
         enableCascadesDebugger();
         boolean queryIsRunning = false;
         Continuation continuation = null;
-        int maxRows = 0;
+        Integer maxRows = null;
         boolean exhausted = false;
         var queryConfigsIterator = queryConfigs.listIterator();
         while (queryConfigsIterator.hasNext()) {
@@ -167,13 +167,13 @@ private void executeInternal(@Nonnull final RelationalConnection connection, boo
                 // Ignore debugger configuration, always set the debugger for plan hashes. Executing the plan hash
                 // can result in the explain plan being put into the plan cache, so running without the debugger
                 // can pollute cache and thus interfere with the explain's results
-                int finalMaxRows = maxRows;
+                Integer finalMaxRows = maxRows;
                 runWithDebugger(() -> executor.execute(connection, null, queryConfig, checkCache, finalMaxRows));
             } else if (QueryConfig.QUERY_CONFIG_EXPLAIN.equals(queryConfig.getConfigName()) || QueryConfig.QUERY_CONFIG_EXPLAIN_CONTAINS.equals(queryConfig.getConfigName())) {
                 Assert.thatUnchecked(!queryIsRunning, "Explain test should not be intermingled with query result tests");
                 // ignore debugger configuration, always set the debugger for explain, so we can always get consistent
                 // results
-                int finalMaxRows1 = maxRows;
+                Integer finalMaxRows1 = maxRows;
                 runWithDebugger(() -> executor.execute(connection, null, queryConfig, checkCache, finalMaxRows1));
             } else if (QueryConfig.QUERY_CONFIG_NO_OP.equals(queryConfig.getConfigName())) {
                 // Do nothing for noop execution.

yaml-tests/src/main/java/com/apple/foundationdb/relational/yamltests/command/QueryConfig.java

@@ -383,6 +383,8 @@ public static QueryConfig getNoOpConfig(int lineNumber, @Nonnull YamlExecutionCo
             @SuppressWarnings("PMD.CloseResource") // lifetime of autocloseable persists beyond method
             @Override
             void checkResultInternal(@Nonnull Object actual, @Nonnull String queryDescription) throws SQLException {
+                // This should not be executed
+                Assertions.fail("NoOp Config should not be executed");
             }
         };
     }