Introduce option to setup transaction before executing queries

scripts/YAML-SQL.xml

@@ -53,8 +53,8 @@
     multi_repetition_parallelized" />
     <keywords4 keywords="connect:;query:;load schema template:;set schema state:;result:;unorderedResult:;explain:;
     explainContains:;count:;error:;planHash:;setup:;schema_template:;supported_version:;test_block:;options:;tests:;
-    maxRows:;initialVersionAtLeast:;initialVersionLessThan:;;mode:;repetition:;seed:;
-    check_cache:;connection_lifecycle:;steps:;preset:;statement_type:;!r;!in;!a;" />
+    maxRows:;initialVersionAtLeast:;initialVersionLessThan:;;mode:;repetition:;seed:;setup:;transaction_setups:;
+    setupReference:;check_cache:;connection_lifecycle:;steps:;preset:;statement_type:;!r;!in;!a;" />
   </highlighting>
   <extensionMap>
     <mapping ext="yamsql" />

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

@@ -25,6 +25,7 @@
 import com.apple.foundationdb.relational.yamltests.block.FileOptions;
 import com.apple.foundationdb.relational.yamltests.block.SetupBlock;
 import com.apple.foundationdb.relational.yamltests.block.TestBlock;
+import com.apple.foundationdb.relational.yamltests.block.TransactionSetupsBlock;
 import com.apple.foundationdb.relational.yamltests.command.Command;
 import com.apple.foundationdb.relational.yamltests.command.QueryConfig;
 import org.yaml.snakeyaml.LoaderOptions;
