Change context options to a type-safe option with generics.

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

@@ -46,6 +46,7 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Comparator;
+import java.util.HashMap;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
@@ -59,10 +60,10 @@
 public final class YamlExecutionContext {
     private static final Logger logger = LogManager.getLogger(YamlRunner.class);
 
-    public static final String OPTION_FORCE_CONTINUATIONS = "optionForceContinuations";
-    public static final String OPTION_CORRECT_EXPLAIN = "optionCorrectExplain";
-    public static final String OPTION_CORRECT_METRICS = "optionCorrectMetrics";
-    public static final String OPTION_SHOW_PLAN_ON_DIFF = "optionShowPlanOnDiff";
+    public static final ContextOption<Boolean> OPTION_FORCE_CONTINUATIONS = new ContextOption<>("optionForceContinuation");
+    public static final ContextOption<Boolean> OPTION_CORRECT_EXPLAIN = new ContextOption<>("optionCorrectExplain");
+    public static final ContextOption<Boolean> OPTION_CORRECT_METRICS = new ContextOption<>("optionCorrectMetrics");
+    public static final ContextOption<Boolean> OPTION_SHOW_PLAN_ON_DIFF = new ContextOption<>("optionShowPlanOnDiff");
 
     @Nonnull final String resourcePath;
     @Nullable
@@ -80,7 +81,7 @@ public final class YamlExecutionContext {
     @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;
+    private final ContextOptions additionalOptions;
 
     public static class YamlExecutionError extends RuntimeException {
 
@@ -92,12 +93,12 @@ public static class YamlExecutionError extends RuntimeException {
     }
 
     YamlExecutionContext(@Nonnull String resourcePath, @Nonnull YamlRunner.YamlConnectionFactory factory,
-                         @Nonnull final Map<String, Object> additionalOptions) throws RelationalException {
+                         @Nonnull final ContextOptions additionalOptions) throws RelationalException {
         this.connectionFactory = factory;
         this.resourcePath = resourcePath;
-        this.editedFileStream = (boolean)additionalOptions.getOrDefault(OPTION_CORRECT_EXPLAIN, false)
+        this.editedFileStream = additionalOptions.getOrDefault(OPTION_CORRECT_EXPLAIN, false)
                                 ? loadYamlResource(resourcePath) : null;
-        this.additionalOptions = Map.copyOf(additionalOptions);
+        this.additionalOptions = additionalOptions;
         this.expectedMetricsMap = loadMetricsResource(resourcePath);
         this.actualMetricsMap = new TreeMap<>(Comparator.comparing(QueryAndLocation::getLineNumber)
                 .thenComparing(QueryAndLocation::getBlockName)
@@ -120,17 +121,17 @@ public YamlRunner.YamlConnectionFactory getConnectionFactory() {
     }
 
     public boolean shouldCorrectExplains() {
-        final var shouldCorrectExplains = (boolean)additionalOptions.getOrDefault(OPTION_CORRECT_EXPLAIN, false);
+        final var shouldCorrectExplains = additionalOptions.getOrDefault(OPTION_CORRECT_EXPLAIN, false);
         Verify.verify(!shouldCorrectExplains || editedFileStream != null);
         return shouldCorrectExplains;
     }
 
     public boolean shouldCorrectMetrics() {
-        return (boolean)additionalOptions.getOrDefault(OPTION_CORRECT_METRICS, false);
+        return additionalOptions.getOrDefault(OPTION_CORRECT_METRICS, false);
     }
 
     public boolean shouldShowPlanOnDiff() {
-        return (boolean)additionalOptions.getOrDefault(OPTION_SHOW_PLAN_ON_DIFF, false);
+        return additionalOptions.getOrDefault(OPTION_SHOW_PLAN_ON_DIFF, false);
     }
 
     public boolean correctExplain(int lineNumber, @Nonnull String actual) {
@@ -303,12 +304,12 @@ 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 option the option to get value for
      * @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);
+    public <T> T getOption(ContextOption<T> option, T defaultValue) {
+        return additionalOptions.getOrDefault(option, defaultValue);
     }
 
     public void saveMetricsAsBinaryProto() {
@@ -490,4 +491,59 @@ public int hashCode() {
             return Objects.hash(blockName, query, lineNumber);
         }
     }
+
+    public static class ContextOptions {
+        public static final ContextOptions EMPTY_OPTIONS = new ContextOptions(Map.of());
+
+        @Nonnull
+        private final Map<ContextOption<?>, Object> map;
+
+        private ContextOptions(final @Nonnull Map<ContextOption<?>, Object> map) {
+            this.map = map;
+        }
+
+        public static <T> ContextOptions of(ContextOption<T> prop, T value) {
+            return new ContextOptions(Map.of(prop, value));
+        }
+
+        public static <T1, T2> ContextOptions of(ContextOption<T1> prop1, T1 value1, ContextOption<T2> prop2, T2 value2) {
+            return new ContextOptions(Map.of(prop1, value1, prop2, value2));
+        }
+
+        public ContextOptions mergeFrom(ContextOptions other) {
+            final Map<ContextOption<?>, Object> newMap = new HashMap<>(map);
+            newMap.putAll(other.map);
+            return new ContextOptions(newMap);
+        }
+
+        @SuppressWarnings("unchecked")
+        public <T> T getOrDefault(ContextOption<T> prop, T defaultValue) {
+            return (T)map.getOrDefault(prop, defaultValue);
+        }
+    }
+
+    public static class ContextOption<T> {
+        private final String name;
+
+        public ContextOption(final String name) {
+            this.name = name;
+        }
+
+        @Override
+        public boolean equals(final Object o) {
+            if (this == o) {
+                return true;
+            }
+            if (!(o instanceof ContextOption)) {
+                return false;
+            }
+            final ContextOption<?> that = (ContextOption<?>)o;
+            return Objects.equals(name, that.name);
+        }
+
+        @Override
+        public int hashCode() {
+            return Objects.hashCode(name);
+        }
+    }
 }

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

@@ -47,7 +47,6 @@
 import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.List;
-import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
@@ -100,7 +99,7 @@ default boolean isMultiServer() {
     }
 
     public YamlRunner(@Nonnull String resourcePath, @Nonnull YamlConnectionFactory factory,
-                      @Nonnull final Map<String, Object> additionalOptions) throws RelationalException {
+                      @Nonnull final YamlExecutionContext.ContextOptions additionalOptions) throws RelationalException {
         this.resourcePath = resourcePath;
         this.executionContext = new YamlExecutionContext(resourcePath, factory, additionalOptions);
     }

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

@@ -48,7 +48,7 @@ boolean filter(final YamlTestConfig config) {
     DO_NOT_FORCE_CONTINUATIONS {
         @Override
         boolean filter(final YamlTestConfig config) {
-            return !(boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_FORCE_CONTINUATIONS, false);
+            return !config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_FORCE_CONTINUATIONS, false);
         }
     },
     /**
@@ -57,8 +57,8 @@ boolean filter(final YamlTestConfig config) {
     CORRECT_EXPLAINS {
         @Override
         boolean filter(final YamlTestConfig config) {
-            return (boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, false) &&
-                    !(boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_METRICS, false);
+            return config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, false) &&
+                    !config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_METRICS, false);
         }
     },
     /**
@@ -67,8 +67,8 @@ boolean filter(final YamlTestConfig config) {
     CORRECT_METRICS {
         @Override
         boolean filter(final YamlTestConfig config) {
-            return !(boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, false) &&
-                    (boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_METRICS, false);
+            return !config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, false) &&
+                    config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_METRICS, false);
         }
     },
     /**
@@ -77,17 +77,8 @@ boolean filter(final YamlTestConfig config) {
     CORRECT_EXPLAIN_AND_METRICS {
         @Override
         boolean filter(final YamlTestConfig config) {
-            return (boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, false) &&
-                    (boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_METRICS, false);
-        }
-    },
-    /**
-     * Used to show the plan diffs graphically.
-     */
-    SHOW_PLAN_ON_DIFF {
-        @Override
-        boolean filter(final YamlTestConfig config) {
-            return (boolean)config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_SHOW_PLAN_ON_DIFF, false);
+            return config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, false) &&
+                    config.getRunnerOptions().getOrDefault(YamlExecutionContext.OPTION_CORRECT_METRICS, false);
         }
     };
 

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

@@ -20,11 +20,10 @@
 
 package com.apple.foundationdb.relational.yamltests.configs;
 
+import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 import com.apple.foundationdb.relational.yamltests.YamlRunner;
 
 import javax.annotation.Nonnull;
-import java.util.HashMap;
-import java.util.Map;
 
 /**
  * An implementation of {@link YamlTestConfig} that sets additional options on top of a base config.
@@ -33,13 +32,11 @@ public class ConfigWithOptions implements YamlTestConfig {
     @Nonnull
     private final YamlTestConfig underlying;
     @Nonnull
-    private final Map<String, Object> runnerOptions;
+    private final YamlExecutionContext.ContextOptions runnerOptions;
 
-    public ConfigWithOptions(@Nonnull final YamlTestConfig underlying, @Nonnull Map<String, Object> newOptions) {
+    public ConfigWithOptions(@Nonnull final YamlTestConfig underlying, @Nonnull YamlExecutionContext.ContextOptions newOptions) {
         this.underlying = underlying;
-        final HashMap<String, Object> options = new HashMap<>(underlying.getRunnerOptions());
-        options.putAll(newOptions);
-        this.runnerOptions = Map.copyOf(options);
+        this.runnerOptions = underlying.getRunnerOptions().mergeFrom(newOptions);
     }
 
 
@@ -49,7 +46,7 @@ public YamlRunner.YamlConnectionFactory createConnectionFactory() {
     }
 
     @Override
-    public @Nonnull Map<String, Object> getRunnerOptions() {
+    public @Nonnull YamlExecutionContext.ContextOptions getRunnerOptions() {
         return runnerOptions;
     }
 

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

@@ -23,7 +23,6 @@
 import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 
 import javax.annotation.Nonnull;
-import java.util.Map;
 
 /**
  * A configuration that runs an underlying configuration, but heals expected explains in YAML files where
@@ -34,6 +33,6 @@
  */
 public class CorrectExplains extends ConfigWithOptions {
     public CorrectExplains(@Nonnull final YamlTestConfig underlying) {
-        super(underlying, Map.of(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, true));
+        super(underlying, YamlExecutionContext.ContextOptions.of(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, true));
     }
 }

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

@@ -23,7 +23,6 @@
 import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 
 import javax.annotation.Nonnull;
-import java.util.Map;
 
 /**
  * A configuration that runs an underlying configuration, but heals expected explains in YAML files where
@@ -34,7 +33,7 @@
  */
 public class CorrectExplainsAndMetrics extends ConfigWithOptions {
     public CorrectExplainsAndMetrics(@Nonnull final YamlTestConfig underlying) {
-        super(underlying, Map.of(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, true,
+        super(underlying, YamlExecutionContext.ContextOptions.of(YamlExecutionContext.OPTION_CORRECT_EXPLAIN, true,
                 YamlExecutionContext.OPTION_CORRECT_METRICS, true));
     }
 }

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

@@ -23,7 +23,6 @@
 import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 
 import javax.annotation.Nonnull;
-import java.util.Map;
 
 /**
  * A configuration that runs an underlying configuration, but heals expected explains in YAML files where
@@ -34,6 +33,6 @@
  */
 public class CorrectMetrics extends ConfigWithOptions {
     public CorrectMetrics(@Nonnull final YamlTestConfig underlying) {
-        super(underlying, Map.of(YamlExecutionContext.OPTION_CORRECT_METRICS, true));
+        super(underlying, YamlExecutionContext.ContextOptions.of(YamlExecutionContext.OPTION_CORRECT_METRICS, true));
     }
 }

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

@@ -22,14 +22,14 @@
 
 import com.apple.foundationdb.relational.api.Options;
 import com.apple.foundationdb.relational.server.FRL;
+import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 import com.apple.foundationdb.relational.yamltests.YamlRunner;
 
 import javax.annotation.Nonnull;
 import java.net.URI;
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.SQLException;
-import java.util.Map;
 import java.util.Set;
 
 /**
@@ -72,8 +72,8 @@ public Set<String> getVersionsUnderTest() {
     }
 
     @Override
-    public @Nonnull Map<String, Object> getRunnerOptions() {
-        return Map.of();
+    public @Nonnull YamlExecutionContext.ContextOptions getRunnerOptions() {
+        return YamlExecutionContext.ContextOptions.EMPTY_OPTIONS;
     }
 
     @Override

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

@@ -23,7 +23,6 @@
 import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 
 import javax.annotation.Nonnull;
-import java.util.Map;
 
 /**
  * A configuration that runs an underlying configuration, but forces every query to be executed with {@code maxRows: 1}.
@@ -33,6 +32,6 @@
  */
 public class ForceContinuations extends ConfigWithOptions {
     public ForceContinuations(@Nonnull final YamlTestConfig underlying) {
-        super(underlying, Map.of(YamlExecutionContext.OPTION_FORCE_CONTINUATIONS, true));
+        super(underlying, YamlExecutionContext.ContextOptions.of(YamlExecutionContext.OPTION_FORCE_CONTINUATIONS, true));
     }
 }

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

@@ -22,6 +22,7 @@
 
 import com.apple.foundationdb.relational.jdbc.JDBCURI;
 import com.apple.foundationdb.relational.server.InProcessRelationalServer;
+import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 import com.apple.foundationdb.relational.yamltests.YamlRunner;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -32,7 +33,6 @@
 import java.sql.Connection;
 import java.sql.DriverManager;
 import java.sql.SQLException;
-import java.util.Map;
 import java.util.Set;
 
 /**
@@ -81,8 +81,8 @@ public Set<String> getVersionsUnderTest() {
     }
 
     @Override
-    public @Nonnull Map<String, Object> getRunnerOptions() {
-        return Map.of();
+    public @Nonnull YamlExecutionContext.ContextOptions getRunnerOptions() {
+        return YamlExecutionContext.ContextOptions.EMPTY_OPTIONS;
     }
 
     @Override

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

@@ -23,7 +23,6 @@
 import com.apple.foundationdb.relational.yamltests.YamlExecutionContext;
 
 import javax.annotation.Nonnull;
-import java.util.Map;
 
 /**
  * A configuration that runs an underlying configuration, but shows the dots of the expected and the actual plan where
@@ -34,6 +33,6 @@
  */
 public class ShowPlanOnDiff extends ConfigWithOptions {
     public ShowPlanOnDiff(@Nonnull final YamlTestConfig underlying) {
-        super(underlying, Map.of(YamlExecutionContext.OPTION_SHOW_PLAN_ON_DIFF, true));
+        super(underlying, YamlExecutionContext.ContextOptions.of(YamlExecutionContext.OPTION_SHOW_PLAN_ON_DIFF, true));
     }
 }