revert: revert check for metric type without query range constraint

* fix: guard against undefined spread in useGetQueryLabels

* chore: add tests

frontend/src/container/DashboardContainer/DashboardDescription/__tests__/DashboardDescription.test.tsx

@@ -193,7 +193,6 @@ describe('Dashboard landing page actions header tests', () => {
 			handleDashboardLockToggle: jest.fn(),
 			dashboardResponse: {} as IDashboardContext['dashboardResponse'],
 			selectedDashboard: (getDashboardById.data as unknown) as Dashboard,
-			dashboardId: '4',
 			layouts: [],
 			panelMap: {},
 			setPanelMap: jest.fn(),
@@ -205,8 +204,6 @@ describe('Dashboard landing page actions header tests', () => {
 			updateLocalStorageDashboardVariables: jest.fn(),
 			dashboardQueryRangeCalled: false,
 			setDashboardQueryRangeCalled: jest.fn(),
-			selectedRowWidgetId: null,
-			setSelectedRowWidgetId: jest.fn(),
 			isDashboardFetching: false,
 			columnWidths: {},
 			setColumnWidths: jest.fn(),

frontend/src/container/DashboardContainer/DashboardDescription/index.tsx

@@ -78,7 +78,6 @@ function DashboardDescription(props: DashboardDescriptionProps): JSX.Element {
 		isDashboardLocked,
 		setSelectedDashboard,
 		handleToggleDashboardSlider,
-		setSelectedRowWidgetId,
 		handleDashboardLockToggle,
 	} = useDashboard();
 
@@ -146,7 +145,6 @@ function DashboardDescription(props: DashboardDescriptionProps): JSX.Element {
 	const [addPanelPermission] = useComponentPermission(permissions, userRole);
 
 	const onEmptyWidgetHandler = useCallback(() => {
-		setSelectedRowWidgetId(null);
 		handleToggleDashboardSlider(true);
 		logEvent('Dashboard Detail: Add new panel clicked', {
 			dashboardId: selectedDashboard?.id,

frontend/src/container/DashboardContainer/DashboardVariablesSelection/useDashboardVariableUpdate.ts

@@ -67,17 +67,18 @@ export const useDashboardVariableUpdate = (): UseDashboardVariableUpdateReturn =
 							const oldVariables = prev?.data.variables;
 							// this is added to handle case where we have two different
 							// schemas for variable response
-							if (oldVariables?.[id]) {
-								oldVariables[id] = {
-									...oldVariables[id],
+							const updatedVariables = { ...oldVariables };
+							if (updatedVariables?.[id]) {
+								updatedVariables[id] = {
+									...updatedVariables[id],
 									selectedValue: value,
 									allSelected,
 									haveCustomValuesSelected,
 								};
 							}
-							if (oldVariables?.[name]) {
-								oldVariables[name] = {
-									...oldVariables[name],
+							if (updatedVariables?.[name]) {
+								updatedVariables[name] = {
+									...updatedVariables[name],
 									selectedValue: value,
 									allSelected,
 									haveCustomValuesSelected,
@@ -87,9 +88,7 @@ export const useDashboardVariableUpdate = (): UseDashboardVariableUpdateReturn =
 								...prev,
 								data: {
 									...prev?.data,
-									variables: {
-										...oldVariables,
-									},
+									variables: updatedVariables,
 								},
 							};
 						}

frontend/src/container/DashboardContainer/visualization/panels/BarPanel/BarPanel.tsx

@@ -6,7 +6,6 @@ import { useResizeObserver } from 'hooks/useDimensions';
 import { LegendPosition } from 'lib/uPlotV2/components/types';
 import ContextMenu from 'periscope/components/ContextMenu';
 import { useTimezone } from 'providers/Timezone';
-import { MetricRangePayloadProps } from 'types/api/metrics/getQueryRange';
 import uPlot from 'uplot';
 import { getTimeRange } from 'utils/getTimeRange';
 
@@ -62,7 +61,7 @@ function BarPanel(props: PanelWrapperProps): JSX.Element {
 			currentQuery: widget.query,
 			onClick: clickHandlerWithContextMenu,
 			onDragSelect,
-			apiResponse: queryResponse?.data?.payload as MetricRangePayloadProps,
+			apiResponse: queryResponse?.data?.payload,
 			timezone,
 			panelMode,
 			minTimeScale: minTimeScale,

frontend/src/container/DashboardContainer/visualization/panels/BarPanel/utils.ts

@@ -11,7 +11,6 @@ import { get } from 'lodash-es';
 import { Widgets } from 'types/api/dashboard/getAll';
 import { MetricRangePayloadProps } from 'types/api/metrics/getQueryRange';
 import { Query } from 'types/api/queryBuilder/queryBuilderData';
-import { QueryData } from 'types/api/widgets/getQuery';
 import { AlignedData } from 'uplot';
 
 import { PanelMode } from '../types';
@@ -44,7 +43,7 @@ export function prepareBarPanelConfig({
 	currentQuery: Query;
 	onClick: OnClickPluginOpts['onClick'];
 	onDragSelect: (startTime: number, endTime: number) => void;
-	apiResponse: MetricRangePayloadProps;
+	apiResponse?: MetricRangePayloadProps;
 	timezone: Timezone;
 	panelMode: PanelMode;
 	minTimeScale?: number;
@@ -76,13 +75,17 @@ export function prepareBarPanelConfig({
 		stepInterval: minStepInterval,
 	});
 
+	if (!(apiResponse && apiResponse?.data?.result)) {
+		// if no data, return the builder without adding any series
+		return builder;
+	}
+
 	if (widget.stackedBarChart) {
-		const seriesCount = (apiResponse?.data?.result?.length ?? 0) + 1; // +1 for 1-based uPlot series indices
+		const seriesCount = (apiResponse.data.result.length ?? 0) + 1; // +1 for 1-based uPlot series indices
 		builder.setBands(getInitialStackedBands(seriesCount));
 	}
 
-	const seriesList: QueryData[] = apiResponse?.data?.result || [];
-	seriesList.forEach((series) => {
+	apiResponse.data.result.forEach((series) => {
 		const baseLabelName = getLabelName(
 			series.metric,
 			series.queryName || '', // query

frontend/src/container/DashboardContainer/visualization/panels/HistogramPanel/HistogramPanel.tsx

@@ -6,7 +6,6 @@ import { useResizeObserver } from 'hooks/useDimensions';
 import { LegendPosition } from 'lib/uPlotV2/components/types';
 import { DashboardCursorSync } from 'lib/uPlotV2/plugins/TooltipPlugin/types';
 import { useTimezone } from 'providers/Timezone';
-import { MetricRangePayloadProps } from 'types/api/metrics/getQueryRange';
 import uPlot from 'uplot';
 
 import Histogram from '../../charts/Histogram/Histogram';
@@ -39,7 +38,7 @@ function HistogramPanel(props: PanelWrapperProps): JSX.Element {
 		return prepareHistogramPanelConfig({
 			widget,
 			isDarkMode,
-			apiResponse: queryResponse?.data?.payload as MetricRangePayloadProps,
+			apiResponse: queryResponse?.data?.payload,
 			panelMode,
 		});
 	}, [widget, isDarkMode, queryResponse?.data?.payload, panelMode]);
@@ -49,7 +48,7 @@ function HistogramPanel(props: PanelWrapperProps): JSX.Element {
 			return [];
 		}
 		return prepareHistogramPanelData({
-			apiResponse: queryResponse?.data?.payload as MetricRangePayloadProps,
+			apiResponse: queryResponse?.data?.payload,
 			bucketWidth: widget?.bucketWidth,
 			bucketCount: widget?.bucketCount,
 			mergeAllActiveQueries: widget?.mergeAllActiveQueries,

frontend/src/container/DashboardContainer/visualization/panels/HistogramPanel/utils.ts

@@ -149,7 +149,7 @@ export function prepareHistogramPanelConfig({
 	isDarkMode,
 }: {
 	widget: Widgets;
-	apiResponse: MetricRangePayloadProps;
+	apiResponse?: MetricRangePayloadProps;
 	panelMode: PanelMode;
 	isDarkMode: boolean;
 }): UPlotConfigBuilder {
@@ -204,7 +204,7 @@ export function prepareHistogramPanelConfig({
 			fillColor: '#4E74F8',
 			isDarkMode,
 		});
-	} else {
+	} else if (apiResponse && apiResponse?.data?.result) {
 		apiResponse.data.result.forEach((series) => {
 			const baseLabelName = getLabelName(
 				series.metric,

frontend/src/container/DashboardContainer/visualization/panels/TimeSeriesPanel/TimeSeriesPanel.tsx

@@ -9,7 +9,6 @@ import { useResizeObserver } from 'hooks/useDimensions';
 import { LegendPosition } from 'lib/uPlotV2/components/types';
 import { ContextMenu } from 'periscope/components/ContextMenu';
 import { useTimezone } from 'providers/Timezone';
-import { MetricRangePayloadProps } from 'types/api/metrics/getQueryRange';
 import uPlot from 'uplot';
 import { getTimeRange } from 'utils/getTimeRange';
 
@@ -68,7 +67,7 @@ function TimeSeriesPanel(props: PanelWrapperProps): JSX.Element {
 			currentQuery: widget.query,
 			onClick: clickHandlerWithContextMenu,
 			onDragSelect,
-			apiResponse: queryResponse?.data?.payload as MetricRangePayloadProps,
+			apiResponse: queryResponse?.data?.payload,
 			timezone,
 			panelMode,
 			minTimeScale: minTimeScale,

frontend/src/container/DashboardContainer/visualization/panels/TimeSeriesPanel/utils.ts

@@ -68,11 +68,12 @@ export const prepareUPlotConfig = ({
 	currentQuery: Query;
 	onClick?: OnClickPluginOpts['onClick'];
 	onDragSelect: (startTime: number, endTime: number) => void;
-	apiResponse: MetricRangePayloadProps;
+	apiResponse?: MetricRangePayloadProps;
 	timezone: Timezone;
 	panelMode: PanelMode;
 	minTimeScale?: number;
 	maxTimeScale?: number;
+	// eslint-disable-next-line sonarjs/cognitive-complexity
 }): UPlotConfigBuilder => {
 	const stepIntervals: ExecStats['stepIntervals'] = get(
 		apiResponse,
@@ -100,7 +101,12 @@ export const prepareUPlotConfig = ({
 		stepInterval: minStepInterval,
 	});
 
-	apiResponse.data?.result?.forEach((series) => {
+	if (!(apiResponse && apiResponse.data.result)) {
+		// if no data, return the builder without adding any series
+		return builder;
+	}
+
+	apiResponse.data.result.forEach((series) => {
 		const hasSingleValidPoint = hasSingleVisiblePointForSeries(series);
 		const baseLabelName = getLabelName(
 			series.metric,

frontend/src/container/DashboardContainer/visualization/panels/utils/baseConfigBuilder.ts

@@ -18,7 +18,7 @@ import { PanelMode } from '../types';
 export interface BaseConfigBuilderProps {
 	id: string;
 	thresholds?: ThresholdProps[];
-	apiResponse: MetricRangePayloadProps;
+	apiResponse?: MetricRangePayloadProps;
 	isDarkMode: boolean;
 	onClick?: OnClickPluginOpts['onClick'];
 	onDragSelect?: (startTime: number, endTime: number) => void;

frontend/src/container/GridCardLayout/DashboardEmptyState/DashboardEmptyState.tsx

@@ -19,7 +19,6 @@ export default function DashboardEmptyState(): JSX.Element {
 		selectedDashboard,
 		isDashboardLocked,
 		handleToggleDashboardSlider,
-		setSelectedRowWidgetId,
 	} = useDashboard();
 
 	const variablesSettingsTabHandle = useRef<VariablesSettingsTab>(null);
@@ -42,7 +41,6 @@ export default function DashboardEmptyState(): JSX.Element {
 	const [addPanelPermission] = useComponentPermission(permissions, userRole);
 
 	const onEmptyWidgetHandler = useCallback(() => {
-		setSelectedRowWidgetId(null);
 		handleToggleDashboardSlider(true);
 		logEvent('Dashboard Detail: Add new panel clicked', {
 			dashboardId: selectedDashboard?.id,

frontend/src/container/GridCardLayout/GridCardLayout.tsx

@@ -71,7 +71,6 @@ function GraphLayout(props: GraphLayoutProps): JSX.Element {
 		isDashboardLocked,
 		dashboardQueryRangeCalled,
 		setDashboardQueryRangeCalled,
-		setSelectedRowWidgetId,
 		isDashboardFetching,
 		columnWidths,
 	} = useDashboard();
@@ -195,7 +194,6 @@ function GraphLayout(props: GraphLayoutProps): JSX.Element {
 
 		updateDashboardMutation.mutate(updatedDashboard, {
 			onSuccess: (updatedDashboard) => {
-				setSelectedRowWidgetId(null);
 				if (updatedDashboard.data) {
 					if (updatedDashboard.data.data.layout) {
 						setLayouts(sortLayout(updatedDashboard.data.data.layout));

frontend/src/container/GridCardLayout/WidgetRow.tsx

@@ -5,6 +5,7 @@ import useComponentPermission from 'hooks/useComponentPermission';
 import { EllipsisIcon, PenLine, Plus, X } from 'lucide-react';
 import { useAppContext } from 'providers/App/App';
 import { useDashboard } from 'providers/Dashboard/Dashboard';
+import { setSelectedRowWidgetId } from 'providers/Dashboard/helpers/selectedRowWidgetIdHelper';
 import { ROLES, USER_ROLES } from 'types/roles';
 import { ComponentTypes } from 'utils/permission';
 
@@ -37,7 +38,6 @@ export function WidgetRowHeader(props: WidgetRowHeaderProps): JSX.Element {
 		handleToggleDashboardSlider,
 		selectedDashboard,
 		isDashboardLocked,
-		setSelectedRowWidgetId,
 	} = useDashboard();
 
 	const permissions: ComponentTypes[] = ['add_panel'];
@@ -81,7 +81,12 @@ export function WidgetRowHeader(props: WidgetRowHeaderProps): JSX.Element {
 							disabled={!editWidget && addPanelPermission && !isDashboardLocked}
 							icon={<Plus size={14} />}
 							onClick={(): void => {
-								setSelectedRowWidgetId(id);
+								// TODO: @AshwinBhatkal Simplify this check in cleanup of https://github.com/SigNoz/engineering-pod/issues/3953
+								if (!selectedDashboard?.id) {
+									return;
+								}
+
+								setSelectedRowWidgetId(selectedDashboard.id, id);
 								handleToggleDashboardSlider(true);
 							}}
 						>

frontend/src/container/NewWidget/index.tsx

@@ -34,6 +34,10 @@ import { cloneDeep, defaultTo, isEmpty, isUndefined } from 'lodash-es';
 import { Check, X } from 'lucide-react';
 import { DashboardWidgetPageParams } from 'pages/DashboardWidget';
 import { useDashboard } from 'providers/Dashboard/Dashboard';
+import {
+	clearSelectedRowWidgetId,
+	getSelectedRowWidgetId,
+} from 'providers/Dashboard/helpers/selectedRowWidgetIdHelper';
 import {
 	getNextWidgets,
 	getPreviousWidgets,
@@ -86,8 +90,6 @@ function NewWidget({
 		selectedDashboard,
 		setSelectedDashboard,
 		setToScrollWidgetId,
-		selectedRowWidgetId,
-		setSelectedRowWidgetId,
 		columnWidths,
 	} = useDashboard();
 
@@ -450,6 +452,8 @@ function NewWidget({
 		const widgetId = query.get('widgetId') || '';
 		let updatedLayout = selectedDashboard.data.layout || [];
 
+		const selectedRowWidgetId = getSelectedRowWidgetId(dashboardId);
+
 		if (isNewDashboard && isEmpty(selectedRowWidgetId)) {
 			const newLayoutItem = placeWidgetAtBottom(widgetId, updatedLayout);
 			updatedLayout = [...updatedLayout, newLayoutItem];
@@ -554,7 +558,6 @@ function NewWidget({
 
 		updateDashboardMutation.mutateAsync(dashboard, {
 			onSuccess: (updatedDashboard) => {
-				setSelectedRowWidgetId(null);
 				setSelectedDashboard(updatedDashboard.data);
 				setToScrollWidgetId(selectedWidget?.id || '');
 				safeNavigate({
@@ -566,7 +569,6 @@ function NewWidget({
 		selectedDashboard,
 		query,
 		isNewDashboard,
-		selectedRowWidgetId,
 		afterWidgets,
 		selectedWidget,
 		selectedTime.enum,
@@ -577,7 +579,6 @@ function NewWidget({
 		widgets,
 		setSelectedDashboard,
 		setToScrollWidgetId,
-		setSelectedRowWidgetId,
 		safeNavigate,
 		dashboardId,
 	]);
@@ -681,6 +682,10 @@ function NewWidget({
 		 * on mount here with the currentQuery in the begining itself
 		 */
 		setSupersetQuery(currentQuery);
+
+		return (): void => {
+			clearSelectedRowWidgetId(dashboardId);
+		};
 		// eslint-disable-next-line react-hooks/exhaustive-deps
 	}, []);
 

frontend/src/hooks/useGetQueryLabels.test.ts

@@ -0,0 +1,129 @@
+import { renderHook } from '@testing-library/react';
+import {
+	IBuilderFormula,
+	IClickHouseQuery,
+	IPromQLQuery,
+	Query,
+} from 'types/api/queryBuilder/queryBuilderData';
+import { EQueryType } from 'types/common/dashboard';
+
+import { useGetQueryLabels } from './useGetQueryLabels';
+
+jest.mock('components/QueryBuilderV2/utils', () => ({
+	getQueryLabelWithAggregation: jest.fn(() => []),
+}));
+
+function buildQuery(overrides: Partial<Query> = {}): Query {
+	return {
+		id: 'test-id',
+		queryType: EQueryType.QUERY_BUILDER,
+		builder: {
+			queryData: [],
+			queryFormulas: [],
+			queryTraceOperator: [],
+		},
+		promql: [],
+		clickhouse_sql: [],
+		...overrides,
+	};
+}
+
+describe('useGetQueryLabels', () => {
+	describe('QUERY_BUILDER type', () => {
+		it('returns empty array when queryFormulas is undefined', () => {
+			const query = buildQuery({
+				queryType: EQueryType.QUERY_BUILDER,
+				builder: {
+					queryData: [],
+					queryFormulas: (undefined as unknown) as IBuilderFormula[],
+					queryTraceOperator: [],
+				},
+			});
+
+			const { result } = renderHook(() => useGetQueryLabels(query));
+
+			expect(result.current).toEqual([]);
+		});
+
+		it('returns formula labels when queryFormulas is populated', () => {
+			const query = buildQuery({
+				queryType: EQueryType.QUERY_BUILDER,
+				builder: {
+					queryData: [],
+					queryFormulas: [
+						({ queryName: 'F1' } as unknown) as IBuilderFormula,
+						({ queryName: 'F2' } as unknown) as IBuilderFormula,
+					],
+					queryTraceOperator: [],
+				},
+			});
+
+			const { result } = renderHook(() => useGetQueryLabels(query));
+
+			expect(result.current).toEqual([
+				{ label: 'F1', value: 'F1' },
+				{ label: 'F2', value: 'F2' },
+			]);
+		});
+	});
+
+	describe('CLICKHOUSE type', () => {
+		it('returns empty array when clickhouse_sql is undefined', () => {
+			const query = buildQuery({
+				queryType: EQueryType.CLICKHOUSE,
+				clickhouse_sql: (undefined as unknown) as IClickHouseQuery[],
+			});
+
+			const { result } = renderHook(() => useGetQueryLabels(query));
+
+			expect(result.current).toEqual([]);
+		});
+
+		it('returns labels from clickhouse_sql when populated', () => {
+			const query = buildQuery({
+				queryType: EQueryType.CLICKHOUSE,
+				clickhouse_sql: [
+					({ name: 'query_a' } as unknown) as IClickHouseQuery,
+					({ name: 'query_b' } as unknown) as IClickHouseQuery,
+				],
+			});
+
+			const { result } = renderHook(() => useGetQueryLabels(query));
+
+			expect(result.current).toEqual([
+				{ label: 'query_a', value: 'query_a' },
+				{ label: 'query_b', value: 'query_b' },
+			]);
+		});
+	});
+
+	describe('PROM type (default)', () => {
+		it('returns empty array when promql is undefined', () => {
+			const query = buildQuery({
+				queryType: EQueryType.PROM,
+				promql: (undefined as unknown) as IPromQLQuery[],
+			});
+
+			const { result } = renderHook(() => useGetQueryLabels(query));
+
+			expect(result.current).toEqual([]);
+		});
+
+		it('returns labels from promql when populated', () => {
+			const query = buildQuery({
+				queryType: EQueryType.PROM,
+				promql: [
+					({ name: 'prom_1' } as unknown) as IPromQLQuery,
+					({ name: 'prom_2' } as unknown) as IPromQLQuery,
+				],
+			});
+
+			const { result } = renderHook(() => useGetQueryLabels(query));
+
+			expect(result.current).toEqual([
+				{ label: 'prom_1', value: 'prom_1' },
+				{ label: 'prom_2', value: 'prom_2' },
+			]);
+		});
+	});
+});

frontend/src/hooks/useGetQueryLabels.ts

@@ -11,7 +11,7 @@ export const useGetQueryLabels = (
 			const queryLabels = getQueryLabelWithAggregation(
 				currentQuery?.builder?.queryData || [],
 			);
-			const formulaLabels = currentQuery?.builder?.queryFormulas?.map(
+			const formulaLabels = (currentQuery?.builder?.queryFormulas ?? []).map(
 				(formula) => ({
 					label: formula.queryName,
 					value: formula.queryName,
@@ -20,10 +20,13 @@ export const useGetQueryLabels = (
 			return [...queryLabels, ...formulaLabels];
 		}
 		if (currentQuery?.queryType === EQueryType.CLICKHOUSE) {
-			return currentQuery?.clickhouse_sql?.map((q) => ({
+			return (currentQuery?.clickhouse_sql ?? []).map((q) => ({
 				label: q.name,
 				value: q.name,
 			}));
 		}
-		return currentQuery?.promql?.map((q) => ({ label: q.name, value: q.name }));
+		return (currentQuery?.promql ?? []).map((q) => ({
+			label: q.name,
+			value: q.name,
+		}));
 	}, [currentQuery]);

frontend/src/providers/Dashboard/Dashboard.tsx

@@ -68,7 +68,6 @@ export const DashboardContext = createContext<IDashboardContext>({
 		APIError
 	>,
 	selectedDashboard: {} as Dashboard,
-	dashboardId: '',
 	layouts: [],
 	panelMap: {},
 	setPanelMap: () => {},
@@ -81,8 +80,6 @@ export const DashboardContext = createContext<IDashboardContext>({
 	updateLocalStorageDashboardVariables: () => {},
 	dashboardQueryRangeCalled: false,
 	setDashboardQueryRangeCalled: () => {},
-	selectedRowWidgetId: '',
-	setSelectedRowWidgetId: () => {},
 	isDashboardFetching: false,
 	columnWidths: {},
 	setColumnWidths: () => {},
@@ -102,10 +99,6 @@ export function DashboardProvider({
 
 	const [isDashboardLocked, setIsDashboardLocked] = useState<boolean>(false);
 
-	const [selectedRowWidgetId, setSelectedRowWidgetId] = useState<string | null>(
-		null,
-	);
-
 	const [
 		dashboardQueryRangeCalled,
 		setDashboardQueryRangeCalled,
@@ -468,8 +461,6 @@ export function DashboardProvider({
 			updateLocalStorageDashboardVariables,
 			dashboardQueryRangeCalled,
 			setDashboardQueryRangeCalled,
-			selectedRowWidgetId,
-			setSelectedRowWidgetId,
 			isDashboardFetching,
 			columnWidths,
 			setColumnWidths,
@@ -488,8 +479,6 @@ export function DashboardProvider({
 			currentDashboard,
 			dashboardQueryRangeCalled,
 			setDashboardQueryRangeCalled,
-			selectedRowWidgetId,
-			setSelectedRowWidgetId,
 			isDashboardFetching,
 			columnWidths,
 			setColumnWidths,

frontend/src/providers/Dashboard/__tests__/Dashboard.test.tsx

@@ -58,8 +58,9 @@ jest.mock('react-redux', () => ({
 jest.mock('uuid', () => ({ v4: jest.fn(() => 'mock-uuid') }));
 
 function TestComponent(): JSX.Element {
-	const { dashboardResponse, dashboardId, selectedDashboard } = useDashboard();
+	const { dashboardResponse, selectedDashboard } = useDashboard();
 	const { dashboardVariables } = useDashboardVariables();
+	const dashboardId = selectedDashboard?.id;
 
 	return (
 		<div>

frontend/src/providers/Dashboard/helpers/selectedRowWidgetIdHelper.ts

@@ -0,0 +1,28 @@
+const PREFIX = 'dashboard_row_widget_';
+
+function getKey(dashboardId: string): string {
+	return `${PREFIX}${dashboardId}`;
+}
+
+export function setSelectedRowWidgetId(
+	dashboardId: string,
+	widgetId: string,
+): void {
+	const key = getKey(dashboardId);
+
+	// remove all other selected widget ids for the dashboard before setting the new one
+	// to ensure only one widget is selected at a time. Helps out in weird navigate and refresh scenarios
+	Object.keys(sessionStorage)
+		.filter((k) => k.startsWith(PREFIX) && k !== key)
+		.forEach((k) => sessionStorage.removeItem(k));
+
+	sessionStorage.setItem(key, widgetId);
+}
+
+export function getSelectedRowWidgetId(dashboardId: string): string | null {
+	return sessionStorage.getItem(getKey(dashboardId));
+}
+
+export function clearSelectedRowWidgetId(dashboardId: string): void {
+	sessionStorage.removeItem(getKey(dashboardId));
+}

frontend/src/providers/Dashboard/types.ts

@@ -15,7 +15,6 @@ export interface IDashboardContext {
 	handleDashboardLockToggle: (value: boolean) => void;
 	dashboardResponse: UseQueryResult<SuccessResponseV2<Dashboard>, unknown>;
 	selectedDashboard: Dashboard | undefined;
-	dashboardId: string;
 	layouts: Layout[];
 	panelMap: Record<string, { widgets: Layout[]; collapsed: boolean }>;
 	setPanelMap: React.Dispatch<React.SetStateAction<Record<string, any>>>;
@@ -40,8 +39,6 @@ export interface IDashboardContext {
 	) => void;
 	dashboardQueryRangeCalled: boolean;
 	setDashboardQueryRangeCalled: (value: boolean) => void;
-	selectedRowWidgetId: string | null;
-	setSelectedRowWidgetId: React.Dispatch<React.SetStateAction<string | null>>;
 	isDashboardFetching: boolean;
 	columnWidths: WidgetColumnWidths;
 	setColumnWidths: React.Dispatch<React.SetStateAction<WidgetColumnWidths>>;

pkg/querier/postprocess.go

@@ -185,6 +185,22 @@ func postProcessMetricQuery(
 	query qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation],
 	req *qbtypes.QueryRangeRequest,
 ) *qbtypes.Result {
+
+	config := query.Aggregations[0]
+	spaceAggOrderBy := fmt.Sprintf("%s(%s)", config.SpaceAggregation.StringValue(), config.MetricName)
+	timeAggOrderBy := fmt.Sprintf("%s(%s)", config.TimeAggregation.StringValue(), config.MetricName)
+	timeSpaceAggOrderBy := fmt.Sprintf("%s(%s(%s))", config.SpaceAggregation.StringValue(), config.TimeAggregation.StringValue(), config.MetricName)
+
+	for idx := range query.Order {
+		if query.Order[idx].Key.Name == spaceAggOrderBy ||
+			query.Order[idx].Key.Name == timeAggOrderBy ||
+			query.Order[idx].Key.Name == timeSpaceAggOrderBy {
+			query.Order[idx].Key.Name = qbtypes.DefaultOrderByKey
+		}
+	}
+
+	result = q.applySeriesLimit(result, query.Limit, query.Order)
+
 	if len(query.Functions) > 0 {
 		step := query.StepInterval.Duration.Milliseconds()
 		functions := q.prepareFillZeroArgsWithStep(query.Functions, req, step)

pkg/querybuilder/fallback_expr.go

@@ -132,14 +132,6 @@ func GroupByKeys(keys []qbtypes.GroupByKey) []string {
 	return k
 }
 
-func OrderByKeys(keys []qbtypes.OrderBy) []string {
-	k := []string{}
-	for _, key := range keys {
-		k = append(k, "`"+key.Key.Name+"`")
-	}
-	return k
-}
-
 func FormatValueForContains(value any) string {
 	if value == nil {
 		return ""

pkg/telemetrymeter/stmt_builder_test.go

@@ -51,7 +51,7 @@ func TestStatementBuilder(t *testing.T) {
 				},
 			},
 			expected: qbtypes.Statement{
-				Query: "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'service.name') AS `service.name`, max(value) AS per_series_value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'service.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) ORDER BY `service.name`, ts ASC",
+				Query: "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'service.name') AS `service.name`, max(value) AS per_series_value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'service.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name`, ts",
 				Args:  []any{"signoz_calls_total", uint64(1747785600000), uint64(1747983420000), "cartservice", "cumulative", 0},
 			},
 			expectedErr: nil,
@@ -84,7 +84,7 @@ func TestStatementBuilder(t *testing.T) {
 				},
 			},
 			expected: qbtypes.Statement{
-				Query: "WITH __spatial_aggregation_cte AS (SELECT toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'service.name') AS `service.name`, sum(value)/86400 AS value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'service.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) ORDER BY `service.name`, ts ASC",
+				Query: "WITH __spatial_aggregation_cte AS (SELECT toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'service.name') AS `service.name`, sum(value)/86400 AS value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'service.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name`, ts",
 				Args:  []any{"signoz_calls_total", uint64(1747872000000), uint64(1747983420000), "cartservice", "delta"},
 			},
 			expectedErr: nil,
@@ -117,7 +117,7 @@ func TestStatementBuilder(t *testing.T) {
 				},
 			},
 			expected: qbtypes.Statement{
-				Query: "WITH __temporal_aggregation_cte AS (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'service.name') AS `service.name`, sum(value)/86400 AS per_series_value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'service.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts), __spatial_aggregation_cte AS (SELECT ts, `service.name`, avg(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) ORDER BY `service.name`, ts ASC",
+				Query: "WITH __temporal_aggregation_cte AS (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'service.name') AS `service.name`, sum(value)/86400 AS per_series_value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'service.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts), __spatial_aggregation_cte AS (SELECT ts, `service.name`, avg(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name`, ts",
 				Args:  []any{"signoz_calls_total", uint64(1747872000000), uint64(1747983420000), "cartservice", "delta", 0},
 			},
 			expectedErr: nil,
@@ -150,7 +150,7 @@ func TestStatementBuilder(t *testing.T) {
 				},
 			},
 			expected: qbtypes.Statement{
-				Query: "WITH __temporal_aggregation_cte AS (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'host.name') AS `host.name`, avg(value) AS per_series_value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'host.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY fingerprint, ts, `host.name` ORDER BY fingerprint, ts), __spatial_aggregation_cte AS (SELECT ts, `host.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `host.name`) SELECT * FROM __spatial_aggregation_cte WHERE (`host.name`) IN (SELECT `host.name` FROM __spatial_aggregation_cte GROUP BY `host.name` ORDER BY avg(value) LIMIT 10) ORDER BY `host.name`, ts ASC",
+				Query: "WITH __temporal_aggregation_cte AS (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(86400)) AS ts, JSONExtractString(labels, 'host.name') AS `host.name`, avg(value) AS per_series_value FROM signoz_meter.distributed_samples AS points WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? AND JSONExtractString(labels, 'host.name') = ? AND LOWER(temporality) LIKE LOWER(?) GROUP BY fingerprint, ts, `host.name` ORDER BY fingerprint, ts), __spatial_aggregation_cte AS (SELECT ts, `host.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `host.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `host.name`, ts",
 				Args:  []any{"system.memory.usage", uint64(1747872000000), uint64(1747983420000), "big-data-node-1", "unspecified", 0},
 			},
 			expectedErr: nil,

pkg/telemetrymetrics/statement_builder.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"log/slog"
-	"strings"
 
 	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/flagger"
@@ -608,73 +607,8 @@ func (b *MetricQueryStatementBuilder) BuildFinalSelect(
 			sb.Where(rewrittenExpr)
 		}
 	}
-	groupByKeys := querybuilder.GroupByKeys(query.GroupBy)
-	hasOrder := len(query.Order) > 0
-	hasLimit := query.Limit > 0
-	hasGroupBy := len(groupByKeys) > 0
-
-	if hasOrder && hasLimit {
-		// order by with limit: add WHERE subquery to restrict to top N distinct key combinations
-		orderByKeys := querybuilder.OrderByKeys(query.Order)
-		var orderByClauses []string
-		for _, o := range query.Order {
-			orderByClauses = append(orderByClauses, fmt.Sprintf("`%s` %s", o.Key.Name, o.Direction.StringValue()))
-		}
-
-		subSb := sqlbuilder.NewSelectBuilder()
-		subSb.Select(fmt.Sprintf("DISTINCT %s", orderByKeys[0]))
-		if len(orderByKeys) > 1 {
-			subSb.SelectMore(orderByKeys[1:]...)
-		}
-		subSb.From("__spatial_aggregation_cte")
-		subSb.OrderBy(orderByClauses...)
-
-		subQ, _ := subSb.BuildWithFlavor(sqlbuilder.ClickHouse)
-		subQ = fmt.Sprintf("%s LIMIT %d", subQ, query.Limit)
-		sb.Where(fmt.Sprintf("(%s) IN (%s)", strings.Join(orderByKeys, ", "), subQ))
-
-		sb.OrderBy(orderByClauses...)
-	} else if hasOrder {
-		// order by without limit: apply order by clauses directly
-		for _, o := range query.Order {
-			key := o.Key.Name
-			if strings.Contains(key, query.Aggregations[0].MetricName) {
-				sb.OrderBy(fmt.Sprintf("avg(value) OVER (PARTITION BY %s) %s", strings.Join(groupByKeys, ", "), o.Direction.StringValue()))
-				continue
-			}
-			sb.OrderBy(fmt.Sprintf("`%s` %s", o.Key.Name, o.Direction.StringValue()))
-		}
-	} else if hasLimit && hasGroupBy {
-		// limit without order by: default ordering by avg(value)
-		subSb := sqlbuilder.NewSelectBuilder()
-		subSb.Select(groupByKeys[0])
-		if len(groupByKeys) > 1 {
-			subSb.SelectMore(groupByKeys[1:]...)
-		}
-		subSb.From("__spatial_aggregation_cte")
-		subSb.GroupBy(groupByKeys...)
-		subSb.OrderBy("avg(value)")
-
-		subQ, _ := subSb.BuildWithFlavor(sqlbuilder.ClickHouse)
-		subQ = fmt.Sprintf("%s LIMIT %d", subQ, query.Limit)
-		sb.Where(fmt.Sprintf("(%s) IN (%s)", strings.Join(groupByKeys, ", "), subQ))
-	} else if hasGroupBy {
-		// grouping without order by or limit: sort by avg(value) DESC with labels as tiebreakers
-		sb.OrderBy(fmt.Sprintf("avg(value) OVER (PARTITION BY %s) DESC", strings.Join(groupByKeys, ", ")))
-	}
-
-	// add any group-by keys not already in the order-by as tiebreakers
-	orderKeySet := make(map[string]struct{})
-	for _, o := range query.Order {
-		orderKeySet[fmt.Sprintf("`%s`", o.Key.Name)] = struct{}{}
-	}
-	for _, g := range groupByKeys {
-		if _, exists := orderKeySet[g]; !exists {
-			sb.OrderBy(g)
-		}
-	}
-
-	sb.OrderBy("ts ASC")
+	sb.OrderBy(querybuilder.GroupByKeys(query.GroupBy)...)
+	sb.OrderBy("ts")
 	if metricType == metrictypes.HistogramType && spaceAgg == metrictypes.SpaceAggregationCount && query.Aggregations[0].ComparisonSpaceAggregationParam == nil {
 		sb.OrderBy("toFloat64(le)")
 	}

pkg/telemetrymetrics/stmt_builder_test.go

@@ -15,17 +15,16 @@ import (
 )
 
 func TestStatementBuilder(t *testing.T) {
-	type baseQuery struct {
-		name     string
-		query    qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]
-		orderKey string
-		args     []any
-		cte      string
-	}
-
-	bases := []baseQuery{
+	cases := []struct {
+		name        string
+		requestType qbtypes.RequestType
+		query       qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]
+		expected    qbtypes.Statement
+		expectedErr error
+	}{
 		{
-			name: "cumulative_rate_sum",
+			name:        "test_cumulative_rate_sum",
+			requestType: qbtypes.RequestTypeTimeSeries,
 			query: qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]{
 				Signal:       telemetrytypes.SignalMetrics,
 				StepInterval: qbtypes.Step{Duration: 30 * time.Second},
@@ -41,16 +40,24 @@ func TestStatementBuilder(t *testing.T) {
 				Filter: &qbtypes.Filter{
 					Expression: "service.name = 'cartservice'",
 				},
+				Limit: 10,
 				GroupBy: []qbtypes.GroupByKey{
-					{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "service.name"}},
+					{
+						TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
+							Name: "service.name",
+						},
+					},
 				},
 			},
-			orderKey: "service.name",
-			args:     []any{"signoz_calls_total", uint64(1747936800000), uint64(1747983420000), "cumulative", false, "cartservice", "signoz_calls_total", uint64(1747947360000), uint64(1747983420000), 0},
-			cte:      "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, max(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'service.name') = ? GROUP BY fingerprint, `service.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`)",
+			expected: qbtypes.Statement{
+				Query: "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, max(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'service.name') = ? GROUP BY fingerprint, `service.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name`, ts",
+				Args:  []any{"signoz_calls_total", uint64(1747936800000), uint64(1747983420000), "cumulative", false, "cartservice", "signoz_calls_total", uint64(1747947360000), uint64(1747983420000), 0},
+			},
+			expectedErr: nil,
 		},
 		{
-			name: "cumulative_rate_sum_with_mat_column",
+			name:        "test_cumulative_rate_sum_with_mat_column",
+			requestType: qbtypes.RequestTypeTimeSeries,
 			query: qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]{
 				Signal:       telemetrytypes.SignalMetrics,
 				StepInterval: qbtypes.Step{Duration: 30 * time.Second},
@@ -66,16 +73,24 @@ func TestStatementBuilder(t *testing.T) {
 				Filter: &qbtypes.Filter{
 					Expression: "materialized.key.name REGEXP 'cartservice' OR service.name = 'cartservice'",
 				},
+				Limit: 10,
 				GroupBy: []qbtypes.GroupByKey{
-					{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "service.name"}},
+					{
+						TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
+							Name: "service.name",
+						},
+					},
 				},
 			},
-			orderKey: "service.name",
-			args:     []any{"signoz_calls_total", uint64(1747936800000), uint64(1747983420000), "cumulative", false, "cartservice", "cartservice", "signoz_calls_total", uint64(1747947360000), uint64(1747983420000), 0},
-			cte:      "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, max(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND (match(JSONExtractString(labels, 'materialized.key.name'), ?) OR JSONExtractString(labels, 'service.name') = ?) GROUP BY fingerprint, `service.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`)",
+			expected: qbtypes.Statement{
+				Query: "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, max(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND (match(JSONExtractString(labels, 'materialized.key.name'), ?) OR JSONExtractString(labels, 'service.name') = ?) GROUP BY fingerprint, `service.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `service.name` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name`, ts",
+				Args:  []any{"signoz_calls_total", uint64(1747936800000), uint64(1747983420000), "cumulative", false, "cartservice", "cartservice", "signoz_calls_total", uint64(1747947360000), uint64(1747983420000), 0},
+			},
+			expectedErr: nil,
 		},
 		{
-			name: "delta_rate_sum",
+			name:        "test_delta_rate_sum",
+			requestType: qbtypes.RequestTypeTimeSeries,
 			query: qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]{
 				Signal:       telemetrytypes.SignalMetrics,
 				StepInterval: qbtypes.Step{Duration: 30 * time.Second},
@@ -91,16 +106,24 @@ func TestStatementBuilder(t *testing.T) {
 				Filter: &qbtypes.Filter{
 					Expression: "service.name = 'cartservice'",
 				},
+				Limit: 10,
 				GroupBy: []qbtypes.GroupByKey{
-					{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "service.name"}},
+					{
+						TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
+							Name: "service.name",
+						},
+					},
 				},
 			},
-			orderKey: "service.name",
-			args:     []any{"signoz_calls_total", uint64(1747936800000), uint64(1747983420000), "delta", false, "cartservice", "signoz_calls_total", uint64(1747947390000), uint64(1747983420000)},
-			cte:      "WITH __spatial_aggregation_cte AS (SELECT toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, sum(value)/30 AS value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'service.name') = ? GROUP BY fingerprint, `service.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY ts, `service.name`)",
+			expected: qbtypes.Statement{
+				Query: "WITH __spatial_aggregation_cte AS (SELECT toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, sum(value)/30 AS value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'service.name') = ? GROUP BY fingerprint, `service.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY ts, `service.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name`, ts",
+				Args:  []any{"signoz_calls_total", uint64(1747936800000), uint64(1747983420000), "delta", false, "cartservice", "signoz_calls_total", uint64(1747947390000), uint64(1747983420000)},
+			},
+			expectedErr: nil,
 		},
 		{
-			name: "histogram_percentile1",
+			name:        "test_histogram_percentile1",
+			requestType: qbtypes.RequestTypeTimeSeries,
 			query: qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]{
 				Signal:       telemetrytypes.SignalMetrics,
 				StepInterval: qbtypes.Step{Duration: 30 * time.Second},
@@ -116,16 +139,24 @@ func TestStatementBuilder(t *testing.T) {
 				Filter: &qbtypes.Filter{
 					Expression: "service.name = 'cartservice'",
 				},
+				Limit: 10,
 				GroupBy: []qbtypes.GroupByKey{
-					{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "service.name"}},
+					{
+						TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
+							Name: "service.name",
+						},
+					},
 				},
 			},
-			orderKey: "service.name",
-			args:     []any{"signoz_latency", uint64(1747936800000), uint64(1747983420000), "delta", false, "cartservice", "signoz_latency", uint64(1747947390000), uint64(1747983420000)},
-			cte:      "WITH __spatial_aggregation_cte AS (SELECT toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, `le`, sum(value)/30 AS value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name`, JSONExtractString(labels, 'le') AS `le` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'service.name') = ? GROUP BY fingerprint, `service.name`, `le`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY ts, `service.name`, `le`)",
+			expected: qbtypes.Statement{
+				Query: "WITH __spatial_aggregation_cte AS (SELECT toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, `le`, sum(value)/30 AS value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name`, JSONExtractString(labels, 'le') AS `le` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'service.name') = ? GROUP BY fingerprint, `service.name`, `le`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY ts, `service.name`, `le`) SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte GROUP BY `service.name`, ts ORDER BY `service.name`, ts",
+				Args:  []any{"signoz_latency", uint64(1747936800000), uint64(1747983420000), "delta", false, "cartservice", "signoz_latency", uint64(1747947390000), uint64(1747983420000)},
+			},
+			expectedErr: nil,
 		},
 		{
-			name: "gauge_avg_sum",
+			name:        "test_gauge_avg_sum",
+			requestType: qbtypes.RequestTypeTimeSeries,
 			query: qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]{
 				Signal:       telemetrytypes.SignalMetrics,
 				StepInterval: qbtypes.Step{Duration: 30 * time.Second},
@@ -141,16 +172,24 @@ func TestStatementBuilder(t *testing.T) {
 				Filter: &qbtypes.Filter{
 					Expression: "host.name = 'big-data-node-1'",
 				},
+				Limit: 10,
 				GroupBy: []qbtypes.GroupByKey{
-					{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "host.name"}},
+					{
+						TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
+							Name: "host.name",
+						},
+					},
 				},
 			},
-			orderKey: "host.name",
-			args:     []any{"system.memory.usage", uint64(1747936800000), uint64(1747983420000), "unspecified", false, "big-data-node-1", "system.memory.usage", uint64(1747947390000), uint64(1747983420000), 0},
-			cte:      "WITH __temporal_aggregation_cte AS (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `host.name`, avg(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'host.name') AS `host.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'host.name') = ? GROUP BY fingerprint, `host.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `host.name` ORDER BY fingerprint, ts), __spatial_aggregation_cte AS (SELECT ts, `host.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `host.name`)",
+			expected: qbtypes.Statement{
+				Query: "WITH __temporal_aggregation_cte AS (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `host.name`, avg(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'host.name') AS `host.name` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? AND JSONExtractString(labels, 'host.name') = ? GROUP BY fingerprint, `host.name`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `host.name` ORDER BY fingerprint, ts), __spatial_aggregation_cte AS (SELECT ts, `host.name`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `host.name`) SELECT * FROM __spatial_aggregation_cte ORDER BY `host.name`, ts",
+				Args:  []any{"system.memory.usage", uint64(1747936800000), uint64(1747983420000), "unspecified", false, "big-data-node-1", "system.memory.usage", uint64(1747947390000), uint64(1747983420000), 0},
+			},
+			expectedErr: nil,
 		},
 		{
-			name: "histogram_percentile2",
+			name:        "test_histogram_percentile2",
+			requestType: qbtypes.RequestTypeTimeSeries,
 			query: qbtypes.QueryBuilderQuery[qbtypes.MetricAggregation]{
 				Signal:       telemetrytypes.SignalMetrics,
 				StepInterval: qbtypes.Step{Duration: 30 * time.Second},
@@ -163,69 +202,23 @@ func TestStatementBuilder(t *testing.T) {
 						SpaceAggregation: metrictypes.SpaceAggregationPercentile95,
 					},
 				},
+				Limit: 10,
 				GroupBy: []qbtypes.GroupByKey{
-					{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: "service.name"}},
+					{
+						TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{
+							Name: "service.name",
+						},
+					},
 				},
 			},
-			orderKey: "service.name",
-			args:     []any{"http_server_duration_bucket", uint64(1747936800000), uint64(1747983420000), "cumulative", false, "http_server_duration_bucket", uint64(1747947360000), uint64(1747983420000), 0},
-			cte:      "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, `le`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, `le`, max(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name`, JSONExtractString(labels, 'le') AS `le` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? GROUP BY fingerprint, `service.name`, `le`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `service.name`, `le` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, `le`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`, `le`)",
+			expected: qbtypes.Statement{
+				Query: "WITH __temporal_aggregation_cte AS (SELECT ts, `service.name`, `le`, multiIf(row_number() OVER rate_window = 1, nan, (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) < 0, per_series_value / (ts - lagInFrame(ts, 1) OVER rate_window), (per_series_value - lagInFrame(per_series_value, 1) OVER rate_window) / (ts - lagInFrame(ts, 1) OVER rate_window)) AS per_series_value FROM (SELECT fingerprint, toStartOfInterval(toDateTime(intDiv(unix_milli, 1000)), toIntervalSecond(30)) AS ts, `service.name`, `le`, max(value) AS per_series_value FROM signoz_metrics.distributed_samples_v4 AS points INNER JOIN (SELECT fingerprint, JSONExtractString(labels, 'service.name') AS `service.name`, JSONExtractString(labels, 'le') AS `le` FROM signoz_metrics.time_series_v4_6hrs WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli <= ? AND LOWER(temporality) LIKE LOWER(?) AND __normalized = ? GROUP BY fingerprint, `service.name`, `le`) AS filtered_time_series ON points.fingerprint = filtered_time_series.fingerprint WHERE metric_name IN (?) AND unix_milli >= ? AND unix_milli < ? GROUP BY fingerprint, ts, `service.name`, `le` ORDER BY fingerprint, ts) WINDOW rate_window AS (PARTITION BY fingerprint ORDER BY fingerprint, ts)), __spatial_aggregation_cte AS (SELECT ts, `service.name`, `le`, sum(per_series_value) AS value FROM __temporal_aggregation_cte WHERE isNaN(per_series_value) = ? GROUP BY ts, `service.name`, `le`) SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte GROUP BY `service.name`, ts ORDER BY `service.name`, ts",
+				Args:  []any{"http_server_duration_bucket", uint64(1747936800000), uint64(1747983420000), "cumulative", false, "http_server_duration_bucket", uint64(1747947360000), uint64(1747983420000), 0},
+			},
+			expectedErr: nil,
 		},
 	}
 
-	type variant struct {
-		name     string
-		limit    int
-		hasOrder bool
-	}
-
-	variants := []variant{
-		{"with_limits", 10, false},
-		{"without_limits", 0, false},
-		{"with_order_by", 0, true},
-		{"with_order_by_and_limits", 10, true},
-	}
-
-	// expectedFinalSelects maps "base/variant" to the final SELECT portion after the CTE.
-	// The full expected query is: base.cte + expectedFinalSelects[name]
-	expectedFinalSelects := map[string]string{
-		// cumulative_rate_sum
-		"cumulative_rate_sum/with_limits":              " SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) ORDER BY `service.name`, ts ASC",
-		"cumulative_rate_sum/without_limits":           " SELECT * FROM __spatial_aggregation_cte ORDER BY avg(value) OVER (PARTITION BY `service.name`) DESC, `service.name`, ts ASC",
-		"cumulative_rate_sum/with_order_by":            " SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name` asc, ts ASC",
-		"cumulative_rate_sum/with_order_by_and_limits": " SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT DISTINCT `service.name` FROM __spatial_aggregation_cte ORDER BY `service.name` asc LIMIT 10) ORDER BY `service.name` asc, ts ASC",
-
-		// cumulative_rate_sum_with_mat_column
-		"cumulative_rate_sum_with_mat_column/with_limits":              " SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) ORDER BY `service.name`, ts ASC",
-		"cumulative_rate_sum_with_mat_column/without_limits":           " SELECT * FROM __spatial_aggregation_cte ORDER BY avg(value) OVER (PARTITION BY `service.name`) DESC, `service.name`, ts ASC",
-		"cumulative_rate_sum_with_mat_column/with_order_by":            " SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name` asc, ts ASC",
-		"cumulative_rate_sum_with_mat_column/with_order_by_and_limits": " SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT DISTINCT `service.name` FROM __spatial_aggregation_cte ORDER BY `service.name` asc LIMIT 10) ORDER BY `service.name` asc, ts ASC",
-
-		// delta_rate_sum
-		"delta_rate_sum/with_limits":              " SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) ORDER BY `service.name`, ts ASC",
-		"delta_rate_sum/without_limits":           " SELECT * FROM __spatial_aggregation_cte ORDER BY avg(value) OVER (PARTITION BY `service.name`) DESC, `service.name`, ts ASC",
-		"delta_rate_sum/with_order_by":            " SELECT * FROM __spatial_aggregation_cte ORDER BY `service.name` asc, ts ASC",
-		"delta_rate_sum/with_order_by_and_limits": " SELECT * FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT DISTINCT `service.name` FROM __spatial_aggregation_cte ORDER BY `service.name` asc LIMIT 10) ORDER BY `service.name` asc, ts ASC",
-
-		// histogram_percentile1
-		"histogram_percentile1/with_limits":              " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) GROUP BY `service.name`, ts ORDER BY `service.name`, ts ASC",
-		"histogram_percentile1/without_limits":           " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte GROUP BY `service.name`, ts ORDER BY avg(value) OVER (PARTITION BY `service.name`) DESC, `service.name`, ts ASC",
-		"histogram_percentile1/with_order_by":            " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte GROUP BY `service.name`, ts ORDER BY `service.name` asc, ts ASC",
-		"histogram_percentile1/with_order_by_and_limits": " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT DISTINCT `service.name` FROM __spatial_aggregation_cte ORDER BY `service.name` asc LIMIT 10) GROUP BY `service.name`, ts ORDER BY `service.name` asc, ts ASC",
-
-		// gauge_avg_sum
-		"gauge_avg_sum/with_limits":              " SELECT * FROM __spatial_aggregation_cte WHERE (`host.name`) IN (SELECT `host.name` FROM __spatial_aggregation_cte GROUP BY `host.name` ORDER BY avg(value) LIMIT 10) ORDER BY `host.name`, ts ASC",
-		"gauge_avg_sum/without_limits":           " SELECT * FROM __spatial_aggregation_cte ORDER BY avg(value) OVER (PARTITION BY `host.name`) DESC, `host.name`, ts ASC",
-		"gauge_avg_sum/with_order_by":            " SELECT * FROM __spatial_aggregation_cte ORDER BY `host.name` asc, ts ASC",
-		"gauge_avg_sum/with_order_by_and_limits": " SELECT * FROM __spatial_aggregation_cte WHERE (`host.name`) IN (SELECT DISTINCT `host.name` FROM __spatial_aggregation_cte ORDER BY `host.name` asc LIMIT 10) ORDER BY `host.name` asc, ts ASC",
-
-		// histogram_percentile2
-		"histogram_percentile2/with_limits":              " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT `service.name` FROM __spatial_aggregation_cte GROUP BY `service.name` ORDER BY avg(value) LIMIT 10) GROUP BY `service.name`, ts ORDER BY `service.name`, ts ASC",
-		"histogram_percentile2/without_limits":           " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte GROUP BY `service.name`, ts ORDER BY avg(value) OVER (PARTITION BY `service.name`) DESC, `service.name`, ts ASC",
-		"histogram_percentile2/with_order_by":            " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte GROUP BY `service.name`, ts ORDER BY `service.name` asc, ts ASC",
-		"histogram_percentile2/with_order_by_and_limits": " SELECT ts, `service.name`, histogramQuantile(arrayMap(x -> toFloat64(x), groupArray(le)), groupArray(value), 0.950) AS value FROM __spatial_aggregation_cte WHERE (`service.name`) IN (SELECT DISTINCT `service.name` FROM __spatial_aggregation_cte ORDER BY `service.name` asc LIMIT 10) GROUP BY `service.name`, ts ORDER BY `service.name` asc, ts ASC",
-	}
-
 	fm := NewFieldMapper()
 	cb := NewConditionBuilder(fm)
 	mockMetadataStore := telemetrytypestest.NewMockMetadataStore()
@@ -234,13 +227,15 @@ func TestStatementBuilder(t *testing.T) {
 		t.Fatalf("failed to load field keys: %v", err)
 	}
 	mockMetadataStore.KeysMap = keys
+	// NOTE: LoadFieldKeysFromJSON doesn't set Materialized field
+	// for keys, so we have to set it manually here for testing
 	if _, ok := mockMetadataStore.KeysMap["materialized.key.name"]; ok {
 		if len(mockMetadataStore.KeysMap["materialized.key.name"]) > 0 {
 			mockMetadataStore.KeysMap["materialized.key.name"][0].Materialized = true
 		}
 	}
 
-	fl, err := flagger.New(context.Background(), instrumentationtest.New().ToProviderSettings(), flagger.Config{}, flagger.MustNewRegistry())
+	flagger, err := flagger.New(context.Background(), instrumentationtest.New().ToProviderSettings(), flagger.Config{}, flagger.MustNewRegistry())
 	if err != nil {
 		t.Fatalf("failed to create flagger: %v", err)
 	}
@@ -250,30 +245,23 @@ func TestStatementBuilder(t *testing.T) {
 		mockMetadataStore,
 		fm,
 		cb,
-		fl,
+		flagger,
 	)
 
-	for _, b := range bases {
-		for _, v := range variants {
-			name := b.name + "/" + v.name
-			t.Run(name, func(t *testing.T) {
-				q := b.query
-				q.Limit = v.limit
-				if v.hasOrder {
-					q.Order = []qbtypes.OrderBy{
-						{
-							Key:       qbtypes.OrderByKey{TelemetryFieldKey: telemetrytypes.TelemetryFieldKey{Name: b.orderKey}},
-							Direction: qbtypes.OrderDirectionAsc,
-						},
-					}
-				}
+	for _, c := range cases {
+		t.Run(c.name, func(t *testing.T) {
 
-				result, err := statementBuilder.Build(context.Background(), 1747947419000, 1747983448000, qbtypes.RequestTypeTimeSeries, q, nil)
+			q, err := statementBuilder.Build(context.Background(), 1747947419000, 1747983448000, c.requestType, c.query, nil)
 
+			if c.expectedErr != nil {
+				require.Error(t, err)
+				require.Contains(t, err.Error(), c.expectedErr.Error())
+			} else {
 				require.NoError(t, err)
-				require.Equal(t, b.cte+expectedFinalSelects[name], result.Query)
-				require.Equal(t, b.args, result.Args)
-			})
-		}
+				require.Equal(t, c.expected.Query, q.Query)
+				require.Equal(t, c.expected.Args, q.Args)
+				require.Equal(t, c.expected.Warnings, q.Warnings)
+			}
+		})
 	}
 }