From 75e603de8f342813df447f649d2a46bf9d2bc565 Mon Sep 17 00:00:00 2001 From: VibeSWE <338779468+VibeSWE@users.noreply.github.com> Date: Sat, 24 Oct 2026 15:11:15 -0500 Subject: [PATCH] Strip trailing semicolons by token so '; -- comment' queries run bounded_tool_sql removed the terminator with trim_end_matches(';'), which left the ';' inside the LIMIT wrapper whenever a comment followed it, so `SELECT a FROM t; -- note` passed validation and then failed to parse. Cut the query at the first ';' of the trailing run of ';'/comment tokens instead. Document that projection limits apply only to the outer query. Co-Authored-By: Claude Opus 5.5 --- docs/tool-api.md | 15 ++++++---- src/tool_executor.rs | 71 +++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 7 deletions(-) diff --git a/docs/tool-api.md b/docs/tool-api.md index fc9cca74f..3113562c1 100644 --- a/docs/tool-api.md +++ b/docs/tool-api.md @@ -135,12 +135,15 @@ and pagination are not identical to this summary tool. `query_sql`, `explain_query` and other tools that run SQL internally parse the statement with DataFusion's SQL parser and accept exactly one `SELECT` or `WITH` -query. A trailing semicolon and trailing comments are allowed. Every `SELECT` -projection, including those inside nested queries and set operations, must -name at most 10 explicit columns. `*` and `table.*` are rejected. Casts to -`INT64`, `UINT64` or `FLOAT64` are rejected with a hint to use `BIGINT` or -`DOUBLE`. The statement is wrapped with an outer `LIMIT 500`. `sample_events` -returns at most 10 columns, choosing common observability fields such as +query. Trailing semicolons and trailing comments are allowed, including +comments after the final semicolon (`SELECT a FROM t; -- note`). The outer +query's `SELECT` projection, including each branch of a set operation and any +parenthesized query, must name at most 10 explicit columns, and `*` and +`table.*` are rejected there. Projections inside CTEs, derived tables and +subquery expressions are not checked; the result width is bounded by the outer +projection. Casts to `INT64`, `UINT64` or `FLOAT64` are rejected with a hint +to use `BIGINT` or `DOUBLE`. The statement is wrapped with an outer +`LIMIT 500`. `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. diff --git a/src/tool_executor.rs b/src/tool_executor.rs index f2c477041..90269c154 100644 --- a/src/tool_executor.rs +++ b/src/tool_executor.rs @@ -14,6 +14,10 @@ use crate::{ DataType as SqlDataType, Expr as SqlExpr, SelectItem, SetExpr, Statement as SqlStatement, visit_expressions, }, + sqlparser::{ + dialect::GenericDialect, + tokenizer::{Token, Tokenizer}, + }, }, handlers::http::{ alerts, cluster, health_check, logstream, @@ -125,11 +129,44 @@ pub fn validate_tool_sql(query: &str) -> Result<(), String> { validate_projection(query.body.as_ref()) } +/// Drops trailing statement terminators, along with any comments after them, +/// so that `SELECT a FROM t; -- note` can be wrapped in a subquery. Comments +/// and semicolons inside the statement (for example in string literals) are +/// kept because the cut is made at a tokenizer `;`, not a textual one. +fn strip_trailing_semicolons(query: &str) -> &str { + let Ok(tokens) = Tokenizer::new(&GenericDialect {}, query).tokenize_with_location() else { + return query; + }; + let Some(semicolon) = tokens + .iter() + .rev() + .take_while(|token| matches!(token.token, Token::Whitespace(_) | Token::SemiColon)) + .filter(|token| token.token == Token::SemiColon) + .last() + else { + return query; + }; + let target = semicolon.span.start; + let (mut line, mut column) = (1, 1); + for (offset, character) in query.char_indices() { + if line == target.line && column == target.column { + return query[..offset].trim_end(); + } + if character == '\n' { + line += 1; + column = 1; + } else { + column += 1; + } + } + query +} + pub fn bounded_tool_sql(query: &str) -> Result { validate_tool_sql(query)?; Ok(format!( "SELECT * FROM ({}\n) AS \"__parseable_tool_result\" LIMIT {MAX_TOOL_ROWS}", - query.trim().trim_end_matches(';') + strip_trailing_semicolons(query.trim()) )) } @@ -1964,6 +2001,38 @@ mod tests { assert!(bounded_tool_sql("SELECT a,b,c,d,e,f,g,h,i,j,k FROM logs").is_err()); let bounded = bounded_tool_sql("SELECT a FROM logs -- recent events").unwrap(); assert!(bounded.contains("-- recent events\n) AS")); + // Only the outer projection is limited; nested queries may use `*`. + assert!(bounded_tool_sql("WITH x AS (SELECT * FROM logs) SELECT a FROM x").is_ok()); + assert!(bounded_tool_sql("SELECT a FROM (SELECT * FROM logs) AS t").is_ok()); + assert!(bounded_tool_sql("SELECT a FROM logs WHERE EXISTS (SELECT * FROM logs)").is_ok()); + assert!(bounded_tool_sql("WITH x AS (SELECT a FROM logs) SELECT * FROM x").is_err()); + } + + #[actix_web::test] + async fn bounded_tool_sql_runs_with_comments_after_the_trailing_semicolon() { + let context = datafusion::prelude::SessionContext::new(); + for (source, expected) in [ + ("SELECT 'x' AS a; -- note", "x"), + ("SELECT 'x' AS a; /* note */", "x"), + ("SELECT 'x' AS a ; ; -- note\n-- another note\n", "x"), + ("SELECT 'x' AS a -- note\n; -- another note", "x"), + ("SELECT 'x;' AS a; -- note", "x;"), + ("SELECT 'x' AS a -- note;", "x"), + ("SELECT 'x' AS a /* ; */ ;", "x"), + ("SELECT 'é' AS a; -- note", "é"), + ] { + let bounded = bounded_tool_sql(source).unwrap(); + let batches = context + .sql(&bounded) + .await + .unwrap_or_else(|error| panic!("{source:?} -> {bounded:?}: {error}")) + .collect() + .await + .unwrap(); + let column = batches[0].column(0); + let value = datafusion::arrow::util::display::array_value_to_string(column, 0).unwrap(); + assert_eq!(value, expected, "{source:?} -> {bounded:?}"); + } } #[actix_web::test]