fix: make DuckDB attachments logic more robust (#509)

Author: sgrebnov

Cargo.lock

@@ -61,6 +61,7 @@ native-tls = { version = "0.2.11", optional = true }
 num-bigint = { workspace = true }
 num-traits = { version = "0.2", optional = true }
 odbc-api = { version = "19.0", optional = true }
+once_cell = "1.21"
 pem = { version = "3.0.4", optional = true }
 postgres-native-tls = { version = "0.5.0", optional = true }
 prost = { workspace = true, optional = true }

core/Cargo.toml

@@ -435,7 +435,10 @@ impl TableProviderFactory for DuckDBTableProviderFactory {
             Mode::File => {
                 let read_pool = pool.clone();
 
-                read_pool.set_attached_databases(&self.attach_databases(&options))
+                read_pool
+                    .set_attached_databases(&self.attach_databases(&options))
+                    .context(DbConnectionPoolSnafu)
+                    .map_err(to_datafusion_error)?
             }
             Mode::Memory => pool.clone(),
         };

core/src/duckdb.rs

@@ -18,6 +18,7 @@ use duckdb::ToSql;
 use duckdb::{Connection, DuckdbConnectionManager};
 use dyn_clone::DynClone;
 use rand::distr::{Alphanumeric, SampleString};
+use once_cell::sync::OnceCell;
 use snafu::{prelude::*, ResultExt};
 use tokio::sync::mpsc::Sender;
 
@@ -63,11 +64,15 @@ impl<T: ToSql + Sync + Send + DynClone> DuckDBSyncParameter for T {
 dyn_clone::clone_trait_object!(DuckDBSyncParameter);
 pub type DuckDBParameter = Box<dyn DuckDBSyncParameter>;
 
+/// Configuration and state for attaching external DuckDB databases.
 #[derive(Debug)]
 pub struct DuckDBAttachments {
     attachments: HashSet<Arc<str>>,
     random_id: String,
     main_db: String,
+    /// Cached search_path after first successful attachments initialization.
+    /// Uses OnceCell to ensure attachment happens exactly once.
+    search_path_cache: OnceCell<Arc<str>>,
 }
 
 impl DuckDBAttachments {
@@ -80,9 +85,16 @@ impl DuckDBAttachments {
             attachments,
             random_id,
             main_db: main_db.to_string(),
+            search_path_cache: OnceCell::new(),
         }
     }
 
+    /// Returns a reference to the set of database paths to attach.
+    #[must_use]
+    pub fn attachments(&self) -> &HashSet<Arc<str>> {
+        &self.attachments
+    }
+
     /// Returns the search path for the given database and attachments.
     /// The given database needs to be included separately, as search path by default do not include the main database.
     /// The `attachments` parameter represents full attachment names, e.g., ["attachment_zCVN0zYJ_0", ...]
@@ -213,6 +225,23 @@ impl DuckDBAttachments {
     fn get_attachment_name(random_id: &str, index: usize) -> String {
         format!("attachment_{random_id}_{index}")
     }
+
+    /// Lazily attaches databases on first call, then applies search_path to the connection.
+    /// Uses cached search_path on subsequent calls.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if attachment or search_path setting fails.
+    pub fn attach_once(&self, conn: &Connection) -> Result<()> {
+        let search_path = self.search_path_cache.get_or_try_init(|| {
+            self.attach(conn)
+        })?;
+
+        conn.execute(&format!("SET search_path = '{}'", search_path), [])
+            .context(DuckDBConnectionSnafu)?;
+
+        Ok(())
+    }
 }
 
 pub struct DuckDbConnection {
@@ -266,7 +295,7 @@ impl DuckDbConnection {
     ) -> &mut r2d2::PooledConnection<DuckdbConnectionManager> {
         &mut self.conn
     }
-
+ 
     #[must_use]
     pub fn with_unsupported_type_action(
         mut self,
@@ -429,7 +458,11 @@ impl SyncDbConnection<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
         let (batch_tx, mut batch_rx) = tokio::sync::mpsc::channel::<RecordBatch>(4);
 
         let conn = self.conn.try_clone()?;
-        Self::attach(&conn, &self.attachments)?;
+
+        if let Some(attachments) = &self.attachments {
+            attachments.attach_once(&conn)?;
+        }
+
         self.apply_connection_setup_queries(&conn)?;
 
         let fetch_schema_sql =

core/src/sql/db_connection_pool/dbconnection/duckdbconn.rs

@@ -1,5 +1,6 @@
 use async_trait::async_trait;
 use duckdb::{vtab::arrow::ArrowVTab, AccessMode, DuckdbConnectionManager};
+use once_cell::sync::OnceCell;
 use snafu::{prelude::*, ResultExt};
 use std::sync::Arc;
 
@@ -119,7 +120,7 @@ impl DuckDbConnectionPoolBuilder {
             path: ":memory:".into(),
             pool,
             join_push_down: JoinPushDown::AllowedFor(":memory:".to_string()),
-            attached_databases: Vec::new(),
+            attached_databases: Arc::new(OnceCell::new()),
             mode: Mode::Memory,
             unsupported_type_action: UnsupportedTypeAction::Error,
             connection_setup_queries: self.connection_setup_queries,
@@ -160,7 +161,7 @@ impl DuckDbConnectionPoolBuilder {
             pool,
             // Allow join-push down for any other instances that connect to the same underlying file.
             join_push_down: JoinPushDown::AllowedFor(self.path),
-            attached_databases: Vec::new(),
+            attached_databases: Arc::new(OnceCell::new()),
             mode: Mode::File,
             unsupported_type_action: UnsupportedTypeAction::Error,
             connection_setup_queries: self.connection_setup_queries,
@@ -180,7 +181,8 @@ pub struct DuckDbConnectionPool {
     path: Arc<str>,
     pool: Arc<r2d2::Pool<DuckdbConnectionManager>>,
     join_push_down: JoinPushDown,
-    attached_databases: Vec<Arc<str>>,
+    /// Shared across clones. Initialized once, first set of attached databases wins.
+    attached_databases: Arc<OnceCell<Arc<DuckDBAttachments>>>,
     mode: Mode,
     unsupported_type_action: UnsupportedTypeAction,
     connection_setup_queries: Vec<Arc<str>>,
@@ -191,7 +193,7 @@ impl std::fmt::Debug for DuckDbConnectionPool {
         f.debug_struct("DuckDbConnectionPool")
             .field("path", &self.path)
             .field("join_push_down", &self.join_push_down)
-            .field("attached_databases", &self.attached_databases)
+            .field("attached_databases", &self.attached_databases.get())
             .field("mode", &self.mode)
             .field("unsupported_type_action", &self.unsupported_type_action)
             .finish()
@@ -254,19 +256,58 @@ impl DuckDbConnectionPool {
         self
     }
 
-    #[must_use]
-    pub fn set_attached_databases(mut self, databases: &[Arc<str>]) -> Self {
-        self.attached_databases = databases.to_vec();
-
+    /// Sets the databases to attach for cross-database queries.
+    /// Attachments are performed lazily on first query using `OnceCell`.
+    ///
+    /// If attachments are already configured with the same databases, this is a no-op.
+    /// If attachments are already configured with different databases, a warning is logged
+    /// and the existing attachments are preserved.
+    ///
+    /// # Errors
+    ///
+    /// Returns an error if unable to extract database name from path.
+    pub fn set_attached_databases(mut self, databases: &[Arc<str>]) -> Result<Self> {
         if !databases.is_empty() {
-            let mut paths = self.attached_databases.clone();
+            let mut paths: Vec<Arc<str>> = databases.to_vec();
             paths.push(Arc::clone(&self.path));
             paths.sort();
-            let push_down_context = paths.join(";");
+            let push_down_context = paths
+                .iter()
+                .map(|p| p.as_ref())
+                .collect::<Vec<_>>()
+                .join(";");
             self.join_push_down = JoinPushDown::AllowedFor(push_down_context);
         }
 
-        self
+        let db_name = extract_db_name(Arc::clone(&self.path))?;
+        let new_set: std::collections::HashSet<Arc<str>> = databases.iter().cloned().collect();
+        let path = Arc::clone(&self.path);
+
+        let existing = self.attached_databases.get_or_init(|| {
+            tracing::debug!(
+                "pool_path = {}, db_name = {}, databases = {:?}, set_attached_databases: creating new DuckDBAttachments",
+                path, db_name, databases
+            );
+            Arc::new(DuckDBAttachments::new(&db_name, databases))
+        });
+
+        // Check if the existing attachments match what was requested
+        let existing_set = existing.attachments();
+        if *existing_set != new_set {
+            tracing::warn!(
+                "Unable to reconfigure DuckDB attachments for database {}: attachments are already configured with a different set of databases. \
+                 Existing: {existing_set:?}, Requested: {new_set:?}. Keeping existing attachments.",
+                self.path
+            );
+        }
+
+        Ok(self)
+    }
+
+    /// Returns the attachments configuration.
+    #[must_use]
+    pub fn get_attachments(&self) -> Option<&Arc<DuckDBAttachments>> {
+        self.attached_databases.get()
     }
 
     #[must_use]
@@ -289,16 +330,14 @@ impl DuckDbConnectionPool {
         let conn: r2d2::PooledConnection<DuckdbConnectionManager> =
             pool.get().context(ConnectionPoolSnafu)?;
 
-        let attachments = self.get_attachments()?;
-
         for query in self.connection_setup_queries.iter() {
             tracing::debug!("DuckDB connection setup: {}", query);
             conn.execute(query, []).context(DuckDBConnectionSnafu)?;
         }
 
         Ok(Box::new(
             DuckDbConnection::new(conn)
-                .with_attachments(attachments)
+                .with_attachments(self.attached_databases.get().cloned())
                 .with_connection_setup_queries(self.connection_setup_queries.clone())
                 .with_unsupported_type_action(self.unsupported_type_action),
         ))
@@ -308,21 +347,6 @@ impl DuckDbConnectionPool {
     pub fn mode(&self) -> Mode {
         self.mode
     }
-
-    pub fn get_attachments(&self) -> Result<Option<Arc<DuckDBAttachments>>> {
-        if self.attached_databases.is_empty() {
-            Ok(None)
-        } else {
-            #[cfg(not(feature = "duckdb-federation"))]
-            return Ok(None);
-
-            #[cfg(feature = "duckdb-federation")]
-            Ok(Some(Arc::new(DuckDBAttachments::new(
-                &extract_db_name(Arc::clone(&self.path))?,
-                &self.attached_databases,
-            ))))
-        }
-    }
 }
 
 #[async_trait]
@@ -338,16 +362,14 @@ impl DbConnectionPool<r2d2::PooledConnection<DuckdbConnectionManager>, DuckDBPar
         let conn: r2d2::PooledConnection<DuckdbConnectionManager> =
             pool.get().context(ConnectionPoolSnafu)?;
 
-        let attachments = self.get_attachments()?;
-
         for query in self.connection_setup_queries.iter() {
             tracing::debug!("DuckDB connection setup: {}", query);
             conn.execute(query, []).context(DuckDBConnectionSnafu)?;
         }
 
         Ok(Box::new(
             DuckDbConnection::new(conn)
-                .with_attachments(attachments)
+                .with_attachments(self.attached_databases.get().cloned())
                 .with_connection_setup_queries(self.connection_setup_queries.clone())
                 .with_unsupported_type_action(self.unsupported_type_action),
         ))
@@ -438,12 +460,14 @@ mod test {
         let db_attached_name = random_db_name();
         let pool = DuckDbConnectionPool::new_file(&db_base_name, &AccessMode::ReadWrite)
             .expect("DuckDB connection pool to be created")
-            .set_attached_databases(&[Arc::from(db_attached_name.as_str())]);
+            .set_attached_databases(&[Arc::from(db_attached_name.as_str())])
+            .expect("Attached databases should be set");
 
         let pool_attached =
             DuckDbConnectionPool::new_file(&db_attached_name, &AccessMode::ReadWrite)
                 .expect("DuckDB connection pool to be created")
-                .set_attached_databases(&[Arc::from(db_base_name.as_str())]);
+                .set_attached_databases(&[Arc::from(db_base_name.as_str())])
+                .expect("Attached databases should be set");
 
         let conn = pool
             .pool