feat: add PromQL to Composer

Author: vadimkibana

src/composer/__tests__/esql.promql.test.ts

@@ -0,0 +1,122 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { esql } from '../esql';
+import { pqlFunc, pqlSel, pqlNum } from '../synth';
+
+test('simple instant vector', () => {
+  const query = esql.promql(pqlSel('up'));
+
+  expect(query.print()).toBe('PROMQL (up)');
+});
+
+test('range vector selector', () => {
+  const query = esql.promql(pqlSel('http_requests_total', '5m'));
+
+  expect(query.print()).toBe('PROMQL (http_requests_total[5m])');
+});
+
+test('from PromQL synth-created node', () => {
+  const query = esql.promql(esql.pql`http_requests_total[5m]`);
+
+  expect(query.print()).toBe('PROMQL (http_requests_total[5m])');
+});
+
+test('from PromQL synth-created node - 2', () => {
+  const query = esql.promql(esql.pql`http_requests_total[5m]`, { params: { a: 'b' } });
+
+  expect(query.print()).toBe('PROMQL a = b (http_requests_total[5m])');
+});
+
+test('from PromQL synth-created node - 3', () => {
+  const query = esql.promql(esql.pql`http_requests_total[5m]`, {
+    params: { a: 'b' },
+    outputName: 'result',
+  });
+
+  expect(query.print()).toBe('PROMQL a = b result = (http_requests_total[5m])');
+});
+
+test('function expression', () => {
+  const query = esql.promql(pqlFunc('rate', pqlSel('http_requests_total', '5m')));
+
+  expect(query.print()).toBe('PROMQL (rate(http_requests_total[5m]))');
+});
+
+test('aggregation with grouping', () => {
+  const expr = pqlFunc('sum', pqlFunc('rate', pqlSel('metric', '5m')), { by: ['job'] });
+  const query = esql.promql(expr);
+
+  expect(query.print()).toBe('PROMQL (sum(rate(metric[5m])) by (job))');
+});
+
+test('accepts a raw PromQL string', () => {
+  const query = esql.promql('rate(http_requests_total[5m])');
+
+  expect(query.print()).toBe('PROMQL (rate(http_requests_total[5m]))');
+});
+
+test('accepts a simple metric string', () => {
+  const query = esql.promql('up');
+
+  expect(query.print()).toBe('PROMQL (up)');
+});
+
+test('accepts a pql-tagged expression', () => {
+  const expr = esql.pql`sum(rate(metric[5m])) by (job)`;
+  const query = esql.promql(expr);
+
+  expect(query.print()).toBe('PROMQL (sum(rate(metric[5m])) by (job))');
+});
+
+test('single command-level param', () => {
+  const query = esql.promql(pqlSel('up'), { params: { index: 'k8s' } });
+
+  expect(query.print()).toBe('PROMQL index = k8s (up)');
+});
+
+test('multiple command-level params', () => {
+  const query = esql.promql(pqlSel('up'), { params: { index: 'k8s', timeout: '10s' } });
+
+  expect(query.print()).toBe('PROMQL index = k8s timeout = 10s (up)');
+});
+
+test('named output column', () => {
+  const query = esql.promql(pqlFunc('sum', pqlSel('metric', '5m')), { outputName: 'result' });
+
+  expect(query.print()).toBe('PROMQL result = (sum(metric[5m]))');
+});
+
+test('params and output name together', () => {
+  const query = esql.promql(pqlSel('up'), {
+    params: { index: 'k8s' },
+    outputName: 'health',
+  });
+
+  expect(query.print()).toBe('PROMQL index = k8s health = (up)');
+});
+
+test('result is a ComposerQuery that can be piped', () => {
+  const query = esql.promql(pqlSel('up')).limit(100);
+
+  expect(query.print()).toBe('PROMQL (up) | LIMIT 100');
+});
+
+test('result can be sorted', () => {
+  const query = esql.promql(pqlFunc('rate', pqlSel('metric', '5m')))
+    .where`value > ${pqlNum(0)}`.limit(50);
+
+  expect(query.print()).toBe('PROMQL (rate(metric[5m])) | WHERE value > 0 | LIMIT 50');
+});
+
+test('first command is a PROMQL command', () => {
+  const query = esql.promql(pqlSel('up'));
+  const [cmd] = query.ast.commands;
+
+  expect(cmd.type).toBe('command');
+  expect(cmd.name).toBe('promql');
+});

src/composer/esql.ts

@@ -19,6 +19,8 @@ import type {
   FromSourcesAndMetadataQueryStarter,
   FromSourcesQueryStarter,
   ParametrizedComposerQueryTag,
+  PromqlQueryStarter,
+  PromqlStarterOpts,
 } from './types';
 import type { ESQLSource } from '../types';
 import { isSource } from '../ast/is';
