diff --git a/object-store-rs/python/object_store_rs/_object_store_rs.pyi b/object-store-rs/python/object_store_rs/_object_store_rs.pyi index a14c41e..52d0fa5 100644 --- a/object-store-rs/python/object_store_rs/_object_store_rs.pyi +++ b/object-store-rs/python/object_store_rs/_object_store_rs.pyi @@ -24,5 +24,5 @@ from ._rename import rename as rename from ._rename import rename_async as rename_async from ._sign import HTTP_METHOD as HTTP_METHOD from ._sign import SignCapableStore as SignCapableStore -from ._sign import sign_url as sign_url -from ._sign import sign_url_async as sign_url_async +from ._sign import sign as sign +from ._sign import sign_async as sign_async diff --git a/object-store-rs/python/object_store_rs/_sign.pyi b/object-store-rs/python/object_store_rs/_sign.pyi index eba353b..9e6dc9b 100644 --- a/object-store-rs/python/object_store_rs/_sign.pyi +++ b/object-store-rs/python/object_store_rs/_sign.pyi @@ -1,5 +1,5 @@ from datetime import timedelta -from typing import Literal +from typing import List, Literal, Sequence, overload from .store import AzureStore, GCSStore, S3Store @@ -11,13 +11,27 @@ HTTP_METHOD = Literal[ SignCapableStore = AzureStore | GCSStore | S3Store """ObjectStore instances that are capable of signing.""" -def sign_url( - store: SignCapableStore, method: HTTP_METHOD, path: str, expires_in: timedelta -) -> str: +@overload +def sign( # type: ignore + store: SignCapableStore, method: HTTP_METHOD, paths: str, expires_in: timedelta +) -> str: ... +@overload +def sign( + store: SignCapableStore, + method: HTTP_METHOD, + paths: Sequence[str], + expires_in: timedelta, +) -> List[str]: ... +def sign( + store: SignCapableStore, + method: HTTP_METHOD, + paths: str | Sequence[str], + expires_in: timedelta, +) -> str | List[str]: """Create a signed URL. - Given the intended [`Method`] and [`Path`] to use and the desired length of time for - which the URL should be valid, return a signed [`Url`] created with the object store + Given the intended `method` and `paths` to use and the desired length of time for + which the URL should be valid, return a signed URL created with the object store implementation's credentials such that the URL can be handed to something that doesn't have access to the object store's credentials, to allow limited access to the object store. @@ -25,17 +39,34 @@ def sign_url( Args: store: The ObjectStore instance to use. method: The HTTP method to use. - path: The path within ObjectStore to retrieve. - expires_in: How long the signed URL should be valid. + paths: The path(s) within ObjectStore to retrieve. If + expires_in: How long the signed URL(s) should be valid. Returns: _description_ """ -async def sign_url_async( - store: SignCapableStore, method: HTTP_METHOD, path: str, expires_in: timedelta -) -> str: - """Call `sign_url` asynchronously. +@overload +async def sign_async( + store: SignCapableStore, + method: HTTP_METHOD, + paths: str, + expires_in: timedelta, +) -> str: ... +@overload +async def sign_async( + store: SignCapableStore, + method: HTTP_METHOD, + paths: Sequence[str], + expires_in: timedelta, +) -> List[str]: ... +async def sign_async( + store: SignCapableStore, + method: HTTP_METHOD, + paths: str | Sequence[str], + expires_in: timedelta, +) -> str | List[str]: + """Call `sign` asynchronously. - Refer to the documentation for [sign_url][object_store_rs.sign_url]. + Refer to the documentation for [sign][object_store_rs.sign]. """ diff --git a/object-store-rs/src/delete.rs b/object-store-rs/src/delete.rs index f770325..9b73a5a 100644 --- a/object-store-rs/src/delete.rs +++ b/object-store-rs/src/delete.rs @@ -1,48 +1,25 @@ use futures::{StreamExt, TryStreamExt}; -use object_store::path::Path; -use pyo3::exceptions::PyTypeError; use pyo3::prelude::*; use pyo3_object_store::error::{PyObjectStoreError, PyObjectStoreResult}; use pyo3_object_store::PyObjectStore; +use crate::path::PyPaths; use crate::runtime::get_runtime; -pub(crate) enum PyLocations { - One(Path), - // TODO: also support an Arrow String Array here. - Many(Vec), -} - -impl<'py> FromPyObject<'py> for PyLocations { - fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { - if let Ok(path) = ob.extract::() { - Ok(Self::One(path.into())) - } else if let Ok(paths) = ob.extract::>() { - Ok(Self::Many( - paths.into_iter().map(|path| path.into()).collect(), - )) - } else { - Err(PyTypeError::new_err( - "Expected string path or sequence of string paths.", - )) - } - } -} - #[pyfunction] pub(crate) fn delete( py: Python, store: PyObjectStore, - locations: PyLocations, + locations: PyPaths, ) -> PyObjectStoreResult<()> { let runtime = get_runtime(py)?; let store = store.into_inner(); py.allow_threads(|| { match locations { - PyLocations::One(path) => { + PyPaths::One(path) => { runtime.block_on(store.delete(&path))?; } - PyLocations::Many(paths) => { + PyPaths::Many(paths) => { // TODO: add option to allow some errors here? let stream = store.delete_stream(futures::stream::iter(paths.into_iter().map(Ok)).boxed()); @@ -57,18 +34,18 @@ pub(crate) fn delete( pub(crate) fn delete_async( py: Python, store: PyObjectStore, - locations: PyLocations, + locations: PyPaths, ) -> PyResult> { let store = store.into_inner(); pyo3_async_runtimes::tokio::future_into_py(py, async move { match locations { - PyLocations::One(path) => { + PyPaths::One(path) => { store .delete(&path) .await .map_err(PyObjectStoreError::ObjectStoreError)?; } - PyLocations::Many(paths) => { + PyPaths::Many(paths) => { // TODO: add option to allow some errors here? let stream = store.delete_stream(futures::stream::iter(paths.into_iter().map(Ok)).boxed()); diff --git a/object-store-rs/src/lib.rs b/object-store-rs/src/lib.rs index f7ab43b..cd75d5b 100644 --- a/object-store-rs/src/lib.rs +++ b/object-store-rs/src/lib.rs @@ -5,6 +5,7 @@ mod delete; mod get; mod head; mod list; +mod path; mod put; mod rename; mod runtime; @@ -44,8 +45,8 @@ fn _object_store_rs(py: Python, m: &Bound) -> PyResult<()> { 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))?; - m.add_wrapped(wrap_pyfunction!(signer::sign_url))?; + m.add_wrapped(wrap_pyfunction!(signer::sign_async))?; + m.add_wrapped(wrap_pyfunction!(signer::sign))?; Ok(()) } diff --git a/object-store-rs/src/path.rs b/object-store-rs/src/path.rs new file mode 100644 index 0000000..1641c73 --- /dev/null +++ b/object-store-rs/src/path.rs @@ -0,0 +1,25 @@ +use object_store::path::Path; +use pyo3::exceptions::PyTypeError; +use pyo3::prelude::*; + +pub(crate) enum PyPaths { + One(Path), + // TODO: also support an Arrow String Array here. + Many(Vec), +} + +impl<'py> FromPyObject<'py> for PyPaths { + fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { + if let Ok(path) = ob.extract::() { + Ok(Self::One(path.into())) + } else if let Ok(paths) = ob.extract::>() { + Ok(Self::Many( + paths.into_iter().map(|path| path.into()).collect(), + )) + } else { + Err(PyTypeError::new_err( + "Expected string path or sequence of string paths.", + )) + } + } +} diff --git a/object-store-rs/src/signer.rs b/object-store-rs/src/signer.rs index 339eb3d..42b5a34 100644 --- a/object-store-rs/src/signer.rs +++ b/object-store-rs/src/signer.rs @@ -1,9 +1,13 @@ use core::time::Duration; +use std::future::Future; +use std::pin::Pin; use std::sync::Arc; +use http::Method; use object_store::aws::AmazonS3; use object_store::azure::MicrosoftAzure; use object_store::gcp::GoogleCloudStorage; +use object_store::path::Path; use object_store::signer::Signer; use pyo3::exceptions::PyValueError; use pyo3::intern; @@ -13,6 +17,7 @@ use pyo3_object_store::error::{PyObjectStoreError, PyObjectStoreResult}; use pyo3_object_store::{PyAzureStore, PyGCSStore, PyS3Store}; use url::Url; +use crate::path::PyPaths; use crate::runtime::get_runtime; #[derive(Debug)] @@ -61,16 +66,10 @@ impl<'py> FromPyObject<'py> for SignCapableStore { impl Signer for SignCapableStore { fn signed_url<'life0, 'life1, 'async_trait>( &'life0 self, - method: http::Method, - path: &'life1 object_store::path::Path, + method: Method, + path: &'life1 Path, expires_in: Duration, - ) -> ::core::pin::Pin< - Box< - dyn ::core::future::Future> - + ::core::marker::Send - + 'async_trait, - >, - > + ) -> Pin> + Send + 'async_trait>> where 'life0: 'async_trait, 'life1: 'async_trait, @@ -82,23 +81,41 @@ impl Signer for SignCapableStore { Self::Azure(inner) => inner.signed_url(method, path, expires_in), } } + + fn signed_urls<'life0, 'life1, 'async_trait>( + &'life0 self, + method: Method, + paths: &'life1 [Path], + expires_in: Duration, + ) -> Pin>> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + match self { + Self::S3(inner) => inner.signed_urls(method, paths, expires_in), + Self::Gcs(inner) => inner.signed_urls(method, paths, expires_in), + Self::Azure(inner) => inner.signed_urls(method, paths, expires_in), + } + } } -pub(crate) struct PyMethod(http::Method); +pub(crate) struct PyMethod(Method); impl<'py> FromPyObject<'py> for PyMethod { fn extract_bound(ob: &Bound<'py, PyAny>) -> PyResult { let s = ob.extract::()?; let method = match s.as_ref() { - "GET" => http::Method::GET, - "PUT" => http::Method::PUT, - "POST" => http::Method::POST, - "HEAD" => http::Method::HEAD, - "PATCH" => http::Method::PATCH, - "TRACE" => http::Method::TRACE, - "DELETE" => http::Method::DELETE, - "OPTIONS" => http::Method::OPTIONS, - "CONNECT" => http::Method::CONNECT, + "GET" => Method::GET, + "PUT" => Method::PUT, + "POST" => Method::POST, + "HEAD" => Method::HEAD, + "PATCH" => Method::PATCH, + "TRACE" => Method::TRACE, + "DELETE" => Method::DELETE, + "OPTIONS" => Method::OPTIONS, + "CONNECT" => Method::CONNECT, other => { return Err(PyValueError::new_err(format!( "Unsupported HTTP method {}", @@ -112,50 +129,86 @@ impl<'py> FromPyObject<'py> for PyMethod { pub(crate) struct PyUrl(url::Url); -impl IntoPy for PyUrl { - fn into_py(self, _py: Python<'_>) -> String { - self.0.into() +impl IntoPy for PyUrl { + fn into_py(self, py: Python<'_>) -> PyObject { + String::from(self.0).into_py(py) } } -impl IntoPy for PyUrl { +pub(crate) struct PyUrls(Vec); + +impl IntoPy for PyUrls { fn into_py(self, py: Python<'_>) -> PyObject { - String::from(self.0).into_py(py) + self.0.into_py(py) + } +} + +pub(crate) enum PySignResult { + One(PyUrl), + Many(PyUrls), +} + +impl IntoPy for PySignResult { + fn into_py(self, py: Python<'_>) -> PyObject { + match self { + Self::One(url) => url.into_py(py), + Self::Many(urls) => urls.into_py(py), + } } } #[pyfunction] -pub(crate) fn sign_url( +pub(crate) fn sign( py: Python, store: SignCapableStore, method: PyMethod, - path: String, + paths: PyPaths, expires_in: Duration, -) -> PyObjectStoreResult { +) -> PyObjectStoreResult { let runtime = get_runtime(py)?; let method = method.0; - let signed_url = py.allow_threads(|| { - let url = runtime.block_on(store.signed_url(method, &path.into(), expires_in))?; - Ok::<_, object_store::Error>(url) - })?; - Ok(signed_url.into()) + py.allow_threads(|| match paths { + PyPaths::One(path) => { + let url = runtime.block_on(store.signed_url(method, &path, expires_in))?; + Ok(PySignResult::One(PyUrl(url))) + } + PyPaths::Many(paths) => { + let urls = runtime.block_on(store.signed_urls(method, &paths, expires_in))?; + Ok(PySignResult::Many(PyUrls( + urls.into_iter().map(PyUrl).collect(), + ))) + } + }) } #[pyfunction] -pub(crate) fn sign_url_async( +pub(crate) fn sign_async( py: Python, store: SignCapableStore, method: PyMethod, - path: String, + paths: PyPaths, expires_in: Duration, -) -> PyResult { - let fut = pyo3_async_runtimes::tokio::future_into_py(py, async move { - let url = store - .signed_url(method.0, &path.into(), expires_in) - .await - .map_err(PyObjectStoreError::ObjectStoreError)?; - Ok(PyUrl(url)) - })?; - Ok(fut.into()) +) -> PyResult> { + let method = method.0; + pyo3_async_runtimes::tokio::future_into_py(py, async move { + match paths { + PyPaths::One(path) => { + let url = store + .signed_url(method, &path, expires_in) + .await + .map_err(PyObjectStoreError::ObjectStoreError)?; + Ok(PySignResult::One(PyUrl(url))) + } + PyPaths::Many(paths) => { + let urls = store + .signed_urls(method, &paths, expires_in) + .await + .map_err(PyObjectStoreError::ObjectStoreError)?; + Ok(PySignResult::Many(PyUrls( + urls.into_iter().map(PyUrl).collect(), + ))) + } + } + }) }