diff --git a/docs/tool-api.md b/docs/tool-api.md index 74d415fdd..26cffe630 100644 --- a/docs/tool-api.md +++ b/docs/tool-api.md @@ -148,7 +148,8 @@ top-level set operation such as `UNION` and to a parenthesized outer query. Projections inside CTE definitions, derived tables and subqueries are not checked. Casts to `INT64`, `UINT64` or `FLOAT64` anywhere in the statement are rejected with a hint to use `BIGINT` or `DOUBLE`. The statement is wrapped with -an outer `LIMIT 500`. `sample_events` +an outer `LIMIT 500`, except in `get_trace_activity`, whose generated SQL carries +its own limit (see below). `sample_events` returns at most 10 columns, choosing common observability fields such as `p_timestamp`, `service.name`, `severity_text`, `body` and `trace_id` before other fields. @@ -243,14 +244,17 @@ dataset's time partition, or `p_timestamp` if none is configured. The dataset must have traces telemetry and contain `span_span_id` and `span_trace_id`. Unknown fields are rejected. `span_count` is `COUNT(DISTINCT span_span_id)` and `trace_count` is `COUNT(DISTINCT span_trace_id)`, so duplicate span rows are -not double counted. Rows are ordered by bucket. The result contains `query`, -`bucketInterval`, `semantics` and up to 500 `rows`. - -`query` is the aggregation SQL before the tool wraps it with `LIMIT 500`. It -has no time predicate: the time range comes from the request's `startTime` and -`endTime`, which the query API rounds down to the minute. To reproduce the -result, for example in a dashboard tile, run `query` through the query API with -the same `startTime` and `endTime`. +not double counted. Rows are ordered by bucket, oldest first. The result +contains `query`, `bucketInterval`, `semantics`, up to 500 `rows` and +`truncated`. + +`query` is the exact SQL the tool runs; it ends in `LIMIT 501` so the tool can +detect overflow. When more than 500 rows match, the newest buckets are dropped +and `truncated` is `true`; reduce `bucketCount`, `groupBy` or the time range to +get the full series. `query` has no time predicate: the time range comes from +the request's `startTime` and `endTime`, which the query API rounds down to the +minute. To reproduce the result, for example in a dashboard tile, run `query` +through the query API with the same `startTime` and `endTime`. ## Argument schema validation diff --git a/src/tool_executor.rs b/src/tool_executor.rs index 360632335..b32a4d299 100644 --- a/src/tool_executor.rs +++ b/src/tool_executor.rs @@ -1400,15 +1400,32 @@ fn trace_activity_query( .first() .map(|_| format!(" WHERE {}", clauses.join(" AND "))) .unwrap_or_default(); + // One row past the tool limit lets the caller detect truncation. format!( - "SELECT {} FROM {}{} GROUP BY {} ORDER BY 1 LIMIT 500", + "SELECT {} FROM {}{} GROUP BY {} ORDER BY 1 LIMIT {}", select.join(", "), quote_identifier(dataset), where_clause, groups, + MAX_TOOL_ROWS + 1, ) } +/// Caps trace activity rows at the tool limit and reports whether any were dropped. +fn truncate_trace_activity_rows(rows: &mut Value) -> bool { + let rows = match rows { + Value::Array(rows) => rows, + Value::Object(object) => match object.get_mut("records").and_then(Value::as_array_mut) { + Some(rows) => rows, + None => return false, + }, + _ => return false, + }; + let truncated = rows.len() > MAX_TOOL_ROWS as usize; + rows.truncate(MAX_TOOL_ROWS as usize); + truncated +} + async fn trace_activity_fields( dataset: &str, context: OssToolContext<'_>, @@ -1532,9 +1549,12 @@ async fn execute_trace_activity_tool( let bucket_count = bounded_integer(arguments, "bucketCount", 60, 1, 100); let interval = time_bucket_interval(start_time, end_time, bucket_count); let query = trace_activity_query(dataset, &time_field, &interval, &group_by, &clauses); + // The generated query carries its own LIMIT; wrapping it in the generic + // 500-row bound would hide the extra row used to detect truncation. + validate_tool_sql(&query)?; let response = query::query_internal( query::Query { - query: bounded_select(&query)?, + query: query.clone(), start_time: start_time.to_owned(), end_time: end_time.to_owned(), send_null: true, @@ -1547,7 +1567,8 @@ async fn execute_trace_activity_tool( ) .await .map_err(|error| error.to_string())?; - let rows = response_value(response).await?; + let mut rows = response_value(response).await?; + let truncated = truncate_trace_activity_rows(&mut rows); Ok(json!({ "dataset": dataset, "queryType": "sql", @@ -1558,6 +1579,7 @@ async fn execute_trace_activity_tool( "trace_count": "COUNT(DISTINCT span_trace_id)" }, "rows": rows, + "truncated": truncated, })) } @@ -2305,21 +2327,8 @@ mod tests { let interval = super::time_bucket_interval("2026-10-22T00:00:00Z", "2026-10-22T00:02:00Z", 2); assert_eq!(interval, "60 seconds"); - let query = bounded_tool_sql(&trace_activity_query( - "traces", - "p_timestamp", - &interval, - group_by, - clauses, - )) - .unwrap(); - let batches = context.sql(&query).await.unwrap().collect().await.unwrap(); - let mut writer = arrow_json::ArrayWriter::new(Vec::new()); - writer - .write_batches(&batches.iter().collect::>()) - .unwrap(); - writer.finish().unwrap(); - let rows: Vec = serde_json::from_slice(&writer.into_inner()).unwrap(); + let query = trace_activity_query("traces", "p_timestamp", &interval, group_by, clauses); + let rows = execute_trace_activity_sql(&context, &query).await; let times = rows .iter() .map(|row| row["time"].as_str().unwrap()) @@ -2331,6 +2340,69 @@ mod tests { rows } + /// Runs trace activity SQL the way the tool does: validated, without the + /// generic 500-row wrapper. + async fn execute_trace_activity_sql( + context: &datafusion::prelude::SessionContext, + query: &str, + ) -> Vec { + super::validate_tool_sql(query).unwrap(); + let batches = context.sql(query).await.unwrap().collect().await.unwrap(); + let mut writer = arrow_json::ArrayWriter::new(Vec::new()); + writer + .write_batches(&batches.iter().collect::>()) + .unwrap(); + writer.finish().unwrap(); + serde_json::from_slice(&writer.into_inner()).unwrap() + } + + #[actix_web::test] + async fn trace_activity_signals_truncation_past_the_row_limit() { + // 20 services x 60 one-minute buckets = 1200 groups x buckets. + let context = datafusion::prelude::SessionContext::new(); + let fixture = context + .sql( + r#"SELECT CAST(TO_TIMESTAMP_SECONDS(1792627200 + (value / 20) * 60) AS TIMESTAMP) + AS p_timestamp, + CONCAT('t', value) AS span_trace_id, CONCAT('s', value) AS span_span_id, + CONCAT('svc-', value % 20) AS "service.name" + FROM range(1200)"#, + ) + .await + .unwrap(); + context + .register_table("traces", fixture.into_view()) + .unwrap(); + let interval = + super::time_bucket_interval("2026-10-22T00:00:00Z", "2026-10-22T01:00:00Z", 60); + assert_eq!(interval, "60 seconds"); + let query = trace_activity_query( + "traces", + "p_timestamp", + &interval, + &["service.name".to_owned()], + &[], + ); + let mut rows = Value::Array(execute_trace_activity_sql(&context, &query).await); + assert!(super::truncate_trace_activity_rows(&mut rows)); + let rows = rows.as_array().unwrap(); + assert_eq!(rows.len(), 500); + // Ascending order keeps the oldest 25 buckets; the newest are the ones dropped. + assert_eq!(rows[499]["time"], "2026-10-22T00:24:00"); + + // A result that fits is returned whole and not flagged. + let query = trace_activity_query("traces", "p_timestamp", &interval, &[], &[]); + let mut rows = Value::Array(execute_trace_activity_sql(&context, &query).await); + assert!(!super::truncate_trace_activity_rows(&mut rows)); + assert_eq!(rows.as_array().unwrap().len(), 60); + + let mut exact = Value::Array(vec![json!({}); 500]); + assert!(!super::truncate_trace_activity_rows(&mut exact)); + let mut records = json!({ "records": vec![json!({}); 501] }); + assert!(super::truncate_trace_activity_rows(&mut records)); + assert_eq!(records["records"].as_array().unwrap().len(), 500); + } + #[actix_web::test] async fn trace_activity_query_executes_with_distinct_bucketed_counts() { let rows = run_trace_activity_fixture(&["service.name".to_owned()], &[], "BIGINT").await;