feat: support aggregation function with values

frontend/src/components/QueryBuilderV2/QueryAggregation/QueryAggregation.styles.scss

@@ -0,0 +1,156 @@
+.query-aggregation-container {
+	display: flex;
+	flex-direction: row;
+	align-items: center;
+	gap: 8px;
+
+	.query-aggregation-options-input {
+		width: 100%;
+		height: 36px;
+		line-height: 36px;
+		border-radius: 2px;
+		border: 1px solid var(--bg-slate-400);
+		box-shadow: 0px 0px 8px 0px rgba(0, 0, 0, 0.1);
+
+		font-family: 'Space Mono', monospace !important;
+
+		&::placeholder {
+			color: var(--bg-vanilla-100);
+			opacity: 0.5;
+		}
+	}
+
+	.query-aggregation-interval {
+		display: flex;
+		flex-direction: row;
+		align-items: center;
+		gap: 8px;
+	}
+}
+
+.query-aggregation-select-container {
+	width: 100%;
+
+	.query-aggregation-select-editor {
+		border-radius: 2px;
+
+		.cm-content {
+			padding: 0;
+		}
+
+		.cm-editor {
+			border-radius: 2px;
+			overflow: hidden;
+			background-color: transparent !important;
+
+			&:focus-within {
+				border-color: var(--bg-robin-500);
+			}
+
+			.cm-content {
+				border-radius: 2px;
+				border: 1px solid var(--Slate-400, #1d212d);
+				padding: 0px !important;
+				background-color: #121317 !important;
+
+				&:focus-within {
+					border-color: var(--bg-ink-200);
+				}
+			}
+
+			.cm-tooltip-autocomplete {
+				background: var(--bg-ink-300) !important;
+				border-radius: 2px !important;
+				font-size: 12px !important;
+				font-weight: 500 !important;
+				margin-top: -2px !important;
+				min-width: 400px !important;
+				position: relative !important;
+				top: 0px !important;
+				left: 0px !important;
+
+				border-radius: 4px;
+				border: 1px solid var(--bg-slate-200, #1d212d);
+				background: linear-gradient(
+					139deg,
+					rgba(18, 19, 23, 0.8) 0%,
+					rgba(18, 19, 23, 0.9) 98.68%
+				) !important;
+				backdrop-filter: blur(20px);
+				box-sizing: border-box;
+				font-family: 'Space Mono', monospace !important;
+
+				ul {
+					width: 100% !important;
+					max-width: 100% !important;
+					font-family: 'Space Mono', monospace !important;
+					min-height: 200px !important;
+
+					&::-webkit-scrollbar {
+						width: 0.3rem;
+					}
+					&::-webkit-scrollbar-corner {
+						background: transparent;
+					}
+					&::-webkit-scrollbar-thumb {
+						background: rgb(136, 136, 136);
+						border-radius: 0.625rem;
+					}
+					&::-webkit-scrollbar-track {
+						background: transparent;
+					}
+
+					li {
+						width: 100% !important;
+						max-width: 100% !important;
+						line-height: 36px !important;
+						height: 36px !important;
+						padding: 4px 8px !important;
+
+						display: flex !important;
+						align-items: center !important;
+						gap: 8px !important;
+						box-sizing: border-box;
+						overflow: hidden;
+
+						font-family: 'Space Mono', monospace !important;
+
+						.cm-completionIcon {
+							display: none !important;
+						}
+
+						&[aria-selected='true'] {
+							// background-color: rgba(78, 116, 248, 0.7) !important;
+							background: rgba(171, 189, 255, 0.04) !important;
+						}
+					}
+				}
+			}
+
+			.cm-gutters {
+				display: none !important;
+			}
+
+			.cm-line {
+				line-height: 34px !important;
+				font-family: 'Space Mono', monospace !important;
+				background-color: #121317 !important;
+
+				::-moz-selection {
+					background: var(--bg-ink-100) !important;
+					opacity: 0.5 !important;
+				}
+
+				::selection {
+					background: var(--bg-ink-100) !important;
+					opacity: 0.5 !important;
+				}
+			}
+
+			.cm-selectionBackground {
+				background: var(--bg-ink-100) !important;
+				opacity: 0.5 !important;
+			}
+		}
+	}
+}

frontend/src/components/QueryBuilderV2/QueryAggregation/QueryAggregation.tsx

@@ -1,15 +1,19 @@
-import './QueryAggregationOptions.styles.scss';
+import './QueryAggregation.styles.scss';
 
-import { Input } from 'antd';
+// import { Input } from 'antd';
 import InputWithLabel from 'components/InputWithLabel/InputWithLabel';
 
+import QueryAggregationSelect from './QueryAggregationSelect';
+
 function QueryAggregationOptions(): JSX.Element {
 	return (
 		<div className="query-aggregation-container">
-			<Input
+			{/* <Input
 				placeholder="Search aggregation options..."
 				className="query-aggregation-options-input"
-			/>
+			/> */}
+
+			<QueryAggregationSelect />
 
 			<div className="query-aggregation-interval">
 				<div className="query-aggregation-interval-label">every</div>

frontend/src/components/QueryBuilderV2/QueryAggregation/QueryAggregationSelect.tsx

@@ -0,0 +1,210 @@
+/* eslint-disable sonarjs/cognitive-complexity */
+import './QueryAggregation.styles.scss';
+
+import {
+	autocompletion,
+	Completion,
+	CompletionContext,
+} from '@codemirror/autocomplete';
+import { javascript } from '@codemirror/lang-javascript';
+import { copilot } from '@uiw/codemirror-theme-copilot';
+import CodeMirror, { EditorView } from '@uiw/react-codemirror';
+import { tracesAggregateOperatorOptions } from 'constants/queryBuilderOperators';
+import { TracesAggregatorOperator } from 'types/common/queryBuilder';
+
+const operatorArgMeta: Record<
+	string,
+	{ acceptsArgs: boolean; multiple: boolean }
+> = {
+	[TracesAggregatorOperator.NOOP]: { acceptsArgs: false, multiple: false },
+	[TracesAggregatorOperator.COUNT]: { acceptsArgs: false, multiple: false },
+	[TracesAggregatorOperator.COUNT_DISTINCT]: {
+		acceptsArgs: true,
+		multiple: true,
+	},
+	[TracesAggregatorOperator.SUM]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.AVG]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.MAX]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.MIN]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P05]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P10]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P20]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P25]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P50]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P75]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P90]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P95]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.P99]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.RATE]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.RATE_SUM]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.RATE_AVG]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.RATE_MIN]: { acceptsArgs: true, multiple: false },
+	[TracesAggregatorOperator.RATE_MAX]: { acceptsArgs: true, multiple: false },
+};
+
+const fieldSuggestions: Completion[] = [
+	{ label: 'duration', type: 'variable' },
+	{ label: 'status_code', type: 'variable' },
+	{ label: 'service_name', type: 'variable' },
+	{ label: 'trace_id', type: 'variable' },
+];
+
+const mapToFunctionCompletions = (
+	operators: typeof tracesAggregateOperatorOptions,
+): Completion[] =>
+	operators.map((op) => ({
+		label: `${op.value}()`,
+		type: 'function',
+		apply: `${op.value}()`,
+	}));
+
+const applyFieldSuggestion = (
+	view: EditorView,
+	suggestion: Completion,
+): void => {
+	const currentText = view.state.sliceDoc(0, view.state.selection.main.from);
+	const endPos = view.state.selection.main.from;
+
+	// Find the last opening parenthesis before the cursor
+	const lastOpenParen = currentText.lastIndexOf('(');
+	if (lastOpenParen === -1) return;
+
+	// Find the last comma after the opening parenthesis
+	const textAfterParen = currentText.slice(lastOpenParen);
+	const lastComma = textAfterParen.lastIndexOf(',');
+
+	// Calculate the start position for insertion
+	const startPos =
+		lastComma === -1 ? lastOpenParen + 1 : lastOpenParen + lastComma + 1;
+
+	// Insert the suggestion
+	view.dispatch({
+		changes: { from: startPos, to: endPos, insert: suggestion.label },
+		selection: { anchor: startPos + suggestion.label.length },
+	});
+};
+
+const applyOperatorSuggestion = (
+	view: EditorView,
+	from: number,
+	label: string,
+): void => {
+	view.dispatch({
+		changes: { from, insert: label },
+		selection: { anchor: from + label.length },
+	});
+};
+
+const getOperatorSuggestions = (
+	from: number,
+	operators: typeof tracesAggregateOperatorOptions,
+): Completion[] =>
+	mapToFunctionCompletions(operators).map((op) => ({
+		...op,
+		apply: (view: EditorView): void =>
+			applyOperatorSuggestion(view, from, op.label),
+	}));
+
+const aggregatorAutocomplete = autocompletion({
+	override: [
+		(context: CompletionContext): any => {
+			const word = context.matchBefore(/[\w\d_\s]*(\()?[^)]*$/);
+			if (!word || (word.from === word.to && !context.explicit)) return null;
+
+			const textBeforeCursor = context.state.sliceDoc(0, context.pos);
+			const functionMatch = textBeforeCursor.match(/(\w+)\(([^)]*)$/);
+			const funcName = functionMatch?.[1]?.toLowerCase();
+
+			// Handle argument suggestions when cursor is inside parentheses
+			if (funcName && operatorArgMeta[funcName]) {
+				const { acceptsArgs, multiple } = operatorArgMeta[funcName];
+
+				if (!acceptsArgs) return null;
+
+				// Get all arguments for the current function
+				const argsMatch = functionMatch?.[2];
+				const argsSoFar =
+					argsMatch
+						?.split(',')
+						.map((arg) => arg.trim())
+						.filter(Boolean) || [];
+
+				if (!multiple && argsSoFar.length >= 1) return null;
+
+				return {
+					from: context.pos,
+					options: fieldSuggestions.map((suggestion) => ({
+						...suggestion,
+						apply: (view: EditorView): void => {
+							applyFieldSuggestion(view, suggestion);
+							// For count_distinct, add a comma after the field
+							if (funcName === TracesAggregatorOperator.COUNT_DISTINCT.toLowerCase()) {
+								const currentPos = view.state.selection.main.from;
+								view.dispatch({
+									changes: { from: currentPos, insert: ', ' },
+									selection: { anchor: currentPos + 2 },
+								});
+							}
+						},
+					})),
+				};
+			}
+
+			// Handle operator suggestions
+			const isAfterCompleteFunction = textBeforeCursor.match(/\w+\([^)]*\)\s*$/);
+			if (isAfterCompleteFunction) {
+				return {
+					from: context.pos,
+					options: getOperatorSuggestions(
+						context.pos,
+						tracesAggregateOperatorOptions,
+					),
+				};
+			}
+
+			// Regular word-based suggestions
+			const wordBeforeCursor = word.text.trim();
+			if (wordBeforeCursor) {
+				const filteredOperators = tracesAggregateOperatorOptions.filter((op) =>
+					op.value.toLowerCase().startsWith(wordBeforeCursor.toLowerCase()),
+				);
+				return {
+					from: word.from,
+					options: getOperatorSuggestions(word.from, filteredOperators),
+				};
+			}
+
+			// Show all options if no word before cursor
+			return {
+				from: word.from,
+				options: getOperatorSuggestions(word.from, tracesAggregateOperatorOptions),
+			};
+		},
+	],
+});
+
+function QueryAggregationSelect(): JSX.Element {
+	return (
+		<div className="query-aggregation-select-container">
+			<CodeMirror
+				className="query-aggregation-select-editor"
+				width="100%"
+				theme={copilot}
+				extensions={[
+					aggregatorAutocomplete,
+					javascript({ jsx: false, typescript: true }),
+				]}
+				placeholder="Type aggregator functions like sum(), count_distinct(...), etc."
+				basicSetup={{
+					lineNumbers: false,
+					closeBrackets: true,
+					autocompletion: true,
+					completionKeymap: true,
+				}}
+				lang="sql"
+			/>
+		</div>
+	);
+}
+
+export default QueryAggregationSelect;

