From 43e2be033310a1bdb640b53c287f6cdf942b48f4 Mon Sep 17 00:00:00 2001
From: Shaheer Kochai <ashaheerki@gmail.com>
Date: Tue, 22 Apr 2025 19:54:03 +0430
Subject: [PATCH] feat: use search v2 component for traces (#7537)

* Revert "fix: display same key with multiple data types in filter suggestions by enhancing the deduping logic (#7255)"

This reverts commit 1e85981a17a8e715e948308d3e85072d976907d3.

* fix: use query search v2 for traces data source to handle multiple data types for the same key

* fix(QueryBuilderSearchV2): add user typed option if it doesn't exist in the payload

* fix(QueryBuilderSearchV2): increase the height of search dropdown for non-logs data sources

* fix: display span scope selector for trace data source

* chore: remove the span scope selector from qb search v1 and move the component to search v2

* fix: write test to ensure that we display span scope selector for traces data source

* fix: limit converting  ->   only to log data source

* fix: don't display empty suggestion if only spaces are typed

* chore: tests for span scope selector

* chore: qb search flow (key, operator, value) test cases

* refactor: fix the Maximum update depth reached issue while running tests

* chore: overall improvements to span scope selector tests
---
 .../QueryBuilder/components/Query/Query.tsx   |   2 +-
 .../filters/QueryBuilderSearch/index.tsx      |   7 -
 .../QueryBuilderSearchDropdown.tsx            |   7 +-
 .../QueryBuilderSearchV2.styles.scss          |   5 +
 .../QueryBuilderSearchV2.tsx                  |  44 +++-
 .../SpanScopeSelector.tsx                     |   1 +
 .../__test__/QueryBuilderSearchV2.test.tsx    | 196 ++++++++++++++++++
 .../__test__/SpanScopeSelector.test.tsx       | 165 +++++++++++++++
 frontend/src/hooks/queryBuilder/useOptions.ts |   6 +-
 9 files changed, 413 insertions(+), 20 deletions(-)
 rename frontend/src/container/QueryBuilder/filters/{QueryBuilderSearch => QueryBuilderSearchV2}/SpanScopeSelector.tsx (98%)
 create mode 100644 frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/QueryBuilderSearchV2.test.tsx
 create mode 100644 frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/SpanScopeSelector.test.tsx

diff --git a/frontend/src/container/QueryBuilder/components/Query/Query.tsx b/frontend/src/container/QueryBuilder/components/Query/Query.tsx
index 23c5f09528..f1ce6aeaa8 100644
--- a/frontend/src/container/QueryBuilder/components/Query/Query.tsx
+++ b/frontend/src/container/QueryBuilder/components/Query/Query.tsx
@@ -453,7 +453,7 @@ export const Query = memo(function Query({
 										</Col>
 									)}
 									<Col flex="1" className="qb-search-container">
-										{query.dataSource === DataSource.LOGS ? (
+										{[DataSource.LOGS, DataSource.TRACES].includes(query.dataSource) ? (
 											<QueryBuilderSearchV2
 												query={query}
 												onChange={handleChangeTagFilters}
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/index.tsx b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/index.tsx
index a781743db5..dc22801cea 100644
--- a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/index.tsx
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/index.tsx
@@ -56,7 +56,6 @@ import { PLACEHOLDER } from './constant';
 import ExampleQueriesRendererForLogs from './ExampleQueriesRendererForLogs';
 import OptionRenderer from './OptionRenderer';
 import OptionRendererForLogs from './OptionRendererForLogs';
-import SpanScopeSelector from './SpanScopeSelector';
 import { StyledCheckOutlined, TypographyText } from './style';
 import {
 	convertExampleQueriesToOptions,
@@ -84,11 +83,6 @@ function QueryBuilderSearch({
 		pathname,
 	]);
 
-	const isTracesExplorerPage = useMemo(
-		() => pathname === ROUTES.TRACES_EXPLORER,
-		[pathname],
-	);
-
 	const [isEditingTag, setIsEditingTag] = useState(false);
 
 	const {
@@ -489,7 +483,6 @@ function QueryBuilderSearch({
 							</Select.Option>
 					  ))}
 			</Select>
-			{isTracesExplorerPage && <SpanScopeSelector queryName={query.queryName} />}
 		</div>
 	);
 }
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchDropdown.tsx b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchDropdown.tsx
index 577b72ff83..c046b1ca4f 100644
--- a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchDropdown.tsx
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchDropdown.tsx
@@ -2,6 +2,7 @@
 import './QueryBuilderSearchV2.styles.scss';
 
 import { Typography } from 'antd';
+import cx from 'classnames';
 import {
 	ArrowDown,
 	ArrowUp,
@@ -25,6 +26,7 @@ interface ICustomDropdownProps {
 	exampleQueries: TagFilter[];
 	onChange: (value: TagFilter) => void;
 	currentFilterItem?: ITag;
+	isLogsDataSource: boolean;
 }
 
 export default function QueryBuilderSearchDropdown(
@@ -38,11 +40,14 @@ export default function QueryBuilderSearchDropdown(
 		exampleQueries,
 		options,
 		onChange,
+		isLogsDataSource,
 	} = props;
 	const userOs = getUserOperatingSystem();
 	return (
 		<>
-			<div className="content">
+			<div
+				className={cx('content', { 'non-logs-data-source': !isLogsDataSource })}
+			>
 				{!currentFilterItem?.key ? (
 					<div className="suggested-filters">Suggested Filters</div>
 				) : !currentFilterItem?.op ? (
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.styles.scss b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.styles.scss
index c7a96456e9..49f04baf0e 100644
--- a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.styles.scss
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.styles.scss
@@ -11,6 +11,11 @@
 			.rc-virtual-list-holder {
 				height: 115px;
 			}
+			&.non-logs-data-source {
+				.rc-virtual-list-holder {
+					height: 256px;
+				}
+			}
 		}
 	}
 
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.tsx b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.tsx
index fa4c9aea0c..1c0cc87704 100644
--- a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.tsx
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/QueryBuilderSearchV2.tsx
@@ -65,6 +65,7 @@ import {
 } from '../QueryBuilderSearch/utils';
 import { filterByOperatorConfig } from '../utils';
 import QueryBuilderSearchDropdown from './QueryBuilderSearchDropdown';
+import SpanScopeSelector from './SpanScopeSelector';
 import Suggestions from './Suggestions';
 
 export interface ITag {
@@ -294,7 +295,8 @@ function QueryBuilderSearchV2(
 				if (
 					isObject(parsedValue) &&
 					parsedValue?.key &&
-					parsedValue?.key?.split(' ').length > 1
+					parsedValue?.key?.split(' ').length > 1 &&
+					isLogsDataSource
 				) {
 					setTags((prev) => [
 						...prev,
@@ -409,7 +411,13 @@ function QueryBuilderSearchV2(
 				}
 			}
 		},
-		[currentFilterItem?.key, currentFilterItem?.op, currentState, searchValue],
+		[
+			currentFilterItem?.key,
+			currentFilterItem?.op,
+			currentState,
+			isLogsDataSource,
+			searchValue,
+		],
 	);
 
 	const handleSearch = useCallback((value: string) => {
@@ -693,12 +701,29 @@ function QueryBuilderSearchV2(
 					})),
 				);
 			} else {
-				setDropdownOptions(
-					data?.payload?.attributeKeys?.map((key) => ({
+				setDropdownOptions([
+					// Add user typed option if it doesn't exist in the payload
+					...(tagKey.trim().length > 0 &&
+					!data?.payload?.attributeKeys?.some((val) => val.key === tagKey)
+						? [
+								{
+									label: tagKey,
+									value: {
+										key: tagKey,
+										dataType: DataTypes.EMPTY,
+										type: '',
+										isColumn: false,
+										isJSON: false,
+									},
+								},
+						  ]
+						: []),
+					// Map existing attribute keys from payload
+					...(data?.payload?.attributeKeys?.map((key) => ({
 						label: key.key,
 						value: key,
-					})) || [],
-				);
+					})) || []),
+				]);
 			}
 		}
 		if (currentState === DropdownState.OPERATOR) {
@@ -911,6 +936,11 @@ function QueryBuilderSearchV2(
 		);
 	};
 
+	const isTracesDataSource = useMemo(
+		() => query.dataSource === DataSource.TRACES,
+		[query.dataSource],
+	);
+
 	return (
 		<div className="query-builder-search-v2">
 			<Select
@@ -968,6 +998,7 @@ function QueryBuilderSearchV2(
 						exampleQueries={suggestionsData?.payload?.example_queries || []}
 						tags={tags}
 						currentFilterItem={currentFilterItem}
+						isLogsDataSource={isLogsDataSource}
 					/>
 				)}
 			>
@@ -994,6 +1025,7 @@ function QueryBuilderSearchV2(
 					);
 				})}
 			</Select>
+			{isTracesDataSource && <SpanScopeSelector queryName={query.queryName} />}
 		</div>
 	);
 }
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/SpanScopeSelector.tsx b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/SpanScopeSelector.tsx
similarity index 98%
rename from frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/SpanScopeSelector.tsx
rename to frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/SpanScopeSelector.tsx
index 7ff6394123..bdf2d16343 100644
--- a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearch/SpanScopeSelector.tsx
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/SpanScopeSelector.tsx
@@ -120,6 +120,7 @@ function SpanScopeSelector({ queryName }: SpanScopeSelectorProps): JSX.Element {
 		<Select
 			value={selectedScope}
 			className="span-scope-selector"
+			data-testid="span-scope-selector"
 			onChange={handleScopeChange}
 			options={SELECT_OPTIONS}
 		/>
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/QueryBuilderSearchV2.test.tsx b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/QueryBuilderSearchV2.test.tsx
new file mode 100644
index 0000000000..1227c7f191
--- /dev/null
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/QueryBuilderSearchV2.test.tsx
@@ -0,0 +1,196 @@
+/* eslint-disable react/jsx-props-no-spreading */
+import {
+	act,
+	fireEvent,
+	render,
+	RenderResult,
+	screen,
+} from '@testing-library/react';
+import {
+	initialQueriesMap,
+	initialQueryBuilderFormValues,
+} from 'constants/queryBuilder';
+import { QueryBuilderContext } from 'providers/QueryBuilder';
+import { QueryClient, QueryClientProvider } from 'react-query';
+import { DataTypes } from 'types/api/queryBuilder/queryAutocompleteResponse';
+import { DataSource } from 'types/common/queryBuilder';
+
+import QueryBuilderSearchV2 from '../QueryBuilderSearchV2';
+
+const queryClient = new QueryClient({
+	defaultOptions: {
+		queries: {
+			refetchOnWindowFocus: false,
+		},
+	},
+});
+
+describe('Span scope selector', () => {
+	it('should render span scope selector when data source is TRACES', () => {
+		const { getByTestId } = render(
+			<QueryClientProvider client={queryClient}>
+				<QueryBuilderSearchV2
+					query={{
+						...initialQueryBuilderFormValues,
+						dataSource: DataSource.TRACES,
+					}}
+					onChange={jest.fn()}
+				/>
+			</QueryClientProvider>,
+		);
+
+		expect(getByTestId('span-scope-selector')).toBeInTheDocument();
+	});
+
+	it('should not render span scope selector for non-TRACES data sources', () => {
+		const { queryByTestId } = render(
+			<QueryClientProvider client={queryClient}>
+				<QueryBuilderSearchV2
+					query={{
+						...initialQueryBuilderFormValues,
+						dataSource: DataSource.METRICS,
+					}}
+					onChange={jest.fn()}
+				/>
+			</QueryClientProvider>,
+		);
+
+		expect(queryByTestId('span-scope-selector')).not.toBeInTheDocument();
+	});
+});
+
+const mockOnChange = jest.fn();
+const mockHandleRunQuery = jest.fn();
+const defaultProps = {
+	query: {
+		...initialQueriesMap.traces.builder.queryData[0],
+		dataSource: DataSource.TRACES,
+		queryName: 'traces_query',
+	},
+	onChange: mockOnChange,
+};
+
+const renderWithContext = (props = {}): RenderResult => {
+	const mergedProps = { ...defaultProps, ...props };
+
+	return render(
+		<QueryClientProvider client={queryClient}>
+			<QueryBuilderContext.Provider
+				value={
+					{
+						currentQuery: initialQueriesMap.traces,
+						handleRunQuery: mockHandleRunQuery,
+					} as any
+				}
+			>
+				<QueryBuilderSearchV2 {...mergedProps} />
+			</QueryBuilderContext.Provider>
+		</QueryClientProvider>,
+	);
+};
+
+const mockAggregateKeysData = {
+	payload: {
+		attributeKeys: [
+			{
+				// eslint-disable-next-line sonarjs/no-duplicate-string
+				key: 'http.status',
+				dataType: DataTypes.String,
+				type: 'tag',
+				isColumn: false,
+				isJSON: false,
+				id: 'http.status--string--tag--false',
+			},
+		],
+	},
+};
+
+jest.mock('hooks/queryBuilder/useGetAggregateKeys', () => ({
+	useGetAggregateKeys: jest.fn(() => ({
+		data: mockAggregateKeysData,
+		isFetching: false,
+	})),
+}));
+
+const mockAggregateValuesData = {
+	payload: {
+		stringAttributeValues: ['200', '404', '500'],
+		numberAttributeValues: [200, 404, 500],
+	},
+};
+
+jest.mock('hooks/queryBuilder/useGetAggregateValues', () => ({
+	useGetAggregateValues: jest.fn(() => ({
+		data: mockAggregateValuesData,
+		isFetching: false,
+	})),
+}));
+
+jest.mock('hooks/useSafeNavigate', () => ({
+	useSafeNavigate: (): any => ({
+		safeNavigate: jest.fn(),
+	}),
+}));
+
+describe('Suggestion Key -> Operator -> Value Flow', () => {
+	beforeEach(() => {
+		jest.clearAllMocks();
+	});
+	it('should complete full flow from key selection to value', async () => {
+		const { container } = renderWithContext();
+
+		// Get the combobox input specifically
+		const combobox = container.querySelector(
+			'.query-builder-search-v2 .ant-select-selection-search-input',
+		) as HTMLInputElement;
+
+		// 1. Focus and type to trigger key suggestions
+		await act(async () => {
+			fireEvent.focus(combobox);
+			fireEvent.change(combobox, { target: { value: 'http.' } });
+		});
+
+		// Wait for dropdown to appear
+		await screen.findByRole('listbox');
+
+		// 2. Select a key from suggestions
+		const statusOption = await screen.findByText('http.status');
+		await act(async () => {
+			fireEvent.click(statusOption);
+		});
+
+		// Should show operator suggestions
+		expect(screen.getByText('=')).toBeInTheDocument();
+		expect(screen.getByText('!=')).toBeInTheDocument();
+
+		// 3. Select an operator
+		const equalsOption = screen.getByText('=');
+		await act(async () => {
+			fireEvent.click(equalsOption);
+		});
+
+		// Should show value suggestions
+		expect(screen.getByText('200')).toBeInTheDocument();
+		expect(screen.getByText('404')).toBeInTheDocument();
+		expect(screen.getByText('500')).toBeInTheDocument();
+
+		// 4. Select a value
+		const valueOption = screen.getByText('200');
+		await act(async () => {
+			fireEvent.click(valueOption);
+		});
+
+		// Verify final filter
+		expect(mockOnChange).toHaveBeenCalledWith(
+			expect.objectContaining({
+				items: expect.arrayContaining([
+					expect.objectContaining({
+						key: expect.objectContaining({ key: 'http.status' }),
+						op: '=',
+						value: '200',
+					}),
+				]),
+			}),
+		);
+	});
+});
diff --git a/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/SpanScopeSelector.test.tsx b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/SpanScopeSelector.test.tsx
new file mode 100644
index 0000000000..180c941814
--- /dev/null
+++ b/frontend/src/container/QueryBuilder/filters/QueryBuilderSearchV2/__test__/SpanScopeSelector.test.tsx
@@ -0,0 +1,165 @@
+import {
+	fireEvent,
+	render,
+	RenderResult,
+	screen,
+} from '@testing-library/react';
+import { initialQueriesMap } from 'constants/queryBuilder';
+import { QueryBuilderContext } from 'providers/QueryBuilder';
+import { Query, TagFilterItem } from 'types/api/queryBuilder/queryBuilderData';
+
+import SpanScopeSelector from '../SpanScopeSelector';
+
+const mockRedirectWithQueryBuilderData = jest.fn();
+
+// Helper to create filter items
+const createSpanScopeFilter = (key: string): TagFilterItem => ({
+	id: 'span-filter',
+	key: {
+		key,
+		type: 'spanSearchScope',
+	},
+	op: '=',
+	value: 'true',
+});
+
+const defaultQuery = {
+	...initialQueriesMap.traces,
+	builder: {
+		...initialQueriesMap.traces.builder,
+		queryData: [
+			{
+				...initialQueriesMap.traces.builder.queryData[0],
+				queryName: 'A',
+			},
+		],
+	},
+};
+
+// Helper to create query with filters
+const createQueryWithFilters = (filters: TagFilterItem[]): Query => ({
+	...defaultQuery,
+	builder: {
+		...defaultQuery.builder,
+		queryData: [
+			{
+				...defaultQuery.builder.queryData[0],
+				filters: {
+					items: filters,
+					op: 'AND',
+				},
+			},
+		],
+	},
+});
+
+const renderWithContext = (
+	queryName = 'A',
+	initialQuery = defaultQuery,
+): RenderResult =>
+	render(
+		<QueryBuilderContext.Provider
+			value={
+				{
+					currentQuery: initialQuery,
+					redirectWithQueryBuilderData: mockRedirectWithQueryBuilderData,
+					// eslint-disable-next-line @typescript-eslint/no-explicit-any
+				} as any
+			}
+		>
+			<SpanScopeSelector queryName={queryName} />
+		</QueryBuilderContext.Provider>,
+	);
+
+describe('SpanScopeSelector', () => {
+	beforeEach(() => {
+		jest.clearAllMocks();
+	});
+
+	it('should render with default ALL_SPANS selected', () => {
+		renderWithContext();
+		expect(screen.getByText('All Spans')).toBeInTheDocument();
+	});
+
+	describe('when selecting different options', () => {
+		const selectOption = (optionText: string): void => {
+			const selector = screen.getByRole('combobox');
+			fireEvent.mouseDown(selector);
+			const option = screen.getByText(optionText);
+			fireEvent.click(option);
+		};
+
+		const assertFilterAdded = (
+			updatedQuery: Query,
+			expectedKey: string,
+		): void => {
+			const filters = updatedQuery.builder.queryData[0].filters.items;
+			expect(filters).toContainEqual(
+				expect.objectContaining({
+					key: expect.objectContaining({
+						key: expectedKey,
+						type: 'spanSearchScope',
+					}),
+					op: '=',
+					value: 'true',
+				}),
+			);
+		};
+
+		it('should remove span scope filters when selecting ALL_SPANS', () => {
+			const queryWithSpanScope = createQueryWithFilters([
+				createSpanScopeFilter('isRoot'),
+			]);
+			renderWithContext('A', queryWithSpanScope);
+
+			selectOption('All Spans');
+
+			expect(mockRedirectWithQueryBuilderData).toHaveBeenCalled();
+			const updatedQuery = mockRedirectWithQueryBuilderData.mock.calls[0][0];
+			const filters = updatedQuery.builder.queryData[0].filters.items;
+			expect(filters).not.toContainEqual(
+				expect.objectContaining({
+					key: expect.objectContaining({ type: 'spanSearchScope' }),
+				}),
+			);
+		});
+
+		it('should add isRoot filter when selecting ROOT_SPANS', async () => {
+			renderWithContext();
+			await selectOption('Root Spans');
+
+			expect(mockRedirectWithQueryBuilderData).toHaveBeenCalled();
+			assertFilterAdded(
+				mockRedirectWithQueryBuilderData.mock.calls[0][0],
+				'isRoot',
+			);
+		});
+
+		it('should add isEntryPoint filter when selecting ENTRYPOINT_SPANS', () => {
+			renderWithContext();
+			selectOption('Entrypoint Spans');
+
+			expect(mockRedirectWithQueryBuilderData).toHaveBeenCalled();
+			assertFilterAdded(
+				mockRedirectWithQueryBuilderData.mock.calls[0][0],
+				'isEntryPoint',
+			);
+		});
+	});
+
+	describe('when initializing with existing filters', () => {
+		it.each([
+			['Root Spans', 'isRoot'],
+			['Entrypoint Spans', 'isEntryPoint'],
+		])(
+			'should initialize with %s selected when %s filter exists',
+			async (expectedText, filterKey) => {
+				const queryWithFilter = createQueryWithFilters([
+					createSpanScopeFilter(filterKey),
+				]);
+				renderWithContext('A', queryWithFilter);
+				expect(await screen.findByText(expectedText)).toBeInTheDocument();
+			},
+		);
+	});
+});
diff --git a/frontend/src/hooks/queryBuilder/useOptions.ts b/frontend/src/hooks/queryBuilder/useOptions.ts
index 7f50f1881b..e990f789de 100644
--- a/frontend/src/hooks/queryBuilder/useOptions.ts
+++ b/frontend/src/hooks/queryBuilder/useOptions.ts
@@ -170,11 +170,7 @@ export const useOptions = (
 					(option, index, self) =>
 						index ===
 							self.findIndex(
-								(o) =>
-									// to remove duplicate & empty options from list
-									o.label === option.label &&
-									o.value === option.value &&
-									o.dataType?.toLowerCase() === option.dataType?.toLowerCase(), // handle case sensitivity
+								(o) => o.label === option.label && o.value === option.value, // to remove duplicate & empty options from list
 							) && option.value !== '',
 				) || []
 			).map((option) => {