diff --git a/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts b/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts index 54b7b5eb04..bf27f299bd 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlExecutionTarget.spec.ts @@ -13,6 +13,12 @@ describe("resolveExecutableSql", () => { expect(resolveExecutableSql(sql, selectedSql, { mode: "current", cursorPos: cursorAfterFirstSemicolon })).toBe("select 2;"); }); + + it("keeps a manually selected proxy directive unchanged", () => { + const selectedSql = "/*proxy*/\nSHOW PROXY STATUS"; + + 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 faaf68f6e9..85d48b7101 100644 --- a/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts +++ b/apps/desktop/src/lib/__tests__/sql/sqlStatementRanges.spec.ts @@ -1199,6 +1199,44 @@ describe("buildExecutionCandidates", () => { expect(splitSqlStatementRanges("/*&tenant:mctest*/", "mysql")).toEqual([]); }); + it("preserves only the exact TDSQL proxy directive for the current statement", () => { + const directedSql = "/*proxy*/ \n\tSHOW PROXY STATUS"; + const sql = `SELECT 1;\n${directedSql};\nSELECT 2;`; + const candidates = buildExecutionCandidates(sql, indexOf(sql, "STATUS"), "mysql"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe(directedSql); + expect(executionCandidateForMode(candidates, "all")?.sql).toBe(sql); + expect(rangeSqlTexts(executableStatementRanges(sql, "mysql"))).toEqual(["SELECT 1", directedSql, "SELECT 2"]); + expect(splitSqlStatementRanges("/*proxy*/SHOW PROXY STATUS", "mysql")[0]?.sql).toBe("/*proxy*/SHOW PROXY STATUS"); + expect(splitSqlStatementRanges("/* ordinary */\n/*proxy*/SHOW PROXY STATUS", "mysql")[0]?.sql).toBe("/*proxy*/SHOW PROXY STATUS"); + expect(splitSqlStatementRanges("/*proxy*/", "mysql")).toEqual([]); + }); + + 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"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe("SHOW PROXY STATUS"); + }); + + it.each(["/*proxy*/\n/* audit */\nSHOW PROXY STATUS", "/*proxy*/\n-- audit\nSHOW PROXY STATUS"])("does not preserve a proxy directive separated from SQL by another comment", (sql) => { + const candidates = buildExecutionCandidates(sql, indexOf(sql, "STATUS"), "mysql"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe("SHOW PROXY STATUS"); + }); + + it("does not preserve the proxy directive for other database types", () => { + const sql = "/*proxy*/\nSHOW PROXY STATUS"; + const candidates = buildExecutionCandidates(sql, indexOf(sql, "STATUS"), "postgres"); + + expect(executionCandidateForMode(candidates, "current")?.sql).toBe("SHOW PROXY STATUS"); + }); + + it("keeps existing hints without allowing them between the proxy directive and SQL", () => { + expect(splitSqlStatementRanges("/*+ route */\n/*proxy*/SHOW PROXY STATUS", "mysql")[0]?.sql).toBe("/*proxy*/SHOW PROXY STATUS"); + expect(splitSqlStatementRanges("/*proxy*/\n/*+ route */\nSHOW PROXY STATUS", "mysql")[0]?.sql).toBe("/*+ route */\nSHOW PROXY STATUS"); + }); + it("uses the cursor statement for the first candidate when there is no selection", () => { const sql = "SELECT *\nFROM users\nWHERE active = 1"; const candidates = buildExecutionCandidates(sql, indexOf(sql, "users")); diff --git a/apps/desktop/src/lib/sql/sqlStatementRanges.ts b/apps/desktop/src/lib/sql/sqlStatementRanges.ts index 74f7bece5a..7d4969640c 100644 --- a/apps/desktop/src/lib/sql/sqlStatementRanges.ts +++ b/apps/desktop/src/lib/sql/sqlStatementRanges.ts @@ -297,6 +297,7 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType let statementEnd = -1; let statementHitStart = 0; let pendingHintStart = -1; + let pendingProxyDirectiveStart = -1; let customDelimiter: string | null = null; let state: QuoteState = "none"; let dollarTag = ""; @@ -308,8 +309,9 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType const markContent = (pos: number) => { if (statementStart === -1) { - statementStart = pendingHintStart === -1 ? pos : pendingHintStart; + statementStart = pendingProxyDirectiveStart !== -1 ? pendingProxyDirectiveStart : pendingHintStart === -1 ? pos : pendingHintStart; pendingHintStart = -1; + pendingProxyDirectiveStart = -1; } statementEnd = pos + 1; }; @@ -318,6 +320,7 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType if (statementStart === -1) { statementEnd = -1; pendingHintStart = -1; + pendingProxyDirectiveStart = -1; postgresDollarQuotedRoutine = false; oraclePlSqlStatementEnd = undefined; return; @@ -329,6 +332,7 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType statementStart = -1; statementEnd = -1; pendingHintStart = -1; + pendingProxyDirectiveStart = -1; postgresDollarQuotedRoutine = false; oraclePlSqlStatementEnd = undefined; }; @@ -433,11 +437,13 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType // Line comments consume up to (and including) the newline. if (ch === "-" && next === "-") { + pendingProxyDirectiveStart = -1; const newline = sql.indexOf("\n", i); i = newline === -1 ? len : newline + 1; continue; } if (startsHashLineComment(sql, i, databaseType, parameterOptions)) { + pendingProxyDirectiveStart = -1; const newline = sql.indexOf("\n", i); i = newline === -1 ? len : newline + 1; continue; @@ -447,6 +453,9 @@ export function splitSqlStatementRanges(sql: string, databaseType?: DatabaseType const hintMarker = sql[i + 2]; 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; + } i = close === -1 ? len : close + 2; continue; } diff --git a/crates/dbx-core/src/sql.rs b/crates/dbx-core/src/sql.rs index 63fa54108f..30be134951 100644 --- a/crates/dbx-core/src/sql.rs +++ b/crates/dbx-core/src/sql.rs @@ -117,6 +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, } impl Default for SqlDialectProfile { @@ -132,12 +133,17 @@ 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, } } } impl SqlDialectProfile { fn for_database_type(db_type: DatabaseType) -> Self { + if matches!(db_type, DatabaseType::Mysql) { + return Self::mysql(); + } + if matches!(db_type, DatabaseType::Gaussdb) { return Self::gaussdb(); } @@ -165,6 +171,10 @@ impl SqlDialectProfile { Self { supports_hash_line_comments: true, supports_mysql_routine_blocks: true, ..Self::default() } } + fn mysql() -> Self { + Self { preserves_tdsql_proxy_directive: true, ..Self::mysql_compatible() } + } + fn oracle_like() -> Self { Self { supports_oracle_plsql_blocks: true, supports_slash_line_block_delimiter: true, ..Self::default() } } @@ -2463,10 +2473,20 @@ fn executable_sql_bounds(statement: &str, options: SqlParsingOptions) -> Option< if executable.is_empty() { return None; } - let start = trimmed.len() - executable.len(); + let executable_start = trimmed.len() - executable.len(); + let start = if options.profile.preserves_tdsql_proxy_directive { + tdsql_proxy_directive_start(trimmed, executable_start).unwrap_or(executable_start) + } else { + executable_start + }; Some((start, trimmed_end)) } +fn tdsql_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()) +} + fn has_executable_sql_with_options(statement: &str, options: SqlParsingOptions) -> bool { let chars = statement.chars().collect::>(); let mut in_line_comment = false; @@ -3190,12 +3210,15 @@ 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); let mysql = SqlDialectProfile::for_database_type(DatabaseType::Mysql); - assert_eq!(mysql, SqlDialectProfile::mysql_compatible()); + assert_eq!(mysql, SqlDialectProfile::mysql()); assert!(mysql.supports_hash_line_comments); assert!(mysql.supports_mysql_routine_blocks); + assert!(mysql.preserves_tdsql_proxy_directive); 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::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); @@ -4069,6 +4092,65 @@ delimiter ;"; ); } + #[test] + fn mysql_current_statement_preserves_only_exact_tdsql_proxy_directive() { + let sql = "SELECT 1;\n/*proxy*/ \n\tSHOW PROXY STATUS;\nSELECT 2;"; + let cursor = sql[..sql.find("STATUS").unwrap()].encode_utf16().count(); + + assert_eq!( + find_statement_at_cursor_for_database(sql, cursor, DatabaseType::Mysql), + "/*proxy*/ \n\tSHOW PROXY STATUS" + ); + assert_eq!( + split_sql_statements_for_database(sql, DatabaseType::Mysql), + vec!["SELECT 1", "/*proxy*/ \n\tSHOW PROXY STATUS", "SELECT 2"] + ); + for inline in [ + "/*proxy*/SHOW PROXY STATUS", + "/* ordinary */\n/*proxy*/SHOW PROXY STATUS", + "/*+ route */\n/*proxy*/SHOW PROXY STATUS", + ] { + let cursor = inline[..inline.find("STATUS").unwrap()].encode_utf16().count(); + assert_eq!( + find_statement_at_cursor_for_database(inline, cursor, DatabaseType::Mysql), + "/*proxy*/SHOW PROXY STATUS" + ); + } + assert!(split_sql_statements_for_database("/*proxy*/", DatabaseType::Mysql).is_empty()); + } + + #[test] + fn mysql_current_statement_strips_other_leading_block_comments() { + for comment in ["/* ordinary */", "/*unknown*/", "/* proxy */", "/*PROXY*/"] { + let sql = format!("{comment}\nSHOW PROXY STATUS"); + let cursor = sql[..sql.find("STATUS").unwrap()].encode_utf16().count(); + + assert_eq!(find_statement_at_cursor_for_database(&sql, cursor, DatabaseType::Mysql), "SHOW PROXY STATUS"); + } + + for separated in ["/*proxy*/\n/* audit */\nSHOW PROXY STATUS", "/*proxy*/\n-- audit\nSHOW PROXY STATUS"] { + let cursor = separated[..separated.find("STATUS").unwrap()].encode_utf16().count(); + assert_eq!( + find_statement_at_cursor_for_database(separated, cursor, DatabaseType::Mysql), + "SHOW PROXY STATUS" + ); + } + + let postgres = "/*proxy*/\nSHOW PROXY STATUS"; + let cursor = postgres[..postgres.find("STATUS").unwrap()].encode_utf16().count(); + assert_eq!( + find_statement_at_cursor_for_database(postgres, cursor, DatabaseType::Postgres), + "SHOW PROXY STATUS" + ); + + let followed_by_hint = "/*proxy*/\n/*+ route */\nSHOW PROXY STATUS"; + let cursor = followed_by_hint[..followed_by_hint.find("STATUS").unwrap()].encode_utf16().count(); + assert_eq!( + find_statement_at_cursor_for_database(followed_by_hint, cursor, DatabaseType::Mysql), + "SHOW PROXY STATUS" + ); + } + #[test] fn mysql_single_statement_with_inline_comment_stays_executable_per_issue_428() { let sql = "SELECT 1 # mysql comment";