@@ -61,6 +62,7 @@ public CustomYamlConstructor(LoaderOptions loaderOptions) {
         requireLineNumber.add(FileOptions.OPTIONS);
         requireLineNumber.add(SetupBlock.SETUP_BLOCK);
         requireLineNumber.add(SetupBlock.SchemaTemplateBlock.SCHEMA_TEMPLATE_BLOCK);
+        requireLineNumber.add(TransactionSetupsBlock.TRANSACTION_SETUP);
         requireLineNumber.add(TestBlock.TEST_BLOCK);
         // commands
         requireLineNumber.add(Command.COMMAND_LOAD_SCHEMA_TEMPLATE);
@@ -78,6 +80,8 @@ public CustomYamlConstructor(LoaderOptions loaderOptions) {
         requireLineNumber.add(QueryConfig.QUERY_CONFIG_SUPPORTED_VERSION);
         requireLineNumber.add(QueryConfig.QUERY_CONFIG_INITIAL_VERSION_LESS_THAN);
         requireLineNumber.add(QueryConfig.QUERY_CONFIG_INITIAL_VERSION_AT_LEAST);
+        requireLineNumber.add(QueryConfig.QUERY_CONFIG_SETUP);
+        requireLineNumber.add(QueryConfig.QUERY_CONFIG_SETUP_REFERENCE);
         requireLineNumber.add(QueryConfig.QUERY_CONFIG_DEBUGGER);
     }
 

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

@@ -24,8 +24,10 @@
 import com.apple.foundationdb.relational.api.RelationalConnection;
 import com.apple.foundationdb.relational.api.RelationalPreparedStatement;
 import com.apple.foundationdb.relational.api.RelationalStatement;
+import com.apple.foundationdb.relational.api.exceptions.RelationalException;
 import com.apple.foundationdb.relational.api.metrics.MetricCollector;
 import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalConnection;
+import com.apple.foundationdb.relational.yamltests.command.SQLFunction;
 import com.apple.foundationdb.relational.yamltests.server.SemanticVersion;
 import com.google.common.collect.Iterables;
 import org.junit.jupiter.api.Assumptions;
@@ -117,6 +119,24 @@ public SemanticVersion getInitialVersion() {
         return versions.get(0);
     }
 
+    @Override
+    public <T> T executeTransactionally(final SQLFunction<YamlConnection, T> transactionalWork) throws SQLException, RelationalException {
+        underlying.setAutoCommit(false);
+        T result;
+        try {
+            result = transactionalWork.apply(this);
+        } catch (final SQLException | RelationalException e) {
+            underlying.rollback();
+            throw e;
+        } finally {
+            // enabling autoCommit will commit if there is outstanding work
+            // It would probably be good to commit earlier, but https://github.com/FoundationDB/fdb-record-layer/pull/3477
+            // causes the setAutoCommit to fail if you do that
+            underlying.setAutoCommit(true);
+        }
+        return result;
+    }
+
     @Override
     public String toString() {
         return connectionLabel;

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

@@ -24,8 +24,10 @@
 import com.apple.foundationdb.relational.api.RelationalConnection;
 import com.apple.foundationdb.relational.api.RelationalPreparedStatement;
 import com.apple.foundationdb.relational.api.RelationalStatement;
+import com.apple.foundationdb.relational.api.exceptions.RelationalException;
 import com.apple.foundationdb.relational.api.metrics.MetricCollector;
 import com.apple.foundationdb.relational.recordlayer.EmbeddedRelationalConnection;
+import com.apple.foundationdb.relational.yamltests.command.SQLFunction;
 import com.apple.foundationdb.relational.yamltests.server.SemanticVersion;
 
 import javax.annotation.Nonnull;
@@ -104,4 +106,6 @@ public interface YamlConnection extends AutoCloseable {
      */
     @Nonnull
     SemanticVersion getInitialVersion();
+
+    <T> T executeTransactionally(SQLFunction<YamlConnection, T> transactionalWork) throws SQLException, RelationalException;
 }

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

@@ -84,6 +84,7 @@ public final class YamlExecutionContext {
     private final List<String> connectionURIs = new ArrayList<>();
     // Additional options that can be set by the runners to impact test execution
     private final ContextOptions additionalOptions;
+    private final Map<String, String> transactionSetups = new HashMap<>();
 
     public static class YamlExecutionError extends RuntimeException {
 
@@ -248,6 +249,20 @@ public URI inferConnectionURI(@Nullable Object connectObject) {
         }
     }
 
+    public void registerTransactionSetup(final String name, final String command) {
+        // Note: at the time of writing, this is only called by code that is iterating over a Map from yaml, so it will
+        // not prevent two entries in the yaml file itself
+        Assert.thatUnchecked(!transactionSetups.containsKey(name), ErrorCode.INTERNAL_ERROR,
+                () -> "Transaction setup " + name + " is defined multiple times.");
+        transactionSetups.put(name, command);
+    }
+
+    public String getTransactionSetup(final Object name) {
+        return Matchers.notNull(
+                transactionSetups.get(Matchers.string(name, "setup reference")),
+                "transaction setup " + name + " is not defined");
+    }
+
     @Nonnull
     public List<Block> getFinalizeBlocks() {
         return finalizeBlocks;

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

@@ -72,7 +72,8 @@ public void run() throws Exception {
             LoaderOptions loaderOptions = new LoaderOptions();
             loaderOptions.setAllowDuplicateKeys(true);
             DumperOptions dumperOptions = new DumperOptions();
-            final var yaml = new Yaml(new CustomYamlConstructor(loaderOptions), new Representer(dumperOptions), new DumperOptions(), loaderOptions, new Resolver());
+            final var yaml = new Yaml(new CustomYamlConstructor(loaderOptions), new Representer(dumperOptions),
+                    new DumperOptions(), loaderOptions, new Resolver());
 
             final var testBlocks = new ArrayList<TestBlock>();
             int blockNumber = 0;

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

@@ -41,6 +41,7 @@
  * </ul>
  */
 public interface Block {
+
     /**
      * Looks at the block to determine if its one of the valid blocks. If it is a valid one, parses it to that. This
      * method dispatches the execution to the right block which takes care of reading from the block and initializing
@@ -62,12 +63,14 @@ static Block parse(@Nonnull Object document, int blockNumber, @Nonnull YamlExecu
             switch (blockKey) {
                 case SetupBlock.SETUP_BLOCK:
                     return SetupBlock.ManualSetupBlock.parse(lineNumber, entry.getValue(), executionContext);
+                case TransactionSetupsBlock.TRANSACTION_SETUP:
+                    return TransactionSetupsBlock.parse(lineNumber, entry.getValue(), executionContext);
                 case TestBlock.TEST_BLOCK:
                     return TestBlock.parse(blockNumber, lineNumber, entry.getValue(), executionContext);
                 case SetupBlock.SchemaTemplateBlock.SCHEMA_TEMPLATE_BLOCK:
                     return SetupBlock.SchemaTemplateBlock.parse(lineNumber, entry.getValue(), executionContext);
                 case FileOptions.OPTIONS:
-                    Assert.thatUnchecked(blockNumber == 0,
+                    Assert.that(blockNumber == 0,
                             "File level options must be the first block, but found one at line " + lineNumber);
                     return FileOptions.parse(lineNumber, entry.getValue(), executionContext);
                 default:

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

@@ -0,0 +1,43 @@
+/*
+ * TransactionSetupsBlock.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.block;
+
+import com.apple.foundationdb.relational.yamltests.Matchers;
+import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
+
+import java.util.Map;
+
+public class TransactionSetupsBlock {
+    public static final String TRANSACTION_SETUP = "transaction_setups";
+
+    public static Block parse(final int lineNumber,
+                              final Object document,
+                              final YamlExecutionContext executionContext) {
+        final Map<?, ?> map = Matchers.map(document);
+        for (final Map.Entry<?, ?> entry : map.entrySet()) {
+            final String transactionSetupName = Matchers.string(entry.getKey(), "transaction setup name");
+            final String transactionSetupCommand = Matchers.string(entry.getValue(), "transaction setup command");
+            executionContext.registerTransactionSetup(transactionSetupName, transactionSetupCommand);
+        }
+        return new NoOpBlock(lineNumber);
+    }
+
+}

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

@@ -207,6 +207,18 @@ private void executeInternal(@Nonnull final YamlConnection connection, boolean c
             } else if (QueryConfig.QUERY_CONFIG_NO_OP.equals(queryConfig.getConfigName())) {
                 // Do nothing for noop execution.
                 continue;
+            } else if (QueryConfig.QUERY_CONFIG_SETUP.equals(queryConfig.getConfigName())) {
+                Assert.that(!queryIsRunning, "Transaction setup should not be intermingled with query results");
+                final String setupStatement = Matchers.notNull(Matchers.string(Matchers.notNull(queryConfig.getVal(),
+                                "Setup Config Val"), "Transaction setup"), "Transaction setup");
+                // we restrict transaction setups to CREATE TEMPORARY FUNCTION, because other mutations could be hard
+                // to reason about when running in a parallel world. It's possible the right answer is that we shouldn't
+                // commit after the query, or that we shouldn't allow things that modify state in the database. This will
+                // become clearer as we have more related tests, and for now, just try to stop people from being confused.
+                final String allowedStatement = "CREATE TEMPORARY FUNCTION";
+                Assert.that(setupStatement.regionMatches(true, 0, allowedStatement, 0, allowedStatement.length()),
+                        "Only \"CREATE TEMPORARY FUNCTION\" is allowed for transaction setups");
+                executor.addSetup(setupStatement);
             } else if (!QueryConfig.QUERY_CONFIG_SUPPORTED_VERSION.equals(queryConfig.getConfigName()) &&
                     !QueryConfig.QUERY_CONFIG_DEBUGGER.equals(queryConfig.getConfigName())) {
                 if (QueryConfig.QUERY_CONFIG_ERROR.equals(queryConfig.getConfigName())) {

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

@@ -85,6 +85,8 @@ public abstract class QueryConfig {
     public static final String QUERY_CONFIG_INITIAL_VERSION_AT_LEAST = "initialVersionAtLeast";
     public static final String QUERY_CONFIG_INITIAL_VERSION_LESS_THAN = "initialVersionLessThan";
     public static final String QUERY_CONFIG_NO_OP = "noOp";
+    public static final String QUERY_CONFIG_SETUP = "setup";
+    public static final String QUERY_CONFIG_SETUP_REFERENCE = "setupReference";
     public static final String QUERY_CONFIG_DEBUGGER = "debugger";
 
     private static final Set<String> RESULT_CONFIGS = ImmutableSet.of(QUERY_CONFIG_ERROR, QUERY_CONFIG_COUNT, QUERY_CONFIG_RESULT, QUERY_CONFIG_UNORDERED_RESULT);
@@ -466,6 +468,16 @@ void checkResultInternal(@Nonnull String currentQuery, @Nonnull Object actual, @
         };
     }
 
+    private static QueryConfig getSetupConfig(final Object value, final int lineNumber, final YamlExecutionContext executionContext) {
+        return new QueryConfig(QUERY_CONFIG_SETUP, value, lineNumber, executionContext) {
+            @Override
+            void checkResultInternal(@Nonnull final String currentQuery, @Nonnull final Object actual,
+                                     @Nonnull final String queryDescription) throws SQLException {
+                Assert.failUnchecked("No results to check on a setup config");
+            }
+        };
+    }
+
     private static QueryConfig getDebuggerConfig(@Nonnull Object value, int lineNumber, @Nonnull YamlExecutionContext executionContext) {
         return new QueryConfig(QUERY_CONFIG_DEBUGGER, DebuggerImplementation.valueOf(((String)value).toUpperCase(Locale.ROOT)),
                 lineNumber, executionContext) {
@@ -590,10 +602,14 @@ private static QueryConfig parseConfig(String blockName, String key, Object valu
             return getCheckResultConfig(false, key, value, lineNumber, executionContext);
         } else if (QUERY_CONFIG_MAX_ROWS.equals(key)) {
             return getMaxRowConfig(value, lineNumber, executionContext);
+        } else if (QUERY_CONFIG_SETUP.equals(key)) {
+            return getSetupConfig(value, lineNumber, executionContext);
+        } else if (QUERY_CONFIG_SETUP_REFERENCE.equals(key)) {
+            return getSetupConfig(executionContext.getTransactionSetup(value), lineNumber, executionContext);
         } else if (QUERY_CONFIG_DEBUGGER.equals(key)) {
             return getDebuggerConfig(value, lineNumber, executionContext);
         } else {
-            throw Assert.failUnchecked("‼️ '%s' is not a valid configuration");
+            throw Assert.failUnchecked("‼️ '" + key + "' is not a valid configuration");
         }
     }
 

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

@@ -24,6 +24,7 @@
 import com.apple.foundationdb.relational.api.RelationalPreparedStatement;
 import com.apple.foundationdb.relational.api.RelationalResultSet;
 import com.apple.foundationdb.relational.api.RelationalResultSetMetaData;
+import com.apple.foundationdb.relational.api.RelationalStatement;
 import com.apple.foundationdb.relational.api.exceptions.RelationalException;
 import com.apple.foundationdb.relational.api.metrics.RelationalMetric;
 import com.apple.foundationdb.relational.recordlayer.ContinuationImpl;
@@ -69,6 +70,8 @@ public class QueryExecutor {
      */
     @Nullable
     private final List<Parameter> parameters;
+    @Nonnull
+    private final List<String> setup = new ArrayList<>();
 
     QueryExecutor(@Nonnull String query, int lineNumber) {
         this(query, lineNumber, null);
@@ -163,23 +166,29 @@ private Continuation executeQuery(@Nonnull YamlConnection connection, @Nonnull Q
         try {
             if (parameters == null) {
                 logger.debug("⏳ Executing query '{}'", this.toString());
-                try (var s = connection.createStatement()) {
-                    final var queryResult = executeStatementAndCheckCacheIfNeeded(s, false, connection, currentQuery, checkCache, maxRows);
-                    config.checkResult(currentQuery, queryResult, this.toString(), connection);
-                    if (queryResult instanceof RelationalResultSet) {
-                        continuationAfter = ((RelationalResultSet) queryResult).getContinuation();
+                continuationAfter = executeWithSetup(connection, singleConnection -> {
+                    try (var s = singleConnection.createStatement()) {
+                        final var queryResult = executeStatementAndCheckCacheIfNeeded(s, false, singleConnection, currentQuery, checkCache, maxRows);
+                        config.checkResult(currentQuery, queryResult, this.toString(), singleConnection);
+                        if (queryResult instanceof RelationalResultSet) {
+                            return ((RelationalResultSet) queryResult).getContinuation();
+                        }
                     }
-                }
+                    return null;
+                });
             } else {
                 logger.debug("⏳ Executing query '{}'", this.toString());
-                try (var s = connection.prepareStatement(currentQuery)) {
-                    setParametersInPreparedStatement(s);
-                    final var queryResult = executeStatementAndCheckCacheIfNeeded(s, true, connection, currentQuery, checkCache, maxRows);
-                    config.checkResult(currentQuery, queryResult, this.toString(), connection);
-                    if (queryResult instanceof RelationalResultSet) {
-                        continuationAfter = ((RelationalResultSet) queryResult).getContinuation();
+                continuationAfter = executeWithSetup(connection, singleConnection -> {
+                    try (var s = singleConnection.prepareStatement(currentQuery)) {
+                        setParametersInPreparedStatement(s);
+                        final var queryResult = executeStatementAndCheckCacheIfNeeded(s, true, singleConnection, currentQuery, checkCache, maxRows);
+                        config.checkResult(currentQuery, queryResult, this.toString(), singleConnection);
+                        if (queryResult instanceof RelationalResultSet) {
+                            return ((RelationalResultSet)queryResult).getContinuation();
+                        }
                     }
-                }
+                    return null;
+                });
             }
             logger.debug("👍 Finished executing query '{}'", this.toString());
         } catch (SQLException sqle) {
@@ -188,6 +197,22 @@ private Continuation executeQuery(@Nonnull YamlConnection connection, @Nonnull Q
         return continuationAfter;
     }
 
+    @Nullable
+    private Continuation executeWithSetup(final @Nonnull YamlConnection connection,
+                                          SQLFunction<YamlConnection, Continuation> execute) throws SQLException, RelationalException {
+        if (!setup.isEmpty()) {
+            return connection.executeTransactionally(singleConnection -> {
+                for (var setupStatement : setup) {
+                    final RelationalStatement statement = singleConnection.createStatement();
+                    statement.execute(setupStatement);
+                }
+                return execute.apply(singleConnection);
+            });
+        } else {
+            return execute.apply(connection);
+        }
+    }
+
     @Nullable
     private Continuation executeContinuation(@Nonnull YamlConnection connection, @Nonnull Continuation continuation,
                                              @Nonnull QueryConfig config, final @Nonnull String currentQuery, @Nullable Integer maxRows) {
@@ -336,4 +361,8 @@ public String toString() {
             return query + " with parameters (" + String.join(", ", paramList) + ")";
         }
     }
+
+    public void addSetup(final String setupStatement) {
+        setup.add(setupStatement);
+    }
 }