From 06e653b0bb7c127218b96a46230486e9eae43e51 Mon Sep 17 00:00:00 2001 From: t8y2 <1156263951@qq.com> Date: Fri, 3 Jul 2026 20:02:19 +0800 Subject: [PATCH] fix(export): stream postgres query result exports --- crates/dbx-core/src/db/postgres.rs | 186 ++++++++++++++++++ crates/dbx-core/src/query.rs | 2 +- crates/dbx-core/src/query_result_export.rs | 153 +++++++++++++- .../live_postgres_query_result_export.rs | 148 ++++++++++++++ 4 files changed, 487 insertions(+), 2 deletions(-) create mode 100644 crates/dbx-core/tests/live_postgres_query_result_export.rs diff --git a/crates/dbx-core/src/db/postgres.rs b/crates/dbx-core/src/db/postgres.rs index 279c3bf50..7a740f36f 100644 --- a/crates/dbx-core/src/db/postgres.rs +++ b/crates/dbx-core/src/db/postgres.rs @@ -712,6 +712,139 @@ async fn execute_select_query( } } +pub enum PostgresQueryStreamItem { + Columns { columns: Vec, column_types: Vec }, + Row(Vec), +} + +enum PostgresQueryStreamError { + Postgres { err: tokio_postgres::Error, emitted: bool }, + Export(String), +} + +impl PostgresQueryStreamError { + fn into_string(self) -> String { + match self { + Self::Postgres { err, .. } => pg_error_to_string(err), + Self::Export(err) => err, + } + } +} + +async fn stream_select_query_prepared( + client: &deadpool_postgres::Client, + sql: &str, + row_limit: Option, + on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, +) -> Result { + let stmt = + client.prepare_cached(sql).await.map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?; + let columns: Vec = stmt.columns().iter().map(|c| c.name().to_string()).collect(); + let column_types: Vec = stmt.columns().iter().map(|c| c.type_().name().to_string()).collect(); + + let params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> = Vec::new(); + let stream = client + .query_raw(&stmt, params) + .await + .map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?; + tokio::pin!(stream); + let mut rows_streamed = 0_u64; + let mut columns_emitted = false; + while let Some(row_result) = stream.next().await { + if row_limit.is_some_and(|limit| rows_streamed as usize >= limit) { + break; + } + let row = row_result + .map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: columns_emitted || rows_streamed > 0 })?; + if !columns_emitted { + on_item(PostgresQueryStreamItem::Columns { columns: columns.clone(), column_types: column_types.clone() }) + .map_err(PostgresQueryStreamError::Export)?; + columns_emitted = true; + } + let values = (0..row.columns().len()) + .map(|i| pg_value_to_json(&row, i, column_types.get(i).map(String::as_str).unwrap_or(""))) + .collect(); + on_item(PostgresQueryStreamItem::Row(values)).map_err(PostgresQueryStreamError::Export)?; + rows_streamed += 1; + } + if !columns_emitted { + on_item(PostgresQueryStreamItem::Columns { columns, column_types }) + .map_err(PostgresQueryStreamError::Export)?; + } + Ok(rows_streamed) +} + +async fn stream_select_query_text( + client: &deadpool_postgres::Client, + sql: &str, + row_limit: Option, + on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, +) -> Result { + let stream = client.simple_query_raw(sql).await.map_err(pg_error_to_string)?; + tokio::pin!(stream); + let mut columns: Vec = Vec::new(); + let mut rows_streamed = 0_u64; + while let Some(message) = stream.next().await { + match message.map_err(pg_error_to_string)? { + SimpleQueryMessage::RowDescription(cols) => { + columns = cols.iter().map(|c| c.name().to_string()).collect(); + on_item(PostgresQueryStreamItem::Columns { columns: columns.clone(), column_types: Vec::new() })?; + } + SimpleQueryMessage::Row(row) => { + if row_limit.is_some_and(|limit| rows_streamed as usize >= limit) { + break; + } + if columns.is_empty() { + columns = row.columns().iter().map(|c| c.name().to_string()).collect(); + on_item(PostgresQueryStreamItem::Columns { columns: columns.clone(), column_types: Vec::new() })?; + } + let mut values = Vec::with_capacity(row.len()); + for i in 0..row.len() { + values.push(match row.try_get(i).map_err(pg_error_to_string)? { + Some(value) => serde_json::Value::String(value.to_string()), + None => serde_json::Value::Null, + }); + } + on_item(PostgresQueryStreamItem::Row(values))?; + rows_streamed += 1; + } + SimpleQueryMessage::CommandComplete(_) => {} + _ => {} + } + } + Ok(rows_streamed) +} + +async fn stream_select_query_inner( + client: &deadpool_postgres::Client, + sql: &str, + row_limit: Option, + on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, +) -> Result { + match stream_select_query_prepared(client, sql, row_limit, on_item).await { + Ok(rows) => Ok(rows), + Err(PostgresQueryStreamError::Postgres { err, emitted: false }) if should_retry_postgres_stale_cache(&err) => { + // The cached prepared statement can become stale after schema changes. + // Evict and retry once, matching the normal query execution path. + log::warn!("[postgres][stream:stale_cache] evicting cached statement: {}", pg_error_to_string(err)); + client.statement_cache.remove(sql, &[]); + match stream_select_query_prepared(client, sql, row_limit, on_item).await { + Ok(rows) => Ok(rows), + Err(PostgresQueryStreamError::Postgres { err, emitted: false }) + if should_retry_postgres_text_query(&err) => + { + stream_select_query_text(client, sql, row_limit, on_item).await + } + Err(err) => Err(err.into_string()), + } + } + Err(PostgresQueryStreamError::Postgres { err, emitted: false }) if should_retry_postgres_text_query(&err) => { + stream_select_query_text(client, sql, row_limit, on_item).await + } + Err(err) => Err(err.into_string()), + } +} + pub async fn stream_query_rows( pool: &Pool, sql: &str, @@ -1963,6 +2096,59 @@ pub async fn execute_query_with_max_rows_and_cancel( .await } +pub async fn stream_select_query_with_cancel( + pool: &Pool, + schema: Option<&str>, + sql: &str, + max_rows: Option, + cancel_token: Option, + budget: DbOperationBudget, + cancel_context: Option, + on_item: impl FnMut(PostgresQueryStreamItem) -> Result<(), String>, +) -> Result { + let start = Instant::now(); + let client = checkout_postgres_client(pool, cancel_token.as_ref(), budget.checkout_timeout).await?; + let mut on_item = on_item; + let row_limit = max_rows.map(|limit| limit.max(1)); + let schema = schema.map(str::trim).filter(|schema| !schema.is_empty()); + let schema_was_set = schema.is_some_and(|_| !is_transaction_recovery_statement(sql)); + + if let Some(schema) = schema.filter(|_| schema_was_set) { + // Match normal query execution: export may reference unqualified names + // in the active schema, so the streaming path must use the same search_path. + execute_postgres_infra_statement( + &client, + &format!("SET search_path TO {}, public", pg_quote_ident(schema)), + budget.recycle_timeout, + "schema.set", + ) + .await?; + } + + let pg_cancel_token = client.cancel_token(); + let result = wait_postgres_query( + pg_cancel_token, + cancel_context, + cancel_token, + budget.query_timeout, + budget.cancel_timeout, + stream_select_query_inner(&client, sql, row_limit, &mut on_item), + ) + .await; + + if schema_was_set { + let reset_result = reset_postgres_search_path(&client, budget.cleanup_timeout, start).await; + match (result, reset_result) { + (Ok(rows), Ok(())) => Ok(rows), + (Err(query_err), Ok(())) => Err(query_err), + (Ok(_), Err(reset_err)) => Err(reset_err), + (Err(query_err), Err(reset_err)) => Err(format!("{query_err}; {reset_err}")), + } + } else { + result + } +} + pub async fn execute_query_with_schema(pool: &Pool, schema: &str, sql: &str) -> Result { execute_query_with_schema_and_max_rows(pool, schema, sql, None).await } diff --git a/crates/dbx-core/src/query.rs b/crates/dbx-core/src/query.rs index ac0dfdfb0..8b410dd6e 100644 --- a/crates/dbx-core/src/query.rs +++ b/crates/dbx-core/src/query.rs @@ -884,7 +884,7 @@ fn resolve_query_timeout(timeout_secs: Option) -> Option { } } -async fn operation_budget_for_pool_key( +pub async fn operation_budget_for_pool_key( state: &AppState, pool_key: &str, query_timeout: Option, diff --git a/crates/dbx-core/src/query_result_export.rs b/crates/dbx-core/src/query_result_export.rs index 7fddc598c..6176ca687 100644 --- a/crates/dbx-core/src/query_result_export.rs +++ b/crates/dbx-core/src/query_result_export.rs @@ -9,7 +9,8 @@ use crate::database_export::is_export_cancelled; pub use crate::database_export::ExportStatus; use crate::models::connection::DatabaseType; use crate::query::{ - canceled_error, close_query_session, execute_sql_statement_with_options, QueryExecutionOptions, QUERY_CANCELED, + canceled_error, close_query_session, execute_sql_statement_with_options, operation_budget_for_pool_key, + QueryExecutionOptions, QUERY_CANCELED, }; use crate::query_result_sql::{ build_query_pagination_execution_plan, QueryPagination, QueryPaginationExecutionPlanOptions, @@ -314,6 +315,10 @@ async fn export_query_result_core_inner( on_progress(progress(request, 0, ExportStatus::Running, None)); + if try_export_postgres_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? { + return Ok(()); + } + if try_export_sqlserver_query_result_stream(state, request, &format, cancel_token.clone(), on_progress).await? { return Ok(()); } @@ -554,6 +559,152 @@ async fn export_query_result_core_inner( Ok(()) } +async fn try_export_postgres_query_result_stream( + state: &AppState, + request: &QueryResultExportRequest, + format: &str, + cancel_token: Option, + on_progress: &impl Fn(TableExportProgress), +) -> Result { + if request.use_agent_cursor + || !crate::sql::starts_with_executable_sql_keyword( + &request.sql, + &["SELECT", "SHOW", "EXPLAIN", "WITH", "TABLE"], + ) + { + return Ok(false); + } + + let database = request.database.trim(); + let pool_key = if database.is_empty() { + state.get_or_create_pool_for_session(&request.connection_id, None, request.client_session_id.as_deref()).await? + } else { + state + .get_or_create_pool_for_session( + &request.connection_id, + Some(database), + request.client_session_id.as_deref(), + ) + .await? + }; + let connections = state.connections.read().await; + let Some(pool) = connections.get(&pool_key).and_then(|pool| match pool { + PoolKind::Postgres(pool) => Some(pool.clone()), + _ => None, + }) else { + return Ok(false); + }; + drop(connections); + + if let Some(execution_id) = request.execution_id.as_deref() { + state.running_queries.set_pool_key(execution_id, pool_key.clone()); + } + state.touch_pool_activity(&pool_key).await; + let _activity_touch = state.pool_activity_touch(&pool_key); + + let xlsx_hard_limit_active = xlsx_hard_limit_active(format, request); + let row_limit = effective_row_limit(format, request); + let stream_row_limit = + if xlsx_hard_limit_active { row_limit.map(|limit| limit.saturating_add(1)) } else { row_limit }; + let progress_row_interval = request.page_size.max(1) as u64; + let mut columns: Vec = Vec::new(); + let mut rows_exported = 0_u64; + let mut last_progress_rows = 0_u64; + let mut last_progress_at = Instant::now(); + let mut csv_file = if format == "csv" { + let mut file = + BufWriter::new(File::create(&request.file_path).map_err(|e| format!("Failed to create file: {e}"))?); + file.write_all(b"\xEF\xBB\xBF").map_err(|e| format!("Failed to write BOM: {e}"))?; + Some(file) + } else { + None + }; + let mut xlsx = None; + let budget = operation_budget_for_pool_key(state, &pool_key, query_export_timeout(request.timeout_secs)).await; + let cancel_context = state.get_postgres_cancel_context(&pool_key).await; + + crate::db::postgres::stream_select_query_with_cancel( + &pool, + request.schema.as_deref(), + &request.sql, + stream_row_limit, + cancel_token, + budget, + cancel_context, + |item| { + match item { + crate::db::postgres::PostgresQueryStreamItem::Columns { columns: stream_columns, .. } => { + columns = stream_columns; + if let Some(file) = csv_file.as_mut() { + let csv = format_query_result_csv(&columns, &[]); + let header = csv.strip_suffix('\n').unwrap_or(&csv); + file.write_all(header.as_bytes()).map_err(|e| format!("Failed to write CSV: {e}"))?; + } else { + let xlsx_file = + File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; + xlsx = + Some(start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?); + } + } + crate::db::postgres::PostgresQueryStreamItem::Row(row) => { + if xlsx_hard_limit_active && rows_exported as usize >= XLSX_MAX_DATA_ROWS { + return Err(XLSX_ROW_LIMIT_ERROR.to_string()); + } + if let Some(file) = csv_file.as_mut() { + let rows_csv = format_query_result_csv_rows(std::slice::from_ref(&row)); + write!(file, "\n{rows_csv}").map_err(|e| format!("Failed to write CSV rows: {e}"))?; + } else if let Some(writer) = xlsx.as_mut() { + writer.write_row(&row).map_err(|e| format!("Failed to write XLSX row: {e}"))?; + } else { + let xlsx_file = + File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; + xlsx = + Some(start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?); + if let Some(writer) = xlsx.as_mut() { + writer.write_row(&row).map_err(|e| format!("Failed to write XLSX row: {e}"))?; + } + } + rows_exported += 1; + let now = Instant::now(); + if should_emit_stream_progress( + rows_exported, + last_progress_rows, + progress_row_interval, + now.duration_since(last_progress_at), + ) { + on_progress(progress(request, rows_exported, ExportStatus::Running, None)); + last_progress_rows = rows_exported; + last_progress_at = now; + } + } + } + Ok(()) + }, + ) + .await?; + + if rows_exported != last_progress_rows { + on_progress(progress(request, rows_exported, ExportStatus::Running, None)); + } + on_progress(progress(request, rows_exported, ExportStatus::Writing, None)); + if let Some(file) = csv_file.as_mut() { + file.flush().map_err(|e| format!("Failed to flush CSV file: {e}"))?; + } + if let Some(writer) = xlsx { + let mut buf = + finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; + buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; + } else if format == "xlsx" { + let xlsx_file = File::create(&request.file_path).map_err(|e| format!("Failed to create XLSX file: {e}"))?; + let writer = start_streaming_xlsx_workbook(BufWriter::new(xlsx_file), Some("Result"), &columns)?; + let mut buf = + finish_streaming_xlsx_workbook(writer).map_err(|e| format!("Failed to finalize XLSX file: {e}"))?; + buf.flush().map_err(|e| format!("Failed to flush XLSX file: {e}"))?; + } + on_progress(progress(request, rows_exported, ExportStatus::Done, None)); + Ok(true) +} + async fn try_export_sqlserver_query_result_stream( state: &AppState, request: &QueryResultExportRequest, diff --git a/crates/dbx-core/tests/live_postgres_query_result_export.rs b/crates/dbx-core/tests/live_postgres_query_result_export.rs new file mode 100644 index 000000000..46b05e146 --- /dev/null +++ b/crates/dbx-core/tests/live_postgres_query_result_export.rs @@ -0,0 +1,148 @@ +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use dbx_core::connection::AppState; +use dbx_core::db::postgres; +use dbx_core::models::connection::{ConnectionConfig, DatabaseType}; +use dbx_core::query_result_export::{export_query_result_core, ExportStatus, QueryResultExportRequest}; +use dbx_core::storage::Storage; + +fn live_postgres_config( + id: &str, + host: &str, + port: u16, + user: &str, + password: &str, + database: &str, +) -> ConnectionConfig { + ConnectionConfig { + id: id.to_string(), + name: id.to_string(), + db_type: DatabaseType::Postgres, + driver_profile: None, + driver_label: None, + url_params: None, + host: host.to_string(), + port, + username: user.to_string(), + password: password.to_string(), + database: Some(database.to_string()), + visible_databases: None, + visible_schemas: None, + attached_databases: Vec::new(), + color: None, + transport_layers: Vec::new(), + connect_timeout_secs: 10, + query_timeout_secs: 30, + idle_timeout_secs: 60, + keepalive_interval_secs: 0, + ssl: false, + ca_cert_path: String::new(), + client_cert_path: String::new(), + client_key_path: String::new(), + sysdba: false, + oracle_connection_type: None, + connection_string: None, + redis_connection_mode: None, + redis_sentinel_master: String::new(), + redis_sentinel_nodes: String::new(), + redis_sentinel_username: String::new(), + redis_sentinel_password: String::new(), + redis_sentinel_tls: false, + redis_cluster_nodes: String::new(), + redis_key_separator: dbx_core::models::connection::default_redis_key_separator(), + redis_scan_page_size: None, + etcd_endpoints: String::new(), + gbase_server: String::new(), + informix_server: String::new(), + external_config: None, + jdbc_driver_class: None, + jdbc_driver_paths: Vec::new(), + one_time: false, + read_only: false, + } +} + +#[tokio::test] +#[ignore = "requires DBX_LIVE_POSTGRES_HOST/PORT/USER/PASSWORD/DATABASE pointing at a writable PostgreSQL database"] +async fn live_postgres_query_result_export_uses_single_streamed_query() { + let host = std::env::var("DBX_LIVE_POSTGRES_HOST").unwrap_or_else(|_| "127.0.0.1".to_string()); + let port = std::env::var("DBX_LIVE_POSTGRES_PORT").ok().and_then(|value| value.parse().ok()).unwrap_or(5432); + let user = std::env::var("DBX_LIVE_POSTGRES_USER").unwrap_or_else(|_| "postgres".to_string()); + let password = std::env::var("DBX_LIVE_POSTGRES_PASSWORD").unwrap_or_default(); + let database = std::env::var("DBX_LIVE_POSTGRES_DATABASE").unwrap_or_else(|_| "postgres".to_string()); + let url = format!("postgresql://{user}:{password}@{host}:{port}/{database}"); + let setup_pool = postgres::connect(&url, Duration::from_secs(10)).await.expect("connect PostgreSQL"); + + let suffix = uuid::Uuid::new_v4().simple().to_string(); + let schema = format!("dbx_query_export_{}", &suffix[..8]); + let setup = vec![ + format!("CREATE SCHEMA \"{schema}\""), + format!( + "CREATE OR REPLACE FUNCTION \"{schema}\".assert_no_limit_offset() RETURNS integer LANGUAGE plpgsql AS $$ \ + DECLARE q text; \ + BEGIN \ + SELECT query INTO q FROM pg_stat_activity WHERE pid = pg_backend_pid(); \ + IF q ~* '\\m(limit|offset)\\M' THEN \ + RAISE EXCEPTION 'query was paginated: %', q; \ + END IF; \ + RETURN 1; \ + END; \ + $$" + ), + ]; + let cleanup = vec![format!("DROP SCHEMA IF EXISTS \"{schema}\" CASCADE")]; + let _ = postgres::execute_batch(&setup_pool, &cleanup).await; + postgres::execute_batch(&setup_pool, &setup).await.expect("create live test schema"); + + let dir = std::env::temp_dir().join(format!("dbx-live-postgres-query-export-{suffix}")); + std::fs::create_dir_all(&dir).unwrap(); + let storage = Storage::open(&dir.join("storage.db")).await.unwrap(); + let state = AppState::new(storage); + let connection_id = "live-postgres-query-export"; + let config = live_postgres_config(connection_id, &host, port, &user, &password, &database); + state.configs.write().await.insert(config.id.clone(), config); + + let file_path = dir.join("result.csv"); + let sql = + format!("SELECT i, \"{schema}\".assert_no_limit_offset() AS marker FROM generate_series(1, 2050) AS s(i)"); + let request = QueryResultExportRequest { + export_id: format!("live-postgres-query-export-{suffix}"), + connection_id: connection_id.to_string(), + database: database.clone(), + schema: Some(schema.clone()), + sql: sql.clone(), + query_base_sql: sql, + database_type: DatabaseType::Postgres, + use_agent_cursor: false, + file_path: file_path.to_string_lossy().to_string(), + format: "csv".to_string(), + page_size: 100, + row_limit: None, + total_rows: None, + timeout_secs: Some(30), + keyset_optimization_enabled: true, + client_session_id: None, + execution_id: Some(format!("live-postgres-query-export-{suffix}")), + }; + let done_seen = AtomicBool::new(false); + let result = export_query_result_core(&state, &request, None, |progress| { + if matches!(progress.status, ExportStatus::Done) { + done_seen.store(true, Ordering::Relaxed); + } + }) + .await; + + let cleanup_result = postgres::execute_batch(&setup_pool, &cleanup).await; + let csv = std::fs::read_to_string(&file_path).unwrap_or_default(); + let _ = std::fs::remove_dir_all(&dir); + + result.expect("export query result"); + cleanup_result.expect("cleanup live test schema"); + assert!(done_seen.load(Ordering::Relaxed)); + assert!(csv.starts_with('\u{feff}')); + assert!(csv.contains("\"i\",\"marker\""), "csv={csv:?}"); + assert!(csv.contains("\"1\",\"1\"")); + assert!(csv.contains("\"2050\",\"1\"")); + assert_eq!(csv.lines().count(), 2051, "unexpected csv row count"); +}