ESQL: Timezone query setting and added to the query config

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/QuerySettings.java

@@ -7,96 +7,67 @@
 
 package org.elasticsearch.xpack.esql.plan;
 
+import org.elasticsearch.core.Nullable;
 import org.elasticsearch.transport.RemoteClusterService;
+import org.elasticsearch.xpack.esql.core.QlIllegalArgumentException;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.type.DataType;
+import org.elasticsearch.xpack.esql.expression.Foldables;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 
-import java.util.function.Predicate;
+import java.time.ZoneId;
+import java.time.ZoneOffset;
+import java.util.function.Function;
 
-public enum QuerySettings {
+public class QuerySettings {
     // TODO check cluster state and see if project routing is allowed
     // see https://github.com/elastic/elasticsearch/pull/134446
     // PROJECT_ROUTING(..., state -> state.getRemoteClusterNames().crossProjectEnabled());
-    PROJECT_ROUTING(
+    public static final QuerySettingDef<String> PROJECT_ROUTING = new QuerySettingDef<>(
         "project_routing",
         DataType.KEYWORD,
         true,
         false,
         true,
         "A project routing expression, "
             + "used to define which projects to route the query to. "
-            + "Only supported if Cross-Project Search is enabled."
-    ),;
+            + "Only supported if Cross-Project Search is enabled.",
+        (value, settings) -> Foldables.stringLiteralValueOf(value, "Unexpected value"),
+        (_rcs) -> null
+    );
 
-    private String settingName;
-    private DataType type;
-    private final boolean serverlessOnly;
-    private final boolean snapshotOnly;
-    private final boolean preview;
-    private final String description;
-    private final Predicate<RemoteClusterService> validator;
-
-    QuerySettings(
-        String name,
-        DataType type,
-        boolean serverlessOnly,
-        boolean preview,
-        boolean snapshotOnly,
-        String description,
-        Predicate<RemoteClusterService> validator
-    ) {
-        this.settingName = name;
-        this.type = type;
-        this.serverlessOnly = serverlessOnly;
-        this.preview = preview;
-        this.snapshotOnly = snapshotOnly;
-        this.description = description;
-        this.validator = validator;
-    }
-
-    QuerySettings(String name, DataType type, boolean serverlessOnly, boolean preview, boolean snapshotOnly, String description) {
-        this(name, type, serverlessOnly, preview, snapshotOnly, description, state -> true);
-    }
-
-    public String settingName() {
-        return settingName;
-    }
-
-    public DataType type() {
-        return type;
-    }
-
-    public boolean serverlessOnly() {
-        return serverlessOnly;
-    }
-
-    public boolean snapshotOnly() {
-        return snapshotOnly;
-    }
-
-    public boolean preview() {
-        return preview;
-    }
-
-    public String description() {
-        return description;
-    }
+    public static final QuerySettingDef<ZoneId> TIME_ZONE = new QuerySettingDef<>(
+        "time_zone",
+        DataType.KEYWORD,
+        false,
+        true,
+        true,
+        "The default timezone to be used in the query, by the functions and commands that require it. Defaults to UTC",
+        (value, _rcs) -> {
+            String timeZone = Foldables.stringLiteralValueOf(value, "Unexpected value");
+            try {
+                return ZoneId.of(timeZone);
+            } catch (Exception exc) {
+                throw new QlIllegalArgumentException("Invalid time zone [" + timeZone + "]");
+            }
+        },
+        (_rcs) -> ZoneOffset.UTC
+    );
 
-    public Predicate<RemoteClusterService> validator() {
-        return validator;
-    }
+    public static final QuerySettingDef<?>[] ALL_SETTINGS = { PROJECT_ROUTING, TIME_ZONE };
 
     public static void validate(EsqlStatement statement, RemoteClusterService clusterService) {
         for (QuerySetting setting : statement.settings()) {
             boolean found = false;
-            for (QuerySettings qs : values()) {
-                if (qs.settingName().equals(setting.name())) {
+            for (QuerySettingDef<?> def : ALL_SETTINGS) {
+                if (def.name().equals(setting.name())) {
                     found = true;
-                    if (setting.value().dataType() != qs.type()) {
-                        throw new ParsingException(setting.source(), "Setting [" + setting.name() + "] must be of type " + qs.type());
+                    if (setting.value().dataType() != def.type()) {
+                        throw new ParsingException(setting.source(), "Setting [" + setting.name() + "] must be of type " + def.type());
                     }
-                    if (qs.validator().test(clusterService) == false) {
-                        throw new ParsingException(setting.source(), "Setting [" + setting.name() + "] is not allowed");
+                    String error = def.validator().validate(setting.value(), clusterService);
+                    if (error != null) {
+                        throw new ParsingException("Error validating setting [" + setting.name() + "]: " + error);
                     }
                     break;
                 }
@@ -106,4 +77,75 @@ public static void validate(EsqlStatement statement, RemoteClusterService cluste
             }
         }
     }
+
+    /**
+     * Definition of a query setting.
+     *
+     * @param name The name to be used when setting it in the query. E.g. {@code SET name=value}
+     * @param type The allowed datatype of the setting.
+     * @param serverlessOnly
+     * @param preview
+     * @param snapshotOnly
+     * @param description The user-facing description of the setting.
+     * @param validator A validation function to check the setting value.
+     *                  Defaults to calling the {@link #parser} and returning the error message of any exception it throws.
+     * @param parser A function to parse the setting value into the final object.
+     * @param defaultValueSupplier A supplier of the default value to be used when the setting is not set.
+     * @param <T> The type of the setting value.
+     */
+    public record QuerySettingDef<T>(
+        String name,
+        DataType type,
+        boolean serverlessOnly,
+        boolean preview,
+        boolean snapshotOnly,
+        String description,
+        Validator validator,
+        Parser<T> parser,
+        Function<RemoteClusterService, T> defaultValueSupplier
+    ) {
+        public QuerySettingDef(
+            String name,
+            DataType type,
+            boolean serverlessOnly,
+            boolean preview,
+            boolean snapshotOnly,
+            String description,
+            Parser<T> parser,
+            Function<RemoteClusterService, T> defaultValueSupplier
+        ) {
+            this(name, type, serverlessOnly, preview, snapshotOnly, description, (value, rcs) -> {
+                try {
+                    parser.parse(value, rcs);
+                    return null;
+                } catch (Exception exc) {
+                    return exc.getMessage();
+                }
+            }, parser, defaultValueSupplier);
+        }
+
+        public T get(Expression value, RemoteClusterService clusterService) {
+            if (value == null) {
+                return defaultValueSupplier.apply(clusterService);
+            }
+            return parser.parse(value, clusterService);
+        }
+
+        @FunctionalInterface
+        public interface Validator {
+            /**
+             * Validates the setting value and returns the error message if there's an error, or null otherwise.
+             */
+            @Nullable
+            String validate(Expression value, RemoteClusterService clusterService);
+        }
+
+        @FunctionalInterface
+        public interface Parser<T> {
+            /**
+             * Parses an already validated expression.
+             */
+            T parse(Expression value, RemoteClusterService clusterService);
+        }
+    }
 }

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/DocsV3Support.java

@@ -42,7 +42,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.NotEquals;
-import org.elasticsearch.xpack.esql.plan.QuerySettings;
+import org.elasticsearch.xpack.esql.plan.QuerySettings.QuerySettingDef;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.session.Configuration;
 
@@ -1081,10 +1081,10 @@ void renderTypes(String name, List<EsqlFunctionRegistry.ArgSignature> args) thro
 
     public static class SettingsDocsSupport extends DocsV3Support {
 
-        private final QuerySettings setting;
+        private final QuerySettingDef<?> setting;
 
-        public SettingsDocsSupport(QuerySettings setting, Class<?> testClass, Callbacks callbacks) {
-            super("settings", setting.settingName(), testClass, Set::of, callbacks);
+        public SettingsDocsSupport(QuerySettingDef<?> setting, Class<?> testClass, Callbacks callbacks) {
+            super("settings", setting.name(), testClass, Set::of, callbacks);
             this.setting = setting;
         }
 

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/plan/QuerySettingsTests.java

@@ -7,56 +7,99 @@
 
 package org.elasticsearch.xpack.esql.plan;
 
-import org.elasticsearch.common.lucene.BytesRefs;
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.esql.core.expression.Alias;
+import org.elasticsearch.xpack.esql.core.expression.Expression;
 import org.elasticsearch.xpack.esql.core.expression.Literal;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.expression.function.DocsV3Support;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
+import org.hamcrest.Matcher;
 import org.junit.AfterClass;
 
+import java.time.ZoneId;
+import java.time.ZoneOffset;
 import java.util.List;
 
+import static org.hamcrest.Matchers.both;
 import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.nullValue;
 
 public class QuerySettingsTests extends ESTestCase {
+    public void testNonExistingSetting() {
+        String settingName = "non_existing";
 
-    public void test() {
+        assertInvalid(settingName, Literal.keyword(Source.EMPTY, "12"), "Unknown setting [" + settingName + "]");
+    }
 
-        QuerySetting project_routing = new QuerySetting(
-            Source.EMPTY,
-            new Alias(Source.EMPTY, "project_routing", new Literal(Source.EMPTY, BytesRefs.toBytesRef("my-project"), DataType.KEYWORD))
-        );
-        QuerySettings.validate(new EsqlStatement(null, List.of(project_routing)), null);
+    public void testProjectRouting() {
+        var setting = QuerySettings.PROJECT_ROUTING;
 
-        QuerySetting wrong_type = new QuerySetting(
-            Source.EMPTY,
-            new Alias(Source.EMPTY, "project_routing", new Literal(Source.EMPTY, 12, DataType.INTEGER))
-        );
-        assertThat(
-            expectThrows(ParsingException.class, () -> QuerySettings.validate(new EsqlStatement(null, List.of(wrong_type)), null))
-                .getMessage(),
-            containsString("Setting [project_routing] must be of type KEYWORD")
+        assertDefault(setting, nullValue());
+        assertValid(setting, Literal.keyword(Source.EMPTY, "my-project"), equalTo("my-project"));
+
+        assertInvalid(
+            setting.name(),
+            new Literal(Source.EMPTY, 12, DataType.INTEGER),
+            "Setting [" + setting.name() + "] must be of type KEYWORD"
         );
+    }
+
+    public void testTimeZone() {
+        var setting = QuerySettings.TIME_ZONE;
+
+        assertDefault(setting, both(equalTo(ZoneId.of("Z"))).and(equalTo(ZoneOffset.UTC)));
+
+        assertValid(setting, Literal.keyword(Source.EMPTY, "UTC"), equalTo(ZoneId.of("UTC")));
+        assertValid(setting, Literal.keyword(Source.EMPTY, "Z"), both(equalTo(ZoneId.of("Z"))).and(equalTo(ZoneOffset.UTC)));
+        assertValid(setting, Literal.keyword(Source.EMPTY, "Europe/Madrid"), equalTo(ZoneId.of("Europe/Madrid")));
+        assertValid(setting, Literal.keyword(Source.EMPTY, "+05:00"), equalTo(ZoneId.of("+05:00")));
+        assertValid(setting, Literal.keyword(Source.EMPTY, "+05"), equalTo(ZoneId.of("+05")));
+        assertValid(setting, Literal.keyword(Source.EMPTY, "+07:15"), equalTo(ZoneId.of("+07:15")));
 
-        QuerySetting non_existing = new QuerySetting(
-            Source.EMPTY,
-            new Alias(Source.EMPTY, "non_existing", new Literal(Source.EMPTY, BytesRefs.toBytesRef("12"), DataType.KEYWORD))
+        assertInvalid(setting.name(), Literal.integer(Source.EMPTY, 12), "Setting [" + setting.name() + "] must be of type KEYWORD");
+        assertInvalid(
+            setting.name(),
+            Literal.keyword(Source.EMPTY, "Europe/New York"),
+            "Error validating setting [" + setting.name() + "]: Invalid time zone [Europe/New York]"
         );
+    }
+
+    private static <T> void assertValid(
+        QuerySettings.QuerySettingDef<T> settingDef,
+        Expression valueExpression,
+        Matcher<T> parsedValueMatcher
+    ) {
+        QuerySetting setting = new QuerySetting(Source.EMPTY, new Alias(Source.EMPTY, settingDef.name(), valueExpression));
+        QuerySettings.validate(new EsqlStatement(null, List.of(setting)), null);
+
+        T value = settingDef.get(valueExpression, null);
+
+        assertThat(value, parsedValueMatcher);
+    }
+
+    private static void assertInvalid(String settingName, Expression valueExpression, String expectedMessage) {
+        QuerySetting setting = new QuerySetting(Source.EMPTY, new Alias(Source.EMPTY, settingName, valueExpression));
         assertThat(
-            expectThrows(ParsingException.class, () -> QuerySettings.validate(new EsqlStatement(null, List.of(non_existing)), null))
+            expectThrows(ParsingException.class, () -> QuerySettings.validate(new EsqlStatement(null, List.of(setting)), null))
                 .getMessage(),
-            containsString("Unknown setting [non_existing]")
+            containsString(expectedMessage)
         );
     }
 
+    private static <T> void assertDefault(QuerySettings.QuerySettingDef<T> settingDef, Matcher<? super T> defaultMatcher) {
+        T value = settingDef.get(null, null);
+
+        assertThat(value, defaultMatcher);
+    }
+
     @AfterClass
     public static void generateDocs() throws Exception {
-        for (QuerySettings value : QuerySettings.values()) {
+        for (QuerySettings.QuerySettingDef<?> def : QuerySettings.ALL_SETTINGS) {
             DocsV3Support.SettingsDocsSupport settingsDocsSupport = new DocsV3Support.SettingsDocsSupport(
-                value,
+                def,
                 QuerySettingsTests.class,
                 DocsV3Support.callbacksFromSystemProperty()
             );