perf(query): optimize result loading and connection reuse

crates/dbx-core/src/connection.rs

@@ -2648,6 +2648,10 @@ fn normalize_client_session_id(client_session_id: Option<&str>) -> Option<String
     client_session_id.map(str::trim).filter(|session| !session.is_empty()).map(|session| session.replace(':', "_"))
 }
 
+pub fn task_client_session_id(task_kind: &str, task_id: &str) -> String {
+    format!("{task_kind}:{task_id}")
+}
+
 fn mysql_pool_max_connections_for_session(client_session_id: Option<&str>) -> usize {
     if normalize_client_session_id(client_session_id).is_some() {
         1
@@ -3078,8 +3082,8 @@ mod tests {
         metadata_connection_config, mysql_metadata_fallback_url, oceanbase_mysql_query_timeout_sql,
         oceanbase_mysql_setup_queries, prestosql_jdbc_config_for_endpoint, redacted_connection_url_for_endpoint,
         redis_sentinel_transport_id, redis_sentinel_transport_prefix, sqlserver_legacy_agent_config,
-        sqlserver_legacy_agent_error, uses_bare_mysql_pool, uses_tcp_probe, validate_h2_database_path, AppState,
-        MysqlMode, PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS,
+        sqlserver_legacy_agent_error, task_client_session_id, uses_bare_mysql_pool, uses_tcp_probe,
+        validate_h2_database_path, AppState, MysqlMode, PoolKind, PRESTOSQL_JDBC_DRIVER_CLASS,
     };
     use crate::agent_connection::{
         agent_connect_params, mongo_legacy_error_with_auth_hint, mongo_uses_legacy_driver,
@@ -3149,6 +3153,13 @@ mod tests {
         }
     }
 
+    #[test]
+    fn task_client_session_ids_are_stable_and_isolated() {
+        assert_eq!(task_client_session_id("table-export", "job-1"), "table-export:job-1");
+        assert_ne!(task_client_session_id("table-export", "job-1"), task_client_session_id("database-export", "job-1"));
+        assert_ne!(task_client_session_id("table-export", "job-1"), task_client_session_id("table-export", "job-2"));
+    }
+
     #[test]
     fn access_agent_connect_timeout_has_longer_default_floor() {
         let mut config = mysql_config(None);

crates/dbx-core/src/database_export.rs

@@ -4,6 +4,7 @@ use std::collections::HashSet;
 use std::io::Write;
 use tokio::sync::RwLock;
 
+use crate::connection::task_client_session_id;
 use crate::models::connection::DatabaseType;
 use crate::object_source_sql::build_export_object_source_sql;
 use crate::sql_dialect::{qualified_table_name, quote_table_identifier, uses_single_row_insert_statements};
@@ -16,6 +17,10 @@ use crate::transfer::{
 static EXPORT_CANCELLED: std::sync::LazyLock<RwLock<HashSet<String>>> =
     std::sync::LazyLock::new(|| RwLock::new(HashSet::new()));
 
+pub fn database_export_client_session_id(export_id: &str) -> String {
+    task_client_session_id("database-export", export_id)
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct DatabaseExportRequest {
@@ -979,7 +984,10 @@ pub async fn export_database_sql_core(
         .ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?;
 
     // 2. Get pool
-    let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?;
+    let client_session_id = database_export_client_session_id(&request.export_id);
+    let pool_key = state
+        .get_or_create_pool_for_session(&request.connection_id, Some(&request.database), Some(&client_session_id))
+        .await?;
 
     // 3. List tables
     let all_tables = crate::schema::list_tables_core(
@@ -1213,7 +1221,7 @@ pub async fn export_database_sql_core(
             if !col_names.is_empty() {
                 // Get row count
                 let count_query = crate::transfer::count_sql(table_name, &request.schema, &db_type);
-                let total_rows = match crate::transfer::execute_on_pool(state, &pool_key, &count_query).await {
+                let total_rows = match crate::transfer::execute_read_on_pool(state, &pool_key, &count_query).await {
                     Ok(result) => result.rows.first().and_then(|r| r.first()).and_then(|v| match v {
                         serde_json::Value::Number(n) => n.as_u64(),
                         serde_json::Value::String(s) => s.parse::<u64>().ok(),
@@ -1251,7 +1259,7 @@ pub async fn export_database_sql_core(
                         batch_size,
                     );
 
-                    let result = match crate::transfer::execute_on_pool(state, &pool_key, &sql).await {
+                    let result = match crate::transfer::execute_read_on_pool(state, &pool_key, &sql).await {
                         Ok(r) => r,
                         Err(e) => {
                             writeln!(file, "-- ERROR exporting data for table {table_name}: {e}")

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

@@ -985,7 +985,13 @@ pub async fn connect(url: &str, fallback_timeout: Duration) -> Result<Pool, Stri
         let pg_config = tokio_postgres::Config::from_str(&postgres_url.url)
             .map_err(|e| format!("Invalid PostgreSQL connection URL: {e}"))?;
 
-        let mgr_config = ManagerConfig { recycling_method: RecyclingMethod::Verified };
+        // Fast recycling only checks whether the connection is already closed
+        // instead of issuing a validation query on every checkout, saving one
+        // round-trip per query. Connections that went stale without being
+        // observed are caught when the query runs and recovered by the
+        // executor's ReconnectAndRetry path (see pool_error_action / do_execute
+        // in query.rs).
+        let mgr_config = ManagerConfig { recycling_method: RecyclingMethod::Fast };
         let tls_config = postgres_tls_config(
             &pg_config,
             &postgres_url.ssl_files,

crates/dbx-core/src/query.rs

@@ -23,6 +23,7 @@ use crate::connection::{AppState, PoolKind, TransactionSession, TxnConnection};
 use crate::database_capabilities;
 use crate::db;
 use crate::models::connection::{ConnectionConfig, DatabaseType};
+use crate::query_execution_sql::is_write_sql;
 #[cfg(feature = "duckdb-bundled")]
 use crate::sql::starts_with_duckdb_result_sql_keyword;
 use crate::sql::{split_sql_batches, split_sql_statements, starts_with_executable_sql_keyword};
@@ -841,6 +842,16 @@ pub fn should_discard_pool_after_error(db_type: Option<DatabaseType>, err: &str)
     matches!(pool_error_action(db_type, err), PoolErrorAction::Discard | PoolErrorAction::ReconnectAndRetry)
 }
 
+fn query_pool_error_action(db_type: Option<DatabaseType>, sql: &str, err: &str) -> PoolErrorAction {
+    match pool_error_action(db_type, err) {
+        // A connection error does not prove that the database did not receive
+        // a write. Only replay statements already accepted by the read-only
+        // protection classifier; writes discard the stale pool without retry.
+        PoolErrorAction::ReconnectAndRetry if is_write_sql(sql) => PoolErrorAction::Discard,
+        action => action,
+    }
+}
+
 fn is_os_connection_error(lower: &str) -> bool {
     let os_error_codes = ["10053", "10054", "10057", "10058", "10060", "10061"];
     if let Some(pos) = lower.find("os error ") {
@@ -1577,7 +1588,7 @@ pub async fn execute_sql_statement_with_options(
         do_execute(state, &pool_key, mysql_dialect, Some(database), sql, schema, cancel_token.clone(), options.clone())
             .await;
 
-    let action = result.as_ref().err().map(|e| pool_error_action(db_type, e));
+    let action = result.as_ref().err().map(|e| query_pool_error_action(db_type, sql, e));
     match action {
         Some(PoolErrorAction::ReconnectAndRetry) if !is_canceled(&cancel_token) => {
             let db_opt = if database.is_empty() { None } else { Some(database) };
@@ -3067,6 +3078,22 @@ mod tests {
         assert!(!is_agent_execute_batch_unsupported("Agent RPC error (-1): unknown method: execute_query"));
     }
 
+    #[test]
+    fn query_pool_error_policy_retries_reads_but_not_writes() {
+        assert_eq!(
+            query_pool_error_action(Some(DatabaseType::Postgres), "SELECT * FROM users", "connection reset by peer"),
+            PoolErrorAction::ReconnectAndRetry
+        );
+        assert_eq!(
+            query_pool_error_action(
+                Some(DatabaseType::Postgres),
+                "UPDATE users SET active = true",
+                "connection reset by peer"
+            ),
+            PoolErrorAction::Discard
+        );
+    }
+
     #[test]
     fn external_driver_method_unsupported_detects_legacy_plugin_errors() {
         assert!(is_external_driver_method_unsupported(

crates/dbx-core/src/table_export.rs

@@ -7,14 +7,14 @@ use std::time::Duration;
 use tokio_util::sync::CancellationToken;
 
 use crate::connection::MysqlMode;
-use crate::connection::{AppState, PoolKind};
+use crate::connection::{task_client_session_id, AppState, PoolKind};
 use crate::csv_export::{escape_csv, format_csv, value_to_csv_text};
 pub use crate::database_export::ExportStatus;
 use crate::database_export::{build_export_insert_statements, is_export_cancelled, BuildExportInsertStatementsOptions};
 use crate::db::agent_driver::AgentTableReadStartParams;
 use crate::models::connection::DatabaseType;
 use crate::transfer::{
-    count_sql_with_where, execute_on_pool, execute_on_pool_with_max_rows, keyset_pagination_sql,
+    count_sql_with_where, execute_read_on_pool, execute_read_on_pool_with_max_rows, keyset_pagination_sql,
     pagination_sql_with_filter_order, qualified_table, quote_identifier,
 };
 use crate::types::QueryResult;
@@ -23,6 +23,10 @@ use crate::xlsx_export::{finish_streaming_xlsx_workbook, start_streaming_xlsx_wo
 const DEFAULT_BATCH_SIZE: usize = 10_000;
 const SQL_INSERT_BATCH_SIZE: usize = 100;
 
+pub fn table_export_client_session_id(export_id: &str) -> String {
+    task_client_session_id("table-export", export_id)
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase")]
 pub struct TableExportRequest {
@@ -353,7 +357,7 @@ async fn fetch_paginated_table_export_batch(
         offset,
         active_batch_size,
     );
-    execute_on_pool_with_max_rows(state, pool_key, &sql, Some(active_batch_size)).await
+    execute_read_on_pool_with_max_rows(state, pool_key, &sql, Some(active_batch_size)).await
 }
 
 async fn close_table_read_session_if_open(
@@ -773,7 +777,10 @@ pub async fn export_table_data_core(
         .ok_or_else(|| format!("Connection config not found: {}", request.connection_id))?;
 
     // 2. Get pool
-    let pool_key = state.get_or_create_pool(&request.connection_id, Some(&request.database)).await?;
+    let client_session_id = table_export_client_session_id(&request.export_id);
+    let pool_key = state
+        .get_or_create_pool_for_session(&request.connection_id, Some(&request.database), Some(&client_session_id))
+        .await?;
 
     // 3. Resolve columns. Data grid exports can provide columns/primary keys
     // directly, which avoids expensive metadata round-trips on JDBC drivers.
@@ -830,7 +837,7 @@ pub async fn export_table_data_core(
             &db_type,
             request.where_input.as_deref(),
         );
-        match execute_on_pool(state, &pool_key, &count_query).await {
+        match execute_read_on_pool(state, &pool_key, &count_query).await {
             Ok(result) => result
                 .rows
                 .first()

crates/dbx-core/src/table_import.rs

@@ -6,7 +6,7 @@ use calamine::{open_workbook_auto, Data, Reader};
 use chrono::{DateTime, NaiveDate, NaiveDateTime};
 use serde::{Deserialize, Serialize};
 
-use crate::connection::AppState;
+use crate::connection::{task_client_session_id, AppState};
 use crate::models::connection::DatabaseType;
 use crate::transfer::{
     execute_on_pool, generate_insert_typed, get_columns_for_transfer, qualified_table, quote_identifier,
@@ -17,6 +17,10 @@ pub const DEFAULT_BATCH_SIZE: usize = 500;
 pub const CREATE_TABLE_INFERENCE_ROWS: usize = 100;
 pub const MAX_NON_STREAMING_IMPORT_BYTES: u64 = 100 * 1024 * 1024;
 
+pub fn table_import_client_session_id(import_id: &str) -> String {
+    task_client_session_id("table-import", import_id)
+}
+
 #[derive(Debug, Clone)]
 pub struct ParsedImportFile {
     pub columns: Vec<String>,

crates/dbx-core/src/transfer.rs

@@ -3,12 +3,12 @@ use serde::{Deserialize, Serialize};
 use std::collections::{HashMap, HashSet};
 use tokio::sync::RwLock;
 
-use crate::connection::{AppState, PoolKind};
+use crate::connection::{config_for_pool_key, AppState, PoolKind};
 use crate::db;
 use crate::db::mongo_driver::MongoDocumentResult;
 use crate::models::connection::DatabaseType;
 use crate::object_source_sql::{build_executable_object_source_statements, EditableObjectSourceSqlInput};
-use crate::query::{agent_execute_query_params, should_discard_pool_after_error, QueryExecutionOptions};
+use crate::query::{agent_execute_query_params, pool_error_action, PoolErrorAction, QueryExecutionOptions};
 use crate::sql::split_sql_statements;
 #[cfg(feature = "duckdb-bundled")]
 use crate::sql::starts_with_executable_sql_keyword;
@@ -2513,7 +2513,20 @@ fn mongo_columns_from_documents(documents: &[serde_json::Value]) -> Vec<db::Colu
 }
 
 pub async fn execute_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result<db::QueryResult, String> {
-    execute_on_pool_with_max_rows(state, pool_key, sql, None).await
+    execute_on_pool_with_options(state, pool_key, sql, None, TransferExecutionSafety::WriteNoReplay).await
+}
+
+pub async fn execute_read_on_pool(state: &AppState, pool_key: &str, sql: &str) -> Result<db::QueryResult, String> {
+    execute_read_on_pool_with_max_rows(state, pool_key, sql, None).await
+}
+
+pub async fn execute_read_on_pool_with_max_rows(
+    state: &AppState,
+    pool_key: &str,
+    sql: &str,
+    max_rows: Option<usize>,
+) -> Result<db::QueryResult, String> {
+    execute_on_pool_with_options(state, pool_key, sql, max_rows, TransferExecutionSafety::ReadOnlyRetryable).await
 }
 
 async fn execute_transfer_ddl_on_pool(
@@ -2596,6 +2609,91 @@ pub async fn execute_on_pool_with_max_rows(
     pool_key: &str,
     sql: &str,
     max_rows: Option<usize>,
+) -> Result<db::QueryResult, String> {
+    execute_on_pool_with_options(state, pool_key, sql, max_rows, TransferExecutionSafety::WriteNoReplay).await
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+enum TransferExecutionSafety {
+    ReadOnlyRetryable,
+    WriteNoReplay,
+}
+
+fn transfer_pool_error_action(
+    safety: TransferExecutionSafety,
+    db_type: Option<DatabaseType>,
+    err: &str,
+) -> PoolErrorAction {
+    match (safety, pool_error_action(db_type, err)) {
+        (TransferExecutionSafety::WriteNoReplay, PoolErrorAction::ReconnectAndRetry) => PoolErrorAction::Discard,
+        (_, action) => action,
+    }
+}
+
+async fn transfer_pool_context(
+    state: &AppState,
+    pool_key: &str,
+) -> (Option<String>, Option<String>, Option<DatabaseType>) {
+    let configs = state.configs.read().await;
+    let config = config_for_pool_key(pool_key, &configs);
+    (
+        config.map(|config| config.id.clone()),
+        database_from_pool_key(pool_key).map(str::to_string),
+        config.map(|config| config.db_type),
+    )
+}
+
+fn client_session_id_from_pool_key(pool_key: &str) -> Option<&str> {
+    pool_key.split_once(":session:").map(|(_, session)| session).filter(|session| !session.is_empty())
+}
+
+async fn execute_on_pool_with_options(
+    state: &AppState,
+    pool_key: &str,
+    sql: &str,
+    max_rows: Option<usize>,
+    safety: TransferExecutionSafety,
+) -> Result<db::QueryResult, String> {
+    let (connection_id, database, db_type) = transfer_pool_context(state, pool_key).await;
+    let client_session_id = client_session_id_from_pool_key(pool_key).map(str::to_string);
+    let mut current_pool_key = pool_key.to_string();
+
+    for attempt in 0..2 {
+        let result = execute_on_pool_once(state, &current_pool_key, sql, max_rows).await;
+        let Some(error) = result.as_ref().err() else {
+            return result;
+        };
+
+        match transfer_pool_error_action(safety, db_type, error) {
+            PoolErrorAction::Keep => return result,
+            PoolErrorAction::Discard => {
+                state.remove_pool_by_key(&current_pool_key).await;
+                return result;
+            }
+            PoolErrorAction::ReconnectAndRetry if attempt == 0 => {
+                let Some(connection_id) = connection_id.as_deref() else {
+                    state.remove_pool_by_key(&current_pool_key).await;
+                    return result;
+                };
+                current_pool_key = state
+                    .reconnect_pool_for_session(connection_id, database.as_deref(), client_session_id.as_deref())
+                    .await?;
+            }
+            PoolErrorAction::ReconnectAndRetry => {
+                state.remove_pool_by_key(&current_pool_key).await;
+                return result;
+            }
+        }
+    }
+
+    unreachable!("transfer pool execution retry loop runs at most twice")
+}
+
+async fn execute_on_pool_once(
+    state: &AppState,
+    pool_key: &str,
+    sql: &str,
+    max_rows: Option<usize>,
 ) -> Result<db::QueryResult, String> {
     // Read-only check: block transfer operations in readonly mode
     crate::query::check_read_only_for_connection(state, pool_key, sql).await?;
@@ -2631,10 +2729,6 @@ pub async fn execute_on_pool_with_max_rows(
             let mut client = client.lock().await;
             let result = db::sqlserver::execute_query_with_max_rows(&mut client, sql, max_rows).await;
             drop(client);
-            if matches!(result.as_ref(), Err(err) if should_discard_pool_after_error(Some(DatabaseType::SqlServer), err))
-            {
-                state.remove_pool_by_key(pool_key).await;
-            }
             result
         }
         PoolKind::Agent(client) => {
@@ -6248,6 +6342,30 @@ SELECT 1 FROM dual"#
         assert_eq!(database_from_pool_key("conn"), None);
     }
 
+    #[test]
+    fn transfer_read_retry_policy_retries_connection_errors() {
+        assert_eq!(
+            transfer_pool_error_action(
+                TransferExecutionSafety::ReadOnlyRetryable,
+                Some(DatabaseType::Postgres),
+                "connection reset by peer"
+            ),
+            PoolErrorAction::ReconnectAndRetry
+        );
+    }
+
+    #[test]
+    fn transfer_write_retry_policy_discards_without_replaying_batch() {
+        assert_eq!(
+            transfer_pool_error_action(
+                TransferExecutionSafety::WriteNoReplay,
+                Some(DatabaseType::Postgres),
+                "connection reset by peer"
+            ),
+            PoolErrorAction::Discard
+        );
+    }
+
     #[test]
     fn map_column_type_preserves_longtext_for_mysql_target() {
         assert_eq!(map_column_type("longtext", &DatabaseType::Mysql, &DatabaseType::Mysql), "longtext");