SQL: Preserve original source for each expression (#36912)

Improve parsing to save the source for each token alongside the location 
 of each Node/Expression for accurate reproducibility of an expression
 name and source

Fix #36894

Author: costin

x-pack/plugin/sql/sql-client/src/main/java/org/elasticsearch/xpack/sql/client/SslConfig.java

@@ -121,19 +121,19 @@ private KeyManager[] loadKeyManagers() throws GeneralSecurityException, IOExcept
     }
 
 
-    private KeyStore loadKeyStore(String location, char[] pass, String keyStoreType) throws GeneralSecurityException, IOException {
+    private KeyStore loadKeyStore(String source, char[] pass, String keyStoreType) throws GeneralSecurityException, IOException {
         KeyStore keyStore = KeyStore.getInstance(keyStoreType);
-        Path path = Paths.get(location);
+        Path path = Paths.get(source);
 
         if (!Files.exists(path)) {
            throw new ClientException(
-                   "Expected to find keystore file at [" + location + "] but was unable to. Make sure you have specified a valid URI.");
+                    "Expected to find keystore file at [" + source + "] but was unable to. Make sure you have specified a valid URI.");
         }
 
-        try (InputStream in = Files.newInputStream(Paths.get(location), StandardOpenOption.READ)) {
+        try (InputStream in = Files.newInputStream(Paths.get(source), StandardOpenOption.READ)) {
             keyStore.load(in, pass);
         } catch (Exception ex) {
-            throw new ClientException("Cannot open keystore [" + location + "] - " + ex.getMessage(), ex);
+            throw new ClientException("Cannot open keystore [" + source + "] - " + ex.getMessage(), ex);
         } finally {
 
         }
@@ -174,6 +174,7 @@ public boolean equals(Object obj) {
                 && Objects.equals(truststoreType, other.truststoreType);
     }
 
+    @Override
     public int hashCode() {
         return getClass().hashCode();
     }

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/AnalysisException.java

@@ -21,8 +21,8 @@ public AnalysisException(Node<?> source, String message, Object... args) {
         super(message, args);
 
         Location loc = Location.EMPTY;
-        if (source != null && source.location() != null) {
-            loc = source.location();
+        if (source != null && source.source() != null) {
+            loc = source.source().source();
         }
         this.line = loc.getLineNumber();
         this.column = loc.getColumnNumber();
@@ -32,8 +32,8 @@ public AnalysisException(Node<?> source, String message, Throwable cause) {
         super(message, cause);
 
         Location loc = Location.EMPTY;
-        if (source != null && source.location() != null) {
-            loc = source.location();
+        if (source != null && source.source() != null) {
+            loc = source.source().source();
         }
         this.line = loc.getLineNumber();
         this.column = loc.getColumnNumber();

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/Analyzer.java

@@ -5,14 +5,14 @@
  */
 package org.elasticsearch.xpack.sql.analysis.analyzer;
 
-import java.util.Collection;
-import java.util.stream.Collectors;
-
 import org.elasticsearch.xpack.sql.analysis.AnalysisException;
 import org.elasticsearch.xpack.sql.analysis.analyzer.Verifier.Failure;
 import org.elasticsearch.xpack.sql.tree.Location;
 import org.elasticsearch.xpack.sql.util.StringUtils;
 
+import java.util.Collection;
+import java.util.stream.Collectors;
+
 
 public class VerificationException extends AnalysisException {
 
@@ -27,7 +27,7 @@ protected VerificationException(Collection<Failure> sources) {
     public String getMessage() {
         return failures.stream()
                 .map(f -> {
-                    Location l = f.source().location();
+                    Location l = f.source().source().source();
                     return "line " + l.getLineNumber() + ":" + l.getColumnNumber() + ": " + f.message();
                 })
                 .collect(Collectors.joining(StringUtils.NEW_LINE, "Found " + failures.size() + " problem(s)\n", StringUtils.EMPTY));

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/analysis/analyzer/VerificationException.java

@@ -281,7 +281,7 @@ private BucketExtractor createExtractor(FieldExtraction ref, BucketExtractor tot
                 // wrap only agg inputs
                 proc = proc.transformDown(l -> {
                     BucketExtractor be = createExtractor(l.context(), totalCount);
-                    return new AggExtractorInput(l.location(), l.expression(), l.action(), be);
+                    return new AggExtractorInput(l.source(), l.expression(), l.action(), be);
                 }, AggPathInput.class);
 
                 return new ComputingExtractor(proc.asProcessor());
@@ -364,7 +364,7 @@ private HitExtractor createExtractor(FieldExtraction ref) {
                         throw new SqlIllegalArgumentException("Multi-level nested fields [{}] not supported yet", hitNames);
                     }
 
-                    return new HitExtractorInput(l.location(), l.expression(), he);
+                    return new HitExtractorInput(l.source(), l.expression(), he);
                 }, ReferenceInput.class);
                 String hitName = null;
                 if (hitNames.size() == 1) {

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/execution/search/Querier.java

@@ -7,7 +7,7 @@
 
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
-import org.elasticsearch.xpack.sql.tree.Location;
+import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.tree.NodeInfo;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.type.EsField;
@@ -36,20 +36,20 @@ public class Alias extends NamedExpression {
      */
     private Attribute lazyAttribute;
 
-    public Alias(Location location, String name, Expression child) {
-        this(location, name, null, child, null);
+    public Alias(Source source, String name, Expression child) {
+        this(source, name, null, child, null);
     }
 
-    public Alias(Location location, String name, String qualifier, Expression child) {
-        this(location, name, qualifier, child, null);
+    public Alias(Source source, String name, String qualifier, Expression child) {
+        this(source, name, qualifier, child, null);
     }
 
-    public Alias(Location location, String name, String qualifier, Expression child, ExpressionId id) {
-        this(location, name, qualifier, child, id, false);
+    public Alias(Source source, String name, String qualifier, Expression child, ExpressionId id) {
+        this(source, name, qualifier, child, id, false);
     }
 
-    public Alias(Location location, String name, String qualifier, Expression child, ExpressionId id, boolean synthetic) {
-        super(location, name, singletonList(child), id, synthetic);
+    public Alias(Source source, String name, String qualifier, Expression child, ExpressionId id, boolean synthetic) {
+        super(source, name, singletonList(child), id, synthetic);
         this.child = child;
         this.qualifier = qualifier;
     }
@@ -64,7 +64,7 @@ public Expression replaceChildren(List<Expression> newChildren) {
         if (newChildren.size() != 1) {
             throw new IllegalArgumentException("expected [1] child but received [" + newChildren.size() + "]");
         }
-        return new Alias(location(), name(), qualifier, newChildren.get(0), id(), synthetic());
+        return new Alias(source(), name(), qualifier, newChildren.get(0), id(), synthetic());
     }
 
     public Expression child() {
@@ -104,17 +104,17 @@ private Attribute createAttribute() {
 
             Attribute attr = Expressions.attribute(c);
             if (attr != null) {
-                return attr.clone(location(), name(), qualifier, child.nullable(), id(), synthetic());
+                return attr.clone(source(), name(), qualifier, child.nullable(), id(), synthetic());
             }
             else {
                 // TODO: WE need to fix this fake Field
-                return new FieldAttribute(location(), null, name(),
+                return new FieldAttribute(source(), null, name(),
                         new EsField(name(), child.dataType(), Collections.emptyMap(), true),
                         qualifier, child.nullable(), id(), synthetic());
             }
         }
 
-        return new UnresolvedAttribute(location(), name(), qualifier);
+        return new UnresolvedAttribute(source(), name(), qualifier);
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Alias.java

@@ -7,7 +7,7 @@
 
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 import org.elasticsearch.xpack.sql.expression.gen.script.ScriptTemplate;
-import org.elasticsearch.xpack.sql.tree.Location;
+import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.tree.NodeInfo;
 
 import java.util.List;
@@ -44,16 +44,16 @@ public abstract class Attribute extends NamedExpression {
     // can the attr be null - typically used in JOINs
     private final boolean nullable;
 
-    public Attribute(Location location, String name, String qualifier, ExpressionId id) {
-        this(location, name, qualifier, true, id);
+    public Attribute(Source source, String name, String qualifier, ExpressionId id) {
+        this(source, name, qualifier, true, id);
     }
 
-    public Attribute(Location location, String name, String qualifier, boolean nullable, ExpressionId id) {
-        this(location, name, qualifier, nullable, id, false);
+    public Attribute(Source source, String name, String qualifier, boolean nullable, ExpressionId id) {
+        this(source, name, qualifier, nullable, id, false);
     }
 
-    public Attribute(Location location, String name, String qualifier, boolean nullable, ExpressionId id, boolean synthetic) {
-        super(location, name, emptyList(), id, synthetic);
+    public Attribute(Source source, String name, String qualifier, boolean nullable, ExpressionId id, boolean synthetic) {
+        super(source, name, emptyList(), id, synthetic);
         this.qualifier = qualifier;
         this.nullable = nullable;
     }
@@ -86,19 +86,19 @@ public AttributeSet references() {
         return new AttributeSet(this);
     }
 
-    public Attribute withLocation(Location location) {
-        return Objects.equals(location(), location) ? this : clone(location, name(), qualifier(), nullable(), id(), synthetic());
+    public Attribute withLocation(Source source) {
+        return Objects.equals(source(), source) ? this : clone(source, name(), qualifier(), nullable(), id(), synthetic());
     }
 
     public Attribute withQualifier(String qualifier) {
-        return Objects.equals(qualifier(), qualifier) ? this : clone(location(), name(), qualifier, nullable(), id(), synthetic());
+        return Objects.equals(qualifier(), qualifier) ? this : clone(source(), name(), qualifier, nullable(), id(), synthetic());
     }
 
     public Attribute withNullability(boolean nullable) {
-        return Objects.equals(nullable(), nullable) ? this : clone(location(), name(), qualifier(), nullable, id(), synthetic());
+        return Objects.equals(nullable(), nullable) ? this : clone(source(), name(), qualifier(), nullable, id(), synthetic());
     }
 
-    protected abstract Attribute clone(Location location, String name, String qualifier, boolean nullable, ExpressionId id,
+    protected abstract Attribute clone(Source source, String name, String qualifier, boolean nullable, ExpressionId id,
                                        boolean synthetic);
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Attribute.java

@@ -6,18 +6,18 @@
 package org.elasticsearch.xpack.sql.expression;
 
 import org.elasticsearch.xpack.sql.plan.logical.LogicalPlan;
-import org.elasticsearch.xpack.sql.tree.Location;
+import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.tree.NodeInfo;
 import org.elasticsearch.xpack.sql.type.DataType;
 
 public class Exists extends SubQueryExpression {
 
-    public Exists(Location location, LogicalPlan query) {
-        this(location, query, null);
+    public Exists(Source source, LogicalPlan query) {
+        this(source, query, null);
     }
 
-    public Exists(Location location, LogicalPlan query, ExpressionId id) {
-        super(location, query, id);
+    public Exists(Source source, LogicalPlan query, ExpressionId id) {
+        super(source, query, id);
     }
 
     @Override
@@ -27,7 +27,7 @@ protected NodeInfo<Exists> info() {
 
     @Override
     protected SubQueryExpression clone(LogicalPlan newQuery) {
-        return new Exists(location(), newQuery);
+        return new Exists(source(), newQuery);
     }
 
     @Override

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Exists.java

@@ -8,7 +8,7 @@
 import org.elasticsearch.xpack.sql.SqlIllegalArgumentException;
 import org.elasticsearch.xpack.sql.capabilities.Resolvable;
 import org.elasticsearch.xpack.sql.capabilities.Resolvables;
-import org.elasticsearch.xpack.sql.tree.Location;
+import org.elasticsearch.xpack.sql.tree.Source;
 import org.elasticsearch.xpack.sql.tree.Node;
 import org.elasticsearch.xpack.sql.type.DataType;
 import org.elasticsearch.xpack.sql.util.StringUtils;
@@ -65,8 +65,8 @@ public String message() {
     private Boolean lazyChildrenResolved = null;
     private Expression lazyCanonical = null;
 
-    public Expression(Location location, List<Expression> children) {
-        super(location, children);
+    public Expression(Source source, List<Expression> children) {
+        super(source, children);
     }
 
     // whether the expression can be evaluated statically (folded) or not

x-pack/plugin/sql/src/main/java/org/elasticsearch/xpack/sql/expression/Expression.java

@@ -35,7 +35,7 @@ public enum ParamOrdinal {
     private Expressions() {}
 
     public static NamedExpression wrapAsNamed(Expression exp) {
-        return exp instanceof NamedExpression ? (NamedExpression) exp : new Alias(exp.location(), exp.nodeName(), exp);
+        return exp instanceof NamedExpression ? (NamedExpression) exp : new Alias(exp.source(), exp.nodeName(), exp);
     }
 
     public static List<Attribute> asAttributes(List<? extends NamedExpression> named) {