feat: allow setting custom Session on ListingDatabase (#2526)

## 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 <noreply@anthropic.com>
This commit is contained in:
Will Jones
2025-07-21 16:28:39 -07:00
committed by GitHub
parent 75edb8756c
commit 0579303602
3 changed files with 36 additions and 2 deletions

View File

@@ -216,6 +216,7 @@ impl Catalog for ListingCatalog {
client_config: Default::default(), client_config: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
options: Default::default(), options: Default::default(),
session: None,
}; };
// Add the db options to the connect request // Add the db options to the connect request
@@ -243,6 +244,7 @@ impl Catalog for ListingCatalog {
client_config: Default::default(), client_config: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
options: Default::default(), options: Default::default(),
session: None,
}; };
// Add the db options to the connect request // Add the db options to the connect request
@@ -312,6 +314,7 @@ mod tests {
client_config: Default::default(), client_config: Default::default(),
options: Default::default(), options: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
session: None,
}; };
let catalog = ListingCatalog::connect(&request).await.unwrap(); let catalog = ListingCatalog::connect(&request).await.unwrap();
@@ -573,6 +576,7 @@ mod tests {
client_config: Default::default(), client_config: Default::default(),
options: Default::default(), options: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
session: None,
}; };
let catalog = ListingCatalog::connect(&request).await.unwrap(); let catalog = ListingCatalog::connect(&request).await.unwrap();
@@ -592,6 +596,7 @@ mod tests {
client_config: Default::default(), client_config: Default::default(),
options: Default::default(), options: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
session: None,
}; };
let catalog = ListingCatalog::connect(&request).await.unwrap(); let catalog = ListingCatalog::connect(&request).await.unwrap();
@@ -608,6 +613,7 @@ mod tests {
client_config: Default::default(), client_config: Default::default(),
options: Default::default(), options: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
session: None,
}; };
let result = ListingCatalog::connect(&request).await; let result = ListingCatalog::connect(&request).await;

View File

@@ -627,6 +627,12 @@ pub struct ConnectRequest {
/// consistency only applies to read operations. Write operations are /// consistency only applies to read operations. Write operations are
/// always consistent. /// always consistent.
pub read_consistency_interval: Option<std::time::Duration>, pub read_consistency_interval: Option<std::time::Duration>,
/// 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<Arc<lance::session::Session>>,
} }
#[derive(Debug)] #[derive(Debug)]
@@ -645,6 +651,7 @@ impl ConnectBuilder {
client_config: Default::default(), client_config: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
options: HashMap::new(), options: HashMap::new(),
session: None,
}, },
embedding_registry: None, embedding_registry: None,
} }
@@ -802,6 +809,20 @@ impl ConnectBuilder {
self 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<lance::session::Session>) -> Self {
self.request.session = Some(session);
self
}
#[cfg(feature = "remote")] #[cfg(feature = "remote")]
fn execute_remote(self) -> Result<Connection> { fn execute_remote(self) -> Result<Connection> {
use crate::remote::db::RemoteDatabaseOptions; use crate::remote::db::RemoteDatabaseOptions;
@@ -884,6 +905,7 @@ impl CatalogConnectBuilder {
client_config: Default::default(), client_config: Default::default(),
read_consistency_interval: None, read_consistency_interval: None,
options: HashMap::new(), options: HashMap::new(),
session: None,
}, },
} }
} }

View File

@@ -265,6 +265,7 @@ impl ListingDatabase {
uri, uri,
request.read_consistency_interval, request.read_consistency_interval,
options.new_table_config, options.new_table_config,
request.session.clone(),
) )
.await .await
} }
@@ -316,7 +317,10 @@ impl ListingDatabase {
let plain_uri = url.to_string(); 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 { let os_params = ObjectStoreParams {
storage_options: Some(options.storage_options.clone()), storage_options: Some(options.storage_options.clone()),
..Default::default() ..Default::default()
@@ -357,6 +361,7 @@ impl ListingDatabase {
uri, uri,
request.read_consistency_interval, request.read_consistency_interval,
options.new_table_config, options.new_table_config,
request.session.clone(),
) )
.await .await
} }
@@ -367,8 +372,9 @@ impl ListingDatabase {
path: &str, path: &str,
read_consistency_interval: Option<std::time::Duration>, read_consistency_interval: Option<std::time::Duration>,
new_table_config: NewTableConfig, new_table_config: NewTableConfig,
session: Option<Arc<lance::session::Session>>,
) -> Result<Self> { ) -> Result<Self> {
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( let (object_store, base_path) = ObjectStore::from_uri_and_params(
session.store_registry(), session.store_registry(),
path, path,