解析PromQL并修改添加Label

最近做的项目中用到了Prometheus做预警服务，其中Prometheus使用promql语言来查询。项目中用户通过UI或者自己手动输入PromQL时候是缺少一些系统参数的，所以需要在用户输入完成以后同步到Prometheus时候将这部分缺失的信息给添加回去，这里就需要修改用户写的PromQL了。

实现思路是通过Antlr4来解析PromQL并修改。

https://github.com/antlr/grammars-v4

上述url是antlr官方提供的各个语言的语法定义文件，其中就包含我需要PromQL，将上述代码中的promql包中的两个g4文件拷贝到自己项目中，我对拷贝的PromQLLexer.g4文件中的最后的空格处做了处理改成如下内容否则重写以后会丢失原语句中的空格。

WS: [\r\t\n ]+ -> channel(HIDDEN);

新建maven项目，此处我就叫promql-parser。

pom部分内容如下：

    <properties><antlr4.version>4.13.1</antlr4.version></properties><dependencies><dependency><groupId>org.antlr</groupId><artifactId>antlr4-runtime</artifactId><version>${antlr4.version}</version></dependency></dependencies><build><plugins><plugin><groupId>org.antlr</groupId><artifactId>antlr4-maven-plugin</artifactId><version>${antlr4.version}</version><executions><execution><phase>generate-sources</phase><goals><goal>antlr4</goal></goals></execution></executions><configuration><visitor>true</visitor><sourceDirectory>src/main/antlr4/com/swan</sourceDirectory><outputDirectory>src/main/java/com/swan</outputDirectory></configuration></plugin></plugins></build>

