ESQL: Time_zone tests

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvSpecReader.java

@@ -13,9 +13,6 @@
 import java.util.function.Function;
 import java.util.regex.Pattern;
 
-import static org.hamcrest.CoreMatchers.is;
-import static org.junit.Assert.assertThat;
-
 public final class CsvSpecReader {
 
     private CsvSpecReader() {}
@@ -25,9 +22,6 @@ public static SpecReader.Parser specParser() {
     }
 
     public static class CsvSpecParser implements SpecReader.Parser {
-        private static final String SCHEMA_PREFIX = "schema::";
-
-        private final StringBuilder earlySchema = new StringBuilder();
         private final StringBuilder query = new StringBuilder();
         private final StringBuilder data = new StringBuilder();
         private final List<String> requiredCapabilities = new ArrayList<>();
@@ -39,21 +33,22 @@ private CsvSpecParser() {}
         public Object parse(String line) {
             // read the query
             if (testCase == null) {
-                if (line.startsWith(SCHEMA_PREFIX)) {
-                    assertThat("Early schema already declared " + earlySchema, earlySchema.length(), is(0));
-                    earlySchema.append(line.substring(SCHEMA_PREFIX.length()).trim());
-                } else if (line.toLowerCase(Locale.ROOT).startsWith("required_capability:")) {
+                if (line.toLowerCase(Locale.ROOT).startsWith("required_capability:")) {
                     requiredCapabilities.add(line.substring("required_capability:".length()).trim());
                 } else {
-                    if (line.endsWith(";")) {
+                    if (line.endsWith("\\;")) {
+                        // SET statement with escaped ";"
+                        var updatedLine = line.substring(0, line.length() - 2);
+                        query.append(updatedLine);
+                        query.append(";");
+                        query.append("\r\n");
+                    } else if (line.endsWith(";")) {
                         // pick up the query
                         testCase = new CsvTestCase();
                         query.append(line.substring(0, line.length() - 1).trim());
                         testCase.query = query.toString();
-                        testCase.earlySchema = earlySchema.toString();
                         testCase.requiredCapabilities = List.copyOf(requiredCapabilities);
                         requiredCapabilities.clear();
-                        earlySchema.setLength(0);
                         query.setLength(0);
                     }
                     // keep reading the query
@@ -109,7 +104,6 @@ private static Pattern warningRegexToPattern(String regex) {
 
     public static class CsvTestCase {
         public String query;
-        public String earlySchema;
         public String expectedResults;
         private final List<String> expectedWarnings = new ArrayList<>();
         private final List<String> expectedWarningsRegexString = new ArrayList<>();

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -78,7 +78,6 @@
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.core.type.DataType;
 import org.elasticsearch.xpack.esql.core.type.EsField;
-import org.elasticsearch.xpack.esql.core.util.DateUtils;
 import org.elasticsearch.xpack.esql.core.util.StringUtils;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.expression.function.scalar.spatial.StGeohash;
@@ -99,6 +98,8 @@
 import org.elasticsearch.xpack.esql.inference.InferenceService;
 import org.elasticsearch.xpack.esql.optimizer.LogicalOptimizerContext;
 import org.elasticsearch.xpack.esql.parser.QueryParam;
+import org.elasticsearch.xpack.esql.plan.EsqlStatement;
+import org.elasticsearch.xpack.esql.plan.QuerySettings;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.EsRelation;
 import org.elasticsearch.xpack.esql.plan.logical.Limit;
@@ -516,9 +517,9 @@ private static ThreadPool createMockThreadPool() {
 
     private EsqlTestUtils() {}
 
-    public static Configuration configuration(QueryPragmas pragmas, String query) {
+    public static Configuration configuration(QueryPragmas pragmas, String query, EsqlStatement statement) {
         return new Configuration(
-            DateUtils.UTC,
+            statement.setting(QuerySettings.TIME_ZONE),
             Locale.US,
             null,
             null,
@@ -535,12 +536,16 @@ public static Configuration configuration(QueryPragmas pragmas, String query) {
         );
     }
 
+    public static Configuration configuration(QueryPragmas pragmas, String query) {
+        return configuration(pragmas, query, new EsqlStatement(null, List.of()));
+    }
+
     public static Configuration configuration(QueryPragmas pragmas) {
         return configuration(pragmas, StringUtils.EMPTY);
     }
 
     public static Configuration configuration(String query) {
-        return configuration(new QueryPragmas(Settings.EMPTY), query);
+        return configuration(QueryPragmas.EMPTY, query);
     }
 
     public static AnalyzerSettings queryClusterSettings() {

x-pack/plugin/esql/qa/testFixtures/src/main/resources/set.csv-spec

@@ -0,0 +1,40 @@
+set
+required_capability: global_timezone_parameter
+
+set time_zone="+02:00"\;
+from employees 
+| sort emp_no 
+| keep emp_no, hire_date
+| eval hour = date_extract("hour_of_day", hire_date) 
+| limit 1;
+
+emp_no:integer | hire_date:date           | hour:long
+10001          | 1986-06-26T00:00:00.000Z | 2
+;
+
+set with foldable
+required_capability: global_timezone_parameter
+
+set time_zone="+02:00"\;
+ROW date = "1986-06-26T00:00:00.000Z"::date
+| eval hour = date_extract("hour_of_day", date)
+;
+
+date:date                | hour:long
+1986-06-26T00:00:00.000Z | 2
+;
+
+last set prevails
+required_capability: global_timezone_parameter
+
+set time_zone="+02:00"\;
+set time_zone="+05:00"\;
+from employees 
+| sort emp_no 
+| keep emp_no, hire_date
+| eval hour = date_extract("hour_of_day", hire_date) 
+| limit 1;
+
+emp_no:integer | hire_date:date           | hour:long
+10001          | 1986-06-26T00:00:00.000Z | 5
+;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/action/EsqlCapabilities.java

@@ -1516,6 +1516,11 @@ public enum Cap {
          */
         FIX_FILTER_ORDINALS,
 
+        /**
+         * "time_zone" parameter in request body and in {@code SET "time_zone"="x"}
+         */
+        GLOBAL_TIMEZONE_PARAMETER(Build.current().isSnapshot()),
+
         ;
 
         private final boolean enabled;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlParser.java

@@ -115,6 +115,11 @@ public LogicalPlan createStatement(String query, QueryParams params, PlanTelemet
         return invokeParser(query, params, metrics, EsqlBaseParser::singleStatement, AstBuilder::plan);
     }
 
+    // testing utility
+    public EsqlStatement createQuery(String query) {
+        return createQuery(query, new QueryParams());
+    }
+
     // testing utility
     public EsqlStatement createQuery(String query, QueryParams params) {
         return createQuery(query, params, new PlanTelemetry(new EsqlFunctionRegistry()));

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/Configuration.java

@@ -202,7 +202,7 @@ public long absoluteStartedTimeInMillis() {
     /**
      * @return Start time of the ESQL query in nanos
      */
-    public long getQueryStartTimeNanos() {
+    public long queryStartTimeNanos() {
         return queryStartTimeNanos;
     }
 

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

@@ -79,6 +79,7 @@
 import org.elasticsearch.xpack.esql.optimizer.LogicalPreOptimizerContext;
 import org.elasticsearch.xpack.esql.optimizer.TestLocalPhysicalPlanOptimizer;
 import org.elasticsearch.xpack.esql.parser.EsqlParser;
+import org.elasticsearch.xpack.esql.plan.EsqlStatement;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.plan.physical.ChangePointExec;
@@ -189,9 +190,13 @@ public class CsvTests extends ESTestCase {
     private final CsvSpecReader.CsvTestCase testCase;
     private final String instructions;
 
-    private final Configuration configuration = EsqlTestUtils.configuration(
-        new QueryPragmas(Settings.builder().put("page_size", randomPageSize()).build())
-    );
+    /**
+     * The configuration to be used in the tests.
+     * <p>
+     *     Initialized in {@link #executePlan}.
+     * </p>
+     */
+    private Configuration configuration;
     private final EsqlFunctionRegistry functionRegistry = new EsqlFunctionRegistry();
     private final EsqlParser parser = new EsqlParser();
     private final Mapper mapper = new Mapper();
@@ -526,6 +531,7 @@ private static EnrichPolicy loadEnrichPolicyMapping(String policyFileName) {
 
     private LogicalPlan analyzedPlan(
         LogicalPlan parsed,
+        Configuration configuration,
         CsvTestsDataLoader.MultiIndexTestDataset datasets,
         TransportVersion minimumVersion
     ) {
@@ -594,11 +600,16 @@ private static TestPhysicalOperationProviders testOperationProviders(
     }
 
     private ActualResults executePlan(BigArrays bigArrays) throws Exception {
-        LogicalPlan parsed = parser.createStatement(testCase.query);
-        var testDatasets = testDatasets(parsed);
+        EsqlStatement statement = parser.createQuery(testCase.query);
+        this.configuration = EsqlTestUtils.configuration(
+            new QueryPragmas(Settings.builder().put("page_size", randomPageSize()).build()),
+            testCase.query,
+            statement
+        );
+        var testDatasets = testDatasets(statement.plan());
         // Specifically use the newest transport version; the csv tests correspond to a single node cluster on the current version.
         TransportVersion minimumVersion = TransportVersion.current();
-        LogicalPlan analyzed = analyzedPlan(parsed, testDatasets, minimumVersion);
+        LogicalPlan analyzed = analyzedPlan(statement.plan(), configuration, testDatasets, minimumVersion);
 
         FoldContext foldCtx = FoldContext.small();
         EsqlSession session = new EsqlSession(

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

@@ -117,6 +117,8 @@ protected static List<TestCaseSupplier> withNoRowsExpectingNull(List<TestCaseSup
                     var newData = testCase.getData().stream().map(td -> td.isMultiRow() ? td.withData(List.of()) : td).toList();
 
                     return new TestCaseSupplier.TestCase(
+                        testCase.getSource(),
+                        testCase.getConfiguration(),
                         newData,
                         testCase.evaluatorToString(),
                         testCase.expectedType(),

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

@@ -179,6 +179,8 @@ protected static List<TestCaseSupplier> anyNullIsNull(
                     }).toList();
                     TestCaseSupplier.TypedData nulledData = oc.getData().get(finalNullPosition);
                     return new TestCaseSupplier.TestCase(
+                        oc.getSource(),
+                        oc.getConfiguration(),
                         data,
                         evaluatorToString.evaluatorToString(finalNullPosition, nulledData, oc.evaluatorToString()),
                         expectedType.expectedType(finalNullPosition, nulledData.type(), oc),
@@ -211,6 +213,8 @@ protected static List<TestCaseSupplier> anyNullIsNull(
                                 )
                                 .toList();
                             return new TestCaseSupplier.TestCase(
+                                oc.getSource(),
+                                oc.getConfiguration(),
                                 data,
                                 equalTo("LiteralsEvaluator[lit=null]"),
                                 expectedType.expectedType(finalNullPosition, DataType.NULL, oc),