From f4566221b1565a87c48365e056c87ef4a2dced4e Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Sun, 17 May 2026 00:10:18 +0800 Subject: [PATCH] feat(query): add paged result fetching --- crates/dbx-core/src/db/clickhouse_driver.rs | 91 ++++++- .../dbx-core/src/db/elasticsearch_driver.rs | 4 + crates/dbx-core/src/db/mysql.rs | 6 + crates/dbx-core/src/db/postgres.rs | 8 + crates/dbx-core/src/db/sqlite.rs | 4 + crates/dbx-core/src/db/sqlserver.rs | 10 + crates/dbx-core/src/query.rs | 256 +++++++++++++++++- crates/dbx-core/src/transfer.rs | 4 + crates/dbx-core/src/types.rs | 4 + src-tauri/src/commands/query.rs | 43 ++- src-tauri/src/lib.rs | 1 + src-web/src/main.rs | 1 + src-web/src/routes/query.rs | 39 ++- src/components/grid/DataGrid.vue | 47 +++- src/components/layout/ContentArea.vue | 4 + src/components/sidebar/TreeItem.vue | 2 + src/composables/useDataGridActions.ts | 22 +- src/composables/useNavigationTargets.ts | 6 + src/lib/api.ts | 1 + src/lib/http.ts | 10 +- src/lib/queryResultPagination.ts | 136 ++++++++++ src/lib/tauri.ts | 10 +- src/stores/queryStore.ts | 87 +++++- src/types/database.ts | 7 + tests/queryResultPagination.test.ts | 141 ++++++++++ 25 files changed, 885 insertions(+), 59 deletions(-) create mode 100644 src/lib/queryResultPagination.ts create mode 100644 tests/queryResultPagination.test.ts diff --git a/crates/dbx-core/src/db/clickhouse_driver.rs b/crates/dbx-core/src/db/clickhouse_driver.rs index aa603505e..b8ee736fb 100644 --- a/crates/dbx-core/src/db/clickhouse_driver.rs +++ b/crates/dbx-core/src/db/clickhouse_driver.rs @@ -3,6 +3,7 @@ use serde::Deserialize; use std::time::Instant; use super::{connection_timeout, with_connection_timeout}; +use crate::query::MAX_ROWS; use crate::sql::starts_with_executable_sql_keyword; use crate::types::{ColumnInfo, DatabaseInfo, QueryResult, TableInfo}; @@ -48,6 +49,22 @@ struct ChColumn { _type: String, } +enum QueryResultLimit { + Unlimited, + Limited(usize), +} + +fn build_query_url(base_url: &str, database: Option<&str>, limit: QueryResultLimit) -> String { + let mut url = format!("{}/?default_format=JSONCompact", base_url); + if let Some(db) = database { + url.push_str(&format!("&database={db}")); + } + if let QueryResultLimit::Limited(max_rows) = limit { + url.push_str(&format!("&max_result_rows={max_rows}&result_overflow_mode=break")); + } + url +} + fn build_request(client: &ChClient, req: reqwest::RequestBuilder) -> reqwest::RequestBuilder { match (&client.username, &client.password) { (Some(u), Some(p)) if !u.is_empty() => req.basic_auth(u, Some(p)), @@ -57,10 +74,16 @@ fn build_request(client: &ChClient, req: reqwest::RequestBuilder) -> reqwest::Re } async fn ch_query(client: &ChClient, sql: &str, database: Option<&str>) -> Result { - let mut url = format!("{}/?default_format=JSONCompact", client.base_url); - if let Some(db) = database { - url.push_str(&format!("&database={}", db)); - } + ch_query_with_limit(client, sql, database, QueryResultLimit::Unlimited).await +} + +async fn ch_query_with_limit( + client: &ChClient, + sql: &str, + database: Option<&str>, + limit: QueryResultLimit, +) -> Result { + let url = build_query_url(&client.base_url, database, limit); log::info!("[clickhouse] query url={url} user={:?} has_pass={}", client.username, client.password.is_some()); let req = build_request(client, client.http.post(&url).body(sql.to_string())); let resp = req.send().await.map_err(|e| format!("ClickHouse request failed: {e}"))?; @@ -73,6 +96,16 @@ async fn ch_query(client: &ChClient, sql: &str, database: Option<&str>) -> Resul resp.json::().await.map_err(|e| format!("ClickHouse parse error: {e}")) } +fn limited_query_result(result: ChJsonResult, execution_time_ms: u128) -> QueryResult { + let columns: Vec = result.meta.iter().map(|c| c.name.clone()).collect(); + let mut rows = result.data; + let truncated = rows.len() > MAX_ROWS; + if truncated { + rows.truncate(MAX_ROWS); + } + QueryResult { columns, rows, affected_rows: 0, execution_time_ms, truncated, session_id: None, has_more: false } +} + pub async fn test_connection(client: &ChClient) -> Result<(), String> { let url = format!("{}/?query=SELECT%201", client.base_url); let req = build_request(client, client.http.get(&url)); @@ -158,17 +191,10 @@ pub async fn execute_query(client: &ChClient, database: &str, sql: &str) -> Resu let start = Instant::now(); if starts_with_executable_sql_keyword(sql, &["SELECT", "SHOW", "DESCRIBE", "EXPLAIN", "WITH"]) { - let result = ch_query(client, sql, Some(database)).await?; - let columns: Vec = result.meta.iter().map(|c| c.name.clone()).collect(); - Ok(QueryResult { - columns, - rows: result.data, - affected_rows: 0, - execution_time_ms: start.elapsed().as_millis(), - truncated: false, - }) + 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())) } else { - let url = format!("{}/?default_format=JSONCompact&database={}", client.base_url, database); + 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())); let resp = req.send().await.map_err(|e| format!("ClickHouse request failed: {e}"))?; if !resp.status().is_success() { @@ -181,6 +207,43 @@ pub async fn execute_query(client: &ChClient, database: &str, sql: &str) -> Resu affected_rows: 0, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn query_url_for_result_sets_adds_row_limit_break_settings() { + let url = build_query_url( + "http://localhost:8123", + Some("analytics"), + QueryResultLimit::Limited(crate::query::MAX_ROWS + 1), + ); + + assert_eq!( + url, + "http://localhost:8123/?default_format=JSONCompact&database=analytics&max_result_rows=10001&result_overflow_mode=break" + ); + } + + #[test] + fn limited_query_result_truncates_extra_probe_row() { + let result = ChJsonResult { + meta: vec![ChColumn { name: "id".to_string(), _type: "UInt64".to_string() }], + data: (0..=crate::query::MAX_ROWS).map(|value| vec![serde_json::Value::Number(value.into())]).collect(), + rows: crate::query::MAX_ROWS + 1, + }; + + let result = limited_query_result(result, 12); + + assert_eq!(result.columns, vec!["id"]); + assert_eq!(result.rows.len(), crate::query::MAX_ROWS); + assert_eq!(result.execution_time_ms, 12); + assert!(result.truncated); + } +} diff --git a/crates/dbx-core/src/db/elasticsearch_driver.rs b/crates/dbx-core/src/db/elasticsearch_driver.rs index 3a7ccaa82..619309d2b 100644 --- a/crates/dbx-core/src/db/elasticsearch_driver.rs +++ b/crates/dbx-core/src/db/elasticsearch_driver.rs @@ -295,6 +295,8 @@ pub async fn execute_rest_query(client: &EsClient, input: &str) -> Result Result Result Result Result Result Result Result Result Result Result Result(), execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } } @@ -478,6 +484,8 @@ pub async fn execute_batch(client: &mut SqlServerClient, sql: &str) -> Result Result, + pub fetch_size: Option, + pub page_size: Option, + pub result_session_id: Option, +} + pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result { let start = std::time::Instant::now(); @@ -54,6 +62,8 @@ pub fn duckdb_execute(con: &duckdb::Connection, sql: &str) -> Result Result db::QueryResult { result } +pub fn agent_execute_query_params( + sql: &str, + schema: Option<&str>, + options: QueryExecutionOptions, +) -> serde_json::Value { + let mut params = serde_json::json!({ + "sql": sql, + "maxRows": options.max_rows.unwrap_or(MAX_ROWS), + }); + if let Some(schema) = schema { + params["schema"] = serde_json::json!(schema); + } + if let Some(fetch_size) = options.fetch_size { + params["fetchSize"] = serde_json::json!(fetch_size); + } + params +} + +pub fn agent_execute_query_page_params( + sql: &str, + schema: Option<&str>, + options: QueryExecutionOptions, +) -> serde_json::Value { + let mut params = serde_json::json!({ + "sql": sql, + "pageSize": options.page_size.unwrap_or(MAX_ROWS), + "maxRows": options.max_rows.unwrap_or(MAX_ROWS), + }); + if let Some(schema) = schema { + params["schema"] = serde_json::json!(schema); + } + if let Some(fetch_size) = options.fetch_size { + params["fetchSize"] = serde_json::json!(fetch_size); + } + params +} + +pub fn agent_fetch_query_page_params(session_id: &str, page_size: usize) -> serde_json::Value { + serde_json::json!({ + "sessionId": session_id, + "pageSize": page_size, + }) +} + +pub fn agent_close_query_session_params(session_id: &str) -> serde_json::Value { + serde_json::json!({ + "sessionId": session_id, + }) +} + pub fn is_connection_error(err: &str) -> bool { let lower = err.to_lowercase(); lower.contains("connection") @@ -141,6 +203,7 @@ pub async fn do_execute( sql: &str, schema: Option<&str>, cancel_token: Option, + options: QueryExecutionOptions, ) -> Result { let connections = state.connections.read().await; let pool = connections.get(pool_key).ok_or("Connection not found")?; @@ -218,11 +281,16 @@ pub async fn do_execute( drop(connections); wait_for_query(cancel_token, async move { let mut client = client.lock().await; - let params = match schema { - Some(s) => serde_json::json!({"sql": sql, "schema": s}), - None => serde_json::json!({"sql": sql}), - }; - client.call("execute_query", params).await + if let Some(session_id) = options.result_session_id.as_deref() { + let params = agent_fetch_query_page_params(session_id, options.page_size.unwrap_or(MAX_ROWS)); + client.call("fetch_query_page", params).await + } else if options.page_size.is_some() { + let params = agent_execute_query_page_params(&sql, schema.as_deref(), options); + client.call("execute_query_page", params).await + } else { + let params = agent_execute_query_params(&sql, schema.as_deref(), options); + client.call("execute_query", params).await + } }) .await .map(truncate_result) @@ -281,6 +349,27 @@ pub async fn execute_sql_statement( sql: &str, schema: Option<&str>, cancel_token: Option, +) -> Result { + execute_sql_statement_with_options( + state, + connection_id, + database, + sql, + schema, + cancel_token, + QueryExecutionOptions::default(), + ) + .await +} + +pub async fn execute_sql_statement_with_options( + state: &AppState, + connection_id: &str, + database: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + options: QueryExecutionOptions, ) -> Result { let pool_key = if database.is_empty() { connection_id.to_string() @@ -292,18 +381,43 @@ pub async fn execute_sql_statement( return Err(canceled_error()); } - let result = do_execute(state, &pool_key, sql, schema, cancel_token.clone()).await; + let result = do_execute(state, &pool_key, sql, schema, cancel_token.clone(), options.clone()).await; match &result { Err(e) if is_connection_error(e) && !is_canceled(&cancel_token) => { let db_opt = if database.is_empty() { None } else { Some(database) }; let new_key = state.reconnect_pool(connection_id, db_opt).await?; - do_execute(state, &new_key, sql, schema, cancel_token).await + do_execute(state, &new_key, sql, schema, cancel_token, options).await } _ => result, } } +pub async fn close_query_session( + state: &AppState, + connection_id: &str, + database: &str, + session_id: &str, +) -> Result { + let pool_key = if database.is_empty() { + connection_id.to_string() + } else { + state.get_or_create_pool(connection_id, Some(database)).await? + }; + + let connections = state.connections.read().await; + let pool = connections.get(&pool_key).ok_or("Connection not found")?; + match pool { + PoolKind::Agent(client) => { + let client = client.clone(); + drop(connections); + let mut client = client.lock().await; + client.call("close_query_session", agent_close_query_session_params(session_id)).await + } + _ => Ok(false), + } +} + pub async fn execute_multi_core( state: &AppState, connection_id: &str, @@ -311,6 +425,27 @@ pub async fn execute_multi_core( sql: &str, schema: Option<&str>, cancel_token: Option, +) -> Result, String> { + execute_multi_core_with_options( + state, + connection_id, + database, + sql, + schema, + cancel_token, + QueryExecutionOptions::default(), + ) + .await +} + +pub async fn execute_multi_core_with_options( + state: &AppState, + connection_id: &str, + database: &str, + sql: &str, + schema: Option<&str>, + cancel_token: Option, + options: QueryExecutionOptions, ) -> Result, String> { let pool_key = if database.is_empty() { connection_id.to_string() @@ -330,7 +465,16 @@ pub async fn execute_multi_core( let statements = split_sql_statements(sql); if statements.len() <= 1 { let single_sql = statements.into_iter().next().unwrap_or_default(); - let result = execute_sql_statement(state, connection_id, database, &single_sql, schema, cancel_token).await?; + let result = execute_sql_statement_with_options( + state, + connection_id, + database, + &single_sql, + schema, + cancel_token, + options, + ) + .await?; return Ok(vec![result]); } @@ -343,6 +487,8 @@ pub async fn execute_multi_core( affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }); break; } @@ -355,6 +501,8 @@ pub async fn execute_multi_core( affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }); } } @@ -380,6 +528,8 @@ async fn execute_multi_sqlserver( affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }); break; } @@ -410,6 +560,8 @@ async fn execute_multi_sqlserver( affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }); } } @@ -422,6 +574,8 @@ async fn execute_multi_sqlserver( affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }); } @@ -445,7 +599,7 @@ pub async fn execute_statements( let start = std::time::Instant::now(); for (i, sql) in statements.iter().enumerate() { - match do_execute(state, &pool_key, sql, schema, None).await { + match do_execute(state, &pool_key, sql, schema, None, QueryExecutionOptions::default()).await { Ok(result) => { total_affected += result.affected_rows; } @@ -470,6 +624,8 @@ pub async fn execute_statements( affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } @@ -566,6 +722,8 @@ async fn exec_tx_pg_inner( affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } @@ -613,6 +771,8 @@ async fn exec_tx_mysql_raw_inner( affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } @@ -640,6 +800,8 @@ async fn exec_tx_sqlite_inner( affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } @@ -662,18 +824,20 @@ async fn exec_tx_explicit_inner( } drop(conns); - do_execute(state, pool_key, "BEGIN", schema, None) + do_execute(state, pool_key, "BEGIN", schema, None, QueryExecutionOptions::default()) .await .map_err(|e| format!("Failed to begin transaction: {}", e))?; let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { - match do_execute(state, pool_key, sql, schema, None).await { + match do_execute(state, pool_key, sql, schema, None, QueryExecutionOptions::default()).await { Ok(result) => { total_affected += result.affected_rows; } Err(e) => { - if let Err(rb_err) = do_execute(state, pool_key, "ROLLBACK", schema, None).await { + if let Err(rb_err) = + do_execute(state, pool_key, "ROLLBACK", schema, None, QueryExecutionOptions::default()).await + { log::error!("ROLLBACK failed after statement {} error: {}", i + 1, rb_err); } return Err(format!("Statement {} failed: {}", i + 1, e)); @@ -681,7 +845,9 @@ async fn exec_tx_explicit_inner( } } - do_execute(state, pool_key, "COMMIT", schema, None).await.map_err(|e| format!("COMMIT failed: {}", e))?; + do_execute(state, pool_key, "COMMIT", schema, None, QueryExecutionOptions::default()) + .await + .map_err(|e| format!("COMMIT failed: {}", e))?; Ok(db::QueryResult { columns: vec![], @@ -689,6 +855,8 @@ async fn exec_tx_explicit_inner( affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } @@ -702,7 +870,7 @@ async fn exec_tx_none_inner( let mut total_affected: u64 = 0; for (i, sql) in statements.iter().enumerate() { log::info!("[query][tx-none:statement:start] index={} sql={}", i + 1, sql); - match do_execute(state, pool_key, sql, schema, None).await { + match do_execute(state, pool_key, sql, schema, None, QueryExecutionOptions::default()).await { Ok(result) => { total_affected += result.affected_rows; log::info!("[query][tx-none:statement:done] index={} affected_rows={}", i + 1, result.affected_rows); @@ -724,6 +892,8 @@ async fn exec_tx_none_inner( affected_rows: total_affected, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } @@ -745,6 +915,8 @@ mod tests { affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }) }) .await; @@ -762,6 +934,8 @@ mod tests { affected_rows: 0, execution_time_ms: 0, truncated: false, + session_id: None, + has_more: false, }) }) .await; @@ -848,4 +1022,58 @@ mod tests { assert_eq!(params["database"], "analytics"); assert_eq!(params["schema"], "app"); } + + #[test] + fn agent_execute_query_params_include_row_and_fetch_limits() { + let params = agent_execute_query_params( + "SELECT * FROM events", + Some("app"), + QueryExecutionOptions { max_rows: Some(500), fetch_size: Some(250), ..Default::default() }, + ); + + assert_eq!(params["sql"], "SELECT * FROM events"); + assert_eq!(params["schema"], "app"); + assert_eq!(params["maxRows"], 500); + assert_eq!(params["fetchSize"], 250); + } + + #[test] + fn agent_execute_query_params_default_to_safety_row_limit() { + let params = agent_execute_query_params("SELECT * FROM events", None, QueryExecutionOptions::default()); + + assert_eq!(params["sql"], "SELECT * FROM events"); + assert!(params.get("schema").is_none()); + assert_eq!(params["maxRows"], MAX_ROWS); + assert!(params.get("fetchSize").is_none()); + } + + #[test] + fn agent_execute_query_page_params_include_page_fetch_and_safety_limits() { + let params = agent_execute_query_page_params( + "SELECT * FROM events", + Some("app"), + QueryExecutionOptions { page_size: Some(500), fetch_size: Some(250), ..Default::default() }, + ); + + assert_eq!(params["sql"], "SELECT * FROM events"); + assert_eq!(params["schema"], "app"); + assert_eq!(params["pageSize"], 500); + assert_eq!(params["fetchSize"], 250); + assert_eq!(params["maxRows"], MAX_ROWS); + } + + #[test] + fn agent_fetch_query_page_params_include_session_and_page_size() { + let params = agent_fetch_query_page_params("session-1", 500); + + assert_eq!(params["sessionId"], "session-1"); + assert_eq!(params["pageSize"], 500); + } + + #[test] + fn agent_close_query_session_params_include_session() { + let params = agent_close_query_session_params("session-1"); + + assert_eq!(params["sessionId"], "session-1"); + } } diff --git a/crates/dbx-core/src/transfer.rs b/crates/dbx-core/src/transfer.rs index 969b53a7f..b5c6f597a 100644 --- a/crates/dbx-core/src/transfer.rs +++ b/crates/dbx-core/src/transfer.rs @@ -548,6 +548,8 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res affected_rows: 0, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } else { let affected = con.execute(&sql, []).map_err(|e| e.to_string())?; @@ -557,6 +559,8 @@ pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Res affected_rows: affected as u64, execution_time_ms: start.elapsed().as_millis(), truncated: false, + session_id: None, + has_more: false, }) } }) diff --git a/crates/dbx-core/src/types.rs b/crates/dbx-core/src/types.rs index 71d59cb81..904312220 100644 --- a/crates/dbx-core/src/types.rs +++ b/crates/dbx-core/src/types.rs @@ -58,6 +58,10 @@ pub struct QueryResult { pub execution_time_ms: u128, #[serde(default)] pub truncated: bool, + #[serde(default)] + pub session_id: Option, + #[serde(default)] + pub has_more: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] diff --git a/src-tauri/src/commands/query.rs b/src-tauri/src/commands/query.rs index b45ca02fd..fdbd2d995 100644 --- a/src-tauri/src/commands/query.rs +++ b/src-tauri/src/commands/query.rs @@ -16,13 +16,25 @@ pub async fn execute_query( sql: String, schema: Option, execution_id: Option, + max_rows: Option, + fetch_size: Option, + page_size: Option, + result_session_id: Option, ) -> Result { let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); let cancel_token = registered_query.as_ref().map(|query| query.token()); - dbx_core::query::execute_sql_statement(&state, &connection_id, &database, &sql, schema.as_deref(), cancel_token) - .await + dbx_core::query::execute_sql_statement_with_options( + &state, + &connection_id, + &database, + &sql, + schema.as_deref(), + cancel_token, + dbx_core::query::QueryExecutionOptions { max_rows, fetch_size, page_size, result_session_id }, + ) + .await } #[tauri::command] @@ -33,6 +45,10 @@ pub async fn execute_multi( sql: String, schema: Option, execution_id: Option, + max_rows: Option, + fetch_size: Option, + page_size: Option, + result_session_id: Option, ) -> Result, String> { let registered_query = execution_id.as_ref().filter(|id| !id.trim().is_empty()).map(|id| state.running_queries.register(id.clone())); @@ -47,9 +63,16 @@ pub async fn execute_multi( sql ); - let result = - dbx_core::query::execute_multi_core(&state, &connection_id, &database, &sql, schema.as_deref(), cancel_token) - .await; + let result = dbx_core::query::execute_multi_core_with_options( + &state, + &connection_id, + &database, + &sql, + schema.as_deref(), + cancel_token, + dbx_core::query::QueryExecutionOptions { max_rows, fetch_size, page_size, result_session_id }, + ) + .await; match &result { Ok(results) => log::info!( "[query][execute_multi:done] trace_id={} result_count={} row_counts={:?}", @@ -67,6 +90,16 @@ pub async fn cancel_query(state: State<'_, Arc>, execution_id: String) Ok(state.running_queries.cancel(&execution_id)) } +#[tauri::command] +pub async fn close_query_session( + state: State<'_, Arc>, + connection_id: String, + database: String, + session_id: String, +) -> Result { + dbx_core::query::close_query_session(&state, &connection_id, &database, &session_id).await +} + #[tauri::command] pub async fn execute_batch( state: State<'_, Arc>, diff --git a/src-tauri/src/lib.rs b/src-tauri/src/lib.rs index cef320a88..5eb29ba33 100644 --- a/src-tauri/src/lib.rs +++ b/src-tauri/src/lib.rs @@ -123,6 +123,7 @@ pub fn run() { commands::query::execute_query, commands::query::execute_multi, commands::query::cancel_query, + commands::query::close_query_session, commands::query::execute_batch, commands::query::execute_script, commands::query::execute_in_transaction, diff --git a/src-web/src/main.rs b/src-web/src/main.rs index 297829d8b..2251600e3 100644 --- a/src-web/src/main.rs +++ b/src-web/src/main.rs @@ -105,6 +105,7 @@ async fn main() { .route("/query/execute-script", post(routes::query::execute_script)) .route("/query/execute-in-transaction", post(routes::query::execute_in_transaction)) .route("/query/cancel", post(routes::query::cancel_query)) + .route("/query/close-session", post(routes::query::close_query_session)) // Redis .route("/redis/list-databases", post(routes::redis::list_databases)) .route("/redis/scan-keys", post(routes::redis::scan_keys)) diff --git a/src-web/src/routes/query.rs b/src-web/src/routes/query.rs index 125868800..dd728c2a7 100644 --- a/src-web/src/routes/query.rs +++ b/src-web/src/routes/query.rs @@ -15,6 +15,10 @@ pub struct ExecuteQueryRequest { pub sql: String, pub schema: Option, pub execution_id: Option, + pub max_rows: Option, + pub fetch_size: Option, + pub page_size: Option, + pub result_session_id: Option, } #[derive(Deserialize)] @@ -23,6 +27,14 @@ pub struct CancelRequest { pub execution_id: String, } +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct CloseSessionRequest { + pub connection_id: String, + pub database: String, + pub session_id: String, +} + #[derive(Deserialize)] #[serde(rename_all = "camelCase")] pub struct ExecuteBatchRequest { @@ -41,13 +53,19 @@ pub async fn execute_query( let registered = state.app.running_queries.register(execution_id); let cancel_token = registered.token(); - let result = dbx_core::query::execute_sql_statement( + let result = dbx_core::query::execute_sql_statement_with_options( &state.app, &req.connection_id, &req.database, &req.sql, req.schema.as_deref(), Some(cancel_token), + dbx_core::query::QueryExecutionOptions { + max_rows: req.max_rows, + fetch_size: req.fetch_size, + page_size: req.page_size, + result_session_id: req.result_session_id, + }, ) .await .map_err(AppError)?; @@ -65,13 +83,19 @@ pub async fn execute_multi( let registered = state.app.running_queries.register(execution_id); let cancel_token = registered.token(); - let result = dbx_core::query::execute_multi_core( + let result = dbx_core::query::execute_multi_core_with_options( &state.app, &req.connection_id, &req.database, &req.sql, req.schema.as_deref(), Some(cancel_token), + dbx_core::query::QueryExecutionOptions { + max_rows: req.max_rows, + fetch_size: req.fetch_size, + page_size: req.page_size, + result_session_id: req.result_session_id, + }, ) .await .map_err(AppError)?; @@ -105,6 +129,17 @@ pub async fn cancel_query( Json(serde_json::json!({ "cancelled": cancelled })) } +pub async fn close_query_session( + State(state): State>, + Json(req): Json, +) -> Result, AppError> { + let closed = dbx_core::query::close_query_session(&state.app, &req.connection_id, &req.database, &req.session_id) + .await + .map_err(AppError)?; + + Ok(Json(serde_json::json!(closed))) +} + pub async fn execute_script( State(state): State>, Json(req): Json, diff --git a/src/components/grid/DataGrid.vue b/src/components/grid/DataGrid.vue index 98f1eac9b..c9ee74c09 100644 --- a/src/components/grid/DataGrid.vue +++ b/src/components/grid/DataGrid.vue @@ -111,6 +111,7 @@ const props = defineProps<{ databaseType?: DatabaseType; connectionId?: string; database?: string; + schema?: string; context?: "results" | "table-data"; sourceColumns?: Array; queryEditabilityReason?: string; @@ -121,6 +122,9 @@ const props = defineProps<{ columns: ColumnInfo[]; primaryKeys: string[]; }; + pageOffset?: number; + pageLimit?: number; + countSql?: string; loading?: boolean; onExecuteSql?: (sql: string) => Promise; customSave?: (changes: { @@ -951,7 +955,16 @@ watch( // --- Pagination --- const pageSize = ref(settingsStore.editorSettings.pageSize); const currentPage = ref(1); -const isFullPage = computed(() => props.result.rows.length >= pageSize.value); +watch( + () => [props.pageOffset, props.pageLimit], + ([offset, limit]) => { + if (typeof offset !== "number" || typeof limit !== "number" || limit <= 0) return; + pageSize.value = limit; + currentPage.value = Math.floor(offset / limit) + 1; + }, +); +const canGoNextPage = computed(() => props.result.has_more === true || props.result.rows.length >= pageSize.value); +const canJumpLastPage = computed(() => canGoNextPage.value && (!!props.tableMeta || !!props.countSql)); const isResultsContext = computed(() => props.context === "results"); const resultEditStatus = computed(() => { if (!isResultsContext.value || !hasData.value) return null; @@ -1036,7 +1049,7 @@ function prevPage() { emit("paginate", (currentPage.value - 1) * pageSize.value, pageSize.value, currentWhereInput(), currentOrderBy()); } function nextPage() { - if (!isFullPage.value) return; + if (!canGoNextPage.value) return; currentPage.value++; resetGridVerticalScroll(true); emit("paginate", (currentPage.value - 1) * pageSize.value, pageSize.value, currentWhereInput(), currentOrderBy()); @@ -1050,17 +1063,23 @@ function changePageSize(size: number) { } async function lastPage() { - if (!props.connectionId || !props.tableMeta) return; - const table = qualifiedTableName({ - databaseType: props.databaseType, - schema: props.tableMeta.schema, - tableName: props.tableMeta.tableName, - }); - const predicate = normalizeWhereInput(currentWhereInput()); - const where = predicate ? ` WHERE (${predicate})` : ""; - const sql = `SELECT COUNT(*) AS cnt FROM ${table}${where}`; + if (!props.connectionId) return; + let sql = props.countSql; + let schema = props.schema; + if (props.tableMeta) { + const table = qualifiedTableName({ + databaseType: props.databaseType, + schema: props.tableMeta.schema, + tableName: props.tableMeta.tableName, + }); + const predicate = normalizeWhereInput(currentWhereInput()); + const where = predicate ? ` WHERE (${predicate})` : ""; + sql = `SELECT COUNT(*) AS cnt FROM ${table}${where}`; + schema = props.tableMeta.schema; + } + if (!sql) return; try { - const result = await api.executeQuery(props.connectionId, props.database ?? "", sql, props.tableMeta.schema); + const result = await api.executeQuery(props.connectionId, props.database ?? "", sql, schema); const total = Number(result.rows?.[0]?.[0] ?? 0); if (total <= 0) return; const lastPageNum = Math.ceil(total / pageSize.value); @@ -3523,10 +3542,10 @@ defineExpose({ {{ currentPage }} - - diff --git a/src/components/layout/ContentArea.vue b/src/components/layout/ContentArea.vue index c52a4d938..7ef82f598 100644 --- a/src/components/layout/ContentArea.vue +++ b/src/components/layout/ContentArea.vue @@ -304,7 +304,11 @@ defineExpose({ focusSearch }); :database-type="activeConnection?.db_type" :connection-id="activeTab.connectionId" :database="activeTab.database" + :schema="activeTab.schema" :table-meta="activeTab.tableMeta" + :page-offset="activeTab.resultPageOffset" + :page-limit="activeTab.resultPageLimit" + :count-sql="activeTab.resultCountSql" :on-execute-sql="async (sql: string) => emit('executeSql', sql)" @reload=" ( diff --git a/src/components/sidebar/TreeItem.vue b/src/components/sidebar/TreeItem.vue index 26d8ef3d5..514a0c491 100644 --- a/src/components/sidebar/TreeItem.vue +++ b/src/components/sidebar/TreeItem.vue @@ -471,12 +471,14 @@ async function openData() { elapsed: elapsed(), }); const pks = editablePrimaryKeys(config.db_type, columns); + const limit = settingsStore.editorSettings.pageSize; const sql = buildTableSelectSql({ databaseType: config.db_type, schema: node.schema, tableName: node.label, columns: columns.map((column) => column.name), primaryKeys: pks, + limit, includeRowId: usesSyntheticRowIdKey(config.db_type, pks), }); console.info("[DBX][openData:sql-built]", { diff --git a/src/composables/useDataGridActions.ts b/src/composables/useDataGridActions.ts index bf0bd35ca..07ae062be 100644 --- a/src/composables/useDataGridActions.ts +++ b/src/composables/useDataGridActions.ts @@ -2,6 +2,7 @@ import { type ComputedRef } from "vue"; import { useI18n } from "vue-i18n"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; +import { useSettingsStore } from "@/stores/settingsStore"; import { buildTableSelectSql, quoteTableIdentifier } from "@/lib/tableSelectSql"; import { editablePrimaryKeys, usesSyntheticRowIdKey } from "@/lib/tableEditing"; import { buildSortedQuerySql } from "@/lib/queryResultSort"; @@ -13,6 +14,7 @@ export function useDataGridActions(activeTab: ComputedRef) const { toast } = useToast(); const connectionStore = useConnectionStore(); const queryStore = useQueryStore(); + const settingsStore = useSettingsStore(); function quoteIdent(tab: QueryTab, name: string): string { const config = connectionStore.getConfig(tab.connectionId); @@ -41,6 +43,7 @@ export function useDataGridActions(activeTab: ComputedRef) primaryKeys, fallbackOrderColumns, includeRowId: useRowId, + limit: options.limit ?? settingsStore.editorSettings.pageSize, ...options, }); } @@ -87,7 +90,24 @@ export function useDataGridActions(activeTab: ComputedRef) async function onPaginate(offset: number, limit: number, whereInput?: string, orderBy?: string) { const tab = activeTab.value; - if (!tab?.tableMeta) return; + if (!tab) return; + if (tab.mode !== "data") { + const baseSql = tab.resultSortedSql ?? tab.resultBaseSql ?? tab.lastExecutedSql ?? tab.sql; + if (!baseSql.trim()) return; + const expectedNextOffset = (tab.resultPageOffset ?? 0) + (tab.resultPageLimit ?? limit); + const sessionId = + tab.result?.has_more && tab.result?.session_id && offset === expectedNextOffset && limit === tab.resultPageLimit + ? tab.result.session_id + : undefined; + await queryStore.executeTabSql(tab.id, baseSql, { + resultBaseSql: tab.resultBaseSql ?? tab.sql, + resultSortedSql: tab.resultSortedSql, + pagination: { offset, limit, sessionId }, + }); + return; + } + + if (!tab.tableMeta) return; tab.whereInput = whereInput ?? ""; const sql = buildTableSql(tab, { limit, offset, whereInput, orderBy }); queryStore.updateSql(tab.id, sql); diff --git a/src/composables/useNavigationTargets.ts b/src/composables/useNavigationTargets.ts index 7c421cef7..96ccead1a 100644 --- a/src/composables/useNavigationTargets.ts +++ b/src/composables/useNavigationTargets.ts @@ -3,6 +3,7 @@ import { buildTableSelectSql } from "@/lib/tableSelectSql"; import { editablePrimaryKeys, usesSyntheticRowIdKey } from "@/lib/tableEditing"; import { useConnectionStore } from "@/stores/connectionStore"; import { useQueryStore } from "@/stores/queryStore"; +import { useSettingsStore } from "@/stores/settingsStore"; export type NavigationTarget = { connectionId: string; @@ -16,6 +17,8 @@ export type NavigationTarget = { async function openTableTarget(target: NavigationTarget) { const connectionStore = useConnectionStore(); const queryStore = useQueryStore(); + const settingsStore = useSettingsStore(); + const pageLimit = settingsStore.editorSettings.pageSize; connectionStore.activeConnectionId = target.connectionId; const config = connectionStore.getConfig(target.connectionId); @@ -37,6 +40,7 @@ async function openTableTarget(target: NavigationTarget) { columns: columns.map((column) => column.name), primaryKeys, whereInput: target.whereInput, + limit: pageLimit, }); queryStore.updateSql(tabId, sql); queryStore.setTableMeta(tabId, { @@ -53,6 +57,7 @@ async function openTableTarget(target: NavigationTarget) { schema: target.schema, tableName: target.tableName, whereInput: target.whereInput, + limit: pageLimit, }); queryStore.updateSql(tabId, sql); queryStore.setTableMeta(tabId, { @@ -83,6 +88,7 @@ async function openTableTarget(target: NavigationTarget) { primaryKeys, columns: columns.map((column) => column.name), includeRowId: true, + limit: pageLimit, }); queryStore.updateSql(tabId, newSql); await queryStore.executeTabSql(tabId, newSql); diff --git a/src/lib/api.ts b/src/lib/api.ts index 133113930..917f5c151 100644 --- a/src/lib/api.ts +++ b/src/lib/api.ts @@ -71,6 +71,7 @@ export const executeBatch = forward("executeBatch"); export const executeScript = forward("executeScript"); export const executeInTransaction = forward("executeInTransaction"); export const cancelQuery = forward("cancelQuery"); +export const closeQuerySession = forward("closeQuerySession"); // AI export const aiComplete = forward("aiComplete"); diff --git a/src/lib/http.ts b/src/lib/http.ts index 3a2ce6e1a..8b845ba32 100644 --- a/src/lib/http.ts +++ b/src/lib/http.ts @@ -254,8 +254,9 @@ export async function executeQuery( sql: string, schema?: string, executionId?: string, + options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, ): Promise { - return post("/api/query/execute", { connectionId, database, sql, schema, executionId }); + return post("/api/query/execute", { connectionId, database, sql, schema, executionId, ...options }); } export async function executeMulti( @@ -264,8 +265,13 @@ export async function executeMulti( sql: string, schema?: string, executionId?: string, + options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, ): Promise { - return post("/api/query/execute-multi", { connectionId, database, sql, schema, executionId }); + return post("/api/query/execute-multi", { connectionId, database, sql, schema, executionId, ...options }); +} + +export async function closeQuerySession(connectionId: string, database: string, sessionId: string): Promise { + return post("/api/query/close-session", { connectionId, database, sessionId }); } export async function executeBatch( diff --git a/src/lib/queryResultPagination.ts b/src/lib/queryResultPagination.ts new file mode 100644 index 000000000..11516d500 --- /dev/null +++ b/src/lib/queryResultPagination.ts @@ -0,0 +1,136 @@ +import type { DatabaseType } from "../types/database.ts"; +import { usesFetchFirst } from "./databaseCapabilities.ts"; +import { quoteTableIdentifier } from "./tableSelectSql.ts"; +import { findStatementAtCursor } from "./sqlStatementSplit.ts"; + +export interface PaginatedQuerySqlResult { + ok: true; + sql: string; +} + +export interface PaginatedQuerySqlError { + ok: false; + reason: "empty" | "multi" | "not_select" | "unsupported"; +} + +export interface QueryPaginationExecutionPlan { + sqlToExecute: string; + pageSql?: string; + pageLimit?: number; + pageOffset?: number; + countSql?: string; + useAgentResultSession: boolean; +} + +const unsupportedPaginationTypes = new Set(["neo4j", "mongodb", "redis", "elasticsearch"]); + +export function buildQueryPaginationExecutionPlan({ + sql, + queryBaseSql, + databaseType, + pagination, + useAgentCursor, +}: { + sql: string; + queryBaseSql: string; + databaseType: DatabaseType | undefined; + pagination: { limit: number; offset: number; sessionId?: string }; + useAgentCursor: boolean; +}): QueryPaginationExecutionPlan { + const plan: QueryPaginationExecutionPlan = { + sqlToExecute: sql, + useAgentResultSession: false, + }; + const counted = buildCountQuerySql(queryBaseSql, databaseType); + if (counted.ok) { + plan.countSql = counted.sql; + } + + if (pagination.sessionId) { + plan.pageLimit = pagination.limit; + plan.pageOffset = pagination.offset; + plan.useAgentResultSession = true; + return plan; + } + + if (useAgentCursor && pagination.offset === 0) { + plan.sqlToExecute = queryBaseSql; + plan.pageLimit = pagination.limit; + plan.pageOffset = pagination.offset; + plan.useAgentResultSession = true; + return plan; + } + + const paginated = buildPaginatedQuerySql(sql, databaseType, pagination.limit, pagination.offset); + if (paginated.ok) { + plan.sqlToExecute = paginated.sql; + plan.pageSql = paginated.sql; + plan.pageLimit = pagination.limit; + plan.pageOffset = pagination.offset; + } + return plan; +} + +export function buildPaginatedQuerySql( + originalSql: string, + databaseType: DatabaseType | undefined, + limit: number, + offset: number, +): PaginatedQuerySqlResult | PaginatedQuerySqlError { + const statement = singleSelectableStatement(originalSql); + if (!statement.ok) return statement; + if (unsupportedPaginationTypes.has(databaseType)) return { ok: false, reason: "unsupported" }; + + const safeLimit = Math.max(1, Math.floor(limit)); + const safeOffset = Math.max(0, Math.floor(offset)); + const alias = quoteTableIdentifier(databaseType, "dbx_page"); + const base = `SELECT * FROM (${statement.sql}) ${alias}`; + + if (databaseType === "sqlserver") { + return { + ok: true, + sql: `${base} ORDER BY (SELECT NULL) OFFSET ${safeOffset} ROWS FETCH NEXT ${safeLimit} ROWS ONLY`, + }; + } + + if (usesFetchFirst(databaseType)) { + const offsetSql = safeOffset ? ` OFFSET ${safeOffset} ROWS` : ""; + return { ok: true, sql: `${base}${offsetSql} FETCH FIRST ${safeLimit} ROWS ONLY` }; + } + + const offsetSql = safeOffset ? ` OFFSET ${safeOffset}` : ""; + return { ok: true, sql: `${base} LIMIT ${safeLimit}${offsetSql};` }; +} + +export function buildCountQuerySql( + originalSql: string, + databaseType: DatabaseType | undefined, +): PaginatedQuerySqlResult | PaginatedQuerySqlError { + const statement = singleSelectableStatement(originalSql); + if (!statement.ok) return statement; + if (unsupportedPaginationTypes.has(databaseType)) return { ok: false, reason: "unsupported" }; + + const alias = quoteTableIdentifier(databaseType, "dbx_count"); + return { ok: true, sql: `SELECT COUNT(*) AS dbx_total_rows FROM (${statement.sql}) ${alias};` }; +} + +function singleSelectableStatement( + originalSql: string, +): { ok: true; sql: string } | Pick { + const baseSql = originalSql.trim(); + if (!baseSql) return { ok: false, reason: "empty" }; + + const statement = findStatementAtCursor(baseSql, 0) + .trim() + .replace(/;+\s*$/, "") + .trim(); + if (!statement) return { ok: false, reason: "empty" }; + if (statement.length !== baseSql.replace(/;+\s*$/, "").trim().length) { + return { ok: false, reason: "multi" }; + } + if (!/^\s*(SELECT|WITH)\b/i.test(statement)) { + return { ok: false, reason: "not_select" }; + } + + return { ok: true, sql: statement }; +} diff --git a/src/lib/tauri.ts b/src/lib/tauri.ts index 7893dba2e..086320a8a 100644 --- a/src/lib/tauri.ts +++ b/src/lib/tauri.ts @@ -181,8 +181,9 @@ export async function executeQuery( sql: string, schema?: string, executionId?: string, + options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, ): Promise { - return invoke("execute_query", { connectionId, database, sql, schema, executionId }); + return invoke("execute_query", { connectionId, database, sql, schema, executionId, ...options }); } export async function executeMulti( @@ -191,14 +192,19 @@ export async function executeMulti( sql: string, schema?: string, executionId?: string, + options?: { maxRows?: number; fetchSize?: number; pageSize?: number; resultSessionId?: string }, ): Promise { - return invoke("execute_multi", { connectionId, database, sql, schema, executionId }); + return invoke("execute_multi", { connectionId, database, sql, schema, executionId, ...options }); } export async function cancelQuery(executionId: string): Promise { return invoke("cancel_query", { executionId }); } +export async function closeQuerySession(connectionId: string, database: string, sessionId: string): Promise { + return invoke("close_query_session", { connectionId, database, sessionId }); +} + export async function executeBatch( connectionId: string, database: string, diff --git a/src/stores/queryStore.ts b/src/stores/queryStore.ts index 566425a65..fb8a2bc55 100644 --- a/src/stores/queryStore.ts +++ b/src/stores/queryStore.ts @@ -14,9 +14,12 @@ import { } from "@/lib/sqlAnalysis"; import { restoreOpenTabsState, serializeOpenTabs } from "@/lib/openTabsPersistence"; import { mongoDocumentsToQueryResult, parseMongoFindCommand } from "@/lib/mongoShellCommand"; +import { buildQueryPaginationExecutionPlan } from "@/lib/queryResultPagination"; +import { AGENT_DRIVER_TYPES } from "@/lib/databaseCapabilities"; import { editablePrimaryKeys } from "@/lib/tableEditing"; import * as api from "@/lib/api"; import { useConnectionStore } from "@/stores/connectionStore"; +import { useSettingsStore } from "@/stores/settingsStore"; import { isTauriRuntime } from "@/lib/tauriRuntime"; import type { SavedSqlFile } from "@/types/database"; @@ -46,6 +49,19 @@ export const useQueryStore = defineStore("query", () => { const activeTabId = ref(restored.activeTabId); const MAX_CACHED_RESULTS = 10; + async function closeResultSession(tab: QueryTab | undefined, preserveSessionId?: string) { + const sessionId = tab?.resultSessionId ?? tab?.result?.session_id; + if (!tab || !sessionId || sessionId === preserveSessionId) return; + try { + await api.closeQuerySession(tab.connectionId, tab.database, sessionId); + } catch (error) { + console.warn("[DBX][query-session:close:error]", { tabId: tab.id, sessionId, error }); + } finally { + if (tab.resultSessionId === sessionId) tab.resultSessionId = undefined; + if (tab.result?.session_id === sessionId) tab.result.session_id = undefined; + } + } + const _persistSnapshot = computed(() => tabs.value.map((t) => ({ id: t.id, @@ -147,6 +163,7 @@ export const useQueryStore = defineStore("query", () => { if (idx < 0) return; if (tabs.value[idx].isExecuting) void cancelTabExecution(id); if (tabs.value[idx].isExplaining) void cancelTabExplain(id); + void closeResultSession(tabs.value[idx]); tabs.value[idx].result = undefined; tabs.value[idx].results = undefined; tabs.value.splice(idx, 1); @@ -158,6 +175,7 @@ export const useQueryStore = defineStore("query", () => { function closeOtherTabs(id: string) { tabs.value.filter((tab) => tab.id !== id && tab.isExecuting).forEach((tab) => void cancelTabExecution(tab.id)); tabs.value.filter((tab) => tab.id !== id && tab.isExplaining).forEach((tab) => void cancelTabExplain(tab.id)); + tabs.value.filter((tab) => tab.id !== id).forEach((tab) => void closeResultSession(tab)); const next = closeOtherTabsState(tabs.value, activeTabId.value, id); tabs.value = next.tabs; activeTabId.value = next.activeTabId; @@ -166,6 +184,7 @@ export const useQueryStore = defineStore("query", () => { function closeAllTabs() { tabs.value.filter((tab) => tab.isExecuting).forEach((tab) => void cancelTabExecution(tab.id)); tabs.value.filter((tab) => tab.isExplaining).forEach((tab) => void cancelTabExplain(tab.id)); + tabs.value.forEach((tab) => void closeResultSession(tab)); const next = closeAllTabsState(tabs.value, activeTabId.value); tabs.value = next.tabs; activeTabId.value = next.activeTabId; @@ -410,7 +429,11 @@ export const useQueryStore = defineStore("query", () => { async function executeTabSql( id: string, sql: string, - options?: { resultBaseSql?: string; resultSortedSql?: string | undefined }, + options?: { + resultBaseSql?: string; + resultSortedSql?: string | undefined; + pagination?: { limit: number; offset: number; sessionId?: string }; + }, ) { const tab = tabs.value.find((t) => t.id === id); if (!tab || !sql.trim()) return; @@ -432,9 +455,35 @@ export const useQueryStore = defineStore("query", () => { schema: tab.schema, sql, }); + const queryBaseSql = options?.resultBaseSql ?? sql; + let sqlToExecute = sql; + let pageSql: string | undefined; + let pageLimit: number | undefined; + let pageOffset: number | undefined; + let countSql: string | undefined; + let useAgentResultSession = false; try { const connStore = useConnectionStore(); const conn = connStore.getConfig(tab.connectionId); + const useAgentCursor = !!conn?.db_type && AGENT_DRIVER_TYPES.has(conn.db_type); + await closeResultSession(tab, options?.pagination?.sessionId); + if (tab.mode === "query") { + const settingsStore = useSettingsStore(); + const pagination = options?.pagination ?? { limit: settingsStore.editorSettings.pageSize, offset: 0 }; + const plan = buildQueryPaginationExecutionPlan({ + sql, + queryBaseSql, + databaseType: conn?.db_type, + pagination, + useAgentCursor, + }); + sqlToExecute = plan.sqlToExecute; + pageSql = plan.pageSql; + pageLimit = plan.pageLimit; + pageOffset = plan.pageOffset; + countSql = plan.countSql; + useAgentResultSession = plan.useAgentResultSession; + } const mongoFind = conn?.db_type === "mongodb" ? parseMongoFindCommand(sql) : null; if (mongoFind) { console.info("[DBX][executeTabSql:mongo-find:start]", { traceId, collection: mongoFind.collection }); @@ -469,7 +518,25 @@ export const useQueryStore = defineStore("query", () => { } console.info("[DBX][executeTabSql:execute-multi:start]", { traceId, elapsed: elapsed() }); - const results = await api.executeMulti(tab.connectionId, tab.database, sql, tab.schema, executionId); + const executionOptions = + typeof pageLimit === "number" + ? useAgentResultSession + ? { + maxRows: 10000, + fetchSize: pageLimit, + pageSize: pageLimit, + resultSessionId: options?.pagination?.sessionId, + } + : { maxRows: pageLimit, fetchSize: pageLimit } + : undefined; + const results = await api.executeMulti( + tab.connectionId, + tab.database, + sqlToExecute, + tab.schema, + executionId, + executionOptions, + ); console.info("[DBX][executeTabSql:execute-multi:done]", { traceId, resultCount: results.length, @@ -488,10 +555,15 @@ export const useQueryStore = defineStore("query", () => { current.activeResultIndex = undefined; current.result = results[0]; } - current.resultBaseSql = options?.resultBaseSql ?? sql; + current.resultBaseSql = queryBaseSql; current.resultSortedSql = options?.resultSortedSql; + current.resultPageSql = pageSql; + current.resultPageLimit = pageLimit; + current.resultPageOffset = pageOffset; + current.resultCountSql = countSql; + current.resultSessionId = current.result?.session_id ?? undefined; console.info("[DBX][executeTabSql:metadata:start]", { traceId, elapsed: elapsed() }); - await analyzeQueryMetadata(current, current.resultBaseSql); + await analyzeQueryMetadata(current, queryBaseSql); console.info("[DBX][executeTabSql:metadata:done]", { traceId, elapsed: elapsed() }); } else { console.warn("[DBX][executeTabSql:stale-result]", { @@ -511,8 +583,13 @@ export const useQueryStore = defineStore("query", () => { current.querySourceColumns = undefined; current.queryEditabilityReason = undefined; if (current.mode !== "data") current.tableMeta = undefined; - current.resultBaseSql = options?.resultBaseSql ?? sql; + current.resultBaseSql = queryBaseSql; current.resultSortedSql = options?.resultSortedSql; + current.resultPageSql = pageSql; + current.resultPageLimit = pageLimit; + current.resultPageOffset = pageOffset; + current.resultCountSql = countSql; + current.resultSessionId = undefined; } } finally { const current = tabs.value.find((t) => t.id === id); diff --git a/src/types/database.ts b/src/types/database.ts index b88d01d31..2117b15ee 100644 --- a/src/types/database.ts +++ b/src/types/database.ts @@ -175,6 +175,8 @@ export interface QueryResult { affected_rows: number; execution_time_ms: number; truncated?: boolean; + session_id?: string | null; + has_more?: boolean; } export type TreeNodeType = @@ -253,6 +255,11 @@ export interface QueryTab { lastExecutedSql?: string; resultBaseSql?: string; resultSortedSql?: string; + resultPageSql?: string; + resultPageLimit?: number; + resultPageOffset?: number; + resultCountSql?: string; + resultSessionId?: string; pinned?: boolean; result?: QueryResult; results?: QueryResult[]; diff --git a/tests/queryResultPagination.test.ts b/tests/queryResultPagination.test.ts new file mode 100644 index 000000000..c7cbe809b --- /dev/null +++ b/tests/queryResultPagination.test.ts @@ -0,0 +1,141 @@ +import assert from "node:assert/strict"; +import test from "node:test"; +import { + buildCountQuerySql, + buildPaginatedQuerySql, + buildQueryPaginationExecutionPlan, +} from "../src/lib/queryResultPagination.ts"; + +test("wraps a single select query with limit and offset", () => { + const result = buildPaginatedQuerySql("SELECT id, name FROM users;", "postgres", 100, 200); + + assert.deepEqual(result, { + ok: true, + sql: 'SELECT * FROM (SELECT id, name FROM users) "dbx_page" LIMIT 100 OFFSET 200;', + }); +}); + +test("uses MySQL style quoting for paginated query alias", () => { + const result = buildPaginatedQuerySql("SELECT id FROM users WHERE active = 1", "mysql", 50, 0); + + assert.deepEqual(result, { + ok: true, + sql: "SELECT * FROM (SELECT id FROM users WHERE active = 1) `dbx_page` LIMIT 50;", + }); +}); + +test("uses SQL Server offset fetch pagination", () => { + const result = buildPaginatedQuerySql("SELECT id FROM users", "sqlserver", 100, 300); + + assert.deepEqual(result, { + ok: true, + sql: "SELECT * FROM (SELECT id FROM users) [dbx_page] ORDER BY (SELECT NULL) OFFSET 300 ROWS FETCH NEXT 100 ROWS ONLY", + }); +}); + +test("uses fetch first pagination for Oracle first page", () => { + const result = buildPaginatedQuerySql("SELECT id FROM users", "oracle", 100, 0); + + assert.deepEqual(result, { + ok: true, + sql: 'SELECT * FROM (SELECT id FROM users) "dbx_page" FETCH FIRST 100 ROWS ONLY', + }); +}); + +test("supports CTE select queries", () => { + const result = buildPaginatedQuerySql("WITH cte AS (SELECT 1 AS id) SELECT * FROM cte", "clickhouse", 100, 0); + + assert.deepEqual(result, { + ok: true, + sql: 'SELECT * FROM (WITH cte AS (SELECT 1 AS id) SELECT * FROM cte) "dbx_page" LIMIT 100;', + }); +}); + +test("rejects multiple statements", () => { + const result = buildPaginatedQuerySql("SELECT 1; SELECT 2;", "postgres", 100, 0); + + assert.deepEqual(result, { ok: false, reason: "multi" }); +}); + +test("rejects non select statements", () => { + const result = buildPaginatedQuerySql("UPDATE users SET name = 'A'", "postgres", 100, 0); + + assert.deepEqual(result, { ok: false, reason: "not_select" }); +}); + +test("wraps a single select query for total row count", () => { + const result = buildCountQuerySql("SELECT id, name FROM users;", "postgres"); + + assert.deepEqual(result, { + ok: true, + sql: 'SELECT COUNT(*) AS dbx_total_rows FROM (SELECT id, name FROM users) "dbx_count";', + }); +}); + +test("uses MySQL style quoting for count query alias", () => { + const result = buildCountQuerySql("WITH cte AS (SELECT 1 AS id) SELECT * FROM cte", "mysql"); + + assert.deepEqual(result, { + ok: true, + sql: "SELECT COUNT(*) AS dbx_total_rows FROM (WITH cte AS (SELECT 1 AS id) SELECT * FROM cte) `dbx_count`;", + }); +}); + +test("rejects count query for unsupported database types", () => { + const result = buildCountQuerySql("SELECT * FROM nodes", "neo4j"); + + assert.deepEqual(result, { ok: false, reason: "unsupported" }); +}); + +test("rejects count query for multiple statements", () => { + const result = buildCountQuerySql("SELECT 1; SELECT 2;", "postgres"); + + assert.deepEqual(result, { ok: false, reason: "multi" }); +}); + +test("uses an agent result session for the first jdbc page", () => { + const plan = buildQueryPaginationExecutionPlan({ + sql: "SELECT * FROM events", + queryBaseSql: "SELECT * FROM events", + databaseType: "oracle", + pagination: { limit: 500, offset: 0 }, + useAgentCursor: true, + }); + + assert.equal(plan.sqlToExecute, "SELECT * FROM events"); + assert.equal(plan.pageLimit, 500); + assert.equal(plan.pageOffset, 0); + assert.equal(plan.pageSql, undefined); + assert.equal(plan.useAgentResultSession, true); +}); + +test("keeps using an agent result session for sequential jdbc pages", () => { + const plan = buildQueryPaginationExecutionPlan({ + sql: "SELECT * FROM events", + queryBaseSql: "SELECT * FROM events", + databaseType: "oracle", + pagination: { limit: 500, offset: 500, sessionId: "session-1" }, + useAgentCursor: true, + }); + + assert.equal(plan.sqlToExecute, "SELECT * FROM events"); + assert.equal(plan.pageLimit, 500); + assert.equal(plan.pageOffset, 500); + assert.equal(plan.useAgentResultSession, true); +}); + +test("uses SQL pagination instead of jdbc cursor for random agent page jumps", () => { + const plan = buildQueryPaginationExecutionPlan({ + sql: "SELECT * FROM events", + queryBaseSql: "SELECT * FROM events", + databaseType: "oracle", + pagination: { limit: 500, offset: 1500 }, + useAgentCursor: true, + }); + + assert.equal(plan.sqlToExecute, 'SELECT * FROM (SELECT * FROM events) "dbx_page" OFFSET 1500 ROWS FETCH FIRST 500 ROWS ONLY'); + assert.equal(plan.pageSql, plan.sqlToExecute); + assert.equal(plan.pageLimit, 500); + assert.equal(plan.pageOffset, 1500); + assert.equal(plan.useAgentResultSession, false); +});