diff --git a/rust/lance/src/dataset/refs.rs b/rust/lance/src/dataset/refs.rs index 0d3f65f7959..4a4845555db 100644 --- a/rust/lance/src/dataset/refs.rs +++ b/rust/lance/src/dataset/refs.rs @@ -573,6 +573,28 @@ impl Branches<'_> { log::warn!("BranchContents of {} does not exist", branch); } + // Tags identify snapshots by (branch, version). Deleting a branch removes its entire version chain, + // so any tag whose branch matches the deletion target blocks the operation, regardless of the tagged + // version. + let referenced_tags = self + .refs + .tags() + .list() + .await? + .into_iter() + .filter_map(|(name, contents)| { + (contents.branch.as_deref() == Some(branch)).then_some((name, contents.version)) + }) + .collect_vec(); + if !referenced_tags.is_empty() { + return Err(Error::RefConflict { + message: format!( + "Branch {} is referenced by tags {:?}, can not delete", + branch, referenced_tags + ), + }); + } + let root_location = self.refs.root()?; let branch_file = branch_contents_path(&root_location.path, branch); if self.object_store().exists(&branch_file).await? { diff --git a/rust/lance/src/dataset/tests/dataset_versioning.rs b/rust/lance/src/dataset/tests/dataset_versioning.rs index 0a65fa01a51..c95e3a6d250 100644 --- a/rust/lance/src/dataset/tests/dataset_versioning.rs +++ b/rust/lance/src/dataset/tests/dataset_versioning.rs @@ -843,6 +843,30 @@ async fn test_create_branch_and_shallow_clone_from_other_branch() { ); } +#[tokio::test] +async fn test_cannot_delete_branch_referenced_by_tag() { + let tempdir = TempDir::default(); + let test_uri = tempdir.path_str(); + let data = gen_batch() + .col("id", array::step::()) + .into_reader_rows(RowCount::from(1), BatchCount::from(1)); + let mut dataset = Dataset::write(data, &test_uri, None).await.unwrap(); + let branch = dataset + .create_branch("dev", /*version=*/ 1, /*store_params=*/ None) + .await + .unwrap(); + dataset + .tags() + .create("keep-dev", ("dev", branch.version().version)) + .await + .unwrap(); + + let error = dataset.delete_branch("dev").await.unwrap_err(); + assert!(matches!(&error, Error::RefConflict { .. })); + assert!(error.to_string().contains("keep-dev")); + dataset.checkout_version("keep-dev").await.unwrap(); +} + #[tokio::test] async fn test_branch() { let tempdir = TempDir::default(); @@ -1174,6 +1198,7 @@ async fn test_branch() { let cleaned_path = Path::parse(format!("{}/tree/feature", test_uri)).unwrap(); assert!(!dataset.object_store.exists(&cleaned_path).await.unwrap()); + dataset.tags().delete("tag1").await.unwrap(); dataset.delete_branch("dev/branch2").await.unwrap(); dataset.delete_branch("branch1").await.unwrap();