Chagne key errors to use python builtins, add simple test for handling (

#365)
earth-mover · Oct 31, 2024 · 22acd16 · 22acd16
1 parent 2c9a311
commit 22acd16
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 31 deletions.
diff --git a/icechunk-python/python/icechunk/__init__.py b/icechunk-python/python/icechunk/__init__.py
@@ -8,7 +8,6 @@
 from zarr.core.sync import SyncMixin
 
 from ._icechunk_python import (
-    KeyNotFound,
     PyIcechunkStore,
     S3Credentials,
     SnapshotMetadata,
@@ -485,7 +484,7 @@ async def get(
 
         try:
             result = await self._store.get(key, byte_range)
-        except KeyNotFound as _e:
+        except KeyError as _e:
             # Zarr python expects None to be returned if the key does not exist
             # but an IcechunkStore returns an error if the key does not exist
             return None

diff --git a/icechunk-python/python/icechunk/_icechunk_python.pyi b/icechunk-python/python/icechunk/_icechunk_python.pyi
@@ -249,12 +249,6 @@ class VirtualRefConfig:
         """
         ...
 
-class KeyNotFound(Exception):
-    def __init__(
-        self,
-        info: Any
-    ): ...
-
 class StoreConfig:
     """Configuration for an IcechunkStore"""
 

diff --git a/icechunk-python/src/errors.rs b/icechunk-python/src/errors.rs
@@ -2,7 +2,7 @@ use icechunk::{
     format::IcechunkFormatError, repository::RepositoryError, zarr::StoreError,
 };
 use pyo3::{
-    exceptions::{PyException, PyValueError},
+    exceptions::{PyKeyError, PyValueError},
     PyErr,
 };
 use thiserror::Error;
@@ -16,17 +16,17 @@ use thiserror::Error;
 #[derive(Debug, Error)]
 #[allow(dead_code)]
 pub(crate) enum PyIcechunkStoreError {
-    #[error("key not found error: {0}")]
-    KeyNotFound(String),
     #[error("store error: {0}")]
     StoreError(#[from] StoreError),
-    #[error("repository Error: {0}")]
+    #[error("repository error: {0}")]
     RepositoryError(#[from] RepositoryError),
     #[error("icechunk format error: {0}")]
     IcechunkFormatError(#[from] IcechunkFormatError),
-    #[error("value error: {0}")]
+    #[error("{0}")]
+    PyKeyError(String),
+    #[error("{0}")]
     PyValueError(String),
-    #[error("error: {0}")]
+    #[error("{0}")]
     PyError(#[from] PyErr),
     #[error("{0}")]
     UnkownError(String),
@@ -35,20 +35,12 @@ pub(crate) enum PyIcechunkStoreError {
 impl From<PyIcechunkStoreError> for PyErr {
     fn from(error: PyIcechunkStoreError) -> Self {
         match error {
-            PyIcechunkStoreError::KeyNotFound(_) => {
-                KeyNotFound::new_err(error.to_string())
-            }
+            PyIcechunkStoreError::PyKeyError(e) => PyKeyError::new_err(e),
+            PyIcechunkStoreError::PyValueError(e) => PyValueError::new_err(e),
             PyIcechunkStoreError::PyError(err) => err,
             _ => PyValueError::new_err(error.to_string()),
         }
     }
 }
 
 pub(crate) type PyIcechunkStoreResult<T> = Result<T, PyIcechunkStoreError>;
-
-pyo3::create_exception!(
-    _icechunk_python,
-    KeyNotFound,
-    PyException,
-    "The key is not present in the repository"
-);
diff --git a/icechunk-python/src/lib.rs b/icechunk-python/src/lib.rs
@@ -21,7 +21,7 @@ use icechunk::{
     Repository, SnapshotMetadata,
 };
 use pyo3::{
-    exceptions::PyValueError,
+    exceptions::{PyKeyError, PyValueError},
     prelude::*,
     types::{PyBytes, PyList, PyNone, PyString},
 };
@@ -32,8 +32,6 @@ use tokio::{
     sync::{Mutex, RwLock},
 };
 
-pub use errors::KeyNotFound;
-
 #[pyclass]
 struct PyIcechunkStore {
     consolidated: ConsolidatedStore,
@@ -674,7 +672,7 @@ impl PyIcechunkStore {
             let byte_range = byte_range.unwrap_or((None, None)).into();
             let data = store.read().await.get(&key, &byte_range).await;
             // We need to distinguish the "safe" case of trying to fetch an uninitialized key
-            // from other types of errors, we use KeyNotFound exception for that
+            // from other types of errors, we use PyKeyError exception for that
             match data {
                 Ok(data) => {
                     let pybytes = Python::with_gil(|py| {
@@ -683,7 +681,7 @@ impl PyIcechunkStore {
                     });
                     Ok(pybytes)
                 }
-                Err(StoreError::NotFound(_)) => Err(KeyNotFound::new_err(key)),
+                Err(StoreError::NotFound(_)) => Err(PyKeyError::new_err(key)),
                 Err(err) => Err(PyIcechunkStoreError::StoreError(err).into()),
             }
         })
@@ -1045,9 +1043,8 @@ async fn do_set_virtual_ref(
 
 /// The icechunk Python module implemented in Rust.
 #[pymodule]
-fn _icechunk_python(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+fn _icechunk_python(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add("__version__", env!("CARGO_PKG_VERSION"))?;
-    m.add("KeyNotFound", py.get_type_bound::<KeyNotFound>())?;
     m.add_class::<PyStorageConfig>()?;
     m.add_class::<PyIcechunkStore>()?;
     m.add_class::<PyS3Credentials>()?;

diff --git a/icechunk-python/tests/test_timetravel.py b/icechunk-python/tests/test_timetravel.py
@@ -1,5 +1,7 @@
 import icechunk
 import zarr
+import zarr.core
+import zarr.core.buffer
 
 
 def test_timetravel():
@@ -83,3 +85,6 @@ async def test_branch_reset():
     assert "a/zarr.json" in keys
     assert "b/zarr.json" not in keys
 
+    assert (
+        await store.get("b/zarr.json", zarr.core.buffer.default_buffer_prototype())
+    ) is None