icechunk 0.2.0__tar.gz → 0.2.2__tar.gz

{icechunk-0.2.0 → icechunk-0.2.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: icechunk
-Version: 0.2.0
+Version: 0.2.2
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "icechunk"
-version = "0.2.0"
+version = "0.2.2"
 description = "Transactional storage engine for Zarr designed for use on cloud object storage"
 readme = "README.md"
 repository = "https://github.com/earth-mover/icechunk"

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/examples/multithreaded_store.rs

@@ -54,6 +54,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
     async fn writer(name: &str, range: Range<u64>, store: &Store) {
         println!("Starting writer {name}.");
         for i in range {
+            #[allow(clippy::dbg_macro)]
             if let Err(err) = store
                 .set(
                     format!("array/c/{i}").as_str(),

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/src/repository.rs

@@ -5,7 +5,9 @@ use std::{
     sync::Arc,
 };
 
+use async_recursion::async_recursion;
 use bytes::Bytes;
+use chrono::{DateTime, Utc};
 use err_into::ErrorInto as _;
 use futures::{
     stream::{FuturesOrdered, FuturesUnordered},
@@ -37,15 +39,13 @@ use crate::{
     Storage, StorageError,
 };
 
-#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+#[derive(Debug, Clone, PartialEq, Eq)]
 #[non_exhaustive]
 pub enum VersionInfo {
-    #[serde(rename = "snapshot_id")]
     SnapshotId(SnapshotId),
-    #[serde(rename = "tag")]
     TagRef(String),
-    #[serde(rename = "branch")]
     BranchTipRef(String),
+    AsOf { branch: String, at: DateTime<Utc> },
 }
 
 #[derive(Debug, Error)]
@@ -60,6 +60,8 @@ pub enum RepositoryErrorKind {
 
     #[error("snapshot not found: `{id}`")]
     SnapshotNotFound { id: SnapshotId },
+    #[error("branch {branch} does not have a snapshots before or at {at}")]
+    InvalidAsOfSpec { branch: String, at: DateTime<Utc> },
     #[error("invalid snapshot id: `{0}`")]
     InvalidSnapshotId(String),
     #[error("tag error: `{0}`")]
@@ -404,11 +406,12 @@ impl Repository {
     }
 
     /// Returns the sequence of parents of the snapshot pointed by the given version
+    #[async_recursion(?Send)]
     #[instrument(skip(self))]
-    pub async fn ancestry(
-        &self,
+    pub async fn ancestry<'a>(
+        &'a self,
         version: &VersionInfo,
-    ) -> RepositoryResult<impl Stream<Item = RepositoryResult<SnapshotInfo>> + '_> {
+    ) -> RepositoryResult<impl Stream<Item = RepositoryResult<SnapshotInfo>> + 'a> {
         let snapshot_id = self.resolve_version(version).await?;
         self.snapshot_ancestry(&snapshot_id).await
     }
@@ -572,6 +575,24 @@ impl Repository {
                 .await?;
                 Ok(ref_data.snapshot)
             }
+            VersionInfo::AsOf { branch, at } => {
+                let tip = VersionInfo::BranchTipRef(branch.clone());
+                let snap = self
+                    .ancestry(&tip)
+                    .await?
+                    .try_skip_while(|parent| ready(Ok(&parent.flushed_at > at)))
+                    .take(1)
+                    .try_collect::<Vec<_>>()
+                    .await?;
+                match snap.into_iter().next() {
+                    Some(snap) => Ok(snap.id),
+                    None => Err(RepositoryErrorKind::InvalidAsOfSpec {
+                        branch: branch.clone(),
+                        at: *at,
+                    }
+                    .into()),
+                }
+            }
         }
     }
 

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/src/storage/mod.rs

@@ -174,6 +174,9 @@ impl ConcurrencySettings {
 #[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Default)]
 pub struct Settings {
     pub concurrency: Option<ConcurrencySettings>,
+    pub unsafe_use_conditional_update: Option<bool>,
+    pub unsafe_use_conditional_create: Option<bool>,
+    pub unsafe_use_metadata: Option<bool>,
 }
 
 static DEFAULT_CONCURRENCY: OnceLock<ConcurrencySettings> = OnceLock::new();
@@ -185,6 +188,18 @@ impl Settings {
             .unwrap_or_else(|| DEFAULT_CONCURRENCY.get_or_init(Default::default))
     }
 
+    pub fn unsafe_use_conditional_create(&self) -> bool {
+        self.unsafe_use_conditional_create.unwrap_or(true)
+    }
+
+    pub fn unsafe_use_conditional_update(&self) -> bool {
+        self.unsafe_use_conditional_update.unwrap_or(true)
+    }
+
+    pub fn unsafe_use_metadata(&self) -> bool {
+        self.unsafe_use_metadata.unwrap_or(true)
+    }
+
     pub fn merge(&self, other: Self) -> Self {
         Self {
             concurrency: match (&self.concurrency, other.concurrency) {
@@ -193,6 +208,33 @@ impl Settings {
                 (Some(c), None) => Some(c.clone()),
                 (Some(mine), Some(theirs)) => Some(mine.merge(theirs)),
             },
+            unsafe_use_conditional_create: match (
+                &self.unsafe_use_conditional_create,
+                other.unsafe_use_conditional_create,
+            ) {
+                (None, None) => None,
+                (None, Some(c)) => Some(c),
+                (Some(c), None) => Some(*c),
+                (Some(_), Some(theirs)) => Some(theirs),
+            },
+            unsafe_use_conditional_update: match (
+                &self.unsafe_use_conditional_update,
+                other.unsafe_use_conditional_update,
+            ) {
+                (None, None) => None,
+                (None, Some(c)) => Some(c),
+                (Some(c), None) => Some(*c),
+                (Some(_), Some(theirs)) => Some(theirs),
+            },
+            unsafe_use_metadata: match (
+                &self.unsafe_use_metadata,
+                other.unsafe_use_metadata,
+            ) {
+                (None, None) => None,
+                (None, Some(c)) => Some(c),
+                (Some(c), None) => Some(*c),
+                (Some(_), Some(theirs)) => Some(theirs),
+            },
         }
     }
 }
@@ -364,7 +406,11 @@ pub trait Storage: fmt::Debug + private::Sealed + Sync + Send {
     ) -> StorageResult<DateTime<Utc>>;
 
     async fn root_is_clean(&self) -> StorageResult<bool> {
-        Ok(self.list_objects(&Settings::default(), "").await?.next().await.is_none())
+        match self.list_objects(&Settings::default(), "").await?.next().await {
+            None => Ok(true),
+            Some(Ok(_)) => Ok(false),
+            Some(Err(err)) => Err(err),
+        }
     }
 
     async fn list_chunks(

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/src/storage/object_store.rs

@@ -158,16 +158,6 @@ impl ObjectStorage {
         self.backend.artificially_sort_refs_in_mem()
     }
 
-    /// We need this because object_store's local file implementation doesn't support metadata.
-    pub fn supports_metadata(&self) -> bool {
-        self.backend.supports_metadata()
-    }
-
-    /// We need this because object_store's local file implementation doesn't support it
-    pub fn supports_conditional_put_updates(&self) -> bool {
-        self.backend.supports_conditional_put_updates()
-    }
-
     /// Return all keys in the store
     ///
     /// Intended for testing and debugging purposes only.
@@ -251,8 +241,12 @@ impl ObjectStorage {
             .compat())
     }
 
-    fn metadata_to_attributes(&self, metadata: Vec<(String, String)>) -> Attributes {
-        if self.supports_metadata() {
+    fn metadata_to_attributes(
+        &self,
+        settings: &Settings,
+        metadata: Vec<(String, String)>,
+    ) -> Attributes {
+        if settings.unsafe_use_metadata() {
             Attributes::from_iter(metadata.into_iter().map(|(key, val)| {
                 (
                     Attribute::Metadata(std::borrow::Cow::Owned(key)),
@@ -270,18 +264,24 @@ impl ObjectStorage {
         Some(parent.as_ref().to_string())
     }
 
-    fn get_put_mode(&self, previous_version: &VersionInfo) -> PutMode {
-        let degrade_to_overwrite =
-            !previous_version.is_create() && !self.supports_conditional_put_updates();
-        if degrade_to_overwrite {
-            PutMode::Overwrite
-        } else if previous_version.is_create() {
-            PutMode::Create
-        } else {
-            PutMode::Update(UpdateVersion {
+    fn get_put_mode(
+        &self,
+        settings: &Settings,
+        previous_version: &VersionInfo,
+    ) -> PutMode {
+        match (
+            previous_version.is_create(),
+            settings.unsafe_use_conditional_create(),
+            settings.unsafe_use_conditional_update(),
+        ) {
+            (true, true, _) => PutMode::Create,
+            (true, false, _) => PutMode::Overwrite,
+
+            (false, _, true) => PutMode::Update(UpdateVersion {
                 e_tag: previous_version.etag().cloned(),
                 version: previous_version.generation().cloned(),
-            })
+            }),
+            (false, _, false) => PutMode::Overwrite,
         }
     }
 }
@@ -317,15 +317,15 @@ impl Storage for ObjectStorage {
             Err(err) => Err(err.into()),
         }
     }
-    #[instrument(skip(self, _settings, config))]
+    #[instrument(skip(self, settings, config))]
     async fn update_config(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         config: Bytes,
         previous_version: &VersionInfo,
     ) -> StorageResult<UpdateConfigResult> {
         let path = self.get_config_path();
-        let attributes = if self.supports_metadata() {
+        let attributes = if settings.unsafe_use_metadata() {
             Attributes::from_iter(vec![(
                 Attribute::ContentType,
                 AttributeValue::from("application/yaml"),
@@ -334,7 +334,7 @@ impl Storage for ObjectStorage {
             Attributes::new()
         };
 
-        let mode = self.get_put_mode(previous_version);
+        let mode = self.get_put_mode(settings, previous_version);
 
         let options = PutOptions { mode, attributes, ..PutOptions::default() };
         let res = self.get_client().await.put_opts(&path, config.into(), options).await;
@@ -394,48 +394,48 @@ impl Storage for ObjectStorage {
         Ok(Box::new(self.get_object_reader(settings, &path).await?))
     }
 
-    #[instrument(skip(self, _settings, metadata, bytes))]
+    #[instrument(skip(self, settings, metadata, bytes))]
     async fn write_snapshot(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: SnapshotId,
         metadata: Vec<(String, String)>,
         bytes: Bytes,
     ) -> StorageResult<()> {
         let path = self.get_snapshot_path(&id);
-        let attributes = self.metadata_to_attributes(metadata);
+        let attributes = self.metadata_to_attributes(settings, metadata);
         let options = PutOptions { attributes, ..PutOptions::default() };
         // FIXME: use multipart
         self.get_client().await.put_opts(&path, bytes.into(), options).await?;
         Ok(())
     }
 
-    #[instrument(skip(self, _settings, metadata, bytes))]
+    #[instrument(skip(self, settings, metadata, bytes))]
     async fn write_manifest(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: ManifestId,
         metadata: Vec<(String, String)>,
         bytes: Bytes,
     ) -> StorageResult<()> {
         let path = self.get_manifest_path(&id);
-        let attributes = self.metadata_to_attributes(metadata);
+        let attributes = self.metadata_to_attributes(settings, metadata);
         let options = PutOptions { attributes, ..PutOptions::default() };
         // FIXME: use multipart
         self.get_client().await.put_opts(&path, bytes.into(), options).await?;
         Ok(())
     }
 
-    #[instrument(skip(self, _settings, metadata, bytes))]
+    #[instrument(skip(self, settings, metadata, bytes))]
     async fn write_transaction_log(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: SnapshotId,
         metadata: Vec<(String, String)>,
         bytes: Bytes,
     ) -> StorageResult<()> {
         let path = self.get_transaction_path(&id);
-        let attributes = self.metadata_to_attributes(metadata);
+        let attributes = self.metadata_to_attributes(settings, metadata);
         let options = PutOptions { attributes, ..PutOptions::default() };
         // FIXME: use multipart
         self.get_client().await.put_opts(&path, bytes.into(), options).await?;
@@ -506,16 +506,16 @@ impl Storage for ObjectStorage {
             .await?)
     }
 
-    #[instrument(skip(self, _settings, bytes))]
+    #[instrument(skip(self, settings, bytes))]
     async fn write_ref(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         ref_key: &str,
         bytes: Bytes,
         previous_version: &VersionInfo,
     ) -> StorageResult<WriteRefResult> {
         let key = self.ref_key(ref_key);
-        let mode = self.get_put_mode(previous_version);
+        let mode = self.get_put_mode(settings, previous_version);
         let opts = PutOptions { mode, ..PutOptions::default() };
 
         match self
@@ -634,19 +634,7 @@ pub trait ObjectStoreBackend: Debug + Sync + Send {
         false
     }
 
-    /// We need this because object_store's local file implementation doesn't support metadata.
-    fn supports_metadata(&self) -> bool {
-        true
-    }
-
-    /// We need this because object_store's local file implementation doesn't support it
-    fn supports_conditional_put_updates(&self) -> bool {
-        true
-    }
-
-    fn default_settings(&self) -> Settings {
-        Settings::default()
-    }
+    fn default_settings(&self) -> Settings;
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -674,6 +662,7 @@ impl ObjectStoreBackend for InMemoryObjectStoreBackend {
                     NonZeroU64::new(1).unwrap_or(NonZeroU64::MIN),
                 ),
             }),
+            ..Default::default()
         }
     }
 }
@@ -706,14 +695,6 @@ impl ObjectStoreBackend for LocalFileSystemObjectStoreBackend {
         true
     }
 
-    fn supports_metadata(&self) -> bool {
-        false
-    }
-
-    fn supports_conditional_put_updates(&self) -> bool {
-        false
-    }
-
     fn default_settings(&self) -> Settings {
         Settings {
             concurrency: Some(ConcurrencySettings {
@@ -724,6 +705,9 @@ impl ObjectStoreBackend for LocalFileSystemObjectStoreBackend {
                     NonZeroU64::new(4 * 1024).unwrap_or(NonZeroU64::MIN),
                 ),
             }),
+            unsafe_use_conditional_update: Some(false),
+            unsafe_use_metadata: Some(false),
+            ..Default::default()
         }
     }
 }
@@ -795,6 +779,10 @@ impl ObjectStoreBackend for S3ObjectStoreBackend {
     fn prefix(&self) -> String {
         self.prefix.clone().unwrap_or("".to_string())
     }
+
+    fn default_settings(&self) -> Settings {
+        Default::default()
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -844,6 +832,10 @@ impl ObjectStoreBackend for AzureObjectStoreBackend {
     fn prefix(&self) -> String {
         self.prefix.clone().unwrap_or("".to_string())
     }
+
+    fn default_settings(&self) -> Settings {
+        Default::default()
+    }
 }
 
 #[derive(Debug, Serialize, Deserialize)]
@@ -910,6 +902,10 @@ impl ObjectStoreBackend for GcsObjectStoreBackend {
     fn prefix(&self) -> String {
         self.prefix.clone().unwrap_or("".to_string())
     }
+
+    fn default_settings(&self) -> Settings {
+        Default::default()
+    }
 }
 
 #[derive(Debug)]

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/src/storage/s3.rs

@@ -196,6 +196,7 @@ impl S3Storage {
         I: IntoIterator<Item = (impl Into<String>, impl Into<String>)>,
     >(
         &self,
+        settings: &Settings,
         key: &str,
         content_type: Option<impl Into<String>>,
         metadata: I,
@@ -204,14 +205,17 @@ impl S3Storage {
         let mut b =
             self.get_client().await.put_object().bucket(self.bucket.clone()).key(key);
 
-        if let Some(ct) = content_type {
-            b = b.content_type(ct)
-        };
-
-        for (k, v) in metadata {
-            b = b.metadata(k, v);
+        if settings.unsafe_use_metadata() {
+            if let Some(ct) = content_type {
+                b = b.content_type(ct)
+            };
         }
 
+        if settings.unsafe_use_metadata() {
+            for (k, v) in metadata {
+                b = b.metadata(k, v);
+            }
+        }
         b.body(bytes.into()).send().await?;
         Ok(())
     }
@@ -296,10 +300,10 @@ impl Storage for S3Storage {
         }
     }
 
-    #[instrument(skip(self, _settings, config))]
+    #[instrument(skip(self, settings, config))]
     async fn update_config(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         config: Bytes,
         previous_version: &VersionInfo,
     ) -> StorageResult<UpdateConfigResult> {
@@ -310,13 +314,20 @@ impl Storage for S3Storage {
             .put_object()
             .bucket(self.bucket.clone())
             .key(key)
-            .content_type("application/yaml")
             .body(config.into());
 
-        if let Some(etag) = previous_version.etag() {
-            req = req.if_match(etag)
-        } else {
-            req = req.if_none_match("*")
+        if settings.unsafe_use_metadata() {
+            req = req.content_type("application/yaml")
+        }
+
+        match (
+            previous_version.etag(),
+            settings.unsafe_use_conditional_create(),
+            settings.unsafe_use_conditional_update(),
+        ) {
+            (None, true, _) => req = req.if_none_match("*"),
+            (Some(etag), _, true) => req = req.if_match(etag),
+            (_, _, _) => {}
         }
 
         let res = req.send().await;
@@ -411,53 +422,67 @@ impl Storage for S3Storage {
             .await
     }
 
-    #[instrument(skip(self, _settings, metadata, bytes))]
+    #[instrument(skip(self, settings, metadata, bytes))]
     async fn write_snapshot(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: SnapshotId,
         metadata: Vec<(String, String)>,
         bytes: Bytes,
     ) -> StorageResult<()> {
         let key = self.get_snapshot_path(&id)?;
-        self.put_object(key.as_str(), None::<String>, metadata, bytes).await
+        self.put_object(settings, key.as_str(), None::<String>, metadata, bytes).await
     }
 
-    #[instrument(skip(self, _settings, metadata, bytes))]
+    #[instrument(skip(self, settings, metadata, bytes))]
     async fn write_manifest(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: ManifestId,
         metadata: Vec<(String, String)>,
         bytes: Bytes,
     ) -> StorageResult<()> {
         let key = self.get_manifest_path(&id)?;
-        self.put_object(key.as_str(), None::<String>, metadata.into_iter(), bytes).await
+        self.put_object(
+            settings,
+            key.as_str(),
+            None::<String>,
+            metadata.into_iter(),
+            bytes,
+        )
+        .await
     }
 
-    #[instrument(skip(self, _settings, metadata, bytes))]
+    #[instrument(skip(self, settings, metadata, bytes))]
     async fn write_transaction_log(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: SnapshotId,
         metadata: Vec<(String, String)>,
         bytes: Bytes,
     ) -> StorageResult<()> {
         let key = self.get_transaction_path(&id)?;
-        self.put_object(key.as_str(), None::<String>, metadata.into_iter(), bytes).await
+        self.put_object(
+            settings,
+            key.as_str(),
+            None::<String>,
+            metadata.into_iter(),
+            bytes,
+        )
+        .await
     }
 
-    #[instrument(skip(self, _settings, bytes))]
+    #[instrument(skip(self, settings, bytes))]
     async fn write_chunk(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         id: ChunkId,
         bytes: bytes::Bytes,
     ) -> Result<(), StorageError> {
         let key = self.get_chunk_path(&id)?;
         //FIXME: use multipart upload
         let metadata: [(String, String); 0] = [];
-        self.put_object(key.as_str(), None::<String>, metadata, bytes).await
+        self.put_object(settings, key.as_str(), None::<String>, metadata, bytes).await
     }
 
     #[instrument(skip(self, _settings))]
@@ -523,10 +548,10 @@ impl Storage for S3Storage {
         Ok(res)
     }
 
-    #[instrument(skip(self, _settings, bytes))]
+    #[instrument(skip(self, settings, bytes))]
     async fn write_ref(
         &self,
-        _settings: &Settings,
+        settings: &Settings,
         ref_key: &str,
         bytes: Bytes,
         previous_version: &VersionInfo,
@@ -539,10 +564,18 @@ impl Storage for S3Storage {
             .bucket(self.bucket.clone())
             .key(key.clone());
 
-        if let Some(etag) = previous_version.etag() {
-            builder = builder.if_match(etag);
-        } else {
-            builder = builder.if_none_match("*");
+        match (
+            previous_version.etag(),
+            settings.unsafe_use_conditional_create(),
+            settings.unsafe_use_conditional_update(),
+        ) {
+            (None, true, _) => {
+                builder = builder.if_none_match("*");
+            }
+            (Some(etag), _, true) => {
+                builder = builder.if_match(etag);
+            }
+            (_, _, _) => {}
         }
 
         let res = builder.body(bytes.into()).send().await;

{icechunk-0.2.0 → icechunk-0.2.2}/icechunk/tests/test_storage.rs

@@ -14,7 +14,7 @@ use icechunk::{
         RefErrorKind,
     },
     storage::{
-        new_in_memory_storage, new_s3_storage, ETag, FetchConfigResult, Generation,
+        self, new_in_memory_storage, new_s3_storage, ETag, FetchConfigResult, Generation,
         StorageResult, UpdateConfigResult, VersionInfo,
     },
     ObjectStorage, Storage,
@@ -488,3 +488,55 @@ pub async fn test_write_config_fails_on_bad_version_when_existing(
     }).await?;
     Ok(())
 }
+
+#[tokio::test]
+#[allow(clippy::panic)]
+pub async fn test_write_config_can_overwrite_with_unsafe_config(
+) -> Result<(), Box<dyn std::error::Error>> {
+    with_storage(|_, storage| async move {
+        let storage_settings = storage::Settings {
+            unsafe_use_conditional_update: Some(false),
+            unsafe_use_conditional_create: Some(false),
+            ..storage.default_settings()
+        };
+
+        // create the initial version
+        let config = Bytes::copy_from_slice(b"hello");
+        match storage
+            .update_config(
+                &storage_settings,
+                config.clone(),
+                &VersionInfo {
+                    etag: Some(ETag("some-bad-etag".to_string())),
+                    generation: Some(Generation("42".to_string())),
+                }
+            )
+            .await?
+        {
+            UpdateConfigResult::Updated { new_version } => new_version,
+            _ => panic!(),
+        };
+
+        // attempt a bad change that should succeed in this config
+        let update_res = storage
+            .update_config(
+                &storage_settings,
+                Bytes::copy_from_slice(b"bye"),
+                &VersionInfo {
+                    etag: Some(ETag("other-bad-etag".to_string())),
+                    generation: Some(Generation("55".to_string())),
+                },
+            )
+            .await?;
+
+        assert!(matches!(update_res, UpdateConfigResult::Updated { .. }));
+
+        let fetch_res = storage.fetch_config(&storage_settings).await?;
+        assert!(
+            matches!(fetch_res, FetchConfigResult::Found{bytes, ..} if bytes.as_ref() == b"bye")
+        );
+        Ok(())
+    })
+    .await?;
+    Ok(())
+}