fix(sql): preserve same-line TDSQL directives

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", () => {

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", () => {

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");

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;

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<usize> {
-    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<usize> {
+    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);
     }
 }

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*/"] {