From 5695ef6e77cf66270b6150d1d4537e948e8293ff Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9D=8E=E5=AD=90=E5=87=A1?= <113765024+xiaou61@users.noreply.github.com> Date: Fri, 14 Aug 2026 22:46:32 +0800 Subject: [PATCH] fix(sql): preserve same-line TDSQL directives --- .../sql/executableStatementRangeCache.spec.ts | 17 +++--- .../__tests__/sql/sqlExecutionTarget.spec.ts | 6 +++ .../__tests__/sql/sqlStatementRanges.spec.ts | 30 +++++++++++ .../desktop/src/lib/sql/sqlStatementRanges.ts | 40 +++++++++++--- crates/dbx-core/src/db/tdsql_mysql.rs | 38 ++++++++----- crates/dbx-core/src/sql.rs | 53 +++++++++++++++---- 6 files changed, 145 insertions(+), 39 deletions(-) diff --git a/apps/desktop/src/lib/__tests__/sql/executableStatementRangeCache.spec.ts b/apps/desktop/src/lib/__tests__/sql/executableStatementRangeCache.spec.ts index 9191aa0c65..7bd6a6269a 100644 --- a/apps/desktop/src/lib/__tests__/sql/executableStatementRangeCache.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/executableStatementRangeCache.spec.ts @@ -118,12 +118,12 @@ describe("executableStatementRangeCacheForDoc", () => { expect(executableStatementRangeStartingAt(cache, doc.line(3).from)).toBeNull(); }); - it("does not resolve gutter run buttons when non-whitespace precedes the statement on the same line", () => { - const doc = Text.of(["/* comment */ SELECT 1;"]); + it("resolves a same-line leading TDSQL directive for gutter execution", () => { + const sql = "/*sets:allsets */ SELECT 1;"; + const doc = Text.of([sql]); const cache = executableStatementRangeCacheForDoc(null, doc, "mysql"); - expect(executableStatementRangeStartingAt(cache, doc.line(1).from)).toBeNull(); - expect(executableStatementRangeStartingAt(cache, doc.toString().indexOf("SELECT"))?.sql).toBe("SELECT 1"); + expect(executableStatementRangeStartingAt(cache, doc.line(1).from)?.sql).toBe(sql.slice(0, -1)); }); it("resolves the current statement from a cursor inside a continuation line", () => { @@ -205,12 +205,13 @@ describe("executableStatementRangeCacheForDoc", () => { expect(executableStatementRangeAtCursor(cache, doc.line(4).from)).toBeNull(); }); - it("resolves SQL after a leading block comment on the same line", () => { - const doc = Text.of(["/* comment */ SELECT 1;"]); + it("preserves a same-line TDSQL directive while ignoring a cursor inside its comment", () => { + const sql = "/*sets:allsets */ SELECT 1;"; + const doc = Text.of([sql]); const cache = executableStatementRangeCacheForDoc(null, doc, "mysql"); - expect(executableStatementRangeAtCursor(cache, doc.toString().indexOf("SELECT"))?.sql).toBe("SELECT 1"); - expect(executableStatementRangeAtCursor(cache, doc.toString().indexOf("comment"))).toBeNull(); + expect(executableStatementRangeAtCursor(cache, sql.indexOf("SELECT"))?.sql).toBe(sql.slice(0, -1)); + expect(executableStatementRangeAtCursor(cache, sql.indexOf("sets"))).toBeNull(); }); it("rebuilds the cache when the document instance changes", () => { diff --git a/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts index bf27f299bd..6b123e8beb 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts @@ -19,6 +19,12 @@ describe("resolveExecutableSql", () => { expect(resolveExecutableSql("SELECT 1;", selectedSql, { mode: "current", cursorPos: 0 })).toBe(selectedSql); }); + + it("keeps a manually selected same-line TDSQL directive unchanged", () => { + const selectedSql = "/*sets:allsets */ SELECT count(*) FROM tenant_table"; + + expect(resolveExecutableSql("SELECT 1;", selectedSql, { mode: "current", cursorPos: 0 })).toBe(selectedSql); + }); }); describe("executionCandidateForMode", () => { diff --git a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts index 1fd007bac4..ff748b4adf 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts @@ -1263,6 +1263,36 @@ describe("buildExecutionCandidates", () => { expect(splitSqlStatementRanges("/*proxy*/", "mysql")).toEqual([]); }); + it.each(["/*sets:allsets */", "/*master*/", "/*slave:set_1781591902_7*/", "/*future-route:anywhere*/"])("preserves a same-line TDSQL directive without relying on a keyword allowlist: %s", (directive) => { + const directedSql = `${directive} SELECT count(*) FROM tenant_table`; + const sql = `SELECT 1;\n${directedSql};\nSELECT 2;`; + const candidates = buildExecutionCandidates(sql, indexOf(sql, "tenant_table"), "mysql"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe(directedSql); + expect(rangeSqlTexts(executableStatementRanges(sql, "mysql"))).toEqual(["SELECT 1", directedSql, "SELECT 2"]); + }); + + it("handles long same-line directive chains without rescanning growing prefixes", () => { + const prefix = "/**/".repeat(80_000); + const sql = `${prefix} SELECT 1`; + + expect(splitSqlStatementRanges(sql, "mysql")[0]?.sql).toBe(sql); + }); + + it("does not preserve a generic TDSQL-style directive on a separate line", () => { + const sql = "/*sets:allsets */\nSELECT count(*) FROM tenant_table"; + const candidates = buildExecutionCandidates(sql, indexOf(sql, "tenant_table"), "mysql"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe("SELECT count(*) FROM tenant_table"); + }); + + it("does not preserve a same-line TDSQL-style directive for other database types", () => { + const sql = "/*sets:allsets */ SELECT count(*) FROM tenant_table"; + const candidates = buildExecutionCandidates(sql, indexOf(sql, "tenant_table"), "postgres"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe("SELECT count(*) FROM tenant_table"); + }); + it.each(["/* ordinary */", "/*unknown*/", "/* proxy */", "/*PROXY*/"])("keeps %s as a non-executable leading comment", (comment) => { const sql = `${comment}\nSHOW PROXY STATUS`; const candidates = buildExecutionCandidates(sql, indexOf(sql, "STATUS"), "mysql"); diff --git a/apps/desktop/src/lib/sql/sqlStatementRanges.ts b/apps/desktop/src/lib/sql/sqlStatementRanges.ts index ede1d8435c..e9f3748ba2 100644 --- a/apps/desktop/src/lib/sql/sqlStatementRanges.ts +++ b/apps/desktop/src/lib/sql/sqlStatementRanges.ts @@ -301,7 +301,8 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType let statementEnd = -1; let statementHitStart = 0; let pendingHintStart = -1; - let pendingProxyDirectiveStart = -1; + let pendingMysqlDirectiveStart = -1; + let pendingMysqlDirectiveLineEnd = -1; let customDelimiter: string | null = null; let state: QuoteState = "none"; let dollarTag = ""; @@ -310,12 +311,23 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType let i = 0; const isWhitespace = (ch: string) => ch === " " || ch === "\t" || ch === "\r" || ch === "\n"; + const nextLineBreak = (from: number) => { + const carriageReturn = sql.indexOf("\r", from); + const lineFeed = sql.indexOf("\n", from); + if (carriageReturn === -1) return lineFeed === -1 ? len : lineFeed; + if (lineFeed === -1) return carriageReturn; + return Math.min(carriageReturn, lineFeed); + }; const markContent = (pos: number) => { if (statementStart === -1) { - statementStart = pendingProxyDirectiveStart !== -1 ? pendingProxyDirectiveStart : pendingHintStart === -1 ? pos : pendingHintStart; + // TDSQL accepts arbitrary leading block directives on the SQL line. The + // exact /*proxy*/ directive also keeps its historical multiline support. + const directiveStart = pendingMysqlDirectiveStart !== -1 && (sql.startsWith("/*proxy*/", pendingMysqlDirectiveStart) || pos < pendingMysqlDirectiveLineEnd) ? pendingMysqlDirectiveStart : -1; + statementStart = directiveStart !== -1 ? directiveStart : pendingHintStart === -1 ? pos : pendingHintStart; pendingHintStart = -1; - pendingProxyDirectiveStart = -1; + pendingMysqlDirectiveStart = -1; + pendingMysqlDirectiveLineEnd = -1; } statementEnd = pos + 1; }; @@ -324,7 +336,8 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType if (statementStart === -1) { statementEnd = -1; pendingHintStart = -1; - pendingProxyDirectiveStart = -1; + pendingMysqlDirectiveStart = -1; + pendingMysqlDirectiveLineEnd = -1; postgresDollarQuotedRoutine = false; oraclePlSqlStatementEnd = undefined; return; @@ -336,7 +349,8 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType statementStart = -1; statementEnd = -1; pendingHintStart = -1; - pendingProxyDirectiveStart = -1; + pendingMysqlDirectiveStart = -1; + pendingMysqlDirectiveLineEnd = -1; postgresDollarQuotedRoutine = false; oraclePlSqlStatementEnd = undefined; }; @@ -441,13 +455,15 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType // Line comments consume up to (and including) the newline. if (ch === "-" && next === "-") { - pendingProxyDirectiveStart = -1; + pendingMysqlDirectiveStart = -1; + pendingMysqlDirectiveLineEnd = -1; const newline = sql.indexOf("\n", i); i = newline === -1 ? len : newline + 1; continue; } if (startsHashLineComment(sql, i, databaseType, parameterOptions)) { - pendingProxyDirectiveStart = -1; + pendingMysqlDirectiveStart = -1; + pendingMysqlDirectiveLineEnd = -1; const newline = sql.indexOf("\n", i); i = newline === -1 ? len : newline + 1; continue; @@ -458,7 +474,15 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType if (statementStart === -1 && pendingHintStart === -1 && (hintMarker === "+" || hintMarker === "@" || hintMarker === "&")) pendingHintStart = i; const close = sql.indexOf("*/", i + 2); if (statementStart === -1) { - pendingProxyDirectiveStart = databaseType === "mysql" && sql.startsWith("/*proxy*/", i) ? i : -1; + if (databaseType === "mysql") { + if (pendingMysqlDirectiveStart === -1 || i > pendingMysqlDirectiveLineEnd) { + pendingMysqlDirectiveStart = i; + pendingMysqlDirectiveLineEnd = nextLineBreak(i); + } + } else { + pendingMysqlDirectiveStart = -1; + pendingMysqlDirectiveLineEnd = -1; + } } i = close === -1 ? len : close + 2; continue; diff --git a/crates/dbx-core/src/db/tdsql_mysql.rs b/crates/dbx-core/src/db/tdsql_mysql.rs index 3dab4815db..90ab5cf7c7 100644 --- a/crates/dbx-core/src/db/tdsql_mysql.rs +++ b/crates/dbx-core/src/db/tdsql_mysql.rs @@ -11,13 +11,21 @@ pub fn is_profile(db_type: &DatabaseType, driver_profile: Option<&str>) -> bool && driver_profile.is_some_and(|profile| profile.eq_ignore_ascii_case(DRIVER_PROFILE)) } -pub(crate) fn preserves_proxy_directive_for_database_type(db_type: DatabaseType) -> bool { +pub(crate) fn preserves_leading_directives_for_database_type(db_type: DatabaseType) -> bool { db_type == DatabaseType::Mysql } -pub(crate) fn proxy_directive_start(statement: &str, executable_start: usize) -> Option { - let prefix = statement.get(..executable_start)?.trim_end(); - prefix.strip_suffix("/*proxy*/").map(|before| before.len()) +pub(crate) fn leading_directive_start(statement: &str, executable_start: usize) -> Option { + let prefix = statement.get(..executable_start)?; + let line_start = prefix.rfind(['\r', '\n']).map_or(0, |index| index + 1); + let line_prefix = &prefix[line_start..]; + let trimmed_line_start = line_prefix.trim_start(); + if trimmed_line_start.starts_with("/*") && trimmed_line_start.trim_end().ends_with("*/") { + return Some(line_start + line_prefix.len() - trimmed_line_start.len()); + } + + let trimmed_prefix = prefix.trim_end(); + trimmed_prefix.strip_suffix("/*proxy*/").map(|before| before.len()) } #[cfg(test)] @@ -32,23 +40,27 @@ mod tests { } #[test] - fn proxy_directive_is_available_only_to_mysql_parsing() { - assert!(preserves_proxy_directive_for_database_type(DatabaseType::Mysql)); - assert!(!preserves_proxy_directive_for_database_type(DatabaseType::Postgres)); + fn leading_directives_are_available_only_to_mysql_parsing() { + assert!(preserves_leading_directives_for_database_type(DatabaseType::Mysql)); + assert!(!preserves_leading_directives_for_database_type(DatabaseType::Postgres)); } #[test] - fn proxy_directive_must_be_exact_and_adjacent_to_executable_sql() { + fn leading_directives_preserve_same_line_and_exact_proxy_behavior() { let exact = "/*proxy*/ \n\tSHOW PROXY STATUS"; let executable_start = exact.find("SHOW").unwrap(); - assert_eq!(proxy_directive_start(exact, executable_start), Some(0)); + assert_eq!(leading_directive_start(exact, executable_start), Some(0)); - let ordinary = "/* proxy */\nSHOW PROXY STATUS"; - let executable_start = ordinary.find("SHOW").unwrap(); - assert_eq!(proxy_directive_start(ordinary, executable_start), None); + let same_line = " /*sets:allsets */ SELECT 1"; + let executable_start = same_line.find("SELECT").unwrap(); + assert_eq!(leading_directive_start(same_line, executable_start), Some(2)); + + let separate_line = "/*sets:allsets */\nSELECT 1"; + let executable_start = separate_line.find("SELECT").unwrap(); + assert_eq!(leading_directive_start(separate_line, executable_start), None); let separated = "/*proxy*/\n/* audit */\nSHOW PROXY STATUS"; let executable_start = separated.find("SHOW").unwrap(); - assert_eq!(proxy_directive_start(separated, executable_start), None); + assert_eq!(leading_directive_start(separated, executable_start), None); } } diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index efe517be5d..76c5b485fd 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -117,7 +117,7 @@ struct SqlDialectProfile { supports_hana_do_blocks: bool, supports_go_batch_separator: bool, keeps_sqlserver_module_batch_at_cursor: bool, - preserves_tdsql_proxy_directive: bool, + preserves_tdsql_leading_directives: bool, requires_whitespace_after_line_comment_dashes: bool, } @@ -134,7 +134,7 @@ impl Default for SqlDialectProfile { supports_hana_do_blocks: false, supports_go_batch_separator: false, keeps_sqlserver_module_batch_at_cursor: false, - preserves_tdsql_proxy_directive: false, + preserves_tdsql_leading_directives: false, requires_whitespace_after_line_comment_dashes: false, } } @@ -144,9 +144,8 @@ impl SqlDialectProfile { fn for_database_type(db_type: DatabaseType) -> Self { if matches!(db_type, DatabaseType::Mysql) { return Self { - preserves_tdsql_proxy_directive: crate::db::tdsql_mysql::preserves_proxy_directive_for_database_type( - db_type, - ), + preserves_tdsql_leading_directives: + crate::db::tdsql_mysql::preserves_leading_directives_for_database_type(db_type), ..Self::mysql_compatible() }; } @@ -2530,8 +2529,8 @@ fn executable_sql_bounds(statement: &str, options: SqlParsingOptions) -> Option< return None; } let executable_start = trimmed.len() - executable.len(); - let start = if options.profile.preserves_tdsql_proxy_directive { - crate::db::tdsql_mysql::proxy_directive_start(trimmed, executable_start).unwrap_or(executable_start) + let start = if options.profile.preserves_tdsql_leading_directives { + crate::db::tdsql_mysql::leading_directive_start(trimmed, executable_start).unwrap_or(executable_start) } else { executable_start }; @@ -3384,14 +3383,14 @@ SELECT 2;"; assert!(!default.supports_slash_line_block_delimiter); assert!(!default.supports_go_batch_separator); assert!(!default.keeps_sqlserver_module_batch_at_cursor); - assert!(!default.preserves_tdsql_proxy_directive); + assert!(!default.preserves_tdsql_leading_directives); let mysql = SqlDialectProfile::for_database_type(DatabaseType::Mysql); assert!(mysql.supports_hash_line_comments); assert!(mysql.supports_mysql_routine_blocks); - assert!(mysql.preserves_tdsql_proxy_directive); + assert!(mysql.preserves_tdsql_leading_directives); assert!(SqlDialectProfile::for_database_type(DatabaseType::Doris).supports_hash_line_comments); - assert!(!SqlDialectProfile::for_database_type(DatabaseType::Doris).preserves_tdsql_proxy_directive); + assert!(!SqlDialectProfile::for_database_type(DatabaseType::Doris).preserves_tdsql_leading_directives); assert!(SqlDialectProfile::for_database_type(DatabaseType::StarRocks).supports_hash_line_comments); assert!(SqlDialectProfile::for_database_type(DatabaseType::ManticoreSearch).supports_hash_line_comments); assert!(SqlDialectProfile::for_database_type(DatabaseType::Goldendb).supports_hash_line_comments); @@ -4292,6 +4291,40 @@ delimiter ;"; assert!(split_sql_statements_for_database("/*proxy*/", DatabaseType::Mysql).is_empty()); } + #[test] + fn mysql_current_statement_preserves_same_line_tdsql_directives_without_an_allowlist() { + for directive in ["/*sets:allsets */", "/*master*/", "/*slave:set_1781591902_7*/", "/*future-route:anywhere*/"] + { + let sql = format!("{directive} SELECT count(*) FROM tenant_table"); + let cursor = sql[..sql.find("tenant_table").unwrap()].encode_utf16().count(); + + assert_eq!(find_statement_at_cursor_for_database(&sql, cursor, DatabaseType::Mysql), sql); + assert_eq!(split_sql_statements_for_database(&sql, DatabaseType::Mysql), vec![sql]); + } + } + + #[test] + fn mysql_current_statement_strips_generic_tdsql_style_directive_on_a_separate_line() { + let sql = "/*sets:allsets */\nSELECT count(*) FROM tenant_table"; + let cursor = sql[..sql.find("tenant_table").unwrap()].encode_utf16().count(); + + assert_eq!( + find_statement_at_cursor_for_database(sql, cursor, DatabaseType::Mysql), + "SELECT count(*) FROM tenant_table" + ); + } + + #[test] + fn other_databases_strip_same_line_tdsql_style_directives() { + let sql = "/*sets:allsets */ SELECT count(*) FROM tenant_table"; + let cursor = sql[..sql.find("tenant_table").unwrap()].encode_utf16().count(); + + assert_eq!( + find_statement_at_cursor_for_database(sql, cursor, DatabaseType::Postgres), + "SELECT count(*) FROM tenant_table" + ); + } + #[test] fn mysql_current_statement_strips_other_leading_block_comments() { for comment in ["/* ordinary */", "/*unknown*/", "/* proxy */", "/*PROXY*/"] {