feat(tracemetrics): Use reference map to render equations

static/app/views/explore/metrics/equationBuilder.spec.tsx

@@ -0,0 +1,61 @@
+import {render, screen, userEvent} from 'sentry-test/reactTestingLibrary';
+
+import {EquationBuilder} from 'sentry/views/explore/metrics/equationBuilder';
+
+describe('EquationBuilder', () => {
+  it('takes an equation and represents the equation using the provided reference map', async () => {
+    const expression = 'count(metricA) + sum(metricB)';
+    const referenceMap = {
+      A: 'count(metricA)',
+      F: 'sum(metricB)',
+    };
+
+    render(
+      <EquationBuilder
+        expression={expression}
+        referenceMap={referenceMap}
+        handleExpressionChange={() => {}}
+      />
+    );
+
+    const tokens = await screen.findAllByRole('row');
+
+    // tokens are flanked by empty text nodes for cursor movement,
+    // and '+' is encoded as 'op:0'
+    expect(tokens.map(token => token.getAttribute('aria-label'))).toEqual([
+      '',
+      'A',
+      '',
+      'op:0',
+      '',
+      'F',
+      '',
+    ]);
+  });
+
+  it('calls the handleExpressionChange callback when the expression changes', async () => {
+    const expression = '';
+
+    const handleExpressionChange = jest.fn();
+
+    render(
+      <EquationBuilder
+        expression={expression}
+        referenceMap={{A: 'count(value,metricA,distribution,none)'}}
+        handleExpressionChange={handleExpressionChange}
+      />
+    );
+
+    // Typing the reference de-references it into the metric call
+    await userEvent.type(
+      await screen.findByRole('combobox', {name: 'Add a term'}),
+      'A * 2'
+    );
+
+    expect(handleExpressionChange).toHaveBeenCalledWith(
+      expect.objectContaining({
+        text: 'count(value,metricA,distribution,none) * 2',
+      })
+    );
+  });
+});

static/app/views/explore/metrics/equationBuilder.tsx

@@ -0,0 +1,137 @@
+import {useCallback, useEffect, useMemo, useRef, useTransition} from 'react';
+import isEqual from 'lodash/isEqual';
+
+import {ArithmeticBuilder} from 'sentry/components/arithmeticBuilder';
+import {Expression} from 'sentry/components/arithmeticBuilder/expression';
+import {
+  isTokenFreeText,
+  isTokenReference,
+} from 'sentry/components/arithmeticBuilder/token';
+import {tokenizeExpression} from 'sentry/components/arithmeticBuilder/tokenizer';
+
+/**
+ * Takes an expression and map of references and returns the internal string representation that uses the references.
+ */
+function unresolveExpression(
+  expression: string,
+  referenceMap: Record<string, string> = {}
+): string {
+  // Reverse the keys and values of the reference map, duplicates keep the first reference found
+  const reversedReferenceMap = Object.entries(referenceMap).reduce(
+    (reversedMap: Record<string, string>, [key, value]) => {
+      if (!reversedMap[value]) {
+        reversedMap[value] = key;
+      }
+      return reversedMap;
+    },
+    {}
+  );
+
+  const tokens = tokenizeExpression(expression);
+  return tokens
+    .map(token => {
+      if (token.text in reversedReferenceMap) {
+        return reversedReferenceMap[token.text] ?? token.text;
+      }
+      return token.text;
+    })
+    .join(' ');
+}
+
+/**
+ * Resolves the expression using references into a format that is compatible with our querying endpoints.
+ */
+function resolveExpression(
+  expression: Expression,
+  referenceMap: Record<string, string> = {}
+): Expression {
+  const newTokens = expression.tokens
+    .map(token => {
+      if (referenceMap && isTokenReference(token) && token.text in referenceMap) {
+        return referenceMap[token.text];
+      }
+      if (!isTokenFreeText(token)) {
+        return token.text;
+      }
+      return null;
+    })
+    .filter(Boolean);
+
+  return new Expression(newTokens.join(' '));
+}
+
+/**
+ * A component that takes an equation in full resolved form and allows
+ * the user to edit it using "references" to refer to the different components
+ * of the equation.
+ *
+ * The references are used to resolve the equation into a format that is
+ * compatible with our querying endpoints.
+ */
+export function EquationBuilder({
+  expression,
+  referenceMap,
+  handleExpressionChange,
+}: {
+  expression: string;
+  handleExpressionChange: (expression: Expression) => void;
+  referenceMap?: Record<string, string>;
+}) {
+  const [_, startTransition] = useTransition();
+  const references = useMemo(
+    () => new Set(Object.keys(referenceMap ?? {})),
+    [referenceMap]
+  );
+
+  // Tracks the reference map that `expression` was last resolved against.
+  // When referenceMap changes externally, expression still contains values
+  // resolved against the previous map until we re-resolve and the parent updates.
+  const expressionMapRef = useRef(referenceMap);
+  const mapChanged = !isEqual(expressionMapRef.current, referenceMap);
+
+  const internalExpression = unresolveExpression(
+    expression,
+    mapChanged ? expressionMapRef.current : referenceMap
+  );
+
+  // When the reference map changes, re-resolve the expression and invoke the callback
+  useEffect(() => {
+    if (!isEqual(expressionMapRef.current, referenceMap)) {
+      const expr = new Expression(
+        internalExpression,
+        new Set(Object.keys(expressionMapRef.current ?? {}))
+      );
+      const resolved = resolveExpression(expr, referenceMap);
+
+      // Check the validity of the internal expression, not the resolved one because
+      // there are issues with validating it with the new _if aggregation format and
+      // this check just needs to ensure the structure is valid.
+      if (expr.isValid) {
+        handleExpressionChange(resolved);
+        expressionMapRef.current = referenceMap;
+      }
+    }
+  }, [referenceMap, internalExpression, handleExpressionChange]);
+
+  const handleInternalExpressionChange = useCallback(
+    (newExpression: Expression) => {
+      startTransition(() => {
+        if (newExpression.isValid) {
+          handleExpressionChange(resolveExpression(newExpression, referenceMap));
+        }
+      });
+    },
+    [handleExpressionChange, referenceMap]
+  );
+
+  return (
+    <ArithmeticBuilder
+      aggregations={[]}
+      expression={internalExpression}
+      functionArguments={[]}
+      getFieldDefinition={() => null}
+      references={references}
+      setExpression={handleInternalExpressionChange}
+    />
+  );
+}

