fix: protect active dataset versions from cleanup

java/src/test/java/org/lance/CleanupTest.java

@@ -114,7 +114,7 @@ public void testCleanupBeforeTimestamp(@TempDir Path tempDir) throws Exception {
 
   @Test
   public void testCleanupTaggedVersion(@TempDir Path tempDir) throws Exception {
-    String datasetPath = tempDir.resolve("test_dataset_for_cleanup").toString();
+    String datasetPath = TestUtils.sharedMemoryUri(tempDir, "test_dataset_for_cleanup");
     try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
       TestUtils.SimpleTestDataset testDataset =
           new TestUtils.SimpleTestDataset(allocator, datasetPath);

java/src/test/java/org/lance/DatasetTest.java

@@ -308,7 +308,7 @@ void testDatasetCheckoutVersion(@TempDir Path tempDir) {
 
   @Test
   void testTags(@TempDir Path tempDir) {
-    String datasetPath = tempDir.resolve("dataset_tags").toString();
+    String datasetPath = TestUtils.sharedMemoryUri(tempDir, "dataset_tags");
     try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
       TestUtils.SimpleTestDataset testDataset =
           new TestUtils.SimpleTestDataset(allocator, datasetPath);
@@ -453,7 +453,7 @@ void testTags(@TempDir Path tempDir) {
 
   @Test
   void testDatasetRestore(@TempDir Path tempDir) {
-    String datasetPath = tempDir.resolve("dataset_restore").toString();
+    String datasetPath = TestUtils.sharedMemoryUri(tempDir, "dataset_restore");
     try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
       TestUtils.SimpleTestDataset testDataset =
           new TestUtils.SimpleTestDataset(allocator, datasetPath);
@@ -1785,9 +1785,9 @@ public ByteBuffer readManifest(Path filePath) throws IOException {
 
   @Test
   void testShallowClone(@TempDir Path tempDir) {
-    String srcPath = tempDir.resolve("shallow_clone_version_src").toString();
-    String dstPathByVersion = tempDir.resolve("shallow_clone_version_dst").toString();
-    String dstPathByTag = tempDir.resolve("shallow_clone_tag_dst").toString();
+    String srcPath = TestUtils.sharedMemoryUri(tempDir, "shallow_clone_version_src");
+    String dstPathByVersion = TestUtils.sharedMemoryUri(tempDir, "shallow_clone_version_dst");
+    String dstPathByTag = TestUtils.sharedMemoryUri(tempDir, "shallow_clone_tag_dst");
 
     try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
       // Prepare a simple source dataset with some rows
@@ -1841,7 +1841,7 @@ void testShallowClone(@TempDir Path tempDir) {
 
   @Test
   void testBranches(@TempDir Path tempDir) {
-    String datasetPath = tempDir.resolve("testBranches").toString();
+    String datasetPath = TestUtils.sharedMemoryUri(tempDir, "testBranches");
     try (RootAllocator allocator = new RootAllocator(Long.MAX_VALUE)) {
       TestUtils.SimpleTestDataset suite = new TestUtils.SimpleTestDataset(allocator, datasetPath);
 

java/src/test/java/org/lance/TestUtils.java

@@ -67,6 +67,10 @@
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 public class TestUtils {
+  public static String sharedMemoryUri(Path tempDir, String datasetName) {
+    return String.format("shared-memory://%s/%s", tempDir.getFileName(), datasetName);
+  }
+
   private abstract static class TestDataset {
     protected final BufferAllocator allocator;
     protected final String datasetPath;

python/python/lance/__init__.py

@@ -32,6 +32,7 @@
     MergeInsertBuilder,
     Session,
     Transaction,
+    VersionLease,
     __version__,
     batch_udf,
     write_dataset,
@@ -107,6 +108,7 @@
     "MergeInsertBuilder",
     "ScanStatistics",
     "Transaction",
+    "VersionLease",
     "__version__",
     "batch_udf",
     "bytes_read_counter",

python/python/lance/dataset.py

@@ -64,6 +64,7 @@
     LanceSchema,
     PySearchFilter,
     ScanStatistics,
+    VersionLease,
     _Dataset,
     _format_field_path,
     _MergeInsertBuilder,
@@ -3089,6 +3090,38 @@ def checkout_version(
             ds._ds = self._ds.checkout_version(version)
         return ds
 
+    def acquire_version_lease(self, ttl: timedelta) -> VersionLease:
+        """Protect this dataset version from cleanup with a renewable lease.
+
+        Acquire the lease before starting a long-running read and renew it before
+        :attr:`VersionLease.expires_at`. Cleanup may remove the version after the
+        lease expires. A context manager releases the lease on exit; a process
+        that exits without releasing it stops protecting the version at expiry.
+
+        Parameters
+        ----------
+        ttl : timedelta
+            How long the lease protects this version. Must be positive.
+
+        Returns
+        -------
+        VersionLease
+            A renewable lease for this dataset version.
+
+        Examples
+        --------
+        >>> import lance
+        >>> import pyarrow as pa
+        >>> from datetime import timedelta
+        >>> from tempfile import TemporaryDirectory
+        >>> with TemporaryDirectory() as tmp:
+        ...     historical = lance.write_dataset(pa.table({"id": [1]}), tmp)
+        ...     with historical.acquire_version_lease(timedelta(minutes=5)) as lease:
+        ...         table = historical.to_table()
+        ...         lease.renew(timedelta(minutes=5))
+        """
+        return self._ds.acquire_version_lease(td_to_micros(ttl))
+
     def restore(self):
         """
         Restore the currently checked out version as the latest version of the dataset.

python/python/lance/lance/__init__.pyi

@@ -12,7 +12,7 @@
 #  See the License for the specific language governing permissions and
 #  limitations under the License.
 
-from datetime import timedelta
+from datetime import datetime, timedelta
 from pathlib import Path
 from typing import (
     Any,
@@ -137,6 +137,16 @@ class CleanupExplanation:
     referenced_branches: List[CleanupReferencedBranch]
     warnings: List[str]
 
+class VersionLease:
+    @property
+    def version(self) -> int: ...
+    @property
+    def expires_at(self) -> datetime: ...
+    def renew(self, ttl: timedelta) -> None: ...
+    def release(self) -> None: ...
+    def __enter__(self) -> VersionLease: ...
+    def __exit__(self, exc_type, exc_value, traceback) -> bool: ...
+
 class LanceFileWriteSummary:
     num_rows: int
     size_bytes: int
@@ -460,6 +470,7 @@ class _Dataset:
     def checkout_version(
         self, version: int | str | Tuple[Optional[str], Optional[int]]
     ) -> _Dataset: ...
+    def acquire_version_lease(self, ttl_micros: int) -> VersionLease: ...
     def checkout_latest(self) -> _Dataset: ...
     def shallow_clone(
         self,

python/python/tests/test_dataset.py

@@ -650,10 +650,10 @@ def test_v2_manifest_paths_migration(tmp_path: Path):
 
 def test_tag(tmp_path: Path):
     table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]})
-    base_dir = tmp_path / "test"
+    base_dir = "shared-memory://test-tag/dataset"
 
-    lance.write_dataset(table, base_dir)
-    ds = lance.write_dataset(table, base_dir, mode="append")
+    ds = lance.write_dataset(table, base_dir)
+    ds = lance.write_dataset(table, ds, mode="append")
 
     assert len(ds.tags.list()) == 0
 
@@ -689,11 +689,11 @@ def test_tag(tmp_path: Path):
 
     assert ds.checkout_version("tag1").version == 1
 
-    ds = lance.dataset(base_dir, "tag1")
+    ds = ds.checkout_version("tag1")
     assert ds.version == 1
 
-    with pytest.raises(ValueError):
-        lance.dataset(base_dir, "missing-tag")
+    with pytest.raises(OSError):
+        ds.checkout_version("missing-tag")
 
     # test tag update
     with pytest.raises(
@@ -712,7 +712,7 @@ def test_tag(tmp_path: Path):
     assert tag1_created_at is not None
     assert tag1_updated_at is not None
     ds.tags.replace_metadata("tag1", {"description": "updated tag"})
-    ds = lance.dataset(base_dir, "tag1")
+    ds = ds.checkout_version("tag1")
     assert ds.version == 1
     replaced_tag1_meta = ds.tags.list()["tag1"]
     assert replaced_tag1_meta["metadata"] == {"description": "updated tag"}
@@ -728,17 +728,17 @@ def test_tag(tmp_path: Path):
     assert updated_tag1_meta["created_at"] == tag1_created_at
     assert updated_tag1_meta["updated_at"] is not None
     assert updated_tag1_meta["updated_at"] >= tag1_updated_at
-    ds = lance.dataset(base_dir, "tag1")
+    ds = ds.checkout_version("tag1")
     assert ds.version == 2
     assert ds.tags.list()["tag1"]["metadata"] == {"owner": "ml-team"}
 
     ds.tags.replace_metadata("tag1", {})
-    ds = lance.dataset(base_dir, "tag1")
+    ds = ds.checkout_version("tag1")
     assert ds.version == 2
     assert ds.tags.list()["tag1"]["metadata"] == {}
 
     ds.tags.update("tag1", 1)
-    ds = lance.dataset(base_dir, "tag1")
+    ds = ds.checkout_version("tag1")
     assert ds.version == 1
     assert ds.tags.list()["tag1"]["metadata"] == {}
 
@@ -775,11 +775,11 @@ def test_tag(tmp_path: Path):
 
 def test_tag_order(tmp_path: Path):
     table = pa.Table.from_pydict({"colA": [1, 2, 3], "colB": [4, 5, 6]})
-    base_dir = tmp_path / "test"
+    base_dir = "shared-memory://test-tag-order/dataset"
 
-    for i in range(3):
-        mode = "append" if i > 0 else "create"
-        ds = lance.write_dataset(table, base_dir, mode=mode)
+    ds = lance.write_dataset(table, base_dir, mode="create")
+    for _ in range(2):
+        ds = lance.write_dataset(table, ds, mode="append")
 
     expected_tags = {"tag3": 3, "tag2": 2, "tag1": 1}
     for name, version in expected_tags.items():
@@ -1592,6 +1592,26 @@ def test_cleanup_old_versions(tmp_path):
     assert stats.old_versions == 1
 
 
+def test_cleanup_retains_active_version_lease(tmp_path: Path):
+    table = pa.table({"value": range(10)})
+    dataset = lance.write_dataset(table, tmp_path / "leased", mode="create")
+    historical = dataset.checkout_version(1)
+    dataset = lance.write_dataset(table, tmp_path / "leased", mode="overwrite")
+
+    with historical.acquire_version_lease(timedelta(minutes=1)) as lease:
+        assert lease.version == 1
+        assert lease.expires_at > datetime.now(tz=lease.expires_at.tzinfo)
+        lease.renew(timedelta(minutes=1))
+        stats = dataset.cleanup_old_versions(retain_versions=1)
+        assert stats.old_versions == 0
+        assert historical.to_table() == table
+
+    stats = dataset.cleanup_old_versions(retain_versions=1)
+    assert stats.old_versions == 1
+    with pytest.raises(OSError):
+        dataset.checkout_version(1)
+
+
 def test_explain_cleanup_old_versions(tmp_path):
     table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
     base_dir = tmp_path / "test"
@@ -1633,14 +1653,12 @@ def test_explain_cleanup_old_versions(tmp_path):
 
 def test_cleanup_error_when_tagged_old_versions(tmp_path):
     table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
-    base_dir = tmp_path / "test"
-    lance.write_dataset(table, base_dir)
-    lance.write_dataset(table, base_dir, mode="overwrite")
+    base_dir = "shared-memory://cleanup-error-tagged/dataset"
+    dataset = lance.write_dataset(table, base_dir)
+    dataset = lance.write_dataset(table, dataset, mode="overwrite")
     time.sleep(0.1)
     moment = datetime.now()
-    lance.write_dataset(table, base_dir, mode="overwrite")
-
-    dataset = lance.dataset(base_dir)
+    dataset = lance.write_dataset(table, dataset, mode="overwrite")
     dataset.tags.create("old-tag", 1)
     dataset.tags.create("another-old-tag", 2)
 
@@ -1662,14 +1680,12 @@ def test_cleanup_error_when_tagged_old_versions(tmp_path):
 
 def test_cleanup_around_tagged_old_versions(tmp_path):
     table = pa.Table.from_pydict({"a": range(100), "b": range(100)})
-    base_dir = tmp_path / "test"
-    lance.write_dataset(table, base_dir)
-    lance.write_dataset(table, base_dir, mode="overwrite")
+    base_dir = "shared-memory://cleanup-around-tagged/dataset"
+    dataset = lance.write_dataset(table, base_dir)
+    dataset = lance.write_dataset(table, dataset, mode="overwrite")
     time.sleep(0.1)
     moment = datetime.now()
-    lance.write_dataset(table, base_dir, mode="overwrite")
-
-    dataset = lance.dataset(base_dir)
+    dataset = lance.write_dataset(table, dataset, mode="overwrite")
     dataset.tags.create("old-tag", 1)
     dataset.tags.create("another-old-tag", 2)
     dataset.tags.create("tag-latest", 3)
@@ -6223,10 +6239,10 @@ def test_update_config_transaction(tmp_path: Path):
 
 
 def test_shallow_clone(tmp_path: Path):
-    """Shallow clone a filesystem dataset by version number and by tag.
+    """Shallow clone a filesystem dataset by version and released tag.
 
     Arrange:
-      - Create a source dataset at a filesystem path with two versions
+      - Create a filesystem source dataset with two versions
         (create v1, then overwrite to v2).
       - Create a tag "v1" pointing to version 1.
     Act:
@@ -6236,19 +6252,27 @@ def test_shallow_clone(tmp_path: Path):
       - Re-open cloned datasets and verify their tables equal the source
         version they were cloned from (schema and record count).
 
-    This test uses pathlib paths and tmp_path for cross-platform compatibility
-    and should not skip on Windows.
+    The tag fixture uses the stable released-client JSON representation.
     """
     # Prepare source dataset with two versions
     src_dir = tmp_path / "shallow_src"
     table_v1 = pa.table({"a": [1, 2, 3], "b": [10, 20, 30]})
-    lance.write_dataset(table_v1, src_dir, mode="create")
+    ds = lance.write_dataset(
+        table_v1, src_dir, mode="create", enable_v2_manifest_paths=False
+    )
 
     table_v2 = pa.table({"a": [4, 5, 6], "b": [40, 50, 60]})
-    ds = lance.write_dataset(table_v2, src_dir, mode="overwrite")
+    ds = lance.write_dataset(table_v2, ds, mode="overwrite")
 
-    # Create a tag pointing to version 1
-    ds.tags.create("v1", 1)
+    # Emulate the stable canonical payload emitted by a released client. Local
+    # stores cannot safely execute current conditional reference deletion.
+    tag_path = src_dir / "_refs" / "tags" / "v1.json"
+    tag_path.parent.mkdir(parents=True)
+    manifest_size = (src_dir / "_versions" / "1.manifest").stat().st_size
+    tag_path.write_text(
+        f'{{"branch":null,"version":1,"manifestSize":{manifest_size},"metadata":{{}}}}',
+        encoding="utf-8",
+    )
 
     # Clone by numeric version (v2) and assert equality
     clone_v2_dir = tmp_path / "clone_v2"
@@ -6262,18 +6286,10 @@ def test_shallow_clone(tmp_path: Path):
     assert ds_clone_v1_tag.to_table() == table_v1
     assert lance.dataset(clone_v1_tag_dir).to_table() == table_v1
 
-    table_v3 = pa.table({"a": [7, 8, 9], "b": [40, 50, 60]})
-    branch = ds.create_branch("branch", 2)
-    lance.write_dataset(table_v3, branch.uri, mode="overwrite")
-    clone_branch_v3 = tmp_path / "clone_branch_v3"
-    cloned_by_branch = branch.shallow_clone(clone_branch_v3, 3)
-    assert cloned_by_branch.to_table() == table_v3
-    assert lance.dataset(clone_branch_v3).to_table() == table_v3
-
 
 def test_branches(tmp_path: Path):
     # Step 1: create branch1 from main → append to branch1 → create branch2 from tag
-    base_dir = tmp_path / "test_branches"
+    base_dir = "shared-memory://branches/dataset"
     main_table = pa.Table.from_pydict({"a": [1, 2, 3], "b": [4, 5, 6]})
     ds_main = lance.write_dataset(main_table, base_dir)