ESQL: Time_zone tests

x-pack/plugin/esql/qa/server/multi-clusters/src/javaRestTest/java/org/elasticsearch/xpack/esql/ccq/MultiClusterSpecIT.java

@@ -304,7 +304,15 @@ static CsvSpecReader.CsvTestCase convertToRemoteIndices(CsvSpecReader.CsvTestCas
         String first = commands[0].trim();
         // If true, we're using *:index, otherwise we're using *:index,index
         boolean onlyRemotes = canUseRemoteIndicesOnly() && randomBoolean();
-        String[] commandParts = first.split("\\s+", 2);
+
+        // Split "SET a=b; FROM x" into "SET a=b" and "FROM x"
+        int lastSetDelimiterPosition = first.lastIndexOf(';');
+        String setStatements = lastSetDelimiterPosition == -1 ? "" : first.substring(0, lastSetDelimiterPosition + 1);
+        String afterSetStatements = lastSetDelimiterPosition == -1 ? first : first.substring(lastSetDelimiterPosition + 1);
+
+        // Split "FROM a, b, c" into "FROM" and "a, b, c"
+        String[] commandParts = afterSetStatements.trim().split("\\s+", 2);
+
         String command = commandParts[0].trim();
         if (command.equalsIgnoreCase("from") || command.equalsIgnoreCase("ts")) {
             String[] indexMetadataParts = commandParts[1].split("(?i)\\bmetadata\\b", 2);
@@ -326,7 +334,7 @@ static CsvSpecReader.CsvTestCase convertToRemoteIndices(CsvSpecReader.CsvTestCas
                     + remoteIndices
                     + " "
                     + (indexMetadataParts.length == 1 ? "" : "metadata " + indexMetadataParts[1]);
-                testCase.query = newFirstCommand + query.substring(first.length());
+                testCase.query = setStatements + newFirstCommand + query.substring(first.length());
             }
         }
 

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

@@ -81,7 +81,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;
@@ -102,7 +101,9 @@
 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.IndexPattern;
+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.Explain;
@@ -527,9 +528,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,
@@ -546,12 +547,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/date.csv-spec

@@ -55,7 +55,6 @@ emp_no:integer | x:keyword                     | y:keyword
 10061          | 1985-09-17T00:00:00.000Z      | 1985-09-17                  
 ;
 
-
 compareToString
 from employees | where hire_date < "1985-03-01T00:00:00Z" | keep emp_no, hire_date;
 ignoreOrder:true
@@ -1085,6 +1084,33 @@ date:date                 | year:long
 2022-05-06T00:00:00.000Z  | 2022
 ;
 
+dateExtractSetTimezoneFrom
+required_capability: global_timezone_parameter
+
+set time_zone="+07:00"\;
+from employees
+| sort hire_date
+| eval hour = date_extract("hour_of_day", hire_date)
+| keep emp_no, hire_date, hour
+| limit 2;
+
+emp_no:integer | hire_date:date           | hour:long
+10009          | 1985-02-18T00:00:00.000Z | 7
+10048          | 1985-02-24T00:00:00.000Z | 7
+;
+
+dateExtractSetTimezoneRow
+required_capability: global_timezone_parameter
+
+set time_zone="+07:00"\;
+ROW hire_date = "2020-02-28T23:00:00.000Z"::date
+| EVAL hour = date_extract("hour_of_day", hire_date)
+| KEEP hour;
+
+hour:long
+6
+;
+
 docsDateExtractBusinessHours
 // tag::docsDateExtractBusinessHours[]
 FROM sample_data
@@ -1971,6 +1997,30 @@ emp_no:integer | birth_date:date                     | hire_date:date
 10040          | null                                | 1993-02-14T00:00:00.000Z   | null  
 ;
 
+dayNameSetTimezoneRow
+required_capability: global_timezone_parameter
+
+set time_zone="-02:00"\;
+row dt = to_datetime("1953-09-02T00:00:00.000Z")
+| eval weekday = day_name(dt);
+
+dt:date   | weekday:keyword
+1953-09-02T00:00:00.000Z | Tuesday
+;
+
+dayNameSetTimezoneFrom
+required_capability: global_timezone_parameter
+
+set time_zone="-02:00"\;
+from employees
+| sort emp_no
+| keep emp_no, hire_date
+| eval day = day_name(hire_date)
+| limit 1;
+
+emp_no:integer | hire_date:date             | day:keyword
+10001          | 1986-06-26T00:00:00.000Z   | Wednesday
+;
 
 monthNameRowTest
 required_capability:fn_month_name
@@ -2068,3 +2118,28 @@ from employees
 emp_no:integer | birth_date:date                     | hire_date:date             | monthName:keyword
 10040          | null                                | 1993-02-14T00:00:00.000Z   | null  
 ;
+
+monthNameSetTimezoneRow
+required_capability: global_timezone_parameter
+
+set time_zone="Europe/Paris"\;
+row dt = to_datetime("1996-01-31T23:00:00.000Z")
+| eval monthName = MONTH_NAME(dt);
+
+dt:date                  | monthName:keyword
+1996-01-31T23:00:00.000Z | February
+;
+
+monthNameSetTimezoneFrom
+required_capability: global_timezone_parameter
+
+set time_zone="-10:00"\;
+from employees 
+| WHERE emp_no == 10004 
+| keep emp_no, hire_date
+| eval monthName = month_name(hire_date) 
+| limit 1;
+
+emp_no:integer | hire_date:date             | monthName:keyword
+10004          | 1986-12-01T00:00:00.000Z   | November
+;

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

@@ -1541,6 +1541,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()),
+
         /**
          * Optional options argument for DATE_PARSE
          */

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

@@ -80,6 +80,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.IndexPattern;
 import org.elasticsearch.xpack.esql.plan.logical.Enrich;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
@@ -191,9 +192,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();
@@ -551,6 +556,7 @@ private static EnrichPolicy loadEnrichPolicyMapping(String policyFileName) {
 
     private LogicalPlan analyzedPlan(
         LogicalPlan parsed,
+        Configuration configuration,
         Map<IndexPattern, CsvTestsDataLoader.MultiIndexTestDataset> datasets,
         TransportVersion minimumVersion
     ) {
@@ -638,11 +644,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

@@ -116,6 +116,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

@@ -180,6 +180,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),
@@ -212,6 +214,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),