use std::collections::HashSet;
use std::fs::create_dir_all;
use std::path::Path;
use std::{collections::HashMap, sync::Arc};
use lance::dataset::refs::Ref;
use lance::dataset::{ReadParams, WriteMode, builder::DatasetBuilder};
use lance::io::{ObjectStore, ObjectStoreParams, WrappingObjectStore};
use lance_datafusion::utils::StreamingWriteSource;
use lance_file::version::LanceFileVersion;
use lance_io::object_store::{StorageOptionsAccessor, StorageOptionsProvider};
use lance_table::io::commit::commit_handler_from_url;
use object_store::local::LocalFileSystem;
use snafu::ResultExt;
use crate::blob::{ensure_blob_storage_version, has_blob_columns};
use crate::connection::ConnectRequest;
use crate::database::ReadConsistency;
use crate::database::namespace::LanceNamespaceDatabase;
use crate::error::{CreateDirSnafu, Error, Result};
use crate::io::object_store::MirroringObjectStoreWrapper;
use crate::table::NativeTable;
use crate::utils::validate_table_name;
use lance_namespace::models::{
CreateNamespaceRequest, CreateNamespaceResponse, DescribeNamespaceRequest,
DescribeNamespaceResponse, DropNamespaceRequest, DropNamespaceResponse, ListNamespacesRequest,
ListNamespacesResponse, ListTablesRequest, ListTablesResponse,
};
use super::{
BaseTable, CloneTableRequest, CreateTableMode, CreateTableRequest, Database, DatabaseOptions,
OpenTableRequest, TableNamesRequest,
};
pub const LANCE_FILE_EXTENSION: &str = "lance";
pub const OPT_NEW_TABLE_STORAGE_VERSION: &str = "new_table_data_storage_version";
pub const OPT_NEW_TABLE_V2_MANIFEST_PATHS: &str = "new_table_enable_v2_manifest_paths";
pub const OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS: &str = "new_table_enable_stable_row_ids";
#[derive(Clone, Debug, Default)]
pub struct NewTableConfig {
pub data_storage_version: Option<LanceFileVersion>,
pub enable_v2_manifest_paths: Option<bool>,
pub enable_stable_row_ids: Option<bool>,
}
#[derive(Debug, Default, Clone)]
pub struct ListingDatabaseOptions {
pub new_table_config: NewTableConfig,
pub storage_options: HashMap<String, String>,
}
impl ListingDatabaseOptions {
pub fn builder() -> ListingDatabaseOptionsBuilder {
ListingDatabaseOptionsBuilder::new()
}
pub(crate) fn parse_from_map(map: &HashMap<String, String>) -> Result<Self> {
let new_table_config = NewTableConfig {
data_storage_version: map
.get(OPT_NEW_TABLE_STORAGE_VERSION)
.map(|s| s.parse())
.transpose()?,
enable_v2_manifest_paths: map
.get(OPT_NEW_TABLE_V2_MANIFEST_PATHS)
.map(|s| {
s.parse::<bool>().map_err(|_| Error::InvalidInput {
message: format!(
"enable_v2_manifest_paths must be a boolean, received {}",
s
),
})
})
.transpose()?,
enable_stable_row_ids: map
.get(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS)
.map(|s| {
s.parse::<bool>().map_err(|_| Error::InvalidInput {
message: format!("enable_stable_row_ids must be a boolean, received {}", s),
})
})
.transpose()?,
};
let storage_options = map
.iter()
.filter(|(key, _)| {
key.as_str() != OPT_NEW_TABLE_STORAGE_VERSION
&& key.as_str() != OPT_NEW_TABLE_V2_MANIFEST_PATHS
&& key.as_str() != OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS
})
.map(|(key, value)| (key.clone(), value.clone()))
.collect();
Ok(Self {
new_table_config,
storage_options,
})
}
}
impl DatabaseOptions for ListingDatabaseOptions {
fn serialize_into_map(&self, map: &mut HashMap<String, String>) {
if let Some(storage_version) = &self.new_table_config.data_storage_version {
map.insert(
OPT_NEW_TABLE_STORAGE_VERSION.to_string(),
storage_version.to_string(),
);
}
if let Some(enable_v2_manifest_paths) = self.new_table_config.enable_v2_manifest_paths {
map.insert(
OPT_NEW_TABLE_V2_MANIFEST_PATHS.to_string(),
enable_v2_manifest_paths.to_string(),
);
}
if let Some(enable_stable_row_ids) = self.new_table_config.enable_stable_row_ids {
map.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
enable_stable_row_ids.to_string(),
);
}
}
}
#[derive(Clone, Debug, Default)]
pub struct ListingDatabaseOptionsBuilder {
options: ListingDatabaseOptions,
}
impl ListingDatabaseOptionsBuilder {
pub fn new() -> Self {
Self {
options: ListingDatabaseOptions::default(),
}
}
}
impl ListingDatabaseOptionsBuilder {
pub fn data_storage_version(mut self, data_storage_version: LanceFileVersion) -> Self {
self.options.new_table_config.data_storage_version = Some(data_storage_version);
self
}
pub fn enable_v2_manifest_paths(mut self, enable_v2_manifest_paths: bool) -> Self {
self.options.new_table_config.enable_v2_manifest_paths = Some(enable_v2_manifest_paths);
self
}
pub fn storage_option(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.options
.storage_options
.insert(key.into(), value.into());
self
}
pub fn storage_options(
mut self,
pairs: impl IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
) -> Self {
for (key, value) in pairs {
self.options
.storage_options
.insert(key.into(), value.into());
}
self
}
pub fn build(self) -> ListingDatabaseOptions {
self.options
}
}
#[derive(Debug)]
pub struct ListingDatabase {
object_store: Arc<ObjectStore>,
query_string: Option<String>,
pub(crate) uri: String,
pub(crate) base_path: object_store::path::Path,
pub(crate) store_wrapper: Option<Arc<dyn WrappingObjectStore>>,
read_consistency_interval: Option<std::time::Duration>,
storage_options: HashMap<String, String>,
pub(crate) storage_options_provider: Option<Arc<dyn StorageOptionsProvider>>,
new_table_config: NewTableConfig,
session: Arc<lance::session::Session>,
namespace_database: Arc<LanceNamespaceDatabase>,
}
impl std::fmt::Display for ListingDatabase {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"ListingDatabase(uri={}, read_consistency_interval={})",
self.uri,
match self.read_consistency_interval {
None => {
"None".to_string()
}
Some(duration) => {
format!("{}s", duration.as_secs_f64())
}
}
)
}
}
const LANCE_EXTENSION: &str = "lance";
fn table_name(location: &object_store::path::Path, dir_suffix: &str) -> Option<String> {
location
.filename()?
.strip_suffix(dir_suffix)
.map(String::from)
.filter(|name| !name.is_empty())
}
struct DirPage {
common_prefixes: Vec<object_store::path::Path>,
page_token: Option<String>,
}
fn relative_key<'a>(prefix: Option<&str>, location: &'a str) -> Option<&'a str> {
let relative = match prefix {
Some(prefix) => location.strip_prefix(prefix)?,
None => location,
};
(!relative.is_empty()).then_some(relative)
}
async fn read_dir_page(
object_store: &ObjectStore,
base_path: &object_store::path::Path,
page_token: Option<String>,
limit: Option<usize>,
) -> Result<DirPage> {
let listed = object_store.list_with_delimiter(Some(base_path)).await?;
let prefix = {
let base = base_path.as_ref();
(!base.is_empty()).then(|| format!("{base}/"))
};
let table_dir_suffix = format!(".{LANCE_EXTENSION}/");
let mut children: Vec<(String, object_store::path::Path)> = listed
.common_prefixes
.into_iter()
.filter_map(|location| {
let key = format!("{}/", relative_key(prefix.as_deref(), location.as_ref())?);
(key.len() > table_dir_suffix.len() && key.ends_with(&table_dir_suffix))
.then_some((key, location))
})
.collect();
children.sort_unstable_by(|(left, _), (right, _)| left.cmp(right));
if let Some(resume) = &page_token {
children.retain(|(key, _)| key > resume);
}
let total = children.len();
children.truncate(limit.unwrap_or(total).min(total));
let page_token = match children.last() {
Some((last, _)) if children.len() < total => Some(last.clone()),
_ => None,
};
Ok(DirPage {
common_prefixes: children.into_iter().map(|(_, location)| location).collect(),
page_token,
})
}
const ENGINE: &str = "engine";
const MIRRORED_STORE: &str = "mirroredStore";
impl ListingDatabase {
pub(crate) fn build_namespace_client_properties(
uri: &str,
storage_options: &HashMap<String, String>,
namespace_client_properties: HashMap<String, String>,
) -> HashMap<String, String> {
let mut properties = namespace_client_properties;
properties.insert("root".to_string(), uri.to_string());
for (key, value) in storage_options {
properties.insert(format!("storage.{}", key), value.clone());
}
properties
}
pub(crate) fn build_manifest_enabled_namespace_client_properties(
uri: &str,
storage_options: &HashMap<String, String>,
namespace_client_properties: HashMap<String, String>,
) -> HashMap<String, String> {
let mut properties = Self::build_namespace_client_properties(
uri,
storage_options,
namespace_client_properties,
);
properties.insert("manifest_enabled".to_string(), "true".to_string());
properties.insert(
"dir_listing_to_manifest_migration_enabled".to_string(),
"true".to_string(),
);
properties
}
async fn connect_namespace_database(
uri: &str,
storage_options: HashMap<String, String>,
namespace_client_properties: HashMap<String, String>,
read_consistency_interval: Option<std::time::Duration>,
session: Arc<lance::session::Session>,
) -> Result<Arc<LanceNamespaceDatabase>> {
let ns_properties = Self::build_namespace_client_properties(
uri,
&storage_options,
namespace_client_properties,
);
Ok(Arc::new(
LanceNamespaceDatabase::connect(
"dir",
ns_properties,
storage_options,
read_consistency_interval,
Some(session),
HashSet::new(),
)
.await?,
))
}
fn storage_base_uri(uri: &str) -> String {
let Ok(mut url) = url::Url::parse(uri) else {
return uri.to_string();
};
url.set_query(None);
let Some((storage_scheme, _commit_scheme)) = url.scheme().split_once('+') else {
return url.to_string();
};
let storage_scheme = storage_scheme.to_string();
let _ = url.set_scheme(&storage_scheme);
url.to_string()
}
async fn prepare_namespace_root(
uri: &str,
storage_options: &HashMap<String, String>,
session: Arc<lance::session::Session>,
) -> Result<String> {
match url::Url::parse(uri) {
Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
let (object_store, _) = ObjectStore::from_uri_and_params(
session.store_registry(),
uri,
&ObjectStoreParams::default(),
)
.await?;
if object_store.is_local() {
Self::try_create_dir(uri).context(CreateDirSnafu { path: uri })?;
}
Ok(uri.to_string())
}
Ok(mut url) => {
if url.scheme().contains('+') {
return Err(Error::NotSupported {
message: "commit engine URI schemes are not supported for manifest-enabled namespace connections".to_string(),
});
}
for (key, value) in url.query_pairs() {
if key == ENGINE {
return Err(Error::NotSupported {
message: format!(
"commit engine '{}' is not supported for manifest-enabled namespace connections",
value
),
});
} else if key == MIRRORED_STORE {
return Err(Error::NotSupported {
message: "mirrored store is not supported for manifest-enabled namespace connections"
.to_string(),
});
}
}
url.set_query(None);
let plain_uri = url.to_string();
let os_params = ObjectStoreParams {
storage_options_accessor: if storage_options.is_empty() {
None
} else {
Some(Arc::new(StorageOptionsAccessor::with_static_options(
storage_options.clone(),
)))
},
..Default::default()
};
let (object_store, _) = ObjectStore::from_uri_and_params(
session.store_registry(),
&plain_uri,
&os_params,
)
.await?;
if object_store.is_local() {
Self::try_create_dir(&plain_uri).context(CreateDirSnafu {
path: plain_uri.clone(),
})?;
}
Ok(plain_uri)
}
Err(_) => {
let (object_store, _) = ObjectStore::from_uri_and_params(
session.store_registry(),
uri,
&ObjectStoreParams::default(),
)
.await?;
if object_store.is_local() {
Self::try_create_dir(uri).context(CreateDirSnafu { path: uri })?;
}
Ok(uri.to_string())
}
}
}
pub(crate) async fn connect_manifest_enabled_namespace_database(
request: &ConnectRequest,
) -> Result<LanceNamespaceDatabase> {
let options = ListingDatabaseOptions::parse_from_map(&request.options)?;
let session = request
.session
.clone()
.unwrap_or_else(|| Arc::new(lance::session::Session::default()));
let namespace_root =
Self::prepare_namespace_root(&request.uri, &options.storage_options, session.clone())
.await?;
let ns_properties = Self::build_manifest_enabled_namespace_client_properties(
&namespace_root,
&options.storage_options,
request.namespace_client_properties.clone(),
);
LanceNamespaceDatabase::connect_with_new_table_config(
"dir",
ns_properties,
options.storage_options,
request.read_consistency_interval,
Some(session),
HashSet::new(),
options.new_table_config,
)
.await
.map(|db| db.with_uri(request.uri.clone()))
}
pub async fn connect_with_options(request: &ConnectRequest) -> Result<Self> {
let uri = &request.uri;
let parse_res = url::Url::parse(uri);
let options = ListingDatabaseOptions::parse_from_map(&request.options)?;
match parse_res {
Ok(url) if url.scheme().len() == 1 && cfg!(windows) => {
Self::open_path(
uri,
request.read_consistency_interval,
options.new_table_config,
request.namespace_client_properties.clone(),
request.session.clone(),
)
.await
}
Ok(mut url) => {
let mut engine = None;
let mut mirrored_store = None;
let mut filtered_querys = vec![];
for (key, value) in url.query_pairs() {
if key == ENGINE {
engine = Some(value.to_string());
} else if key == MIRRORED_STORE {
if cfg!(windows) {
return Err(Error::NotSupported {
message: "mirrored store is not supported on windows".into(),
});
}
mirrored_store = Some(value.to_string());
} else {
filtered_querys.push((key.to_string(), value.to_string()));
}
}
url.query_pairs_mut().clear();
url.query_pairs_mut().extend_pairs(filtered_querys);
let query_string = url.query().filter(|q| !q.is_empty()).map(|s| s.to_string());
url.set_query(None);
let storage_base_uri = Self::storage_base_uri(url.as_str());
let table_base_uri = if let Some(store) = engine {
static WARN_ONCE: std::sync::Once = std::sync::Once::new();
WARN_ONCE.call_once(|| {
log::warn!("Specifying engine is not a publicly supported feature in lancedb yet. THE API WILL CHANGE");
});
let old_scheme = url.scheme().to_string();
let new_scheme = format!("{}+{}", old_scheme, store);
url.to_string().replacen(&old_scheme, &new_scheme, 1)
} else {
url.to_string()
};
let session = request
.session
.clone()
.unwrap_or_else(|| Arc::new(lance::session::Session::default()));
let os_params = ObjectStoreParams {
storage_options_accessor: if options.storage_options.is_empty() {
None
} else {
Some(Arc::new(StorageOptionsAccessor::with_static_options(
options.storage_options.clone(),
)))
},
..Default::default()
};
let (object_store, base_path) = ObjectStore::from_uri_and_params(
session.store_registry(),
&storage_base_uri,
&os_params,
)
.await?;
if object_store.is_local() {
Self::try_create_dir(&storage_base_uri).context(CreateDirSnafu {
path: storage_base_uri.clone(),
})?;
}
let write_store_wrapper = match mirrored_store {
Some(path) => {
let mirrored_store = Arc::new(LocalFileSystem::new_with_prefix(path)?);
let wrapper = MirroringObjectStoreWrapper::new(mirrored_store);
Some(Arc::new(wrapper) as Arc<dyn WrappingObjectStore>)
}
None => None,
};
let namespace_database = Self::connect_namespace_database(
&storage_base_uri,
options.storage_options.clone(),
request.namespace_client_properties.clone(),
request.read_consistency_interval,
session.clone(),
)
.await?;
Ok(Self {
uri: table_base_uri,
query_string,
base_path,
object_store,
store_wrapper: write_store_wrapper,
read_consistency_interval: request.read_consistency_interval,
storage_options: options.storage_options,
storage_options_provider: None,
new_table_config: options.new_table_config,
session,
namespace_database,
})
}
Err(_) => {
Self::open_path(
uri,
request.read_consistency_interval,
options.new_table_config,
request.namespace_client_properties.clone(),
request.session.clone(),
)
.await
}
}
}
async fn open_path(
path: &str,
read_consistency_interval: Option<std::time::Duration>,
new_table_config: NewTableConfig,
namespace_client_properties: HashMap<String, String>,
session: Option<Arc<lance::session::Session>>,
) -> Result<Self> {
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,
&ObjectStoreParams::default(),
)
.await?;
if object_store.is_local() {
Self::try_create_dir(path).context(CreateDirSnafu { path })?;
}
let namespace_database = Self::connect_namespace_database(
path,
HashMap::new(),
namespace_client_properties,
read_consistency_interval,
session.clone(),
)
.await?;
Ok(Self {
uri: path.to_string(),
query_string: None,
base_path,
object_store,
store_wrapper: None,
read_consistency_interval,
storage_options: HashMap::new(),
storage_options_provider: None,
new_table_config,
session,
namespace_database,
})
}
fn try_create_dir(path: &str) -> core::result::Result<(), std::io::Error> {
let fs_path = if let Some(stripped) = path.strip_prefix("file://") {
stripped
} else if let Some(stripped) = path.strip_prefix("file:") {
stripped
} else {
path
};
let path = Path::new(fs_path);
if !path.try_exists()? {
create_dir_all(path)?;
}
Ok(())
}
fn table_uri(&self, name: &str) -> Result<String> {
validate_table_name(name)?;
let mut uri = self.uri.clone();
let has_scheme = uri.contains("://");
let ends_with_separator = uri.ends_with('/') || uri.ends_with('\\');
if !ends_with_separator {
if has_scheme {
uri.push('/');
} else {
uri.push(std::path::MAIN_SEPARATOR);
}
}
uri.push_str(&format!("{}.{}", name, LANCE_FILE_EXTENSION));
if let Some(query) = self.query_string.as_ref() {
uri.push('?');
uri.push_str(query.as_str());
}
Ok(uri)
}
fn namespace_database(&self) -> Arc<LanceNamespaceDatabase> {
self.namespace_database.clone()
}
async fn drop_tables(&self, names: Vec<String>) -> Result<()> {
let object_store_params = ObjectStoreParams {
storage_options_accessor: if self.storage_options.is_empty() {
None
} else {
Some(Arc::new(StorageOptionsAccessor::with_static_options(
self.storage_options.clone(),
)))
},
..Default::default()
};
let mut uri = self.uri.clone();
if let Some(query_string) = &self.query_string {
uri.push_str(&format!("?{}", query_string));
}
let commit_handler = commit_handler_from_url(&uri, &Some(object_store_params)).await?;
for name in names {
let dir_name = format!("{}.{}", name, LANCE_EXTENSION);
let full_path = self.base_path.clone().join(dir_name.clone());
commit_handler.delete(&full_path).await?;
self.object_store
.remove_dir_all(full_path.clone())
.await
.map_err(|err| match err {
lance::Error::NotFound { .. } => Error::TableNotFound {
name: name.clone(),
source: Box::new(err),
},
_ => Error::from(err),
})?;
}
Ok(())
}
fn inherit_storage_options(&self, target: &mut HashMap<String, String>) {
for (key, value) in self.storage_options.iter() {
if !target.contains_key(key) {
target.insert(key.clone(), value.clone());
}
}
}
fn prepare_write_params(
&self,
request: &CreateTableRequest,
mut write_params: lance::dataset::WriteParams,
overrides: NewTableConfig,
) -> lance::dataset::WriteParams {
if !self.storage_options.is_empty() || self.storage_options_provider.is_some() {
let store_params = write_params
.store_params
.get_or_insert_with(Default::default);
let mut storage_options = store_params.storage_options().cloned().unwrap_or_default();
if !self.storage_options.is_empty() {
self.inherit_storage_options(&mut storage_options);
}
let accessor = if let Some(ref provider) = self.storage_options_provider {
StorageOptionsAccessor::with_initial_and_provider(storage_options, provider.clone())
} else {
StorageOptionsAccessor::with_static_options(storage_options)
};
store_params.storage_options_accessor = Some(Arc::new(accessor));
}
write_params.data_storage_version = overrides
.data_storage_version
.or(write_params.data_storage_version)
.or(self.new_table_config.data_storage_version);
if let Some(enable_v2_manifest_paths) = overrides
.enable_v2_manifest_paths
.or(self.new_table_config.enable_v2_manifest_paths)
{
write_params.enable_v2_manifest_paths = enable_v2_manifest_paths;
}
let data_schema = request.data.arrow_schema();
if let Some(enable_stable_row_ids) = overrides
.enable_stable_row_ids
.or(self.new_table_config.enable_stable_row_ids)
.or(has_blob_columns(&data_schema).then_some(true))
{
write_params.enable_stable_row_ids = enable_stable_row_ids;
}
ensure_blob_storage_version(&data_schema, &mut write_params);
if matches!(&request.mode, CreateTableMode::Overwrite) {
write_params.mode = WriteMode::Overwrite;
}
write_params.session = Some(self.session.clone());
write_params
}
async fn handle_table_exists(
&self,
table_name: &str,
namespace_path: Vec<String>,
mode: CreateTableMode,
data_schema: &arrow_schema::Schema,
) -> Result<Arc<dyn BaseTable>> {
match mode {
CreateTableMode::Create => Err(Error::TableAlreadyExists {
name: table_name.to_string(),
}),
CreateTableMode::ExistOk(callback) => {
let req = OpenTableRequest {
name: table_name.to_string(),
namespace_path: namespace_path.clone(),
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
};
let req = (callback)(req);
let table = self.open_table(req).await?;
let table_schema = table.schema().await?;
if table_schema.as_ref() != data_schema {
return Err(Error::Schema {
message: "Provided schema does not match existing table schema".to_string(),
});
}
Ok(table)
}
CreateTableMode::Overwrite => unreachable!(),
}
}
}
#[async_trait::async_trait]
impl Database for ListingDatabase {
async fn list_namespaces(
&self,
request: ListNamespacesRequest,
) -> Result<ListNamespacesResponse> {
self.namespace_database().list_namespaces(request).await
}
fn uri(&self) -> &str {
&self.uri
}
async fn read_consistency(&self) -> Result<ReadConsistency> {
if let Some(read_consistency_inverval) = self.read_consistency_interval {
if read_consistency_inverval.is_zero() {
Ok(ReadConsistency::Strong)
} else {
Ok(ReadConsistency::Eventual(read_consistency_inverval))
}
} else {
Ok(ReadConsistency::Manual)
}
}
async fn create_namespace(
&self,
request: CreateNamespaceRequest,
) -> Result<CreateNamespaceResponse> {
self.namespace_database().create_namespace(request).await
}
async fn drop_namespace(&self, request: DropNamespaceRequest) -> Result<DropNamespaceResponse> {
self.namespace_database().drop_namespace(request).await
}
async fn describe_namespace(
&self,
request: DescribeNamespaceRequest,
) -> Result<DescribeNamespaceResponse> {
self.namespace_database().describe_namespace(request).await
}
#[allow(deprecated)]
async fn table_names(&self, request: TableNamesRequest) -> Result<Vec<String>> {
if !request.namespace_path.is_empty() {
return self.namespace_database().table_names(request).await;
}
let mut f = self
.object_store
.read_dir(self.base_path.clone())
.await?
.iter()
.map(Path::new)
.filter(|path| {
let is_lance = path
.extension()
.and_then(|e| e.to_str())
.map(|e| e == LANCE_EXTENSION);
is_lance.unwrap_or(false)
})
.filter_map(|p| p.file_stem().and_then(|s| s.to_str().map(String::from)))
.collect::<Vec<String>>();
f.sort();
if let Some(start_after) = request.start_after {
let index = f
.iter()
.position(|name| name.as_str() > start_after.as_str())
.unwrap_or(f.len());
f.drain(0..index);
}
if let Some(limit) = request.limit {
f.truncate(limit as usize);
}
Ok(f)
}
async fn list_tables(&self, request: ListTablesRequest) -> Result<ListTablesResponse> {
if request.id.as_ref().map(|v| !v.is_empty()).unwrap_or(false) {
return self.namespace_database().list_tables(request).await;
}
let limit = request.limit.map(|limit| limit.max(0) as usize);
let dir_suffix = format!(".{LANCE_EXTENSION}");
let mut tables = Vec::new();
let mut page_token = request.page_token.filter(|token| !token.is_empty());
if limit == Some(0) {
return Ok(ListTablesResponse {
context: None,
tables,
page_token: None,
});
}
let page = read_dir_page(
&self.object_store,
&self.base_path,
page_token.take(),
limit,
)
.await?;
page_token = page.page_token;
tables.extend(
page.common_prefixes
.iter()
.filter_map(|location| table_name(location, &dir_suffix)),
);
Ok(ListTablesResponse {
context: None,
tables,
page_token,
})
}
async fn create_table(&self, request: CreateTableRequest) -> Result<Arc<dyn BaseTable>> {
if !request.namespace_path.is_empty() {
return self.namespace_database().create_table(request).await;
}
let table_uri = request
.location
.clone()
.unwrap_or_else(|| self.table_uri(&request.name).unwrap());
let mut write_params = request
.write_options
.lance_write_params
.clone()
.unwrap_or_default();
let overrides = take_request_creation_overrides(&mut write_params)?;
let write_params = self.prepare_write_params(&request, write_params, overrides);
let data_schema = request.data.arrow_schema();
match NativeTable::create(
&table_uri,
&request.name,
request.namespace_path.clone(),
request.data,
self.store_wrapper.clone(),
Some(write_params),
self.read_consistency_interval,
request.namespace_client,
HashSet::new(), )
.await
{
Ok(table) => Ok(Arc::new(table)),
Err(Error::TableAlreadyExists { .. }) => {
self.handle_table_exists(
&request.name,
request.namespace_path.clone(),
request.mode,
&data_schema,
)
.await
}
Err(err) => Err(err),
}
}
async fn clone_table(&self, request: CloneTableRequest) -> Result<Arc<dyn BaseTable>> {
if !request.target_namespace_path.is_empty() {
return Err(Error::NotSupported {
message: "Namespace parameter is not supported for listing database. Only root namespace is supported.".into(),
});
}
if !request.is_shallow {
return Err(Error::NotSupported {
message: "Deep clone is not yet implemented".to_string(),
});
}
validate_table_name(&request.target_table_name)?;
let storage_params = ObjectStoreParams {
storage_options_accessor: if self.storage_options.is_empty() {
None
} else {
Some(Arc::new(StorageOptionsAccessor::with_static_options(
self.storage_options.clone(),
)))
},
..Default::default()
};
let read_params = ReadParams {
store_options: Some(storage_params.clone()),
session: Some(self.session.clone()),
..Default::default()
};
let mut source_dataset = DatasetBuilder::from_uri(&request.source_uri)
.with_read_params(read_params.clone())
.load()
.await
.map_err(|e| -> Error { e.into() })?;
let version_ref = match (request.source_version, request.source_tag) {
(Some(v), None) => Ok(Ref::Version(None, Some(v))),
(None, Some(tag)) => Ok(Ref::Tag(tag)),
(None, None) => Ok(Ref::Version(None, Some(source_dataset.version().version))),
_ => Err(Error::InvalidInput {
message: "Cannot specify both source_version and source_tag".to_string(),
}),
}?;
let target_uri = self.table_uri(&request.target_table_name)?;
source_dataset
.shallow_clone(&target_uri, version_ref, Some(storage_params))
.await
.map_err(|e| -> Error { e.into() })?;
let cloned_table = NativeTable::open_with_params(
&target_uri,
&request.target_table_name,
request.target_namespace_path,
self.store_wrapper.clone(),
None,
self.read_consistency_interval,
request.namespace_client,
HashSet::new(), None, )
.await?;
Ok(Arc::new(cloned_table))
}
async fn open_table(&self, mut request: OpenTableRequest) -> Result<Arc<dyn BaseTable>> {
if !request.namespace_path.is_empty() {
return self.namespace_database().open_table(request).await;
}
let table_uri = request
.location
.clone()
.unwrap_or_else(|| self.table_uri(&request.name).unwrap());
if !self.storage_options.is_empty() || self.storage_options_provider.is_some() {
let store_params = request
.lance_read_params
.get_or_insert_with(Default::default)
.store_options
.get_or_insert_with(Default::default);
let mut storage_options = store_params.storage_options().cloned().unwrap_or_default();
if !self.storage_options.is_empty() {
self.inherit_storage_options(&mut storage_options);
}
let request_provider = store_params
.storage_options_accessor
.as_ref()
.and_then(|a| a.provider().cloned());
let provider = self.storage_options_provider.clone().or(request_provider);
let accessor = if let Some(provider) = provider {
StorageOptionsAccessor::with_initial_and_provider(storage_options, provider)
} else {
StorageOptionsAccessor::with_static_options(storage_options)
};
store_params.storage_options_accessor = Some(Arc::new(accessor));
}
let mut read_params = request.lance_read_params.unwrap_or_else(|| {
let mut default_params = ReadParams::default();
if let Some(index_cache_size) = request.index_cache_size {
#[allow(deprecated)]
default_params.index_cache_size(index_cache_size as usize);
}
default_params
});
read_params.session(self.session.clone());
let native_table = Arc::new(
NativeTable::open_with_params(
&table_uri,
&request.name,
request.namespace_path,
self.store_wrapper.clone(),
Some(read_params),
self.read_consistency_interval,
request.namespace_client,
HashSet::new(), request.managed_versioning, )
.await?,
);
Ok(native_table)
}
async fn rename_table(
&self,
_cur_name: &str,
_new_name: &str,
cur_namespace_path: &[String],
new_namespace_path: &[String],
) -> Result<()> {
if !cur_namespace_path.is_empty() {
return Err(Error::NotSupported {
message: "Namespace parameter is not supported for listing database.".into(),
});
}
if !new_namespace_path.is_empty() {
return Err(Error::NotSupported {
message: "Namespace parameter is not supported for listing database.".into(),
});
}
Err(Error::NotSupported {
message: "rename_table is not supported in LanceDB OSS".into(),
})
}
async fn drop_table(&self, name: &str, namespace_path: &[String]) -> Result<()> {
if !namespace_path.is_empty() {
return self
.namespace_database()
.drop_table(name, namespace_path)
.await;
}
self.drop_tables(vec![name.to_string()]).await
}
#[allow(deprecated)]
async fn drop_all_tables(&self, namespace_path: &[String]) -> Result<()> {
if !namespace_path.is_empty() {
return self
.namespace_database()
.drop_all_tables(namespace_path)
.await;
}
let tables = self.table_names(TableNamesRequest::default()).await?;
self.drop_tables(tables).await
}
fn as_any(&self) -> &dyn std::any::Any {
self
}
async fn namespace_client(&self) -> Result<Arc<dyn lance_namespace::LanceNamespace>> {
self.namespace_database.namespace_client().await
}
async fn namespace_client_config(&self) -> Result<(String, HashMap<String, String>)> {
self.namespace_database.namespace_client_config().await
}
}
pub(crate) fn take_request_creation_overrides(
params: &mut lance::dataset::WriteParams,
) -> Result<NewTableConfig> {
let storage_options = params
.store_params
.as_ref()
.and_then(|sp| sp.storage_options());
let overrides = NewTableConfig {
data_storage_version: storage_options
.and_then(|opts| opts.get(OPT_NEW_TABLE_STORAGE_VERSION))
.map(|s| s.parse::<LanceFileVersion>())
.transpose()?,
enable_v2_manifest_paths: storage_options
.and_then(|opts| opts.get(OPT_NEW_TABLE_V2_MANIFEST_PATHS))
.map(|s| s.parse::<bool>())
.transpose()
.map_err(|_| Error::InvalidInput {
message: "enable_v2_manifest_paths must be a boolean".to_string(),
})?,
enable_stable_row_ids: storage_options
.and_then(|opts| opts.get(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS))
.map(|s| s.parse::<bool>())
.transpose()
.map_err(|_| Error::InvalidInput {
message: "enable_stable_row_ids must be a boolean".to_string(),
})?,
};
if let Some(store_params) = params.store_params.as_mut() {
strip_new_table_creation_keys(store_params);
}
Ok(overrides)
}
fn strip_new_table_creation_keys(store_params: &mut ObjectStoreParams) {
let mut options = store_params.storage_options().cloned().unwrap_or_default();
let mut removed = false;
for key in [
OPT_NEW_TABLE_STORAGE_VERSION,
OPT_NEW_TABLE_V2_MANIFEST_PATHS,
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS,
] {
removed |= options.remove(key).is_some();
}
if !removed {
return;
}
let provider = store_params
.storage_options_accessor
.as_ref()
.and_then(|accessor| accessor.provider().cloned());
store_params.storage_options_accessor = match (options.is_empty(), provider) {
(true, None) => None,
(true, Some(provider)) => Some(Arc::new(StorageOptionsAccessor::with_provider(provider))),
(false, Some(provider)) => Some(Arc::new(
StorageOptionsAccessor::with_initial_and_provider(options, provider),
)),
(false, None) => Some(Arc::new(StorageOptionsAccessor::with_static_options(
options,
))),
};
}
#[cfg(test)]
mod tests {
#[tokio::test]
async fn request_level_creation_keys_do_not_fork_the_store() {
use crate::query::ExecutableQuery;
use futures::TryStreamExt;
let db = crate::connect("memory://").execute().await.unwrap();
let batch = arrow_array::record_batch!(("x", Int32, [1, 2])).unwrap();
let store_params = ObjectStoreParams {
storage_options_accessor: Some(Arc::new(StorageOptionsAccessor::with_static_options(
HashMap::from([(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
)]),
))),
..Default::default()
};
db.create_table("t", batch)
.write_options(crate::table::WriteOptions {
lance_write_params: Some(lance::dataset::WriteParams {
store_params: Some(store_params),
..Default::default()
}),
})
.execute()
.await
.unwrap();
let table = db.open_table("t").execute().await.unwrap();
let rows: usize = table
.query()
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap()
.iter()
.map(|b| b.num_rows())
.sum();
assert_eq!(rows, 2, "the table must live in the session's store");
}
mod strip_new_table_creation_keys {
use super::super::*;
#[derive(Debug)]
struct EmptyProvider;
#[async_trait::async_trait]
impl StorageOptionsProvider for EmptyProvider {
async fn fetch_storage_options(
&self,
) -> lance_core::Result<Option<HashMap<String, String>>> {
Ok(Some(HashMap::new()))
}
fn provider_id(&self) -> String {
"empty-test-provider".into()
}
}
fn params_with_static(options: &[(&str, &str)]) -> ObjectStoreParams {
ObjectStoreParams {
storage_options_accessor: Some(Arc::new(
StorageOptionsAccessor::with_static_options(
options
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect(),
),
)),
..Default::default()
}
}
#[test]
fn creation_keys_are_removed_and_store_keys_kept() {
let mut params = params_with_static(&[
("region", "us-west-2"),
(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS, "true"),
]);
strip_new_table_creation_keys(&mut params);
let options = params.storage_options().cloned().unwrap();
assert_eq!(options.get("region").map(String::as_str), Some("us-west-2"));
assert!(!options.contains_key(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS));
let mut params = params_with_static(&[(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS, "true")]);
strip_new_table_creation_keys(&mut params);
assert!(params.storage_options_accessor.is_none());
}
#[test]
fn provider_accessors_survive_the_strip() {
let accessor = Arc::new(StorageOptionsAccessor::with_provider(Arc::new(
EmptyProvider,
)));
let mut params = ObjectStoreParams {
storage_options_accessor: Some(accessor.clone()),
..Default::default()
};
strip_new_table_creation_keys(&mut params);
assert!(Arc::ptr_eq(
params.storage_options_accessor.as_ref().unwrap(),
&accessor
));
let mut params = ObjectStoreParams {
storage_options_accessor: Some(Arc::new(
StorageOptionsAccessor::with_initial_and_provider(
HashMap::from([
("region".to_string(), "us-west-2".to_string()),
(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
),
]),
Arc::new(EmptyProvider),
),
)),
..Default::default()
};
strip_new_table_creation_keys(&mut params);
let accessor = params.storage_options_accessor.unwrap();
assert!(accessor.has_provider());
assert_eq!(
accessor
.initial_storage_options()
.and_then(|o| o.get("region").cloned())
.as_deref(),
Some("us-west-2")
);
let mut params = ObjectStoreParams {
storage_options_accessor: Some(Arc::new(
StorageOptionsAccessor::with_initial_and_provider(
HashMap::from([(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
)]),
Arc::new(EmptyProvider),
),
)),
..Default::default()
};
strip_new_table_creation_keys(&mut params);
let accessor = params.storage_options_accessor.unwrap();
assert!(accessor.has_provider());
assert!(accessor.initial_storage_options().is_none());
}
}
use super::*;
use crate::Table;
use crate::arrow::{SendableRecordBatchStream, SimpleRecordBatchStream};
use crate::connection::ConnectRequest;
use crate::data::scannable::Scannable;
use crate::database::{CreateTableMode, CreateTableRequest};
use crate::query::QueryRequest;
use crate::table::{AnyQuery, WriteOptions};
use arrow_array::{Int32Array, RecordBatch, StringArray};
use arrow_schema::{DataType, Field, Schema, SchemaRef};
use futures::{TryStreamExt, future::try_join_all, stream::once};
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;
use tempfile::tempdir;
use tokio::sync::Barrier;
use tokio::time::timeout;
async fn create_tables(db: &ListingDatabase, names: &[&str]) {
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
for name in names {
db.create_table(CreateTableRequest {
name: name.to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema.clone())) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
}
}
async fn walk(db: &ListingDatabase, limit: Option<i32>) -> Vec<String> {
let mut seen = Vec::new();
let mut page_token = None;
loop {
let page = db
.list_tables(ListTablesRequest {
limit,
page_token,
..Default::default()
})
.await
.unwrap();
seen.extend(page.tables);
page_token = page.page_token;
if page_token.is_none() {
return seen;
}
assert!(
seen.len() < 100,
"the walk is serving tables more than once"
);
}
}
#[rstest::rstest]
#[tokio::test]
async fn test_list_tables_pages_over_every_table_once(#[values(1, 2, 3, 5, 10)] limit: i32) {
let (_tempdir, db) = setup_database().await;
create_tables(&db, &["a", "b", "c", "d", "e"]).await;
assert_eq!(walk(&db, Some(limit)).await, vec!["a", "b", "c", "d", "e"]);
}
#[tokio::test]
async fn test_the_page_token_is_not_a_table_name() {
let (_tempdir, db) = setup_database().await;
create_tables(&db, &["a", "b", "c"]).await;
let page = db
.list_tables(ListTablesRequest {
limit: Some(1),
..Default::default()
})
.await
.unwrap();
assert_eq!(page.tables, vec!["a"]);
let token = page.page_token.expect("two tables are still to come");
assert_ne!(token, "a");
let rest = db
.list_tables(ListTablesRequest {
page_token: Some(token),
..Default::default()
})
.await
.unwrap();
assert_eq!(rest.tables, vec!["b", "c"]);
}
#[tokio::test]
async fn test_a_listing_that_runs_out_has_no_token() {
let (_tempdir, db) = setup_database().await;
create_tables(&db, &["a", "b"]).await;
let page = db
.list_tables(ListTablesRequest {
limit: Some(10),
..Default::default()
})
.await
.unwrap();
assert_eq!(page.tables, vec!["a", "b"]);
assert_eq!(page.page_token, None);
}
#[tokio::test]
async fn test_an_empty_page_token_lists_from_the_start() {
let (_tempdir, db) = setup_database().await;
create_tables(&db, &["a", "b"]).await;
let page = db
.list_tables(ListTablesRequest {
page_token: Some(String::new()),
..Default::default()
})
.await
.unwrap();
assert_eq!(page.tables, vec!["a", "b"]);
}
#[tokio::test]
async fn test_listing_order_follows_the_store_not_the_table_name() {
let (_tempdir, db) = setup_database().await;
create_tables(&db, &["users", "users-archive", "users.old"]).await;
assert_eq!(
walk(&db, None).await,
vec!["users-archive", "users", "users.old"]
);
assert_eq!(
walk(&db, Some(1)).await,
vec!["users-archive", "users", "users.old"]
);
}
#[tokio::test]
async fn test_listing_ignores_non_table_children() {
let (tempdir, db) = setup_database().await;
create_tables(&db, &["real"]).await;
std::fs::write(tempdir.path().join("aaa-loose.lance"), b"not a table").unwrap();
create_dir_all(tempdir.path().join("aaa-scratch")).unwrap();
let page = db
.list_tables(ListTablesRequest {
limit: Some(1),
..Default::default()
})
.await
.unwrap();
assert_eq!(page.tables, vec!["real"]);
}
#[tokio::test]
async fn test_one_full_listing_per_public_page() {
use crate::io::object_store::io_tracking::IoStatsHolder;
use lance_io::object_store::WrappingObjectStore;
let (tempdir, mut db) = setup_database().await;
create_tables(&db, &["real"]).await;
std::fs::write(tempdir.path().join("aaa-loose.lance"), b"not a table").unwrap();
create_dir_all(tempdir.path().join("aaa-scratch")).unwrap();
let io_stats = IoStatsHolder::default();
let mut tracked_store = (*db.object_store).clone();
tracked_store.inner =
io_stats.wrap(&tracked_store.store_prefix, tracked_store.inner.clone());
db.object_store = Arc::new(tracked_store);
let page = db
.list_tables(ListTablesRequest {
limit: Some(1),
..Default::default()
})
.await
.unwrap();
assert_eq!(page.tables, vec!["real"]);
assert_eq!(io_stats.incremental_stats().read_iops, 1);
}
#[tokio::test]
async fn listing_ignores_empty_table_name() {
let (tempdir, db) = setup_database().await;
create_dir_all(tempdir.path().join(".lance")).unwrap();
let page = db.list_tables(ListTablesRequest::default()).await.unwrap();
assert!(
page.tables.is_empty(),
"invalid empty table name was listed"
);
}
async fn setup_database() -> (tempfile::TempDir, ListingDatabase) {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
(tempdir, db)
}
struct BarrierScannable {
batch: RecordBatch,
barrier: Arc<Barrier>,
}
impl Scannable for BarrierScannable {
fn schema(&self) -> SchemaRef {
self.batch.schema()
}
fn scan_as_stream(&mut self) -> SendableRecordBatchStream {
let batch = self.batch.clone();
let schema = batch.schema();
let barrier = self.barrier.clone();
Box::pin(SimpleRecordBatchStream {
schema,
stream: once(async move {
barrier.wait().await;
Ok(batch)
}),
})
}
}
fn create_request(name: &str, data: Box<dyn Scannable>) -> CreateTableRequest {
CreateTableRequest {
name: name.to_string(),
namespace_path: vec![],
data,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
}
}
#[tokio::test]
async fn test_create_ignores_uncommitted_storage_without_manifest() {
let (tmp_dir, db) = setup_database().await;
let data_dir = tmp_dir.path().join("test.lance/data");
std::fs::create_dir_all(&data_dir).unwrap();
std::fs::write(data_dir.join("orphan.lance"), b"uncommitted").unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch =
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]).unwrap();
let table = db
.create_table(create_request("test", Box::new(batch)))
.await
.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 1);
}
#[tokio::test]
async fn test_concurrent_create_is_arbitrated_by_manifest_commit() {
let uri = format!("memory:///concurrent-create-{}", uuid::Uuid::new_v4());
let db = crate::connect(&uri).execute().await.unwrap();
let store: Arc<dyn object_store::ObjectStore> =
Arc::new(object_store::memory::InMemory::new());
let table_url = url::Url::parse("memory:///database/test.lance").unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch =
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]).unwrap();
let barrier = Arc::new(Barrier::new(2));
#[allow(deprecated)]
let request = |batch, barrier| {
let mut request = create_request("test", Box::new(BarrierScannable { batch, barrier }));
request.write_options = WriteOptions {
lance_write_params: Some(lance::dataset::WriteParams {
store_params: Some(ObjectStoreParams {
object_store: Some((store.clone(), table_url.clone())),
..Default::default()
}),
commit_handler: Some(Arc::new(
lance_table::io::commit::ConditionalPutCommitHandler,
)),
..Default::default()
}),
};
request
};
let left = db
.database()
.create_table(request(batch.clone(), barrier.clone()));
let right = db.database().create_table(request(batch, barrier));
let (left, right) = timeout(Duration::from_secs(30), async { tokio::join!(left, right) })
.await
.expect("concurrent creates deadlocked");
let results = [left, right];
assert_eq!(
results.iter().filter(|result| result.is_ok()).count(),
1,
"expected one successful create, got {results:?}"
);
assert_eq!(
results
.iter()
.filter(|result| matches!(result, Err(Error::TableAlreadyExists { .. })))
.count(),
1,
"expected one manifest conflict, got {results:?}"
);
}
#[tokio::test]
async fn test_concurrent_open_table_reuses_connection_object_store() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let session = Arc::new(lance::session::Session::default());
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: Some(session.clone()),
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "test".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let before = session.store_registry().stats();
let opened_tables = try_join_all((0..32).map(|_| {
db.open_table(OpenTableRequest {
name: "test".to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
})
}))
.await
.unwrap();
let after = session.store_registry().stats();
assert_eq!(opened_tables.len(), 32);
assert_eq!(after.misses, before.misses);
assert_eq!(after.active_stores, before.active_stores);
assert!(after.hits >= before.hits + 32);
}
#[tokio::test]
async fn test_listing_database_root_ops_do_not_create_manifest() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
assert!(!tempdir.path().join("__manifest").exists());
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "root_table".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
db.open_table(OpenTableRequest {
name: "root_table".to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
})
.await
.unwrap();
#[allow(deprecated)]
let table_names = db.table_names(TableNamesRequest::default()).await.unwrap();
assert_eq!(table_names, vec!["root_table".to_string()]);
assert!(!tempdir.path().join("__manifest").exists());
}
#[tokio::test]
async fn test_open_table_reuses_connection_object_store() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let registry = Arc::new(lance_io::object_store::ObjectStoreRegistry::default());
let session = Arc::new(lance::session::Session::new(16, 16, registry.clone()));
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: Some(session),
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "test".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let before_open = registry.stats();
for _ in 0..3 {
let table = db
.open_table(OpenTableRequest {
name: "test".to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
})
.await
.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 0);
}
let after_open = registry.stats();
assert_eq!(after_open.misses, before_open.misses);
assert!(after_open.hits >= before_open.hits + 3);
}
#[cfg(unix)]
#[tokio::test]
async fn test_open_table_follows_hugging_face_symlinks() {
let (tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "test".to_string(),
namespace_path: vec![],
data: Box::new(
RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1, 2, 3]))])
.unwrap(),
) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let table_dir = tempdir.path().join("test.lance");
let versions_dir = table_dir.join("_versions");
let manifest_path = std::fs::read_dir(&versions_dir)
.unwrap()
.map(|entry| entry.unwrap().path())
.find(|path| path.extension().is_some_and(|ext| ext == "manifest"))
.unwrap();
let data_path = std::fs::read_dir(table_dir.join("data"))
.unwrap()
.map(|entry| entry.unwrap().path())
.find(|path| path.extension().is_some_and(|ext| ext == "lance"))
.unwrap();
let blobs_dir = tempdir.path().join("blobs");
std::fs::create_dir(&blobs_dir).unwrap();
let manifest_blob = "9b603c63d0e692e05d58be25605f2f2064cc781e5ff94fe983a405059547b816";
let data_blob = "be64f20e5723bd0a27cfdbdb41cf7d6fad94cd572a71973b717fb8340f4310c5";
std::fs::rename(&manifest_path, blobs_dir.join(manifest_blob)).unwrap();
std::fs::rename(&data_path, blobs_dir.join(data_blob)).unwrap();
std::os::unix::fs::symlink(Path::new("../../blobs").join(manifest_blob), &manifest_path)
.unwrap();
std::os::unix::fs::symlink(Path::new("../../blobs").join(data_blob), &data_path).unwrap();
let symlink_len = std::fs::symlink_metadata(&manifest_path).unwrap().len();
let target_len = std::fs::metadata(&manifest_path).unwrap().len();
assert_ne!(symlink_len, target_len);
drop(db);
let db = ListingDatabase::connect_with_options(&ConnectRequest {
uri: tempdir.path().to_str().unwrap().to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
})
.await
.unwrap();
let table = db
.open_table(OpenTableRequest {
name: "test".to_string(),
namespace_path: vec![],
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
})
.await
.unwrap();
let batches = table
.query(
&AnyQuery::Query(QueryRequest::default()),
Default::default(),
)
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.unwrap();
assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::<usize>(), 3);
}
#[tokio::test]
async fn test_clone_table_basic() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
]));
let source_table = db
.create_table(CreateTableRequest {
name: "source_table".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema.clone())) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source_table").unwrap();
let cloned_table = db
.clone_table(CloneTableRequest {
target_table_name: "cloned_table".to_string(),
target_namespace_path: vec![],
source_uri: source_uri.clone(),
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await
.unwrap();
#[allow(deprecated)]
let table_names = db.table_names(TableNamesRequest::default()).await.unwrap();
assert!(table_names.contains(&"source_table".to_string()));
assert!(table_names.contains(&"cloned_table".to_string()));
assert_eq!(
source_table.schema().await.unwrap(),
cloned_table.schema().await.unwrap()
);
}
#[tokio::test]
async fn test_clone_table_with_data() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2, 3])),
Arc::new(StringArray::from(vec!["a", "b", "c"])),
],
)
.unwrap();
let source_table = db
.create_table(CreateTableRequest {
name: "source_with_data".to_string(),
namespace_path: vec![],
data: Box::new(batch) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source_with_data").unwrap();
let cloned_table = db
.clone_table(CloneTableRequest {
target_table_name: "cloned_with_data".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await
.unwrap();
let source_count = source_table.count_rows(None).await.unwrap();
let cloned_count = cloned_table.count_rows(None).await.unwrap();
assert_eq!(source_count, cloned_count);
assert_eq!(source_count, 3);
}
#[tokio::test]
async fn test_clone_table_with_storage_options() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let mut options = HashMap::new();
options.insert("test_option".to_string(), "test_value".to_string());
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: options.clone(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "source".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source").unwrap();
let cloned = db
.clone_table(CloneTableRequest {
target_table_name: "cloned".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await;
assert!(cloned.is_ok());
}
#[tokio::test]
async fn test_clone_table_deep_not_supported() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "source".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source").unwrap();
let result = db
.clone_table(CloneTableRequest {
target_table_name: "cloned".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: None,
is_shallow: false, namespace_client: None,
})
.await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
Error::NotSupported { message } if message.contains("Deep clone")
));
}
#[tokio::test]
async fn test_clone_table_with_namespace_not_supported() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "source".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source").unwrap();
let result = db
.clone_table(CloneTableRequest {
target_table_name: "cloned".to_string(),
target_namespace_path: vec!["namespace".to_string()], source_uri,
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
Error::NotSupported { message } if message.contains("Namespace parameter is not supported")
));
}
#[tokio::test]
async fn test_clone_table_invalid_target_name() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "source".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source").unwrap();
let result = db
.clone_table(CloneTableRequest {
target_table_name: "invalid/name".to_string(), target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_clone_table_source_not_found() {
let (_tempdir, db) = setup_database().await;
let result = db
.clone_table(CloneTableRequest {
target_table_name: "cloned".to_string(),
target_namespace_path: vec![],
source_uri: "/nonexistent/table.lance".to_string(),
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await;
assert!(result.is_err());
}
#[tokio::test]
async fn test_clone_table_with_version_and_tag_error() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
db.create_table(CreateTableRequest {
name: "source".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("source").unwrap();
let result = db
.clone_table(CloneTableRequest {
target_table_name: "cloned".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: Some(1),
source_tag: Some("v1.0".to_string()),
is_shallow: true,
namespace_client: None,
})
.await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
Error::InvalidInput { message } if message.contains("Cannot specify both source_version and source_tag")
));
}
#[tokio::test]
async fn test_clone_table_with_specific_version() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("value", DataType::Utf8, false),
]));
let batch1 = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec!["a", "b"])),
],
)
.unwrap();
let source_table = db
.create_table(CreateTableRequest {
name: "versioned_source".to_string(),
namespace_path: vec![],
data: Box::new(batch1) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let initial_version = source_table.version().await.unwrap();
let batch2 = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![3, 4])),
Arc::new(StringArray::from(vec!["c", "d"])),
],
)
.unwrap();
let db = Arc::new(db);
let source_table_obj = Table::new(source_table.clone(), db.clone());
source_table_obj.add(batch2).execute().await.unwrap();
assert_eq!(source_table.count_rows(None).await.unwrap(), 4);
let source_uri = db.table_uri("versioned_source").unwrap();
let cloned_table = db
.clone_table(CloneTableRequest {
target_table_name: "cloned_from_version".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: Some(initial_version),
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(cloned_table.count_rows(None).await.unwrap(), 2);
assert_eq!(source_table.count_rows(None).await.unwrap(), 4);
}
#[tokio::test]
async fn test_clone_table_with_tag() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("value", DataType::Utf8, false),
]));
let batch1 = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec!["a", "b"])),
],
)
.unwrap();
let source_table = db
.create_table(CreateTableRequest {
name: "tagged_source".to_string(),
namespace_path: vec![],
data: Box::new(batch1),
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let db = Arc::new(db);
let source_table_obj = Table::new(source_table.clone(), db.clone());
let mut tags = source_table_obj.tags().await.unwrap();
tags.create("v1.0", source_table.version().await.unwrap())
.await
.unwrap();
let batch2 = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![3, 4])),
Arc::new(StringArray::from(vec!["c", "d"])),
],
)
.unwrap();
let source_table_obj = Table::new(source_table.clone(), db.clone());
source_table_obj.add(batch2).execute().await.unwrap();
assert_eq!(source_table.count_rows(None).await.unwrap(), 4);
let source_uri = db.table_uri("tagged_source").unwrap();
let cloned_table = db
.clone_table(CloneTableRequest {
target_table_name: "cloned_from_tag".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: Some("v1.0".to_string()),
is_shallow: true,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(cloned_table.count_rows(None).await.unwrap(), 2);
}
#[tokio::test]
async fn test_cloned_tables_evolve_independently() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("value", DataType::Utf8, false),
]));
let batch1 = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![1, 2])),
Arc::new(StringArray::from(vec!["a", "b"])),
],
)
.unwrap();
let source_table = db
.create_table(CreateTableRequest {
name: "independent_source".to_string(),
namespace_path: vec![],
data: Box::new(batch1),
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let source_uri = db.table_uri("independent_source").unwrap();
let cloned_table = db
.clone_table(CloneTableRequest {
target_table_name: "independent_clone".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(source_table.count_rows(None).await.unwrap(), 2);
assert_eq!(cloned_table.count_rows(None).await.unwrap(), 2);
let batch_clone = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![3, 4, 5])),
Arc::new(StringArray::from(vec!["c", "d", "e"])),
],
)
.unwrap();
let db = Arc::new(db);
let cloned_table_obj = Table::new(cloned_table.clone(), db.clone());
cloned_table_obj.add(batch_clone).execute().await.unwrap();
let batch_source = RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(Int32Array::from(vec![10, 11])),
Arc::new(StringArray::from(vec!["x", "y"])),
],
)
.unwrap();
let source_table_obj = Table::new(source_table.clone(), db);
source_table_obj.add(batch_source).execute().await.unwrap();
assert_eq!(source_table.count_rows(None).await.unwrap(), 4); assert_eq!(cloned_table.count_rows(None).await.unwrap(), 5); }
#[tokio::test]
async fn test_clone_latest_version() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch1 =
RecordBatch::try_new(schema.clone(), vec![Arc::new(Int32Array::from(vec![1, 2]))])
.unwrap();
let source_table = db
.create_table(CreateTableRequest {
name: "latest_version_source".to_string(),
namespace_path: vec![],
data: Box::new(batch1),
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let db = Arc::new(db);
for i in 0..3 {
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![i * 10, i * 10 + 1]))],
)
.unwrap();
let source_table_obj = Table::new(source_table.clone(), db.clone());
source_table_obj.add(batch).execute().await.unwrap();
}
let source_count = source_table.count_rows(None).await.unwrap();
assert_eq!(source_count, 8);
let source_uri = db.table_uri("latest_version_source").unwrap();
let cloned_table = db
.clone_table(CloneTableRequest {
target_table_name: "cloned_latest".to_string(),
target_namespace_path: vec![],
source_uri,
source_version: None,
source_tag: None,
is_shallow: true,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(cloned_table.count_rows(None).await.unwrap(), 8);
}
#[tokio::test]
async fn test_create_table_with_stable_row_ids_connection_level() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let mut options = HashMap::new();
options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
);
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options,
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
assert_eq!(db.new_table_config.enable_stable_row_ids, Some(true));
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let table = db
.create_table(CreateTableRequest {
name: "test_stable".to_string(),
namespace_path: vec![],
data: Box::new(batch),
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 3);
}
#[tokio::test]
async fn test_create_table_with_stable_row_ids_table_level() {
let (_tempdir, db) = setup_database().await;
assert_eq!(db.new_table_config.enable_stable_row_ids, None);
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let mut storage_options = HashMap::new();
storage_options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
);
let write_options = WriteOptions {
lance_write_params: Some(lance::dataset::WriteParams {
store_params: Some(lance::io::ObjectStoreParams {
storage_options_accessor: Some(Arc::new(
StorageOptionsAccessor::with_static_options(storage_options),
)),
..Default::default()
}),
..Default::default()
}),
};
let table = db
.create_table(CreateTableRequest {
name: "test_stable_table_level".to_string(),
namespace_path: vec![],
data: Box::new(batch),
mode: CreateTableMode::Create,
write_options,
location: None,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 3);
}
#[tokio::test]
async fn test_create_table_stable_row_ids_table_overrides_connection() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let mut options = HashMap::new();
options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
);
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options,
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
assert_eq!(db.new_table_config.enable_stable_row_ids, Some(true));
let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));
let batch = RecordBatch::try_new(
schema.clone(),
vec![Arc::new(Int32Array::from(vec![1, 2, 3]))],
)
.unwrap();
let mut storage_options = HashMap::new();
storage_options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"false".to_string(),
);
let write_options = WriteOptions {
lance_write_params: Some(lance::dataset::WriteParams {
store_params: Some(lance::io::ObjectStoreParams {
storage_options_accessor: Some(Arc::new(
StorageOptionsAccessor::with_static_options(storage_options),
)),
..Default::default()
}),
..Default::default()
}),
};
let table = db
.create_table(CreateTableRequest {
name: "test_override".to_string(),
namespace_path: vec![],
data: Box::new(batch),
mode: CreateTableMode::Create,
write_options,
location: None,
namespace_client: None,
})
.await
.unwrap();
assert_eq!(table.count_rows(None).await.unwrap(), 3);
}
#[tokio::test]
async fn test_stable_row_ids_invalid_value() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let mut options = HashMap::new();
options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"not_a_boolean".to_string(),
);
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options,
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let result = ListingDatabase::connect_with_options(&request).await;
assert!(result.is_err());
assert!(matches!(
result.unwrap_err(),
Error::InvalidInput { message } if message.contains("enable_stable_row_ids must be a boolean")
));
}
#[test]
fn test_stable_row_ids_config_serialization() {
let mut options = HashMap::new();
options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"true".to_string(),
);
let db_options = ListingDatabaseOptions::parse_from_map(&options).unwrap();
assert_eq!(
db_options.new_table_config.enable_stable_row_ids,
Some(true)
);
let mut serialized = HashMap::new();
db_options.serialize_into_map(&mut serialized);
assert_eq!(
serialized.get(OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS),
Some(&"true".to_string())
);
}
#[test]
fn test_stable_row_ids_config_parse_false() {
let mut options = HashMap::new();
options.insert(
OPT_NEW_TABLE_ENABLE_STABLE_ROW_IDS.to_string(),
"false".to_string(),
);
let db_options = ListingDatabaseOptions::parse_from_map(&options).unwrap();
assert_eq!(
db_options.new_table_config.enable_stable_row_ids,
Some(false)
);
}
#[test]
fn test_stable_row_ids_config_not_set() {
let options = HashMap::new();
let db_options = ListingDatabaseOptions::parse_from_map(&options).unwrap();
assert_eq!(db_options.new_table_config.enable_stable_row_ids, None);
}
#[tokio::test]
async fn test_table_uri() {
let (_tempdir, mut db) = setup_database().await;
let mut pb = PathBuf::new();
pb.push(db.uri.clone());
pb.push("test.lance");
let expected = pb.to_str().unwrap();
let uri = db.table_uri("test").ok().unwrap();
assert_eq!(uri, expected);
for base_uri in ["az://container/prefix", "az://container/prefix/"] {
db.uri = base_uri.to_string();
assert_eq!(
db.table_uri("test").unwrap(),
"az://container/prefix/test.lance"
);
}
}
#[tokio::test]
async fn test_table_uri_uses_forward_slashes_for_azure() {
let (_tempdir, mut db) = setup_database().await;
db.uri = "az://test/db/test".to_string();
let uri = db.table_uri("test").unwrap();
assert_eq!(uri, "az://test/db/test/test.lance");
}
fn capture_query_like_connect(input_uri: &str) -> Option<String> {
let mut url = url::Url::parse(input_uri).unwrap();
let mut filtered_querys = Vec::new();
for (key, value) in url.query_pairs() {
if key == ENGINE || key == MIRRORED_STORE {
continue;
}
filtered_querys.push((key.to_string(), value.to_string()));
}
url.query_pairs_mut().clear();
url.query_pairs_mut().extend_pairs(filtered_querys);
url.query().filter(|q| !q.is_empty()).map(|s| s.to_string())
}
#[test]
fn test_capture_query_treats_empty_as_none() {
assert_eq!(
capture_query_like_connect("s3://bucket/prefix/"),
None,
"empty query after mutation must be treated as no query"
);
assert_eq!(
capture_query_like_connect("s3://bucket/prefix/?foo=bar"),
Some("foo=bar".to_string())
);
assert_eq!(
capture_query_like_connect(&format!("s3://bucket/prefix/?{}=mem", ENGINE)),
None
);
let captured =
capture_query_like_connect(&format!("s3://bucket/prefix/?{}=mem&foo=bar", ENGINE));
assert_eq!(captured.as_deref(), Some("foo=bar"));
}
#[test]
fn test_storage_base_uri_strips_commit_engine_scheme() {
assert_eq!(
ListingDatabase::storage_base_uri("s3+ddb://bucket/prefix?ddbTableName=commit_table"),
"s3://bucket/prefix"
);
assert_eq!(
ListingDatabase::storage_base_uri("s3://bucket/prefix?foo=bar"),
"s3://bucket/prefix"
);
assert_eq!(
ListingDatabase::storage_base_uri("/tmp/lancedb"),
"/tmp/lancedb"
);
}
#[cfg(not(windows))]
#[tokio::test]
async fn test_table_uri_url_path_has_no_trailing_question_mark() {
let tempdir = tempdir().unwrap();
let uri = format!("file://{}", tempdir.path().to_str().unwrap());
let request = ConnectRequest {
uri: uri.clone(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
assert_eq!(
db.query_string, None,
"no input query → no captured query_string"
);
let table_uri = db.table_uri("test").unwrap();
assert!(
!table_uri.ends_with('?'),
"table_uri must not have a trailing `?`: {}",
table_uri
);
assert_eq!(table_uri, format!("{}/test.lance", uri));
let with_query = format!("{}?foo=bar", uri);
let request_with_query = ConnectRequest {
uri: with_query,
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties: Default::default(),
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db_with_query = ListingDatabase::connect_with_options(&request_with_query)
.await
.unwrap();
assert_eq!(db_with_query.query_string.as_deref(), Some("foo=bar"));
let table_uri = db_with_query.table_uri("test").unwrap();
assert_eq!(table_uri, format!("{}/test.lance?foo=bar", uri));
}
#[tokio::test]
async fn test_namespace_client() {
let (_tempdir, db) = setup_database().await;
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
]));
db.create_table(CreateTableRequest {
name: "table1".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema.clone())) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
db.create_table(CreateTableRequest {
name: "table2".to_string(),
namespace_path: vec![],
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let namespace_client = db.namespace_client().await;
assert!(namespace_client.is_ok());
let namespace_client = namespace_client.unwrap();
let list_result = namespace_client
.list_tables(lance_namespace::models::ListTablesRequest {
id: Some(vec![]),
..Default::default()
})
.await;
assert!(
list_result.is_ok(),
"list_tables failed: {:?}",
list_result.err()
);
let tables = list_result.unwrap().tables;
assert_eq!(tables.len(), 2);
assert!(tables.contains(&"table1".to_string()));
assert!(tables.contains(&"table2".to_string()));
}
#[tokio::test]
async fn test_listing_database_namespace_operations() {
let (_tempdir, db) = setup_database().await;
db.create_namespace(CreateNamespaceRequest {
id: Some(vec!["parent".to_string()]),
..Default::default()
})
.await
.unwrap();
db.create_namespace(CreateNamespaceRequest {
id: Some(vec!["parent".to_string(), "child".to_string()]),
..Default::default()
})
.await
.unwrap();
let root_namespaces = db
.list_namespaces(ListNamespacesRequest {
id: Some(vec![]),
..Default::default()
})
.await
.unwrap();
assert!(root_namespaces.namespaces.contains(&"parent".to_string()));
let child_namespaces = db
.list_namespaces(ListNamespacesRequest {
id: Some(vec!["parent".to_string()]),
..Default::default()
})
.await
.unwrap();
assert!(child_namespaces.namespaces.contains(&"child".to_string()));
db.describe_namespace(DescribeNamespaceRequest {
id: Some(vec!["parent".to_string(), "child".to_string()]),
..Default::default()
})
.await
.unwrap();
}
#[tokio::test]
#[cfg(not(windows))] async fn test_listing_database_with_namespace_client_properties() {
let tempdir = tempdir().unwrap();
let uri = tempdir.path().to_str().unwrap();
let mut namespace_client_properties = HashMap::new();
namespace_client_properties.insert(
"table_version_tracking_enabled".to_string(),
"true".to_string(),
);
namespace_client_properties.insert("manifest_enabled".to_string(), "true".to_string());
let request = ConnectRequest {
uri: uri.to_string(),
#[cfg(feature = "remote")]
client_config: Default::default(),
options: Default::default(),
namespace_client_properties,
manifest_enabled: false,
read_consistency_interval: None,
session: None,
};
let db = ListingDatabase::connect_with_options(&request)
.await
.unwrap();
let namespace_path = vec!["test_ns".to_string()];
db.create_namespace(CreateNamespaceRequest {
id: Some(namespace_path.clone()),
..Default::default()
})
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
]));
db.create_table(CreateTableRequest {
name: "managed_table".to_string(),
namespace_path: namespace_path.clone(),
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let namespace_client = db.namespace_client().await.unwrap();
let describe = namespace_client
.describe_table(lance_namespace::models::DescribeTableRequest {
id: Some(vec!["test_ns".to_string(), "managed_table".to_string()]),
..Default::default()
})
.await
.unwrap();
assert_eq!(describe.managed_versioning, Some(true));
}
#[tokio::test]
async fn test_listing_database_nested_namespace_table_ops() {
let (_tempdir, db) = setup_database().await;
let namespace_path = vec!["parent".to_string(), "child".to_string()];
db.create_namespace(CreateNamespaceRequest {
id: Some(vec!["parent".to_string()]),
..Default::default()
})
.await
.unwrap();
db.create_namespace(CreateNamespaceRequest {
id: Some(namespace_path.clone()),
..Default::default()
})
.await
.unwrap();
let schema = Arc::new(Schema::new(vec![
Field::new("id", DataType::Int32, false),
Field::new("name", DataType::Utf8, false),
]));
db.create_table(CreateTableRequest {
name: "nested_table".to_string(),
namespace_path: namespace_path.clone(),
data: Box::new(RecordBatch::new_empty(schema)) as Box<dyn Scannable>,
mode: CreateTableMode::Create,
write_options: Default::default(),
location: None,
namespace_client: None,
})
.await
.unwrap();
let namespace_client = db.namespace_client().await.unwrap();
let describe = namespace_client
.describe_table(lance_namespace::models::DescribeTableRequest {
id: Some(vec![
"parent".to_string(),
"child".to_string(),
"nested_table".to_string(),
]),
..Default::default()
})
.await
.unwrap();
assert!(describe.location.is_some());
let table = db
.open_table(OpenTableRequest {
name: "nested_table".to_string(),
namespace_path: namespace_path.clone(),
index_cache_size: None,
lance_read_params: None,
location: None,
namespace_client: None,
managed_versioning: None,
})
.await
.unwrap();
assert_eq!(table.name(), "nested_table");
#[allow(deprecated)]
let table_names = db
.table_names(TableNamesRequest {
namespace_path: namespace_path.clone(),
start_after: None,
limit: None,
})
.await
.unwrap();
assert_eq!(table_names, vec!["nested_table".to_string()]);
let list_tables = db
.list_tables(ListTablesRequest {
id: Some(namespace_path.clone()),
..Default::default()
})
.await
.unwrap();
assert_eq!(list_tables.tables, vec!["nested_table".to_string()]);
db.drop_table("nested_table", &namespace_path)
.await
.unwrap();
let post_drop = db
.list_tables(ListTablesRequest {
id: Some(namespace_path),
..Default::default()
})
.await
.unwrap();
assert!(post_drop.tables.is_empty());
}
}