From 05793036026b3ccc7e137ce80393842b5f959964 Mon Sep 17 00:00:00 2001 From: Will Jones Date: Mon, 21 Jul 2025 16:28:39 -0700 Subject: [PATCH] feat: allow setting custom Session on ListingDatabase (#2526) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Summary Add support for providing a custom `Session` when connecting to a `ListingDatabase`. This allows users to configure object store registries, caching, and other session-related settings while maintaining full backward compatibility. ## Usage Example ```rust use std::sync::Arc; use lancedb::connect; let custom_session = Arc::new(lance::session::Session::default()); let db = connect("/path/to/database") .session(custom_session) .execute() .await?; ``` 🤖 Generated with [Claude Code](https://claude.ai/code) Co-authored-by: Claude --- rust/lancedb/src/catalog/listing.rs | 6 ++++++ rust/lancedb/src/connection.rs | 22 ++++++++++++++++++++++ rust/lancedb/src/database/listing.rs | 10 ++++++++-- 3 files changed, 36 insertions(+), 2 deletions(-) diff --git a/rust/lancedb/src/catalog/listing.rs b/rust/lancedb/src/catalog/listing.rs index efbe87c6..c78336f0 100644 --- a/rust/lancedb/src/catalog/listing.rs +++ b/rust/lancedb/src/catalog/listing.rs @@ -216,6 +216,7 @@ impl Catalog for ListingCatalog { client_config: Default::default(), read_consistency_interval: None, options: Default::default(), + session: None, }; // Add the db options to the connect request @@ -243,6 +244,7 @@ impl Catalog for ListingCatalog { client_config: Default::default(), read_consistency_interval: None, options: Default::default(), + session: None, }; // Add the db options to the connect request @@ -312,6 +314,7 @@ mod tests { client_config: Default::default(), options: Default::default(), read_consistency_interval: None, + session: None, }; let catalog = ListingCatalog::connect(&request).await.unwrap(); @@ -573,6 +576,7 @@ mod tests { client_config: Default::default(), options: Default::default(), read_consistency_interval: None, + session: None, }; let catalog = ListingCatalog::connect(&request).await.unwrap(); @@ -592,6 +596,7 @@ mod tests { client_config: Default::default(), options: Default::default(), read_consistency_interval: None, + session: None, }; let catalog = ListingCatalog::connect(&request).await.unwrap(); @@ -608,6 +613,7 @@ mod tests { client_config: Default::default(), options: Default::default(), read_consistency_interval: None, + session: None, }; let result = ListingCatalog::connect(&request).await; diff --git a/rust/lancedb/src/connection.rs b/rust/lancedb/src/connection.rs index aae34204..68e2e98e 100644 --- a/rust/lancedb/src/connection.rs +++ b/rust/lancedb/src/connection.rs @@ -627,6 +627,12 @@ pub struct ConnectRequest { /// consistency only applies to read operations. Write operations are /// always consistent. pub read_consistency_interval: Option, + + /// Optional session for object stores and caching + /// + /// If provided, this session will be used instead of creating a default one. + /// This allows for custom configuration of object store registries, caching, etc. + pub session: Option>, } #[derive(Debug)] @@ -645,6 +651,7 @@ impl ConnectBuilder { client_config: Default::default(), read_consistency_interval: None, options: HashMap::new(), + session: None, }, embedding_registry: None, } @@ -802,6 +809,20 @@ impl ConnectBuilder { self } + /// Set a custom session for object stores and caching. + /// + /// By default, a new session with default configuration will be created. + /// This method allows you to provide a custom session with your own + /// configuration for object store registries, caching, etc. + /// + /// # Arguments + /// + /// * `session` - A custom session to use for this connection + pub fn session(mut self, session: Arc) -> Self { + self.request.session = Some(session); + self + } + #[cfg(feature = "remote")] fn execute_remote(self) -> Result { use crate::remote::db::RemoteDatabaseOptions; @@ -884,6 +905,7 @@ impl CatalogConnectBuilder { client_config: Default::default(), read_consistency_interval: None, options: HashMap::new(), + session: None, }, } } diff --git a/rust/lancedb/src/database/listing.rs b/rust/lancedb/src/database/listing.rs index 1e016a1a..d63e46f5 100644 --- a/rust/lancedb/src/database/listing.rs +++ b/rust/lancedb/src/database/listing.rs @@ -265,6 +265,7 @@ impl ListingDatabase { uri, request.read_consistency_interval, options.new_table_config, + request.session.clone(), ) .await } @@ -316,7 +317,10 @@ impl ListingDatabase { let plain_uri = url.to_string(); - let session = Arc::new(lance::session::Session::default()); + let session = request + .session + .clone() + .unwrap_or_else(|| Arc::new(lance::session::Session::default())); let os_params = ObjectStoreParams { storage_options: Some(options.storage_options.clone()), ..Default::default() @@ -357,6 +361,7 @@ impl ListingDatabase { uri, request.read_consistency_interval, options.new_table_config, + request.session.clone(), ) .await } @@ -367,8 +372,9 @@ impl ListingDatabase { path: &str, read_consistency_interval: Option, new_table_config: NewTableConfig, + session: Option>, ) -> Result { - let session = Arc::new(lance::session::Session::default()); + let session = session.unwrap_or_else(|| Arc::new(lance::session::Session::default())); let (object_store, base_path) = ObjectStore::from_uri_and_params( session.store_registry(), path,