fix(export): stream postgres query result exports

crates/dbx-core/src/db/postgres.rs

@@ -712,6 +712,139 @@ async fn execute_select_query(
     }
 }
 
+pub enum PostgresQueryStreamItem {
+    Columns { columns: Vec<String>, column_types: Vec<String> },
+    Row(Vec<serde_json::Value>),
+}
+
+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<usize>,
+    on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>,
+) -> Result<u64, PostgresQueryStreamError> {
+    let stmt =
+        client.prepare_cached(sql).await.map_err(|err| PostgresQueryStreamError::Postgres { err, emitted: false })?;
+    let columns: Vec<String> = stmt.columns().iter().map(|c| c.name().to_string()).collect();
+    let column_types: Vec<String> = 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<usize>,
+    on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>,
+) -> Result<u64, String> {
+    let stream = client.simple_query_raw(sql).await.map_err(pg_error_to_string)?;
+    tokio::pin!(stream);
+    let mut columns: Vec<String> = 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<usize>,
+    on_item: &mut impl FnMut(PostgresQueryStreamItem) -> Result<(), String>,
+) -> Result<u64, String> {
+    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<usize>,
+    cancel_token: Option<CancellationToken>,
+    budget: DbOperationBudget,
+    cancel_context: Option<PostgresCancelContext>,
+    on_item: impl FnMut(PostgresQueryStreamItem) -> Result<(), String>,
+) -> Result<u64, String> {
+    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<QueryResult, String> {
     execute_query_with_schema_and_max_rows(pool, schema, sql, None).await
 }

crates/dbx-core/src/query.rs

@@ -884,7 +884,7 @@ fn resolve_query_timeout(timeout_secs: Option<u64>) -> Option<Duration> {
     }
 }
 
-async fn operation_budget_for_pool_key(
+pub async fn operation_budget_for_pool_key(
     state: &AppState,
     pool_key: &str,
     query_timeout: Option<Duration>,

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<CancellationToken>,
+    on_progress: &impl Fn(TableExportProgress),
+) -> Result<bool, String> {
+    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<String> = 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,

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");
+}