feat(explore): Map input to known attribute in equations (#95589)

When users type the name of an attribute without the explicit tag
syntax, we want to map it to the explicit tag syntax so it's a valid
column whenever possible.

Author: Zylphrex

static/app/components/arithmeticBuilder/context.tsx

@@ -14,6 +14,7 @@ interface ArithmeticBuilderContextData {
   focusOverride: FocusOverride | null;
   functionArguments: FunctionArgument[];
   getFieldDefinition: (key: string) => FieldDefinition | null;
+  getSuggestedKey?: (key: string) => string | null;
 }
 
 export const ArithmeticBuilderContext = createContext<ArithmeticBuilderContextData>({

static/app/components/arithmeticBuilder/index.tsx

@@ -18,16 +18,26 @@ interface ArithmeticBuilderProps {
   functionArguments: FunctionArgument[];
   getFieldDefinition: (key: string) => FieldDefinition | null;
   className?: string;
+  'data-test-id'?: string;
   disabled?: boolean;
+  /**
+   * This is used when a user types in a search key and submits the token.
+   * The submission happens when the user types a colon or presses enter.
+   * When this happens, this function is used to try to map the user input
+   * to a known column.
+   */
+  getSuggestedKey?: (key: string) => string | null;
   setExpression?: (expression: Expression) => void;
 }
 
 export function ArithmeticBuilder({
+  'data-test-id': dataTestId,
   expression,
   setExpression,
   aggregations,
   functionArguments,
   getFieldDefinition,
+  getSuggestedKey,
   className,
   disabled,
 }: ArithmeticBuilderProps) {
@@ -45,16 +55,24 @@ export function ArithmeticBuilder({
       }),
       functionArguments,
       getFieldDefinition,
+      getSuggestedKey,
     };
-  }, [state, dispatch, aggregations, functionArguments, getFieldDefinition]);
+  }, [
+    state,
+    dispatch,
+    aggregations,
+    functionArguments,
+    getFieldDefinition,
+    getSuggestedKey,
+  ]);
 
   return (
     <PanelProvider>
       <ArithmeticBuilderContext value={contextValue}>
         <Wrapper
           className={className}
           aria-disabled={disabled}
-          data-test-id="arithmetic-builder"
+          data-test-id={dataTestId ?? 'arithmetic-builder'}
           state={state.expression.isValid ? 'valid' : 'invalid'}
         >
           <TokenGrid tokens={state.expression.tokens} />

static/app/components/arithmeticBuilder/token/function.tsx

@@ -120,7 +120,8 @@ function InternalInput({
     updateSelectionIndex();
   }, [updateSelectionIndex]);
 
-  const {dispatch, functionArguments, getFieldDefinition} = useArithmeticBuilder();
+  const {dispatch, functionArguments, getFieldDefinition, getSuggestedKey} =
+    useArithmeticBuilder();
 
   const parameterDefinition = useMemo(
     () => getFieldDefinition(token.function)?.parameters?.[argumentIndex],
@@ -182,7 +183,16 @@ function InternalInput({
   );
 
   const onInputCommit = useCallback(() => {
-    const value = inputValue.trim() || attribute.attribute;
+    let value = inputValue.trim() || attribute.attribute;
+
+    if (
+      defined(getSuggestedKey) &&
+      parameterDefinition &&
+      parameterDefinition.kind === 'column'
+    ) {
+      value = getSuggestedKey(value) ?? value;
+    }
+
     dispatch({
       text: `${token.function}(${value})`,
       type: 'REPLACE_TOKEN',
@@ -192,7 +202,16 @@ function InternalInput({
       },
     });
     resetInputValue();
-  }, [dispatch, state, token, attribute, inputValue, resetInputValue]);
+  }, [
+    dispatch,
+    state,
+    token,
+    attribute,
+    inputValue,
+    resetInputValue,
+    getSuggestedKey,
+    parameterDefinition,
+  ]);
 
   const onInputEscape = useCallback(() => {
     resetInputValue();

static/app/components/arithmeticBuilder/token/index.spec.tsx

@@ -36,6 +36,18 @@ const getSpanFieldDefinition = (key: string) => {
   return getFieldDefinition(key, 'span', argument?.kind);
 };
 
+const getSuggestedKey = (key: string) => {
+  switch (key) {
+    case 'duration':
+    case 'self_time':
+    case 'op':
+    case 'description':
+      return `span.${key}`;
+    default:
+      return null;
+  }
+};
+
 interface TokensProp {
   expression: string;
   dispatch?: Dispatch<ArithmeticBuilderAction>;
@@ -62,6 +74,7 @@ function Tokens(props: TokensProp) {
         aggregations,
         functionArguments,
         getFieldDefinition: getSpanFieldDefinition,
+        getSuggestedKey,
       }}
     >
       <TokenGrid tokens={state.expression.tokens} />
@@ -518,6 +531,38 @@ describe('token', function () {
       ).toBeInTheDocument();
     });
 
+    it('maps key to suggested key on enter', async function () {
+      render(<Tokens expression="avg(span.duration)" />);
+
+      expect(
+        await screen.findByRole('row', {
+          name: 'avg(span.duration)',
+        })
+      ).toBeInTheDocument();
+
+      const input = screen.getByRole('combobox', {
+        name: 'Select an attribute',
+      });
+      expect(input).toBeInTheDocument();
+
+      await userEvent.click(input);
+      expect(input).toHaveFocus();
+      expect(input).toHaveAttribute('placeholder', 'span.duration');
+      expect(input).toHaveValue('');
+
+      await userEvent.type(input, 'self_time{Enter}');
+
+      const lastInput = getLastInput();
+      await waitFor(() => expect(lastInput).toHaveFocus());
+      await userEvent.type(lastInput, '{Escape}');
+
+      expect(
+        await screen.findByRole('row', {
+          name: 'avg(span.self_time)',
+        })
+      ).toBeInTheDocument();
+    });
+
     it('doesnt change argument on enter if input is empty', async function () {
       render(<Tokens expression="avg(span.duration)" />);
 

static/app/views/explore/components/traceItemSearchQueryBuilder.tsx

@@ -1,4 +1,4 @@
-import {useCallback, useMemo} from 'react';
+import {useMemo} from 'react';
 
 import {getHasTag} from 'sentry/components/events/searchBar';
 import type {EAPSpanSearchQueryBuilderProps} from 'sentry/components/performance/spanSearchQueryBuilder';
@@ -7,6 +7,7 @@ import {t} from 'sentry/locale';
 import {SavedSearchType, type TagCollection} from 'sentry/types/group';
 import type {AggregationKey} from 'sentry/utils/fields';
 import {FieldKind, getFieldDefinition} from 'sentry/utils/fields';
+import {useExploreSuggestedAttribute} from 'sentry/views/explore/hooks/useExploreSuggestedAttribute';
 import {useGetTraceItemAttributeValues} from 'sentry/views/explore/hooks/useGetTraceItemAttributeValues';
 import {LOGS_FILTER_KEY_SECTIONS} from 'sentry/views/explore/logs/constants';
 import {TraceItemDataset} from 'sentry/views/explore/types';
@@ -74,24 +75,10 @@ export function useSearchQueryBuilderProps({
     projectIds: projects,
   });
 
-  const getSuggestedFilterKey = useCallback(
-    (key: string) => {
-      // prioritize exact matches first
-      if (filterTags.hasOwnProperty(key)) {
-        return key;
-      }
-
-      // try to see if there's numeric attribute by the same name
-      const explicitNumberTag = `tags[${key},number]`;
-      if (filterTags.hasOwnProperty(explicitNumberTag)) {
-        return explicitNumberTag;
-      }
-
-      // give up, and fall back to the default behaviour
-      return null;
-    },
-    [filterTags]
-  );
+  const getSuggestedAttribute = useExploreSuggestedAttribute({
+    numberAttributes,
+    stringAttributes,
+  });
 
   return {
     placeholder: placeholderText,
@@ -104,7 +91,7 @@ export function useSearchQueryBuilderProps({
     getFilterTokenWarning,
     searchSource,
     filterKeySections,
-    getSuggestedFilterKey,
+    getSuggestedFilterKey: getSuggestedAttribute,
     getTagValues: getTraceItemAttributeValues,
     disallowUnsupportedFilters: true,
     recentSearches: itemTypeToRecentSearches(itemType),

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

@@ -0,0 +1,38 @@
+import {useCallback} from 'react';
+
+import type {TagCollection} from 'sentry/types/group';
+
+interface UseExploreSuggestedAttributeOptions {
+  numberAttributes: TagCollection;
+  stringAttributes: TagCollection;
+}
+
+export function useExploreSuggestedAttribute({
+  numberAttributes,
+  stringAttributes,
+}: UseExploreSuggestedAttributeOptions) {
+  return useCallback(
+    (key: string): string | null => {
+      if (stringAttributes.hasOwnProperty(key)) {
+        return key;
+      }
+
+      if (numberAttributes.hasOwnProperty(key)) {
+        return key;
+      }
+
+      const explicitStringAttribute = `tags[${key},string]`;
+      if (stringAttributes.hasOwnProperty(explicitStringAttribute)) {
+        return explicitStringAttribute;
+      }
+
+      const explicitNumberAttribute = `tags[${key},number]`;
+      if (numberAttributes.hasOwnProperty(explicitNumberAttribute)) {
+        return explicitNumberAttribute;
+      }
+
+      return null;
+    },
+    [numberAttributes, stringAttributes]
+  );
+}

static/app/views/explore/tables/aggregateColumnEditorModal.spec.tsx

@@ -1,3 +1,4 @@
+import {initializeOrg} from 'sentry-test/initializeOrg';
 import {
   act,
   renderGlobalModal,
@@ -57,6 +58,11 @@ const numberTags: TagCollection = {
     name: 'span.self_time',
     kind: FieldKind.MEASUREMENT,
   },
+  'tags[foo,number]': {
+    key: 'tags[foo,number]',
+    name: 'foo',
+    kind: FieldKind.MEASUREMENT,
+  },
 };
 
 describe('AggregateColumnEditorModal', function () {
@@ -263,6 +269,47 @@ describe('AggregateColumnEditorModal', function () {
       {yAxes: ['count(span.duration)']},
     ]);
   });
+
+  it('allows adding an equation', async function () {
+    const {organization} = initializeOrg({
+      organization: {
+        features: ['visibility-explore-equations'],
+      },
+    });
+
+    const onColumnsChange = jest.fn();
+
+    renderGlobalModal({organization});
+
+    act(() => {
+      openModal(
+        modalProps => (
+          <AggregateColumnEditorModal
+            {...modalProps}
+            columns={[{groupBy: 'geo.country'}, new Visualize(DEFAULT_VISUALIZATION)]}
+            onColumnsChange={onColumnsChange}
+            stringTags={stringTags}
+            numberTags={numberTags}
+          />
+        ),
+        {onClose: jest.fn()}
+      );
+    });
+
+    await userEvent.click(screen.getByRole('button', {name: 'Add a Column'}));
+    await userEvent.click(screen.getByRole('menuitemradio', {name: 'Equation'}));
+
+    await userEvent.click(screen.getByRole('combobox', {name: 'Add a term'}));
+
+    await userEvent.keyboard('avg(foo{Enter}*5{Escape}');
+
+    await userEvent.click(screen.getByRole('button', {name: 'Apply'}));
+    expect(onColumnsChange).toHaveBeenCalledWith([
+      {groupBy: 'geo.country'},
+      {yAxes: ['count(span.duration)']},
+      {yAxes: ['equation|avg(tags[foo,number]) * 5']},
+    ]);
+  });
 });
 
 function expectRows(rows: HTMLElement[]) {

static/app/views/explore/tables/aggregateColumnEditorModal.tsx

@@ -48,6 +48,7 @@ import {
   Visualize,
 } from 'sentry/views/explore/contexts/pageParamsContext/visualizes';
 import type {Column} from 'sentry/views/explore/hooks/useDragNDropColumns';
+import {useExploreSuggestedAttribute} from 'sentry/views/explore/hooks/useExploreSuggestedAttribute';
 import {useGroupByFields} from 'sentry/views/explore/hooks/useGroupByFields';
 import {useVisualizeFields} from 'sentry/views/explore/hooks/useVisualizeFields';
 
@@ -425,13 +426,20 @@ function EquationSelector({
     [onChange, visualize]
   );
 
+  const getSuggestedAttribute = useExploreSuggestedAttribute({
+    numberAttributes: numberTags,
+    stringAttributes: stringTags,
+  });
+
   return (
     <ArithmeticBuilder
+      data-test-id="editor-visualize-equation"
       aggregations={ALLOWED_EXPLORE_VISUALIZE_AGGREGATES}
       functionArguments={functionArguments}
       getFieldDefinition={getSpanFieldDefinition}
       expression={expression}
       setExpression={handleExpressionChange}
+      getSuggestedKey={getSuggestedAttribute}
     />
   );
 }

static/app/views/explore/toolbar/toolbarVisualize.tsx

@@ -33,6 +33,7 @@ import {
   Visualize,
 } from 'sentry/views/explore/contexts/pageParamsContext/visualizes';
 import {useTraceItemTags} from 'sentry/views/explore/contexts/spanTagsContext';
+import {useExploreSuggestedAttribute} from 'sentry/views/explore/hooks/useExploreSuggestedAttribute';
 import {useVisualizeFields} from 'sentry/views/explore/hooks/useVisualizeFields';
 
 import {
@@ -206,6 +207,11 @@ function VisualizeEquation({
     [group, setVisualizes, visualizes]
   );
 
+  const getSuggestedAttribute = useExploreSuggestedAttribute({
+    numberAttributes: numberTags,
+    stringAttributes: stringTags,
+  });
+
   return (
     <ToolbarRow>
       <ArithmeticBuilder
@@ -214,6 +220,7 @@ function VisualizeEquation({
         getFieldDefinition={getSpanFieldDefinition}
         expression={expression}
         setExpression={handleExpressionChange}
+        getSuggestedKey={getSuggestedAttribute}
       />
       <Button
         borderless