github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-21 19:23:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(core): 支持 MySQL 和 Postgres 查询行数限制

Browse Source 
- 新增 MySQL/Postgres limit-aware 查询执行函数

- 保持旧执行 API 默认转调 MAX_ROWS

- 让 query.rs 对 MySQL/Postgres 传递 row_limit

- 补充 limit-aware 驱动调用回归测试
This commit is contained in:
Illuminated2020
2026-05-12 16:49:29 +08:00
parent 4267c6f1e0
commit 8a81bdfdd5
3 changed files with 73 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/mysql.rs
+21 -6
View File
@@ -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<QueryResult, String> {
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<QueryResult, String> {
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<Qu
.map(|i| mysql_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;
}
}
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<Qu
.map(|i| mysql_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;
}
}
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));
crates/dbx-core/src/db/postgres.rs
+36 -6
View File
@@ -281,7 +281,12 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result<Vec
}
pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, String> {
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<QueryResult, 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(pool);
@@ -301,7 +306,7 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, Stri
.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;
}
}
@@ -311,9 +316,9 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, Stri
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 {
@@ -337,11 +342,21 @@ pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, Stri
}
pub async fn execute_query_with_schema(pool: &PgPool, schema: &str, sql: &str) -> Result<QueryResult, String> {
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<QueryResult, String> {
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<Ve
.collect())
}
#[cfg(test)]
mod tests {
use super::*;
#[allow(dead_code)]
fn postgres_limit_aware_execute_query_api_compiles(pool: &PgPool, sql: &str) {
let _ = execute_query_with_row_limit(pool, sql, 7);
}
#[allow(dead_code)]
fn postgres_schema_limit_aware_execute_query_api_compiles(pool: &PgPool, schema: &str, sql: &str) {
let _ = execute_query_with_schema_and_row_limit(pool, schema, sql, 7);
}
}
pub async fn list_foreign_keys(pool: &PgPool, schema: &str, table: &str) -> Result<Vec<ForeignKeyInfo>, String> {
let rows: Vec<PgRow> = sqlx::query(
"SELECT kcu.constraint_name, kcu.column_name, \
crates/dbx-core/src/query.rs
+16 -3
View File
@@ -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();