diff --git a/Cargo.lock b/Cargo.lock index e79211b5436..b707f99260a 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4950,10 +4950,12 @@ dependencies = [ "defmt", "jiff-static", "jiff-tzdb-platform", + "js-sys", "log", "portable-atomic", "portable-atomic-util", "serde_core", + "wasm-bindgen", "windows-link", ] @@ -5973,6 +5975,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "mea" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2640d335e7273dacdcf51044026139b2e269c3bb0dfc3f8cb3496b85e3f6a42c" +dependencies = [ + "slab", +] + [[package]] name = "memchr" version = "2.8.2" @@ -6428,6 +6439,23 @@ dependencies = [ "web-time", ] +[[package]] +name = "object_store_opendal" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eb12a624a41fce745838d0ef3701ff6c47797c13cd18ad3612fd2a3134fdbd8" +dependencies = [ + "async-trait", + "bytes", + "chrono", + "futures", + "mea", + "object_store", + "opendal", + "pin-project", + "tokio", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -6465,6 +6493,61 @@ version = "11.1.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" +[[package]] +name = "opendal" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c9c85ce253ff87225e7669979d877a20c98a06604ec9d6dd5f4473e08f1ae1" +dependencies = [ + "opendal-core", + "opendal-service-cos", +] + +[[package]] +name = "opendal-core" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4f8607c90e2c963a91467f50fb49fbc7fb3d573f88cea219ca59ccd3740b309" +dependencies = [ + "anyhow", + "base64", + "bytes", + "futures", + "http", + "http-body", + "jiff", + "log", + "md-5 0.11.0", + "mea", + "percent-encoding", + "quick-xml", + "reqsign-core", + "reqwest 0.13.4", + "serde", + "serde_json", + "tokio", + "url", + "uuid", + "web-time", +] + +[[package]] +name = "opendal-service-cos" +version = "0.57.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa8cafe9729213375c7331019b0cb756ad3e1aff7f45cd32c45eae91ebde8901" +dependencies = [ + "bytes", + "http", + "log", + "opendal-core", + "quick-xml", + "reqsign-core", + "reqsign-file-read-tokio", + "reqsign-tencent-cos", + "serde", +] + [[package]] name = "openssl-probe" version = "0.2.1" @@ -7250,8 +7333,6 @@ dependencies = [ [[package]] name = "pyo3-object_store" version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "13aaa9b876f9b02773cc18fd03dd359bddf3ca21978068024511f166c9c0c6a5" dependencies = [ "async-trait", "bytes", @@ -7740,6 +7821,54 @@ dependencies = [ "bytecheck", ] +[[package]] +name = "reqsign-core" +version = "3.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "514a1e0b4aa288652a3fdbda4f0a610f379cdf5374e55a37c9edd03d57ed856b" +dependencies = [ + "anyhow", + "base64", + "bytes", + "form_urlencoded", + "futures", + "hex", + "hmac", + "http", + "jiff", + "log", + "percent-encoding", + "sha1", + "sha2 0.11.0", + "windows-sys 0.61.2", +] + +[[package]] +name = "reqsign-file-read-tokio" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b472a8d1f2e5a4be8ce13bb7bdf4b59e9bee613ce124aca23959ddb42176b39" +dependencies = [ + "anyhow", + "reqsign-core", + "tokio", +] + +[[package]] +name = "reqsign-tencent-cos" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f5c353c69bdbd6e3deb8b067a5547504cc155b7f024998fef005e37d65d646c" +dependencies = [ + "anyhow", + "http", + "log", + "percent-encoding", + "reqsign-core", + "serde", + "serde_json", +] + [[package]] name = "reqwest" version = "0.12.28" @@ -10339,6 +10468,7 @@ dependencies = [ "url", "vortex", "vortex-arrow", + "vortex-object-store-opendal", "vortex-parquet-variant", ] @@ -10433,6 +10563,17 @@ dependencies = [ "vortex-cuda-macros", ] +[[package]] +name = "vortex-object-store-opendal" +version = "0.1.0" +dependencies = [ + "object_store", + "object_store_opendal", + "opendal", + "url", + "vortex-utils", +] + [[package]] name = "vortex-onpair" version = "0.1.0" @@ -10520,8 +10661,10 @@ dependencies = [ "vortex", "vortex-array", "vortex-arrow", + "vortex-object-store-opendal", "vortex-python-abi", "vortex-tui", + "vortex-utils", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 68be0c50653..74254e95b74 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ members = [ "vortex-flatbuffers", "vortex-metrics", "vortex-io", + "vortex-object-store-opendal", "vortex-proto", "vortex-array", "vortex-arrow", @@ -429,3 +430,12 @@ incremental = false # This improved build times significantly for default common cases that we use locally [profile.dev.package.vortex-fastlanes] debug = false + +# Vendored fork of `pyo3-object_store` that adds `From> for +# PyObjectStore`, allowing any pre-built `object_store::ObjectStore` (e.g. the OpenDAL-backed +# COS/OSS store) to be passed to Vortex's `read_url(store=...)` / `write(store=...)` APIs. +# +# Once https://github.com/developmentseed/obstore (pyo3-object_store) ships this constructor +# upstream, this patch can be removed. +[patch.crates-io] +pyo3-object_store = { path = "vendor/pyo3-object_store" } diff --git a/docs/api/python/store.rst b/docs/api/python/store.rst index 305657c65f1..ce4a6f75b2d 100644 --- a/docs/api/python/store.rst +++ b/docs/api/python/store.rst @@ -13,6 +13,7 @@ Vortex arrays support reading and writing to many object storage systems: store/http store/local store/memory + store/opendal store/config .. autofunction:: vortex.store.from_url diff --git a/docs/api/python/store/opendal.rst b/docs/api/python/store/opendal.rst new file mode 100644 index 00000000000..96f23d51622 --- /dev/null +++ b/docs/api/python/store/opendal.rst @@ -0,0 +1,62 @@ +======== +OpenDAL (COS) +======== + +Vortex can read from and write to Tencent Cloud COS through +`OpenDAL `_, which provides native service support. + +This store is available only when Vortex is built with the ``opendal`` feature +(e.g. ``maturin develop --features opendal`` or ``cargo build -p vortex-jni --features opendal``). + +.. autoclass:: vortex.store.CosStore + :members: + +Reading from COS +================ + +Pass a ``cos://`` URL directly. Credentials and the endpoint are picked up from the environment +variables OpenDAL's COS builder reads (``TENCENTCLOUD_SECRET_ID``, ``TENCENTCLOUD_SECRET_KEY`` and +``COS_ENDPOINT``): + +.. code-block:: python + + import vortex as vx + + a = vx.io.read_url("cos://my-bucket/path/to/dataset.vortex") + +Or configure explicitly with :class:`~vortex.store.CosStore` before reading: + +.. code-block:: python + + from vortex.store import CosStore + + CosStore( + bucket="my-bucket", + endpoint="https://cos.ap-guangzhou.myqcloud.com", + secret_id="AKID...", + secret_key="...", + ).apply() + + a = vx.io.read_url("cos://my-bucket/path/to/dataset.vortex") + +Passing a store object directly +=============================== + +``CosStore`` is a concrete store object. You can build one once and pass it +directly to :func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument, +exactly like the built-in S3/Azure/GCS stores: + +.. code-block:: python + + from vortex.io import read_url + from vortex.store import CosStore + + store = CosStore( + bucket="my-bucket", + endpoint="https://cos.ap-guangzhou.myqcloud.com", + secret_id="AKID...", + secret_key="...", + ) + + a = read_url("cos://my-bucket/path/to/dataset.vortex", store=store) + diff --git a/vendor/pyo3-object_store/Cargo.toml b/vendor/pyo3-object_store/Cargo.toml new file mode 100644 index 00000000000..d27644b1605 --- /dev/null +++ b/vendor/pyo3-object_store/Cargo.toml @@ -0,0 +1,47 @@ +[package] +name = "pyo3-object_store" +version = "0.11.0" +authors = ["Kyle Barron "] +edition = "2021" +description = "object_store integration for pyo3." +readme = "README.md" +repository = "https://github.com/developmentseed/obstore" +license = "MIT OR Apache-2.0" +keywords = [] +categories = [] +rust-version = "1.75" +# Include the Python type hints as part of the cargo distribution +include = ["src", "type-hints", "README.md", "LICENSE"] + +[features] +default = ["external-store-warning"] +external-store-warning = [] + +[dependencies] +async-trait = "0.1.85" +bytes = "1" +chrono = "0.4" +futures = "0.3" +# This is already an object_store dependency +humantime = "2.1" +# This is already an object_store dependency +http = "1" +# This is already an object_store dependency +itertools = "0.14.0" +object_store = { version = "0.13.0", features = [ + "aws", + "azure", + "gcp", + "http", +] } +# This is already an object_store dependency +percent-encoding = "2.1" +pyo3 = { version = "0.29", features = ["chrono", "indexmap"] } +pyo3-async-runtimes = { version = "0.29", features = ["tokio-runtime"] } +serde = "1" +thiserror = "1" +tokio = { version = "1.40", features = ["rt-multi-thread"] } +url = "2" + +[lib] +crate-type = ["rlib"] diff --git a/vendor/pyo3-object_store/LICENSE b/vendor/pyo3-object_store/LICENSE new file mode 100644 index 00000000000..7bca320a280 --- /dev/null +++ b/vendor/pyo3-object_store/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Development Seed + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/vendor/pyo3-object_store/README.md b/vendor/pyo3-object_store/README.md new file mode 100644 index 00000000000..f6177685ea7 --- /dev/null +++ b/vendor/pyo3-object_store/README.md @@ -0,0 +1,79 @@ +# pyo3-object_store + +Integration between [`object_store`](https://docs.rs/object_store) and [`pyo3`](https://github.com/PyO3/pyo3). + +This provides Python builder classes so that Python users can easily create [`Arc`][object_store::ObjectStore] instances, which can then be used in pure-Rust code. + +## Usage + +1. Register the builders. + + ```rs + #[pymodule] + fn python_module(py: Python, m: &Bound) -> PyResult<()> { + pyo3_object_store::register_store_module(py, m, "python_module", "store")?; + pyo3_object_store::register_exceptions_module(py, m, "python_module", "exceptions")?; + } + ``` + + This exports the underlying Python classes from your own Rust-Python library. + + Refer to [`register_store_module`] and [`register_exceptions_module`] for more information. + +2. Accept [`PyObjectStore`] as a parameter in your function exported to Python. Its [`into_dyn`][PyObjectStore::into_dyn] method (or `Into` impl) gives you an [`Arc`][object_store::ObjectStore]. + + ```rs + #[pyfunction] + pub fn use_object_store(store: PyObjectStore) { + let store: Arc = store.into_dyn(); + } + ``` + + You can also accept [`AnyObjectStore`] as a parameter, which wraps [`PyObjectStore`] and [`PyExternalObjectStore`]. This allows you to seamlessly recreate `ObjectStore` instances that users pass in from other Python libraries (like [`obstore`][obstore]) that themselves export `pyo3-object_store` builders. + + Note however that due to lack of [ABI stability](#abi-stability), `ObjectStore` instances will be **recreated**, and so there will be no connection pooling across the external store. + +## Example + +The [`obstore`][obstore] Python library gives a full real-world example of using `pyo3-object_store`, exporting a Python API that mimics the Rust [`ObjectStore`][object_store::ObjectStore] API. + +[obstore]: https://developmentseed.org/obstore/latest/ + +## Feature flags + +- `external-store-warning` (enabled by default): Emit a user warning when constructing a `PyExternalObjectStore` (or `AnyObjectStore::PyExternalObjectStore`), to inform users that there may be performance implications due to lack of connection pooling across separately-compiled Python libraries. Disable this feature if you don't want the warning. + +## ABI stability + +It's [not currently possible](https://github.com/PyO3/pyo3/issues/1444) to share a `#[pyclass]` across multiple Python libraries, except in special cases where the underlying data has a stable ABI. + +As `object_store` does not currently have a stable ABI, we can't share `PyObjectStore` instances across multiple separately-compiled Python libraries. + +We have two ways to get around this: + +- Export your own Python classes so that users can construct `ObjectStore` instances that were compiled _with your library_. See [`register_store_module`]. +- Accept [`AnyObjectStore`] or [`PyExternalObjectStore`] as a parameter, which allows for seamlessly **reconstructing** stores from an external Python library, like [`obstore`][obstore]. This has some overhead and removes any possibility of connection pooling across the two instances. + +Note about not being able to use these across Python packages. It has to be used with the exported classes from your own library. + +## Python Type hints + +We don't yet have a _great_ solution here for reusing the store builder type hints in your own library. Type hints are shipped with the cargo dependency. Or, you can use a submodule on the `obstore` repo. See [`async-tiff` for an example](https://github.com/developmentseed/async-tiff/blob/35eaf116d9b1ab31232a1e23298b3102d2879e9c/python/python/async_tiff/store). + +## Version compatibility + +| pyo3-object_store | pyo3 | object_store | +| ----------------- | ------------------ | ------------------ | +| 0.1.x | 0.23 | 0.12 | +| 0.2.x | 0.24 | 0.12 | +| 0.3.x | **0.23** :warning: | **0.11** :warning: | +| 0.4.x | 0.24 | **0.11** :warning: | +| 0.5.x | 0.25 | 0.12 | +| 0.6.x | 0.26 | 0.12 | +| 0.7.x | 0.27 | 0.12 | +| 0.8.x | 0.27 | 0.13 | +| 0.9.x | 0.28 | 0.13 | +| 0.10.x | 0.28 | **0.12** :warning: | +| 0.11.x | 0.29 | 0.13 | + +Note that 0.3.x, 0.4.x, and 0.10.x are compatibility releases to use `pyo3-object_store` with older versions of `pyo3` and `object_store`. diff --git a/vendor/pyo3-object_store/src/api.rs b/vendor/pyo3-object_store/src/api.rs new file mode 100644 index 00000000000..ad22d32f566 --- /dev/null +++ b/vendor/pyo3-object_store/src/api.rs @@ -0,0 +1,181 @@ +use pyo3::intern; +use pyo3::prelude::*; + +use crate::error::*; +use crate::{ + from_url, PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store, +}; + +/// Export the default Python API as a submodule named `store` within the given parent module +/// +/// The following will add a `store` submodule within a Python top-level module called `"python_module"`. +/// +/// Args: +/// +/// - [`Python`][pyo3::prelude::Python] token +/// - parent_module: [`PyModule`][pyo3::prelude::PyModule] object +/// - parent_module_str: the string name of the parent Python module for how this is exported. +/// - sub_module_str: the string name of **this** Python module. Usually `store` but `_store` may be preferred in some cases. +/// +/// ```notest +/// #[pymodule] +/// fn rust_module(py: Python, m: &Bound) -> PyResult<()> { +/// pyo3_object_store::register_store_module(py, m, "python_module")?; +/// } +/// ``` +/// +/// Or as another example, in the `obstore` Python-facing API, this is exported as +/// +/// ```notest +/// #[pymodule] +/// fn _obstore(py: Python, m: &Bound) -> PyResult<()> { +/// pyo3_object_store::register_store_module(py, m, "obstore")?; +/// } +/// ``` +/// +/// Then `obstore._obstore.*` is re-exported at the top-level at `obstore.*`. So this means the +/// store will be available at `obstore.store`. +/// +// https://github.com/PyO3/pyo3/issues/1517#issuecomment-808664021 +// https://github.com/PyO3/pyo3/issues/759#issuecomment-977835119 +pub fn register_store_module( + py: Python<'_>, + parent_module: &Bound<'_, PyModule>, + parent_module_str: &str, + sub_module_str: &str, +) -> PyResult<()> { + let full_module_string = format!("{parent_module_str}.{sub_module_str}"); + + let child_module = PyModule::new(parent_module.py(), sub_module_str)?; + + child_module.add_wrapped(wrap_pyfunction!(from_url))?; + child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_class::()?; + child_module.add_class::()?; + + // Set the value of `__module__` correctly on each publicly exposed function or class + let __module__ = intern!(py, "__module__"); + child_module + .getattr("from_url")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("AzureStore")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("GCSStore")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("HTTPStore")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("LocalStore")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("MemoryStore")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("S3Store")? + .setattr(__module__, &full_module_string)?; + + // Add the child module to the parent module + parent_module.add_submodule(&child_module)?; + + py.import(intern!(py, "sys"))? + .getattr(intern!(py, "modules"))? + .set_item(&full_module_string, &child_module)?; + + // needs to be set *after* `add_submodule()` + child_module.setattr("__name__", full_module_string)?; + + Ok(()) +} + +/// Export exceptions as a submodule named `exceptions` within the given parent module +// https://github.com/PyO3/pyo3/issues/1517#issuecomment-808664021 +// https://github.com/PyO3/pyo3/issues/759#issuecomment-977835119 +pub fn register_exceptions_module( + py: Python<'_>, + parent_module: &Bound<'_, PyModule>, + parent_module_str: &str, + sub_module_str: &str, +) -> PyResult<()> { + let full_module_string = format!("{parent_module_str}.{sub_module_str}"); + + let child_module = PyModule::new(parent_module.py(), sub_module_str)?; + + child_module.add("BaseError", py.get_type::())?; + child_module.add("GenericError", py.get_type::())?; + child_module.add("NotFoundError", py.get_type::())?; + child_module.add("InvalidPathError", py.get_type::())?; + child_module.add("JoinError", py.get_type::())?; + child_module.add("NotSupportedError", py.get_type::())?; + child_module.add("AlreadyExistsError", py.get_type::())?; + child_module.add("PreconditionError", py.get_type::())?; + child_module.add("NotModifiedError", py.get_type::())?; + child_module.add( + "PermissionDeniedError", + py.get_type::(), + )?; + child_module.add( + "UnauthenticatedError", + py.get_type::(), + )?; + child_module.add( + "UnknownConfigurationKeyError", + py.get_type::(), + )?; + + // Set the value of `__module__` correctly on each publicly exposed function or class + let __module__ = intern!(py, "__module__"); + child_module + .getattr("BaseError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("GenericError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("NotFoundError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("InvalidPathError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("JoinError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("NotSupportedError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("AlreadyExistsError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("PreconditionError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("NotModifiedError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("PermissionDeniedError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("UnauthenticatedError")? + .setattr(__module__, &full_module_string)?; + child_module + .getattr("UnknownConfigurationKeyError")? + .setattr(__module__, &full_module_string)?; + + // Add the child module to the parent module + parent_module.add_submodule(&child_module)?; + + py.import(intern!(py, "sys"))? + .getattr(intern!(py, "modules"))? + .set_item(full_module_string.as_str(), &child_module)?; + + // needs to be set *after* `add_submodule()` + child_module.setattr("__name__", full_module_string)?; + + Ok(()) +} diff --git a/vendor/pyo3-object_store/src/aws/credentials.rs b/vendor/pyo3-object_store/src/aws/credentials.rs new file mode 100644 index 00000000000..11dcc0b4bdd --- /dev/null +++ b/vendor/pyo3-object_store/src/aws/credentials.rs @@ -0,0 +1,216 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use object_store::aws::AwsCredential; +use object_store::CredentialProvider; +use pyo3::exceptions::PyTypeError; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +use crate::aws::store::PyAmazonS3Config; +use crate::credentials::{is_awaitable, TemporaryToken, TokenCache}; + +/// A wrapper around an [AwsCredential] that includes an optional expiry timestamp. +struct PyAwsCredential { + credential: AwsCredential, + expires_at: Option>, +} + +impl<'py> FromPyObject<'_, 'py> for PyAwsCredential { + type Error = PyErr; + + /// Converts from a Python dictionary of the form + /// + /// ```py + /// class S3Credential(TypedDict): + /// access_key_id: str + /// secret_access_key: str + /// token: str | None + /// expires_at: datetime | None + /// ``` + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let py = obj.py(); + let key_id = obj.get_item(intern!(py, "access_key_id"))?.extract()?; + let secret_key = obj.get_item(intern!(py, "secret_access_key"))?.extract()?; + let token = if let Ok(token) = obj.get_item(intern!(py, "token")) { + token.extract()? + } else { + // Allow the dictionary not having a `token` key (so `get_item` will `Err` above) + None + }; + let credential = AwsCredential { + key_id, + secret_key, + token, + }; + let expires_at = obj.get_item(intern!(py, "expires_at"))?.extract()?; + Ok(Self { + credential, + expires_at, + }) + } +} + +// TODO: don't use a cache for static credentials where `expires_at` is `None` +// (so you don't need to access a mutex) +#[derive(Debug)] +pub struct PyAWSCredentialProvider { + /// The provided user callback to manage credential refresh + user_callback: Py, + cache: TokenCache>, + /// An optional config passed down from the credential provider class + config: Option, +} + +impl PyAWSCredentialProvider { + /// Access the S3 config passed down from the credential provider + pub(crate) fn config(&self) -> Option<&PyAmazonS3Config> { + self.config.as_ref() + } + + fn equals(&self, py: Python, other: &Self) -> PyResult { + self.user_callback + .call_method1(py, "__eq__", PyTuple::new(py, vec![&other.user_callback])?)? + .extract(py) + } +} + +impl Clone for PyAWSCredentialProvider { + fn clone(&self) -> Self { + let cloned_callback = Python::attach(|py| self.user_callback.clone_ref(py)); + Self { + user_callback: cloned_callback, + cache: self.cache.clone(), + config: self.config.clone(), + } + } +} + +impl PartialEq for PyAWSCredentialProvider { + fn eq(&self, other: &Self) -> bool { + Python::attach(|py| self.equals(py, other)).unwrap_or(false) + } +} + +impl<'py> FromPyObject<'_, 'py> for PyAWSCredentialProvider { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if !obj.hasattr(intern!(obj.py(), "__call__"))? { + return Err(PyTypeError::new_err( + "Expected callable object for credential_provider.", + )); + } + let mut cache = TokenCache::default(); + if let Ok(refresh_threshold) = obj.getattr(intern!(obj.py(), "refresh_threshold")) { + cache = cache.with_min_ttl(refresh_threshold.extract()?); + } + + let config = if let Ok(config) = obj.getattr(intern!(obj.py(), "config")) { + config.extract()? + } else { + // Allow not having a `config` attribute + None + }; + + Ok(Self { + user_callback: obj.as_unbound().clone_ref(obj.py()), + cache, + config, + }) + } +} + +impl<'py> IntoPyObject<'py> for PyAWSCredentialProvider { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +impl<'py> IntoPyObject<'py> for &PyAWSCredentialProvider { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.user_callback.bind(py).clone()) + } +} + +/// Note: This is copied across providers at the moment +enum PyCredentialProviderResult { + Async(Py), + Sync(PyAwsCredential), +} + +impl PyCredentialProviderResult { + async fn resolve(self) -> PyResult { + match self { + Self::Sync(credentials) => Ok(credentials), + Self::Async(coroutine) => { + let future = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(coroutine.bind(py).clone()) + })?; + let result = future.await?; + Python::attach(|py| result.extract(py)) + } + } + } +} + +impl<'py> FromPyObject<'_, 'py> for PyCredentialProviderResult { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if is_awaitable(&obj)? { + Ok(Self::Async(obj.as_unbound().clone_ref(obj.py()))) + } else { + Ok(Self::Sync(obj.extract()?)) + } + } +} + +impl PyAWSCredentialProvider { + /// Call the user-provided callback and extract to a token. + /// + /// This is separate from `fetch_token` below so that it can return a `PyResult`. + async fn call(&self) -> PyResult { + let call_result = Python::attach(|py| { + self.user_callback + .call0(py)? + .extract::(py) + })?; + call_result.resolve().await + } + + /// Call the user-provided callback + async fn fetch_token(&self) -> object_store::Result>> { + let credential = self + .call() + .await + .map_err(|err| object_store::Error::Unauthenticated { + path: "External AWS credential provider".to_string(), + source: Box::new(err), + })?; + + Ok(TemporaryToken { + token: Arc::new(credential.credential), + expiry: credential.expires_at, + }) + } +} + +#[async_trait] +impl CredentialProvider for PyAWSCredentialProvider { + type Credential = AwsCredential; + + async fn get_credential(&self) -> object_store::Result> { + self.cache.get_or_insert_with(|| self.fetch_token()).await + } +} diff --git a/vendor/pyo3-object_store/src/aws/mod.rs b/vendor/pyo3-object_store/src/aws/mod.rs new file mode 100644 index 00000000000..a4871db182e --- /dev/null +++ b/vendor/pyo3-object_store/src/aws/mod.rs @@ -0,0 +1,4 @@ +mod credentials; +mod store; + +pub use store::PyS3Store; diff --git a/vendor/pyo3-object_store/src/aws/store.rs b/vendor/pyo3-object_store/src/aws/store.rs new file mode 100644 index 00000000000..f8041c4a08d --- /dev/null +++ b/vendor/pyo3-object_store/src/aws/store.rs @@ -0,0 +1,430 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use itertools::Itertools; +use object_store::aws::{AmazonS3, AmazonS3Builder, AmazonS3ConfigKey}; +use object_store::ObjectStoreScheme; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::{PyDict, PyString, PyTuple, PyType}; +use pyo3::{intern, IntoPyObjectExt}; +use url::Url; + +use crate::aws::credentials::PyAWSCredentialProvider; +use crate::client::PyClientOptions; +use crate::config::PyConfigValue; +use crate::error::{GenericError, ParseUrlError, PyObjectStoreError, PyObjectStoreResult}; +use crate::path::PyPath; +use crate::prefix::MaybePrefixedStore; +use crate::retry::PyRetryConfig; +use crate::PyUrl; + +#[derive(Debug, Clone, PartialEq)] +struct S3Config { + prefix: Option, + config: PyAmazonS3Config, + client_options: Option, + retry_config: Option, + credential_provider: Option, +} + +impl S3Config { + fn bucket(&self) -> &str { + self.config + .0 + .get(&PyAmazonS3ConfigKey(AmazonS3ConfigKey::Bucket)) + .expect("bucket should always exist in the config") + .as_ref() + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let args = PyTuple::empty(py).into_bound_py_any(py)?; + let kwargs = PyDict::new(py); + + if let Some(prefix) = &self.prefix { + kwargs.set_item(intern!(py, "prefix"), prefix.as_ref().as_ref())?; + } + kwargs.set_item(intern!(py, "config"), &self.config)?; + if let Some(client_options) = &self.client_options { + kwargs.set_item(intern!(py, "client_options"), client_options)?; + } + if let Some(retry_config) = &self.retry_config { + kwargs.set_item(intern!(py, "retry_config"), retry_config)?; + } + if let Some(credential_provider) = &self.credential_provider { + kwargs.set_item("credential_provider", credential_provider)?; + } + + PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) + } +} + +/// A Python-facing wrapper around an [`AmazonS3`]. +#[derive(Debug, Clone)] +#[pyclass(name = "S3Store", frozen, subclass, from_py_object)] +pub struct PyS3Store { + store: Arc>, + /// A config used for pickling. This must stay in sync with the underlying store's config. + config: S3Config, +} + +impl AsRef>> for PyS3Store { + fn as_ref(&self) -> &Arc> { + &self.store + } +} + +impl PyS3Store { + /// Consume self and return the underlying [`AmazonS3`]. + pub fn into_inner(self) -> Arc> { + self.store + } +} + +#[pymethods] +impl PyS3Store { + // Create from parameters + #[new] + #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + fn new( + bucket: Option, + prefix: Option, + config: Option, + client_options: Option, + retry_config: Option, + credential_provider: Option, + kwargs: Option, + ) -> PyObjectStoreResult { + let mut builder = AmazonS3Builder::from_env(); + let mut config = config.unwrap_or_default(); + + if let Some(bucket) = bucket { + // Note: we apply the bucket to the config, not directly to the builder, so they stay + // in sync. + config.insert_raising_if_exists(AmazonS3ConfigKey::Bucket, bucket)?; + } + + let mut combined_config = combine_config_kwargs(config, kwargs)?; + + if let Some(client_options) = client_options.clone() { + builder = builder.with_client_options(client_options.into()) + } + if let Some(retry_config) = retry_config.clone() { + builder = builder.with_retry(retry_config.into()) + } + + if let Some(credential_provider) = credential_provider.clone() { + // Apply credential provider config onto main config + if let Some(credential_config) = credential_provider.config() { + for (key, val) in credential_config.0.iter() { + // Give precedence to passed-in config values + combined_config.insert_if_not_exists(key.clone(), val.clone()); + } + } + builder = builder.with_credentials(Arc::new(credential_provider)); + } + + builder = combined_config.clone().apply_config(builder); + + Ok(Self { + store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), + config: S3Config { + prefix, + config: combined_config, + client_options, + retry_config, + credential_provider, + }, + }) + } + + #[classmethod] + #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + pub(crate) fn from_url<'py>( + cls: &Bound<'py, PyType>, + url: PyUrl, + config: Option, + client_options: Option, + retry_config: Option, + credential_provider: Option, + kwargs: Option, + ) -> PyObjectStoreResult> { + // We manually parse the URL to find the prefix because `with_url` does not apply the + // prefix. + let (_, prefix) = + ObjectStoreScheme::parse(url.as_ref()).map_err(object_store::Error::from)?; + let prefix: Option = if prefix.parts().count() != 0 { + Some(prefix.into()) + } else { + None + }; + let config = parse_url(config, url.as_ref())?; + + // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the + // subclass + let kwargs = kwargs.unwrap_or_default().into_pyobject(cls.py())?; + kwargs.set_item("prefix", prefix)?; + kwargs.set_item("config", config)?; + kwargs.set_item("client_options", client_options)?; + kwargs.set_item("retry_config", retry_config)?; + kwargs.set_item("credential_provider", credential_provider)?; + Ok(cls.call((), Some(&kwargs))?) + } + + fn __eq__(&self, other: &Bound) -> bool { + // Ensure we never error on __eq__ by returning false if the other object is not an S3Store + other + .cast::() + .map(|other| self.config == other.get().config) + .unwrap_or(false) + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + self.config.__getnewargs_ex__(py) + } + + fn __repr__(&self) -> String { + let bucket = self.config.bucket(); + if let Some(prefix) = &self.config.prefix { + format!( + "S3Store(bucket=\"{}\", prefix=\"{}\")", + bucket, + prefix.as_ref() + ) + } else { + format!("S3Store(bucket=\"{bucket}\")") + } + } + + #[getter] + fn prefix(&self) -> Option<&PyPath> { + self.config.prefix.as_ref() + } + + #[getter] + fn config(&self) -> &PyAmazonS3Config { + &self.config.config + } + + #[getter] + fn client_options(&self) -> Option<&PyClientOptions> { + self.config.client_options.as_ref() + } + + #[getter] + fn credential_provider(&self) -> Option<&PyAWSCredentialProvider> { + self.config.credential_provider.as_ref() + } + + #[getter] + fn retry_config(&self) -> Option<&PyRetryConfig> { + self.config.retry_config.as_ref() + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct PyAmazonS3ConfigKey(AmazonS3ConfigKey); + +impl<'py> FromPyObject<'_, 'py> for PyAmazonS3ConfigKey { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let s = obj.extract::()?.to_lowercase(); + let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?; + Ok(Self(key)) + } +} + +impl AsRef for PyAmazonS3ConfigKey { + fn as_ref(&self) -> &str { + self.0.as_ref() + } +} + +impl<'py> IntoPyObject<'py> for &PyAmazonS3ConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let s = self + .0 + .as_ref() + .strip_prefix("aws_") + .expect("Expected config prefix to start with aws_"); + Ok(PyString::new(py, s)) + } +} + +impl<'py> IntoPyObject<'py> for PyAmazonS3ConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +impl From for PyAmazonS3ConfigKey { + fn from(value: AmazonS3ConfigKey) -> Self { + Self(value) + } +} + +impl From for AmazonS3ConfigKey { + fn from(value: PyAmazonS3ConfigKey) -> Self { + value.0 + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, IntoPyObject, IntoPyObjectRef)] +pub struct PyAmazonS3Config(HashMap); + +// Note: we manually impl FromPyObject instead of deriving it so that we can raise an +// UnknownConfigurationKeyError instead of a `TypeError` on invalid config keys. +// +// We also manually impl this so that we can raise on duplicate keys. +impl<'py> FromPyObject<'_, 'py> for PyAmazonS3Config { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let mut slf = Self::new(); + for (key, val) in obj.extract::>()?.iter() { + slf.insert_raising_if_exists( + key.extract::()?, + val.extract::()?, + )?; + } + Ok(slf) + } +} + +impl PyAmazonS3Config { + fn new() -> Self { + Self(HashMap::new()) + } + + fn apply_config(self, mut builder: AmazonS3Builder) -> AmazonS3Builder { + for (key, value) in self.0.into_iter() { + builder = builder.with_config(key.0, value.0); + } + builder + } + + fn merge(mut self, other: PyAmazonS3Config) -> PyObjectStoreResult { + for (key, val) in other.0.into_iter() { + self.insert_raising_if_exists(key, val)?; + } + + Ok(self) + } + + fn insert_raising_if_exists( + &mut self, + key: impl Into, + val: impl Into, + ) -> PyObjectStoreResult<()> { + let key = key.into(); + let old_value = self.0.insert(key.clone(), PyConfigValue::new(val.into())); + if old_value.is_some() { + return Err(GenericError::new_err(format!( + "Duplicate key {} provided", + key.0.as_ref() + )) + .into()); + } + + Ok(()) + } + + /// Insert a key only if it does not already exist. + /// + /// This is used for URL parsing, where any parts of the URL **do not** override any + /// configuration keys passed manually. + fn insert_if_not_exists( + &mut self, + key: impl Into, + val: impl Into, + ) { + self.0.entry(key.into()).or_insert(PyConfigValue::new(val)); + } +} + +fn combine_config_kwargs( + config: PyAmazonS3Config, + kwargs: Option, +) -> PyObjectStoreResult { + if let Some(kwargs) = kwargs { + config.merge(kwargs) + } else { + Ok(config) + } +} + +/// Sets properties on a configuration based on a URL +/// +/// This is vendored from +/// https://github.com/apache/arrow-rs/blob/f7263e253655b2ee613be97f9d00e063444d3df5/object_store/src/aws/builder.rs#L600-L647 +/// +/// We do our own URL parsing so that we can keep our own config in sync with what is passed to the +/// underlying ObjectStore builder. Passing the URL on verbatim makes it hard because the URL +/// parsing only happens in `build()`. Then the config parameters we have don't include any config +/// applied from the URL. +fn parse_url( + config: Option, + parsed: &Url, +) -> object_store::Result { + let host = parsed + .host_str() + .ok_or_else(|| ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + })?; + let mut config = config.unwrap_or_default(); + + match parsed.scheme() { + "s3" | "s3a" => { + config.insert_if_not_exists(AmazonS3ConfigKey::Bucket, host); + } + "https" => match host.splitn(4, '.').collect_tuple() { + Some(("s3", region, "amazonaws", "com")) => { + config.insert_if_not_exists(AmazonS3ConfigKey::Region, region); + let bucket = parsed.path_segments().into_iter().flatten().next(); + if let Some(bucket) = bucket { + config.insert_if_not_exists(AmazonS3ConfigKey::Bucket, bucket); + } + } + Some((bucket, "s3", "amazonaws", "com")) => { + config.insert_if_not_exists(AmazonS3ConfigKey::Bucket, bucket); + config.insert_if_not_exists(AmazonS3ConfigKey::VirtualHostedStyleRequest, "true"); + } + Some((bucket, "s3", region, "amazonaws.com")) => { + config.insert_if_not_exists(AmazonS3ConfigKey::Bucket, bucket); + config.insert_if_not_exists(AmazonS3ConfigKey::Region, region); + config.insert_if_not_exists(AmazonS3ConfigKey::VirtualHostedStyleRequest, "true"); + } + Some((account, "r2", "cloudflarestorage", "com")) => { + config.insert_if_not_exists(AmazonS3ConfigKey::Region, "auto"); + let endpoint = format!("https://{account}.r2.cloudflarestorage.com"); + config.insert_if_not_exists(AmazonS3ConfigKey::Endpoint, endpoint); + + let bucket = parsed.path_segments().into_iter().flatten().next(); + if let Some(bucket) = bucket { + config.insert_if_not_exists(AmazonS3ConfigKey::Bucket, bucket); + } + } + _ => { + return Err(ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + } + .into()) + } + }, + scheme => { + let scheme = scheme.into(); + return Err(ParseUrlError::UnknownUrlScheme { scheme }.into()); + } + }; + + Ok(config) +} diff --git a/vendor/pyo3-object_store/src/azure/credentials.rs b/vendor/pyo3-object_store/src/azure/credentials.rs new file mode 100644 index 00000000000..48890fab920 --- /dev/null +++ b/vendor/pyo3-object_store/src/azure/credentials.rs @@ -0,0 +1,313 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, Utc}; +use object_store::azure::{AzureAccessKey, AzureCredential}; +use object_store::CredentialProvider; +use percent_encoding::percent_decode_str; +use pyo3::exceptions::{PyTypeError, PyValueError}; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::PyTuple; + +use crate::azure::error::Error; +use crate::azure::store::PyAzureConfig; +use crate::credentials::{is_awaitable, TemporaryToken, TokenCache}; +use crate::path::PyPath; +use crate::PyObjectStoreError; + +struct PyAzureAccessKey { + access_key: AzureAccessKey, + expires_at: Option>, +} + +// Extract the dict {"access_key": str} +impl<'py> FromPyObject<'_, 'py> for PyAzureAccessKey { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let s = obj + .get_item(intern!(obj.py(), "access_key"))? + .extract::()?; + let access_key = + AzureAccessKey::try_new(&s).map_err(|err| PyValueError::new_err(err.to_string()))?; + let expires_at = obj.get_item(intern!(obj.py(), "expires_at"))?.extract()?; + Ok(Self { + access_key, + expires_at, + }) + } +} + +struct PyAzureSASToken { + sas_token: Vec<(String, String)>, + expires_at: Option>, +} + +// Extract the dict {"sas_token": str | list[tuple[str, str]]} +impl<'py> FromPyObject<'_, 'py> for PyAzureSASToken { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let expires_at = obj.get_item(intern!(obj.py(), "expires_at"))?.extract()?; + let py_sas_token = obj.get_item(intern!(obj.py(), "sas_token"))?; + if let Ok(sas_token_str) = py_sas_token.extract::() { + Ok(Self { + sas_token: split_sas(&sas_token_str).map_err(PyObjectStoreError::from)?, + expires_at, + }) + } else if let Ok(sas_token_list) = py_sas_token.extract() { + Ok(Self { + sas_token: sas_token_list, + expires_at, + }) + } else { + Err(PyTypeError::new_err( + "Expected a string or list[tuple[str, str]]", + )) + } + } +} + +struct PyBearerToken { + token: String, + expires_at: Option>, +} + +// Extract the dict {"token": str} +impl<'py> FromPyObject<'_, 'py> for PyBearerToken { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let token = obj.get_item(intern!(obj.py(), "token"))?.extract()?; + let expires_at = obj.get_item(intern!(obj.py(), "expires_at"))?.extract()?; + Ok(Self { token, expires_at }) + } +} + +#[derive(FromPyObject)] +enum PyAzureCredential { + AccessKey(PyAzureAccessKey), + SASToken(PyAzureSASToken), + BearerToken(PyBearerToken), +} + +impl PyAzureCredential { + fn into_temporary_token(self) -> TemporaryToken> { + let (credential, expiry) = match self { + Self::AccessKey(key) => (AzureCredential::AccessKey(key.access_key), key.expires_at), + Self::SASToken(token) => (AzureCredential::SASToken(token.sas_token), token.expires_at), + Self::BearerToken(token) => { + (AzureCredential::BearerToken(token.token), token.expires_at) + } + }; + + TemporaryToken { + token: Arc::new(credential), + expiry, + } + } +} + +impl From for AzureCredential { + fn from(value: PyAzureCredential) -> Self { + match value { + PyAzureCredential::AccessKey(key) => Self::AccessKey(key.access_key), + PyAzureCredential::SASToken(token) => Self::SASToken(token.sas_token), + PyAzureCredential::BearerToken(token) => Self::BearerToken(token.token), + } + } +} + +// Vendored from upstream +// https://github.com/apache/arrow-rs/blob/92cfd99e9ab4a6c54500ec65252027b9edf1ee55/object_store/src/azure/builder.rs#L1055-L1072 +fn split_sas(sas: &str) -> Result, object_store::Error> { + let sas = percent_decode_str(sas) + .decode_utf8() + .map_err(|source| Error::DecodeSasKey { source })?; + let kv_str_pairs = sas + .trim_start_matches('?') + .split('&') + .filter(|s| !s.chars().all(char::is_whitespace)); + let mut pairs = Vec::new(); + for kv_pair_str in kv_str_pairs { + let (k, v) = kv_pair_str + .trim() + .split_once('=') + .ok_or(Error::MissingSasComponent {})?; + pairs.push((k.into(), v.into())) + } + Ok(pairs) +} + +#[derive(Debug)] +pub struct PyAzureCredentialProvider { + /// The provided user callback to manage credential refresh + user_callback: Py, + cache: TokenCache>, + /// An optional config passed down from the credential provider class + config: Option, + /// An optional prefix passed down from the credential provider class + prefix: Option, +} + +impl PyAzureCredentialProvider { + /// Access the Azure config passed down from the credential provider + pub(crate) fn config(&self) -> Option<&PyAzureConfig> { + self.config.as_ref() + } + + /// Access the store prefix passed down from the credential provider + pub(crate) fn prefix(&self) -> Option<&PyPath> { + self.prefix.as_ref() + } + + fn equals(&self, py: Python, other: &Self) -> PyResult { + self.user_callback + .call_method1(py, "__eq__", PyTuple::new(py, vec![&other.user_callback])?)? + .extract(py) + } +} + +impl Clone for PyAzureCredentialProvider { + fn clone(&self) -> Self { + let cloned_callback = Python::attach(|py| self.user_callback.clone_ref(py)); + Self { + user_callback: cloned_callback, + cache: self.cache.clone(), + config: self.config.clone(), + prefix: self.prefix.clone(), + } + } +} + +impl PartialEq for PyAzureCredentialProvider { + fn eq(&self, other: &Self) -> bool { + Python::attach(|py| self.equals(py, other)).unwrap_or(false) + } +} + +impl<'py> FromPyObject<'_, 'py> for PyAzureCredentialProvider { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if !obj.hasattr(intern!(obj.py(), "__call__"))? { + return Err(PyTypeError::new_err( + "Expected callable object for credential_provider.", + )); + } + + let mut cache = TokenCache::default(); + if let Ok(refresh_threshold) = obj.getattr(intern!(obj.py(), "refresh_threshold")) { + cache = cache.with_min_ttl(refresh_threshold.extract()?); + } + + let config = if let Ok(config) = obj.getattr(intern!(obj.py(), "config")) { + config.extract()? + } else { + // Allow not having a `config` attribute + None + }; + + let prefix = if let Ok(prefix) = obj.getattr(intern!(obj.py(), "prefix")) { + prefix.extract()? + } else { + // Allow not having a `prefix` attribute + None + }; + + Ok(Self { + user_callback: obj.as_unbound().clone_ref(obj.py()), + cache, + config, + prefix, + }) + } +} + +impl<'py> IntoPyObject<'py> for &PyAzureCredentialProvider { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.user_callback.bind(py).clone()) + } +} + +impl<'py> IntoPyObject<'py> for PyAzureCredentialProvider { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +enum PyCredentialProviderResult { + Async(Py), + Sync(PyAzureCredential), +} + +impl PyCredentialProviderResult { + async fn resolve(self) -> PyResult { + match self { + Self::Sync(credentials) => Ok(credentials), + Self::Async(coroutine) => { + let future = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(coroutine.bind(py).clone()) + })?; + let result = future.await?; + Python::attach(|py| result.extract(py)) + } + } + } +} + +impl<'py> FromPyObject<'_, 'py> for PyCredentialProviderResult { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if is_awaitable(&obj)? { + Ok(Self::Async(obj.as_unbound().clone_ref(obj.py()))) + } else { + Ok(Self::Sync(obj.extract()?)) + } + } +} + +impl PyAzureCredentialProvider { + async fn call(&self) -> PyResult { + let call_result = Python::attach(|py| { + self.user_callback + .call0(py)? + .extract::(py) + })?; + call_result.resolve().await + } + + /// Call the user-provided callback + async fn fetch_token(&self) -> object_store::Result>> { + let credential = self + .call() + .await + .map_err(|err| object_store::Error::Unauthenticated { + path: "External Azure credential provider".to_string(), + source: Box::new(err), + })?; + + Ok(credential.into_temporary_token()) + } +} + +// TODO: store expiration time and only call the external Python function as needed +#[async_trait] +impl CredentialProvider for PyAzureCredentialProvider { + type Credential = AzureCredential; + + async fn get_credential(&self) -> object_store::Result> { + self.cache.get_or_insert_with(|| self.fetch_token()).await + } +} diff --git a/vendor/pyo3-object_store/src/azure/error.rs b/vendor/pyo3-object_store/src/azure/error.rs new file mode 100644 index 00000000000..8b39dacb659 --- /dev/null +++ b/vendor/pyo3-object_store/src/azure/error.rs @@ -0,0 +1,21 @@ +const STORE: &str = "MicrosoftAzure"; + +// Vendored from upstream +/// A specialized `Error` for Azure builder-related errors +#[derive(Debug, thiserror::Error)] +pub(crate) enum Error { + #[error("Failed parsing an SAS key")] + DecodeSasKey { source: std::str::Utf8Error }, + + #[error("Missing component in SAS query pair")] + MissingSasComponent {}, +} + +impl From for object_store::Error { + fn from(source: Error) -> Self { + Self::Generic { + store: STORE, + source: Box::new(source), + } + } +} diff --git a/vendor/pyo3-object_store/src/azure/mod.rs b/vendor/pyo3-object_store/src/azure/mod.rs new file mode 100644 index 00000000000..b413a914c97 --- /dev/null +++ b/vendor/pyo3-object_store/src/azure/mod.rs @@ -0,0 +1,5 @@ +mod credentials; +mod error; +mod store; + +pub use store::PyAzureStore; diff --git a/vendor/pyo3-object_store/src/azure/store.rs b/vendor/pyo3-object_store/src/azure/store.rs new file mode 100644 index 00000000000..97420416e12 --- /dev/null +++ b/vendor/pyo3-object_store/src/azure/store.rs @@ -0,0 +1,489 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::azure::{AzureConfigKey, MicrosoftAzure, MicrosoftAzureBuilder}; +use object_store::ObjectStoreScheme; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::{PyDict, PyString, PyTuple, PyType}; +use pyo3::{intern, IntoPyObjectExt}; +use url::Url; + +use crate::azure::credentials::PyAzureCredentialProvider; +use crate::client::PyClientOptions; +use crate::config::PyConfigValue; +use crate::error::{GenericError, ParseUrlError, PyObjectStoreError, PyObjectStoreResult}; +use crate::path::PyPath; +use crate::retry::PyRetryConfig; +use crate::{MaybePrefixedStore, PyUrl}; + +#[derive(Debug, Clone, PartialEq)] +struct AzureConfig { + prefix: Option, + config: PyAzureConfig, + client_options: Option, + retry_config: Option, + credential_provider: Option, +} + +impl AzureConfig { + fn account_name(&self) -> &str { + self.config + .0 + .get(&PyAzureConfigKey(AzureConfigKey::AccountName)) + .expect("Account name should always exist in the config") + .as_ref() + } + + fn container_name(&self) -> &str { + self.config + .0 + .get(&PyAzureConfigKey(AzureConfigKey::ContainerName)) + .expect("Container should always exist in the config") + .as_ref() + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let args = PyTuple::empty(py).into_bound_py_any(py)?; + let kwargs = PyDict::new(py); + + if let Some(prefix) = &self.prefix { + kwargs.set_item(intern!(py, "prefix"), prefix.as_ref().as_ref())?; + } + kwargs.set_item(intern!(py, "config"), &self.config)?; + if let Some(client_options) = &self.client_options { + kwargs.set_item(intern!(py, "client_options"), client_options)?; + } + if let Some(retry_config) = &self.retry_config { + kwargs.set_item(intern!(py, "retry_config"), retry_config)?; + } + if let Some(credential_provider) = &self.credential_provider { + kwargs.set_item("credential_provider", credential_provider)?; + } + + PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) + } +} + +/// A Python-facing wrapper around a [`MicrosoftAzure`]. +#[derive(Debug, Clone)] +#[pyclass(name = "AzureStore", frozen, subclass, from_py_object)] +pub struct PyAzureStore { + store: Arc>, + /// A config used for pickling. This must stay in sync with the underlying store's config. + config: AzureConfig, +} + +impl AsRef>> for PyAzureStore { + fn as_ref(&self) -> &Arc> { + &self.store + } +} + +impl PyAzureStore { + /// Consume self and return the underlying [`MicrosoftAzure`]. + pub fn into_inner(self) -> Arc> { + self.store + } +} + +#[pymethods] +impl PyAzureStore { + // Create from parameters + #[new] + #[pyo3(signature = (container_name=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + fn new( + container_name: Option, + mut prefix: Option, + config: Option, + client_options: Option, + retry_config: Option, + credential_provider: Option, + kwargs: Option, + ) -> PyObjectStoreResult { + let mut builder = MicrosoftAzureBuilder::from_env(); + let mut config = config.unwrap_or_default(); + + if let Some(container_name) = container_name { + // Note: we apply the bucket to the config, not directly to the builder, so they stay + // in sync. + config.insert_raising_if_exists(AzureConfigKey::ContainerName, container_name)?; + } + + let mut combined_config = combine_config_kwargs(Some(config), kwargs)?; + + if let Some(client_options) = client_options.clone() { + builder = builder.with_client_options(client_options.into()) + } + if let Some(retry_config) = retry_config.clone() { + builder = builder.with_retry(retry_config.into()) + } + + if let Some(credential_provider) = credential_provider.clone() { + // Apply credential provider config onto main config + if let Some(credential_config) = credential_provider.config() { + for (key, val) in credential_config.0.iter() { + // Give precedence to passed-in config values + combined_config.insert_if_not_exists(key.clone(), val.clone()); + } + } + + if let Some(passed_down_prefix) = credential_provider.prefix() { + // Don't override a prefix manually passed in to AzureStore + // + // If a user wishes to override a prefix passed down by a credential provider, they + // can pass `prefix=""` or `prefix="/"` to the AzureStore. + if prefix.is_none() { + prefix = Some(passed_down_prefix.clone()); + } + } + + builder = builder.with_credentials(Arc::new(credential_provider)); + } + + builder = combined_config.clone().apply_config(builder); + + Ok(Self { + store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), + config: AzureConfig { + prefix, + config: combined_config, + client_options, + retry_config, + credential_provider, + }, + }) + } + + #[classmethod] + #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + pub(crate) fn from_url<'py>( + cls: &Bound<'py, PyType>, + url: PyUrl, + config: Option, + client_options: Option, + retry_config: Option, + credential_provider: Option, + kwargs: Option, + ) -> PyObjectStoreResult> { + // We manually parse the URL to find the prefix because `parse_url` does not apply the + // prefix. + let (_, prefix) = + ObjectStoreScheme::parse(url.as_ref()).map_err(object_store::Error::from)?; + let prefix: Option = if prefix.parts().count() != 0 { + Some(prefix.into()) + } else { + None + }; + let config = parse_url(config, url.as_ref())?; + + // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the + // subclass + let kwargs = kwargs.unwrap_or_default().into_pyobject(cls.py())?; + kwargs.set_item("prefix", prefix)?; + kwargs.set_item("config", config)?; + kwargs.set_item("client_options", client_options)?; + kwargs.set_item("retry_config", retry_config)?; + kwargs.set_item("credential_provider", credential_provider)?; + Ok(cls.call((), Some(&kwargs))?) + } + + fn __eq__(&self, other: &Bound) -> bool { + // Ensure we never error on __eq__ by returning false if the other object is not the same + // type + other + .cast::() + .map(|other| self.config == other.get().config) + .unwrap_or(false) + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + self.config.__getnewargs_ex__(py) + } + + fn __repr__(&self) -> String { + let account_name = self.config.account_name(); + let container_name = self.config.container_name(); + if let Some(prefix) = &self.config.prefix { + format!( + "AzureStore(container_name=\"{}\", account_name=\"{}\", prefix=\"{}\")", + container_name, + account_name, + prefix.as_ref() + ) + } else { + format!( + "AzureStore(container_name=\"{container_name}\", account_name=\"{account_name}\")" + ) + } + } + + #[getter] + fn prefix(&self) -> Option<&PyPath> { + self.config.prefix.as_ref() + } + + #[getter] + fn config(&self) -> &PyAzureConfig { + &self.config.config + } + + #[getter] + fn client_options(&self) -> Option<&PyClientOptions> { + self.config.client_options.as_ref() + } + + #[getter] + fn credential_provider(&self) -> Option<&PyAzureCredentialProvider> { + self.config.credential_provider.as_ref() + } + + #[getter] + fn retry_config(&self) -> Option<&PyRetryConfig> { + self.config.retry_config.as_ref() + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct PyAzureConfigKey(AzureConfigKey); + +impl<'py> FromPyObject<'_, 'py> for PyAzureConfigKey { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let s = obj.extract::()?.to_lowercase(); + let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?; + Ok(Self(key)) + } +} + +impl AsRef for PyAzureConfigKey { + fn as_ref(&self) -> &str { + self.0.as_ref() + } +} + +impl<'py> IntoPyObject<'py> for &PyAzureConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let s = self.0.as_ref(); + // Anything with an `azure_storage_` prefix we can fully strip + if let Some(stripped) = s.strip_prefix("azure_storage_") { + return Ok(PyString::new(py, stripped)); + } + Ok(PyString::new( + py, + s.strip_prefix("azure_") + .expect("Expected config prefix to start with azure_"), + )) + } +} + +impl<'py> IntoPyObject<'py> for PyAzureConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +impl From for PyAzureConfigKey { + fn from(value: AzureConfigKey) -> Self { + Self(value) + } +} + +impl From for AzureConfigKey { + fn from(value: PyAzureConfigKey) -> Self { + value.0 + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, IntoPyObject, IntoPyObjectRef)] +pub struct PyAzureConfig(HashMap); + +// Note: we manually impl FromPyObject instead of deriving it so that we can raise an +// UnknownConfigurationKeyError instead of a `TypeError` on invalid config keys. +// +// We also manually impl this so that we can raise on duplicate keys. +impl<'py> FromPyObject<'_, 'py> for PyAzureConfig { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let mut slf = Self::new(); + for (key, val) in obj.extract::>()?.iter() { + slf.insert_raising_if_exists( + key.extract::()?, + val.extract::()?, + )?; + } + Ok(slf) + } +} + +impl PyAzureConfig { + fn new() -> Self { + Self(HashMap::new()) + } + + fn apply_config(self, mut builder: MicrosoftAzureBuilder) -> MicrosoftAzureBuilder { + for (key, value) in self.0.into_iter() { + builder = builder.with_config(key.0, value.0); + } + builder + } + + fn merge(mut self, other: PyAzureConfig) -> PyObjectStoreResult { + for (key, val) in other.0.into_iter() { + self.insert_raising_if_exists(key, val)?; + } + + Ok(self) + } + + fn insert_raising_if_exists( + &mut self, + key: impl Into, + val: impl Into, + ) -> PyObjectStoreResult<()> { + let key = key.into(); + let old_value = self.0.insert(key.clone(), PyConfigValue::new(val.into())); + if old_value.is_some() { + return Err(GenericError::new_err(format!( + "Duplicate key {} provided", + key.0.as_ref() + )) + .into()); + } + + Ok(()) + } + + /// Insert a key only if it does not already exist. + /// + /// This is used for URL parsing, where any parts of the URL **do not** override any + /// configuration keys passed manually. + fn insert_if_not_exists(&mut self, key: impl Into, val: impl Into) { + self.0.entry(key.into()).or_insert(PyConfigValue::new(val)); + } +} + +fn combine_config_kwargs( + config: Option, + kwargs: Option, +) -> PyObjectStoreResult { + match (config, kwargs) { + (None, None) => Ok(Default::default()), + (Some(x), None) | (None, Some(x)) => Ok(x), + (Some(config), Some(kwargs)) => Ok(config.merge(kwargs)?), + } +} + +/// Sets properties on this builder based on a URL +/// +/// This is vendored from +/// https://github.com/apache/arrow-rs/blob/f7263e253655b2ee613be97f9d00e063444d3df5/object_store/src/azure/builder.rs#L639-L705 +/// +/// We do our own URL parsing so that we can keep our own config in sync with what is passed to the +/// underlying ObjectStore builder. Passing the URL on verbatim makes it hard because the URL +/// parsing only happens in `build()`. Then the config parameters we have don't include any config +/// applied from the URL. +fn parse_url(config: Option, parsed: &Url) -> object_store::Result { + let host = parsed + .host_str() + .ok_or_else(|| ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + })?; + let mut config = config.unwrap_or_default(); + + let validate = |s: &str| match s.contains('.') { + true => Err(ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + }), + false => Ok(s.to_string()), + }; + + match parsed.scheme() { + "adl" | "azure" => { + config.insert_if_not_exists(AzureConfigKey::ContainerName, validate(host)?); + } + "az" | "abfs" | "abfss" => { + // abfs(s) might refer to the fsspec convention abfs:/// + // or the convention for the hadoop driver abfs[s]://@.dfs.core.windows.net/ + if parsed.username().is_empty() { + config.insert_if_not_exists(AzureConfigKey::ContainerName, validate(host)?); + } else { + match host.split_once('.') { + Some((a, "dfs.core.windows.net")) | Some((a, "blob.core.windows.net")) => { + config.insert_if_not_exists(AzureConfigKey::AccountName, validate(a)?); + config.insert_if_not_exists( + AzureConfigKey::ContainerName, + validate(parsed.username())?, + ); + } + Some((a, "dfs.fabric.microsoft.com")) + | Some((a, "blob.fabric.microsoft.com")) => { + config.insert_if_not_exists(AzureConfigKey::AccountName, validate(a)?); + config.insert_if_not_exists( + AzureConfigKey::ContainerName, + validate(parsed.username())?, + ); + config.insert_if_not_exists(AzureConfigKey::UseFabricEndpoint, "true"); + } + _ => { + return Err(ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + } + .into()) + } + } + } + } + "https" => match host.split_once('.') { + Some((a, "dfs.core.windows.net")) | Some((a, "blob.core.windows.net")) => { + config.insert_if_not_exists(AzureConfigKey::AccountName, validate(a)?); + let container = + parsed.path_segments().unwrap().next().expect( + "iterator always contains at least one string (which may be empty)", + ); + if !container.is_empty() { + config + .insert_if_not_exists(AzureConfigKey::ContainerName, validate(container)?); + } + } + Some((a, "dfs.fabric.microsoft.com")) | Some((a, "blob.fabric.microsoft.com")) => { + config.insert_if_not_exists(AzureConfigKey::AccountName, validate(a)?); + // Attempt to infer the container name from the URL + // - https://onelake.dfs.fabric.microsoft.com///Files/test.csv + // - https://onelake.dfs.fabric.microsoft.com//.// + // + // See + let workspace = + parsed.path_segments().unwrap().next().expect( + "iterator always contains at least one string (which may be empty)", + ); + if !workspace.is_empty() { + config.insert_if_not_exists(AzureConfigKey::ContainerName, workspace); + } + config.insert_if_not_exists(AzureConfigKey::UseFabricEndpoint, "true"); + } + _ => { + return Err(ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + } + .into()) + } + }, + scheme => { + let scheme = scheme.into(); + return Err(ParseUrlError::UnknownUrlScheme { scheme }.into()); + } + } + + Ok(config) +} diff --git a/vendor/pyo3-object_store/src/client.rs b/vendor/pyo3-object_store/src/client.rs new file mode 100644 index 00000000000..9fe159fb0de --- /dev/null +++ b/vendor/pyo3-object_store/src/client.rs @@ -0,0 +1,180 @@ +use std::collections::HashMap; +use std::str::FromStr; + +use http::{HeaderMap, HeaderName, HeaderValue}; +use object_store::{ClientConfigKey, ClientOptions}; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::pybacked::{PyBackedBytes, PyBackedStr}; +use pyo3::types::{PyDict, PyString}; + +use crate::config::PyConfigValue; +use crate::error::PyObjectStoreError; + +/// A wrapper around `ClientConfigKey` that implements [`FromPyObject`]. +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct PyClientConfigKey(ClientConfigKey); + +impl<'py> FromPyObject<'_, 'py> for PyClientConfigKey { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let s = obj.extract::()?.to_lowercase(); + let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?; + Ok(Self(key)) + } +} + +impl<'py> IntoPyObject<'py> for PyClientConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_ref())) + } +} + +impl<'py> IntoPyObject<'py> for &PyClientConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_ref())) + } +} + +/// A wrapper around `ClientOptions` that implements [`FromPyObject`]. +#[derive(Clone, Debug, PartialEq)] +pub struct PyClientOptions { + string_options: HashMap, + default_headers: Option, +} + +impl<'py> FromPyObject<'_, 'py> for PyClientOptions { + type Error = PyErr; + + // Need custom extraction because default headers is not a string value + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let dict = obj.extract::>()?; + let mut string_options = HashMap::new(); + let mut default_headers = None; + + for (key, value) in dict.iter() { + if let Ok(key) = key.extract::() { + string_options.insert(key, value.extract::()?); + } else { + let key = key.extract::()?; + if &key == "default_headers" { + default_headers = Some(value.extract::()?); + } else { + return Err(PyValueError::new_err(format!("Invalid key: {key}."))); + } + } + } + + Ok(Self { + string_options, + default_headers, + }) + } +} + +impl<'py> IntoPyObject<'py> for PyClientOptions { + type Target = PyDict; + type Output = Bound<'py, PyDict>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let dict = self.string_options.into_pyobject(py)?; + if let Some(headers) = self.default_headers { + dict.set_item("default_headers", headers)?; + } + Ok(dict) + } +} + +impl<'py> IntoPyObject<'py> for &PyClientOptions { + type Target = PyDict; + type Output = Bound<'py, PyDict>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let dict = (&self.string_options).into_pyobject(py)?; + if let Some(headers) = &self.default_headers { + dict.set_item("default_headers", headers)?; + } + Ok(dict.clone()) + } +} + +impl From for ClientOptions { + fn from(value: PyClientOptions) -> Self { + let mut options = ClientOptions::new(); + for (key, value) in value.string_options.into_iter() { + options = options.with_config(key.0, value.0); + } + + if let Some(headers) = value.default_headers { + options = options.with_default_headers(headers.0); + } + + options + } +} + +#[derive(Clone, Debug, PartialEq)] +struct PyHeaderMap(HeaderMap); + +impl<'py> FromPyObject<'_, 'py> for PyHeaderMap { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let dict = obj.extract::>()?; + let mut header_map = HeaderMap::with_capacity(dict.len()); + for (key, value) in dict.iter() { + let key = HeaderName::from_str(&key.extract::()?) + .map_err(|err| PyValueError::new_err(err.to_string()))?; + + // HTTP Header values can have non-ascii bytes, so first try to extract as bytes. + let value = if let Ok(value_bytes) = value.extract::() { + HeaderValue::from_bytes(&value_bytes) + } else { + HeaderValue::from_str(&value.extract::()?) + } + .map_err(|err| PyValueError::new_err(err.to_string()))?; + + header_map.insert(key, value); + } + Ok(Self(header_map)) + } +} + +impl<'py> IntoPyObject<'py> for PyHeaderMap { + type Target = PyDict; + type Output = Bound<'py, PyDict>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let dict = PyDict::new(py); + for (key, value) in self.0.iter() { + dict.set_item(key.as_str(), value.as_bytes())?; + } + Ok(dict) + } +} + +impl<'py> IntoPyObject<'py> for &PyHeaderMap { + type Target = PyDict; + type Output = Bound<'py, PyDict>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let dict = PyDict::new(py); + for (key, value) in self.0.iter() { + dict.set_item(key.as_str(), value.as_bytes())?; + } + Ok(dict) + } +} diff --git a/vendor/pyo3-object_store/src/config.rs b/vendor/pyo3-object_store/src/config.rs new file mode 100644 index 00000000000..29c2acd6a9f --- /dev/null +++ b/vendor/pyo3-object_store/src/config.rs @@ -0,0 +1,58 @@ +use std::time::Duration; + +use humantime::format_duration; +use pyo3::prelude::*; + +/// A wrapper around `String` used to store values for config values. +/// +/// Supported Python input: +/// +/// - `True` and `False` (becomes `"true"` and `"false"`) +/// - `timedelta` +/// - `str` +#[derive(Clone, Debug, PartialEq, Eq, Hash, IntoPyObject, IntoPyObjectRef)] +pub struct PyConfigValue(pub String); + +impl PyConfigValue { + pub(crate) fn new(val: impl Into) -> Self { + Self(val.into()) + } +} + +impl AsRef for PyConfigValue { + fn as_ref(&self) -> &str { + &self.0 + } +} + +impl<'py> FromPyObject<'_, 'py> for PyConfigValue { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if let Ok(val) = obj.extract::() { + Ok(val.into()) + } else if let Ok(duration) = obj.extract::() { + Ok(duration.into()) + } else { + Ok(Self(obj.extract()?)) + } + } +} + +impl From for String { + fn from(value: PyConfigValue) -> Self { + value.0 + } +} + +impl From for PyConfigValue { + fn from(value: bool) -> Self { + Self(value.to_string()) + } +} + +impl From for PyConfigValue { + fn from(value: Duration) -> Self { + Self(format_duration(value).to_string()) + } +} diff --git a/vendor/pyo3-object_store/src/credentials.rs b/vendor/pyo3-object_store/src/credentials.rs new file mode 100644 index 00000000000..2b743fb8f65 --- /dev/null +++ b/vendor/pyo3-object_store/src/credentials.rs @@ -0,0 +1,101 @@ +use chrono::Utc; +use chrono::{DateTime, TimeDelta}; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyTuple; +use std::future::Future; +use tokio::sync::Mutex; + +/// A temporary authentication token with an associated expiry +#[derive(Debug, Clone)] +pub(crate) struct TemporaryToken { + /// The temporary credential + pub token: T, + /// The instant at which this credential is no longer valid + /// None means the credential does not expire + pub expiry: Option>, +} + +/// Provides [`TokenCache::get_or_insert_with`] which can be used to cache a +/// [`TemporaryToken`] based on its expiry +#[derive(Debug)] +pub(crate) struct TokenCache { + /// A temporary token and the instant at which it was fetched + cache: Mutex, DateTime)>>, + min_ttl: TimeDelta, + /// How long to wait before re-attempting a token fetch after receiving one that + /// is still within the min-ttl + fetch_backoff: TimeDelta, +} + +impl Default for TokenCache { + fn default() -> Self { + Self { + cache: Default::default(), + min_ttl: TimeDelta::seconds(300), + fetch_backoff: TimeDelta::milliseconds(100), + } + } +} + +impl Clone for TokenCache { + /// Cloning the token cache invalidates the cache. + fn clone(&self) -> Self { + Self { + cache: Default::default(), + min_ttl: self.min_ttl, + fetch_backoff: self.fetch_backoff, + } + } +} + +impl TokenCache { + /// Override the minimum remaining TTL for a cached token to be used + pub(crate) fn with_min_ttl(self, min_ttl: TimeDelta) -> Self { + Self { min_ttl, ..self } + } + + pub(crate) async fn get_or_insert_with(&self, f: F) -> Result + where + F: FnOnce() -> Fut + Send, + Fut: Future, E>> + Send, + { + // let now = Instant::now(); + let now = Utc::now(); + + let mut locked = self.cache.lock().await; + + if let Some((cached, fetched_at)) = locked.as_ref() { + match cached.expiry { + Some(expiry_time) => { + // let x = ttl - now; + // let x = ttl.signed_duration_since(now); + // let x = expiry_time - now > self.min_ttl.into(); + if expiry_time - now > self.min_ttl || + // if we've recently attempted to fetch this token and it's not actually + // expired, we'll wait to re-fetch it and return the cached one + (Utc::now() - fetched_at < self.fetch_backoff && expiry_time - now > TimeDelta::zero()) + { + return Ok(cached.token.clone()); + } + } + None => return Ok(cached.token.clone()), + } + } + + let cached = f().await?; + let token = cached.token.clone(); + *locked = Some((cached, Utc::now())); + + Ok(token) + } +} + +/// Check whether a Python object is awaitable +pub(crate) fn is_awaitable(ob: &Bound) -> PyResult { + let py = ob.py(); + let inspect_mod = py.import(intern!(py, "inspect"))?; + inspect_mod + .call_method1(intern!(py, "isawaitable"), PyTuple::new(py, [ob])?)? + .extract::() +} diff --git a/vendor/pyo3-object_store/src/error.rs b/vendor/pyo3-object_store/src/error.rs new file mode 100644 index 00000000000..4226d241964 --- /dev/null +++ b/vendor/pyo3-object_store/src/error.rs @@ -0,0 +1,193 @@ +//! Contains the [`PyObjectStoreError`], the error enum returned by all fallible functions in this +//! crate. + +use pyo3::exceptions::{PyFileNotFoundError, PyIOError, PyNotImplementedError, PyValueError}; +use pyo3::prelude::*; +use pyo3::{create_exception, CastError}; +use thiserror::Error; + +// Base exception +// Note that this is named `BaseError` instead of `ObstoreError` to not leak the name "obstore" to +// other Rust-Python libraries using pyo3-object_store. +create_exception!( + pyo3_object_store, + BaseError, + pyo3::exceptions::PyException, + "The base Python-facing exception from which all other errors subclass." +); + +// Subclasses from base exception +create_exception!( + pyo3_object_store, + GenericError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::Generic]." +); +create_exception!( + pyo3_object_store, + NotFoundError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::NotFound]." +); +create_exception!( + pyo3_object_store, + InvalidPathError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::InvalidPath]." +); +create_exception!( + pyo3_object_store, + JoinError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::JoinError]." +); +create_exception!( + pyo3_object_store, + NotSupportedError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::NotSupported]." +); +create_exception!( + pyo3_object_store, + AlreadyExistsError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::AlreadyExists]." +); +create_exception!( + pyo3_object_store, + PreconditionError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::Precondition]." +); +create_exception!( + pyo3_object_store, + NotModifiedError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::NotModified]." +); +create_exception!( + pyo3_object_store, + PermissionDeniedError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::PermissionDenied]." +); +create_exception!( + pyo3_object_store, + UnauthenticatedError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::Unauthenticated]." +); +create_exception!( + pyo3_object_store, + UnknownConfigurationKeyError, + BaseError, + "A Python-facing exception wrapping [object_store::Error::UnknownConfigurationKey]." +); + +/// The Error variants returned by this crate. +#[derive(Error, Debug)] +#[non_exhaustive] +pub enum PyObjectStoreError { + /// A wrapped [object_store::Error] + #[error(transparent)] + ObjectStoreError(#[from] object_store::Error), + + /// A wrapped [PyErr] + #[error(transparent)] + PyErr(#[from] PyErr), + + /// A wrapped [std::io::Error] + #[error(transparent)] + IOError(#[from] std::io::Error), +} + +impl From for PyErr { + fn from(error: PyObjectStoreError) -> Self { + match error { + PyObjectStoreError::PyErr(err) => err, + PyObjectStoreError::ObjectStoreError(ref err) => match err { + object_store::Error::Generic { + store: _, + source: _, + } => GenericError::new_err(print_with_debug(err)), + object_store::Error::NotFound { path: _, source: _ } => { + PyFileNotFoundError::new_err(print_with_debug(err)) + } + object_store::Error::InvalidPath { source: _ } => { + InvalidPathError::new_err(print_with_debug(err)) + } + object_store::Error::JoinError { source: _ } => { + JoinError::new_err(print_with_debug(err)) + } + object_store::Error::NotSupported { source: _ } => { + NotSupportedError::new_err(print_with_debug(err)) + } + object_store::Error::AlreadyExists { path: _, source: _ } => { + AlreadyExistsError::new_err(print_with_debug(err)) + } + object_store::Error::Precondition { path: _, source: _ } => { + PreconditionError::new_err(print_with_debug(err)) + } + object_store::Error::NotModified { path: _, source: _ } => { + NotModifiedError::new_err(print_with_debug(err)) + } + object_store::Error::NotImplemented { + operation: _, + implementer: _, + } => PyNotImplementedError::new_err(print_with_debug(err)), + object_store::Error::PermissionDenied { path: _, source: _ } => { + PermissionDeniedError::new_err(print_with_debug(err)) + } + object_store::Error::Unauthenticated { path: _, source: _ } => { + UnauthenticatedError::new_err(print_with_debug(err)) + } + object_store::Error::UnknownConfigurationKey { store: _, key: _ } => { + UnknownConfigurationKeyError::new_err(print_with_debug(err)) + } + _ => GenericError::new_err(print_with_debug(err)), + }, + PyObjectStoreError::IOError(err) => PyIOError::new_err(err), + } + } +} + +fn print_with_debug(err: &object_store::Error) -> String { + // #? gives "pretty-printing" for debug + // https://doc.rust-lang.org/std/fmt/trait.Debug.html + format!("{err}\n\nDebug source:\n{err:#?}") +} + +impl<'a, 'py> From> for PyObjectStoreError { + fn from(other: CastError<'a, 'py>) -> Self { + Self::PyErr(PyValueError::new_err(format!( + "Could not downcast: {other}", + ))) + } +} + +/// A type wrapper around `Result`. +pub type PyObjectStoreResult = Result; + +/// A specialized `Error` for object store-related errors +/// +/// Vendored from upstream to handle our vendored URL parsing +#[derive(Debug, thiserror::Error)] +pub(crate) enum ParseUrlError { + #[error( + "Unknown url scheme cannot be parsed into storage location: {}", + scheme + )] + UnknownUrlScheme { scheme: String }, + + #[error("URL did not match any known pattern for scheme: {}", url)] + UrlNotRecognised { url: String }, +} + +impl From for object_store::Error { + fn from(source: ParseUrlError) -> Self { + Self::Generic { + store: "S3", + source: Box::new(source), + } + } +} diff --git a/vendor/pyo3-object_store/src/gcp/credentials.rs b/vendor/pyo3-object_store/src/gcp/credentials.rs new file mode 100644 index 00000000000..06c8c1cb90d --- /dev/null +++ b/vendor/pyo3-object_store/src/gcp/credentials.rs @@ -0,0 +1,193 @@ +use std::sync::Arc; + +use async_trait::async_trait; +use chrono::{DateTime, TimeDelta, Utc}; +use object_store::gcp::GcpCredential; +use object_store::CredentialProvider; +use pyo3::exceptions::PyTypeError; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyTuple; + +use crate::credentials::{is_awaitable, TemporaryToken, TokenCache}; + +/// Ref https://github.com/apache/arrow-rs/pull/6638 +const DEFAULT_GCP_MIN_TTL: TimeDelta = TimeDelta::minutes(4); + +/// A wrapper around a [GcpCredential] that includes an optional expiry timestamp. +struct PyGcpCredential { + credential: GcpCredential, + expires_at: Option>, +} + +impl<'py> FromPyObject<'_, 'py> for PyGcpCredential { + type Error = PyErr; + + /// Converts from a Python dictionary of the form + /// + /// ```py + /// class GCSCredential(TypedDict): + /// token: str + /// expires_at: datetime | None + /// ``` + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let bearer = obj + .get_item(intern!(obj.py(), "token"))? + .extract::()?; + let credential = GcpCredential { bearer }; + let expires_at = obj.get_item(intern!(obj.py(), "expires_at"))?.extract()?; + Ok(Self { + credential, + expires_at, + }) + } +} + +// TODO: don't use a cache for static credentials where `expires_at` is `None` +// (so you don't need to access a mutex) +#[derive(Debug)] +pub struct PyGcpCredentialProvider { + /// The provided user callback to manage credential refresh + user_callback: Py, + /// The provided user callback to manage credential refresh + cache: TokenCache>, +} + +impl PyGcpCredentialProvider { + fn equals(&self, py: Python, other: &Self) -> PyResult { + self.user_callback + .call_method1(py, "__eq__", PyTuple::new(py, vec![&other.user_callback])?)? + .extract(py) + } +} + +impl Clone for PyGcpCredentialProvider { + fn clone(&self) -> Self { + let cloned_callback = Python::attach(|py| self.user_callback.clone_ref(py)); + Self { + user_callback: cloned_callback, + cache: self.cache.clone(), + } + } +} + +impl PartialEq for PyGcpCredentialProvider { + fn eq(&self, other: &Self) -> bool { + Python::attach(|py| self.equals(py, other)).unwrap_or(false) + } +} + +impl<'py> FromPyObject<'_, 'py> for PyGcpCredentialProvider { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if !obj.hasattr(intern!(obj.py(), "__call__"))? { + return Err(PyTypeError::new_err( + "Expected callable object for credential_provider.", + )); + } + let min_ttl = + if let Ok(refresh_threshold) = obj.getattr(intern!(obj.py(), "refresh_threshold")) { + refresh_threshold.extract()? + } else { + DEFAULT_GCP_MIN_TTL + }; + let cache = TokenCache::default().with_min_ttl(min_ttl); + Ok(Self { + user_callback: obj.as_unbound().clone_ref(obj.py()), + cache, + }) + } +} + +impl<'py> IntoPyObject<'py> for &PyGcpCredentialProvider { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(self.user_callback.bind(py).clone()) + } +} + +impl<'py> IntoPyObject<'py> for PyGcpCredentialProvider { + type Target = PyAny; + type Output = Bound<'py, PyAny>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +/// Note: This is copied across providers at the moment +enum PyCredentialProviderResult { + Async(Py), + Sync(PyGcpCredential), +} + +impl PyCredentialProviderResult { + async fn resolve(self) -> PyResult { + match self { + Self::Sync(credentials) => Ok(credentials), + Self::Async(coroutine) => { + let future = Python::attach(|py| { + pyo3_async_runtimes::tokio::into_future(coroutine.bind(py).clone()) + })?; + let result = future.await?; + Python::attach(|py| result.extract(py)) + } + } + } +} + +impl<'py> FromPyObject<'_, 'py> for PyCredentialProviderResult { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if is_awaitable(&obj)? { + Ok(Self::Async(obj.as_unbound().clone_ref(obj.py()))) + } else { + Ok(Self::Sync(obj.extract()?)) + } + } +} + +impl PyGcpCredentialProvider { + /// Call the user-provided callback and extract to a token. + /// + /// This is separate from `fetch_token` below so that it can return a `PyResult`. + async fn call(&self) -> PyResult { + let call_result = Python::attach(|py| { + self.user_callback + .call0(py)? + .extract::(py) + })?; + call_result.resolve().await + } + + /// Call the user-provided callback + async fn fetch_token(&self) -> object_store::Result>> { + let credential = self + .call() + .await + .map_err(|err| object_store::Error::Unauthenticated { + path: "External GCP credential provider".to_string(), + source: Box::new(err), + })?; + + Ok(TemporaryToken { + token: Arc::new(credential.credential), + expiry: credential.expires_at, + }) + } +} + +#[async_trait] +impl CredentialProvider for PyGcpCredentialProvider { + type Credential = GcpCredential; + + async fn get_credential(&self) -> object_store::Result> { + self.cache.get_or_insert_with(|| self.fetch_token()).await + } +} diff --git a/vendor/pyo3-object_store/src/gcp/mod.rs b/vendor/pyo3-object_store/src/gcp/mod.rs new file mode 100644 index 00000000000..78774709f24 --- /dev/null +++ b/vendor/pyo3-object_store/src/gcp/mod.rs @@ -0,0 +1,4 @@ +mod credentials; +mod store; + +pub use store::PyGCSStore; diff --git a/vendor/pyo3-object_store/src/gcp/store.rs b/vendor/pyo3-object_store/src/gcp/store.rs new file mode 100644 index 00000000000..260708dfaf8 --- /dev/null +++ b/vendor/pyo3-object_store/src/gcp/store.rs @@ -0,0 +1,380 @@ +use std::collections::HashMap; +use std::sync::Arc; + +use object_store::gcp::{GoogleCloudStorage, GoogleCloudStorageBuilder, GoogleConfigKey}; +use object_store::ObjectStoreScheme; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::{PyDict, PyString, PyTuple, PyType}; +use pyo3::{intern, IntoPyObjectExt}; +use url::Url; + +use crate::client::PyClientOptions; +use crate::config::PyConfigValue; +use crate::error::{GenericError, ParseUrlError, PyObjectStoreError, PyObjectStoreResult}; +use crate::gcp::credentials::PyGcpCredentialProvider; +use crate::path::PyPath; +use crate::retry::PyRetryConfig; +use crate::{MaybePrefixedStore, PyUrl}; + +#[derive(Debug, Clone, PartialEq)] +struct GCSConfig { + prefix: Option, + config: PyGoogleConfig, + client_options: Option, + retry_config: Option, + credential_provider: Option, +} + +impl GCSConfig { + fn bucket(&self) -> &str { + self.config + .0 + .get(&PyGoogleConfigKey(GoogleConfigKey::Bucket)) + .expect("Bucket should always exist in the config") + .as_ref() + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let args = PyTuple::empty(py).into_bound_py_any(py)?; + let kwargs = PyDict::new(py); + + if let Some(prefix) = &self.prefix { + kwargs.set_item(intern!(py, "prefix"), prefix.as_ref().as_ref())?; + } + kwargs.set_item(intern!(py, "config"), &self.config)?; + if let Some(client_options) = &self.client_options { + kwargs.set_item(intern!(py, "client_options"), client_options)?; + } + if let Some(retry_config) = &self.retry_config { + kwargs.set_item(intern!(py, "retry_config"), retry_config)?; + } + if let Some(credential_provider) = &self.credential_provider { + kwargs.set_item("credential_provider", credential_provider)?; + } + + PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) + } +} + +/// A Python-facing wrapper around a [`GoogleCloudStorage`]. +#[derive(Debug, Clone)] +#[pyclass(name = "GCSStore", frozen, subclass, from_py_object)] +pub struct PyGCSStore { + store: Arc>, + /// A config used for pickling. This must stay in sync with the underlying store's config. + config: GCSConfig, +} + +impl AsRef>> for PyGCSStore { + fn as_ref(&self) -> &Arc> { + &self.store + } +} + +impl PyGCSStore { + /// Consume self and return the underlying [`GoogleCloudStorage`]. + pub fn into_inner(self) -> Arc> { + self.store + } +} + +#[pymethods] +impl PyGCSStore { + // Create from parameters + #[new] + #[pyo3(signature = (bucket=None, *, prefix=None, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + fn new( + bucket: Option, + prefix: Option, + config: Option, + client_options: Option, + retry_config: Option, + credential_provider: Option, + kwargs: Option, + ) -> PyObjectStoreResult { + let mut builder = GoogleCloudStorageBuilder::from_env(); + let mut config = config.unwrap_or_default(); + if let Some(bucket) = bucket.clone() { + // Note: we apply the bucket to the config, not directly to the builder, so they stay + // in sync. + config.insert_raising_if_exists(GoogleConfigKey::Bucket, bucket)?; + } + let combined_config = combine_config_kwargs(Some(config), kwargs)?; + builder = combined_config.clone().apply_config(builder); + if let Some(client_options) = client_options.clone() { + builder = builder.with_client_options(client_options.into()) + } + if let Some(retry_config) = retry_config.clone() { + builder = builder.with_retry(retry_config.into()) + } + if let Some(credential_provider) = credential_provider.clone() { + builder = builder.with_credentials(Arc::new(credential_provider)); + } + Ok(Self { + store: Arc::new(MaybePrefixedStore::new(builder.build()?, prefix.clone())), + config: GCSConfig { + prefix, + config: combined_config, + client_options, + retry_config, + credential_provider, + }, + }) + } + + #[classmethod] + #[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] + pub(crate) fn from_url<'py>( + cls: &Bound<'py, PyType>, + url: PyUrl, + config: Option, + client_options: Option, + retry_config: Option, + credential_provider: Option, + kwargs: Option, + ) -> PyObjectStoreResult> { + // We manually parse the URL to find the prefix because `parse_url` does not apply the + // prefix. + let (_, prefix) = + ObjectStoreScheme::parse(url.as_ref()).map_err(object_store::Error::from)?; + let prefix: Option = if prefix.parts().count() != 0 { + Some(prefix.into()) + } else { + None + }; + let config = parse_url(config, url.as_ref())?; + + // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the + // subclass + let kwargs = kwargs.unwrap_or_default().into_pyobject(cls.py())?; + kwargs.set_item("prefix", prefix)?; + kwargs.set_item("config", config)?; + kwargs.set_item("client_options", client_options)?; + kwargs.set_item("retry_config", retry_config)?; + kwargs.set_item("credential_provider", credential_provider)?; + Ok(cls.call((), Some(&kwargs))?) + } + + fn __eq__(&self, other: &Bound) -> bool { + // Ensure we never error on __eq__ by returning false if the other object is not the same + // type + other + .cast::() + .map(|other| self.config == other.get().config) + .unwrap_or(false) + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + self.config.__getnewargs_ex__(py) + } + + fn __repr__(&self) -> String { + let bucket = self.config.bucket(); + if let Some(prefix) = &self.config.prefix { + format!( + "GCSStore(bucket=\"{}\", prefix=\"{}\")", + bucket, + prefix.as_ref() + ) + } else { + format!("GCSStore(bucket=\"{bucket}\")") + } + } + + #[getter] + fn prefix(&self) -> Option<&PyPath> { + self.config.prefix.as_ref() + } + + #[getter] + fn config(&self) -> &PyGoogleConfig { + &self.config.config + } + + #[getter] + fn client_options(&self) -> Option<&PyClientOptions> { + self.config.client_options.as_ref() + } + + #[getter] + fn credential_provider(&self) -> Option<&PyGcpCredentialProvider> { + self.config.credential_provider.as_ref() + } + + #[getter] + fn retry_config(&self) -> Option<&PyRetryConfig> { + self.config.retry_config.as_ref() + } +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub struct PyGoogleConfigKey(GoogleConfigKey); + +impl<'py> FromPyObject<'_, 'py> for PyGoogleConfigKey { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let s = obj.extract::()?.to_lowercase(); + // https://github.com/apache/arrow-rs-object-store/pull/467 + if s == "application_credentials" { + return Ok(Self(GoogleConfigKey::ApplicationCredentials)); + } + + let key = s.parse().map_err(PyObjectStoreError::ObjectStoreError)?; + Ok(Self(key)) + } +} + +impl AsRef for PyGoogleConfigKey { + fn as_ref(&self) -> &str { + self.0.as_ref() + } +} + +impl<'py> IntoPyObject<'py> for PyGoogleConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + (&self).into_pyobject(py) + } +} + +impl<'py> IntoPyObject<'py> for &PyGoogleConfigKey { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + let s = self + .0 + .as_ref() + .strip_prefix("google_") + .expect("Expected config prefix to start with google_"); + Ok(PyString::new(py, s)) + } +} + +impl From for PyGoogleConfigKey { + fn from(value: GoogleConfigKey) -> Self { + Self(value) + } +} + +impl From for GoogleConfigKey { + fn from(value: PyGoogleConfigKey) -> Self { + value.0 + } +} + +#[derive(Clone, Debug, Default, PartialEq, Eq, IntoPyObject, IntoPyObjectRef)] +pub struct PyGoogleConfig(HashMap); + +// Note: we manually impl FromPyObject instead of deriving it so that we can raise an +// UnknownConfigurationKeyError instead of a `TypeError` on invalid config keys. +// +// We also manually impl this so that we can raise on duplicate keys. +impl<'py> FromPyObject<'_, 'py> for PyGoogleConfig { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let mut slf = Self::new(); + for (key, val) in obj.extract::>()?.iter() { + slf.insert_raising_if_exists( + key.extract::()?, + val.extract::()?, + )?; + } + Ok(slf) + } +} + +impl PyGoogleConfig { + fn new() -> Self { + Self(HashMap::new()) + } + + fn apply_config(self, mut builder: GoogleCloudStorageBuilder) -> GoogleCloudStorageBuilder { + for (key, value) in self.0.into_iter() { + builder = builder.with_config(key.0, value.0); + } + builder + } + + fn merge(mut self, other: PyGoogleConfig) -> PyObjectStoreResult { + for (key, val) in other.0.into_iter() { + self.insert_raising_if_exists(key, val)?; + } + + Ok(self) + } + + fn insert_raising_if_exists( + &mut self, + key: impl Into, + val: impl Into, + ) -> PyObjectStoreResult<()> { + let key = key.into(); + let old_value = self.0.insert(key.clone(), PyConfigValue::new(val.into())); + if old_value.is_some() { + return Err(GenericError::new_err(format!( + "Duplicate key {} provided", + key.0.as_ref() + )) + .into()); + } + + Ok(()) + } + + /// Insert a key only if it does not already exist. + /// + /// This is used for URL parsing, where any parts of the URL **do not** override any + /// configuration keys passed manually. + fn insert_if_not_exists(&mut self, key: impl Into, val: impl Into) { + self.0.entry(key.into()).or_insert(PyConfigValue::new(val)); + } +} + +fn combine_config_kwargs( + config: Option, + kwargs: Option, +) -> PyObjectStoreResult { + match (config, kwargs) { + (None, None) => Ok(Default::default()), + (Some(x), None) | (None, Some(x)) => Ok(x), + (Some(config), Some(kwargs)) => Ok(config.merge(kwargs)?), + } +} + +/// Sets properties on this builder based on a URL +/// +/// This is vendored from +/// https://github.com/apache/arrow-rs/blob/f7263e253655b2ee613be97f9d00e063444d3df5/object_store/src/gcp/builder.rs#L316-L338 +/// +/// We do our own URL parsing so that we can keep our own config in sync with what is passed to the +/// underlying ObjectStore builder. Passing the URL on verbatim makes it hard because the URL +/// parsing only happens in `build()`. Then the config parameters we have don't include any config +/// applied from the URL. +fn parse_url(config: Option, parsed: &Url) -> object_store::Result { + let host = parsed + .host_str() + .ok_or_else(|| ParseUrlError::UrlNotRecognised { + url: parsed.as_str().to_string(), + })?; + let mut config = config.unwrap_or_default(); + + match parsed.scheme() { + "gs" => { + config.insert_if_not_exists(GoogleConfigKey::Bucket, host); + } + scheme => { + let scheme = scheme.to_string(); + return Err(ParseUrlError::UnknownUrlScheme { scheme }.into()); + } + } + + Ok(config) +} diff --git a/vendor/pyo3-object_store/src/http.rs b/vendor/pyo3-object_store/src/http.rs new file mode 100644 index 00000000000..10d4a7f8237 --- /dev/null +++ b/vendor/pyo3-object_store/src/http.rs @@ -0,0 +1,134 @@ +use std::sync::Arc; + +use object_store::http::{HttpBuilder, HttpStore}; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple, PyType}; +use pyo3::{intern, IntoPyObjectExt}; + +use crate::error::PyObjectStoreResult; +use crate::retry::PyRetryConfig; +use crate::{PyClientOptions, PyUrl}; + +#[derive(Debug, Clone, PartialEq)] +struct HTTPConfig { + url: PyUrl, + client_options: Option, + retry_config: Option, +} + +impl HTTPConfig { + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let args = PyTuple::new(py, vec![self.url.clone()])?.into_bound_py_any(py)?; + let kwargs = PyDict::new(py); + + if let Some(client_options) = &self.client_options { + kwargs.set_item(intern!(py, "client_options"), client_options.clone())?; + } + if let Some(retry_config) = &self.retry_config { + kwargs.set_item(intern!(py, "retry_config"), retry_config.clone())?; + } + + PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) + } +} + +/// A Python-facing wrapper around a [`HttpStore`]. +#[derive(Debug, Clone)] +#[pyclass(name = "HTTPStore", frozen, subclass, from_py_object)] +pub struct PyHttpStore { + // Note: we don't need to wrap this in a MaybePrefixedStore because the HttpStore manages its + // own prefix. + store: Arc, + /// A config used for pickling. This must stay in sync with the underlying store's config. + config: HTTPConfig, +} + +impl AsRef> for PyHttpStore { + fn as_ref(&self) -> &Arc { + &self.store + } +} + +impl PyHttpStore { + /// Consume self and return the underlying [`HttpStore`]. + pub fn into_inner(self) -> Arc { + self.store + } +} + +#[pymethods] +impl PyHttpStore { + #[new] + #[pyo3(signature = (url, *, client_options=None, retry_config=None))] + fn new( + url: PyUrl, + client_options: Option, + retry_config: Option, + ) -> PyObjectStoreResult { + let mut builder = HttpBuilder::new().with_url(url.clone()); + if let Some(client_options) = client_options.clone() { + builder = builder.with_client_options(client_options.into()) + } + if let Some(retry_config) = retry_config.clone() { + builder = builder.with_retry(retry_config.into()) + } + Ok(Self { + store: Arc::new(builder.build()?), + config: HTTPConfig { + url, + client_options, + retry_config, + }, + }) + } + + #[classmethod] + #[pyo3(signature = (url, *, client_options=None, retry_config=None))] + pub(crate) fn from_url<'py>( + cls: &Bound<'py, PyType>, + py: Python<'py>, + url: PyUrl, + client_options: Option, + retry_config: Option, + ) -> PyObjectStoreResult> { + // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the + // subclass + let kwargs = PyDict::new(py); + kwargs.set_item("url", url)?; + kwargs.set_item("client_options", client_options)?; + kwargs.set_item("retry_config", retry_config)?; + Ok(cls.call((), Some(&kwargs))?) + } + + fn __eq__(&self, other: &Bound) -> bool { + // Ensure we never error on __eq__ by returning false if the other object is not the same + // type + other + .cast::() + .map(|other| self.config == other.get().config) + .unwrap_or(false) + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + self.config.__getnewargs_ex__(py) + } + + fn __repr__(&self) -> String { + format!("HTTPStore(\"{}\")", &self.config.url.as_ref()) + } + + #[getter] + fn url(&self) -> &PyUrl { + &self.config.url + } + + #[getter] + fn client_options(&self) -> Option { + self.config.client_options.clone() + } + + #[getter] + fn retry_config(&self) -> Option { + self.config.retry_config.clone() + } +} diff --git a/vendor/pyo3-object_store/src/lib.rs b/vendor/pyo3-object_store/src/lib.rs new file mode 100644 index 00000000000..0d7d7fdfdee --- /dev/null +++ b/vendor/pyo3-object_store/src/lib.rs @@ -0,0 +1,35 @@ +#![doc = include_str!("../README.md")] +#![warn(missing_docs)] + +mod api; +mod aws; +mod azure; +mod client; +mod config; +mod credentials; +pub(crate) mod error; +mod gcp; +mod http; +mod local; +mod memory; +mod path; +mod prefix; +mod retry; +mod simple; +mod store; +mod url; + +pub use api::{register_exceptions_module, register_store_module}; +pub use aws::PyS3Store; +pub use azure::PyAzureStore; +pub use client::{PyClientConfigKey, PyClientOptions}; +pub use error::{PyObjectStoreError, PyObjectStoreResult}; +pub use gcp::PyGCSStore; +pub use http::PyHttpStore; +pub use local::PyLocalStore; +pub use memory::PyMemoryStore; +pub use path::PyPath; +pub use prefix::MaybePrefixedStore; +pub use simple::from_url; +pub use store::{AnyObjectStore, PyExternalObjectStore, PyObjectStore}; +pub use url::PyUrl; diff --git a/vendor/pyo3-object_store/src/local.rs b/vendor/pyo3-object_store/src/local.rs new file mode 100644 index 00000000000..fc36df91b33 --- /dev/null +++ b/vendor/pyo3-object_store/src/local.rs @@ -0,0 +1,143 @@ +use std::fs::create_dir_all; +use std::sync::Arc; + +use object_store::local::LocalFileSystem; +use object_store::ObjectStoreScheme; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyTuple, PyType}; +use pyo3::{intern, IntoPyObjectExt}; + +use crate::error::PyObjectStoreResult; +use crate::PyUrl; + +#[derive(Clone, Debug, PartialEq)] +struct LocalConfig { + prefix: Option, + automatic_cleanup: bool, + mkdir: bool, +} + +impl LocalConfig { + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + let args = PyTuple::new(py, vec![self.prefix.clone()])?.into_bound_py_any(py)?; + let kwargs = PyDict::new(py); + kwargs.set_item(intern!(py, "automatic_cleanup"), self.automatic_cleanup)?; + kwargs.set_item(intern!(py, "mkdir"), self.mkdir)?; + PyTuple::new(py, [args, kwargs.into_bound_py_any(py)?]) + } +} + +/// A Python-facing wrapper around a [`LocalFileSystem`]. +#[derive(Debug, Clone)] +#[pyclass(name = "LocalStore", frozen, subclass, from_py_object)] +pub struct PyLocalStore { + store: Arc, + config: LocalConfig, +} + +impl AsRef> for PyLocalStore { + fn as_ref(&self) -> &Arc { + &self.store + } +} + +impl PyLocalStore { + /// Consume self and return the underlying [`LocalFileSystem`]. + pub fn into_inner(self) -> Arc { + self.store + } +} + +#[pymethods] +impl PyLocalStore { + #[new] + #[pyo3(signature = (prefix=None, *, automatic_cleanup=false, mkdir=false))] + fn new( + prefix: Option, + automatic_cleanup: bool, + mkdir: bool, + ) -> PyObjectStoreResult { + let fs = if let Some(prefix) = &prefix { + if mkdir { + create_dir_all(prefix)?; + } + LocalFileSystem::new_with_prefix(prefix)? + } else { + LocalFileSystem::new() + }; + let fs = fs.with_automatic_cleanup(automatic_cleanup); + Ok(Self { + store: Arc::new(fs), + config: LocalConfig { + prefix, + automatic_cleanup, + mkdir, + }, + }) + } + + #[classmethod] + #[pyo3(signature = (url, *, automatic_cleanup=false, mkdir=false))] + pub(crate) fn from_url<'py>( + cls: &Bound<'py, PyType>, + url: PyUrl, + automatic_cleanup: bool, + mkdir: bool, + ) -> PyObjectStoreResult> { + let url = url.into_inner(); + let (scheme, path) = ObjectStoreScheme::parse(&url).map_err(object_store::Error::from)?; + + if !matches!(scheme, ObjectStoreScheme::Local) { + return Err(PyValueError::new_err("Not a `file://` URL").into()); + } + + // The path returned by `ObjectStoreScheme::parse` strips the initial `/`, so we join it + // onto a root + // Hopefully this also works on Windows. + let root = std::path::Path::new("/"); + let full_path = root.join(path.as_ref()); + + // Note: we pass **back** through Python so that if cls is a subclass, we instantiate the + // subclass + let kwargs = PyDict::new(cls.py()); + kwargs.set_item("prefix", full_path)?; + kwargs.set_item("automatic_cleanup", automatic_cleanup)?; + kwargs.set_item("mkdir", mkdir)?; + Ok(cls.call((), Some(&kwargs))?) + } + + fn __eq__(&self, other: &Bound) -> bool { + // Ensure we never error on __eq__ by returning false if the other object is not the same + // type + other + .cast::() + .map(|other| self.config == other.get().config) + .unwrap_or(false) + } + + fn __getnewargs_ex__<'py>(&'py self, py: Python<'py>) -> PyResult> { + self.config.__getnewargs_ex__(py) + } + + fn __repr__(&self) -> String { + if let Some(prefix) = &self.config.prefix { + format!("LocalStore(\"{}\")", prefix.display()) + } else { + "LocalStore".to_string() + } + } + + #[getter] + fn prefix<'py>(&'py self, py: Python<'py>) -> PyResult> { + // Note: returning a std::path::Path or std::path::PathBuf converts back to a Python _str_ + // not a Python _pathlib.Path_. + // So we manually convert to a pathlib.Path + if let Some(prefix) = &self.config.prefix { + let pathlib_mod = py.import(intern!(py, "pathlib"))?; + pathlib_mod.call_method1(intern!(py, "Path"), PyTuple::new(py, vec![prefix])?) + } else { + py.None().into_bound_py_any(py) + } + } +} diff --git a/vendor/pyo3-object_store/src/memory.rs b/vendor/pyo3-object_store/src/memory.rs new file mode 100644 index 00000000000..d726a6c3001 --- /dev/null +++ b/vendor/pyo3-object_store/src/memory.rs @@ -0,0 +1,47 @@ +use std::sync::Arc; + +use object_store::memory::InMemory; +use pyo3::intern; +use pyo3::prelude::*; +use pyo3::types::PyString; + +/// A Python-facing wrapper around an [`InMemory`]. +#[derive(Debug, Clone)] +#[pyclass(name = "MemoryStore", frozen, subclass, from_py_object)] +pub struct PyMemoryStore(Arc); + +impl AsRef> for PyMemoryStore { + fn as_ref(&self) -> &Arc { + &self.0 + } +} + +impl From> for PyMemoryStore { + fn from(value: Arc) -> Self { + Self(value) + } +} + +impl<'py> PyMemoryStore { + /// Consume self and return the underlying [`InMemory`]. + pub fn into_inner(self) -> Arc { + self.0 + } + + fn __repr__(&'py self, py: Python<'py>) -> &'py Bound<'py, PyString> { + intern!(py, "MemoryStore") + } +} + +#[pymethods] +impl PyMemoryStore { + #[new] + fn py_new() -> Self { + Self(Arc::new(InMemory::new())) + } + + fn __eq__(slf: Py, other: &Bound) -> bool { + // Two memory stores are equal only if they are the same object + slf.is(other) + } +} diff --git a/vendor/pyo3-object_store/src/path.rs b/vendor/pyo3-object_store/src/path.rs new file mode 100644 index 00000000000..55c373214c5 --- /dev/null +++ b/vendor/pyo3-object_store/src/path.rs @@ -0,0 +1,64 @@ +use object_store::path::Path; +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::PyString; + +/// A Python-facing wrapper around a [`Path`]. +#[derive(Clone, Debug, Default, PartialEq)] +pub struct PyPath(Path); + +impl<'py> FromPyObject<'_, 'py> for PyPath { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> Result { + let path = Path::parse(obj.extract::()?) + .map_err(|err| PyValueError::new_err(format!("Could not parse path: {err}")))?; + Ok(Self(path)) + } +} + +impl PyPath { + /// Consume self and return the underlying [`Path`]. + pub fn into_inner(self) -> Path { + self.0 + } +} + +impl<'py> IntoPyObject<'py> for PyPath { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_ref())) + } +} + +impl<'py> IntoPyObject<'py> for &PyPath { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = PyErr; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_ref())) + } +} + +impl AsRef for PyPath { + fn as_ref(&self) -> &Path { + &self.0 + } +} + +impl From for Path { + fn from(value: PyPath) -> Self { + value.0 + } +} + +impl From for PyPath { + fn from(value: Path) -> Self { + Self(value) + } +} diff --git a/vendor/pyo3-object_store/src/prefix.rs b/vendor/pyo3-object_store/src/prefix.rs new file mode 100644 index 00000000000..db84807a93e --- /dev/null +++ b/vendor/pyo3-object_store/src/prefix.rs @@ -0,0 +1,260 @@ +//! An object store wrapper handling a constant path prefix +//! This was originally vendored from https://github.com/apache/arrow-rs/blob/3bf29a2c7474e59722d885cd11fafd0dca13a28e/object_store/src/prefix.rs#L4 so that we can access the raw `T` underlying the MaybePrefixedStore. +//! It was further edited to use an `Option` internally so that we can apply a +//! `MaybePrefixedStore` to all store classes. + +use bytes::Bytes; +use futures::{stream::BoxStream, StreamExt, TryStreamExt}; +use http::Method; +use object_store::signer::Signer; +use std::borrow::Cow; +use std::future::Future; +use std::ops::Range; +use std::pin::Pin; +use std::sync::OnceLock; +use std::time::Duration; +use url::Url; + +use object_store::{path::Path, CopyOptions}; +use object_store::{ + GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, + PutMultipartOptions, PutOptions, PutPayload, PutResult, Result, +}; + +static DEFAULT_PATH: OnceLock = OnceLock::new(); + +/// Store wrapper that applies a constant prefix to all paths handled by the store. +#[derive(Debug, Clone)] +pub struct MaybePrefixedStore { + prefix: Option, + inner: T, +} + +impl std::fmt::Display for MaybePrefixedStore { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + if let Some(prefix) = self.prefix.as_ref() { + write!(f, "PrefixObjectStore({prefix})") + } else { + write!(f, "ObjectStore") + } + } +} + +impl MaybePrefixedStore { + /// Create a new instance of [`MaybePrefixedStore`] + pub fn new(store: T, prefix: Option>) -> Self { + Self { + prefix: prefix.map(|x| x.into()), + inner: store, + } + } + + /// Access the underlying T under the MaybePrefixedStore + pub fn inner(&self) -> &T { + &self.inner + } + + /// Create the full path from a path relative to prefix + fn full_path<'a>(&'a self, location: &'a Path) -> Cow<'a, Path> { + if let Some(prefix) = &self.prefix { + Cow::Owned(prefix.parts().chain(location.parts()).collect()) + } else { + Cow::Borrowed(location) + } + } + + /// Strip the constant prefix from a given path + fn strip_prefix(&self, path: Path) -> Path { + if let Some(prefix) = &self.prefix { + // Note cannot use match because of borrow checker + if let Some(suffix) = path.prefix_match(prefix) { + return suffix.collect(); + } + path + } else { + path + } + } + + /// Strip the constant prefix from a given ObjectMeta + fn strip_meta(&self, meta: ObjectMeta) -> ObjectMeta { + ObjectMeta { + last_modified: meta.last_modified, + size: meta.size, + location: self.strip_prefix(meta.location), + e_tag: meta.e_tag, + version: None, + } + } +} + +// Note: This is a relative hack to move these two functions to pure functions so they don't rely +// on the `self` lifetime. Expected to be cleaned up before merge. +// +/// Create the full path from a path relative to prefix +fn full_path<'a>(prefix: Option<&'a Path>, location: &'a Path) -> Cow<'a, Path> { + if let Some(prefix) = prefix { + Cow::Owned(prefix.parts().chain(location.parts()).collect()) + } else { + Cow::Borrowed(location) + } +} + +/// Strip the constant prefix from a given path +fn strip_prefix(prefix: Option<&Path>, path: Path) -> Path { + if let Some(prefix) = &prefix { + // Note cannot use match because of borrow checker + if let Some(suffix) = path.prefix_match(prefix) { + return suffix.collect(); + } + path + } else { + path + } +} + +/// Strip the constant prefix from a given ObjectMeta +fn strip_meta(prefix: Option<&Path>, meta: ObjectMeta) -> ObjectMeta { + ObjectMeta { + last_modified: meta.last_modified, + size: meta.size, + location: strip_prefix(prefix, meta.location), + e_tag: meta.e_tag, + version: None, + } +} +#[async_trait::async_trait] +impl ObjectStore for MaybePrefixedStore { + async fn put_opts( + &self, + location: &Path, + payload: PutPayload, + opts: PutOptions, + ) -> Result { + let full_path = self.full_path(location); + self.inner.put_opts(&full_path, payload, opts).await + } + + // Remove when updating to object_store 0.13 + #[allow(deprecated)] + async fn put_multipart_opts( + &self, + location: &Path, + opts: PutMultipartOptions, + ) -> Result> { + let full_path = self.full_path(location); + self.inner.put_multipart_opts(&full_path, opts).await + } + + async fn get_opts(&self, location: &Path, options: GetOptions) -> Result { + let full_path = self.full_path(location); + self.inner.get_opts(&full_path, options).await + } + + async fn get_ranges(&self, location: &Path, ranges: &[Range]) -> Result> { + let full_path = self.full_path(location); + self.inner.get_ranges(&full_path, ranges).await + } + + fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result> { + let prefix = self.full_path(prefix.unwrap_or(DEFAULT_PATH.get_or_init(Path::default))); + let s = self.inner.list(Some(&prefix)); + let slf_prefix = self.prefix.clone(); + s.map_ok(move |meta| strip_meta(slf_prefix.as_ref(), meta)) + .boxed() + } + + fn list_with_offset( + &self, + prefix: Option<&Path>, + offset: &Path, + ) -> BoxStream<'static, Result> { + let offset = self.full_path(offset); + let prefix = self.full_path(prefix.unwrap_or(DEFAULT_PATH.get_or_init(Path::default))); + let s = self.inner.list_with_offset(Some(&prefix), &offset); + let slf_prefix = self.prefix.clone(); + s.map_ok(move |meta| strip_meta(slf_prefix.as_ref(), meta)) + .boxed() + } + + async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result { + let prefix = self.full_path(prefix.unwrap_or(DEFAULT_PATH.get_or_init(Path::default))); + self.inner + .list_with_delimiter(Some(&prefix)) + .await + .map(|lst| ListResult { + common_prefixes: lst + .common_prefixes + .into_iter() + .map(|p| self.strip_prefix(p)) + .collect(), + objects: lst + .objects + .into_iter() + .map(|meta| self.strip_meta(meta)) + .collect(), + }) + } + + async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> { + let from_full = self.full_path(from); + let to_full = self.full_path(to); + self.inner.copy_opts(&from_full, &to_full, options).await + } + + fn delete_stream( + &self, + locations: BoxStream<'static, Result>, + ) -> BoxStream<'static, Result> { + let prefix_owned = self.prefix.clone(); + let locations = locations + .map(move |location| { + location.map(|loc| full_path(prefix_owned.as_ref(), &loc).into_owned()) + }) + .boxed(); + let prefix = self.prefix.clone(); + self.inner + .delete_stream(locations) + .map(move |location| location.map(|loc| strip_prefix(prefix.as_ref(), loc))) + .boxed() + } +} + +impl Signer for MaybePrefixedStore { + fn signed_url<'life0, 'life1, 'async_trait>( + &'life0 self, + method: Method, + path: &'life1 Path, + expires_in: Duration, + ) -> Pin> + Send + 'async_trait>> + where + 'life0: 'async_trait, + 'life1: 'async_trait, + Self: 'async_trait, + { + let full = full_path(self.prefix.as_ref(), path).into_owned(); + Box::pin(async move { self.inner.signed_url(method, &full, expires_in).await }) + } + + 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, + { + let full_paths = paths + .iter() + .map(|path| full_path(self.prefix.as_ref(), path).into_owned()) + .collect::>(); + Box::pin(async move { + self.inner + .signed_urls(method, &full_paths, expires_in) + .await + }) + } +} diff --git a/vendor/pyo3-object_store/src/retry.rs b/vendor/pyo3-object_store/src/retry.rs new file mode 100644 index 00000000000..e7079ed0313 --- /dev/null +++ b/vendor/pyo3-object_store/src/retry.rs @@ -0,0 +1,103 @@ +use std::time::Duration; + +use object_store::{BackoffConfig, RetryConfig}; +use pyo3::intern; +use pyo3::prelude::*; + +#[derive(Clone, Debug, IntoPyObject, IntoPyObjectRef, PartialEq)] +pub struct PyBackoffConfig { + #[pyo3(item)] + init_backoff: Duration, + #[pyo3(item)] + max_backoff: Duration, + #[pyo3(item)] + base: f64, +} + +impl<'py> FromPyObject<'_, 'py> for PyBackoffConfig { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let mut backoff_config = BackoffConfig::default(); + let py = obj.py(); + if let Ok(init_backoff) = obj.get_item(intern!(py, "init_backoff")) { + backoff_config.init_backoff = init_backoff.extract()?; + } + if let Ok(max_backoff) = obj.get_item(intern!(py, "max_backoff")) { + backoff_config.max_backoff = max_backoff.extract()?; + } + if let Ok(base) = obj.get_item(intern!(py, "base")) { + backoff_config.base = base.extract()?; + } + Ok(backoff_config.into()) + } +} + +impl From for BackoffConfig { + fn from(value: PyBackoffConfig) -> Self { + BackoffConfig { + init_backoff: value.init_backoff, + max_backoff: value.max_backoff, + base: value.base, + } + } +} + +impl From for PyBackoffConfig { + fn from(value: BackoffConfig) -> Self { + PyBackoffConfig { + init_backoff: value.init_backoff, + max_backoff: value.max_backoff, + base: value.base, + } + } +} + +#[derive(Clone, Debug, IntoPyObject, IntoPyObjectRef, PartialEq)] +pub struct PyRetryConfig { + #[pyo3(item)] + backoff: PyBackoffConfig, + #[pyo3(item)] + max_retries: usize, + #[pyo3(item)] + retry_timeout: Duration, +} + +impl<'py> FromPyObject<'_, 'py> for PyRetryConfig { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let mut retry_config = RetryConfig::default(); + let py = obj.py(); + if let Ok(backoff) = obj.get_item(intern!(py, "backoff")) { + retry_config.backoff = backoff.extract::()?.into(); + } + if let Ok(max_retries) = obj.get_item(intern!(py, "max_retries")) { + retry_config.max_retries = max_retries.extract()?; + } + if let Ok(retry_timeout) = obj.get_item(intern!(py, "retry_timeout")) { + retry_config.retry_timeout = retry_timeout.extract()?; + } + Ok(retry_config.into()) + } +} + +impl From for RetryConfig { + fn from(value: PyRetryConfig) -> Self { + RetryConfig { + backoff: value.backoff.into(), + max_retries: value.max_retries, + retry_timeout: value.retry_timeout, + } + } +} + +impl From for PyRetryConfig { + fn from(value: RetryConfig) -> Self { + PyRetryConfig { + backoff: value.backoff.into(), + max_retries: value.max_retries, + retry_timeout: value.retry_timeout, + } + } +} diff --git a/vendor/pyo3-object_store/src/simple.rs b/vendor/pyo3-object_store/src/simple.rs new file mode 100644 index 00000000000..1966756bad1 --- /dev/null +++ b/vendor/pyo3-object_store/src/simple.rs @@ -0,0 +1,112 @@ +use std::sync::Arc; + +use object_store::memory::InMemory; +use object_store::ObjectStoreScheme; +use pyo3::prelude::*; +use pyo3::types::{PyDict, PyType}; +use pyo3::{intern, IntoPyObjectExt}; + +use crate::error::GenericError; +use crate::retry::PyRetryConfig; +use crate::url::PyUrl; +use crate::{ + PyAzureStore, PyClientOptions, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, + PyObjectStoreResult, PyS3Store, +}; + +/// Simple construction of stores by url. +// Note: We don't extract the PyObject in the function signature because it's possible that +// AWS/Azure/Google config keys could overlap. And so we don't want to accidentally parse a config +// as an AWS config before knowing that the URL scheme is AWS. +#[pyfunction] +#[pyo3(signature = (url, *, config=None, client_options=None, retry_config=None, credential_provider=None, **kwargs))] +pub fn from_url<'py>( + py: Python<'py>, + url: PyUrl, + config: Option>, + client_options: Option, + retry_config: Option, + credential_provider: Option>, + kwargs: Option>, +) -> PyObjectStoreResult> { + let (scheme, _) = ObjectStoreScheme::parse(url.as_ref()).map_err(object_store::Error::from)?; + match scheme { + ObjectStoreScheme::AmazonS3 => PyS3Store::from_url( + &PyType::new::(py), + url, + config.map(|x| x.extract()).transpose()?, + client_options, + retry_config, + credential_provider.map(|x| x.extract()).transpose()?, + kwargs.map(|x| x.extract()).transpose()?, + ), + ObjectStoreScheme::GoogleCloudStorage => PyGCSStore::from_url( + &PyType::new::(py), + url, + config.map(|x| x.extract()).transpose()?, + client_options, + retry_config, + credential_provider.map(|x| x.extract()).transpose()?, + kwargs.map(|x| x.extract()).transpose()?, + ), + ObjectStoreScheme::MicrosoftAzure => PyAzureStore::from_url( + &PyType::new::(py), + url, + config.map(|x| x.extract()).transpose()?, + client_options, + retry_config, + credential_provider.map(|x| x.extract()).transpose()?, + kwargs.map(|x| x.extract()).transpose()?, + ), + ObjectStoreScheme::Http => { + raise_if_config_passed(config, kwargs, "http")?; + PyHttpStore::from_url( + &PyType::new::(py), + py, + url, + client_options, + retry_config, + ) + } + ObjectStoreScheme::Local => { + let mut automatic_cleanup = false; + let mut mkdir = false; + if let Some(kwargs) = kwargs { + let kwargs = kwargs.cast::()?; + if let Some(val) = kwargs.get_item(intern!(py, "automatic_cleanup"))? { + automatic_cleanup = val.extract()?; + } + if let Some(val) = kwargs.get_item(intern!(py, "mkdir"))? { + mkdir = val.extract()?; + } + } + + PyLocalStore::from_url( + &PyType::new::(py), + url, + automatic_cleanup, + mkdir, + ) + } + ObjectStoreScheme::Memory => { + raise_if_config_passed(config, kwargs, "memory")?; + let store: PyMemoryStore = Arc::new(InMemory::new()).into(); + Ok(store.into_bound_py_any(py)?) + } + scheme => Err(GenericError::new_err(format!("Unknown URL scheme {scheme:?}")).into()), + } +} + +fn raise_if_config_passed( + config: Option>, + kwargs: Option>, + scheme: &str, +) -> PyObjectStoreResult<()> { + if config.is_some() || kwargs.is_some() { + return Err(GenericError::new_err(format!( + "Cannot pass config or keyword parameters for scheme {scheme:?}" + )) + .into()); + } + Ok(()) +} diff --git a/vendor/pyo3-object_store/src/store.rs b/vendor/pyo3-object_store/src/store.rs new file mode 100644 index 00000000000..bdb48c5576e --- /dev/null +++ b/vendor/pyo3-object_store/src/store.rs @@ -0,0 +1,276 @@ +use std::sync::Arc; + +use object_store::ObjectStore; +use pyo3::exceptions::{PyRuntimeWarning, PyValueError}; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::{PyDict, PyTuple}; +use pyo3::{intern, PyTypeInfo}; + +use crate::{PyAzureStore, PyGCSStore, PyHttpStore, PyLocalStore, PyMemoryStore, PyS3Store}; + +/// A wrapper around a Rust ObjectStore instance that allows any rust-native implementation of +/// ObjectStore. +/// +/// This will only accept ObjectStore instances created from the same library. See +/// [register_store_module][crate::register_store_module]. +#[derive(Debug, Clone)] +pub struct PyObjectStore(Arc); + +impl<'py> FromPyObject<'_, 'py> for PyObjectStore { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) + } else if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) + } else if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) + } else if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) + } else if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) + } else if let Ok(store) = obj.cast::() { + Ok(Self(store.get().as_ref().clone())) + } else { + let py = obj.py(); + // Check for object-store instance from other library + let cls_name = obj + .getattr(intern!(py, "__class__"))? + .getattr(intern!(py, "__name__"))? + .extract::()?; + if [ + PyAzureStore::type_object(py).name()?.to_str()?, + PyGCSStore::type_object(py).name()?.to_str()?, + PyHttpStore::type_object(py).name()?.to_str()?, + PyLocalStore::type_object(py).name()?.to_str()?, + PyMemoryStore::type_object(py).name()?.to_str()?, + PyS3Store::type_object(py).name()?.to_str()?, + ] + .contains(&cls_name.as_str()) + { + return Err(PyValueError::new_err("You must use an object store instance exported from **the same library** as this function. They cannot be used across libraries.\nThis is because object store instances are compiled with a specific version of Rust and Python." )); + } + + Err(PyValueError::new_err(format!( + "Expected an object store instance, got {}", + obj.repr()? + ))) + } + } +} + +impl AsRef> for PyObjectStore { + fn as_ref(&self) -> &Arc { + &self.0 + } +} + +impl From for Arc { + fn from(value: PyObjectStore) -> Self { + value.0 + } +} + +impl PyObjectStore { + /// Consume self and return the underlying [`ObjectStore`]. + pub fn into_inner(self) -> Arc { + self.0 + } + + /// Consume self and return a reference-counted [`ObjectStore`]. + pub fn into_dyn(self) -> Arc { + self.0 + } + + /// Create a [`PyObjectStore`] from any pre-built [`ObjectStore`] implementation. + /// + /// This allows downstream crates (e.g. Vortex) to wrap an `object_store::ObjectStore` + /// that is not one of the built-in classes (S3/Azure/GCS/HTTP/Local/Memory) and pass + /// it to any API accepting `PyObjectStore`, such as `read_url(store=...)`. + pub fn from_store(store: Arc) -> Self { + Self(store) + } +} + +impl From> for PyObjectStore { + fn from(store: Arc) -> Self { + Self::from_store(store) + } +} + +#[derive(Debug, Clone)] +struct PyExternalObjectStoreInner(Arc); + +impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStoreInner { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + let py = obj.py(); + // Check for object-store instance from other library + let cls_name = obj + .getattr(intern!(py, "__class__"))? + .getattr(intern!(py, "__name__"))? + .extract::()?; + + if cls_name.as_str() == PyAzureStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyAzureStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + return Ok(Self(store.into_inner())); + } + + if cls_name.as_str() == PyGCSStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyGCSStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + return Ok(Self(store.into_inner())); + } + + if cls_name.as_str() == PyHttpStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyHttpStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + return Ok(Self(store.into_inner())); + } + + if cls_name.as_str() == PyLocalStore::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyLocalStore::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + return Ok(Self(store.into_inner())); + } + + if cls_name.as_str() == PyS3Store::type_object(py).name()? { + let (args, kwargs): (Bound, Bound) = obj + .call_method0(intern!(py, "__getnewargs_ex__"))? + .extract()?; + let store = PyS3Store::type_object(py) + .call(args, Some(&kwargs))? + .cast::()? + .get() + .clone(); + return Ok(Self(store.into_inner())); + } + + Err(PyValueError::new_err(format!( + "Expected an object store-compatible instance, got {}", + obj.repr()? + ))) + } +} + +/// A wrapper around a Rust [ObjectStore] instance that will extract and recreate an ObjectStore +/// instance out of a Python object. +/// +/// This will accept [ObjectStore] instances from **any** Python library exporting store classes +/// from `pyo3-object_store`. +/// +/// ## Caveats +/// +/// - This will extract the configuration of the store and **recreate** the store instance in the +/// current module. This means that no connection pooling will be reused from the original +/// library. Also, there is a slight overhead to this as configuration parsing will need to +/// happen from scratch. +/// +/// This will work best when the store is created once and used multiple times. +/// +/// - This relies on the public Python API (`__getnewargs_ex__` and `__init__`) of the store +/// classes to extract the configuration. If the public API changes in a non-backwards compatible +/// way, this store conversion may fail. +/// +/// - While this reuses `__getnewargs_ex__` (from the pickle implementation) to extract arguments +/// to pass into `__init__`, it does not actually _use_ pickle, and so even non-pickleable +/// credential providers should work. +/// +/// - This will not work for `PyMemoryStore` because we can't clone the internal state of the +/// store. +#[derive(Debug, Clone)] +pub struct PyExternalObjectStore(PyExternalObjectStoreInner); + +impl From for Arc { + fn from(value: PyExternalObjectStore) -> Self { + value.0 .0 + } +} + +impl PyExternalObjectStore { + /// Consume self and return a reference-counted [`ObjectStore`]. + pub fn into_dyn(self) -> Arc { + self.into() + } +} + +impl<'py> FromPyObject<'_, 'py> for PyExternalObjectStore { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, pyo3::PyAny>) -> PyResult { + match obj.extract() { + Ok(inner) => { + #[cfg(feature = "external-store-warning")] + { + let py = obj.py(); + + let warnings_mod = py.import(intern!(py, "warnings"))?; + let warning = PyRuntimeWarning::new_err( + "Successfully reconstructed a store defined in another Python module. Connection pooling will not be shared across store instances.", + ); + let args = PyTuple::new(py, vec![warning])?; + warnings_mod.call_method1(intern!(py, "warn"), args)?; + } + Ok(Self(inner)) + } + Err(err) => Err(err), + } + } +} + +/// A convenience wrapper around native and external ObjectStore instances. +/// +/// Note that there may be performance differences between accepted variants here. If you wish to +/// only permit the highest-performance stores, use [`PyObjectStore`] directly as the parameter in +/// your signature. +#[derive(FromPyObject)] +pub enum AnyObjectStore { + /// A wrapper around a [`PyObjectStore`]. + PyObjectStore(PyObjectStore), + /// A wrapper around a [`PyExternalObjectStore`]. + PyExternalObjectStore(PyExternalObjectStore), +} + +impl From for Arc { + fn from(value: AnyObjectStore) -> Self { + match value { + AnyObjectStore::PyObjectStore(store) => store.into(), + AnyObjectStore::PyExternalObjectStore(store) => store.into(), + } + } +} + +impl AnyObjectStore { + /// Consume self and return a reference-counted [`ObjectStore`]. + pub fn into_dyn(self) -> Arc { + self.into() + } +} diff --git a/vendor/pyo3-object_store/src/url.rs b/vendor/pyo3-object_store/src/url.rs new file mode 100644 index 00000000000..a67b0000630 --- /dev/null +++ b/vendor/pyo3-object_store/src/url.rs @@ -0,0 +1,64 @@ +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3::pybacked::PyBackedStr; +use pyo3::types::PyString; +use pyo3::FromPyObject; +use url::Url; + +/// A wrapper around [`url::Url`] that implements [`FromPyObject`]. +#[derive(Debug, Clone, PartialEq)] +pub struct PyUrl(Url); + +impl PyUrl { + /// Create a new PyUrl from a [Url] + pub fn new(url: Url) -> Self { + Self(url) + } + + /// Consume self and return the underlying [Url] + pub fn into_inner(self) -> Url { + self.0 + } +} + +impl<'py> FromPyObject<'_, 'py> for PyUrl { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> Result { + let s = obj.extract::()?; + let url = Url::parse(&s).map_err(|err| PyValueError::new_err(err.to_string()))?; + Ok(Self(url)) + } +} + +impl<'py> IntoPyObject<'py> for PyUrl { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = std::convert::Infallible; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_str())) + } +} + +impl<'py> IntoPyObject<'py> for &PyUrl { + type Target = PyString; + type Output = Bound<'py, PyString>; + type Error = std::convert::Infallible; + + fn into_pyobject(self, py: Python<'py>) -> Result { + Ok(PyString::new(py, self.0.as_str())) + } +} + +impl AsRef for PyUrl { + fn as_ref(&self) -> &Url { + &self.0 + } +} + +impl From for String { + fn from(value: PyUrl) -> Self { + value.0.into() + } +} diff --git a/vendor/pyo3-object_store/type-hints/__init__.pyi b/vendor/pyo3-object_store/type-hints/__init__.pyi new file mode 100644 index 00000000000..39eb60ec542 --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/__init__.pyi @@ -0,0 +1,218 @@ +# TODO: move to reusable types package +import sys +from collections.abc import Callable +from pathlib import Path +from typing import Any, overload + +from ._aws import S3Config as S3Config +from ._aws import S3Credential as S3Credential +from ._aws import S3CredentialProvider as S3CredentialProvider +from ._aws import S3Store as S3Store +from ._azure import AzureAccessKey as AzureAccessKey +from ._azure import AzureBearerToken as AzureBearerToken +from ._azure import AzureConfig as AzureConfig +from ._azure import AzureCredential as AzureCredential +from ._azure import AzureCredentialProvider as AzureCredentialProvider +from ._azure import AzureSASToken as AzureSASToken +from ._azure import AzureStore as AzureStore +from ._client import ClientConfig as ClientConfig +from ._gcs import GCSConfig as GCSConfig +from ._gcs import GCSCredential as GCSCredential +from ._gcs import GCSCredentialProvider as GCSCredentialProvider +from ._gcs import GCSStore as GCSStore +from ._http import HTTPStore as HTTPStore +from ._retry import BackoffConfig as BackoffConfig +from ._retry import RetryConfig as RetryConfig + +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import TypeAlias + +if sys.version_info >= (3, 11): + from typing import Self, Unpack +else: + from typing_extensions import Self, Unpack + +@overload +def from_url( + url: str, + *, + config: S3Config | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: S3CredentialProvider | None = None, + **kwargs: Unpack[S3Config], +) -> ObjectStore: ... +@overload +def from_url( + url: str, + *, + config: GCSConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: GCSCredentialProvider | None = None, + **kwargs: Unpack[GCSConfig], +) -> ObjectStore: ... +@overload +def from_url( + url: str, + *, + config: AzureConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: AzureCredentialProvider | None = None, + **kwargs: Unpack[AzureConfig], +) -> ObjectStore: ... +@overload +def from_url( + url: str, + *, + config: None = None, + client_options: None = None, + retry_config: None = None, + automatic_cleanup: bool = False, + mkdir: bool = False, +) -> ObjectStore: ... +def from_url( # type: ignore[misc] # docstring in pyi file + url: str, + *, + config: S3Config | GCSConfig | AzureConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: Callable | None = None, + **kwargs: Any, +) -> ObjectStore: + """Easy construction of store by URL, identifying the relevant store. + + This will defer to a store-specific `from_url` constructor based on the provided + `url`. E.g. passing `"s3://bucket/path"` will defer to + [`S3Store.from_url`][obstore.store.S3Store.from_url]. + + Supported formats: + + - `file:///path/to/my/file` -> [`LocalStore`][obstore.store.LocalStore] + - `memory:///` -> [`MemoryStore`][obstore.store.MemoryStore] + - `s3://bucket/path` -> [`S3Store`][obstore.store.S3Store] (also supports `s3a`) + - `gs://bucket/path` -> [`GCSStore`][obstore.store.GCSStore] + - `az://account/container/path` -> [`AzureStore`][obstore.store.AzureStore] (also + supports `adl`, `azure`, `abfs`, `abfss`) + - `http://mydomain/path` -> [`HTTPStore`][obstore.store.HTTPStore] + - `https://mydomain/path` -> [`HTTPStore`][obstore.store.HTTPStore] + + There are also special cases for AWS and Azure for `https://{host?}/path` paths: + + - `dfs.core.windows.net`, `blob.core.windows.net`, `dfs.fabric.microsoft.com`, + `blob.fabric.microsoft.com` -> [`AzureStore`][obstore.store.AzureStore] + - `amazonaws.com` -> [`S3Store`][obstore.store.S3Store] + - `r2.cloudflarestorage.com` -> [`S3Store`][obstore.store.S3Store] + + !!! note + For best static typing, use the constructors on individual store classes + directly. + + Args: + url: well-known storage URL. + + Keyword Args: + config: per-store Configuration. Values in this config will override values + inferred from the url. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom credentials to the underlying store classes. + kwargs: per-store configuration passed down to store-specific builders. + + """ + +class LocalStore: + """An ObjectStore interface to local filesystem storage. + + Can optionally be created with a directory prefix. + + ```py + from pathlib import Path + + store = LocalStore() + store = LocalStore(prefix="/path/to/directory") + store = LocalStore(prefix=Path(".")) + ``` + """ + + def __init__( + self, + prefix: str | Path | None = None, + *, + automatic_cleanup: bool = False, + mkdir: bool = False, + ) -> None: + """Create a new LocalStore. + + Args: + prefix: Use the specified prefix applied to all paths. Defaults to `None`. + + Keyword Args: + automatic_cleanup: if `True`, enables automatic cleanup of empty directories + when deleting files. Defaults to False. + mkdir: if `True` and `prefix` is not `None`, the directory at `prefix` will + attempt to be created. Note that this root directory will not be cleaned + up, even if `automatic_cleanup` is `True`. + + """ + @classmethod + def from_url( + cls, + url: str, + *, + automatic_cleanup: bool = False, + mkdir: bool = False, + ) -> Self: + """Construct a new LocalStore from a `file://` URL. + + **Examples:** + + Construct a new store pointing to the root of your filesystem: + ```py + url = "file:///" + store = LocalStore.from_url(url) + ``` + + Construct a new store with a directory prefix: + ```py + url = "file:///Users/kyle/" + store = LocalStore.from_url(url) + ``` + """ + + def __eq__(self, value: object) -> bool: ... + def __getnewargs_ex__(self): ... + @property + def prefix(self) -> Path | None: + """Get the prefix applied to all operations in this store, if any.""" + +class MemoryStore: + """A fully in-memory implementation of ObjectStore. + + Create a new in-memory store: + ```py + store = MemoryStore() + ``` + """ + + def __init__(self) -> None: ... + +ObjectStore: TypeAlias = ( + AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore +) +"""All supported ObjectStore implementations. + +!!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import ObjectStore + ``` +""" diff --git a/vendor/pyo3-object_store/type-hints/_aws.pyi b/vendor/pyo3-object_store/type-hints/_aws.pyi new file mode 100644 index 00000000000..eb625a793bd --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/_aws.pyi @@ -0,0 +1,623 @@ +import sys +from collections.abc import Coroutine +from datetime import datetime +from typing import Any, Literal, Protocol, TypedDict + +from ._client import ClientConfig +from ._retry import RetryConfig + +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import TypeAlias + +if sys.version_info >= (3, 11): + from typing import NotRequired, Self, Unpack +else: + from typing_extensions import NotRequired, Self, Unpack + +# To update s3 region list: +# import pandas as pd # noqa: ERA001 +# url = "https://docs.aws.amazon.com/AmazonRDS/latest/UserGuide/Concepts.RegionsAndAvailabilityZones.html" # noqa: ERA001 +# result = pd.read_html(url) # noqa: ERA001 +# sorted(result[0]["Region"]) # noqa: ERA001 +S3Regions: TypeAlias = Literal[ + "af-south-1", + "ap-east-1", + "ap-northeast-1", + "ap-northeast-2", + "ap-northeast-3", + "ap-south-1", + "ap-south-2", + "ap-southeast-1", + "ap-southeast-2", + "ap-southeast-3", + "ap-southeast-4", + "ap-southeast-5", + "ap-southeast-7", + "ca-central-1", + "ca-west-1", + "eu-central-1", + "eu-central-2", + "eu-north-1", + "eu-south-1", + "eu-south-2", + "eu-west-1", + "eu-west-2", + "eu-west-3", + "il-central-1", + "me-central-1", + "me-south-1", + "mx-central-1", + "sa-east-1", + "us-east-1", + "us-east-2", + "us-gov-east-1", + "us-gov-west-1", + "us-west-1", + "us-west-2", +] +"""AWS regions.""" + +S3ChecksumAlgorithm: TypeAlias = Literal["SHA256", "sha256", "CRC64NVME", "crc64nvme"] +"""S3 Checksum algorithms + +From https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html#using-additional-checksums +""" + +S3EncryptionAlgorithm: TypeAlias = Literal[ + "AES256", + "aws:kms", + "aws:kms:dsse", + "sse-c", +] + +class S3Config(TypedDict, total=False): + """Configuration parameters for S3Store. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import S3Config + ``` + """ + + access_key_id: str + """AWS Access Key. + + **Environment variable**: `AWS_ACCESS_KEY_ID`. + """ + + bucket: str + """Bucket name (required). + + **Environment variables**: + + - `AWS_BUCKET` + - `AWS_BUCKET_NAME` + """ + + checksum_algorithm: S3ChecksumAlgorithm | str + """ + Sets the [checksum algorithm] which has to be used for object integrity check during upload. + + [checksum algorithm]: https://docs.aws.amazon.com/AmazonS3/latest/userguide/checking-object-integrity.html + + **Environment variable**: `AWS_CHECKSUM_ALGORITHM`. + """ + + conditional_put: str + """Configure how to provide conditional put support + + Supported values: + + - `"etag"` (default): Supported for S3-compatible stores that support conditional + put using the standard [HTTP precondition] headers `If-Match` and + `If-None-Match`. + + [HTTP precondition]: https://datatracker.ietf.org/doc/html/rfc9110#name-preconditions + + **Environment variable**: `AWS_CONDITIONAL_PUT`. + """ + + container_credentials_relative_uri: str + """Set the container credentials relative URI when used in ECS. + + + + **Environment variable**: `AWS_CONTAINER_CREDENTIALS_RELATIVE_URI`. + + Example: `/v2/credentials/abc123` + """ + + container_credentials_full_uri: str + """Set the container credentials full URI when used in EKS. + + + + **Environment variable**: `AWS_CONTAINER_CREDENTIALS_FULL_URI`. + + Example: `http://169.254.170.2/v2/credentials/abc123` + """ + + container_authorization_token_file: str + """Set the authorization token in plain text when used in EKS to authenticate with ContainerCredentialsFullUri. + + + + **Environment variable**: `AWS_CONTAINER_AUTHORIZATION_TOKEN_FILE`. + + Example: `/var/run/secrets/eks.amazonaws.com/serviceaccount/token` + """ + + copy_if_not_exists: Literal["multipart"] | str + """Configure how to provide "copy if not exists". + + Supported values: + + - `"multipart"`: + + Native Amazon S3 supports copy if not exists through a multipart upload + where the upload copies an existing object and is completed only if the + new object does not already exist. + + !!! warning + When using this mode, `copy_if_not_exists` does not copy tags + or attributes from the source object. + + !!! warning + When using this mode, `copy_if_not_exists` makes only a best + effort attempt to clean up the multipart upload if the copy operation + fails. Consider using a lifecycle rule to automatically clean up + abandoned multipart uploads. + + - `"header::"`: + + Some S3-compatible stores, such as Cloudflare R2, support copy if not exists + semantics through custom headers. + + If set, `copy_if_not_exists` will perform a normal copy operation with the + provided header pair, and expect the store to fail with `412 Precondition + Failed` if the destination file already exists. + + For example `header: cf-copy-destination-if-none-match: *`, would set + the header `cf-copy-destination-if-none-match` to `*`. + + - `"header-with-status:::"`: + + The same as the header variant above but allows custom status code checking, for + object stores that return values other than 412. + + **Environment variable**: `AWS_COPY_IF_NOT_EXISTS`. + """ + + default_region: S3Regions | str + """Default region. + + **Environment variable**: `AWS_DEFAULT_REGION`. + """ + + disable_tagging: bool + """Disable tagging objects. This can be desirable if not supported by the backing store. + + **Environment variable**: `AWS_DISABLE_TAGGING`. + """ + + endpoint: str + """The endpoint for communicating with AWS S3. + + Defaults to the [region endpoint]. + + For example, this might be set to `"http://localhost:4566:` for testing against a + localstack instance. + + The `endpoint` field should be consistent with `with_virtual_hosted_style_request`, + i.e. if `virtual_hosted_style_request` is set to `True` then `endpoint` should have + the bucket name included. + + By default, only HTTPS schemes are enabled. To connect to an HTTP endpoint, enable + `allow_http` in the client options. + + [region endpoint]: https://docs.aws.amazon.com/general/latest/gr/s3.html + + **Environment variables**: + + - `AWS_ENDPOINT` + - `AWS_ENDPOINT_URL` + - `AWS_ENDPOINT_URL_S3`. + """ + + imdsv1_fallback: bool + """Fall back to ImdsV1. + + By default instance credentials will only be fetched over [IMDSv2], as AWS + recommends against having IMDSv1 enabled on EC2 instances as it is vulnerable to + [SSRF attack] + + However, certain deployment environments, such as those running old versions of + kube2iam, may not support IMDSv2. This option will enable automatic fallback to + using IMDSv1 if the token endpoint returns a 403 error indicating that IMDSv2 is not + supported. + + This option has no effect if not using instance credentials. + + [IMDSv2]: https://docs.aws.amazon.com/AWSEC2/latest/UserGuide/configuring-instance-metadata-service.html + [SSRF attack]: https://aws.amazon.com/blogs/security/defense-in-depth-open-firewalls-reverse-proxies-ssrf-vulnerabilities-ec2-instance-metadata-service/ + + **Environment variable**: `AWS_IMDSV1_FALLBACK`. + """ + + metadata_endpoint: str + """Set the [instance metadata endpoint], used primarily within AWS EC2. + + This defaults to the IPv4 endpoint: `http://169.254.169.254`. One can alternatively + use the IPv6 endpoint `http://fd00:ec2::254`. + + **Environment variable**: `AWS_METADATA_ENDPOINT`. + """ + + region: S3Regions | str + """The region, defaults to `us-east-1` + + **Environment variable**: `AWS_REGION`. + """ + + request_payer: bool | Literal["requester", "REQUESTER"] + """Enable operations on requester-pays buckets. + + Config option can be set to `True` or `"requester"` (case insensitive). + + + + **Environment variable**: `AWS_REQUEST_PAYER` (can be set to either `"requester"` or `"true"`). + """ + role_arn: str + """Role ARN to assume when using web identity token. + + **Environment variable**: `AWS_ROLE_ARN`. + + Example: `arn:aws:iam::123456789012:role/MyWebIdentityRole` + """ + + role_session_name: str + """Session name for web identity role assumption. + + **Environment variable**: `AWS_ROLE_SESSION_NAME`. + """ + + s3_express: bool + """Enable Support for S3 Express One Zone. + + **Environment variable**: `AWS_S3_EXPRESS`. + """ + + secret_access_key: str + """Secret Access Key. + + **Environment variable**: `AWS_SECRET_ACCESS_KEY`. + """ + + server_side_encryption: S3EncryptionAlgorithm | str + """Type of encryption to use. + + If set, must be one of: + + - `"AES256"` (SSE-S3) + - `"aws:kms"` (SSE-KMS) + - `"aws:kms:dsse"` (DSSE-KMS) + - `"sse-c"` + + **Environment variable**: `AWS_SERVER_SIDE_ENCRYPTION`. + """ + + session_token: str + """Token to use for requests (passed to underlying provider). + + **Environment variables**: + + - `AWS_SESSION_TOKEN` + - `AWS_TOKEN` + """ + + skip_signature: bool + """If `True`, S3Store will not fetch credentials and will not sign requests. + + This can be useful when interacting with public S3 buckets that deny authorized requests. + + **Environment variable**: `AWS_SKIP_SIGNATURE`. + """ + + sse_bucket_key_enabled: bool + """Set whether to enable bucket key for server side encryption. + + This overrides the bucket default setting for bucket keys. + + - When `False`, each object is encrypted with a unique data key. + - When `True`, a single data key is used for the entire bucket, + reducing overhead of encryption. + + **Environment variable**: `AWS_SSE_BUCKET_KEY_ENABLED`. + """ + + sse_customer_key_base64: str + """ + The base64 encoded, 256-bit customer encryption key to use for server-side + encryption. If set, the server side encryption config value must be `"sse-c"`. + + **Environment variable**: `AWS_SSE_CUSTOMER_KEY_BASE64`. + """ + + sse_kms_key_id: str + """ + The KMS key ID to use for server-side encryption. + + If set, the server side encryption config value must be `"aws:kms"` or `"aws:kms:dsse"`. + + **Environment variable**: `AWS_SSE_KMS_KEY_ID`. + + Example: `arn:aws:kms:us-east-1:123456789012:key/abcd-1234-efgh-5678` + """ + + sts_endpoint: str + """Custom STS endpoint for web identity token exchange. + + Defaults to `https://sts.{region}.amazonaws.com`. + + **Environment variable**: `AWS_STS_ENDPOINT`. + """ + + unsigned_payload: bool + """Avoid computing payload checksum when calculating signature. + + See [unsigned payload option](https://docs.aws.amazon.com/AmazonS3/latest/API/sig-v4-header-based-auth.html). + + - `False` (default): Signed payload option is used, where the checksum for the request body is computed and included when constructing a canonical request. + - `True`: Unsigned payload option is used. `UNSIGNED-PAYLOAD` literal is included when constructing a canonical request, + + **Environment variable**: `AWS_UNSIGNED_PAYLOAD`. + """ + + virtual_hosted_style_request: bool + """If virtual hosted style request has to be used. + + If `virtual_hosted_style_request` is: + + - `False` (default): Path style request is used + - `True`: Virtual hosted style request is used + + If the `endpoint` is provided then it should be consistent with + `virtual_hosted_style_request`. i.e. if `virtual_hosted_style_request` is set to + `True` then `endpoint` should have bucket name included. + + **Environment variable**: `AWS_VIRTUAL_HOSTED_STYLE_REQUEST`. + """ + + web_identity_token_file: str + """Web identity token file path for AssumeRoleWithWebIdentity. + + **Environment variable**: `AWS_WEB_IDENTITY_TOKEN_FILE`. + + Example: `/var/run/secrets/eks.amazonaws.com/serviceaccount/token` + """ + +class S3Credential(TypedDict): + """An S3 credential. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import S3Credential + ``` + """ + + access_key_id: str + """AWS access key ID.""" + + secret_access_key: str + """AWS secret access key""" + + token: NotRequired[str | None] + """AWS token.""" + + expires_at: datetime | None + """Expiry datetime of credential. The datetime should have time zone set. + + If None, the credential will never expire. + """ + +class S3CredentialProvider(Protocol): + """A type hint for a synchronous or asynchronous callback to provide custom S3 credentials. + + This should be passed into the `credential_provider` parameter of `S3Store`. + + **Examples:** + + Return static credentials that don't expire: + ```py + def get_credentials() -> S3Credential: + return { + "access_key_id": "...", + "secret_access_key": "...", + "token": None, + "expires_at": None, + } + ``` + + Return static credentials that are valid for 5 minutes: + ```py + from datetime import datetime, timedelta, UTC + + async def get_credentials() -> S3Credential: + return { + "access_key_id": "...", + "secret_access_key": "...", + "token": None, + "expires_at": datetime.now(UTC) + timedelta(minutes=5), + } + ``` + + A class-based credential provider with state: + + ```py + from __future__ import annotations + + from typing import TYPE_CHECKING + + import boto3 + import botocore.credentials + + if TYPE_CHECKING: + from obstore.store import S3Credential + + + class Boto3CredentialProvider: + credentials: botocore.credentials.Credentials + + def __init__(self, session: boto3.session.Session) -> None: + credentials = session.get_credentials() + if credentials is None: + raise ValueError("Received None from session.get_credentials") + + self.credentials = credentials + + def __call__(self) -> S3Credential: + frozen_credentials = self.credentials.get_frozen_credentials() + return { + "access_key_id": frozen_credentials.access_key, + "secret_access_key": frozen_credentials.secret_key, + "token": frozen_credentials.token, + "expires_at": None, + } + ``` + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import S3CredentialProvider + ``` + """ + + def __call__(self) -> S3Credential | Coroutine[Any, Any, S3Credential]: + """Return an `S3Credential`.""" + +class S3Store: + """Interface to an Amazon S3 bucket. + + All constructors will check for environment variables. Refer to + [`S3Config`][obstore.store.S3Config] for valid environment variables. + + **Examples**: + + **Using requester-pays buckets**: + + Pass `request_payer=True` as a keyword argument or have `AWS_REQUESTER_PAYS=True` + set in the environment. + + **Anonymous requests**: + + Pass `skip_signature=True` as a keyword argument or have `AWS_SKIP_SIGNATURE=True` + set in the environment. + """ + + def __init__( # type: ignore[misc] # Overlap between argument names and ** TypedDict items: "bucket" + self, + bucket: str | None = None, + *, + prefix: str | None = None, + config: S3Config | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: S3CredentialProvider | None = None, + **kwargs: Unpack[S3Config], # type: ignore # noqa: PGH003 (bucket key overlaps with positional arg) + ) -> None: + """Create a new S3Store. + + Args: + bucket: The AWS bucket to use. + + Keyword Args: + prefix: A prefix within the bucket to use for all operations. + config: AWS configuration. Values in this config will override values inferred from the environment. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom S3 credentials. + kwargs: AWS configuration values. Supports the same values as `config`, but as named keyword args. + + Returns: + S3Store + + """ + @classmethod + def from_url( + cls, + url: str, + *, + config: S3Config | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: S3CredentialProvider | None = None, + **kwargs: Unpack[S3Config], + ) -> Self: + """Parse available connection info from a well-known storage URL. + + Any path on the URL will be assigned as the `prefix` for the store. So if you + pass `s3://bucket/path/to/directory`, the store will be created with a prefix of + `path/to/directory`, and all further operations will use paths relative to that + prefix. + + The supported url schemes are: + + - `s3:///` + - `s3a:///` + - `https://s3..amazonaws.com/` + - `https://.s3..amazonaws.com` + - `https://ACCOUNT_ID.r2.cloudflarestorage.com/bucket` + + Args: + url: well-known storage URL. + + Keyword Args: + config: AWS Configuration. Values in this config will override values inferred from the url. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom S3 credentials. + kwargs: AWS configuration values. Supports the same values as `config`, but as named keyword args. + + + Returns: + S3Store + + """ + + def __eq__(self, value: object) -> bool: ... + def __getnewargs_ex__(self): ... + @property + def prefix(self) -> str | None: + """Get the prefix applied to all operations in this store, if any.""" + @property + def config(self) -> S3Config: + """Get the underlying S3 config parameters.""" + @property + def client_options(self) -> ClientConfig | None: + """Get the store's client configuration.""" + @property + def credential_provider(self) -> S3CredentialProvider | None: + """Get the store's credential provider.""" + @property + def retry_config(self) -> RetryConfig | None: + """Get the store's retry configuration.""" diff --git a/vendor/pyo3-object_store/type-hints/_azure.pyi b/vendor/pyo3-object_store/type-hints/_azure.pyi new file mode 100644 index 00000000000..c27421551d8 --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/_azure.pyi @@ -0,0 +1,422 @@ +import sys +from collections.abc import Coroutine +from datetime import datetime +from typing import Any, Protocol, TypedDict + +from ._client import ClientConfig +from ._retry import RetryConfig + +if sys.version_info >= (3, 10): + from typing import TypeAlias +else: + from typing_extensions import TypeAlias + +if sys.version_info >= (3, 11): + from typing import Self, Unpack +else: + from typing_extensions import Self, Unpack + +class AzureConfig(TypedDict, total=False): + """Configuration parameters for AzureStore. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import AzureConfig + ``` + """ + + account_name: str + """The name of the azure storage account. (Required.) + + **Environment variable**: `AZURE_STORAGE_ACCOUNT_NAME`. + """ + account_key: str + """Master key for accessing storage account. + + **Environment variables**: + + - `AZURE_STORAGE_ACCOUNT_KEY` + - `AZURE_STORAGE_ACCESS_KEY` + - `AZURE_STORAGE_MASTER_KEY` + """ + client_id: str + """The client id for use in client secret or k8s federated credential flow. + + **Environment variables**: + + - `AZURE_STORAGE_CLIENT_ID` + - `AZURE_CLIENT_ID` + """ + client_secret: str + """The client secret for use in client secret flow. + + **Environment variables**: + + - `AZURE_STORAGE_CLIENT_SECRET` + - `AZURE_CLIENT_SECRET` + """ + tenant_id: str + """The tenant id for use in client secret or k8s federated credential flow. + + **Environment variables**: + + - `AZURE_STORAGE_TENANT_ID` + - `AZURE_STORAGE_AUTHORITY_ID` + - `AZURE_TENANT_ID` + - `AZURE_AUTHORITY_ID` + """ + authority_host: str + """Sets an alternative authority host for OAuth based authorization. + + Defaults to `https://login.microsoftonline.com`. + + Common hosts for azure clouds are: + + - Azure China: `"https://login.chinacloudapi.cn"` + - Azure Germany: `"https://login.microsoftonline.de"` + - Azure Government: `"https://login.microsoftonline.us"` + - Azure Public: `"https://login.microsoftonline.com"` + + **Environment variables**: + + - `AZURE_STORAGE_AUTHORITY_HOST` + - `AZURE_AUTHORITY_HOST` + """ + sas_key: str + """ + Shared access signature. + + The signature is expected to be percent-encoded, `much `like they are provided in + the azure storage explorer or azure portal. + + **Environment variables**: + + - `AZURE_STORAGE_SAS_KEY` + - `AZURE_STORAGE_SAS_TOKEN` + """ + token: str + """A static bearer token to be used for authorizing requests. + + **Environment variable**: `AZURE_STORAGE_TOKEN`. + """ + use_emulator: bool + """Set if the Azure emulator should be used (defaults to `False`). + + **Environment variable**: `AZURE_STORAGE_USE_EMULATOR`. + """ + use_fabric_endpoint: bool + """Set if Microsoft Fabric url scheme should be used (defaults to `False`). + + When disabled the url scheme used is `https://{account}.blob.core.windows.net`. + When enabled the url scheme used is `https://{account}.dfs.fabric.microsoft.com`. + + !!! note + + `endpoint` will take precedence over this option. + """ + endpoint: str + """Override the endpoint used to communicate with blob storage. + + Defaults to `https://{account}.blob.core.windows.net`. + + By default, only HTTPS schemes are enabled. To connect to an HTTP endpoint, enable + `allow_http` in the client options. + + **Environment variables**: + + - `AZURE_STORAGE_ENDPOINT` + - `AZURE_ENDPOINT` + """ + msi_endpoint: str + """Endpoint to request a imds managed identity token. + + **Environment variables**: + + - `AZURE_MSI_ENDPOINT` + - `AZURE_IDENTITY_ENDPOINT` + """ + object_id: str + """Object id for use with managed identity authentication. + + **Environment variable**: `AZURE_OBJECT_ID`. + """ + msi_resource_id: str + """Msi resource id for use with managed identity authentication. + + **Environment variable**: `AZURE_MSI_RESOURCE_ID`. + """ + federated_token_file: str + """Sets a file path for acquiring azure federated identity token in k8s. + + Requires `client_id` and `tenant_id` to be set. + + **Environment variable**: `AZURE_FEDERATED_TOKEN_FILE`. + """ + use_azure_cli: bool + """Set if the Azure Cli should be used for acquiring access token. + + . + + **Environment variable**: `AZURE_USE_AZURE_CLI`. + """ + skip_signature: bool + """If enabled, `AzureStore` will not fetch credentials and will not sign requests. + + This can be useful when interacting with public containers. + + **Environment variable**: `AZURE_SKIP_SIGNATURE`. + """ + container_name: str + """Container name. + + **Environment variable**: `AZURE_CONTAINER_NAME`. + """ + disable_tagging: bool + """If set to `True` will ignore any tags provided to uploads. + + **Environment variable**: `AZURE_DISABLE_TAGGING`. + """ + fabric_token_service_url: str + """Service URL for Fabric OAuth2 authentication. + + **Environment variable**: `AZURE_FABRIC_TOKEN_SERVICE_URL`. + """ + fabric_workload_host: str + """Workload host for Fabric OAuth2 authentication. + + **Environment variable**: `AZURE_FABRIC_WORKLOAD_HOST`. + """ + fabric_session_token: str + """Session token for Fabric OAuth2 authentication. + + **Environment variable**: `AZURE_FABRIC_SESSION_TOKEN`. + """ + fabric_cluster_identifier: str + """Cluster identifier for Fabric OAuth2 authentication. + + **Environment variable**: `AZURE_FABRIC_CLUSTER_IDENTIFIER`. + """ + +class AzureAccessKey(TypedDict): + """A shared Azure Storage Account Key. + + + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import AzureAccessKey + ``` + """ + + access_key: str + """Access key value.""" + + expires_at: datetime | None + """Expiry datetime of credential. The datetime should have time zone set. + + If None, the credential will never expire. + """ + +class AzureSASToken(TypedDict): + """A shared access signature. + + + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import AzureSASToken + ``` + """ + + sas_token: str | list[tuple[str, str]] + """SAS token.""" + + expires_at: datetime | None + """Expiry datetime of credential. The datetime should have time zone set. + + If None, the credential will never expire. + """ + +class AzureBearerToken(TypedDict): + """An authorization token. + + + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import AzureBearerToken + ``` + """ + + token: str + """Bearer token.""" + + expires_at: datetime | None + """Expiry datetime of credential. The datetime should have time zone set. + + If None, the credential will never expire. + """ + +AzureCredential: TypeAlias = AzureAccessKey | AzureSASToken | AzureBearerToken +"""A type alias for supported azure credentials to be returned from `AzureCredentialProvider`. + +!!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import AzureCredential + ``` +""" + +class AzureCredentialProvider(Protocol): + """A type hint for a synchronous or asynchronous callback to provide custom Azure credentials. + + This should be passed into the `credential_provider` parameter of `AzureStore`. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import AzureCredentialProvider + ``` + """ + + def __call__(self) -> AzureCredential | Coroutine[Any, Any, AzureCredential]: + """Return an `AzureCredential`.""" + +class AzureStore: + """Interface to a Microsoft Azure Blob Storage container. + + All constructors will check for environment variables. Refer to + [`AzureConfig`][obstore.store.AzureConfig] for valid environment variables. + """ + + def __init__( # type: ignore[misc] # Overlap between argument names and ** TypedDict items: "container_name" + self, + container_name: str | None = None, + *, + prefix: str | None = None, + config: AzureConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: AzureCredentialProvider | None = None, + **kwargs: Unpack[AzureConfig], # type: ignore # noqa: PGH003 (container_name key overlaps with positional arg) + ) -> None: + """Construct a new AzureStore. + + Args: + container_name: the name of the container. + + Keyword Args: + prefix: A prefix within the bucket to use for all operations. + config: Azure Configuration. Values in this config will override values inferred from the url. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom Azure credentials. + kwargs: Azure configuration values. Supports the same values as `config`, but as named keyword args. + + Returns: + AzureStore + + """ + + @classmethod + def from_url( + cls, + url: str, + *, + prefix: str | None = None, + config: AzureConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: AzureCredentialProvider | None = None, + **kwargs: Unpack[AzureConfig], + ) -> Self: + """Construct a new AzureStore with values populated from a well-known storage URL. + + Any path on the URL will be assigned as the `prefix` for the store. So if you + pass `https://.blob.core.windows.net//path/to/directory`, + the store will be created with a prefix of `path/to/directory`, and all further + operations will use paths relative to that prefix. + + The supported url schemes are: + + - `abfs[s]:///` (according to [fsspec](https://github.com/fsspec/adlfs)) + - `abfs[s]://@.dfs.core.windows.net/` + - `abfs[s]://@.dfs.fabric.microsoft.com/` + - `az:///` (according to [fsspec](https://github.com/fsspec/adlfs)) + - `adl:///` (according to [fsspec](https://github.com/fsspec/adlfs)) + - `azure:///` (custom) + - `https://.dfs.core.windows.net` + - `https://.blob.core.windows.net` + - `https://.blob.core.windows.net/` + - `https://.dfs.fabric.microsoft.com` + - `https://.dfs.fabric.microsoft.com/` + - `https://.blob.fabric.microsoft.com` + - `https://.blob.fabric.microsoft.com/` + + Args: + url: well-known storage URL. + + Keyword Args: + prefix: A prefix within the bucket to use for all operations. + config: Azure Configuration. Values in this config will override values inferred from the url. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom Azure credentials. + kwargs: Azure configuration values. Supports the same values as `config`, but as named keyword args. + + Returns: + AzureStore + + """ + + def __eq__(self, value: object) -> bool: ... + def __getnewargs_ex__(self): ... + @property + def prefix(self) -> str | None: + """Get the prefix applied to all operations in this store, if any.""" + @property + def config(self) -> AzureConfig: + """Get the underlying Azure config parameters.""" + @property + def client_options(self) -> ClientConfig | None: + """Get the store's client configuration.""" + @property + def credential_provider(self) -> AzureCredentialProvider | None: + """Get the store's credential provider.""" + @property + def retry_config(self) -> RetryConfig | None: + """Get the store's retry configuration.""" diff --git a/vendor/pyo3-object_store/type-hints/_client.pyi b/vendor/pyo3-object_store/type-hints/_client.pyi new file mode 100644 index 00000000000..6f1f0593893 --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/_client.pyi @@ -0,0 +1,137 @@ +from datetime import timedelta +from typing import TypedDict + +class ClientConfig(TypedDict, total=False): + """HTTP client configuration. + + For timeout values (`connect_timeout`, `http2_keep_alive_timeout`, + `pool_idle_timeout`, and `timeout`), values can either be Python `timedelta` + objects, or they can be "human-readable duration strings". + + The human-readable duration string is a concatenation of time spans. Where each time + span is an integer number and a suffix. Supported suffixes: + + - `nsec`, `ns` -- nanoseconds + - `usec`, `us` -- microseconds + - `msec`, `ms` -- milliseconds + - `seconds`, `second`, `sec`, `s` + - `minutes`, `minute`, `min`, `m` + - `hours`, `hour`, `hr`, `h` + - `days`, `day`, `d` + - `weeks`, `week`, `w` + - `months`, `month`, `M` -- defined as 30.44 days + - `years`, `year`, `y` -- defined as 365.25 days + + For example: + + - `"2h 37min"` + - `"32ms"` + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import ClientConfig + ``` + """ + + allow_http: bool + """Allow non-TLS, i.e. non-HTTPS connections.""" + + allow_invalid_certificates: bool + """Skip certificate validation on https connections. + + !!! warning + + You should think very carefully before using this method. If + invalid certificates are trusted, *any* certificate for *any* site + will be trusted for use. This includes expired certificates. This + introduces significant vulnerabilities, and should only be used + as a last resort or for testing + """ + + connect_timeout: str | timedelta + """Set a timeout for only the connect phase of a Client. + + This is the time allowed for the client to establish a connection + and if the connection is not established within this time, + the client returns a timeout error. + + Timeout errors are retried, subject to the `RetryConfig`. + + Default is 5 seconds. + """ + + default_content_type: str + """Default [`CONTENT_TYPE`] for uploads. + + [`CONTENT_TYPE`]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/Content-Type + """ + default_headers: dict[str, str] | dict[str, bytes] + """Default headers to be sent with each request""" + http1_only: bool + """Only use HTTP/1 connections (default). + + Set to `false` to allow HTTP/2 connections. + """ + http2_keep_alive_interval: str + """Interval for HTTP/2 Ping frames should be sent to keep a connection alive.""" + http2_keep_alive_timeout: str | timedelta + """Timeout for receiving an acknowledgement of the keep-alive ping.""" + http2_keep_alive_while_idle: str + """Enable HTTP/2 keep alive pings for idle connections""" + http2_only: bool + """Only use HTTP/2 connections""" + pool_idle_timeout: str | timedelta + """The pool max idle timeout. + + This is the length of time an idle connection will be kept alive. + """ + pool_max_idle_per_host: str + """Maximum number of idle connections per host.""" + proxy_url: str + + randomize_addresses: bool + """Randomize order addresses that the DNS resolution yields. + + This will spread the connections across more servers. + + !!! warning + + This will override the DNS resolver configured by `reqwest`. + + """ + + read_timeout: str | timedelta + """Read timeout. + + The timeout applies to each read operation, and resets after a + successful read. This is useful for detecting stalled connections + when the size of the response is not known beforehand. + + Timeout errors are retried, subject to the `RetryConfig`. + + Default is disabled (no read timeout). + """ + + """HTTP proxy to use for requests.""" + timeout: str | timedelta + """Set timeout for the overall request + + The timeout starts from when the request starts connecting until the + response body has finished. If the request does not complete within the + timeout, the client returns a timeout error. + + Timeout errors are retried, subject to the `RetryConfig`. + + Default is 30 seconds. + """ + user_agent: str + """[User-Agent] header to be used by this client. + + [User-Agent]: https://developer.mozilla.org/en-US/docs/Web/HTTP/Reference/Headers/User-Agent + """ diff --git a/vendor/pyo3-object_store/type-hints/_gcs.pyi b/vendor/pyo3-object_store/type-hints/_gcs.pyi new file mode 100644 index 00000000000..089f34d75c7 --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/_gcs.pyi @@ -0,0 +1,235 @@ +import sys +from collections.abc import Coroutine +from datetime import datetime +from typing import Any, Protocol, TypedDict + +from ._client import ClientConfig +from ._retry import RetryConfig + +if sys.version_info >= (3, 11): + from typing import Self, Unpack +else: + from typing_extensions import Self, Unpack + +class GCSConfig(TypedDict, total=False): + """Configuration parameters for GCSStore. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import GCSConfig + ``` + """ + + service_account: str + """Path to the service account file. + + This or `service_account_key` must be set. + + Example value `"/tmp/gcs.json"`. Example contents of `gcs.json`: + + ```json + { + "gcs_base_url": "https://localhost:4443", + "disable_oauth": true, + "client_email": "", + "private_key": "" + } + ``` + + **Environment variables**: + + - `GOOGLE_SERVICE_ACCOUNT` + - `GOOGLE_SERVICE_ACCOUNT_PATH` + """ + + service_account_key: str + """The serialized service account key. + + The service account must be in the JSON format. This or `with_service_account_path` + must be set. + + **Environment variable**: `GOOGLE_SERVICE_ACCOUNT_KEY`. + """ + + base_url: str + """Sets the base URL for communicating with GCS. + + If not explicitly set, it will be: + + 1. Derived from the service account credentials, if provided + 2. Otherwise, uses the default GCS endpoint + + **Environment variable**: `GOOGLE_BASE_URL`. + """ + + bucket: str + """Bucket name. (required) + + **Environment variables**: + + - `GOOGLE_BUCKET` + - `GOOGLE_BUCKET_NAME` + """ + + application_credentials: str + """Application credentials path. + + See . + + **Environment variable**: `GOOGLE_APPLICATION_CREDENTIALS`. + """ + + skip_signature: bool + """If `True`, GCSStore will not fetch credentials and will not sign requests. + + This can be useful when interacting with public GCS buckets that deny authorized requests. + + **Environment variable**: `GOOGLE_SKIP_SIGNATURE`. + """ + +class GCSCredential(TypedDict): + """A Google Cloud Storage Credential. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import GCSCredential + ``` + """ + + token: str + """An HTTP bearer token.""" + + expires_at: datetime | None + """Expiry datetime of credential. The datetime should have time zone set. + + If None, the credential will never expire. + """ + +class GCSCredentialProvider(Protocol): + """A type hint for a synchronous or asynchronous callback to provide custom Google Cloud Storage credentials. + + This should be passed into the `credential_provider` parameter of `GCSStore`. + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import GCSCredentialProvider + ``` + """ + + def __call__(self) -> GCSCredential | Coroutine[Any, Any, GCSCredential]: + """Return a `GCSCredential`.""" + +class GCSStore: + """Interface to Google Cloud Storage. + + All constructors will check for environment variables. Refer to + [`GCSConfig`][obstore.store.GCSConfig] for valid environment variables. + + If no credentials are explicitly provided, they will be sourced from the environment + as documented + [here](https://cloud.google.com/docs/authentication/application-default-credentials). + """ + + def __init__( # type: ignore[misc] # Overlap between argument names and ** TypedDict items: "bucket" + self, + bucket: str | None = None, + *, + prefix: str | None = None, + config: GCSConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: GCSCredentialProvider | None = None, + **kwargs: Unpack[GCSConfig], # type: ignore # noqa: PGH003 (bucket key overlaps with positional arg) + ) -> None: + """Construct a new GCSStore. + + Args: + bucket: The GCS bucket to use. + + Keyword Args: + prefix: A prefix within the bucket to use for all operations. + config: GCS Configuration. Values in this config will override values inferred from the environment. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom Google credentials. + kwargs: GCS configuration values. Supports the same values as `config`, but as named keyword args. + + Returns: + GCSStore + + """ + + @classmethod + def from_url( + cls, + url: str, + *, + prefix: str | None = None, + config: GCSConfig | None = None, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + credential_provider: GCSCredentialProvider | None = None, + **kwargs: Unpack[GCSConfig], + ) -> Self: + """Construct a new GCSStore with values populated from a well-known storage URL. + + Any path on the URL will be assigned as the `prefix` for the store. So if you + pass `gs:///path/to/directory`, the store will be created with a prefix + of `path/to/directory`, and all further operations will use paths relative to + that prefix. + + The supported url schemes are: + + - `gs:///` + + Args: + url: well-known storage URL. + + Keyword Args: + prefix: A prefix within the bucket to use for all operations. + config: GCS Configuration. Values in this config will override values inferred from the url. Defaults to None. + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + credential_provider: A callback to provide custom Google credentials. + kwargs: GCS configuration values. Supports the same values as `config`, but as named keyword args. + + Returns: + GCSStore + + """ + + def __eq__(self, value: object) -> bool: ... + def __getnewargs_ex__(self): ... + @property + def prefix(self) -> str | None: + """Get the prefix applied to all operations in this store, if any.""" + @property + def config(self) -> GCSConfig: + """Get the underlying GCS config parameters.""" + @property + def client_options(self) -> ClientConfig | None: + """Get the store's client configuration.""" + @property + def credential_provider(self) -> GCSCredentialProvider | None: + """Get the store's credential provider.""" + @property + def retry_config(self) -> RetryConfig | None: + """Get the store's retry configuration.""" diff --git a/vendor/pyo3-object_store/type-hints/_http.pyi b/vendor/pyo3-object_store/type-hints/_http.pyi new file mode 100644 index 00000000000..3558709d73a --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/_http.pyi @@ -0,0 +1,63 @@ +import sys + +from ._client import ClientConfig +from ._retry import RetryConfig + +if sys.version_info >= (3, 11): + from typing import Self +else: + from typing_extensions import Self + +class HTTPStore: + """Configure a connection to a generic HTTP server.""" + + def __init__( + self, + url: str, + *, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + ) -> None: + """Construct a new HTTPStore from a URL. + + Any path on the URL will be assigned as the `prefix` for the store. So if you + pass `https://example.com/path/to/directory`, the store will be created with a + prefix of `path/to/directory`, and all further operations will use paths + relative to that prefix. + + Args: + url: The base URL to use for the store. + + Keyword Args: + client_options: HTTP Client options. Defaults to None. + retry_config: Retry configuration. Defaults to None. + + Returns: + HTTPStore + + """ + + @classmethod + def from_url( + cls, + url: str, + *, + client_options: ClientConfig | None = None, + retry_config: RetryConfig | None = None, + ) -> Self: + """Construct a new HTTPStore from a URL. + + This is an alias of [`HTTPStore.__init__`][obstore.store.HTTPStore.__init__]. + """ + + def __eq__(self, value: object) -> bool: ... + def __getnewargs_ex__(self): ... + @property + def url(self) -> str: + """Get the base url of this store.""" + @property + def client_options(self) -> ClientConfig | None: + """Get the store's client configuration.""" + @property + def retry_config(self) -> RetryConfig | None: + """Get the store's retry configuration.""" diff --git a/vendor/pyo3-object_store/type-hints/_retry.pyi b/vendor/pyo3-object_store/type-hints/_retry.pyi new file mode 100644 index 00000000000..bfe4068325b --- /dev/null +++ b/vendor/pyo3-object_store/type-hints/_retry.pyi @@ -0,0 +1,97 @@ +from datetime import timedelta +from typing import TypedDict + +class BackoffConfig(TypedDict, total=False): + """Exponential backoff with jitter. + + See + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import BackoffConfig + ``` + """ + + init_backoff: timedelta + """The initial backoff duration. + + Defaults to 100 milliseconds. + """ + + max_backoff: timedelta + """The maximum backoff duration. + + Defaults to 15 seconds. + """ + + base: int | float + """The base of the exponential to use. + + Defaults to `2`. + """ + +class RetryConfig(TypedDict, total=False): + """The configuration for how to respond to request errors. + + The following categories of error will be retried: + + * 5xx server errors + * Connection errors + * Dropped connections + * Timeouts for [safe] / read-only requests + + Requests will be retried up to some limit, using exponential + backoff with jitter. See [`BackoffConfig`][obstore.store.BackoffConfig] for + more information + + [safe]: https://datatracker.ietf.org/doc/html/rfc7231#section-4.2.1 + + !!! warning "Not importable at runtime" + + To use this type hint in your code, import it within a `TYPE_CHECKING` block: + + ```py + from __future__ import annotations + from typing import TYPE_CHECKING + if TYPE_CHECKING: + from obstore.store import RetryConfig + ``` + """ + + backoff: BackoffConfig + """The backoff configuration. + + Defaults to the values listed above if not provided. + """ + + max_retries: int + """ + The maximum number of times to retry a request + + Set to 0 to disable retries. + + Defaults to 10. + """ + + retry_timeout: timedelta + """ + The maximum length of time from the initial request + after which no further retries will be attempted + + This not only bounds the length of time before a server + error will be surfaced to the application, but also bounds + the length of time a request's credentials must remain valid. + + As requests are retried without renewing credentials or + regenerating request payloads, this number should be kept + below 5 minutes to avoid errors due to expired credentials + and/or request payloads. + + Defaults to 3 minutes. + """ diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index 6bb50cf097d..544a9cffd44 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -34,10 +34,16 @@ url = { workspace = true } vortex = { workspace = true, features = ["object_store", "files"] } vortex-arrow = { workspace = true } vortex-parquet-variant = { workspace = true } +vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } [dev-dependencies] jni = { workspace = true, features = ["invocation"] } +[features] +# Enable OpenDAL-backed object stores (Tencent COS, Alibaba OSS) for `cos://` and `oss://` URLs. +# This pulls in the `opendal` dependency, so it is opt-in. +opendal = ["dep:vortex-object-store-opendal"] + [lib] crate-type = ["cdylib"] diff --git a/vortex-jni/src/object_store.rs b/vortex-jni/src/object_store.rs index cfc29f04d20..127b4b05731 100644 --- a/vortex-jni/src/object_store.rs +++ b/vortex-jni/src/object_store.rs @@ -38,16 +38,28 @@ pub(crate) fn object_store_fs( Ok(Arc::new(ObjectStoreFileSystem::new(object_store, handle))) } +/// Process-wide cache of constructed object stores, keyed by URL + properties so that repeated +/// requests against the same bucket/configuration share a single client. +static OBJECT_STORES: LazyLock>>> = + LazyLock::new(|| Mutex::new(HashMap::new())); + #[expect(clippy::cognitive_complexity)] pub(crate) fn make_object_store( url: &Url, properties: &HashMap, ) -> VortexResult> { - static OBJECT_STORES: LazyLock>>> = - LazyLock::new(|| Mutex::new(HashMap::new())); - let start = std::time::Instant::now(); + // OpenDAL-backed stores (Tencent COS, Alibaba OSS) use schemes that `object_store` does not + // recognize natively. Handle them first via the optional `opendal` feature so they can be + // resolved without falling through to the `Unsupported store scheme` error below. + #[cfg(feature = "opendal")] + if matches!(url.scheme(), "cos" | "oss") { + let store = vortex_object_store_opendal::make_opendal_store(url, properties) + .map_err(|e| VortexError::from(object_store::Error::from(e)))?; + return cache_and_return(store, url, properties, &start); + } + let (scheme, _) = ObjectStoreScheme::parse(url) .map_err(|error| VortexError::from(object_store::Error::from(error)))?; @@ -143,6 +155,18 @@ pub(crate) fn make_object_store( } }; + cache_and_return(store, url, properties, &start) +} + +/// Insert the built store into the process-wide cache (keyed by URL + properties) and return it, +/// logging the construction latency. +fn cache_and_return( + store: Arc, + url: &Url, + properties: &HashMap, + start: &std::time::Instant, +) -> VortexResult> { + let cache_key = url_cache_key(url, properties); OBJECT_STORES.lock().insert(cache_key, Arc::clone(&store)); let duration = start.elapsed(); diff --git a/vortex-object-store-opendal/Cargo.toml b/vortex-object-store-opendal/Cargo.toml new file mode 100644 index 00000000000..10ffec397e8 --- /dev/null +++ b/vortex-object-store-opendal/Cargo.toml @@ -0,0 +1,26 @@ +[package] +name = "vortex-object-store-opendal" +description = "OpenDAL-backed object store (Tencent COS) for Vortex" +publish = false +version = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +authors = { workspace = true } +license = { workspace = true } +keywords = { workspace = true } +include = { workspace = true } +edition = { workspace = true } +rust-version = { workspace = true } +categories = { workspace = true } + +[dependencies] +object_store = { workspace = true, features = ["cloud"] } +opendal = { version = "0.57.0", default-features = false, features = [ + "services-cos", +] } +object_store_opendal = "0.57.0" +url = { workspace = true } +vortex-utils = { workspace = true } + +[lints] +workspace = true diff --git a/vortex-object-store-opendal/src/lib.rs b/vortex-object-store-opendal/src/lib.rs new file mode 100644 index 00000000000..69e446b59c8 --- /dev/null +++ b/vortex-object-store-opendal/src/lib.rs @@ -0,0 +1,160 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! OpenDAL-backed [`object_store::ObjectStore`] implementations for cloud providers that are not +//! natively supported by the `object_store` crate, such as Tencent Cloud COS. +//! +//! OpenDAL exposes each service as an `Operator`. We adapt an `Operator` into an +//! `object_store::ObjectStore` via the `object_store_opendal::OpendalStore` bridge, which is built +//! against the same `object_store 0.13.x` version the rest of Vortex uses. This lets Vortex consume +//! COS through its existing `ObjectStoreFileSystem` abstraction without any changes to +//! `vortex-io`. + +use std::sync::Arc; + +use object_store::ObjectStore; +use object_store_opendal::OpendalStore; +use opendal::Operator; +use opendal::services; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +/// Error type for building an OpenDAL-backed object store. +#[derive(Debug)] +pub enum OpenDALStoreError { + /// The URL scheme is not one this crate handles (e.g. `s3`, `gs`, ...). + UnsupportedScheme(String), + /// A required configuration value (bucket and/or endpoint) was missing. + MissingConfig(&'static str), + /// The OpenDAL builder rejected the provided configuration. + Build(opendal::Error), +} + +impl std::fmt::Display for OpenDALStoreError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + OpenDALStoreError::UnsupportedScheme(s) => { + write!(f, "unsupported OpenDAL scheme: {s}") + } + OpenDALStoreError::MissingConfig(k) => { + write!(f, "missing required OpenDAL store configuration: {k}") + } + OpenDALStoreError::Build(e) => write!(f, "failed to build OpenDAL store: {e}"), + } + } +} + +impl std::error::Error for OpenDALStoreError {} + +impl From for object_store::Error { + fn from(e: OpenDALStoreError) -> Self { + object_store::Error::Generic { + store: "OpenDAL", + source: Box::new(e), + } + } +} + +/// Schemes handled by this crate. +pub const COS_SCHEME: &str = "cos"; + +/// Build an [`object_store::ObjectStore`] for a COS URL. +/// +/// `properties` are per-request configuration overrides (matching the `HashMap` +/// passed through the JNI/Python layers). Missing values fall back to environment variables that +/// OpenDAL's builders read automatically (e.g. `TENCENTCLOUD_SECRET_ID`). +/// +/// Returns [`OpenDALStoreError::UnsupportedScheme`] if `url` does not use `cos://`. +pub fn make_opendal_store( + url: &Url, + properties: &HashMap, +) -> Result, OpenDALStoreError> { + match url.scheme() { + COS_SCHEME => make_cos_store(url, properties), + other => Err(OpenDALStoreError::UnsupportedScheme(other.to_string())), + } +} + +fn make_cos_store( + url: &Url, + properties: &HashMap, +) -> Result, OpenDALStoreError> { + let bucket = properties + .get("bucket") + .cloned() + .or_else(|| url.host_str().map(str::to_string)) + .ok_or(OpenDALStoreError::MissingConfig("bucket"))?; + let endpoint = properties + .get("endpoint") + .cloned() + .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; + + let mut builder = services::Cos::default().bucket(&bucket).endpoint(&endpoint); + + if let Some(root) = properties.get("root") { + builder = builder.root(root); + } + if let Some(secret_id) = properties.get("secret_id") { + builder = builder.secret_id(secret_id); + } + if let Some(secret_key) = properties.get("secret_key") { + builder = builder.secret_key(secret_key); + } + if properties.get("disable_config_load").map(String::as_str) == Some("true") { + builder = builder.disable_config_load(); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +fn build_operator(builder: B) -> Result +where + B: opendal::Builder, +{ + let operator: Operator = Operator::new(builder) + .map_err(OpenDALStoreError::Build)? + .finish(); + Ok(operator) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn rejects_unknown_scheme() { + let url = Url::parse("s3://bucket/path").unwrap(); + let props = HashMap::new(); + assert!(matches!( + make_opendal_store(&url, &props), + Err(OpenDALStoreError::UnsupportedScheme(_)) + )); + } + + #[test] + fn cos_requires_endpoint() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let props = HashMap::new(); + assert!(matches!( + make_cos_store(&url, &props), + Err(OpenDALStoreError::MissingConfig("endpoint")) + )); + } + + #[test] + fn cos_derives_bucket_from_host() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let mut props = HashMap::new(); + props.insert( + "endpoint".to_string(), + "https://cos.ap-guangzhou.myqcloud.com".to_string(), + ); + // Missing secret_id/secret_key -> build fails before we can assert bucket derivation, + // but we should get past the MissingConfig checks (i.e. not a MissingConfig error). + assert!(!matches!( + make_cos_store(&url, &props), + Err(OpenDALStoreError::MissingConfig(_)) + )); + } +} diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index a87548d69a9..896519d16c2 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -28,6 +28,9 @@ default = ["extension-module", "tui"] # duplicate PyInit__lib symbols. extension-module = [] tui = ["dep:tokio", "dep:vortex-tui"] +# OpenDAL-backed object stores (Tencent COS, Alibaba OSS) for `cos://` and `oss://` URLs. +# Pulls in `opendal`, so it is opt-in. +opendal = ["dep:vortex-object-store-opendal"] [dependencies] arrow-array = { workspace = true } @@ -53,7 +56,9 @@ tokio = { workspace = true, features = ["rt-multi-thread"], optional = true } url = { workspace = true } vortex = { workspace = true, features = ["object_store"] } vortex-arrow = { workspace = true } +vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } vortex-python-abi = { path = "../vortex-python-abi" } +vortex-utils = { workspace = true } vortex-tui = { workspace = true, optional = true } [dev-dependencies] diff --git a/vortex-python/python/vortex/_lib/store/__init__.pyi b/vortex-python/python/vortex/_lib/store/__init__.pyi index b394b118b80..2e6965872ff 100644 --- a/vortex-python/python/vortex/_lib/store/__init__.pyi +++ b/vortex-python/python/vortex/_lib/store/__init__.pyi @@ -89,6 +89,12 @@ def from_url( # type: ignore[misc] # docstring in pyi file - `gs://bucket/path` -> [`GCSStore`][vortex.store.GCSStore] - `az://account/container/path` -> [`AzureStore`][vortex.store.AzureStore] (also supports `adl`, `azure`, `abfs`, `abfss`) + - `cos://bucket/path` -> OpenDAL-backed Tencent Cloud COS store (requires the + `opendal` feature; configure via environment variables such as + `TENCENTCLOUD_SECRET_ID` / `TENCENTCLOUD_SECRET_KEY` and `COS_ENDPOINT`) + - `oss://bucket/path` -> OpenDAL-backed Alibaba Cloud OSS store (requires the + `opendal` feature; configure via `ALIBABACLOUD_ACCESS_KEY_ID` / + `ALIBABACLOUD_ACCESS_KEY_SECRET` and `OSS_ENDPOINT`) - `http://mydomain/path` -> [`HTTPStore`][vortex.store.HTTPStore] - `https://mydomain/path` -> [`HTTPStore`][vortex.store.HTTPStore] diff --git a/vortex-python/python/vortex/store/__init__.py b/vortex-python/python/vortex/store/__init__.py index 043f0bf9ab6..050d14766ff 100644 --- a/vortex-python/python/vortex/store/__init__.py +++ b/vortex-python/python/vortex/store/__init__.py @@ -16,13 +16,22 @@ AzureStore, ) from ._client import ClientConfig +from ._cos import CosStore from ._gcs import GCSConfig, GCSCredential, GCSCredentialProvider, GCSStore from ._http import HTTPStore from ._local import LocalStore from ._memory import MemoryStore from ._retry import BackoffConfig, RetryConfig -ObjectStore: TypeAlias = AzureStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore +ObjectStore: TypeAlias = ( + AzureStore + | CosStore + | GCSStore + | HTTPStore + | S3Store + | LocalStore + | MemoryStore +) """All supported ObjectStore implementations.""" @@ -89,6 +98,12 @@ def from_url( # type: ignore[misc] # docstring in pyi file - ``gs://bucket/path`` -> :class:`~vortex.store.GCSStore` - ``az://account/container/path`` -> :class:`~vortex.store.AzureStore` (also supports ``adl``, ``azure``, ``abfs``, ``abfss``) + - ``cos://bucket/path`` -> OpenDAL-backed Tencent Cloud COS store (requires the + ``opendal`` feature; configure via environment variables such as + ``TENCENTCLOUD_SECRET_ID`` / ``TENCENTCLOUD_SECRET_KEY`` and ``COS_ENDPOINT``) + - ``oss://bucket/path`` -> OpenDAL-backed Alibaba Cloud OSS store (requires the + ``opendal`` feature; configure via ``ALIBABACLOUD_ACCESS_KEY_ID`` / + ``ALIBABACLOUD_ACCESS_KEY_SECRET`` and ``OSS_ENDPOINT``) - ``http://mydomain/path`` -> :class:`~vortex.store.HTTPStore` - ``https://mydomain/path`` -> :class:`~vortex.store.HTTPStore` @@ -139,6 +154,8 @@ def from_url( # type: ignore[misc] # docstring in pyi file "BackoffConfig", "ClientConfig", "RetryConfig", + # COS (OpenDAL-backed) + "CosStore", # GCS "GCSConfig", "GCSCredential", diff --git a/vortex-python/python/vortex/store/_cos.py b/vortex-python/python/vortex/store/_cos.py new file mode 100644 index 00000000000..8d69f5e4cb7 --- /dev/null +++ b/vortex-python/python/vortex/store/_cos.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright the Vortex contributors + +"""OpenDAL-backed object store for Tencent Cloud COS. + +This store is only available when Vortex is built with the ``opendal`` feature. Unlike the +URL-based resolution (``cos://``), this class is a concrete, standalone +``ObjectStore`` object that you can build once and pass directly to +:func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument:: + + from vortex.store import CosStore + from vortex.io import read_url + + store = CosStore( + bucket="my-bucket", + endpoint="https://cos.ap-guangzhou.myqcloud.com", + secret_id="AKID...", + secret_key="...", + ) + array = read_url("cos://my-bucket/data.vortex", store=store) + +Credentials may also be supplied via the environment variables that OpenDAL's builders read +automatically (``TENCENTCLOUD_SECRET_ID`` / ``TENCENTCLOUD_SECRET_KEY`` and ``COS_ENDPOINT``); +any value passed to the constructor takes precedence over the environment. +""" + +from __future__ import annotations + +from vortex._lib import CosStore + +__all__ = ["CosStore"] + + +class _OpenDALStore: + """Base helper that projects keyword configuration into OpenDAL environment variables.""" + + # (our kwarg name) -> (env var OpenDAL's builder reads) + _ENV_MAP: dict[str, str] + + def __init__(self, **kwargs: Any) -> None: + self._config: dict[str, str] = {} + for key, value in kwargs.items(): + if value is None: + continue + self._config[key] = str(value) + + def apply(self) -> None: + """Export this store's configuration into the process environment. + + After calling :meth:`apply`, ``cos://`` / ``oss://`` URLs resolve using these values. + """ + for key, env_var in self._ENV_MAP.items(): + if key in self._config: + os.environ[env_var] = self._config[key] + + +class CosStore(_OpenDALStore): + """Configuration helper for Tencent Cloud COS, resolved from ``cos://`` URLs. + + Keyword Args: + bucket: COS bucket name. + endpoint: COS endpoint, e.g. ``https://cos.ap-guangzhou.myqcloud.com``. + secret_id: Tencent Cloud secret id (mapped to ``TENCENTCLOUD_SECRET_ID``). + secret_key: Tencent Cloud secret key (mapped to ``TENCENTCLOUD_SECRET_KEY``). + root: Optional root prefix applied to all operations. + """ + + _ENV_MAP = { + "secret_id": "TENCENTCLOUD_SECRET_ID", + "secret_key": "TENCENTCLOUD_SECRET_KEY", + "endpoint": "COS_ENDPOINT", + } diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 182ab7fbe74..1332a063f4c 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -1,6 +1,8 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors +use std::sync::Arc; + use arrow_array::RecordBatchReader; use arrow_array::ffi_stream::ArrowArrayStreamReader; use async_fs::File; @@ -41,6 +43,8 @@ use crate::install_module; use crate::iter::PyArrayIterator; use crate::object_store::resolve::ResolvedStore; use crate::object_store::resolve::resolve_store; +#[cfg(feature = "opendal")] +use crate::opendal_store::CosStore; use crate::session::session; pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { @@ -62,7 +66,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { /// ---------- /// url : str /// The URL to read from. -/// store : vortex.store.AzureStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None +/// store : vortex.store.AzureStore | vortex.store.CosStore | vortex.store.GCSStore | vortex.store.HTTPStore | vortex.store.LocalStore | vortex.store.MemoryStore | vortex.store.S3Store | None /// Pre-configured object store with credentials and settings. /// If provided, uses this store's configuration. /// If None, checks session registry for matching URL pattern. @@ -115,6 +119,7 @@ pub(crate) fn init(py: Python, parent: &Bound) -> PyResult<()> { /// ... secret_access_key="..." /// ... ) /// >>> a = vx.io.read_url("s3://my-bucket/data.vortex", store=store) # doctest: +SKIP + #[pyfunction] #[pyo3(signature = (url, *, store = None, projection = None, row_filter = None, indices = None, row_range = None))] pub fn read_url<'py>( @@ -127,7 +132,7 @@ pub fn read_url<'py>( row_range: Option<(u64, u64)>, ) -> PyVortexResult { let store_arc = if let Some(store_obj) = store { - let py_store: PyObjectStore = store_obj.extract()?; + let py_store: AnyVortexStore = store_obj.extract()?; Some(py_store.into_inner()) } else { None @@ -137,6 +142,46 @@ pub fn read_url<'py>( dataset.to_array_inner(py, projection, row_filter, indices, row_range) } +/// A store object accepted by `read_url` / `write`. +/// +/// This recognizes both the built-in `pyo3-object_store` classes (S3, Azure, GCS, HTTP, +/// Local, Memory) and Vortex's own OpenDAL-backed classes (`CosStore`). +pub(crate) enum AnyVortexStore { + /// A store extracted from one of the built-in `pyo3-object_store` classes. + Builtin(PyObjectStore), + /// Vortex's OpenDAL-backed COS store. + #[cfg(feature = "opendal")] + Cos(CosStore), +} + +impl AnyVortexStore { + /// Consume self and return the underlying `Arc`. + fn into_inner(self) -> Arc { + match self { + AnyVortexStore::Builtin(s) => s.into_inner(), + #[cfg(feature = "opendal")] + AnyVortexStore::Cos(s) => s.to_arc(), + } + } +} + +impl<'py> FromPyObject<'_, 'py> for AnyVortexStore { + type Error = PyErr; + + fn extract(obj: Borrowed<'_, 'py, PyAny>) -> PyResult { + if let Ok(builtin) = obj.extract::() { + return Ok(AnyVortexStore::Builtin(builtin)); + } + #[cfg(feature = "opendal")] + if let Ok(cos) = obj.extract::() { + return Ok(AnyVortexStore::Cos(cos)); + } + Err(PyTypeError::new_err( + "Expected an object store instance (S3/Azure/GCS/HTTP/Local/Memory/COS/OSS store)", + )) + } +} + /// Write an array to a Vortex file. /// /// Parameters @@ -190,7 +235,7 @@ pub fn write( py: Python, iter: PyIntoArrayIterator, path: &str, - store: Option, + store: Option, ) -> PyVortexResult<()> { let session = session(); py.detach(|| { @@ -319,7 +364,7 @@ impl PyVortexWriteOptions { py: Python, iter: PyIntoArrayIterator, path: &str, - store: Option, + store: Option, ) -> PyVortexResult<()> { let session = session(); py.detach(|| { diff --git a/vortex-python/src/lib.rs b/vortex-python/src/lib.rs index 6582a255955..0912426d10b 100644 --- a/vortex-python/src/lib.rs +++ b/vortex-python/src/lib.rs @@ -25,6 +25,8 @@ mod file; mod io; mod iter; mod object_store; +#[cfg(feature = "opendal")] +mod opendal_store; mod python_repr; mod registry; mod runtime; @@ -80,6 +82,8 @@ fn _lib(py: Python, m: &Bound) -> PyResult<()> { file::init(py, m)?; io::init(py, m)?; iter::init(py, m)?; + #[cfg(feature = "opendal")] + opendal_store::init(py, m)?; runtime::init(py, m)?; store::init(py, m)?; registry::init(py, m)?; diff --git a/vortex-python/src/object_store/registry.rs b/vortex-python/src/object_store/registry.rs index c9544e0558e..17975b1e00e 100644 --- a/vortex-python/src/object_store/registry.rs +++ b/vortex-python/src/object_store/registry.rs @@ -19,6 +19,8 @@ use object_store::path::PathPart; use object_store::registry::ObjectStoreRegistry; use parking_lot::RwLock; use url::Url; +#[cfg(feature = "opendal")] +use vortex::utils::aliases::hash_map::HashMap as VortexHashMap; #[derive(Debug, Default)] struct PathEntry { @@ -75,6 +77,27 @@ impl ObjectStoreRegistry for Registry { fn resolve(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { let key = url_key(to_resolve); + + // OpenDAL-backed stores (Tencent COS) use schemes that the `object_store` + // crate does not recognize. Resolve them via the optional `opendal` feature, relying on + // OpenDAL's environment-variable configuration (e.g. `TENCENTCLOUD_SECRET_ID`). + #[cfg(feature = "opendal")] + if matches!(to_resolve.scheme(), "cos") { + let store = + vortex_object_store_opendal::make_opendal_store(to_resolve, &VortexHashMap::new()) + .map_err(|e| object_store::Error::Generic { + store: "OpenDAL", + source: Box::new(e), + })?; + let path = Path::from_url_path(to_resolve.path()).map_err(|e| { + object_store::Error::Generic { + store: "OpenDAL", + source: Box::new(e), + } + })?; + return Ok((store, path)); + } + { let map = self.map.read(); diff --git a/vortex-python/src/opendal_store.rs b/vortex-python/src/opendal_store.rs new file mode 100644 index 00000000000..63acc4c5544 --- /dev/null +++ b/vortex-python/src/opendal_store.rs @@ -0,0 +1,126 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Python-facing [`PyObjectStore`](pyo3_object_store::PyObjectStore) wrapper for the +//! OpenDAL-backed Tencent Cloud COS object store. +//! +//! This lets callers build a concrete store object in Python and pass it to +//! `vortex.io.read_url(store=...)` / `vortex.io.write(..., store=...)` exactly like any +//! of the built-in store classes (S3, Azure, ...). The store is constructed from the same +//! configuration the `cos://` URL registry uses, but materialized eagerly so it +//! can be handed around as a first-class value. +//! +//! This module is only compiled when the `opendal` feature is enabled. + +use std::sync::Arc; + +use pyo3::exceptions::PyValueError; +use pyo3::prelude::*; +use pyo3_object_store::PyObjectStore; +use url::Url; +use vortex_utils::aliases::hash_map::HashMap; + +/// Build a [`PyObjectStore`] for the given `cos://` / `oss://` URL and properties. +fn build_store(url: &str, properties: HashMap) -> PyResult { + let url = Url::parse(url) + .map_err(|e| PyValueError::new_err(format!("invalid store URL {url}: {e}")))?; + let store: Arc = + vortex_object_store_opendal::make_opendal_store(&url, &properties) + .map_err(|e| PyValueError::new_err(e.to_string()))?; + Ok(PyObjectStore::from_store(store)) +} + +/// A Tencent Cloud COS object store, backed by OpenDAL. +/// +/// Construct it with explicit configuration and pass it to +/// ``vortex.io.read_url(url, store=cos_store)`` / ``vortex.io.write(arrays, path, store=cos_store)``. +#[pyclass(name = "CosStore", module = "vortex._lib", from_py_object)] +#[derive(Clone)] +pub struct CosStore { + store: PyObjectStore, +} + +impl CosStore { + /// Clone the underlying object store as an `Arc`. + pub fn to_arc(&self) -> Arc { + self.store.clone().into_inner() + } +} + +impl From for PyObjectStore { + fn from(value: CosStore) -> Self { + value.store + } +} + +#[pymethods] +impl CosStore { + #[new] + #[pyo3(signature = ( + bucket, + endpoint, + *, + secret_id = None, + secret_key = None, + root = None, + disable_config_load = false, + ))] + fn new( + bucket: String, + endpoint: String, + secret_id: Option, + secret_key: Option, + root: Option, + disable_config_load: bool, + ) -> PyResult { + let mut properties = HashMap::new(); + properties.insert("bucket".to_string(), bucket); + properties.insert("endpoint".to_string(), endpoint); + if let Some(v) = secret_id { + properties.insert("secret_id".to_string(), v); + } + if let Some(v) = secret_key { + properties.insert("secret_key".to_string(), v); + } + if let Some(v) = root { + properties.insert("root".to_string(), v); + } + if disable_config_load { + properties.insert("disable_config_load".to_string(), "true".to_string()); + } + let store = build_store("cos://bucket", properties)?; + Ok(Self { store }) + } +} + +/// Register the OpenDAL-backed store classes on the `vortex._lib` module. +#[cfg(feature = "opendal")] +pub(crate) fn init(_py: Python, parent: &Bound) -> PyResult<()> { + parent.add_class::()?; + Ok(()) +} + +#[cfg(all(test, feature = "opendal"))] +mod tests { + use super::*; + + /// A `CosStore` built in Rust must yield a non-null `Arc` that can be + /// handed to `read_url`/`write` (this is the path previously blocked by the missing + /// `From> for PyObjectStore` upstream constructor). + #[test] + fn cos_store_builds_object_store() { + let store = CosStore::new( + "my-bucket".to_string(), + "https://cos.ap-guangzhou.myqcloud.com".to_string(), + Some("AKID".to_string()), + Some("secret".to_string()), + None, + false, + ) + .expect("cos store should build"); + // The store must be usable as an `Arc` (this is the path previously + // blocked by the missing `From> for PyObjectStore` upstream). + let arc: Arc = store.to_arc(); + assert!(Arc::strong_count(&arc) >= 1); + } +}