From dca0f4871bf6a5c66b63bbbcf24d9589ebbac764 Mon Sep 17 00:00:00 2001
From: Nar Saynorath <nar.saynorath@sentry.io>
Date: Wed, 15 Jan 2025 12:04:51 -0500
Subject: [PATCH] feat(widget-builder): Allow other functions in sort for spans
 (#83505)

A user can add a new series to the explore UI that sorts on a different
field than it's showing. This allows us to transfer the query over to
the dashboard by applying the sorting to the fields and then allowing
the sort by selectors to show that option (i.e. unlike tables that need
you to choose a column that's represented)
---
 .../views/dashboards/datasetConfig/spans.tsx   | 18 +++++++++++++++---
 .../buildSteps/sortByStep/sortBySelectors.tsx  |  6 ++++--
 .../widgetBuilder/widgetBuilderSortBy.spec.tsx | 16 ++++++++--------
 .../views/explore/hooks/useAddToDashboard.tsx  |  4 +++-
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/static/app/views/dashboards/datasetConfig/spans.tsx b/static/app/views/dashboards/datasetConfig/spans.tsx
index 636c2b62865c77..9b6c5b4250a0ef 100644
--- a/static/app/views/dashboards/datasetConfig/spans.tsx
+++ b/static/app/views/dashboards/datasetConfig/spans.tsx
@@ -30,6 +30,7 @@ import {
 } from 'sentry/views/dashboards/datasetConfig/base';
 import {
   getTableSortOptions,
+  getTimeseriesSortOptions,
   transformEventsResponseToSeries,
   transformEventsResponseToTable,
 } from 'sentry/views/dashboards/datasetConfig/errorsAndTransactions';
@@ -73,9 +74,6 @@ const EAP_AGGREGATIONS = ALLOWED_EXPLORE_VISUALIZE_AGGREGATES.reduce((acc, aggre
   return acc;
 }, {});
 
-// getTimeseriesSortOptions is undefined because we want to restrict the
-// sort options to the same behaviour as tables. i.e. we are only able
-// to sort by fields that have already been selected
 export const SpansConfig: DatasetConfig<
   EventsStats | MultiSeriesEventsStats,
   TableData | EventsTableData
@@ -86,8 +84,11 @@ export const SpansConfig: DatasetConfig<
   SearchBar: SpansSearchBar,
   filterYAxisAggregateParams: () => filterAggregateParams,
   filterYAxisOptions,
+  filterSeriesSortOptions,
   getTableFieldOptions: getPrimaryFieldOptions,
   getTableSortOptions,
+  getTimeseriesSortOptions: (organization, widgetQuery, tags) =>
+    getTimeseriesSortOptions(organization, widgetQuery, tags, getPrimaryFieldOptions),
   getGroupByFieldOptions,
   handleOrderByReset,
   supportedDisplayTypes: [
@@ -281,3 +282,14 @@ function getSeriesRequest(
 
   return doEventsRequest<true>(api, requestData);
 }
+
+// Filters the primary options in the sort by selector
+export function filterSeriesSortOptions(columns: Set<string>) {
+  return (option: FieldValueOption) => {
+    if (option.value.kind === FieldValueKind.FUNCTION) {
+      return true;
+    }
+
+    return columns.has(option.value.meta.name);
+  };
+}
diff --git a/static/app/views/dashboards/widgetBuilder/buildSteps/sortByStep/sortBySelectors.tsx b/static/app/views/dashboards/widgetBuilder/buildSteps/sortByStep/sortBySelectors.tsx
index 1de4cd9fe4b163..dc8ec643dcbc3d 100644
--- a/static/app/views/dashboards/widgetBuilder/buildSteps/sortByStep/sortBySelectors.tsx
+++ b/static/app/views/dashboards/widgetBuilder/buildSteps/sortByStep/sortBySelectors.tsx
@@ -55,6 +55,7 @@ export function SortBySelectors({
   disableSortDirection,
   widgetQuery,
   displayType,
+  tags,
 }: Props) {
   const datasetConfig = getDatasetConfig(widgetType);
   const organization = useOrganization();
@@ -103,7 +104,7 @@ export function SortBySelectors({
         title={disableSortReason}
         disabled={!disableSort || (disableSortDirection && disableSort)}
       >
-        {displayType === DisplayType.TABLE || widgetType === WidgetType.SPANS ? (
+        {displayType === DisplayType.TABLE ? (
           <SelectControl
             name="sortBy"
             aria-label={t('Sort by')}
@@ -135,7 +136,8 @@ export function SortBySelectors({
             }
             fieldOptions={datasetConfig.getTimeseriesSortOptions!(
               organization,
-              widgetQuery
+              widgetQuery,
+              tags
             )}
             filterPrimaryOptions={
               datasetConfig.filterSeriesSortOptions
diff --git a/static/app/views/dashboards/widgetBuilder/widgetBuilderSortBy.spec.tsx b/static/app/views/dashboards/widgetBuilder/widgetBuilderSortBy.spec.tsx
index 694d53e55ee367..fbaa58f8976105 100644
--- a/static/app/views/dashboards/widgetBuilder/widgetBuilderSortBy.spec.tsx
+++ b/static/app/views/dashboards/widgetBuilder/widgetBuilderSortBy.spec.tsx
@@ -937,7 +937,7 @@ describe('WidgetBuilder', function () {
   });
 
   describe('spans dataset timeseries', function () {
-    it('returns only the selected aggregates and group by as options', async function () {
+    it('returns group by and aggregate as primary options', async function () {
       const widget: Widget = {
         id: '1',
         title: 'Test Widget',
@@ -967,16 +967,16 @@ describe('WidgetBuilder', function () {
       });
 
       await screen.findByText('Sort by a y-axis');
-      await selectEvent.openMenu(await screen.findByText('count(span.duration)'));
+      await selectEvent.openMenu(screen.getAllByText('count(\u2026)')[1]!);
 
-      // 3 options in the dropdown
-      expect(screen.queryAllByTestId('menu-list-item-label')).toHaveLength(3);
+      // 12 options in the dropdown
+      expect(screen.queryAllByTestId('menu-list-item-label')).toHaveLength(12);
 
-      // Appears once in the dropdown and once in the sort by field
-      expect(await screen.findAllByText('count(span.duration)')).toHaveLength(2);
+      // Appears once in the y-axis section, dropdown, and in the sort by field
+      expect(await screen.findAllByText('count(\u2026)')).toHaveLength(3);
 
-      // Appears once in the dropdown
-      expect(await screen.findAllByText('avg(span.duration)')).toHaveLength(1);
+      // Appears once in the y-axis section and in the dropdown
+      expect(await screen.findAllByText('avg(\u2026)')).toHaveLength(2);
 
       // Appears once in the dropdown and once in the group by field
       expect(await screen.findAllByText('transaction')).toHaveLength(2);
diff --git a/static/app/views/explore/hooks/useAddToDashboard.tsx b/static/app/views/explore/hooks/useAddToDashboard.tsx
index 8de959c957b43b..d53f4fd1066e83 100644
--- a/static/app/views/explore/hooks/useAddToDashboard.tsx
+++ b/static/app/views/explore/hooks/useAddToDashboard.tsx
@@ -54,7 +54,9 @@ export function useAddToDashboard() {
           fields = sampleFields.filter(Boolean);
         }
       } else {
-        fields = [...groupBys, ...yAxes].filter(Boolean);
+        fields = [
+          ...new Set([...groupBys, ...yAxes, ...sortBys.map(sort => sort.field)]),
+        ].filter(Boolean);
       }
 
       const search = new MutableSearch(query);