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/metricPanel/index.tsx

@@ -34,10 +34,11 @@ const TWO_MINUTE_DELAY = 120;
 
 interface MetricPanelProps {
   queryIndex: number;
+  references: Set<string>;
   traceMetric: TraceMetric;
 }
 
-export function MetricPanel({traceMetric, queryIndex}: MetricPanelProps) {
+export function MetricPanel({traceMetric, queryIndex, references}: MetricPanelProps) {
   const organization = useOrganization();
   const {
     orientation,
@@ -94,7 +95,11 @@ export function MetricPanel({traceMetric, queryIndex}: MetricPanelProps) {
         <PanelBody>
           <Stack gap="sm">
             <Container paddingBottom={visualize.visible ? undefined : 'sm'}>
-              <MetricToolbar traceMetric={traceMetric} queryIndex={queryIndex} />
+              <MetricToolbar
+                traceMetric={traceMetric}
+                queryIndex={queryIndex}
+                references={references}
+              />
             </Container>
             {visualize.visible ? (
               <SideBySideOrientation

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 {Flex, Grid} 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,13 +42,28 @@ 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 (
       <Grid
         width="100%"
         align="center"
         gap="md"
-        columns={`auto 2fr 3fr 6fr ${canRemoveMetric ? '24px' : '0'}`}
+        columns={
+          isVisualizeFunction(visualize)
+            ? `auto 2fr 3fr 6fr ${canRemoveMetric ? '24px' : '0'}`
+            : `auto 1fr ${canRemoveMetric ? '24px' : '0'}`
+        }
         data-test-id="metric-toolbar"
         paddingLeft="lg"
         paddingRight="lg"
@@ -51,20 +74,33 @@ export function MetricToolbar({traceMetric, queryIndex}: MetricToolbarProps) {
           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>
     );
@@ -75,28 +111,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>
   );