Add SET CSV tests, and adapted CsvTests to work with SET statements

ivancea · ivancea · commit 569d9e530816 · 2025-10-17T14:44:58.000+02:00
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/CsvSpecReader.java b/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,16 @@ 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(";")) {
                         // 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 +98,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<>();
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java b/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,6 +536,10 @@ 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);
     }
diff --git a/x-pack/plugin/esql/qa/testFixtures/src/main/resources/set.csv-spec b/x-pack/plugin/esql/qa/testFixtures/src/main/resources/set.csv-spec
@@ -0,0 +1,36 @@
+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
+;
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlParser.java b/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()));
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/Configuration.java b/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;
     }
 
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/CsvTests.java b/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(

Original file line number	Diff line number	Diff line change
`@@ -202,7 +202,7 @@ public long absoluteStartedTimeInMillis() {`
`202`	`202`	`/**`
`203`	`203`	`* @return Start time of the ESQL query in nanos`
`204`	`204`	`*/`
`205`		`- public long getQueryStartTimeNanos() {`
	`205`	`+ public long queryStartTimeNanos() {`
`206`	`206`	`return queryStartTimeNanos;`
`207`	`207`	`}`
`208`	`208`