Support for non-multipart put uploads (#20)

* Support for non-multipart put uploads

* rename put_file to put

* Add put test

Author: kylebarron

README.md

@@ -52,7 +52,7 @@ import object_store_rs as obs
 
 store = obs.store.MemoryStore()
 
-obs.put_file(store, "file.txt", b"hello world!")
+obs.put(store, "file.txt", b"hello world!")
 response = obs.get(store, "file.txt")
 response.meta
 # {'size': 12,
@@ -77,7 +77,7 @@ import object_store_rs as obs
 
 store = obs.store.MemoryStore()
 
-await obs.put_file_async(store, "file.txt", b"hello world!")
+await obs.put_async(store, "file.txt", b"hello world!")
 response = await obs.get_async(store, "file.txt")
 response.meta
 # {

docs/api/put.md

@@ -1,4 +1,4 @@
 # Put
 
-::: object_store_rs.put_file
-::: object_store_rs.put_file_async
+::: object_store_rs.put
+::: object_store_rs.put_async

object-store-rs/python/object_store_rs/_object_store_rs.pyi

@@ -1,7 +1,5 @@
 from ._copy import copy as copy
 from ._copy import copy_async as copy_async
-from ._copy import copy_if_not_exists as copy_if_not_exists
-from ._copy import copy_if_not_exists_async as copy_if_not_exists_async
 from ._delete import delete as delete
 from ._delete import delete_async as delete_async
 from ._get import GetOptions as GetOptions
@@ -20,12 +18,10 @@ from ._list import list as list
 from ._list import list_async as list_async
 from ._list import list_with_delimiter as list_with_delimiter
 from ._list import list_with_delimiter_async as list_with_delimiter_async
-from ._put import put_file as put_file
-from ._put import put_file_async as put_file_async
+from ._put import put as put
+from ._put import put_async as put_async
 from ._rename import rename as rename
 from ._rename import rename_async as rename_async
-from ._rename import rename_if_not_exists as rename_if_not_exists
-from ._rename import rename_if_not_exists_async as rename_if_not_exists_async
 from ._sign import HTTP_METHOD as HTTP_METHOD
 from ._sign import SignCapableStore as SignCapableStore
 from ._sign import sign_url as sign_url

object-store-rs/python/object_store_rs/_put.pyi

@@ -1,42 +1,60 @@
 from pathlib import Path
-from typing import IO
+from typing import IO, TypedDict
 
 from .store import ObjectStore
 
-def put_file(
+class PutResult(TypedDict):
+    """
+    Result for a put request.
+    """
+
+    e_tag: str | None
+    """
+    The unique identifier for the newly created object
+
+    <https://datatracker.ietf.org/doc/html/rfc9110#name-etag>
+    """
+
+    version: str | None
+    """A version indicator for the newly created object."""
+
+def put(
     store: ObjectStore,
     location: str,
     file: IO[bytes] | Path | bytes,
     *,
-    chunk_size: int = 5 * 1024,
+    use_multipart: bool | None = None,
+    chunk_size: int = 5 * 1024 * 1024,
     max_concurrency: int = 12,
-) -> None:
+) -> PutResult:
     """Save the provided bytes to the specified location
 
     The operation is guaranteed to be atomic, it will either successfully write the
     entirety of `file` to `location`, or fail. No clients should be able to observe a
     partially written object.
 
-    This will use a multipart upload under the hood.
-
     Args:
         store: The ObjectStore instance to use.
         location: The path within ObjectStore for where to save the file.
         file: The object to upload. Can either be file-like, a `Path` to a local file,
             or a `bytes` object.
+
+    Keyword args:
+        use_multipart: Whether to use a multipart upload under the hood. Defaults using a multipart upload if the length of the file is greater than `chunk_size`.
         chunk_size: The size of chunks to use within each part of the multipart upload. Defaults to 5 MB.
         max_concurrency: The maximum number of chunks to upload concurrently. Defaults to 12.
     """
 
-async def put_file_async(
+async def put_async(
     store: ObjectStore,
     location: str,
     file: IO[bytes] | Path | bytes,
     *,
-    chunk_size: int = 5 * 1024,
+    use_multipart: bool | None = None,
+    chunk_size: int = 5 * 1024 * 1024,
     max_concurrency: int = 12,
-) -> None:
-    """Call `put_file` asynchronously.
+) -> PutResult:
+    """Call `put` asynchronously.
 
-    Refer to the documentation for [put_file][object_store_rs.put_file].
+    Refer to the documentation for [put][object_store_rs.put].
     """

object-store-rs/src/lib.rs

@@ -40,8 +40,8 @@ fn _object_store_rs(py: Python, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(list::list_with_delimiter_async))?;
     m.add_wrapped(wrap_pyfunction!(list::list_with_delimiter))?;
     m.add_wrapped(wrap_pyfunction!(list::list))?;
-    m.add_wrapped(wrap_pyfunction!(put::put_file_async))?;
-    m.add_wrapped(wrap_pyfunction!(put::put_file))?;
+    m.add_wrapped(wrap_pyfunction!(put::put_async))?;
+    m.add_wrapped(wrap_pyfunction!(put::put))?;
     m.add_wrapped(wrap_pyfunction!(rename::rename_async))?;
     m.add_wrapped(wrap_pyfunction!(rename::rename))?;
     m.add_wrapped(wrap_pyfunction!(signer::sign_url_async))?;

object-store-rs/src/list.rs

@@ -21,7 +21,7 @@ impl PyObjectMeta {
 
 impl IntoPy<PyObject> for PyObjectMeta {
     fn into_py(self, py: Python<'_>) -> PyObject {
-        let mut dict = HashMap::new();
+        let mut dict = HashMap::with_capacity(5);
         dict.insert("location", self.0.location.as_ref().into_py(py));
         dict.insert("last_modified", self.0.last_modified.into_py(py));
         dict.insert("size", self.0.size.into_py(py));
@@ -35,7 +35,7 @@ pub(crate) struct PyListResult(ListResult);
 
 impl IntoPy<PyObject> for PyListResult {
     fn into_py(self, py: Python<'_>) -> PyObject {
-        let mut dict = HashMap::new();
+        let mut dict = HashMap::with_capacity(2);
         dict.insert(
             "common_prefixes",
             self.0

object-store-rs/src/put.rs

@@ -1,11 +1,11 @@
+use std::collections::HashMap;
 use std::fs::File;
-use std::io::{BufReader, Cursor, Read};
+use std::io::{BufReader, Cursor, Read, Seek, SeekFrom};
 use std::path::PathBuf;
 use std::sync::Arc;
 
 use object_store::path::Path;
-use object_store::{ObjectStore, WriteMultipart};
-use pyo3::exceptions::PyIOError;
+use object_store::{ObjectStore, PutPayload, PutResult, WriteMultipart};
 use pyo3::prelude::*;
 use pyo3::pybacked::PyBackedBytes;
 use pyo3_file::PyFileLikeObject;
@@ -22,6 +22,21 @@ pub(crate) enum MultipartPutInput {
     Buffer(Cursor<PyBackedBytes>),
 }
 
+impl MultipartPutInput {
+    /// Number of bytes in the file-like object
+    fn nbytes(&mut self) -> PyObjectStoreResult<usize> {
+        let origin_pos = self.stream_position()?;
+        let size = self.seek(SeekFrom::End(0))?;
+        self.seek(SeekFrom::Start(origin_pos))?;
+        Ok(size.try_into().unwrap())
+    }
+
+    /// Whether to use multipart uploads.
+    fn use_multipart(&mut self, chunk_size: usize) -> PyObjectStoreResult<bool> {
+        Ok(self.nbytes()? > chunk_size)
+    }
+}
+
 impl<'py> FromPyObject<'py> for MultipartPutInput {
     fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult<Self> {
         let py = ob.py();
@@ -51,69 +66,120 @@ impl Read for MultipartPutInput {
     }
 }
 
+impl Seek for MultipartPutInput {
+    fn seek(&mut self, pos: std::io::SeekFrom) -> std::io::Result<u64> {
+        match self {
+            Self::File(f) => f.seek(pos),
+            Self::FileLike(f) => f.seek(pos),
+            Self::Buffer(f) => f.seek(pos),
+        }
+    }
+}
+
+pub(crate) struct PyPutResult(PutResult);
+
+impl IntoPy<PyObject> for PyPutResult {
+    fn into_py(self, py: Python<'_>) -> PyObject {
+        let mut dict = HashMap::with_capacity(2);
+        dict.insert("e_tag", self.0.e_tag.into_py(py));
+        dict.insert("version", self.0.version.into_py(py));
+        dict.into_py(py)
+    }
+}
+
 #[pyfunction]
-#[pyo3(signature = (store, location, file, *, chunk_size = 5120, max_concurrency = 12))]
-pub(crate) fn put_file(
+#[pyo3(signature = (store, location, file, *, use_multipart = None, chunk_size = 5242880, max_concurrency = 12))]
+pub(crate) fn put(
     py: Python,
     store: PyObjectStore,
     location: String,
-    file: MultipartPutInput,
+    mut file: MultipartPutInput,
+    use_multipart: Option<bool>,
     chunk_size: usize,
     max_concurrency: usize,
-) -> PyObjectStoreResult<()> {
+) -> PyObjectStoreResult<PyPutResult> {
+    let use_multipart = if let Some(use_multipart) = use_multipart {
+        use_multipart
+    } else {
+        file.use_multipart(chunk_size)?
+    };
     let runtime = get_runtime(py)?;
-    runtime.block_on(put_multipart_inner(
-        store.into_inner(),
-        &location.into(),
-        file,
-        chunk_size,
-        max_concurrency,
-    ))
+    if use_multipart {
+        runtime.block_on(put_multipart_inner(
+            store.into_inner(),
+            &location.into(),
+            file,
+            chunk_size,
+            max_concurrency,
+        ))
+    } else {
+        runtime.block_on(put_inner(store.into_inner(), &location.into(), file))
+    }
 }
 
 #[pyfunction]
-#[pyo3(signature = (store, location, file, *, chunk_size = 5120, max_concurrency = 12))]
-pub(crate) fn put_file_async(
+#[pyo3(signature = (store, location, file, *, use_multipart = None, chunk_size = 5242880, max_concurrency = 12))]
+pub(crate) fn put_async(
     py: Python,
     store: PyObjectStore,
     location: String,
-    file: MultipartPutInput,
+    mut file: MultipartPutInput,
+    use_multipart: Option<bool>,
     chunk_size: usize,
     max_concurrency: usize,
 ) -> PyResult<Bound<PyAny>> {
+    let use_multipart = if let Some(use_multipart) = use_multipart {
+        use_multipart
+    } else {
+        file.use_multipart(chunk_size)?
+    };
     pyo3_async_runtimes::tokio::future_into_py(py, async move {
-        Ok(put_multipart_inner(
-            store.into_inner(),
-            &location.into(),
-            file,
-            chunk_size,
-            max_concurrency,
-        )
-        .await?)
+        let result = if use_multipart {
+            put_multipart_inner(
+                store.into_inner(),
+                &location.into(),
+                file,
+                chunk_size,
+                max_concurrency,
+            )
+            .await?
+        } else {
+            put_inner(store.into_inner(), &location.into(), file).await?
+        };
+        Ok(result)
     })
 }
 
+async fn put_inner(
+    store: Arc<dyn ObjectStore>,
+    location: &Path,
+    mut reader: MultipartPutInput,
+) -> PyObjectStoreResult<PyPutResult> {
+    let nbytes = reader.nbytes()?;
+    let mut buffer = Vec::with_capacity(nbytes);
+    reader.read_to_end(&mut buffer)?;
+    let payload = PutPayload::from_bytes(buffer.into());
+    Ok(PyPutResult(store.put(location, payload).await?))
+}
+
 async fn put_multipart_inner<R: Read>(
     store: Arc<dyn ObjectStore>,
     location: &Path,
     mut reader: R,
     chunk_size: usize,
     max_concurrency: usize,
-) -> PyObjectStoreResult<()> {
+) -> PyObjectStoreResult<PyPutResult> {
     let upload = store.put_multipart(location).await?;
     let mut write = WriteMultipart::new(upload);
     let mut scratch_buffer = vec![0; chunk_size];
     loop {
-        let read_size = reader
-            .read(&mut scratch_buffer)
-            .map_err(|err| PyIOError::new_err(err.to_string()))?;
+        let read_size = reader.read(&mut scratch_buffer)?;
         if read_size == 0 {
             break;
         } else {
             write.wait_for_capacity(max_concurrency).await?;
             write.write(&scratch_buffer[0..read_size]);
         }
     }
-    write.finish().await?;
-    Ok(())
+    Ok(PyPutResult(write.finish().await?))
 }

pyo3-object_store/src/error.rs

@@ -1,7 +1,9 @@
 //! Contains the [`PyObjectStoreError`], the Error returned by most fallible functions in this
 //! crate.
 
-use pyo3::exceptions::{PyException, PyValueError};
+use pyo3::exceptions::{
+    PyException, PyFileNotFoundError, PyIOError, PyNotImplementedError, PyValueError,
+};
 use pyo3::prelude::*;
 use pyo3::DowncastError;
 use thiserror::Error;
@@ -17,13 +19,26 @@ pub enum PyObjectStoreError {
     /// A wrapped [PyErr]
     #[error(transparent)]
     PyErr(#[from] PyErr),
+
+    /// A wrapped [std::io::Error]
+    #[error(transparent)]
+    IOError(#[from] std::io::Error),
 }
 
 impl From<PyObjectStoreError> for PyErr {
     fn from(error: PyObjectStoreError) -> Self {
         match error {
             PyObjectStoreError::PyErr(err) => err,
-            PyObjectStoreError::ObjectStoreError(err) => PyException::new_err(err.to_string()),
+            PyObjectStoreError::ObjectStoreError(ref err) => match err {
+                object_store::Error::NotFound { path: _, source: _ } => {
+                    PyFileNotFoundError::new_err(err.to_string())
+                }
+                object_store::Error::NotImplemented => {
+                    PyNotImplementedError::new_err(err.to_string())
+                }
+                _ => PyException::new_err(err.to_string()),
+            },
+            PyObjectStoreError::IOError(err) => PyIOError::new_err(err.to_string()),
         }
     }
 }