static/app/views/explore/metrics/hooks/useMetricReferences.spec.tsx

@@ -0,0 +1,78 @@
+import type {ReactNode} from 'react';
+
+import {renderHookWithProviders} from 'sentry-test/reactTestingLibrary';
+
+import {Mode} from 'sentry/views/explore/contexts/pageParamsContext/mode';
+import {encodeMetricQueryParams} from 'sentry/views/explore/metrics/metricQuery';
+import {
+  MultiMetricsQueryParamsProvider,
+  useMultiMetricsQueryParams,
+} from 'sentry/views/explore/metrics/multiMetricsQueryParams';
+import {ReadableQueryParams} from 'sentry/views/explore/queryParams/readableQueryParams';
+import {VisualizeFunction} from 'sentry/views/explore/queryParams/visualize';
+
+import {useMetricReferences} from './useMetricReferences';
+
+function Wrapper({children}: {children: ReactNode}) {
+  return <MultiMetricsQueryParamsProvider>{children}</MultiMetricsQueryParamsProvider>;
+}
+
+describe('useMetricReferences', () => {
+  it('returns _if form for metrics with a filter and plain yAxis for metrics without', () => {
+    const metricWithFilter = encodeMetricQueryParams({
+      metric: {name: 'metric_a', type: 'counter', unit: 'none'},
+      queryParams: new ReadableQueryParams({
+        extrapolate: true,
+        mode: Mode.SAMPLES,
+        query: 'status:ok',
+        cursor: '',
+        fields: ['id', 'timestamp'],
+        sortBys: [{field: 'timestamp', kind: 'desc'}],
+        aggregateCursor: '',
+        aggregateFields: [new VisualizeFunction('count(value,metric_a,counter,none)')],
+        aggregateSortBys: [{field: 'count(value,metric_a,counter,none)', kind: 'desc'}],
+      }),
+    });
+
+    const metricWithoutFilter = encodeMetricQueryParams({
+      metric: {name: 'metric_b', type: 'counter', unit: 'none'},
+      queryParams: new ReadableQueryParams({
+        extrapolate: true,
+        mode: Mode.SAMPLES,
+        query: '',
+        cursor: '',
+        fields: ['id', 'timestamp'],
+        sortBys: [{field: 'timestamp', kind: 'desc'}],
+        aggregateCursor: '',
+        aggregateFields: [new VisualizeFunction('sum(value,metric_b,counter,none)')],
+        aggregateSortBys: [{field: 'sum(value,metric_b,counter,none)', kind: 'desc'}],
+      }),
+    });
+
+    const {result} = renderHookWithProviders(
+      () => ({
+        references: useMetricReferences(),
+        queries: useMultiMetricsQueryParams(),
+      }),
+      {
+        additionalWrapper: Wrapper,
+        initialRouterConfig: {
+          location: {
+            pathname: '/',
+            query: {
+              metric: [metricWithFilter, metricWithoutFilter],
+            },
+          },
+        },
+      }
+    );
+
+    // Metric A has a filter, so it gets the _if form
+    expect(result.current.references.A).toBe(
+      'count_if(`status:ok`,value,metric_a,counter,none)'
+    );
+
+    // Metric B has no filter, so the yAxis is returned as-is
+    expect(result.current.references.B).toBe('sum(value,metric_b,counter,none)');
+  });
+});

