Adjust source with offsets (#137260)

Author: felixbarny

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/optimizer/LogicalPlanOptimizer.java

@@ -118,10 +118,10 @@ public LogicalPlanOptimizer(LogicalOptimizerContext optimizerContext) {
     public LogicalPlan optimize(LogicalPlan verified) {
         var optimized = execute(verified);
 
-//        Failures failures = verifier.verify(optimized, verified.output());
-//        if (failures.hasFailures()) {
-//            throw new VerificationException(failures);
-//        }
+        // Failures failures = verifier.verify(optimized, verified.output());
+        // if (failures.hasFailures()) {
+        // throw new VerificationException(failures);
+        // }
         optimized.setOptimized();
         return optimized;
     }

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

@@ -8,6 +8,7 @@
 package org.elasticsearch.xpack.esql.parser;
 
 import org.antlr.v4.runtime.ParserRuleContext;
+import org.antlr.v4.runtime.Token;
 import org.antlr.v4.runtime.tree.ParseTree;
 import org.apache.lucene.util.BytesRef;
 import org.elasticsearch.Build;
@@ -52,6 +53,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.logical.Not;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.Equals;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.EsqlBinaryComparison;
+import org.elasticsearch.xpack.esql.parser.promql.ParsingUtils;
 import org.elasticsearch.xpack.esql.plan.EsqlStatement;
 import org.elasticsearch.xpack.esql.plan.IndexPattern;
 import org.elasticsearch.xpack.esql.plan.QuerySetting;
@@ -1266,46 +1268,33 @@ public PlanFactory visitPromqlCommand(EsqlBaseParser.PromqlCommandContext ctx) {
 
         // TODO: Perform type and value validation
         var queryCtx = ctx.promqlQueryPart();
+        if (queryCtx == null || queryCtx.isEmpty()) {
+            throw new ParsingException(source, "PromQL expression cannot be empty");
+        }
 
-        String promqlQuery = queryCtx == null || queryCtx.isEmpty()
-            ? StringUtils.EMPTY
-            // copy the query verbatim to avoid missing tokens interpreted by the enclosing lexer
-            : source(queryCtx.get(0).start, queryCtx.get(queryCtx.size() - 1).stop).text();
-
-        if (promqlQuery.isEmpty()) {
+        Token startToken = queryCtx.getFirst().start;
+        Token stopToken = queryCtx.getLast().stop;
+        // copy the query verbatim to avoid missing tokens interpreted by the enclosing lexer
+        String promqlQuery = source(startToken, stopToken).text();
+        if (promqlQuery.isBlank()) {
             throw new ParsingException(source, "PromQL expression cannot be empty");
         }
 
-        int promqlStartLine = source.source().getLineNumber();
-        int promqlStartColumn = queryCtx != null && !queryCtx.isEmpty()
-            ? queryCtx.get(0).start.getCharPositionInLine()
-            : source.source().getColumnNumber();
+        int promqlStartLine = startToken.getLine();
+        int promqlStartColumn = startToken.getCharPositionInLine();
 
         PromqlParser promqlParser = new PromqlParser();
         LogicalPlan promqlPlan;
         try {
             // The existing PromqlParser is used to parse the inner query
-            promqlPlan = promqlParser.createStatement(promqlQuery);
+            promqlPlan = promqlParser.createStatement(promqlQuery, null, null, promqlStartLine, promqlStartColumn);
         } catch (ParsingException pe) {
-            throw getParsingException(pe, promqlStartLine, promqlStartColumn);
+            throw ParsingUtils.adjustParsingException(pe, promqlStartLine, promqlStartColumn);
         }
 
         return plan -> time != null
             ? new PromqlCommand(source, plan, promqlPlan, time)
             : new PromqlCommand(source, plan, promqlPlan, start, end, step);
     }
 
-    private static ParsingException getParsingException(ParsingException pe, int promqlStartLine, int promqlStartColumn) {
-        int adjustedLine = promqlStartLine + (pe.getLineNumber() - 1);
-        int adjustedColumn = (pe.getLineNumber() == 1 ? promqlStartColumn + pe.getColumnNumber() : pe.getColumnNumber()) - 1;
-
-        ParsingException adjusted = new ParsingException(
-            pe.getErrorMessage(),
-            pe.getCause() instanceof Exception ? (Exception) pe.getCause() : null,
-            adjustedLine,
-            adjustedColumn
-        );
-        adjusted.setStackTrace(pe.getStackTrace());
-        return adjusted;
-    }
 }

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

@@ -46,24 +46,26 @@ public class PromqlParser {
      * Parses an PromQL expression into execution plan
      */
     public LogicalPlan createStatement(String query) {
-        return createStatement(query, null, null);
+        return createStatement(query, null, null, 0, 0);
     }
 
-    public LogicalPlan createStatement(String query, Instant start, Instant end) {
+    public LogicalPlan createStatement(String query, Instant start, Instant end, int startLine, int startColumn) {
         if (log.isDebugEnabled()) {
             log.debug("Parsing as expression: {}", query);
         }
 
         if (start == null) {
             start = Instant.now(UTC);
         }
-        return invokeParser(query, start, end, PromqlBaseParser::singleStatement, PromqlAstBuilder::plan);
+        return invokeParser(query, start, end, startLine, startColumn, PromqlBaseParser::singleStatement, PromqlAstBuilder::plan);
     }
 
     private <T> T invokeParser(
         String query,
         Instant start,
         Instant end,
+        int startLine,
+        int startColumn,
         Function<PromqlBaseParser, ParserRuleContext> parseFunction,
         BiFunction<PromqlAstBuilder, ParserRuleContext, T> visitor
     ) {
@@ -97,7 +99,7 @@ private <T> T invokeParser(
             if (log.isTraceEnabled()) {
                 log.trace("Parse tree: {}", tree.toStringTree());
             }
-            return visitor.apply(new PromqlAstBuilder(start, end), tree);
+            return visitor.apply(new PromqlAstBuilder(start, end, startLine, startColumn), tree);
         } catch (StackOverflowError e) {
             throw new ParsingException(
                 "PromQL statement is too large, causing stack overflow when generating the parsing tree: [{}]",

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

@@ -7,16 +7,23 @@
 
 package org.elasticsearch.xpack.esql.parser.promql;
 
+import org.antlr.v4.runtime.ParserRuleContext;
+import org.antlr.v4.runtime.Token;
 import org.antlr.v4.runtime.tree.ParseTree;
 import org.antlr.v4.runtime.tree.TerminalNode;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.parser.ParserUtils;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 import org.elasticsearch.xpack.esql.parser.PromqlBaseParserBaseVisitor;
 
-import static org.elasticsearch.xpack.esql.parser.ParserUtils.source;
-
 class AbstractBuilder extends PromqlBaseParserBaseVisitor<Object> {
+    private final int startLine, startColumn;
+
+    AbstractBuilder(int startLine, int startColumn) {
+        this.startLine = startLine;
+        this.startColumn = startColumn;
+    }
+
     @Override
     public Object visit(ParseTree tree) {
         return ParserUtils.visit(super::visit, tree);
@@ -25,7 +32,7 @@ public Object visit(ParseTree tree) {
     /**
      * Extracts the actual unescaped string (literal) value of a terminal node.
      */
-    static String string(TerminalNode node) {
+    String string(TerminalNode node) {
         return node == null ? null : unquote(source(node));
     }
 
@@ -42,4 +49,21 @@ public Object visitTerminal(TerminalNode node) {
         Source source = source(node);
         throw new ParsingException(source, "Does not know how to handle {}", source.text());
     }
+
+    protected Source source(TerminalNode terminalNode) {
+        return ParsingUtils.adjustSource(ParserUtils.source(terminalNode), startLine, startColumn);
+    }
+
+    protected Source source(ParserRuleContext parserRuleContext) {
+        return ParsingUtils.adjustSource(ParserUtils.source(parserRuleContext), startLine, startColumn);
+    }
+
+    protected Source source(Token token) {
+        return ParsingUtils.adjustSource(ParserUtils.source(token), startLine, startColumn);
+    }
+
+    protected Source source(Token start, Token stop) {
+        return ParsingUtils.adjustSource(ParserUtils.source(start, stop), startLine, startColumn);
+    }
+
 }

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

@@ -50,7 +50,6 @@
 import java.util.concurrent.TimeUnit;
 
 import static java.util.Collections.emptyList;
-import static org.elasticsearch.xpack.esql.parser.ParserUtils.source;
 import static org.elasticsearch.xpack.esql.parser.ParserUtils.typedParsing;
 import static org.elasticsearch.xpack.esql.parser.ParserUtils.visitList;
 import static org.elasticsearch.xpack.esql.parser.PromqlBaseParser.AND;
@@ -80,10 +79,11 @@ class ExpressionBuilder extends IdentifierBuilder {
     protected final Instant start, end;
 
     ExpressionBuilder() {
-        this(null, null);
+        this(null, null, 0, 0);
     }
 
-    ExpressionBuilder(Instant start, Instant end) {
+    ExpressionBuilder(Instant start, Instant end, int startLine, int startColumn) {
+        super(startLine, startColumn);
         Instant now = null;
         if (start == null || end == null) {
             now = DateUtils.nowWithMillisResolution().toInstant();

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

@@ -11,6 +11,10 @@
 
 abstract class IdentifierBuilder extends AbstractBuilder {
 
+    IdentifierBuilder(int startLine, int startColumn) {
+        super(startLine, startColumn);
+    }
+
     @Override
     public String visitIdentifier(IdentifierContext ctx) {
         return ctx == null ? null : ctx.getText();

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

@@ -9,6 +9,7 @@
 
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.core.Tuple;
+import org.elasticsearch.xpack.esql.core.tree.Location;
 import org.elasticsearch.xpack.esql.core.tree.Source;
 import org.elasticsearch.xpack.esql.parser.ParsingException;
 
@@ -224,4 +225,49 @@ private static String fromRadix(Source source, char[] chars, int offset, int cou
 
         return String.valueOf(Character.toChars(code));
     }
+
+    /**
+     * Adjusts the location of the source by the line and column offsets.
+     * @see #adjustLocation(Location, int, int)
+     */
+    public static Source adjustSource(Source source, int startLine, int startColumn) {
+        return new Source(adjustLocation(source.source(), startLine, startColumn), source.text());
+    }
+
+    /**
+     * Adjusts the location by the line and column offsets.
+     * The PromQL query inside the PROMQL command is parsed separately,
+     * so its line and column numbers need to be adjusted to match their
+     * position inside the full ES|QL query.
+     */
+    public static Location adjustLocation(Location location, int startLine, int startColumn) {
+        return new Location(adjustLine(location.getLineNumber(), startLine), adjustColumn(location.getColumnNumber(), startColumn));
+    }
+
+    /**
+     * Adjusts the line and column numbers of the given {@code ParsingException}
+     * by the provided offsets.
+     * The PromQL query inside the PROMQL command is parsed separately,
+     * so its line and column numbers need to be adjusted to match their
+     * position inside the full ES|QL query.
+     */
+    public static ParsingException adjustParsingException(ParsingException pe, int promqlStartLine, int promqlStartColumn) {
+        ParsingException adjusted = new ParsingException(
+            pe.getErrorMessage(),
+            pe.getCause() instanceof Exception ? (Exception) pe.getCause() : null,
+            adjustLine(pe.getLineNumber(), promqlStartLine),
+            adjustColumn(pe.getColumnNumber(), promqlStartColumn)
+        );
+        adjusted.setStackTrace(pe.getStackTrace());
+        return adjusted;
+    }
+
+    private static int adjustLine(int lineNumber, int startLine) {
+        return lineNumber + startLine - 1;
+    }
+
+    private static int adjustColumn(int columnNumber, int startColumn) {
+        // the column offset only applies to the first line of the PROMQL command
+        return columnNumber == 1 ? columnNumber + startColumn - 1 : columnNumber;
+    }
 }

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

@@ -20,11 +20,11 @@ public class PromqlAstBuilder extends PromqlLogicalPlanBuilder {
     private int expressionDepth = 0;
 
     public PromqlAstBuilder() {
-        this(null, null);
+        this(null, null, 0, 0);
     }
 
-    public PromqlAstBuilder(Instant start, Instant end) {
-        super(start, end);
+    public PromqlAstBuilder(Instant start, Instant end, int startLine, int startColumn) {
+        super(start, end, startLine, startColumn);
     }
 
     public LogicalPlan plan(ParseTree ctx) {

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

@@ -47,7 +47,6 @@
 import static java.util.Collections.emptyList;
 import static org.elasticsearch.xpack.esql.expression.promql.function.FunctionType.ACROSS_SERIES_AGGREGATION;
 import static org.elasticsearch.xpack.esql.expression.promql.function.FunctionType.WITHIN_SERIES_AGGREGATION;
-import static org.elasticsearch.xpack.esql.parser.ParserUtils.source;
 import static org.elasticsearch.xpack.esql.parser.ParserUtils.typedParsing;
 import static org.elasticsearch.xpack.esql.parser.ParserUtils.visitList;
 import static org.elasticsearch.xpack.esql.parser.PromqlBaseParser.AND;
@@ -70,11 +69,11 @@
 public class PromqlLogicalPlanBuilder extends ExpressionBuilder {
 
     PromqlLogicalPlanBuilder() {
-        this(null, null);
+        this(null, null, 0, 0);
     }
 
-    PromqlLogicalPlanBuilder(Instant start, Instant end) {
-        super(start, end);
+    PromqlLogicalPlanBuilder(Instant start, Instant end, int startLine, int startColumn) {
+        super(start, end, startLine, startColumn);
     }
 
     protected LogicalPlan plan(ParseTree ctx) {

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plan/logical/promql/PromqlCommand.java

@@ -102,9 +102,9 @@ public int hashCode() {
     public boolean equals(Object obj) {
         if (super.equals(obj)) {
 
-        PromqlCommand other = (PromqlCommand) obj;
-        return Objects.equals(child(), other.child()) && Objects.equals(promqlPlan, other.promqlPlan);
-    }
+            PromqlCommand other = (PromqlCommand) obj;
+            return Objects.equals(child(), other.child()) && Objects.equals(promqlPlan, other.promqlPlan);
+        }
 
         return false;
     }