fix(mysql): avoid Aliyun ADB timeout for bounded text results

crates/dbx-core/src/db/mysql.rs

@@ -1434,6 +1434,104 @@ fn query_result_row_limit(max_rows: Option<usize>) -> usize {
     max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
 }
 
+fn should_collect_text_result_set(sql: &str, row_limit: usize, max_rows: Option<usize>) -> bool {
+    max_rows.is_some_and(|_| mysql_top_level_limit(sql).is_some_and(|limit| limit <= row_limit))
+}
+
+fn mysql_top_level_limit(sql: &str) -> Option<usize> {
+    let sql = sql.trim().trim_end_matches(';');
+    let bytes = sql.as_bytes();
+    let mut depth = 0usize;
+    let mut i = 0;
+
+    while i < bytes.len() {
+        i = skip_sql_whitespace_and_comments(bytes, i);
+        if i >= bytes.len() {
+            break;
+        }
+
+        let ch = bytes[i];
+        if matches!(ch, b'\'' | b'"' | b'`') {
+            i = skip_mysql_quoted(sql, i, ch);
+            continue;
+        }
+        if ch == b'(' {
+            depth += 1;
+            i += 1;
+            continue;
+        }
+        if ch == b')' {
+            depth = depth.saturating_sub(1);
+            i += 1;
+            continue;
+        }
+        if depth == 0 && mysql_keyword_at(sql, i, "LIMIT") {
+            return parse_mysql_limit_value(sql, i + "LIMIT".len());
+        }
+        i += 1;
+    }
+
+    None
+}
+
+fn parse_mysql_limit_value(sql: &str, start: usize) -> Option<usize> {
+    let bytes = sql.as_bytes();
+    let mut i = skip_sql_whitespace_and_comments(bytes, start);
+    let first = parse_usize_token(sql, &mut i)?;
+    i = skip_sql_whitespace_and_comments(bytes, i);
+
+    if i < bytes.len() && bytes[i] == b',' {
+        i = skip_sql_whitespace_and_comments(bytes, i + 1);
+        return parse_usize_token(sql, &mut i);
+    }
+
+    Some(first)
+}
+
+fn parse_usize_token(sql: &str, i: &mut usize) -> Option<usize> {
+    let bytes = sql.as_bytes();
+    let start = *i;
+    while *i < bytes.len() && bytes[*i].is_ascii_digit() {
+        *i += 1;
+    }
+    if *i == start {
+        return None;
+    }
+    sql[start..*i].parse().ok()
+}
+
+fn mysql_keyword_at(sql: &str, i: usize, keyword: &str) -> bool {
+    let end = i + keyword.len();
+    end <= sql.len()
+        && sql[i..end].eq_ignore_ascii_case(keyword)
+        && (i == 0 || !is_mysql_identifier_byte(sql.as_bytes()[i - 1]))
+        && (end == sql.len() || !is_mysql_identifier_byte(sql.as_bytes()[end]))
+}
+
+fn is_mysql_identifier_byte(byte: u8) -> bool {
+    byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'$')
+}
+
+fn skip_mysql_quoted(sql: &str, start: usize, quote: u8) -> usize {
+    let bytes = sql.as_bytes();
+    let mut i = start + 1;
+    while i < bytes.len() {
+        if bytes[i] == quote {
+            if i + 1 < bytes.len() && bytes[i + 1] == quote {
+                i += 2;
+                continue;
+            }
+            return i + 1;
+        }
+        if quote == b'\'' && bytes[i] == b'\\' {
+            i = (i + 2).min(bytes.len());
+            continue;
+        }
+        i += 1;
+    }
+    bytes.len()
+}
+
 /// Get a connection from the pool with a health check. If the connection is dead
 /// (e.g. after app was backgrounded), it tries again with a fresh connection.
 pub async fn get_conn_with_health_check(pool: &MySqlPool) -> Result<mysql_async::Conn, String> {
@@ -1451,6 +1549,7 @@ async fn execute_result_set_with_text_protocol_on_conn(
     conn: &mut mysql_async::Conn,
     sql: &str,
     row_limit: usize,
+    max_rows: Option<usize>,
     start: Instant,
 ) -> Result<QueryResult, String> {
     let mut result = conn.query_iter(sql).await.map_err(|e| e.to_string())?;
@@ -1458,6 +1557,28 @@ async fn execute_result_set_with_text_protocol_on_conn(
     let column_types: Vec<String> =
         result.columns_ref().iter().map(|c| mysql_column_type_name(c.column_type())).collect();
 
+    if should_collect_text_result_set(sql, row_limit, max_rows) {
+        let rows: Vec<mysql_async::Row> = result.collect_and_drop().await.map_err(|e| e.to_string())?;
+        let truncated = rows.len() > row_limit;
+        let result_rows = rows
+            .iter()
+            .take(row_limit)
+            .map(|row| (0..row.len()).map(|i| mysql_value_to_json(row, i)).collect())
+            .collect();
+
+        return Ok(QueryResult {
+            columns,
+            column_types,
+            column_sortables: vec![],
+            rows: result_rows,
+            affected_rows: 0,
+            execution_time_ms: start.elapsed().as_millis(),
+            truncated,
+            session_id: None,
+            has_more: false,
+        });
+    }
+
     let mut result_rows: Vec<Vec<serde_json::Value>> = Vec::new();
     let mut stream = result
         .stream::<mysql_async::Row>()
@@ -1564,12 +1685,12 @@ pub async fn execute_query_on_conn_with_max_rows(
 
     if is_result_set_query(sql, dialect) {
         if bare || prefers_text_protocol_query(sql, dialect) {
-            execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, start).await
+            execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, max_rows, start).await
         } else {
             match execute_result_set_with_prepared_protocol_on_conn(conn, sql, row_limit, start).await {
                 Ok(result) => Ok(result),
                 Err(err) if mysql_error_should_retry_with_text_protocol(&err) => {
-                    execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, start).await
+                    execute_result_set_with_text_protocol_on_conn(conn, sql, row_limit, max_rows, start).await
                 }
                 Err(err) => Err(err),
             }
@@ -2041,6 +2162,21 @@ mod tests {
         assert!(!prefers_text_protocol_query("UPDATE users SET name = 'Ada' WHERE id = 1", dialect));
     }
 
+    #[test]
+    fn mysql_text_result_sets_use_buffered_collection_for_bounded_page_queries() {
+        assert!(should_collect_text_result_set("SELECT * FROM users LIMIT 100;", 100, Some(100)));
+        assert!(should_collect_text_result_set("SELECT * FROM users ORDER BY id LIMIT 25 OFFSET 50;", 100, Some(100)));
+        assert!(should_collect_text_result_set("SELECT * FROM users LIMIT 20, 50;", 100, Some(100)));
+    }
+
+    #[test]
+    fn mysql_text_result_sets_keep_streaming_when_unbounded_or_too_large() {
+        assert!(!should_collect_text_result_set("SELECT * FROM users", 100, Some(100)));
+        assert!(!should_collect_text_result_set("SELECT * FROM users LIMIT 1000000", 100, Some(100)));
+        assert!(!should_collect_text_result_set("SELECT * FROM users LIMIT 100", 100, None));
+        assert!(!should_collect_text_result_set("SELECT * FROM (SELECT * FROM audit LIMIT 100) t", 100, Some(100)));
+    }
+
     #[test]
     fn mysql_binary_decode_parse_errors_retry_with_text_protocol() {
         assert!(mysql_error_should_retry_with_text_protocol(

crates/dbx-core/tests/live_mysql57.rs

@@ -17,3 +17,19 @@ async fn live_mysql57_text_protocol_select_succeeds() {
     assert_eq!(result.columns, vec!["id", "label"]);
     assert_eq!(result.rows, vec![vec![serde_json::json!("1"), serde_json::json!("mysql57")]]);
 }
+
+#[tokio::test]
+#[ignore = "requires a remote MySQL-compatible endpoint with a limited result-set query"]
+async fn live_mysql_compatible_limited_text_protocol_query_succeeds() {
+    let url = std::env::var("DBX_LIVE_MYSQL_COMPAT_URL").expect("DBX_LIVE_MYSQL_COMPAT_URL");
+    let sql = std::env::var("DBX_LIVE_MYSQL_COMPAT_SQL").expect("DBX_LIVE_MYSQL_COMPAT_SQL");
+
+    let pool = dbx_core::db::mysql::connect(&url, std::time::Duration::from_secs(10)).await.unwrap();
+    let result = dbx_core::db::mysql::execute_query_with_max_rows(&pool, &sql, false, Some(100), Default::default())
+        .await
+        .unwrap();
+
+    assert!(!result.columns.is_empty());
+    assert!(!result.rows.is_empty());
+    assert!(result.rows.len() <= 100);
+}