github/dbx
4
0
Fork 0
mirror of https://github.com/t8y2/dbx.git synced 2026-10-28 00:06:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(grid): respect page size row limits

Browse Source
This commit is contained in:
t8y2
2026-05-19 11:31:38 +08:00
parent 95d70eede6
commit 261d43c1b7
13 changed files with 302 additions and 102 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/dbx-core/src/db/clickhouse_driver.rs
+21 -6
View File
@@ -96,12 +96,17 @@ async fn ch_query_with_limit(
resp.json::<ChJsonResult>().await.map_err(|e| format!("ClickHouse parse error: {e}"))
}
fn limited_query_result(result: ChJsonResult, execution_time_ms: u128) -> QueryResult {
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(MAX_ROWS).max(1)
}
fn limited_query_result(result: ChJsonResult, execution_time_ms: u128, max_rows: Option<usize>) -> QueryResult {
let columns: Vec<String> = result.meta.iter().map(|c| c.name.clone()).collect();
let mut rows = result.data;
let truncated = rows.len() > MAX_ROWS;
let row_limit = query_result_row_limit(max_rows);
let truncated = rows.len() > row_limit;
if truncated {
rows.truncate(MAX_ROWS);
rows.truncate(row_limit);
}
QueryResult { columns, rows, affected_rows: 0, execution_time_ms, truncated, session_id: None, has_more: false }
}
@@ -188,11 +193,21 @@ pub async fn get_columns(client: &ChClient, database: &str, table: &str) -> Resu
}
pub async fn execute_query(client: &ChClient, database: &str, sql: &str) -> Result<QueryResult, String> {
execute_query_with_max_rows(client, database, sql, None).await
}
pub async fn execute_query_with_max_rows(
client: &ChClient,
database: &str,
sql: &str,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH"]) {
let result = ch_query_with_limit(client, sql, Some(database), QueryResultLimit::Limited(MAX_ROWS + 1)).await?;
Ok(limited_query_result(result, start.elapsed().as_millis()))
let result = ch_query_with_limit(client, sql, Some(database), QueryResultLimit::Limited(row_limit + 1)).await?;
Ok(limited_query_result(result, start.elapsed().as_millis(), Some(row_limit)))
} else {
let url = build_query_url(&client.base_url, Some(database), QueryResultLimit::Unlimited);
let req = build_request(client, client.http.post(&url).body(sql.to_string()));
@@ -239,7 +254,7 @@ mod tests {
rows: crate::query::MAX_ROWS + 1,
};
let result = limited_query_result(result, 12);
let result = limited_query_result(result, 12, None);
assert_eq!(result.columns, vec!["id"]);
assert_eq!(result.rows.len(), crate::query::MAX_ROWS);
crates/dbx-core/src/db/mysql.rs
+20 -6
View File
@@ -314,8 +314,22 @@ pub async fn get_columns(pool: &MySqlPool, database: &str, table: &str) -> Resul
.collect())
}
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
}
pub async fn execute_query(pool: &MySqlPool, sql: &str, bare: bool) -> Result<QueryResult, String> {
execute_query_with_max_rows(pool, sql, bare, None).await
}
pub async fn execute_query_with_max_rows(
pool: &MySqlPool,
sql: &str,
bare: bool,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
if is_result_set_query(sql) {
if bare || requires_text_protocol_query(sql) {
@@ -335,14 +349,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 {
@@ -369,14 +383,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 {
crates/dbx-core/src/db/postgres.rs
+29 -6
View File
@@ -333,8 +333,21 @@ pub async fn get_columns(pool: &PgPool, schema: &str, table: &str) -> Result<Vec
.collect())
}
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
}
pub async fn execute_query(pool: &PgPool, sql: &str) -> Result<QueryResult, String> {
execute_query_with_max_rows(pool, sql, None).await
}
pub async fn execute_query_with_max_rows(
pool: &PgPool,
sql: &str,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) {
let mut stream = sqlx::query(sql).persistent(false).fetch(pool);
@@ -354,7 +367,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;
}
}
@@ -364,9 +377,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 {
@@ -394,11 +407,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_max_rows(pool, schema, sql, None).await
}
pub async fn execute_query_with_schema_and_max_rows(
pool: &PgPool,
schema: &str,
sql: &str,
max_rows: Option<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 = query_result_row_limit(max_rows);
if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"]) {
let mut stream = sqlx::query(sql).persistent(false).fetch(&mut *conn);
@@ -418,7 +441,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;
}
}
@@ -429,9 +452,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 {
crates/dbx-core/src/db/sqlite.rs
+16 -3
View File
@@ -162,7 +162,20 @@ pub async fn list_triggers(pool: &SqlitePool, _schema: &str, table: &str) -> Res
}
pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result<QueryResult, String> {
execute_query_with_max_rows(pool, sql, None).await
}
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(crate::query::MAX_ROWS).max(1)
}
pub async fn execute_query_with_max_rows(
pool: &SqlitePool,
sql: &str,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
let row_limit = query_result_row_limit(max_rows);
if starts_with_executable_sql_keyword(sql, &["SELECT", "PRAGMA", "EXPLAIN", "WITH"]) {
let desc = pool.describe(sql).await.map_err(|e| e.to_string())?;
@@ -191,14 +204,14 @@ pub async fn execute_query(pool: &SqlitePool, sql: &str) -> Result<QueryResult,
})
.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 {
crates/dbx-core/src/db/sqlserver.rs
+37 -7
View File
@@ -13,6 +13,10 @@ use crate::types::{ColumnInfo, DatabaseInfo, ForeignKeyInfo, IndexInfo, QueryRes
pub type SqlServerClient = Client<Compat<TcpStream>>;
const SIMPLE_QUERY_MODULE_KEYWORDS: &[&str] = &["FUNCTION", "PROC", "PROCEDURE", "TRIGGER", "VIEW"];
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(MAX_ROWS).max(1)
}
pub async fn connect(
host: &str,
port: u16,
@@ -64,7 +68,12 @@ fn columns_from_metadata(metadata: &tiberius::ResultMetadata) -> Vec<String> {
metadata.columns().iter().map(|c| c.name().to_string()).collect()
}
async fn collect_first_result_limited(mut stream: QueryStream<'_>, start: Instant) -> Result<QueryResult, String> {
async fn collect_first_result_limited(
mut stream: QueryStream<'_>,
start: Instant,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let row_limit = query_result_row_limit(max_rows);
let mut columns: Vec<String> = vec![];
let mut rows: Vec<Vec<serde_json::Value>> = Vec::new();
let mut truncated = false;
@@ -76,7 +85,7 @@ async fn collect_first_result_limited(mut stream: QueryStream<'_>, start: Instan
}
QueryItem::Metadata(_) => {}
QueryItem::Row(row) if row.result_index() == 0 => {
if rows.len() < MAX_ROWS {
if rows.len() < row_limit {
rows.push(row_to_json(&row));
} else {
truncated = true;
@@ -120,7 +129,12 @@ fn push_sqlserver_result_set(results: &mut Vec<QueryResult>, result: Option<SqlS
}
}
async fn collect_result_sets_limited(mut stream: QueryStream<'_>, start: Instant) -> Result<Vec<QueryResult>, String> {
async fn collect_result_sets_limited(
mut stream: QueryStream<'_>,
start: Instant,
max_rows: Option<usize>,
) -> Result<Vec<QueryResult>, String> {
let row_limit = query_result_row_limit(max_rows);
let mut results = Vec::new();
let mut current: Option<SqlServerResultSet> = None;
@@ -140,7 +154,7 @@ async fn collect_result_sets_limited(mut stream: QueryStream<'_>, start: Instant
rows: Vec::new(),
truncated: false,
});
if result.rows.len() < MAX_ROWS {
if result.rows.len() < row_limit {
result.rows.push(row_to_json(&row));
} else {
result.truncated = true;
@@ -527,14 +541,22 @@ pub async fn list_triggers(
}
pub async fn execute_query(client: &mut SqlServerClient, sql: &str) -> Result<QueryResult, String> {
execute_query_with_max_rows(client, sql, None).await
}
pub async fn execute_query_with_max_rows(
client: &mut SqlServerClient,
sql: &str,
max_rows: Option<usize>,
) -> Result<QueryResult, String> {
let start = Instant::now();
if starts_with_executable_sql_keyword(sql, &["SELECT", "EXEC", "WITH", "TABLE"]) {
let stream = client.query(sql, &[]).await.map_err(|e| e.to_string())?;
collect_first_result_limited(stream, start).await
collect_first_result_limited(stream, start, max_rows).await
} else if requires_simple_query_batch(sql) {
let stream = client.simple_query(sql).await.map_err(|e| e.to_string())?;
let _ = collect_result_sets_limited(stream, start).await?;
let _ = collect_result_sets_limited(stream, start, max_rows).await?;
Ok(QueryResult {
columns: vec![],
rows: vec![],
@@ -559,9 +581,17 @@ pub async fn execute_query(client: &mut SqlServerClient, sql: &str) -> Result<Qu
}
pub async fn execute_batch(client: &mut SqlServerClient, sql: &str) -> Result<Vec<QueryResult>, String> {
execute_batch_with_max_rows(client, sql, None).await
}
pub async fn execute_batch_with_max_rows(
client: &mut SqlServerClient,
sql: &str,
max_rows: Option<usize>,
) -> Result<Vec<QueryResult>, String> {
let start = Instant::now();
let stream = client.simple_query(sql).await.map_err(|e| e.to_string())?;
let mut results = collect_result_sets_limited(stream, start).await?;
let mut results = collect_result_sets_limited(stream, start, max_rows).await?;
if results.is_empty() {
results.push(QueryResult {
crates/dbx-core/src/query.rs
+66 -23
View File
@@ -19,8 +19,21 @@ pub struct QueryExecutionOptions {
pub result_session_id: Option<String>,
}
fn query_result_row_limit(max_rows: Option<usize>) -> usize {
max_rows.unwrap_or(MAX_ROWS).max(1)
}
pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result<db::QueryResult, String> {
duckdb_execute_with_max_rows(con, sql, None)
}
pub fn duckdb_execute_with_max_rows(
con: &duckdb::Connection,
sql: &str,
max_rows: Option<usize>,
) -> Result<db::QueryResult, String> {
let start = std::time::Instant::now();
let row_limit = query_result_row_limit(max_rows);
if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH", "PRAGMA"]) {
let mut stmt = con.prepare(sql).map_err(|e| e.to_string())?;
@@ -33,9 +46,6 @@ pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result<db::QueryRe
let mut result_rows = Vec::new();
while let Some(row) = rows.next().map_err(|e| e.to_string())? {
if result_rows.len() >= MAX_ROWS {
break;
}
let vals: Vec<serde_json::Value> = (0..col_count)
.map(|i| {
row.get::<_, String>(i)
@@ -53,9 +63,15 @@ pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result<db::QueryRe
})
.collect();
result_rows.push(vals);
if result_rows.len() > row_limit {
break;
}
}
let truncated = result_rows.len() >= MAX_ROWS;
let truncated = result_rows.len() > row_limit;
if truncated {
result_rows.truncate(row_limit);
}
Ok(db::QueryResult {
columns,
rows: result_rows,
@@ -84,6 +100,7 @@ fn duckdb_execute_for_database(
attached_names: &[String],
database: Option<&str>,
sql: &str,
max_rows: Option<usize>,
) -> Result<db::QueryResult, String> {
if let Some(database) = database.map(str::trim).filter(|database| !database.is_empty()) {
let catalog = if database == "main" {
@@ -93,16 +110,21 @@ fn duckdb_execute_for_database(
};
con.execute_batch(&format!("USE {}", duckdb_quote_ident(&catalog))).map_err(|e| e.to_string())?;
}
duckdb_execute(con, sql)
duckdb_execute_with_max_rows(con, sql, max_rows)
}
fn duckdb_quote_ident(value: &str) -> String {
format!("\"{}\"", value.replace('"', "\"\""))
}
pub fn truncate_result(mut result: db::QueryResult) -> db::QueryResult {
if result.rows.len() > MAX_ROWS {
result.rows.truncate(MAX_ROWS);
pub fn truncate_result(result: db::QueryResult) -> db::QueryResult {
truncate_result_with_max_rows(result, None)
}
pub fn truncate_result_with_max_rows(mut result: db::QueryResult, max_rows: Option<usize>) -> db::QueryResult {
let row_limit = query_result_row_limit(max_rows);
if result.rows.len() > row_limit {
result.rows.truncate(row_limit);
result.truncated = true;
}
result
@@ -243,11 +265,12 @@ pub async fn do_execute(
let sql = sql.to_string();
let database = database.map(str::to_string);
let attached_names = duckdb_attached_names;
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, async move {
let task = tokio::task::spawn_blocking(move || {
let con = con.lock().map_err(|e| e.to_string())?;
duckdb_execute_for_database(&con, &attached_names, database.as_deref(), &sql)
duckdb_execute_for_database(&con, &attached_names, database.as_deref(), &sql, max_rows)
});
task.await.map_err(|e| e.to_string())?
})
@@ -256,34 +279,46 @@ pub async fn do_execute(
PoolKind::Mysql(p, mode) => {
let p = p.clone();
let bare = *mode == crate::connection::MysqlMode::Bare;
let max_rows = options.max_rows;
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_max_rows(&p, sql, bare, max_rows)).await
}
PoolKind::Postgres(p) => {
let p = p.clone();
let schema = schema.map(|s| s.to_string());
let max_rows = options.max_rows;
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_max_rows(&p, &schema, sql, max_rows),
)
.await
} else {
wait_for_query(cancel_token, db::postgres::execute_query(&p, sql)).await
wait_for_query(cancel_token, db::postgres::execute_query_with_max_rows(&p, sql, max_rows)).await
}
}
PoolKind::Sqlite(p) => {
let p = p.clone();
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, db::sqlite::execute_query(&p, sql)).await
wait_for_query(cancel_token, db::sqlite::execute_query_with_max_rows(&p, sql, max_rows)).await
}
PoolKind::ClickHouse(client) => {
let client = client.clone();
let database = pool_key.split(':').nth(1).unwrap_or("default").to_string();
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, db::clickhouse_driver::execute_query(&client, &database, sql))
.await
.map(truncate_result)
wait_for_query(
cancel_token,
db::clickhouse_driver::execute_query_with_max_rows(&client, &database, sql, max_rows),
)
.await
.map(|result| truncate_result_with_max_rows(result, max_rows))
}
PoolKind::SqlServer(client) => {
let client = client.clone();
let max_rows = options.max_rows;
drop(connections);
let mut client = match cancel_token.as_ref() {
Some(token) => tokio::select! {
@@ -293,15 +328,18 @@ pub async fn do_execute(
},
None => client.lock().await,
};
wait_for_query(cancel_token, db::sqlserver::execute_query(&mut client, sql)).await.map(truncate_result)
wait_for_query(cancel_token, db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows))
.await
.map(|result| truncate_result_with_max_rows(result, max_rows))
}
PoolKind::Elasticsearch(client) => {
let client = client.clone();
let sql = sql.to_string();
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, db::elasticsearch_driver::execute_rest_query(&client, &sql))
.await
.map(truncate_result)
.map(|result| truncate_result_with_max_rows(result, max_rows))
}
PoolKind::Redis(_) => Err("Use Redis-specific commands".to_string()),
PoolKind::MongoDb(_) => Err("Use MongoDB-specific commands".to_string()),
@@ -309,6 +347,7 @@ pub async fn do_execute(
let client = client.clone();
let sql = sql.to_string();
let schema = schema.map(|s| s.to_string());
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, async move {
let mut client = client.lock().await;
@@ -324,7 +363,7 @@ pub async fn do_execute(
}
})
.await
.map(truncate_result)
.map(|result| truncate_result_with_max_rows(result, max_rows))
}
PoolKind::ExternalTabular(ext_pool) => {
if !starts_with_executable_sql_keyword(sql, &["SELECT", "WITH", "SHOW", "DESCRIBE", "EXPLAIN", "PRAGMA"]) {
@@ -332,11 +371,12 @@ pub async fn do_execute(
}
let con = ext_pool.cache.clone();
let sql = sql.to_string();
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, async move {
let task = tokio::task::spawn_blocking(move || {
let con = con.lock().map_err(|e| e.to_string())?;
duckdb_execute(&con, &sql)
duckdb_execute_with_max_rows(&con, &sql, max_rows)
});
task.await.map_err(|e| e.to_string())?
})
@@ -348,13 +388,14 @@ pub async fn do_execute(
let sql = sql.to_string();
let schema = schema.map(str::to_string);
let database = config.effective_database().unwrap_or("").to_string();
let max_rows = options.max_rows;
drop(connections);
wait_for_query(cancel_token, async move {
let params = external_driver_query_params(&config, &sql, &database, schema.as_deref());
session.invoke::<db::QueryResult>("executeQuery", params).await
})
.await
.map(truncate_result)
.map(|result| truncate_result_with_max_rows(result, max_rows))
}
}
}
@@ -490,7 +531,7 @@ pub async fn execute_multi_core_with_options(
};
if is_sqlserver {
return execute_multi_sqlserver(state, &pool_key, sql, cancel_token).await;
return execute_multi_sqlserver(state, &pool_key, sql, cancel_token, options).await;
}
let statements = split_sql_statements(sql);
@@ -547,9 +588,11 @@ async fn execute_multi_sqlserver(
pool_key: &str,
sql: &str,
cancel_token: Option<CancellationToken>,
options: QueryExecutionOptions,
) -> Result<Vec<db::QueryResult>, String> {
let batches = split_sql_batches(sql);
let mut all_results = Vec::new();
let max_rows = options.max_rows;
for batch in &batches {
if is_canceled(&cancel_token) {
@@ -582,7 +625,7 @@ async fn execute_multi_sqlserver(
None => client.lock().await,
};
match db::sqlserver::execute_batch(&mut client, batch).await {
match db::sqlserver::execute_batch_with_max_rows(&mut client, batch, max_rows).await {
Ok(results) => all_results.extend(results),
Err(e) => {
all_results.push(db::QueryResult {