From 8a81bdfdd56bb4fa1688f694aad9590eb25082b8 Mon Sep 17 00:00:00 2001 From: Illuminated2020 <2357303264@qq.com> Date: Sun, 10 May 2026 20:00:52 +0800 Subject: [PATCH] =?UTF-8?q?feat(core):=20=E6=94=AF=E6=8C=81=20MySQL=20?= =?UTF-8?q?=E5=92=8C=20Postgres=20=E6=9F=A5=E8=AF=A2=E8=A1=8C=E6=95=B0?= =?UTF-8?q?=E9=99=90=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 MySQL/Postgres limit-aware 查询执行函数 - 保持旧执行 API 默认转调 MAX_ROWS - 让 query.rs 对 MySQL/Postgres 传递 row_limit - 补充 limit-aware 驱动调用回归测试 --- crates/dbx-core/src/db/mysql.rs | 27 ++++++++++++++----- crates/dbx-core/src/db/postgres.rs | 42 +++++++++++++++++++++++++----- crates/dbx-core/src/query.rs | 19 +++++++++++--- 3 files changed, 73 insertions(+), 15 deletions(-) diff --git a/crates/dbx-core/src/db/mysql.rs b/crates/dbx-core/src/db/mysql.rs index 5946f0436e..3fa8585919 100644 --- a/crates/dbx-core/src/db/mysql.rs +++ b/crates/dbx-core/src/db/mysql.rs @@ -249,7 +249,17 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul } pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result { + execute_query_with_row_limit(pool, sql, bare, crate::query::MAX_ROWS).await +} + +pub async fn execute_query_with_row_limit( + pool: &MySqlPool, + sql: &str, + bare: bool, + row_limit: usize, +) -> Result { let start = Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN"]) { if bare { @@ -269,14 +279,14 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -301,14 +311,14 @@ pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -413,6 +423,11 @@ pub async fn list_triggers(pool: &MySqlPool, database: &str, table: &str) -> Res mod tests { use super::*; + #[allow(dead_code)] + fn mysql_limit_aware_execute_query_api_compiles(pool: &MySqlPool, sql: &str, bare: bool) { + let _ = execute_query_with_row_limit(pool, sql, bare, 7); + } + #[test] fn numeric_metadata_accepts_unsigned_information_schema_values() { assert_eq!(numeric_metadata_u64_to_i32(Some(65)), Some(65)); diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index bec71c4fda..c7891fe165 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -281,7 +281,12 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result Result { + execute_query_with_row_limit(pool, sql, crate::query::MAX_ROWS).await +} + +pub async fn execute_query_with_row_limit(pool: &PgPool, sql: &str, row_limit: usize) -> Result { let start = Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(pool); @@ -301,7 +306,7 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } @@ -311,9 +316,9 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -337,11 +342,21 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result Result { + execute_query_with_schema_and_row_limit(pool, schema, sql, crate::query::MAX_ROWS).await +} + +pub async fn execute_query_with_schema_and_row_limit( + pool: &PgPool, + schema: &str, + sql: &str, + row_limit: usize, +) -> Result { let mut conn = pool.acquire().await.map_err(|e| e.to_string())?; let set_path = format!("SET search_path TO \"{}\", public", schema); sqlx::query(&set_path).execute(&mut *conn).await.map_err(|e| e.to_string())?; let start = Instant::now(); + let row_limit = row_limit.max(1); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) { let mut stream = sqlx::query(sql).persistent(false).fetch(&mut *conn); @@ -361,7 +376,7 @@ pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) - .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) .collect(), ); - if result_rows.len() > crate::query::MAX_ROWS { + if result_rows.len() > row_limit { break; } } @@ -372,9 +387,9 @@ pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) - columns = desc.columns().iter().map(|c| c.name().to_string()).collect(); } - let truncated = result_rows.len() > crate::query::MAX_ROWS; + let truncated = result_rows.len() > row_limit; if truncated { - result_rows.truncate(crate::query::MAX_ROWS); + result_rows.truncate(row_limit); } Ok(QueryResult { @@ -446,6 +461,21 @@ pub async fn list_indexes(pool: &PgPool, schema: &str, table: &str) -> Result Result, String> { let rows: Vec = sqlx::query( "SELECT kcu.constraint_name, kcu.column_name, \ diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index 41663707ec..d007e9403a 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -192,16 +192,20 @@ pub async fn do_execute_with_row_limit( let p = p.clone(); let bare = *mode == crate::connection::MysqlMode::Bare; drop(connections); - wait_for_query(cancel_token, db::mysql::execute_query(&p, sql, bare)).await + wait_for_query(cancel_token, db::mysql::execute_query_with_row_limit(&p, sql, bare, row_limit)).await } PoolKind::Postgres(p) => { let p = p.clone(); let schema = schema.map(|s| s.to_string()); drop(connections); if let Some(schema) = schema { - wait_for_query(cancel_token, db::postgres::execute_query_with_schema(&p, &schema, sql)).await + wait_for_query( + cancel_token, + db::postgres::execute_query_with_schema_and_row_limit(&p, &schema, sql, row_limit), + ) + .await } else { - wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)).await + wait_for_query(cancel_token, db::postgres::execute_query_with_row_limit(&p, sql, row_limit)).await } } PoolKind::Sqlite(p) => { @@ -834,6 +838,15 @@ mod tests { assert!(result.truncated); } + #[test] + fn mysql_and_postgres_pool_branches_call_limit_aware_drivers() { + let source = include_str!("query.rs"); + + assert!(source.contains("db::mysql::execute_query_with_row_limit(&p, sql, bare, row_limit)")); + assert!(source.contains("db::postgres::execute_query_with_schema_and_row_limit(&p, &schema, sql, row_limit)")); + assert!(source.contains("db::postgres::execute_query_with_row_limit(&p, sql, row_limit)")); + } + #[tokio::test] async fn wait_for_query_returns_cancelled_when_token_is_cancelled() { let token = CancellationToken::new();