From 021e1682c3cc728d1e82a39c99cbd4ab5e08878f Mon Sep 17 00:00:00 2001 From: luoianun Date: Wed, 17 Jun 2026 14:53:13 +0800 Subject: [PATCH] fix(mysql): avoid Aliyun ADB timeout for bounded text results --- crates/dbx-core/src/db/mysql.rs | 140 +++++++++++++++++++++++++- crates/dbx-core/tests/live_mysql57.rs | 16 +++ 2 files changed, 154 insertions(+), 2 deletions(-) diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 6a3ea6846..7e5e5989a 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -1434,6 +1434,104 @@ fn query_result_row_limit(max_rows: Option) -> 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) -> 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 { + 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 { + 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 { + 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 { @@ -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, start: Instant, ) -> Result { 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 = 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 = 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::new(); let mut stream = result .stream::() @@ -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( diff --git a/crates/dbx-core/tests/live_mysql57.rs b/crates/dbx-core/tests/live_mysql57.rs index 1a8f35473..f05b5a3bf 100644 --- a/crates/dbx-core/tests/live_mysql57.rs +++ b/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); +}