static/app/views/explore/metrics/hooks/useMetricReferences.tsx

@@ -1,19 +1,36 @@
 import {useMemo} from 'react';
 
+import {parseFunction} from 'sentry/utils/discover/fields';
 import {useMultiMetricsQueryParams} from 'sentry/views/explore/metrics/multiMetricsQueryParams';
+import type {ReadableQueryParams} from 'sentry/views/explore/queryParams/readableQueryParams';
 import {isVisualizeFunction} from 'sentry/views/explore/queryParams/visualize';
 import {getVisualizeLabel} from 'sentry/views/explore/toolbar/toolbarVisualize';
 
+function resolveMetricReference(metricConfig: ReadableQueryParams): string {
+  if (metricConfig.query) {
+    // convert the aggregate to an "if" format
+    const visualize = metricConfig.visualizes[0]!;
+    const parsed = parseFunction(visualize.yAxis);
+    if (parsed) {
+      return `${parsed.name}_if(\`${metricConfig.query}\`,${parsed.arguments.join(',')})`;
+    }
+  }
+  return metricConfig.visualizes[0]?.yAxis!;
+}
+
 export function useMetricReferences() {
   const metricQueries = useMultiMetricsQueryParams();
 
   return useMemo(() => {
-    return new Set(
+    return Object.fromEntries(
       metricQueries
         .filter(metricQuery =>
           metricQuery.queryParams.visualizes.some(isVisualizeFunction)
         )
-        .map((metricQuery, index) => metricQuery.label ?? getVisualizeLabel(index, false))
+        .map((metricQuery, index) => [
+          metricQuery.label ?? getVisualizeLabel(index, false),
+          resolveMetricReference(metricQuery.queryParams),
+        ])
     );
   }, [metricQueries]);
 }

static/app/views/explore/metrics/hooks/useMetricTimeseries.tsx

@@ -16,6 +16,7 @@ import {
   useQueryParamsGroupBys,
   useQueryParamsSearch,
 } from 'sentry/views/explore/queryParams/context';
+import {isVisualizeEquation} from 'sentry/views/explore/queryParams/visualize';
 import {useSortedTimeSeries} from 'sentry/views/insights/common/queries/useSortedTimeSeries';
 
 interface UseMetricTimeseriesOptions {
@@ -68,7 +69,12 @@ function useMetricTimeseriesImpl({
       yAxis,
       interval,
       fields: [...groupBys, ...yAxis],
-      enabled: enabled && Boolean(traceMetric.name),
+      enabled:
+        enabled &&
+        (Boolean(traceMetric.name) ||
+          visualizes.some(
+            visualize => isVisualizeEquation(visualize) && visualize.expression.text
+          )),
       topEvents,
       orderby: sortBys.map(formatSort),
       ...queryExtras,

static/app/views/explore/metrics/hooks/useSaveMetricsMultiQuery.tsx

@@ -13,7 +13,10 @@ import {getTitleFromLocation} from 'sentry/views/explore/contexts/pageParamsCont
 import {useInvalidateSavedQueries} from 'sentry/views/explore/hooks/useGetSavedQueries';
 import {useMultiMetricsQueryParams} from 'sentry/views/explore/metrics/multiMetricsQueryParams';
 import {isGroupBy} from 'sentry/views/explore/queryParams/groupBy';
-import {isVisualize} from 'sentry/views/explore/queryParams/visualize';
+import {
+  isVisualize,
+  isVisualizeFunction,
+} from 'sentry/views/explore/queryParams/visualize';
 
 const METRICS_DATASET = 'metrics';
 
@@ -70,7 +73,7 @@ export function useSaveMetricsMultiQuery() {
               ...groupBys,
               ...(yAxes.length > 0 ? [{yAxes, chartType}] : []),
             ],
-            metric: metricQuery.metric,
+            ...(isVisualizeFunction(visualize) ? {metric: metricQuery.metric} : {}),
             fields: metricQuery.queryParams.fields,
             orderby: metricQuery.queryParams.sortBys[0]
               ? encodeSort(metricQuery.queryParams.sortBys[0])