test(querybuilder): cover aggregation arguments with hostile names

Assisted-by: Claude Opus 5.5

pkg/querybuilder/agg_rewrite.go

@@ -5,17 +5,22 @@ import (
 	"fmt"
 	"log/slog"
 	"strings"
+	"unicode"
 
 	chparser "github.com/AfterShip/clickhouse-sql-parser/parser"
 	"github.com/SigNoz/signoz/pkg/errors"
 	"github.com/SigNoz/signoz/pkg/factory"
 	"github.com/SigNoz/signoz/pkg/flagger"
+	grammar "github.com/SigNoz/signoz/pkg/parser/filterquery/grammar"
 	qbtypes "github.com/SigNoz/signoz/pkg/types/querybuildertypes/querybuildertypesv5"
 	"github.com/SigNoz/signoz/pkg/types/telemetrytypes"
 	"github.com/SigNoz/signoz/pkg/valuer"
+	"github.com/antlr4-go/antlr/v4"
 	"github.com/huandu/go-sqlbuilder"
 )
 
+const argumentPlaceholderPrefix = "__signoz_argument_"
+
 type aggExprRewriter struct {
 	logger         *slog.Logger
 	fullTextColumn *telemetrytypes.TelemetryFieldKey
@@ -57,12 +62,13 @@ func (r *aggExprRewriter) Rewrite(
 	keys map[string][]*telemetrytypes.TelemetryFieldKey,
 ) (string, []any, error) {
 
-	wrapped := fmt.Sprintf("SELECT %s", expr)
+	arguments := map[string]string{}
+	wrapped := fmt.Sprintf("SELECT %s", extractArguments(expr, arguments))
 	p := chparser.NewParser(wrapped)
 	stmts, err := p.ParseStmts()
 
 	if err != nil {
-		return "", nil, errors.WrapInternalf(err, errors.CodeInternal, "failed to parse aggregation expression %q", expr)
+		return "", nil, errors.WrapInvalidInputf(err, errors.CodeInvalidInput, "failed to parse aggregation expression %q", expr)
 	}
 
 	if len(stmts) == 0 {
@@ -85,6 +91,7 @@ func (r *aggExprRewriter) Rewrite(
 		fieldKeys:      keys,
 		fullTextColumn: r.fullTextColumn,
 		storage:        r.storage,
+		arguments:      arguments,
 	}
 	// Rewrite the first select item (our expression)
 	if err := sel.SelectItems[0].Accept(visitor); err != nil {
@@ -136,6 +143,7 @@ type exprVisitor struct {
 	fieldKeys      map[string][]*telemetrytypes.TelemetryFieldKey
 	fullTextColumn *telemetrytypes.TelemetryFieldKey
 	storage        qbtypes.Storage
+	arguments      map[string]string
 	Modified       bool
 	chArgs         []any
 	isRate         bool
@@ -155,10 +163,16 @@ func (v *exprVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
 		args = fn.Params.Items.Items
 	}
 
+	// count(*) is count()
+	if len(args) == 1 && aggFunc.MinArgs == 0 && v.arguments[chparser.Format(args[0])] == "*" {
+		fn.Params.Items.Items = nil
+		args = nil
+	}
+
 	// if we know aggregation function, we must ensure that the number of arguments is correct
 	if aggFunc.RequireArgs {
 		if len(args) < aggFunc.MinArgs || len(args) > aggFunc.MaxArgs {
-			return errors.NewInternalf(errors.CodeInternal, "invalid number of arguments for %q: %d", name, len(args))
+			return errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid number of arguments for %q: %d", name, len(args))
 		}
 	}
 	fn.Name.Name = aggFunc.FuncName
@@ -174,7 +188,10 @@ func (v *exprVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
 	// Handle *If functions with predicate + values
 	if aggFunc.FuncCombinator {
 		// Map the predicate (last argument)
-		origPred := chparser.Format(args[len(args)-1])
+		origPred, err := v.argumentText(name, args[len(args)-1])
+		if err != nil {
+			return err
+		}
 		whereClause, err := PrepareWhereClause(
 			origPred,
 			FilterExprVisitorOpts{
@@ -189,6 +206,10 @@ func (v *exprVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
 		if err != nil {
 			return err
 		}
+		// The search scan budget applies only to the filter.
+		if whereClause.RequiresCostGuard {
+			return errors.NewInvalidInputf(errors.CodeInvalidInput, "function `search` is not supported inside an aggregation").WithSuggestions("use `search` in the filter")
+		}
 		// not possible for whereClause to be empty here but still adding a check.
 		if whereClause.IsEmpty() {
 			return errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid predicate argument for %q: %q", name, origPred)
@@ -202,12 +223,18 @@ func (v *exprVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
 		}
 		args[len(args)-1] = parsedPred
 		v.Modified = true
-		v.chArgs = chArgs
+		v.chArgs = append(v.chArgs, chArgs...)
 
 		// Map each value column argument
 		for i := 0; i < len(args)-1; i++ {
-			origVal := chparser.Format(args[i])
-			fieldKey := telemetrytypes.GetFieldKeyFromKeyText(origVal)
+			origVal, err := v.argumentText(name, args[i])
+			if err != nil {
+				return err
+			}
+			fieldKey, err := fieldKeyFromArgument(name, origVal)
+			if err != nil {
+				return err
+			}
 			expr, err := ResolveColumn(v.ctx, v.query, v.storage, &fieldKey, dataType, v.fieldKeys)
 			if err != nil {
 				return errors.WrapInvalidInputf(err, errors.CodeInvalidInput, "failed to get table field name for %q", origVal)
@@ -223,8 +250,14 @@ func (v *exprVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
 	} else {
 		// Non-If functions: map every argument as a column/value
 		for i, arg := range args {
-			orig := chparser.Format(arg)
-			fieldKey := telemetrytypes.GetFieldKeyFromKeyText(orig)
+			orig, err := v.argumentText(name, arg)
+			if err != nil {
+				return err
+			}
+			fieldKey, err := fieldKeyFromArgument(name, orig)
+			if err != nil {
+				return err
+			}
 			expr, err := ResolveColumn(v.ctx, v.query, v.storage, &fieldKey, dataType, v.fieldKeys)
 			if err != nil {
 				return err
@@ -245,6 +278,14 @@ func (v *exprVisitor) VisitFunctionExpr(fn *chparser.FunctionExpr) error {
 	return nil
 }
 
+func (v *exprVisitor) argumentText(name string, arg chparser.Expr) (string, error) {
+	text, ok := v.arguments[chparser.Format(arg)]
+	if !ok {
+		return "", errors.NewInternalf(errors.CodeInternal, "argument of %q was not extracted before parsing", name)
+	}
+	return text, nil
+}
+
 // parseFragment parses a SQL expression fragment by wrapping in SELECT.
 func parseFragment(sql string) (chparser.Expr, error) {
 	wrapped := fmt.Sprintf("SELECT %s", sql)
@@ -262,3 +303,134 @@ func parseFragment(sql string) (chparser.Expr, error) {
 	}
 	return sel.SelectItems[0].Expr, nil
 }
+
+// extractArguments replaces every argument of an aggregation call with a
+// placeholder kept in arguments, because the ClickHouse parser cannot read
+// the predicates such as EXISTS or CONTAINS, and field keys
+// with hyphens or slashes. Unbalanced input is returned unchanged for the
+// parser to reject.
+func extractArguments(expr string, arguments map[string]string) string {
+	var out strings.Builder
+	written := 0
+	for i := 0; i < len(expr); i++ {
+		c := expr[i]
+		if c == '\'' || c == '"' || c == '`' {
+			end := skipQuoted(expr, i)
+			if end < 0 {
+				return expr
+			}
+			i = end
+			continue
+		}
+		if !isIdentifierStart(c) || (i > 0 && (isIdentifierPart(expr[i-1]) || expr[i-1] == '.')) {
+			continue
+		}
+		nameEnd := i
+		for nameEnd < len(expr) && isIdentifierPart(expr[nameEnd]) {
+			nameEnd++
+		}
+		open := nameEnd
+		for open < len(expr) && strings.ContainsRune(" \t\r\n", rune(expr[open])) {
+			open++
+		}
+		if _, ok := AggreFuncMap[valuer.NewString(strings.ToLower(expr[i:nameEnd]))]; !ok || open == len(expr) || expr[open] != '(' {
+			i = nameEnd - 1
+			continue
+		}
+		commas, closing := scanArguments(expr, open)
+		if closing < 0 {
+			return expr
+		}
+		out.WriteString(expr[written : open+1])
+		bounds := append(append([]int{open}, commas...), closing)
+		for j := 1; j < len(bounds); j++ {
+			if j > 1 {
+				out.WriteString(", ")
+			}
+			text := strings.TrimSpace(expr[bounds[j-1]+1 : bounds[j]])
+			// DISTINCT modifies the whole argument list, so it stays for the parser.
+			if end := strings.IndexFunc(text, unicode.IsSpace); j == 1 && end > 0 && strings.EqualFold(text[:end], "DISTINCT") {
+				out.WriteString("DISTINCT ")
+				text = strings.TrimSpace(text[end:])
+			}
+			if text == "" {
+				continue
+			}
+			placeholder := fmt.Sprintf("%s%d", argumentPlaceholderPrefix, len(arguments))
+			arguments[placeholder] = text
+			out.WriteString(placeholder)
+		}
+		written = closing
+		i = closing
+	}
+	out.WriteString(expr[written:])
+	return out.String()
+}
+
+// scanArguments returns the top-level commas and the closing parenthesis of
+// the argument list that opens at open, with -1 as the closing index when the
+// list is unterminated.
+func scanArguments(expr string, open int) ([]int, int) {
+	var commas []int
+	depth := 0
+	for i := open + 1; i < len(expr); i++ {
+		switch expr[i] {
+		case '\'', '"', '`':
+			end := skipQuoted(expr, i)
+			if end < 0 {
+				return nil, -1
+			}
+			i = end
+		case '(', '[':
+			depth++
+		case ']':
+			depth--
+		case ')':
+			if depth == 0 {
+				return commas, i
+			}
+			depth--
+		case ',':
+			if depth == 0 {
+				commas = append(commas, i)
+			}
+		}
+	}
+	return nil, -1
+}
+
+// fieldKeyFromArgument reads a value argument as a single key of the filter
+// grammar.
+func fieldKeyFromArgument(name string, text string) (telemetrytypes.TelemetryFieldKey, error) {
+	lexer := grammar.NewFilterQueryLexer(antlr.NewInputStream(text))
+	lexer.RemoveErrorListeners()
+	if lexer.NextToken().GetTokenType() != grammar.FilterQueryLexerKEY || lexer.NextToken().GetTokenType() != antlr.TokenEOF {
+		return telemetrytypes.TelemetryFieldKey{}, errors.NewInvalidInputf(errors.CodeInvalidInput, "invalid argument for %q: %q is not a field name", name, text)
+	}
+	return telemetrytypes.GetFieldKeyFromKeyText(text), nil
+}
+
+// skipQuoted returns the index of the quote closing the one at start, honoring
+// backslash escapes and doubled quotes, or -1 when it is unterminated.
+func skipQuoted(expr string, start int) int {
+	quote := expr[start]
+	for i := start + 1; i < len(expr); i++ {
+		switch {
+		case expr[i] == '\\':
+			i++
+		case expr[i] == quote && i+1 < len(expr) && expr[i+1] == quote:
+			i++
+		case expr[i] == quote:
+			return i
+		}
+	}
+	return -1
+}
+
+func isIdentifierStart(c byte) bool {
+	return c == '_' || (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z')
+}
+
+func isIdentifierPart(c byte) bool {
+	return isIdentifierStart(c) || (c >= '0' && c <= '9')
+}

tests/integration/tests/queriercommon/08_field_name_quoting.py

@@ -137,6 +137,46 @@ def test_logs_filter_expression_with_dollar_in_key(
         assert {row["data"]["body"] for row in get_rows(response)} == expected_bodies, f"{expression}: {response.text}"
 
 
+def test_logs_aggregation_with_dollar_in_key(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    insert_logs: Callable[[list[Logs]], None],
+) -> None:
+    now = datetime.now(tz=UTC)
+    insert_logs(
+        [
+            Logs(timestamp=now - timedelta(seconds=3), resources={"service.name": SERVICE}, attributes={"at$0dollar": "alpha-dollar", "at${x}": "a$0b", "n$0um": 5}, body="alpha body"),
+            Logs(timestamp=now - timedelta(seconds=2), resources={"service.name": SERVICE}, attributes={"at$0dollar": "beta-dollar", "at${x}": "it's", "n$0um": 7}, body="beta body"),
+            Logs(timestamp=now - timedelta(seconds=1), resources={"service.name": SERVICE}, attributes={"plain": "gamma"}, body="gamma body"),
+        ]
+    )
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+    start_ms = int((now - timedelta(minutes=2)).timestamp() * 1000)
+    end_ms = int((now + timedelta(seconds=1)).timestamp() * 1000)
+
+    for expression, expected in [
+        ("count_distinct(at$0dollar)", 2),
+        ("count(DISTINCT at$0dollar)", 2),
+        ("sum(n$0um)", 12),
+        ("countIf(at$0dollar EXISTS)", 2),
+        ("count_distinct(at${x})", 2),
+        ("countIf(at${x} = 'a$0b')", 1),
+        ("countIf(at${x} = 'it\\'s')", 1),
+        ("sumIf(n$0um, at$0dollar = 'beta-dollar')", 7),
+    ]:
+        response = make_query_request(
+            signoz,
+            token,
+            start_ms,
+            end_ms,
+            [build_scalar_query("A", "logs", [build_aggregation(expression)], filter_expression=f"service.name = '{SERVICE}'")],
+            request_type="scalar",
+        )
+        assert response.status_code == HTTPStatus.OK, f"{expression}: {response.text}"
+        assert get_scalar_table_data(response.json()) == [[expected]], f"{expression}: {response.text}"
+
+
 @pytest.mark.parametrize("key_id", list(RESOURCE_KEYS), ids=list(RESOURCE_KEYS))
 def test_logs_group_by_hostile_resource_attribute(
     signoz: types.SigNoz,

tests/integration/tests/querierlogs/11_aggregation_functions.py

@@ -128,6 +128,172 @@ def test_logs_aggregate_functions(
         assert by_service[service] == pytest.approx(expected), f"{service}: {by_service[service]} != {expected}"
 
 
+def test_logs_aggregate_if_predicates_use_filter_syntax(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    insert_logs: Callable[[list[Logs]], None],
+) -> None:
+    """
+    Setup:
+    svc-a: 3 logs, two with a k8s.pod.name and an x-request-id header, two
+    with "timeout" in the body; svc-b: 1 log without either, with "timeout".
+
+    Tests:
+    countIf predicates accept the filter-only operators EXISTS, NOT EXISTS,
+    CONTAINS and NOT CONTAINS, a combination of them, and a key with hyphens;
+    two predicates in one expression each keep their own values.
+    """
+    now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
+    # (service, k8s.pod.name, x-request-id header, body)
+    specs = [
+        ("svc-a", "pod-1", "abc", "request timeout"),
+        ("svc-a", None, None, "request ok"),
+        ("svc-a", "pod-2", "def", "upstream timeout"),
+        ("svc-b", None, None, "timeout"),
+    ]
+    logs = [
+        Logs(
+            timestamp=now - timedelta(seconds=i + 1),
+            resources={"service.name": service} | ({"k8s.pod.name": pod} if pod else {}),
+            attributes={"http.request.header.x-request-id": request_id} if request_id else {},
+            body=body,
+        )
+        for i, (service, pod, request_id, body) in enumerate(specs)
+    ]
+    insert_logs(logs)
+
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    query = build_scalar_query(
+        name="A",
+        signal="logs",
+        aggregations=[
+            build_aggregation("count()", "cnt"),
+            build_aggregation("countIf(k8s.pod.name EXISTS)", "with_pod"),
+            build_aggregation("countIf(k8s.pod.name NOT EXISTS)", "without_pod"),
+            build_aggregation("countIf(body CONTAINS 'timeout')", "timeouts"),
+            build_aggregation("countIf(body NOT CONTAINS 'timeout')", "not_timeouts"),
+            build_aggregation("countIf(body CONTAINS 'timeout' AND k8s.pod.name EXISTS)", "pod_timeouts"),
+            build_aggregation("countIf(http.request.header.x-request-id = 'abc')", "request_abc"),
+            build_aggregation("countIf(k8s.pod.name = 'pod-1') + countIf(body CONTAINS 'ok')", "pod_1_or_ok"),
+        ],
+        group_by=[build_group_by_field("service.name", "string", "resource")],
+        order=[build_order_by("count()", "desc")],
+    )
+    response = make_scalar_query_request(signoz, token, now, [query])
+
+    assert response.status_code == HTTPStatus.OK, response.text
+    data = get_scalar_table_data(response.json())
+    assert [row[0] for row in data] == ["svc-a", "svc-b"]
+    by_service = {row[0]: row[1:] for row in data}
+
+    for service in ["svc-a", "svc-b"]:
+        rows = [(pod, request_id, body) for svc, pod, request_id, body in specs if svc == service]
+        expected = [
+            len(rows),
+            sum(1 for pod, _, _ in rows if pod),
+            sum(1 for pod, _, _ in rows if not pod),
+            sum(1 for _, _, body in rows if "timeout" in body),
+            sum(1 for _, _, body in rows if "timeout" not in body),
+            sum(1 for pod, _, body in rows if "timeout" in body and pod),
+            sum(1 for _, request_id, _ in rows if request_id == "abc"),
+            sum(1 for pod, _, _ in rows if pod == "pod-1") + sum(1 for _, _, body in rows if "ok" in body),
+        ]
+        assert by_service[service] == expected, f"{service}: {by_service[service]} != {expected}"
+
+
+def test_logs_aggregate_arguments_are_field_names(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    insert_logs: Callable[[list[Logs]], None],
+) -> None:
+    """
+    Setup:
+    svc-a: 3 logs, svc-b: 1 log, each with a body size under a key with a
+    hyphen and an app label under a key with slashes.
+
+    Tests:
+    Aggregation arguments are read as field names, so keys with hyphens,
+    slashes and a type suffix resolve, count(*) counts every log, and DISTINCT
+    keeps its meaning.
+    """
+    now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
+    # (service, http.response.body-size, k8s.pod.labels.app.kubernetes.io/name)
+    specs = [
+        ("svc-a", 100, "web"),
+        ("svc-a", 250, "web"),
+        ("svc-a", 400, "api"),
+        ("svc-b", 70, "worker"),
+    ]
+    logs = [
+        Logs(
+            timestamp=now - timedelta(seconds=i + 1),
+            resources={"service.name": service},
+            attributes={"http.response.body-size": size, "k8s.pod.labels.app.kubernetes.io/name": app},
+            body=f"{service} log {i}",
+        )
+        for i, (service, size, app) in enumerate(specs)
+    ]
+    insert_logs(logs)
+
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    query = build_scalar_query(
+        name="A",
+        signal="logs",
+        aggregations=[
+            build_aggregation("count(*)", "cnt"),
+            build_aggregation("sum(http.response.body-size)", "body_size"),
+            build_aggregation("max(http.response.body-size:float64)", "max_body_size"),
+            build_aggregation("count_distinct(k8s.pod.labels.app.kubernetes.io/name)", "apps"),
+            build_aggregation("count(DISTINCT k8s.pod.labels.app.kubernetes.io/name)", "distinct_apps"),
+        ],
+        group_by=[build_group_by_field("service.name", "string", "resource")],
+        order=[build_order_by("count(*)", "desc")],
+    )
+    response = make_scalar_query_request(signoz, token, now, [query])
+
+    assert response.status_code == HTTPStatus.OK, response.text
+    data = get_scalar_table_data(response.json())
+    assert [row[0] for row in data] == ["svc-a", "svc-b"]
+    by_service = {row[0]: row[1:] for row in data}
+
+    for service in ["svc-a", "svc-b"]:
+        sizes = [size for svc, size, _ in specs if svc == service]
+        apps = {app for svc, _, app in specs if svc == service}
+        expected = [len(sizes), sum(sizes), max(sizes), len(apps), len(apps)]
+        assert by_service[service] == expected, f"{service}: {by_service[service]} != {expected}"
+
+
+@pytest.mark.parametrize(
+    "expression",
+    [
+        pytest.param("countIf(k8s.pod.name EXISTS", id="unclosed_parenthesis"),
+        pytest.param("countIf()", id="missing_predicate"),
+        pytest.param("countIf(latency_ms >= 25, latency_ms < 100)", id="extra_argument"),
+        pytest.param("sum(1)", id="literal_argument"),
+        pytest.param("sum(latency_ms * 2)", id="arithmetic_in_argument"),
+        pytest.param("sum(attributes_number['latency_ms'])", id="column_expression_argument"),
+        pytest.param("countIf(search('timeout'))", id="search_in_predicate"),
+    ],
+)
+def test_logs_aggregate_invalid_expression_is_bad_request(
+    signoz: types.SigNoz,
+    create_user_admin: None,  # pylint: disable=unused-argument
+    get_token: Callable[[str, str], str],
+    expression: str,
+) -> None:
+    now = datetime.now(tz=UTC).replace(second=0, microsecond=0)
+    token = get_token(USER_ADMIN_EMAIL, USER_ADMIN_PASSWORD)
+
+    query = build_scalar_query(name="A", signal="logs", aggregations=[build_aggregation(expression)])
+    response = make_scalar_query_request(signoz, token, now, [query])
+
+    assert response.status_code == HTTPStatus.BAD_REQUEST, response.text
+
+
 @pytest.mark.parametrize(
     "having_expression,aggregations,expected_services",
     [