@@ -177,6 +179,28 @@ const createFromSourceAndMetadataCommandStarter =
       : esql`${synth.kwd(cmd)} ${sourceNodes}`;
   };
 
+const createPromqlStarter = (): PromqlQueryStarter => (expression, opts?: PromqlStarterOpts) => {
+  const expr = typeof expression === 'string' ? synth.pql(expression) : expression;
+
+  const paramsPart = opts?.params
+    ? Object.entries(opts.params)
+        .map(([k, v]) => `${k}=${v}`)
+        .join(' ')
+    : '';
+
+  const outputName = opts?.outputName;
+
+  if (paramsPart && outputName) {
+    return esql`PROMQL ${synth.kwd(paramsPart)} ${synth.col(outputName)} = (${expr as synth.SynthTemplateHole})`;
+  } else if (paramsPart) {
+    return esql`PROMQL ${synth.kwd(paramsPart)} (${expr as synth.SynthTemplateHole})`;
+  } else if (outputName) {
+    return esql`PROMQL ${synth.col(outputName)} = (${expr as synth.SynthTemplateHole})`;
+  } else {
+    return esql`PROMQL (${expr as synth.SynthTemplateHole})`;
+  }
+};
+
 /**
  * ESQL query composer tag function.
  *
@@ -229,6 +253,8 @@ export const esql: ComposerQueryTag &
 
     ts: createFromLikeStarter('TS'),
 
+    promql: createPromqlStarter(),
+
     get nop() {
       return synth.cmd`WHERE TRUE`;
     },

src/composer/index.ts

@@ -11,6 +11,6 @@ export { ComposerQuery } from './composer_query';
 export { ParameterHole } from './parameter_hole';
 
 export * as synth from './synth';
-export { qry, cmd, exp } from './synth';
+export { qry, cmd, exp, pql } from './synth';
 
 export { EsqlQuery } from './query';

src/composer/synth/__tests__/promql.test.ts

@@ -0,0 +1,78 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+import { BasicPrettyPrinter } from '../../../pretty_print';
+import { PromQLBasicPrettyPrinter } from '../../../embedded_languages/promql/pretty_print';
+import { pql, promqlExpression } from '../promql';
+import { cmd } from '../command';
+import { qry } from '../query';
+
+test('creates a PromQL expression node from a tagged template', () => {
+  const node = pql`rate(http_requests_total[5m])`;
+
+  expect(node).toBeDefined();
+  expect(node.type).toBe('query');
+  expect(node.dialect).toBe('promql');
+});
+
+test('pql is an alias for promqlExpression', () => {
+  expect(pql).toBe(promqlExpression);
+});
+
+test('serializes a PromQL expression back to text', () => {
+  const node = pql`rate(http_requests_total[5m])`;
+  const text = PromQLBasicPrettyPrinter.print(node);
+
+  expect(text).toBe('rate(http_requests_total[5m])');
+});
+
+test('serializes a metric selector', () => {
+  const node = pql`http_requests_total{job="api",status="200"}`;
+  const text = PromQLBasicPrettyPrinter.print(node);
+
+  expect(text).toBe('http_requests_total{job="api", status="200"}');
+});
+
+test('can be used as a function call', () => {
+  const node = pql('up == 1');
+  const text = PromQLBasicPrettyPrinter.print(node);
+
+  expect(text).toBe('up == 1');
+});
+
+test('can compose PromQL expressions', () => {
+  const time = pql`5m`;
+  const selector = pql`some_metric{job="api"}[${time}]`;
+  const number = pql`456`;
+  const func = pql`func(${selector}, 123, ${number})`;
+  const text = func + '';
+
+  expect(text).toBe('func(some_metric{job="api"}[5m], 123, 456)');
+});
+
+test('can be interpolated as a hole in a PROMQL command', () => {
+  const expr = pql`bytes_in{job="prometheus"}`;
+  const node = cmd`PROMQL (${expr})`;
+  const text = BasicPrettyPrinter.command(node);
+
+  expect(text).toBe('PROMQL (bytes_in{job="prometheus"})');
+});
+
+test('can build a full PROMQL query via qry', () => {
+  const expr = pql`rate(http_requests_total[5m])`;
+  const node = qry`PROMQL (${expr})`;
+  const text = BasicPrettyPrinter.print(node);
+
+  expect(text).toBe('PROMQL (rate(http_requests_total[5m]))');
+});
+
+test('trims whitespace from source', () => {
+  const node = pql`  up  `;
+  const text = PromQLBasicPrettyPrinter.print(node);
+
+  expect(text).toBe('up');
+});