frontend/src/components/QueryBuilderV2/QueryAggregationOptions/QueryAggregationOptions.styles.scss

@@ -1,29 +0,0 @@
-.query-aggregation-container {
-	display: flex;
-	flex-direction: row;
-	align-items: center;
-	gap: 8px;
-
-	.query-aggregation-options-input {
-		width: 100%;
-		height: 36px;
-		line-height: 36px;
-		border-radius: 2px;
-		border: 1px solid var(--bg-slate-400);
-		box-shadow: 0px 0px 8px 0px rgba(0, 0, 0, 0.1);
-
-		font-family: 'Space Mono', monospace !important;
-
-		&::placeholder {
-			color: var(--bg-vanilla-100);
-			opacity: 0.5;
-		}
-	}
-
-	.query-aggregation-interval {
-		display: flex;
-		flex-direction: row;
-		align-items: center;
-		gap: 8px;
-	}
-}

frontend/src/components/QueryBuilderV2/QueryBuilderV2.tsx

@@ -3,7 +3,7 @@ import './QueryBuilderV2.styles.scss';
 import { useQueryBuilder } from 'hooks/queryBuilder/useQueryBuilder';
 
 import QueryAddOns from './QueryAddOns/QueryAddOns';
-import QueryAggregationOptions from './QueryAggregationOptions/QueryAggregationOptions';
+import QueryAggregation from './QueryAggregation/QueryAggregation';
 import QuerySearch from './QuerySearch/QuerySearch';
 
 function QueryBuilderV2(): JSX.Element {
@@ -12,7 +12,7 @@ function QueryBuilderV2(): JSX.Element {
 	return (
 		<div className="query-builder-v2">
 			<QuerySearch />
-			<QueryAggregationOptions />
+			<QueryAggregation />
 			<QueryAddOns
 				query={currentQuery.builder.queryData[0]}
 				version="v3"

frontend/src/components/QueryBuilderV2/QuerySearch/QuerySearch.tsx

@@ -21,6 +21,7 @@ import CodeMirror, { EditorView, Extension } from '@uiw/react-codemirror';
 import { Card, Collapse, Space, Tag, Typography } from 'antd';
 import { getValueSuggestions } from 'api/querySuggestions/getValueSuggestion';
 import { useGetQueryKeySuggestions } from 'hooks/querySuggestions/useGetQueryKeySuggestions';
+import cloneDeep from 'lodash-es/cloneDeep';
 import { useCallback, useEffect, useRef, useState } from 'react';
 import {
 	IDetailedError,
@@ -29,6 +30,7 @@ import {
 } from 'types/antlrQueryTypes';
 import { QueryKeySuggestionsProps } from 'types/api/querySuggestions/types';
 import { queryOperatorSuggestions, validateQuery } from 'utils/antlrQueryUtils';
+import { detectContext } from 'utils/antlrQueryUtils2';
 import { getQueryContextAtCursor } from 'utils/queryContextUtils';
 
 const { Text } = Typography;
@@ -647,7 +649,11 @@ function QuerySearch(): JSX.Element {
 		// Get the query context at the cursor position
 		const queryContext = getQueryContextAtCursor(query, cursorPos.ch);
 
+		// Get the query context at the cursor position
+		const queryContext2 = detectContext(query, cursorPos.ch);
+
 		console.log('queryContext', queryContext);
+		console.log('queryContext2', cloneDeep(queryContext2));
 
 		// Define autocomplete options based on the context
 		let options: {
@@ -1037,7 +1043,7 @@ function QuerySearch(): JSX.Element {
 						override: [myCompletions],
 						defaultKeymap: true,
 						closeOnBlur: false,
-						activateOnTyping: true,
+						// activateOnTyping: true,
 						maxRenderedOptions: 50,
 					}),
 					javascript({ jsx: false, typescript: false }),

frontend/src/utils/antlrQueryUtils2.ts

@@ -0,0 +1,159 @@
+/* eslint-disable sonarjs/cognitive-complexity */
+/* eslint-disable no-restricted-syntax */
+import antlr4, { CharStreams } from 'antlr4';
+import cloneDeep from 'lodash-es/cloneDeep';
+import FilterQueryLexer from 'parser/FilterQueryLexer';
+
+export enum CursorContext {
+	Key,
+	Operator,
+	Value,
+	NoFilter,
+	FullText,
+}
+
+const contextNames = ['Key', 'Operator', 'Value', 'NoFilter', 'FullText'];
+
+export function contextToString(context: CursorContext): string {
+	return contextNames[context];
+}
+
+export interface ContextInfo {
+	context: CursorContext;
+	key?: string;
+	token?: antlr4.Token;
+	operator?: string;
+}
+
+export function detectContext(
+	query: string,
+	cursorOffset: number,
+): ContextInfo {
+	console.log('query', query);
+	console.log('cursorOffset', cursorOffset);
+
+	const chars = CharStreams.fromString(query);
+	const lexer = new FilterQueryLexer(chars);
+	const tokens = new antlr4.CommonTokenStream(lexer);
+	tokens.fill();
+
+	enum State {
+		ExpectKey,
+		ExpectOperator,
+		ExpectValue,
+	}
+
+	let state = State.ExpectKey;
+	let parens = 0;
+	let array = 0;
+	let lastKey: antlr4.Token | undefined;
+	let lastOperator: antlr4.Token | undefined;
+	let cursorTok: antlr4.Token | undefined;
+	let pos = 0;
+
+	for (const tok of tokens.tokens) {
+		const text = tok.text || '';
+
+		if (
+			tok.channel === antlr4.Token.DEFAULT_CHANNEL &&
+			pos <= cursorOffset &&
+			cursorOffset <= pos + text.length
+		) {
+			cursorTok = tok;
+			break;
+		}
+
+		switch (tok.type) {
+			case FilterQueryLexer.LPAREN:
+				parens++;
+				state = State.ExpectKey;
+				break;
+			case FilterQueryLexer.RPAREN:
+				if (parens > 0) parens--;
+				state = State.ExpectOperator;
+				break;
+			case FilterQueryLexer.LBRACK:
+				array++;
+				state = State.ExpectValue;
+				break;
+			case FilterQueryLexer.RBRACK:
+				if (array > 0) array--;
+				state = State.ExpectOperator;
+				break;
+			case FilterQueryLexer.COMMA:
+				if (array > 0) state = State.ExpectValue;
+				break;
+			case FilterQueryLexer.KEY:
+				if (state === State.ExpectKey) {
+					lastKey = tok;
+					state = State.ExpectOperator;
+				}
+				break;
+			case FilterQueryLexer.QUOTED_TEXT:
+			case FilterQueryLexer.NUMBER:
+			case FilterQueryLexer.BOOL:
+				if (state === State.ExpectValue) {
+					state = State.ExpectOperator;
+				}
+				break;
+			default:
+				if (
+					tok.type >= FilterQueryLexer.EQUALS &&
+					tok.type <= FilterQueryLexer.CONTAINS
+				) {
+					state = State.ExpectValue;
+				}
+				break;
+		}
+
+		pos += text.length;
+	}
+
+	console.log('cursorTok', cursorTok);
+
+	const out: ContextInfo = { context: CursorContext.NoFilter };
+
+	if (cursorTok) {
+		out.token = cursorTok;
+	}
+
+	console.log('out', cloneDeep(out));
+	console.log('state', cloneDeep(state));
+
+	switch (state) {
+		case State.ExpectKey:
+			out.context = CursorContext.Key;
+			break;
+		case State.ExpectOperator:
+			out.context = CursorContext.Operator;
+			if (lastKey) out.key = lastKey.text;
+			break;
+		case State.ExpectValue:
+			out.context = CursorContext.Value;
+			if (lastKey) out.key = lastKey.text;
+
+			if (lastOperator) {
+				out.operator = lastOperator.text;
+			}
+			break;
+		default:
+			out.context = CursorContext.NoFilter;
+			break;
+	}
+
+	console.log('out', cloneDeep(out));
+
+	if (
+		cursorTok &&
+		cursorTok.type === FilterQueryLexer.QUOTED_TEXT &&
+		(out.context === CursorContext.Key || out.context === CursorContext.NoFilter)
+	) {
+		out.context = CursorContext.FullText;
+	}
+
+	// if (!cursorTok || cursorTok.type === antlr4.Token.EOF) {
+	// 	out.context = CursorContext.NoFilter;
+	// }
+
+	return out;
+}