Add addExplain annotation to Yamsql framework

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

@@ -94,6 +94,7 @@ public final class YamlExecutionContext {
     public static final ContextOption<Boolean> OPTION_CORRECT_METRICS = new ContextOption<>("optionCorrectMetrics");
     public static final ContextOption<Boolean> OPTION_CORRECT_RESULT_METADATA = new ContextOption<>("optionCorrectResultMetadata");
     public static final ContextOption<Boolean> OPTION_ADD_RESULT_METADATA = new ContextOption<>("optionAddResultMetadata");
+    public static final ContextOption<Boolean> OPTION_ADD_EXPLAIN = new ContextOption<>("optionAddExplain");
     public static final ContextOption<Boolean> OPTION_SHOW_PLAN_ON_DIFF = new ContextOption<>("optionShowPlanOnDiff");
 
     private static final URI SYSTEM_CATALOG_ADDRESS = URI.create("jdbc:embed:/__SYS?schema=CATALOG");
@@ -139,7 +140,7 @@ public static class YamlExecutionError extends RuntimeException {
         this.connectionFactory = factory;
         this.topLevelResource = topLevelResource;
         this.additionalOptions = additionalOptions;
-        if (isInCI() && (shouldCorrectExplains() || shouldCorrectMetrics() || shouldCorrectResultMetadata() || shouldAddResultMetadata())) {
+        if (isInCI() && (shouldCorrectExplains() || shouldCorrectMetrics() || shouldCorrectResultMetadata() || shouldAddResultMetadata() || shouldAddExplains())) {
             logger.error("‼️ Yamsql files cannot be modified during CI runs.");
             Assertions.fail("‼️ Yamsql files cannot be modified during CI runs. " +
                     "Make sure maintenance annotations have not been checked in.");
@@ -155,7 +156,7 @@ public void registerResource(@Nonnull final YamlReference.YamlResource resource)
         if (registeredResources.contains(resource)) {
             throw new RuntimeException("The resource " + resource + " is already registered.");
         }
-        if (shouldCorrectExplains() || shouldCorrectResultMetadata() || shouldAddResultMetadata()) {
+        if (shouldCorrectExplains() || shouldCorrectResultMetadata() || shouldAddResultMetadata() || shouldAddExplains()) {
             this.editedFileStream.put(resource, loadYamlResource(resource));
         }
         if (this.expectedMetricsMap == null) {
@@ -196,6 +197,10 @@ public boolean shouldAddResultMetadata() {
         return additionalOptions.getOrDefault(OPTION_ADD_RESULT_METADATA, false);
     }
 
+    public boolean shouldAddExplains() {
+        return additionalOptions.getOrDefault(OPTION_ADD_EXPLAIN, false);
+    }
+
     public boolean correctResultMetadata(@Nonnull final YamlReference reference,
                                          @Nonnull final List<CheckResultMetadataConfig.ColumnDescriptor> actualColumns) {
         if (!shouldCorrectResultMetadata()) {
@@ -240,7 +245,7 @@ public boolean addResultMetadata(@Nonnull final YamlReference queryReference,
     }
 
     public boolean correctExplain(@Nonnull final YamlReference reference, @Nonnull String actual) {
-        if (!shouldCorrectExplains()) {
+        if (!shouldCorrectExplains() && !shouldAddExplains()) {
             return false;
         }
         if (editedFileStream.get(reference.getResource()) == null) {
@@ -472,6 +477,83 @@ public void apply(@Nonnull final List<String> lines) {
         }
     }
 
+    public boolean addExplain(@Nonnull final YamlReference queryReference, @Nonnull String actual) {
+        if (!shouldAddExplains()) {
+            return false;
+        }
+        if (editedFileStream.get(queryReference.getResource()) == null) {
+            return false;
+        }
+        synchronized (this) {
+            final List<YamlCorrection> corrections = pendingCorrections
+                    .computeIfAbsent(queryReference.getResource(), k -> new ArrayList<>());
+            final int lineNumber = queryReference.getLineNumber();
+            final boolean alreadyPending = corrections.stream()
+                    .anyMatch(c -> c instanceof AddExplainCorrection && c.getLineNumber() == lineNumber);
+            if (!alreadyPending) {
+                corrections.add(new AddExplainCorrection(queryReference, actual));
+                isDirty.put(queryReference.getResource(), true);
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Inserts a new {@code explain:} line into the YAMSQL source file immediately before the first config entry
+     * that follows the {@code query:} line. Used when {@link #OPTION_ADD_EXPLAIN} is set and the query had no
+     * {@code explain:} block.
+     */
+    public static final class AddExplainCorrection implements YamlCorrection {
+        @Nonnull
+        private final YamlReference queryReference;
+        @Nonnull
+        private final String actual;
+
+        public AddExplainCorrection(@Nonnull final YamlReference queryReference, @Nonnull final String actual) {
+            this.queryReference = queryReference;
+            this.actual = actual;
+        }
+
+        @Override
+        public int getLineNumber() {
+            return queryReference.getLineNumber();
+        }
+
+        @Override
+        public void apply(@Nonnull final List<String> lines) {
+            final int queryLineIdx = queryReference.getLineNumber() - 1; // 1-based → 0-based
+            if (queryLineIdx < 0 || queryLineIdx >= lines.size()) {
+                return;
+            }
+            final String queryLine = lines.get(queryLineIdx);
+
+            // Determine indentation from the "- query:" line
+            int indent = 0;
+            while (indent < queryLine.length() && queryLine.charAt(indent) == ' ') {
+                indent++;
+            }
+
+            final String itemPrefix = " ".repeat(indent);
+            final String explainLine = itemPrefix + "- explain: \"" + actual + "\"";
+
+            // Scan forward past any query-string continuation lines to find the first config entry
+            // at the same indentation level, and insert the explain line before it.
+            int insertIdx = queryLineIdx + 1;
+            for (int i = queryLineIdx + 1; i < lines.size(); i++) {
+                final String line = lines.get(i);
+                if (line.startsWith(itemPrefix + "- ")) {
+                    insertIdx = i;
+                    break;
+                }
+                // Stop if we've moved to the next test item (lower indentation)
+                if (indent > 0 && line.length() >= indent && !line.startsWith(itemPrefix)) {
+                    break;
+                }
+            }
+            lines.add(insertIdx, explainLine);
+        }
+    }
+
     @Nullable
     public ImmutableMap<PlannerMetricsProto.Identifier, PlannerMetricsProto.Info> getMetricsMap() {
         return expectedMetricsMap;

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

@@ -108,6 +108,15 @@ boolean filter(final YamlTestConfig config) {
         boolean filter(final YamlTestConfig config) {
             return config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_ADD_RESULT_METADATA, false);
         }
+    },
+    /**
+     * Used to add {@code explain} blocks to queries that don't yet have one, and to correct existing ones.
+     */
+    ADD_EXPLAINS {
+        @Override
+        boolean filter(final YamlTestConfig config) {
+            return config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_ADD_EXPLAIN, false);
+        }
     };
 
     /**

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

@@ -20,6 +20,7 @@
 
 package com.apple.foundationdb.relational.yamltests;
 
+import com.apple.foundationdb.relational.yamltests.configs.AddExplains;
 import com.apple.foundationdb.relational.yamltests.configs.AddResultMetadata;
 import com.apple.foundationdb.annotation.API;
 import com.apple.foundationdb.relational.yamltests.configs.CorrectExplains;
@@ -114,6 +115,7 @@ public void beforeAll(final ExtensionContext context) throws Exception {
                 new CorrectExpectations(new EmbeddedConfig(clusterFiles)),
                 new CorrectResultMetadata(new EmbeddedConfig(clusterFiles)),
                 new AddResultMetadata(new EmbeddedConfig(clusterFiles)),
+                new AddExplains(new EmbeddedConfig(clusterFiles)),
                 new ShowPlanOnDiff(new EmbeddedConfig(clusterFiles))
         );
         if (Boolean.parseBoolean(System.getProperty("tests.runQuick", "false"))) {

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

@@ -27,6 +27,7 @@
 import com.apple.foundationdb.relational.recordlayer.util.ExceptionUtil;
 import com.apple.foundationdb.relational.util.Assert;
 import com.apple.foundationdb.relational.util.Environment;
+import com.apple.foundationdb.relational.yamltests.command.queryconfigs.CheckExplainConfig;
 import com.apple.foundationdb.relational.yamltests.command.queryconfigs.CheckResultMetadataConfig;
 import com.apple.foundationdb.relational.yamltests.CustomYamlConstructor;
 import com.apple.foundationdb.relational.yamltests.Matchers;
@@ -83,61 +84,80 @@
 
     @Nonnull
     public static Command parse(@Nonnull final YamlReference.YamlResource resource, @Nonnull final Object object, @Nonnull final String blockName, @Nonnull final YamlExecutionContext executionContext) {
         final var queryCommand = Matchers.firstEntry(Matchers.first(Matchers.arrayList(object, "query command")), "query command");
         final var linedObject = CustomYamlConstructor.LinedObject.cast(queryCommand.getKey(), () -> "Invalid command key-value pair: " + queryCommand);
         final var reference = resource.withLineNumber(Matchers.notNull(linedObject, "query").getLineNumber());
         try {
             if (Debugger.getDebugger() != null) {
                 Debugger.getDebugger().onSetup(); // clean all symbols before the next query.
             }
             final var queryString = Matchers.notNull(Matchers.string(queryCommand.getValue(), "query string"), "query string");
             final var queryInterpreter = QueryInterpreter.withQueryString(reference, queryString, executionContext);
             final List<?> queryConfigsWithValueList = Matchers.arrayList(object).stream().skip(1).collect(Collectors.toList());
             List<QueryConfig> configs = queryConfigsWithValueList.isEmpty() ?
                     List.of(QueryConfig.getNoCheckConfig(reference)) :
                     QueryConfig.parseConfigs(blockName, reference, queryConfigsWithValueList, executionContext);
 
             final List<QueryConfig> skipConfigs = configs.stream().filter(config -> config instanceof QueryConfig.SkipConfig)
                     .collect(Collectors.toList());
             Assert.thatUnchecked(skipConfigs.size() < 2, "Query should not have more than one skip");
             if (!skipConfigs.isEmpty()) {
                 return new SkippedCommand(reference, executionContext,
                         ((QueryConfig.SkipConfig)skipConfigs.get(0)).getMessage(),
                         queryInterpreter.getQuery());
             }
 
             // When OPTION_ADD_RESULT_METADATA is set, inject a synthetic resultMetadata config for queries
             // that have a result/unorderedResult config but no resultMetadata block yet.
             // Only inject for result-returning queries (those with result: or unorderedResult:) to avoid
             // re-executing INSERT/UPDATE/DELETE statements that have no result config.
             // Insert it before the first result/unorderedResult config so it runs while queryIsRunning is
             // false (independent query execution), matching the expected ordering in YAMSQL files.
             if (executionContext.shouldAddResultMetadata()
                     && configs.stream().noneMatch(c -> QueryConfig.QUERY_CONFIG_RESULT_METADATA.equals(c.getConfigName()))
                     && configs.stream().anyMatch(c -> QueryConfig.QUERY_CONFIG_RESULT.equals(c.getConfigName())
                             || QueryConfig.QUERY_CONFIG_UNORDERED_RESULT.equals(c.getConfigName()))) {
                 final List<QueryConfig> mutableConfigs = new ArrayList<>(configs);
                 int insertIdx = mutableConfigs.size();
                 for (int i = 0; i < mutableConfigs.size(); i++) {
                     final String name = mutableConfigs.get(i).getConfigName();
                     if (QueryConfig.QUERY_CONFIG_RESULT.equals(name) || QueryConfig.QUERY_CONFIG_UNORDERED_RESULT.equals(name)) {
                         insertIdx = i;
                         break;
                     }
                 }
                 mutableConfigs.add(insertIdx, new CheckResultMetadataConfig(
                         QueryConfig.QUERY_CONFIG_RESULT_METADATA, null, reference, executionContext));
                 configs = mutableConfigs;
             }
 
+            // When OPTION_ADD_EXPLAIN is set, inject a synthetic explain config for queries that have a
+            // result/unorderedResult config but no explain block yet, so the actual plan is captured and
+            // written into the YAMSQL source file.
+            if (executionContext.shouldAddExplains()
+                    && QueryConfig.shouldExecuteExplain(executionContext)
+                    && configs.stream().noneMatch(c -> QueryConfig.QUERY_CONFIG_EXPLAIN.equals(c.getConfigName())
+                            || QueryConfig.QUERY_CONFIG_EXPLAIN_CONTAINS.equals(c.getConfigName()))
+                    && configs.stream().anyMatch(c -> QueryConfig.QUERY_CONFIG_RESULT.equals(c.getConfigName())
+                            || QueryConfig.QUERY_CONFIG_UNORDERED_RESULT.equals(c.getConfigName()))) {
+                final List<QueryConfig> mutableConfigs = new ArrayList<>(configs);
+                // Insert after supported_version if it is the first config, otherwise at the front.
+                int insertIdx = (!mutableConfigs.isEmpty()
+                        && QueryConfig.QUERY_CONFIG_SUPPORTED_VERSION.equals(mutableConfigs.get(0).getConfigName()))
+                        ? 1 : 0;
+                mutableConfigs.add(insertIdx, new CheckExplainConfig(
+                        QueryConfig.QUERY_CONFIG_EXPLAIN, null, reference, executionContext, true, blockName));
+                configs = mutableConfigs;
+            }
+
             final boolean hasDebuggerConfig =
                     configs.stream()
                             .anyMatch(config -> Objects.equals(config.getConfigName(), QueryConfig.QUERY_CONFIG_DEBUGGER));
             return new QueryCommand(reference, queryInterpreter, configs, executionContext, hasDebuggerConfig);
         } catch (Exception e) {
             throw executionContext.wrapContext(e,
                     () -> "‼️ Error parsing query command at " + reference, "query", reference);
         }
     }
 
     public static QueryCommand withQueryString(@Nonnull final YamlReference reference, @Nonnull String singleExecutableCommand,

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

@@ -581,7 +581,7 @@ public SemanticVersion getMaxVersion() {
         }
     }
 
-    private static boolean shouldExecuteExplain(final YamlExecutionContext executionContext) {
+    static boolean shouldExecuteExplain(final YamlExecutionContext executionContext) {
         return (! executionContext.getConnectionFactory().isMultiServer());
     }
 

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

@@ -97,7 +97,8 @@
         checkExplain(queryDescription, actualPlan, actualDot, expectedDot);
 
         final var actualPlannerMetrics = resultSet.getStruct(6);
-        if (isExact && actualPlannerMetrics != null) {
+        if (isExact && getVal() != null && actualPlannerMetrics != null
+                && (expectedPlannerMetricsInfo != null || executionContext.shouldCorrectMetrics())) {
             Objects.requireNonNull(actualDot);
             checkMetrics(currentQuery, setups, actualPlannerMetrics, expectedPlannerMetricsInfo, actualPlan, actualDot);
         }
@@ -107,49 +108,61 @@
                               final @Nonnull String actualPlan,
                               final @Nullable String actualDot,
                               final @Nullable String expectedDot) {
+        // Synthetic explain config (value==null): add the actual plan to the file without comparing.
+        if (getVal() == null) {
+            if (isExact && executionContext.shouldAddExplains()) {
+                if (!executionContext.addExplain(getReference(), actualPlan)) {
+                    QueryCommand.reportTestFailure("‼️ Cannot add explain plan at " + getReference());
+                } else {
+                    logger.debug(() -> "⭐️ Successfully added plan at " + getReference());
+                }
+            }
+            return;
+        }
+
         var success = isExact ? getVal().equals(actualPlan) : actualPlan.contains((String) getVal());
         if (success) {
             logger.debug("✅️ plan match!");
         } else {
             if (executionContext.shouldShowPlanOnDiff() &&
                     actualDot != null && expectedDot != null) {
                 BrowserHelper.browse("/showPlanDiff.html",
                         ImmutableMap.of("$SQL", queryDescription,
                                 "$DOT_EXPECTED", expectedDot,
                                 "$DOT_ACTUAL", actualDot));
             }
 
             final var expectedPlan = getValueString();
             final var diffGenerator = DiffRowGenerator.create()
                     .showInlineDiffs(true)
                     .inlineDiffByWord(true)
                     .newTag(f -> f ? CommandUtil.Color.RED.toString() : CommandUtil.Color.RESET.toString())
                     .oldTag(f -> f ? CommandUtil.Color.GREEN.toString() : CommandUtil.Color.RESET.toString())
                     .build();
             final List<DiffRow> diffRows = diffGenerator.generateDiffRows(
                     Collections.singletonList(expectedPlan),
                     Collections.singletonList(actualPlan));
             final var planDiffs = new StringBuilder();
             for (final var diffRow : diffRows) {
                 planDiffs.append(diffRow.getOldLine()).append('\n').append(diffRow.getNewLine()).append('\n');
             }
-            if (isExact && executionContext.shouldCorrectExplains()) {
+            if (isExact && (executionContext.shouldCorrectExplains() || executionContext.shouldAddExplains())) {
                 if (!executionContext.correctExplain(getReference(), actualPlan)) {
                     QueryCommand.reportTestFailure("‼️ Cannot correct explain plan at " + getReference());
                 } else {
                     logger.debug(() -> "⭐️ Successfully replaced plan at " + getReference());
                 }
             } else {
                 final var diffMessage = String.format(Locale.ROOT, "‼️ plan mismatch at %s:%n" +
                         "⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤%n%s" +
                         "⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤%n" +
                         "↪ expected plan %s:%n%s%n" +
                         "⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤⏤%n" +
                         "↩ actual plan:%n%s",
                         getReference(), planDiffs, (!isExact ? "fragment" : ""), getValueString(), actualPlan);
                 QueryCommand.reportTestFailure(diffMessage);
             }
         }
     }
 
     private void checkMetrics(final @Nonnull String currentQuery,

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

@@ -0,0 +1,45 @@
+/*
+ * AddExplains.java
+ *
+ * This source file is part of the FoundationDB open source project
+ *
+ * Copyright 2015-2026 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.configs;
+
+import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
+
+import javax.annotation.Nonnull;
+
+/**
+ * A configuration that runs an underlying configuration and adds {@code explain} blocks to all
+ * queries in YAMSQL files that do not already have one, and corrects existing {@code explain} blocks
+ * whose values have changed.
+ * <p>
+ *     For each query without an {@code explain:} block, the actual query plan is written into the YAMSQL
+ *     source file immediately after the {@code query:} line (before any other config entries). For queries
+ *     that already have an {@code explain:} block whose value differs from the actual plan, the block is
+ *     updated in place.
+ * </p>
+ * <p>
+ *     See {@link YamlExecutionContext#OPTION_ADD_EXPLAIN}.
+ * </p>
+ */
+public class AddExplains extends ConfigWithOptions {
+    public AddExplains(@Nonnull final YamlTestConfig underlying) {
+        super(underlying, YamlExecutionContext.ContextOptions.of(YamlExecutionContext.OPTION_ADD_EXPLAIN, true));
+    }
+}