将g4文件放置到项目的src/main/antlr4/packages中，这样后续生成的代码也自动会生成到src/main/java/packages中，`替换为自己的java包名，如下图：

执行生成代码：

生成的java代码，没有产生包名，需要手动修改一下代码，把包路径添加上。

现在开始编写解析PromQL并修改PromQL的代码：

public class ParserUtil {public static String addLabels(String promQL, Map<String, String> labels) {PromQLLexer lexer = new PromQLLexer(CharStreams.fromString(promQL));CommonTokenStream tokenStream = new CommonTokenStream(lexer);final TokenStreamRewriter rewriter = new TokenStreamRewriter(tokenStream);PromQLParser parser = new PromQLParser(tokenStream);ParseTreeWalker.DEFAULT.walk(new LabelMatcherListener(rewriter, labels), parser.expression());return rewriter.getText();}private static class LabelMatcherListener extends PromQLParserBaseListener {private final TokenStreamRewriter rewriter;private final Map<String, String> labels;public LabelMatcherListener(TokenStreamRewriter rewriter, Map<String, String> labels) {this.rewriter = rewriter;this.labels = new LinkedHashMap<>(labels);}@Overridepublic void enterInstantSelector(PromQLParser.InstantSelectorContext ctx) {if (labels.isEmpty()) {return;}PromQLParser.LabelMatcherListContext matcherListContext = ctx.labelMatcherList();if (null == matcherListContext) {if (null == ctx.LEFT_BRACE()) {rewriter.insertAfter(ctx.METRIC_NAME().getSymbol(),labels.entrySet().stream().map(e -> String.format("%s=\"%s\"", e.getKey(), e.getValue())).collect(Collectors.joining(",", "{", "}")));} else {rewriter.insertAfter(ctx.LEFT_BRACE().getSymbol(),labels.entrySet().stream().map(e -> String.format("%s=\"%s\"", e.getKey(), e.getValue())).collect(Collectors.joining(",")));}} else {for (PromQLParser.LabelMatcherContext matcherContext :matcherListContext.labelMatcher()) {String labelName = matcherContext.labelName().getText();if (labels.containsKey(labelName)) {rewriter.replace(matcherContext.STRING().getSymbol(),String.format("\"%s\"", labels.get(labelName)));labels.remove(labelName);}}if (!labels.isEmpty()) {rewriter.insertBefore(ctx.RIGHT_BRACE().getSymbol(),labels.entrySet().stream().map(e -> String.format("%s=\"%s\"", e.getKey(), e.getValue())).collect(Collectors.joining(",", ",", "")));}}}}
}

上述代码对应的单元测试代码：

public class ParserUtilTest {@Testpublic void test1() {Assertions.assertEquals("hello{a=\"b\"}",ParserUtil.addLabels("hello", Collections.singletonMap("a", "b")));}@Testpublic void test1_1() {HashMap<String, String> map = new LinkedHashMap<>();map.put("a", "b");map.put("c", "d");Assertions.assertEquals("hello{a=\"b\",c=\"d\"}", ParserUtil.addLabels("hello", map));}@Testpublic void test2() {Assertions.assertEquals("hello{a=\"b\"}",ParserUtil.addLabels("hello{}", Collections.singletonMap("a", "b")));}@Testpublic void test2_1() {HashMap<String, String> map = new LinkedHashMap<>();map.put("a", "b");map.put("c", "d");Assertions.assertEquals("hello{a=\"b\",c=\"d\"}", ParserUtil.addLabels("hello{}", map));}@Testpublic void test3() {Assertions.assertEquals("hello{a=\"b\",c=\"d\"}",ParserUtil.addLabels("hello{a=\"b\"}", Collections.singletonMap("c", "d")));}@Testpublic void test3_1() {HashMap<String, String> map = new LinkedHashMap<>();map.put("a", "b1");map.put("c", "d");Assertions.assertEquals("hello{a=\"b1\",c=\"d\"}", ParserUtil.addLabels("hello{a=\"b\"}", map));}
}

检查promQL表达式是否正确

 /*** 检查表示式是否正确** @param promql* @return*/public static boolean validatePromQL(String promql) {CharStream input = CharStreams.fromString(promql);PromQLLexer lexer = new PromQLLexer(input);CommonTokenStream tokens = new CommonTokenStream(lexer);PromQLParser parser = new PromQLParser(tokens);// 使用正确的入口规则进行解析ParserRuleContext tree = parser.expression();// 如果没有遇到任何语法错误，则认为查询是合法的return parser.getNumberOfSyntaxErrors() == 0;}

提取promeQL的指标名称

@Slf4j
public class MetricNameExtractor extends PromQLParserBaseListener {private Set<String> metricNames = new HashSet<>();@Overridepublic void enterEveryRule(ParserRuleContext ctx) {// 获取当前规则的名字String ruleName = PromQLParser.ruleNames[ctx.getRuleIndex()];if ("instantSelector".equals(ruleName)) {// 假设 metric_name 规则只有一个子节点，即标识符TerminalNode terminalNode = (TerminalNode) ctx.getChild(0);String metricName = terminalNode.getText();metricNames.add(metricName);log.info("enterEveryRule Found: {}==={}", metricName, terminalNode.getClass().getSimpleName());}}@Overridepublic void enterInstantSelector(PromQLParser.InstantSelectorContext ctx) {ParseTree parseTree = ctx.getChild(0);if (parseTree != null) {String metricName = parseTree.getText();metricNames.add(metricName);log.info("enterInstantSelector Found: {}==={}", metricName, parseTree.getClass().getSimpleName());}}@Overridepublic void enterMatrixSelector(PromQLParser.MatrixSelectorContext ctx) {ParseTree parseTree = ctx.getChild(0);if (parseTree != null) {int childCount = parseTree.getChildCount();if (childCount > 0) {parseTree = parseTree.getChild(0);String metricName = parseTree.getText();log.info("enterMatrixSelector Found: {}==={}", metricName, parseTree.getClass().getSimpleName());metricNames.add(metricName);}}}public Set<String> getMetricNames() {return metricNames;}}

public static Set<String> findPromQLMetricNames(String promql){// 初始化词法分析器和解析器PromQLLexer lexer = new PromQLLexer(CharStreams.fromString(promql));PromQLParser parser = new PromQLParser(new CommonTokenStream(lexer));// 解析输入并获取解析树ParseTree tree = parser.expression();// 创建监听器实例并遍历解析树MetricNameExtractor extractor = new MetricNameExtractor();ParseTreeWalker walker = new ParseTreeWalker();walker.walk(extractor, tree);// 输出所有找到的指标名称return extractor.getMetricNames();
}

https://blog.csdn.net/catcher92/article/details/122205513