align AST and pretty printer

Author: bartoval

src/ast/walker/__tests__/walker_promql.test.ts

@@ -417,6 +417,23 @@ describe('Walker PromQL support', () => {
     });
   });
 
+  describe('label param literal traversal', () => {
+    test('visitPromqlLiteral is called for a named param label value via Walker', () => {
+      const query = EsqlQuery.fromSrc('PROMQL bytes_in{job=?job}');
+      const literals: PromQLLiteral[] = [];
+
+      Walker.walk(query.ast, {
+        promql: {
+          visitPromqlLiteral: (node) => literals.push(node),
+        },
+      });
+
+      expect(literals).toHaveLength(1);
+      expect(literals[0].literalType).toBe('param');
+      expect(literals[0].value).toBe('job');
+    });
+  });
+
   describe('abort functionality with PromQL', () => {
     test('can abort PromQL traversal', () => {
       const query = EsqlQuery.fromSrc('PROMQL rate(sum(http_requests_total[5m]))');

src/embedded_languages/promql/ast/builder/builder.ts

@@ -17,6 +17,8 @@ import type {
   PromQLNumericLiteral,
   PromQLStringLiteral,
   PromQLTimeValue,
+  PromQLParamLiteral,
+  PromQLLabelValue,
   PromQLIdentifier,
   PromQLLabel,
   PromQLLabelMap,
@@ -271,6 +273,23 @@ export namespace PromQLBuilder {
           ...PromQLBuilder.parserFields(fromParser),
         };
       };
+
+      export const param = (
+        value: string | number,
+        fromParser?: Partial<AstNodeParserFields>
+      ): PromQLParamLiteral => {
+        const isPositional = typeof value === 'number';
+        return {
+          dialect: 'promql',
+          type: 'literal',
+          literalType: 'param',
+          paramKind: '?',
+          paramType: isPositional ? 'positional' : 'named',
+          name: '',
+          value,
+          ...PromQLBuilder.parserFields(fromParser),
+        };
+      };
     }
   }
 
@@ -302,7 +321,7 @@ export namespace PromQLBuilder {
   export const label = (
     labelName: PromQLLabelName,
     operator: PromQLLabelMatchOperator,
-    value: PromQLStringLiteral | undefined,
+    value: PromQLLabelValue | undefined,
     fromParser?: Partial<AstNodeParserFields>
   ): PromQLLabel => {
     return {

src/embedded_languages/promql/ast/walker/__tests__/promql_walker.test.ts

@@ -338,6 +338,21 @@ describe('Walker PromQL support', () => {
     });
   });
 
+  describe('label param literal traversal', () => {
+    test('visitPromqlLiteral is called for a named param label value', () => {
+      const query = PromQLParser.parse('bytes_in{job=?job}');
+      const literals: PromQLLiteral[] = [];
+
+      PromqlWalker.walk(query.root, {
+        visitPromqlLiteral: (node) => literals.push(node),
+      });
+
+      expect(literals).toHaveLength(1);
+      expect(literals[0].literalType).toBe('param');
+      expect(literals[0].value).toBe('job');
+    });
+  });
+
   describe('abort functionality with PromQL', () => {
     test('can abort PromQL traversal', () => {
       const query = PromQLParser.parse('rate(sum(http_requests_total[5m]))');

src/embedded_languages/promql/parser/__tests__/parser.test.ts

@@ -8,7 +8,9 @@
 import type {
   PromQLBinaryExpression,
   PromQLFunction,
+  PromQLLabel,
   PromQLLiteral,
+  PromQLParamLiteral,
   PromQLParens,
   PromQLSelector,
   PromQLStringLiteral,
@@ -234,6 +236,61 @@ describe('PromQL Parser', () => {
     });
   });
 
+  describe('label parameter values', () => {
+    it('parses a named label parameter (?job)', () => {
+      const result = PromQLParser.parse('http_requests_total{job=?job}');
+
+      expect(result.errors).toHaveLength(0);
+
+      const selector = result.root.expression as PromQLSelector;
+      expect(selector.type).toBe('selector');
+
+      const label = selector.labelMap?.args[0] as PromQLLabel;
+      expect(label.incomplete).toBe(false);
+
+      const value = label.value as PromQLParamLiteral;
+      expect(value.literalType).toBe('param');
+      expect(value.paramType).toBe('named');
+      expect(value.paramKind).toBe('?');
+      expect(value.value).toBe('job');
+      expect(value.text).toBe('?job');
+    });
+
+    it('parses a positional label parameter (?1)', () => {
+      const result = PromQLParser.parse('http_requests_total{job=?1}');
+
+      expect(result.errors).toHaveLength(0);
+
+      const selector = result.root.expression as PromQLSelector;
+      const label = selector.labelMap?.args[0] as PromQLLabel;
+      const value = label.value as PromQLParamLiteral;
+
+      expect(value.literalType).toBe('param');
+      expect(value.paramType).toBe('positional');
+      expect(value.paramKind).toBe('?');
+      expect(value.value).toBe(1);
+      expect(value.text).toBe('?1');
+    });
+
+    it('parses a regex matcher with a named parameter (?pattern)', () => {
+      const result = PromQLParser.parse('http_requests_total{job=~?pattern}');
+
+      expect(result.errors).toHaveLength(0);
+
+      const selector = result.root.expression as PromQLSelector;
+      const label = selector.labelMap?.args[0] as PromQLLabel;
+
+      expect(label.operator).toBe('=~');
+
+      const value = label.value as PromQLParamLiteral;
+      expect(value.literalType).toBe('param');
+      expect(value.paramType).toBe('named');
+      expect(value.paramKind).toBe('?');
+      expect(value.value).toBe('pattern');
+      expect(value.text).toBe('?pattern');
+    });
+  });
+
   describe('error handling', () => {
     it('returns errors for invalid syntax', () => {
       const result = PromQLParser.parse('rate(');

src/embedded_languages/promql/parser/cst_to_ast_converter.ts

@@ -381,12 +381,8 @@ export class PromQLCstToAstConverter {
       }
     }
 
-    const stringToken = ctx.STRING();
-    let value: ast.PromQLStringLiteral | undefined;
-
-    if (stringToken) {
-      value = this.fromStringToken(stringToken.symbol);
-    }
+    const labelValueCtx = ctx.labelValue();
+    const value = this.fromLabelValue(labelValueCtx);
 
     const node = PromQLBuilder.label(labelName, operator, value, this.getParserFields(ctx));
 
@@ -398,9 +394,29 @@ export class PromQLCstToAstConverter {
       node.incomplete = true;
     }
 
+    if (labelValueCtx?.exception) {
+      node.incomplete = true;
+    }
+
     return node;
   }
 
+  private fromLabelValue(ctx: cst.LabelValueContext | undefined): ast.PromQLLabelValue | undefined {
+    if (!ctx) return undefined;
+
+    const stringToken = ctx.STRING();
+    if (stringToken) {
+      return this.fromStringToken(stringToken.symbol);
+    }
+
+    const paramToken = ctx.NAMED_OR_POSITIONAL_PARAM();
+    if (paramToken) {
+      return this.fromNamedOrPositionalParamToken(paramToken.symbol);
+    }
+
+    return undefined;
+  }
+
   private fromLabelName(ctx: cst.LabelNameContext): ast.PromQLLabelName | undefined {
     const identCtx = ctx.identifier();
     if (identCtx) {
@@ -802,6 +818,16 @@ export class PromQLCstToAstConverter {
     );
   }
 
+  private fromNamedOrPositionalParamToken(token: antlr.Token): ast.PromQLParamLiteral {
+    const text = token.text ?? '';
+    const paramValue = text.slice(1);
+    const valueAsNumber = Number(paramValue);
+    const isPositional = String(valueAsNumber) === paramValue;
+    const value = isPositional ? valueAsNumber : paramValue;
+
+    return PromQLBuilder.expression.literal.param(value, this.createParserFieldsFromToken(token));
+  }
+
   private unquoteString(text: string): string {
     if (text.length < 2) return text;
 

src/embedded_languages/promql/pretty_print/__tests__/basic_pretty_printer.test.ts

@@ -58,6 +58,14 @@ describe('parsed PromQL BasicPrettyPrinter', () => {
       test('multiple label matchers', () => {
         assertReprint('{job="api", instance=~"localhost:.*", status!="error"}');
       });
+
+      test('label with named parameter value', () => {
+        assertReprint('http_requests_total{job=?job}');
+      });
+
+      test('label with regex matcher and named parameter', () => {
+        assertReprint('http_requests_total{job=~?pattern}');
+      });
     });
 
     describe('range vector selectors', () => {

src/embedded_languages/promql/pretty_print/__tests__/basic_pretty_printer_builder.test.ts

@@ -419,6 +419,18 @@ describe('synthetic AST (Builder-constructed)', () => {
     });
   });
 
+  describe('param label values', () => {
+    test('named param as label value', () => {
+      const metric = PromQLBuilder.identifier('http_requests_total');
+      const labelName = PromQLBuilder.identifier('job');
+      const labelValue = PromQLBuilder.expression.literal.param('job');
+      const label = PromQLBuilder.label(labelName, '=', labelValue);
+      const labelMap = PromQLBuilder.labelMap([label]);
+      const selector = PromQLBuilder.expression.selector.node({ metric, labelMap });
+      expect(PromQLBasicPrettyPrinter.expression(selector)).toBe('http_requests_total{job=?job}');
+    });
+  });
+
   describe('special cases', () => {
     test('unknown node', () => {
       const unknown = PromQLBuilder.unknown();

src/embedded_languages/promql/pretty_print/basic_pretty_printer.ts

@@ -19,6 +19,7 @@ import type {
   PromQLLiteral,
   PromQLModifier,
   PromQLNumericLiteral,
+  PromQLParamLiteral,
   PromQLParens,
   PromQLSelector,
   PromQLStringLiteral,
@@ -352,6 +353,8 @@ export class PromQLBasicPrettyPrinter {
         return this.printStringLiteral(node as PromQLStringLiteral);
       case 'time':
         return this.printTimeLiteral(node as PromQLTimeValue);
+      case 'param':
+        return this.printParamLiteral(node as PromQLParamLiteral);
       default:
         return String((node as PromQLNumericLiteral).value);
     }
@@ -408,6 +411,10 @@ export class PromQLBasicPrettyPrinter {
     return node.value;
   }
 
+  protected printParamLiteral(node: PromQLParamLiteral): string {
+    return `${node.paramKind}${node.value}`;
+  }
+
   protected printIdentifier(node: PromQLIdentifier): string {
     return node.name;
   }

src/embedded_languages/promql/pretty_print/wrapping_pretty_printer.ts

@@ -445,6 +445,11 @@ export class PromQLWrappingPrettyPrinter {
       case 'time':
         doc = printer.text((node as promql.PromQLTimeValue).value);
         break;
+      case 'param':
+        doc = printer.text(
+          `${(node as promql.PromQLParamLiteral).paramKind}${(node as promql.PromQLParamLiteral).value}`
+        );
+        break;
       default:
         doc = printer.text(String((node as promql.PromQLNumericLiteral).value));
     }

src/embedded_languages/promql/types.ts

@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import type { ESQLLocation, EditorError, ESQLAstBaseItem } from '../../types';
+import type { ESQLLocation, EditorError, ESQLAstBaseItem, ESQLParamLiteral } from '../../types';
 
 /**
  * All PromQL AST nodes have a `dialect: 'promql'` property to distinguish them
@@ -131,9 +131,9 @@ export interface PromQLLabel extends PromQLAstNodeBase {
   operator: PromQLLabelMatchOperator;
 
   /**
-   * The label value (string literal).
+   * The label value (string literal or named/positional parameter).
    */
-  value?: PromQLStringLiteral;
+  value?: PromQLLabelValue;
 }
 
 export type PromQLLabelMatchOperator = '=' | '!=' | '=~' | '!~';
@@ -359,7 +359,11 @@ export interface PromQLParens extends PromQLAstNodeBase<''> {
 
 // ------------------------------------------------------------------- literals
 
-export type PromQLLiteral = PromQLNumericLiteral | PromQLStringLiteral | PromQLTimeValue;
+export type PromQLLiteral =
+  | PromQLNumericLiteral
+  | PromQLStringLiteral
+  | PromQLTimeValue
+  | PromQLParamLiteral;
 
 /**
  * Represents a numeric literal (integer, decimal, or hexadecimal).
@@ -380,6 +384,28 @@ export interface PromQLStringLiteral extends PromQLAstNodeBase {
   valueUnquoted: string;
 }
 
+/**
+ * Represents an ES|QL named or positional parameter used as a PromQL label value.
+ *
+ * ```promql
+ * {job=?job}
+ * {job=?1}
+ * {job=~?pattern}
+ * ```
+ */
+export interface PromQLParamLiteral
+  extends PromQLAstNodeBase, ESQLParamLiteral<'named' | 'positional', '?'> {
+  type: 'literal';
+  literalType: 'param';
+  paramKind: '?';
+  paramType: 'named' | 'positional';
+}
+
+/**
+ * The value of a label matcher: either a string literal or a named/positional parameter.
+ */
+export type PromQLLabelValue = PromQLStringLiteral | PromQLParamLiteral;
+
 /**
  * Represents a time/duration value.
  *