diff --git a/crates/dbx-core/src/query_result_sql.rs b/crates/dbx-core/src/query_result_sql.rs index 455035bd9c..713b349dfc 100644 --- a/crates/dbx-core/src/query_result_sql.rs +++ b/crates/dbx-core/src/query_result_sql.rs @@ -224,18 +224,19 @@ pub fn build_count_query_sql(options: CountQuerySqlOptions) -> QuerySqlBuildResu if unsupported_pagination_type(options.database_type) { return err("unsupported"); } + let (execution_hint, statement) = split_leading_execution_hint(&statement); // A locking clause does not affect cardinality and cannot appear inside // every dialect's derived-table count query. PostgreSQL permits pagination // after the lock clause; decline counting that uncommon order rather than // accidentally dropping the user's explicit LIMIT/OFFSET. - let tokens = top_level_sql_tokens(&statement); + let tokens = top_level_sql_tokens(statement); let statement = if let Some(index) = locking_clause_index(&tokens) { if has_pagination_clause_after(&tokens, index) { return err("locking"); } statement[..index].trim_end().to_string() } else { - statement + statement.to_string() }; // ES SQL can't wrap a SELECT in `SELECT COUNT(*) FROM (...)` — the // driver already reports the true match count via affected_rows. @@ -243,7 +244,9 @@ pub fn build_count_query_sql(options: CountQuerySqlOptions) -> QuerySqlBuildResu return err("unsupported"); } if options.database_type == Some(DatabaseType::SqlServer) { - return sql_server_count_sql(&statement).map(ok).unwrap_or_else(|| err("unsupported")); + return sql_server_count_sql(&statement) + .map(|sql| ok(format!("{execution_hint}{sql}"))) + .unwrap_or_else(|| err("unsupported")); } let alias = quote_table_identifier(options.database_type, "dbx_count"); @@ -251,7 +254,10 @@ pub fn build_count_query_sql(options: CountQuerySqlOptions) -> QuerySqlBuildResu Some(DatabaseType::Iris) => iris_statement_for_derived_table(&statement), _ => statement, }; - ok(derived_table_sql("SELECT COUNT(*) AS dbx_total_rows FROM", &wrapped_sql, &format!("{alias};"))) + ok(format!( + "{execution_hint}{}", + derived_table_sql("SELECT COUNT(*) AS dbx_total_rows FROM", &wrapped_sql, &format!("{alias};")) + )) } pub fn build_sorted_query_sql(options: SortedQuerySqlOptions) -> QuerySqlBuildResult { @@ -354,13 +360,14 @@ fn single_selectable_statement(original_sql: &str) -> Result { return Err(()); } - let statement = find_statement_at_cursor(base_sql, 0).trim().trim_end_matches(';').trim().to_string(); - if statement.is_empty() { + let extracted = find_statement_at_cursor(base_sql, 0).trim().trim_end_matches(';').trim().to_string(); + if extracted.is_empty() { return Err(()); } - if !single_statement_matches_base_sql(&statement, base_sql) { + if !single_statement_matches_base_sql(&extracted, base_sql) { return Err(()); } + let statement = restore_leading_execution_hint(base_sql, &extracted); let statement_without_leading_comments = strip_leading_statement_comments(statement.trim_start_matches(';').trim_start()); let upper = statement_without_leading_comments.to_ascii_uppercase(); @@ -573,6 +580,56 @@ fn skip_leading_sql_comments(sql: &str, mut index: usize) -> usize { } } +fn restore_leading_execution_hint(original_sql: &str, statement: &str) -> String { + if leading_execution_hint_prefix(statement).is_some() { + return statement.to_string(); + } + let normalized = original_sql.trim().trim_end_matches(';').trim(); + match leading_execution_hint_prefix(normalized) { + Some(prefix) => format!("{prefix}{statement}"), + None => statement.to_string(), + } +} + +fn split_leading_execution_hint(sql: &str) -> (&str, &str) { + match leading_execution_hint_prefix(sql) { + Some(prefix) => (prefix, sql[prefix.len()..].trim_start()), + None => ("", sql), + } +} + +fn leading_execution_hint_prefix(sql: &str) -> Option<&str> { + let mut index = 0; + let mut hint_start = None; + loop { + index = skip_sql_whitespace(sql, index); + let rest = &sql[index..]; + if rest.starts_with("--") { + index += 2; + while index < sql.len() && next_char(sql, index) != '\n' { + index += next_char(sql, index).len_utf8(); + } + continue; + } + if !rest.starts_with("/*") { + return hint_start.map(|start| &sql[start..index]); + } + if hint_start.is_none() && matches!(next_char_at(sql, index + 2), Some('+' | '@' | '&')) { + hint_start = Some(index); + } + index += 2; + while index < sql.len() { + let ch = next_char(sql, index); + let next = next_char_at(sql, index + ch.len_utf8()); + index += ch.len_utf8(); + if ch == '*' && next == Some('/') { + index += 1; + break; + } + } + } +} + fn strip_leading_statement_comments(sql: &str) -> &str { &sql[skip_leading_sql_comments(sql, 0)..] } @@ -2606,6 +2663,61 @@ WHERE u.id = picked.id; assert_eq!(result.sql.unwrap(), "SELECT id FROM users WHERE active = 1 LIMIT 50;"); } + #[test] + fn mysql_pagination_preserves_leading_ampersand_routing_hint() { + let result = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: "/*& tenant:'test' */\nSELECT id FROM users".to_string(), + database_type: Some(DatabaseType::Mysql), + limit: 50, + offset: 0, + }); + + assert_eq!(result.sql.unwrap(), "/*& tenant:'test' */\nSELECT id FROM users LIMIT 50;"); + } + + #[test] + fn mysql_pagination_preserves_supported_leading_execution_hints_only() { + for hint in ["/*+ MAX_EXECUTION_TIME(1000) */", "/*@global:true*/", "/*& tenant:'test' */"] { + let result = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: format!("{hint}\nSELECT id FROM users"), + database_type: Some(DatabaseType::Mysql), + limit: 50, + offset: 0, + }); + + assert_eq!(result.sql.unwrap(), format!("{hint}\nSELECT id FROM users LIMIT 50;")); + } + + let ordinary_comment = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: "/* report query */\nSELECT id FROM users".to_string(), + database_type: Some(DatabaseType::Mysql), + limit: 50, + offset: 0, + }); + assert_eq!(ordinary_comment.sql.unwrap(), "SELECT id FROM users LIMIT 50;"); + + let ordinary_comment_before_hint = build_paginated_query_sql(PaginatedQuerySqlOptions { + original_sql: "-- report query\n/*& tenant:'test' */\nSELECT id FROM users".to_string(), + database_type: Some(DatabaseType::Mysql), + limit: 50, + offset: 0, + }); + assert_eq!(ordinary_comment_before_hint.sql.unwrap(), "/*& tenant:'test' */\nSELECT id FROM users LIMIT 50;"); + } + + #[test] + fn mysql_count_preserves_leading_ampersand_routing_hint() { + let result = build_count_query_sql(CountQuerySqlOptions { + original_sql: "/*& tenant:'test' */\nSELECT id FROM users".to_string(), + database_type: Some(DatabaseType::Mysql), + }); + + assert_eq!( + result.sql.unwrap(), + "/*& tenant:'test' */\nSELECT COUNT(*) AS dbx_total_rows FROM (SELECT id FROM users) `dbx_count`;" + ); + } + #[test] fn mysql_pagination_does_not_wrap_duplicate_result_columns() { let result = build_paginated_query_sql(PaginatedQuerySqlOptions {