fix(sql): preserve TDSQL proxy directive

Closes #6043

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<usize> {
+    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::<Vec<_>>();
     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";