feat(tracemetrics): Add skeleton for equation builder in Explore

static/app/views/explore/hooks/useAnalytics.tsx

@@ -39,7 +39,11 @@ import {
   useQueryParamsVisualizes,
 } from 'sentry/views/explore/queryParams/context';
 import type {ReadableQueryParams} from 'sentry/views/explore/queryParams/readableQueryParams';
-import {Visualize} from 'sentry/views/explore/queryParams/visualize';
+import {
+  isVisualizeEquation,
+  isVisualizeFunction,
+  Visualize,
+} from 'sentry/views/explore/queryParams/visualize';
 import {useSpansDataset} from 'sentry/views/explore/spans/spansQueryParams';
 import {
   combineConfidenceForSeries,
@@ -841,7 +845,13 @@ export function useMetricsPanelAnalytics({
   const query = useQueryParamsQuery();
   const groupBys = useQueryParamsGroupBys();
   const visualize = useMetricVisualize();
-  const aggregateFunctionBox = useBox(visualize.parsedFunction?.name ?? '');
+  const aggregateFunctionBox = useBox(
+    isVisualizeFunction(visualize)
+      ? (visualize.parsedFunction?.name ?? '')
+      : isVisualizeEquation(visualize)
+        ? 'equation'
+        : ''
+  );
 
   const tableError =
     mode === Mode.AGGREGATE

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

@@ -0,0 +1,6 @@
+import {useOrganization} from 'sentry/utils/useOrganization';
+
+export function useHasMetricEquations() {
+  const organization = useOrganization();
+  return organization.features.includes('tracemetrics-equations-in-explore');
+}

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

@@ -1,10 +1,15 @@
+import {EQUATION_PREFIX} from 'sentry/utils/discover/fields';
 import {Mode} from 'sentry/views/explore/contexts/pageParamsContext/mode';
 import {
   decodeMetricsQueryParams,
+  defaultMetricQuery,
   encodeMetricQueryParams,
 } from 'sentry/views/explore/metrics/metricQuery';
 import {ReadableQueryParams} from 'sentry/views/explore/queryParams/readableQueryParams';
-import {VisualizeFunction} from 'sentry/views/explore/queryParams/visualize';
+import {
+  VisualizeEquation,
+  VisualizeFunction,
+} from 'sentry/views/explore/queryParams/visualize';
 
 describe('decodeMetricsQueryParams', () => {
   it('parses all visualizes', () => {
@@ -114,3 +119,41 @@ describe('decodeMetricsQueryParams', () => {
     );
   });
 });
+
+describe('defaultMetricQuery', () => {
+  it('returns a default metric query', () => {
+    const result = defaultMetricQuery();
+    expect(result).toEqual({
+      metric: {name: '', type: ''},
+      queryParams: new ReadableQueryParams({
+        extrapolate: true,
+        mode: Mode.SAMPLES,
+        query: '',
+        cursor: '',
+        fields: ['id', 'timestamp'],
+        sortBys: [{field: 'timestamp', kind: 'desc'}],
+        aggregateCursor: '',
+        aggregateFields: [new VisualizeFunction('sum(value)')],
+        aggregateSortBys: [{field: 'sum(value)', kind: 'desc'}],
+      }),
+    });
+  });
+
+  it('returns a default metric query with an equation', () => {
+    const result = defaultMetricQuery({type: 'equation'});
+    expect(result).toEqual({
+      metric: {name: '', type: ''},
+      queryParams: new ReadableQueryParams({
+        extrapolate: true,
+        mode: Mode.SAMPLES,
+        query: '',
+        cursor: '',
+        fields: ['id', 'timestamp'],
+        sortBys: [{field: 'timestamp', kind: 'desc'}],
+        aggregateCursor: '',
+        aggregateFields: [new VisualizeEquation(EQUATION_PREFIX)],
+        aggregateSortBys: [{field: EQUATION_PREFIX, kind: 'desc'}],
+      }),
+    });
+  });
+});

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

@@ -1,7 +1,7 @@
 import type {Location} from 'history';
 
 import {defined} from 'sentry/utils';
-import type {Sort} from 'sentry/utils/discover/fields';
+import {EQUATION_PREFIX, type Sort} from 'sentry/utils/discover/fields';
 import {Mode} from 'sentry/views/explore/contexts/pageParamsContext/mode';
 import type {AggregateField} from 'sentry/views/explore/queryParams/aggregateField';
 import {validateAggregateSort} from 'sentry/views/explore/queryParams/aggregateSortBy';
@@ -11,6 +11,7 @@ import {
   isBaseVisualize,
   isVisualize,
   Visualize,
+  VisualizeEquation,
   VisualizeFunction,
 } from 'sentry/views/explore/queryParams/visualize';
 
@@ -103,7 +104,9 @@ export function encodeMetricQueryParams(metricQuery: BaseMetricQuery): string {
   });
 }
 
-export function defaultMetricQuery(): BaseMetricQuery {
+export function defaultMetricQuery({
+  type = 'aggregate',
+}: {type?: 'aggregate' | 'equation'} = {}): BaseMetricQuery {
   return {
     metric: {name: '', type: ''},
     queryParams: new ReadableQueryParams({
@@ -116,8 +119,11 @@ export function defaultMetricQuery(): BaseMetricQuery {
       sortBys: defaultSortBys(defaultFields()),
 
       aggregateCursor: '',
-      aggregateFields: defaultAggregateFields(),
-      aggregateSortBys: defaultAggregateSortBys(defaultAggregateFields()),
+      aggregateFields:
+        type === 'equation' ? [defaultAggregateEquation()] : defaultAggregateFields(),
+      aggregateSortBys: defaultAggregateSortBys(
+        type === 'equation' ? [defaultAggregateEquation()] : defaultAggregateFields()
+      ),
     }),
   };
 }
@@ -164,6 +170,10 @@ export function defaultAggregateFields(): AggregateField[] {
   return [defaultVisualize(), ...defaultGroupBys()];
 }
 
+function defaultAggregateEquation() {
+  return new VisualizeEquation(EQUATION_PREFIX);
+}
+
 export function defaultAggregateSortBys(aggregateFields: AggregateField[]): Sort[] {
   const visualize = aggregateFields.find(isVisualize);
   if (!defined(visualize)) {

static/app/views/explore/metrics/metricToolbar/aggregateDropdown.tsx

@@ -20,6 +20,7 @@ import {
   useSetMetricVisualizes,
 } from 'sentry/views/explore/metrics/metricsQueryParams';
 import {updateVisualizeYAxis} from 'sentry/views/explore/metrics/utils';
+import {isVisualizeFunction} from 'sentry/views/explore/queryParams/visualize';
 
 const MULTI_SELECT_GROUP_KEYS = new Set(['percentiles', 'stats']);
 
@@ -28,21 +29,30 @@ export function AggregateDropdown({traceMetric}: {traceMetric: TraceMetric}) {
   const visualizes = useMetricVisualizes();
   const setMetricVisualizes = useSetMetricVisualizes();
 
+  if (!isVisualizeFunction(visualize)) {
+    return null;
+  }
+
+  const narrowedVisualize = visualize;
   const groups = GROUPED_OPTIONS_BY_TYPE[traceMetric.type] ?? [];
-  const selectedNames = new Set(visualizes.map(v => v.parsedFunction?.name ?? ''));
+  const selectedNames = new Set(
+    visualizes.map(v => (isVisualizeFunction(v) ? (v.parsedFunction?.name ?? '') : ''))
+  );
 
   function handleChange(selectedOptions: Array<SelectOption<string>>) {
     if (selectedOptions.length === 0) {
       setMetricVisualizes([
         updateVisualizeYAxis(
-          visualize,
+          narrowedVisualize,
           DEFAULT_YAXIS_BY_TYPE[traceMetric.type]!,
           traceMetric
         ),
       ]);
     } else {
       setMetricVisualizes(
-        selectedOptions.map(o => updateVisualizeYAxis(visualize, o.value, traceMetric))
+        selectedOptions.map(o =>
+          updateVisualizeYAxis(narrowedVisualize, o.value, traceMetric)
+        )
       );
     }
   }

static/app/views/explore/metrics/metricToolbar/index.tsx

@@ -1,7 +1,10 @@
-import {useCallback} from 'react';
+import {Fragment, useCallback} from 'react';
 
 import {Container, Flex, Grid, Stack} from '@sentry/scraps/layout';
 
+import {ArithmeticBuilder} from 'sentry/components/arithmeticBuilder';
+import type {Expression} from 'sentry/components/arithmeticBuilder/expression';
+import {EQUATION_PREFIX} from 'sentry/utils/discover/fields';
 import {useOrganization} from 'sentry/utils/useOrganization';
 import {type TraceMetric} from 'sentry/views/explore/metrics/metricQuery';
 import {canUseMetricsUIRefresh} from 'sentry/views/explore/metrics/metricsFlags';
@@ -17,13 +20,18 @@ import {GroupBySelector} from 'sentry/views/explore/metrics/metricToolbar/groupB
 import {MetricSelector} from 'sentry/views/explore/metrics/metricToolbar/metricSelector';
 import {VisualizeLabel} from 'sentry/views/explore/metrics/metricToolbar/visualizeLabel';
 import {useMultiMetricsQueryParams} from 'sentry/views/explore/metrics/multiMetricsQueryParams';
+import {
+  isVisualizeEquation,
+  isVisualizeFunction,
+} from 'sentry/views/explore/queryParams/visualize';
 
 interface MetricToolbarProps {
   queryIndex: number;
+  references: Set<string>;
   traceMetric: TraceMetric;
 }
 
-export function MetricToolbar({traceMetric, queryIndex}: MetricToolbarProps) {
+export function MetricToolbar({traceMetric, queryIndex, references}: MetricToolbarProps) {
   const organization = useOrganization();
   const metricQueries = useMultiMetricsQueryParams();
   const visualize = useMetricVisualize();
@@ -34,6 +42,17 @@ export function MetricToolbar({traceMetric, queryIndex}: MetricToolbarProps) {
   const setTraceMetric = useSetTraceMetric();
   const canRemoveMetric = metricQueries.length > 1;
 
+  const handleExpressionChange = useCallback(
+    (newExpression: Expression) => {
+      const isValid = newExpression.isValid;
+      if (!isValid) {
+        return;
+      }
+      setVisualize(visualize.replace({yAxis: `${EQUATION_PREFIX}${newExpression.text}`}));
+    },
+    [setVisualize, visualize]
+  );
+
   if (canUseMetricsUIRefresh(organization)) {
     return (
       <Flex width="100%" align="start" gap="md" data-test-id="metric-toolbar">
@@ -42,23 +61,39 @@ export function MetricToolbar({traceMetric, queryIndex}: MetricToolbarProps) {
           visualize={visualize}
           onClick={toggleVisibility}
         />
-        <Stack flex="1" minWidth={0} gap="sm" width="100%">
-          <Flex minWidth={0} gap="xs" align="center">
-            <Container width="100%" maxWidth={canRemoveMetric ? '225px' : undefined}>
-              <MetricSelector traceMetric={traceMetric} onChange={setTraceMetric} />
-            </Container>
+        {isVisualizeFunction(visualize) ? (
+          <Fragment>
+            <Stack flex="1" minWidth={0} gap="sm" width="100%">
+              <Flex minWidth={0} gap="xs" align="center">
+                <Container width="100%" maxWidth={canRemoveMetric ? '225px' : undefined}>
+                  <MetricSelector traceMetric={traceMetric} onChange={setTraceMetric} />
+                </Container>
+                {canRemoveMetric && <DeleteMetricButton />}
+              </Flex>
+              <Flex flex="2 1 0" minWidth={0}>
+                <AggregateDropdown traceMetric={traceMetric} />
+              </Flex>
+              <Flex flex="3 1 0" minWidth={0}>
+                <GroupBySelector traceMetric={traceMetric} />
+              </Flex>
+              <Flex minWidth={0} width="100%">
+                <Filter traceMetric={traceMetric} />
+              </Flex>
+            </Stack>
+          </Fragment>
+        ) : isVisualizeEquation(visualize) ? (
+          <Flex direction="row" gap="sm" align="center" minWidth={0} width="100%">
+            <ArithmeticBuilder
+              aggregations={[]}
+              expression={visualize.expression.text}
+              functionArguments={[]}
+              getFieldDefinition={() => null}
+              references={references}
+              setExpression={handleExpressionChange}
+            />
             {canRemoveMetric && <DeleteMetricButton />}
           </Flex>
-          <Flex flex="2 1 0" minWidth={0}>
-            <AggregateDropdown traceMetric={traceMetric} />
-          </Flex>
-          <Flex flex="3 1 0" minWidth={0}>
-            <GroupBySelector traceMetric={traceMetric} />
-          </Flex>
-          <Flex minWidth={0} width="100%">
-            <Filter traceMetric={traceMetric} />
-          </Flex>
-        </Stack>
+        ) : null}
       </Flex>
     );
   }
@@ -68,28 +103,45 @@ export function MetricToolbar({traceMetric, queryIndex}: MetricToolbarProps) {
       width="100%"
       align="center"
       gap="md"
-      columns={`34px 2fr 3fr 6fr ${canRemoveMetric ? '40px' : '0'}`}
+      columns={
+        isVisualizeFunction(visualize)
+          ? `34px 2fr 3fr 6fr ${canRemoveMetric ? '40px' : '0'}`
+          : `34px 1fr ${canRemoveMetric ? '40px' : '0'}`
+      }
       data-test-id="metric-toolbar"
     >
       <VisualizeLabel
         index={queryIndex}
         visualize={visualize}
         onClick={toggleVisibility}
       />
-      <Flex minWidth={0}>
-        <MetricSelector traceMetric={traceMetric} onChange={setTraceMetric} />
-      </Flex>
-      <Flex gap="md" minWidth={0}>
-        <Flex flex="2 1 0" minWidth={0}>
-          <AggregateDropdown traceMetric={traceMetric} />
-        </Flex>
-        <Flex flex="3 1 0" minWidth={0}>
-          <GroupBySelector traceMetric={traceMetric} />
-        </Flex>
-      </Flex>
-      <Flex minWidth={0}>
-        <Filter traceMetric={traceMetric} />
-      </Flex>
+      {isVisualizeFunction(visualize) ? (
+        <Fragment>
+          <Flex minWidth={0}>
+            <MetricSelector traceMetric={traceMetric} onChange={setTraceMetric} />
+          </Flex>
+          <Flex gap="md" minWidth={0}>
+            <Flex flex="2 1 0" minWidth={0}>
+              <AggregateDropdown traceMetric={traceMetric} />
+            </Flex>
+            <Flex flex="3 1 0" minWidth={0}>
+              <GroupBySelector traceMetric={traceMetric} />
+            </Flex>
+          </Flex>
+          <Flex minWidth={0}>
+            <Filter traceMetric={traceMetric} />
+          </Flex>
+        </Fragment>
+      ) : isVisualizeEquation(visualize) ? (
+        <ArithmeticBuilder
+          aggregations={[]}
+          expression={visualize.expression.text}
+          functionArguments={[]}
+          getFieldDefinition={() => null}
+          references={references}
+          setExpression={handleExpressionChange}
+        />
+      ) : null}
       {canRemoveMetric && <DeleteMetricButton />}
     </Grid>
   );