Merge branch 'main' into icechunk-js

Author: mpiannucci

icechunk-python/python/icechunk/_icechunk_python.pyi

@@ -1223,6 +1223,7 @@ class RepositoryConfig:
         virtual_chunk_containers: dict[str, VirtualChunkContainer] | None = None,
         manifest: ManifestConfig | None = None,
         repo_update_retries: RepoUpdateRetryConfig | None = None,
+        num_updates_per_repo_info_file: int | None = None,
     ) -> RepositoryConfig:
         """
         Create a new `RepositoryConfig` object
@@ -1439,6 +1440,12 @@ class RepositoryConfig:
         Clear all virtual chunk containers from the repository.
         """
         ...
+    @property
+    def num_updates_per_repo_info_file(self) -> int | None:
+        """The number of updates per repo info file."""
+        ...
+    @num_updates_per_repo_info_file.setter
+    def num_updates_per_repo_info_file(self, value: int | None) -> None: ...
     def merge(self, other: RepositoryConfig) -> RepositoryConfig:
         """
         Merge another RepositoryConfig with this one.

icechunk-python/python/icechunk/testing/strategies.py

@@ -54,3 +54,24 @@ def splitting_configs(
                 key: draw(st.integers(min_value=1, max_value=size + 10))
             }
     return ic.ManifestSplittingConfig.from_dict(config_dict)
+
+
+@st.composite
+def repository_configs(
+    draw: st.DrawFn,
+    num_updates_per_repo_info_file: st.SearchStrategy[int] = st.integers(  # noqa: B008
+        min_value=1, max_value=50
+    ),
+    inline_chunk_threshold_bytes: st.SearchStrategy[int] | None = None,
+    splitting: st.SearchStrategy[ic.ManifestSplittingConfig] | None = None,
+) -> ic.RepositoryConfig:
+    manifest = None
+    if splitting is not None:
+        manifest = ic.ManifestConfig(splitting=draw(splitting))
+    return ic.RepositoryConfig(
+        num_updates_per_repo_info_file=draw(num_updates_per_repo_info_file),
+        inline_chunk_threshold_bytes=draw(inline_chunk_threshold_bytes)
+        if inline_chunk_threshold_bytes is not None
+        else None,
+        manifest=manifest,
+    )

icechunk-python/src/config.rs

@@ -1529,6 +1529,8 @@ pub struct PyRepositoryConfig {
     pub previous_file: Option<String>,
     #[pyo3(get, set)]
     pub repo_update_retries: Option<Py<PyRepoUpdateRetryConfig>>,
+    #[pyo3(get, set)]
+    pub num_updates_per_repo_info_file: Option<u16>,
 }
 
 impl PartialEq for PyRepositoryConfig {
@@ -1567,6 +1569,7 @@ impl TryFrom<&PyRepositoryConfig> for RepositoryConfig {
                     .repo_update_retries
                     .as_ref()
                     .map(|r| (&*r.borrow(py)).into()),
+                num_updates_per_repo_info_file: value.num_updates_per_repo_info_file,
             })
         })
     }
@@ -1604,6 +1607,7 @@ impl From<RepositoryConfig> for PyRepositoryConfig {
                 Py::new(py, Into::<PyRepoUpdateRetryConfig>::into(r))
                     .expect("Cannot create instance of RepoUpdateRetryConfig")
             }),
+            num_updates_per_repo_info_file: value.num_updates_per_repo_info_file,
         })
     }
 }
@@ -1617,7 +1621,7 @@ impl PyRepositoryConfig {
     }
 
     #[new]
-    #[pyo3(signature = (inline_chunk_threshold_bytes = None, get_partial_values_concurrency = None, compression = None, max_concurrent_requests = None, caching = None, storage = None, virtual_chunk_containers = None, manifest = None, repo_update_retries = None))]
+    #[pyo3(signature = (inline_chunk_threshold_bytes = None, get_partial_values_concurrency = None, compression = None, max_concurrent_requests = None, caching = None, storage = None, virtual_chunk_containers = None, manifest = None, repo_update_retries = None, num_updates_per_repo_info_file = None))]
     #[allow(clippy::too_many_arguments)]
     pub fn new(
         inline_chunk_threshold_bytes: Option<u16>,
@@ -1629,6 +1633,7 @@ impl PyRepositoryConfig {
         virtual_chunk_containers: Option<HashMap<String, PyVirtualChunkContainer>>,
         manifest: Option<Py<PyManifestConfig>>,
         repo_update_retries: Option<Py<PyRepoUpdateRetryConfig>>,
+        num_updates_per_repo_info_file: Option<u16>,
     ) -> Self {
         Self {
             inline_chunk_threshold_bytes,
@@ -1641,6 +1646,7 @@ impl PyRepositoryConfig {
             manifest,
             previous_file: None,
             repo_update_retries,
+            num_updates_per_repo_info_file,
         }
     }
 

icechunk-python/src/repository.rs

@@ -2131,9 +2131,10 @@ impl PyRepository {
         py.detach(move || {
             let result =
                 pyo3_async_runtimes::tokio::get_runtime().block_on(async move {
-                    let asset_manager = {
+                    let (asset_manager, num_updates) = {
                         let lock = self.0.read().await;
-                        Arc::clone(lock.asset_manager())
+                        let num_updates = lock.config().num_updates_per_repo_info_file();
+                        (Arc::clone(lock.asset_manager()), num_updates)
                     };
 
                     let result = expire(
@@ -2150,6 +2151,7 @@ impl PyRepository {
                             ExpiredRefAction::Ignore
                         },
                         None,
+                        num_updates,
                     )
                     .await
                     .map_err(PyIcechunkStoreError::GCError)?;
@@ -2176,9 +2178,10 @@ impl PyRepository {
     ) -> PyResult<Bound<'py, PyAny>> {
         let repository = self.0.clone();
         pyo3_async_runtimes::tokio::future_into_py::<_, HashSet<String>>(py, async move {
-            let asset_manager = {
+            let (asset_manager, num_updates) = {
                 let lock = repository.read().await;
-                Arc::clone(lock.asset_manager())
+                let num_updates = lock.config().num_updates_per_repo_info_file();
+                (Arc::clone(lock.asset_manager()), num_updates)
             };
 
             let result = expire(
@@ -2195,6 +2198,7 @@ impl PyRepository {
                     ExpiredRefAction::Ignore
                 },
                 None,
+                num_updates,
             )
             .await
             .map_err(PyIcechunkStoreError::GCError)?;
@@ -2224,13 +2228,15 @@ impl PyRepository {
                         max_concurrent_manifest_fetches,
                         dry_run,
                     );
-                    let asset_manager = {
+                    let (asset_manager, num_updates) = {
                         let lock = self.0.read().await;
-                        Arc::clone(lock.asset_manager())
+                        let num_updates = lock.config().num_updates_per_repo_info_file();
+                        (Arc::clone(lock.asset_manager()), num_updates)
                     };
-                    let result = garbage_collect(asset_manager, &gc_config, None)
-                        .await
-                        .map_err(PyIcechunkStoreError::GCError)?;
+                    let result =
+                        garbage_collect(asset_manager, &gc_config, None, num_updates)
+                            .await
+                            .map_err(PyIcechunkStoreError::GCError)?;
                     Ok::<_, PyIcechunkStoreError>(result.into())
                 })?;
 
@@ -2258,11 +2264,12 @@ impl PyRepository {
                 max_concurrent_manifest_fetches,
                 dry_run,
             );
-            let asset_manager = {
+            let (asset_manager, num_updates) = {
                 let lock = repository.read().await;
-                Arc::clone(lock.asset_manager())
+                let num_updates = lock.config().num_updates_per_repo_info_file();
+                (Arc::clone(lock.asset_manager()), num_updates)
             };
-            let result = garbage_collect(asset_manager, &gc_config, None)
+            let result = garbage_collect(asset_manager, &gc_config, None, num_updates)
                 .await
                 .map_err(PyIcechunkStoreError::GCError)?;
             Ok(result.into())

icechunk-python/tests/test_stateful_repo_ops.py

@@ -39,10 +39,12 @@
 from icechunk import (
     IcechunkError,
     Repository,
+    RepositoryConfig,
     SnapshotInfo,
     Storage,
     in_memory_storage,
 )
+from icechunk.testing.strategies import repository_configs
 from zarr.testing.stateful import SyncStoreWrapper
 
 # JSON file contents, keep it simple
@@ -110,7 +112,10 @@ class TagModel:
 
 class Model:
     def __init__(self, **kwargs: Any) -> None:
-        self.store: dict[str, Any] = {}  #
+        self.store: dict[str, Any] = {}
+
+        self.spec_version = 1  # will be overwritten on `@initialize`
+        self.num_updates: int = 0
 
         self.initial_snapshot_id: str | None = None
         self.changes_made: bool = False
@@ -151,6 +156,9 @@ def __setitem__(self, key: str, value: Buffer) -> None:
     def __getitem__(self, key: str) -> Buffer:
         return cast(Buffer, self.store[key])
 
+    def upgrade(self) -> None:
+        self.num_updates += 1
+
     @property
     def has_commits(self) -> bool:
         return bool(self.commits)
@@ -171,10 +179,11 @@ def commit(self, snap: SnapshotInfo) -> None:
 
         assert self.branch is not None
         self.branch_heads[self.branch] = ref
+        self.num_updates += 1
 
     def amend(self, snap: SnapshotInfo) -> None:
         """Amend the HEAD commit."""
-        # this is simpe because we aren't modeling the branch as a list of commits
+        # this is simple because we aren't modeling the branch as a list of commits
         self.commit(snap)
 
     def checkout_commit(self, ref: str) -> None:
@@ -189,6 +198,7 @@ def checkout_commit(self, ref: str) -> None:
     def create_branch(self, name: str, commit: str) -> None:
         assert commit in self.commits
         self.branch_heads[name] = commit
+        self.num_updates += 1
 
     def checkout_branch(self, ref: str) -> None:
         self.checkout_commit(self.branch_heads[ref])
@@ -198,17 +208,27 @@ def checkout_branch(self, ref: str) -> None:
     def reset_branch(self, branch: str, commit: str) -> None:
         assert commit in self.commits
         self.branch_heads[branch] = commit
+        self.num_updates += 1
 
     def delete_branch(self, branch_name: str) -> None:
+        self._delete_branch(branch_name)
+        self.num_updates += 1
+
+    def _delete_branch(self, branch_name: str) -> None:
         del self.branch_heads[branch_name]
 
     def delete_tag(self, tag: str) -> None:
+        self._delete_tag(tag)
+        self.num_updates += 1
+
+    def _delete_tag(self, tag: str) -> None:
         del self.tags[tag]
 
     def create_tag(self, tag_name: str, commit_id: str) -> None:
         assert commit_id in self.commits
         self.tags[tag_name] = TagModel(commit_id=str(commit_id))
         self.created_tags.add(tag_name)
+        self.num_updates += 1
 
     def checkout_tag(self, ref: str) -> None:
         self.checkout_commit(self.tags[str(ref)].commit_id)
@@ -261,7 +281,7 @@ def expire_snapshots(
             }
             note(f"deleting tags {tags_to_delete=!r}")
             for tag in tags_to_delete:
-                self.delete_tag(tag)
+                self._delete_tag(tag)
         else:
             tags_to_delete = set()
 
@@ -274,10 +294,12 @@ def expire_snapshots(
             note(f"deleting branches {branches_to_delete=!r}")
             for branch in branches_to_delete:
                 note(f"deleting {branch=!r}, {self.branch_heads[branch]=!r}")
-                self.delete_branch(branch)
+                self._delete_branch(branch)
         else:
             branches_to_delete = set()
 
+        self.num_updates += 1
+
         return ExpireInfo(
             expired_snapshots=expired_snaps,
             deleted_branches=branches_to_delete,
@@ -301,6 +323,7 @@ def garbage_collect(self, older_than: datetime.datetime) -> set[str]:
                 self.ondisk_snaps.pop(k, None)
                 deleted.add(k)
         note(f"Deleted snapshots in model: {deleted!r}")
+        self.num_updates += 1
         return deleted
 
 
@@ -325,7 +348,14 @@ def __init__(self) -> None:
     @initialize(data=st.data(), target=branches, spec_version=st.sampled_from([1, 2]))
     def initialize(self, data: st.DataObject, spec_version: Literal[1, 2]) -> str:
         self.storage = in_memory_storage()
-        self.repo = Repository.create(self.storage, spec_version=spec_version)
+        config = data.draw(repository_configs())
+        self.model.spec_version = spec_version
+
+        self.repo = Repository.create(
+            self.storage,
+            spec_version=spec_version,
+            config=config,
+        )
         self.session = self.repo.writable_session(DEFAULT_BRANCH)
 
         snap = next(iter(self.repo.ancestry(branch=DEFAULT_BRANCH)))
@@ -338,6 +368,9 @@ def initialize(self, data: st.DataObject, spec_version: Literal[1, 2]) -> str:
         self.model.HEAD = HEAD
         self.model.create_branch(DEFAULT_BRANCH, HEAD)
         self.model.checkout_branch(DEFAULT_BRANCH)
+        # RepoInitializedUpdate includes the initial branch creation,
+        # so reset to 1 after create_branch incremented it.
+        self.model.num_updates = 1
 
         # initialize with some data always
         # TODO: always setting array metadata, since we cannot overwrite an existing group's zarr.json
@@ -371,17 +404,22 @@ def upgrade_spec_version(self) -> None:
         # don't test simple cases of catching error upgradging a v2 spec
         # that should be covered in unit tests
         icechunk.upgrade_icechunk_repository(self.repo)
+        self.model.upgrade()
         # TODO: remove the reopen after https://github.com/earth-mover/icechunk/issues/1521
-        self.reopen_repository()
+        self._reopen_repository()
 
-    @rule()
-    def reopen_repository(self) -> None:
+    @rule(data=st.data())
+    def reopen_repository(self, data: st.DataObject) -> None:
+        config = data.draw(repository_configs())
+        self._reopen_repository(config)
+
+    def _reopen_repository(self, config: RepositoryConfig | None = None) -> None:
         """Reopen the repository from storage to get fresh state.
 
         This discards any uncommitted changes.
         """
         assert self.storage is not None, "storage must be initialized"
-        self.repo = Repository.open(self.storage)
+        self.repo = Repository.open(self.storage, config=config)
         note(f"Reopened repository (spec_version={self.repo.spec_version})")
 
         # Reopening discards uncommitted changes - reset model to last committed state
@@ -688,6 +726,17 @@ def check_commit(self, commit: str) -> None:
         # even after expiration, written_at is unmodified
         assert actual.written_at == expected.written_at
 
+    def check_ops_log(self) -> None:
+        if self.model.spec_version == 1:
+            return
+        actual_ops = list(self.repo.ops_log())
+        assert len(actual_ops) == self.model.num_updates, (
+            actual_ops,
+            self.model.num_updates,
+            actual_ops,
+        )
+        assert isinstance(actual_ops[-1], icechunk.RepoInitializedUpdate)
+
     @invariant()
     def checks(self) -> None:
         # this method only exists to reduce verbosity of hypothesis output
@@ -697,6 +746,7 @@ def checks(self) -> None:
         self.check_tags()
         self.check_branches()
         self.check_ancestry()
+        self.check_ops_log()
 
     def check_list_prefix_from_root(self) -> None:
         model_list = self.model.list_prefix("")

icechunk-python/tests/test_zarr/test_stateful.py

@@ -109,11 +109,13 @@ def reopen_with_config(self, data: st.DataObject) -> None:
             st.lists(st.sampled_from(sorted(self.all_arrays)), max_size=3, unique=True)
         )
         arrays = tuple(zarr.open_array(self.model, path=path) for path in array_paths)
-        sconfig = data.draw(icst.splitting_configs(arrays=arrays))
-        config = ic.RepositoryConfig(
-            inline_chunk_threshold_bytes=0, manifest=ic.ManifestConfig(splitting=sconfig)
+        config = data.draw(
+            icst.repository_configs(
+                inline_chunk_threshold_bytes=st.just(0),
+                splitting=icst.splitting_configs(arrays=arrays),
+            )
         )
-        note(f"reopening with splitting config {sconfig=!r}")
+        note(f"reopening with config {config!r}")
         self.repo = Repository.open(self.storage, config=config)
         if data.draw(st.booleans()):
             self.repo.save_config()