From f85d3d2e1efb87ffa38a0fa13c6592d05caf0efd Mon Sep 17 00:00:00 2001 From: forwardxu Date: Sun, 19 Jul 2026 10:37:32 +0800 Subject: [PATCH 1/8] feat(vortex-python): Add OpenDAL-backed CosStore to the Python object-store API Signed-off-by: forwardxu --- Cargo.lock | 153 ++++- Cargo.toml | 10 + docs/api/python/store.rst | 1 + docs/api/python/store/opendal.rst | 62 ++ vendor/pyo3-object_store/Cargo.toml | 47 ++ vendor/pyo3-object_store/LICENSE | 21 + vendor/pyo3-object_store/README.md | 79 +++ vendor/pyo3-object_store/src/api.rs | 181 +++++ .../pyo3-object_store/src/aws/credentials.rs | 216 ++++++ vendor/pyo3-object_store/src/aws/mod.rs | 4 + vendor/pyo3-object_store/src/aws/store.rs | 430 ++++++++++++ .../src/azure/credentials.rs | 313 +++++++++ vendor/pyo3-object_store/src/azure/error.rs | 21 + vendor/pyo3-object_store/src/azure/mod.rs | 5 + vendor/pyo3-object_store/src/azure/store.rs | 489 ++++++++++++++ vendor/pyo3-object_store/src/client.rs | 180 +++++ vendor/pyo3-object_store/src/config.rs | 58 ++ vendor/pyo3-object_store/src/credentials.rs | 101 +++ vendor/pyo3-object_store/src/error.rs | 193 ++++++ .../pyo3-object_store/src/gcp/credentials.rs | 193 ++++++ vendor/pyo3-object_store/src/gcp/mod.rs | 4 + vendor/pyo3-object_store/src/gcp/store.rs | 380 +++++++++++ vendor/pyo3-object_store/src/http.rs | 134 ++++ vendor/pyo3-object_store/src/lib.rs | 35 + vendor/pyo3-object_store/src/local.rs | 143 ++++ vendor/pyo3-object_store/src/memory.rs | 47 ++ vendor/pyo3-object_store/src/path.rs | 64 ++ vendor/pyo3-object_store/src/prefix.rs | 260 ++++++++ vendor/pyo3-object_store/src/retry.rs | 103 +++ vendor/pyo3-object_store/src/simple.rs | 112 ++++ vendor/pyo3-object_store/src/store.rs | 276 ++++++++ vendor/pyo3-object_store/src/url.rs | 64 ++ .../pyo3-object_store/type-hints/__init__.pyi | 218 ++++++ vendor/pyo3-object_store/type-hints/_aws.pyi | 623 ++++++++++++++++++ .../pyo3-object_store/type-hints/_azure.pyi | 422 ++++++++++++ .../pyo3-object_store/type-hints/_client.pyi | 137 ++++ vendor/pyo3-object_store/type-hints/_gcs.pyi | 235 +++++++ vendor/pyo3-object_store/type-hints/_http.pyi | 63 ++ .../pyo3-object_store/type-hints/_retry.pyi | 97 +++ vortex-jni/Cargo.toml | 6 + vortex-jni/src/object_store.rs | 30 +- vortex-object-store-opendal/Cargo.toml | 26 + vortex-object-store-opendal/src/lib.rs | 160 +++++ vortex-python/Cargo.toml | 5 + .../python/vortex/_lib/store/__init__.pyi | 6 + vortex-python/python/vortex/store/__init__.py | 19 +- vortex-python/python/vortex/store/_cos.py | 72 ++ vortex-python/src/io.rs | 53 +- vortex-python/src/lib.rs | 4 + vortex-python/src/object_store/registry.rs | 23 + vortex-python/src/opendal_store.rs | 126 ++++ 51 files changed, 6694 insertions(+), 10 deletions(-) create mode 100644 docs/api/python/store/opendal.rst create mode 100644 vendor/pyo3-object_store/Cargo.toml create mode 100644 vendor/pyo3-object_store/LICENSE create mode 100644 vendor/pyo3-object_store/README.md create mode 100644 vendor/pyo3-object_store/src/api.rs create mode 100644 vendor/pyo3-object_store/src/aws/credentials.rs create mode 100644 vendor/pyo3-object_store/src/aws/mod.rs create mode 100644 vendor/pyo3-object_store/src/aws/store.rs create mode 100644 vendor/pyo3-object_store/src/azure/credentials.rs create mode 100644 vendor/pyo3-object_store/src/azure/error.rs create mode 100644 vendor/pyo3-object_store/src/azure/mod.rs create mode 100644 vendor/pyo3-object_store/src/azure/store.rs create mode 100644 vendor/pyo3-object_store/src/client.rs create mode 100644 vendor/pyo3-object_store/src/config.rs create mode 100644 vendor/pyo3-object_store/src/credentials.rs create mode 100644 vendor/pyo3-object_store/src/error.rs create mode 100644 vendor/pyo3-object_store/src/gcp/credentials.rs create mode 100644 vendor/pyo3-object_store/src/gcp/mod.rs create mode 100644 vendor/pyo3-object_store/src/gcp/store.rs create mode 100644 vendor/pyo3-object_store/src/http.rs create mode 100644 vendor/pyo3-object_store/src/lib.rs create mode 100644 vendor/pyo3-object_store/src/local.rs create mode 100644 vendor/pyo3-object_store/src/memory.rs create mode 100644 vendor/pyo3-object_store/src/path.rs create mode 100644 vendor/pyo3-object_store/src/prefix.rs create mode 100644 vendor/pyo3-object_store/src/retry.rs create mode 100644 vendor/pyo3-object_store/src/simple.rs create mode 100644 vendor/pyo3-object_store/src/store.rs create mode 100644 vendor/pyo3-object_store/src/url.rs create mode 100644 vendor/pyo3-object_store/type-hints/__init__.pyi create mode 100644 vendor/pyo3-object_store/type-hints/_aws.pyi create mode 100644 vendor/pyo3-object_store/type-hints/_azure.pyi create mode 100644 vendor/pyo3-object_store/type-hints/_client.pyi create mode 100644 vendor/pyo3-object_store/type-hints/_gcs.pyi create mode 100644 vendor/pyo3-object_store/type-hints/_http.pyi create mode 100644 vendor/pyo3-object_store/type-hints/_retry.pyi create mode 100644 vortex-object-store-opendal/Cargo.toml create mode 100644 vortex-object-store-opendal/src/lib.rs create mode 100644 vortex-python/python/vortex/store/_cos.py create mode 100644 vortex-python/src/opendal_store.rs diff --git a/Cargo.lock b/Cargo.lock index dd35d2f4725..c6a9d69ed4c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4358,10 +4358,12 @@ dependencies = [ "jiff-core", "jiff-static", "jiff-tzdb-platform", + "js-sys", "log", "portable-atomic", "portable-atomic-util", "serde_core", + "wasm-bindgen", "windows-link", ] @@ -5487,6 +5489,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.3" @@ -5941,6 +5952,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" @@ -5978,6 +6006,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" @@ -6765,8 +6848,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", @@ -7242,6 +7323,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" @@ -9860,7 +9989,14 @@ dependencies = [ "url", "vortex", "vortex-arrow", + "vortex-object-store-opendal", +>>>>>>> 64006a996 (feat(vortex-python): Add OpenDAL-backed CosStore to the Python object-store API) +======= "vortex-geo", + "vortex-object-store-opendal", +======= + "vortex-object-store-opendal", +>>>>>>> 64006a996 (feat(vortex-python): Add OpenDAL-backed CosStore to the Python object-store API) "vortex-parquet-variant", ] @@ -9955,6 +10091,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" @@ -10042,8 +10189,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 c2a3d1c20e8..eb95f0efc8d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,6 +12,7 @@ members = [ "vortex-flatbuffers", "vortex-metrics", "vortex-io", + "vortex-object-store-opendal", "vortex-proto", "vortex-array", "vortex-arrow", @@ -431,3 +432,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 50e30bb5aab..44078eaae4c 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -35,10 +35,16 @@ vortex = { workspace = true, features = ["object_store", "files"] } vortex-arrow = { workspace = true } vortex-geo = { 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); + } +} From 0111496b0ef4eb5c3769b2f5255249f83ebecc8a Mon Sep 17 00:00:00 2001 From: forwardxu Date: Sat, 25 Jul 2026 08:24:15 +0800 Subject: [PATCH 2/8] refactor(vortex-python): drop vendored pyo3-object_store fork The OpenDAL-backed CosStore now holds Arc directly instead of relying on the vendored fork's From> for PyObjectStore constructor. This removes the [patch.crates-io] fork of pyo3-object_store entirely; the upstream 0.11.0 crate (matched to the workspace's object_store 0.13.x) is used for the built-in store classes, preserving obstore interoperability. - Remove vendor/pyo3-object_store (forked source + type hints) - Remove [patch.crates-io] section and its comment from Cargo.toml - Pin pyo3-object_store to 0.11.0 (object_store 0.13.x, no trait conflict) - CosStore stores Arc; drop PyObjectStore dep - Fix clippy clone_on_ref_ptr in to_arc Signed-off-by: forwardxu --- Cargo.lock | 2 + Cargo.toml | 9 +- vendor/pyo3-object_store/Cargo.toml | 47 -- vendor/pyo3-object_store/LICENSE | 21 - vendor/pyo3-object_store/README.md | 79 --- vendor/pyo3-object_store/src/api.rs | 181 ----- .../pyo3-object_store/src/aws/credentials.rs | 216 ------ vendor/pyo3-object_store/src/aws/mod.rs | 4 - vendor/pyo3-object_store/src/aws/store.rs | 430 ------------ .../src/azure/credentials.rs | 313 --------- vendor/pyo3-object_store/src/azure/error.rs | 21 - vendor/pyo3-object_store/src/azure/mod.rs | 5 - vendor/pyo3-object_store/src/azure/store.rs | 489 -------------- vendor/pyo3-object_store/src/client.rs | 180 ----- vendor/pyo3-object_store/src/config.rs | 58 -- vendor/pyo3-object_store/src/credentials.rs | 101 --- vendor/pyo3-object_store/src/error.rs | 193 ------ .../pyo3-object_store/src/gcp/credentials.rs | 193 ------ vendor/pyo3-object_store/src/gcp/mod.rs | 4 - vendor/pyo3-object_store/src/gcp/store.rs | 380 ----------- vendor/pyo3-object_store/src/http.rs | 134 ---- vendor/pyo3-object_store/src/lib.rs | 35 - vendor/pyo3-object_store/src/local.rs | 143 ---- vendor/pyo3-object_store/src/memory.rs | 47 -- vendor/pyo3-object_store/src/path.rs | 64 -- vendor/pyo3-object_store/src/prefix.rs | 260 -------- vendor/pyo3-object_store/src/retry.rs | 103 --- vendor/pyo3-object_store/src/simple.rs | 112 ---- vendor/pyo3-object_store/src/store.rs | 276 -------- vendor/pyo3-object_store/src/url.rs | 64 -- .../pyo3-object_store/type-hints/__init__.pyi | 218 ------ vendor/pyo3-object_store/type-hints/_aws.pyi | 623 ------------------ .../pyo3-object_store/type-hints/_azure.pyi | 422 ------------ .../pyo3-object_store/type-hints/_client.pyi | 137 ---- vendor/pyo3-object_store/type-hints/_gcs.pyi | 235 ------- vendor/pyo3-object_store/type-hints/_http.pyi | 63 -- .../pyo3-object_store/type-hints/_retry.pyi | 97 --- vortex-python/src/opendal_store.rs | 29 +- 38 files changed, 14 insertions(+), 5974 deletions(-) delete mode 100644 vendor/pyo3-object_store/Cargo.toml delete mode 100644 vendor/pyo3-object_store/LICENSE delete mode 100644 vendor/pyo3-object_store/README.md delete mode 100644 vendor/pyo3-object_store/src/api.rs delete mode 100644 vendor/pyo3-object_store/src/aws/credentials.rs delete mode 100644 vendor/pyo3-object_store/src/aws/mod.rs delete mode 100644 vendor/pyo3-object_store/src/aws/store.rs delete mode 100644 vendor/pyo3-object_store/src/azure/credentials.rs delete mode 100644 vendor/pyo3-object_store/src/azure/error.rs delete mode 100644 vendor/pyo3-object_store/src/azure/mod.rs delete mode 100644 vendor/pyo3-object_store/src/azure/store.rs delete mode 100644 vendor/pyo3-object_store/src/client.rs delete mode 100644 vendor/pyo3-object_store/src/config.rs delete mode 100644 vendor/pyo3-object_store/src/credentials.rs delete mode 100644 vendor/pyo3-object_store/src/error.rs delete mode 100644 vendor/pyo3-object_store/src/gcp/credentials.rs delete mode 100644 vendor/pyo3-object_store/src/gcp/mod.rs delete mode 100644 vendor/pyo3-object_store/src/gcp/store.rs delete mode 100644 vendor/pyo3-object_store/src/http.rs delete mode 100644 vendor/pyo3-object_store/src/lib.rs delete mode 100644 vendor/pyo3-object_store/src/local.rs delete mode 100644 vendor/pyo3-object_store/src/memory.rs delete mode 100644 vendor/pyo3-object_store/src/path.rs delete mode 100644 vendor/pyo3-object_store/src/prefix.rs delete mode 100644 vendor/pyo3-object_store/src/retry.rs delete mode 100644 vendor/pyo3-object_store/src/simple.rs delete mode 100644 vendor/pyo3-object_store/src/store.rs delete mode 100644 vendor/pyo3-object_store/src/url.rs delete mode 100644 vendor/pyo3-object_store/type-hints/__init__.pyi delete mode 100644 vendor/pyo3-object_store/type-hints/_aws.pyi delete mode 100644 vendor/pyo3-object_store/type-hints/_azure.pyi delete mode 100644 vendor/pyo3-object_store/type-hints/_client.pyi delete mode 100644 vendor/pyo3-object_store/type-hints/_gcs.pyi delete mode 100644 vendor/pyo3-object_store/type-hints/_http.pyi delete mode 100644 vendor/pyo3-object_store/type-hints/_retry.pyi diff --git a/Cargo.lock b/Cargo.lock index c6a9d69ed4c..dc277010e4d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -6848,6 +6848,8 @@ 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", diff --git a/Cargo.toml b/Cargo.toml index eb95f0efc8d..040aa520992 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -433,11 +433,4 @@ incremental = false [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/vendor/pyo3-object_store/Cargo.toml b/vendor/pyo3-object_store/Cargo.toml deleted file mode 100644 index d27644b1605..00000000000 --- a/vendor/pyo3-object_store/Cargo.toml +++ /dev/null @@ -1,47 +0,0 @@ -[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 deleted file mode 100644 index 7bca320a280..00000000000 --- a/vendor/pyo3-object_store/LICENSE +++ /dev/null @@ -1,21 +0,0 @@ -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 deleted file mode 100644 index f6177685ea7..00000000000 --- a/vendor/pyo3-object_store/README.md +++ /dev/null @@ -1,79 +0,0 @@ -# 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 deleted file mode 100644 index ad22d32f566..00000000000 --- a/vendor/pyo3-object_store/src/api.rs +++ /dev/null @@ -1,181 +0,0 @@ -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 deleted file mode 100644 index 11dcc0b4bdd..00000000000 --- a/vendor/pyo3-object_store/src/aws/credentials.rs +++ /dev/null @@ -1,216 +0,0 @@ -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 deleted file mode 100644 index a4871db182e..00000000000 --- a/vendor/pyo3-object_store/src/aws/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -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 deleted file mode 100644 index f8041c4a08d..00000000000 --- a/vendor/pyo3-object_store/src/aws/store.rs +++ /dev/null @@ -1,430 +0,0 @@ -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 deleted file mode 100644 index 48890fab920..00000000000 --- a/vendor/pyo3-object_store/src/azure/credentials.rs +++ /dev/null @@ -1,313 +0,0 @@ -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 deleted file mode 100644 index 8b39dacb659..00000000000 --- a/vendor/pyo3-object_store/src/azure/error.rs +++ /dev/null @@ -1,21 +0,0 @@ -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 deleted file mode 100644 index b413a914c97..00000000000 --- a/vendor/pyo3-object_store/src/azure/mod.rs +++ /dev/null @@ -1,5 +0,0 @@ -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 deleted file mode 100644 index 97420416e12..00000000000 --- a/vendor/pyo3-object_store/src/azure/store.rs +++ /dev/null @@ -1,489 +0,0 @@ -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 deleted file mode 100644 index 9fe159fb0de..00000000000 --- a/vendor/pyo3-object_store/src/client.rs +++ /dev/null @@ -1,180 +0,0 @@ -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 deleted file mode 100644 index 29c2acd6a9f..00000000000 --- a/vendor/pyo3-object_store/src/config.rs +++ /dev/null @@ -1,58 +0,0 @@ -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 deleted file mode 100644 index 2b743fb8f65..00000000000 --- a/vendor/pyo3-object_store/src/credentials.rs +++ /dev/null @@ -1,101 +0,0 @@ -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 deleted file mode 100644 index 4226d241964..00000000000 --- a/vendor/pyo3-object_store/src/error.rs +++ /dev/null @@ -1,193 +0,0 @@ -//! 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 deleted file mode 100644 index 06c8c1cb90d..00000000000 --- a/vendor/pyo3-object_store/src/gcp/credentials.rs +++ /dev/null @@ -1,193 +0,0 @@ -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 deleted file mode 100644 index 78774709f24..00000000000 --- a/vendor/pyo3-object_store/src/gcp/mod.rs +++ /dev/null @@ -1,4 +0,0 @@ -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 deleted file mode 100644 index 260708dfaf8..00000000000 --- a/vendor/pyo3-object_store/src/gcp/store.rs +++ /dev/null @@ -1,380 +0,0 @@ -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 deleted file mode 100644 index 10d4a7f8237..00000000000 --- a/vendor/pyo3-object_store/src/http.rs +++ /dev/null @@ -1,134 +0,0 @@ -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 deleted file mode 100644 index 0d7d7fdfdee..00000000000 --- a/vendor/pyo3-object_store/src/lib.rs +++ /dev/null @@ -1,35 +0,0 @@ -#![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 deleted file mode 100644 index fc36df91b33..00000000000 --- a/vendor/pyo3-object_store/src/local.rs +++ /dev/null @@ -1,143 +0,0 @@ -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 deleted file mode 100644 index d726a6c3001..00000000000 --- a/vendor/pyo3-object_store/src/memory.rs +++ /dev/null @@ -1,47 +0,0 @@ -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 deleted file mode 100644 index 55c373214c5..00000000000 --- a/vendor/pyo3-object_store/src/path.rs +++ /dev/null @@ -1,64 +0,0 @@ -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 deleted file mode 100644 index db84807a93e..00000000000 --- a/vendor/pyo3-object_store/src/prefix.rs +++ /dev/null @@ -1,260 +0,0 @@ -//! 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 deleted file mode 100644 index e7079ed0313..00000000000 --- a/vendor/pyo3-object_store/src/retry.rs +++ /dev/null @@ -1,103 +0,0 @@ -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 deleted file mode 100644 index 1966756bad1..00000000000 --- a/vendor/pyo3-object_store/src/simple.rs +++ /dev/null @@ -1,112 +0,0 @@ -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 deleted file mode 100644 index bdb48c5576e..00000000000 --- a/vendor/pyo3-object_store/src/store.rs +++ /dev/null @@ -1,276 +0,0 @@ -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 deleted file mode 100644 index a67b0000630..00000000000 --- a/vendor/pyo3-object_store/src/url.rs +++ /dev/null @@ -1,64 +0,0 @@ -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 deleted file mode 100644 index 39eb60ec542..00000000000 --- a/vendor/pyo3-object_store/type-hints/__init__.pyi +++ /dev/null @@ -1,218 +0,0 @@ -# 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 deleted file mode 100644 index eb625a793bd..00000000000 --- a/vendor/pyo3-object_store/type-hints/_aws.pyi +++ /dev/null @@ -1,623 +0,0 @@ -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 deleted file mode 100644 index c27421551d8..00000000000 --- a/vendor/pyo3-object_store/type-hints/_azure.pyi +++ /dev/null @@ -1,422 +0,0 @@ -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 deleted file mode 100644 index 6f1f0593893..00000000000 --- a/vendor/pyo3-object_store/type-hints/_client.pyi +++ /dev/null @@ -1,137 +0,0 @@ -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 deleted file mode 100644 index 089f34d75c7..00000000000 --- a/vendor/pyo3-object_store/type-hints/_gcs.pyi +++ /dev/null @@ -1,235 +0,0 @@ -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 deleted file mode 100644 index 3558709d73a..00000000000 --- a/vendor/pyo3-object_store/type-hints/_http.pyi +++ /dev/null @@ -1,63 +0,0 @@ -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 deleted file mode 100644 index bfe4068325b..00000000000 --- a/vendor/pyo3-object_store/type-hints/_retry.pyi +++ /dev/null @@ -1,97 +0,0 @@ -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-python/src/opendal_store.rs b/vortex-python/src/opendal_store.rs index 63acc4c5544..34f572f50de 100644 --- a/vortex-python/src/opendal_store.rs +++ b/vortex-python/src/opendal_store.rs @@ -1,8 +1,7 @@ // 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. +//! Python-facing 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 @@ -16,18 +15,20 @@ 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 { +/// Build an `object_store::ObjectStore` 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)) + Ok(store) } /// A Tencent Cloud COS object store, backed by OpenDAL. @@ -37,19 +38,13 @@ fn build_store(url: &str, properties: HashMap) -> PyResult, } 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 + Arc::clone(&self.store) } } @@ -105,8 +100,7 @@ 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). + /// handed to `read_url`/`write`. #[test] fn cos_store_builds_object_store() { let store = CosStore::new( @@ -118,8 +112,7 @@ mod tests { 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). + // The store must be usable as an `Arc`. let arc: Arc = store.to_arc(); assert!(Arc::strong_count(&arc) >= 1); } From 188f621a77d6f1ce9911c0fac34d84358dcd3288 Mon Sep 17 00:00:00 2001 From: forwardxu Date: Sat, 25 Jul 2026 09:16:21 +0800 Subject: [PATCH 3/8] style: run taplo fmt on touched Cargo.toml files Fix taplo format check failures in CI: sort dependency tables alphabetically and drop trailing blank lines in Cargo.toml, vortex-jni/Cargo.toml, vortex-object-store-opendal/Cargo.toml, and vortex-python/Cargo.toml. Signed-off-by: forwardxu --- Cargo.toml | 2 -- vortex-jni/Cargo.toml | 2 +- vortex-object-store-opendal/Cargo.toml | 2 +- vortex-python/Cargo.toml | 2 +- 4 files changed, 3 insertions(+), 5 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 040aa520992..a9d8bfefb85 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -432,5 +432,3 @@ incremental = false # This improved build times significantly for default common cases that we use locally [profile.dev.package.vortex-fastlanes] debug = false - - diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index 44078eaae4c..6e50588b1ca 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -34,8 +34,8 @@ url = { workspace = true } vortex = { workspace = true, features = ["object_store", "files"] } vortex-arrow = { workspace = true } vortex-geo = { workspace = true } -vortex-parquet-variant = { workspace = true } vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } +vortex-parquet-variant = { workspace = true } [dev-dependencies] jni = { workspace = true, features = ["invocation"] } diff --git a/vortex-object-store-opendal/Cargo.toml b/vortex-object-store-opendal/Cargo.toml index 10ffec397e8..184c712c0b6 100644 --- a/vortex-object-store-opendal/Cargo.toml +++ b/vortex-object-store-opendal/Cargo.toml @@ -15,10 +15,10 @@ categories = { workspace = true } [dependencies] object_store = { workspace = true, features = ["cloud"] } +object_store_opendal = "0.57.0" opendal = { version = "0.57.0", default-features = false, features = [ "services-cos", ] } -object_store_opendal = "0.57.0" url = { workspace = true } vortex-utils = { workspace = true } diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index 896519d16c2..34a28b3b568 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -58,8 +58,8 @@ 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 } +vortex-utils = { workspace = true } [dev-dependencies] vortex-array = { path = "../vortex-array", features = ["_test-harness"] } From ebf0a94198484dcb8626a3c3fa1e81ada7597d64 Mon Sep 17 00:00:00 2001 From: forwardxu Date: Sat, 25 Jul 2026 10:18:57 +0800 Subject: [PATCH 4/8] refactor(opendal): harden CosStore API and OpenDAL COS client caching Prefer CosConfig over URL/property round-trips, cache COS clients in JNI/Python registries, and document store= usage after dropping CosStore.apply(). Signed-off-by: forwardxu --- Cargo.lock | 1 + docs/api/python/store/opendal.rst | 21 +- vortex-jni/Cargo.toml | 4 +- vortex-jni/src/object_store.rs | 29 +-- vortex-object-store-opendal/Cargo.toml | 1 + vortex-object-store-opendal/src/lib.rs | 204 +++++++++++++++--- vortex-python/Cargo.toml | 8 +- .../python/vortex/_lib/store/__init__.pyi | 3 - vortex-python/python/vortex/store/__init__.py | 13 +- vortex-python/python/vortex/store/_cos.py | 76 ++----- vortex-python/src/io.rs | 4 +- vortex-python/src/object_store/registry.rs | 44 ++-- vortex-python/src/opendal_store.rs | 96 ++++----- 13 files changed, 308 insertions(+), 196 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index dc277010e4d..8ba8297a1d0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10100,6 +10100,7 @@ dependencies = [ "object_store", "object_store_opendal", "opendal", + "tracing", "url", "vortex-utils", ] diff --git a/docs/api/python/store/opendal.rst b/docs/api/python/store/opendal.rst index 96f23d51622..8ec0a3b8fe2 100644 --- a/docs/api/python/store/opendal.rst +++ b/docs/api/python/store/opendal.rst @@ -1,6 +1,6 @@ -======== +============= OpenDAL (COS) -======== +============= Vortex can read from and write to Tencent Cloud COS through `OpenDAL `_, which provides native service support. @@ -24,20 +24,24 @@ variables OpenDAL's COS builder reads (``TENCENTCLOUD_SECRET_ID``, ``TENCENTCLOU a = vx.io.read_url("cos://my-bucket/path/to/dataset.vortex") -Or configure explicitly with :class:`~vortex.store.CosStore` before reading: +Or configure explicitly with :class:`~vortex.store.CosStore` and pass it to +:func:`vortex.io.read_url` via ``store=``: .. code-block:: python + from vortex.io import read_url from vortex.store import CosStore - CosStore( + store = 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") + # When `store=` is supplied, the path is a key within the store, so the scheme and + # bucket are not part of the path passed to read_url. + a = read_url("path/to/dataset.vortex", store=store) Passing a store object directly =============================== @@ -58,5 +62,6 @@ exactly like the built-in S3/Azure/GCS stores: secret_key="...", ) - a = read_url("cos://my-bucket/path/to/dataset.vortex", store=store) - + # When `store=` is supplied, the path is resolved as a key within the store, so the scheme + # and bucket are not part of the path passed to read_url. + a = read_url("path/to/dataset.vortex", store=store) diff --git a/vortex-jni/Cargo.toml b/vortex-jni/Cargo.toml index 6e50588b1ca..80f46a67b82 100644 --- a/vortex-jni/Cargo.toml +++ b/vortex-jni/Cargo.toml @@ -41,8 +41,8 @@ vortex-parquet-variant = { workspace = true } 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. +# Enable OpenDAL-backed object stores (Tencent COS) for `cos://` URLs. This pulls in the +# `opendal` dependency, so it is opt-in. opendal = ["dep:vortex-object-store-opendal"] [lib] diff --git a/vortex-jni/src/object_store.rs b/vortex-jni/src/object_store.rs index 127b4b05731..8e4938d661a 100644 --- a/vortex-jni/src/object_store.rs +++ b/vortex-jni/src/object_store.rs @@ -50,19 +50,9 @@ pub(crate) fn make_object_store( ) -> VortexResult> { 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)))?; - + // The cache key depends only on the URL authority + sorted properties, so we can hoist it + // above the scheme dispatch. This lets every store (including OpenDAL-backed ones) share + // a single client across repeated requests against the same bucket/configuration. let cache_key = url_cache_key(url, properties); { @@ -72,6 +62,19 @@ pub(crate) fn make_object_store( // guard dropped at close of scope } + // OpenDAL-backed stores (Tencent COS) use schemes that `object_store` does not recognize + // natively. Resolve them via the optional `opendal` feature, and cache the result so + // subsequent calls for the same URL share a single client. + #[cfg(feature = "opendal")] + if url.scheme() == "cos" { + 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)))?; + // Configure extra properties on that scheme instead. let store: Arc = match scheme { ObjectStoreScheme::Local => { diff --git a/vortex-object-store-opendal/Cargo.toml b/vortex-object-store-opendal/Cargo.toml index 184c712c0b6..79005633382 100644 --- a/vortex-object-store-opendal/Cargo.toml +++ b/vortex-object-store-opendal/Cargo.toml @@ -19,6 +19,7 @@ object_store_opendal = "0.57.0" opendal = { version = "0.57.0", default-features = false, features = [ "services-cos", ] } +tracing = { workspace = true } url = { workspace = true } vortex-utils = { workspace = true } diff --git a/vortex-object-store-opendal/src/lib.rs b/vortex-object-store-opendal/src/lib.rs index 69e446b59c8..1aaf0dba1d6 100644 --- a/vortex-object-store-opendal/src/lib.rs +++ b/vortex-object-store-opendal/src/lib.rs @@ -9,6 +9,14 @@ //! 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`. +//! +//! # Limitations +//! +//! The [`OpendalStore`] bridge owns its own HTTP request client. Configuration that the JNI/Python +//! layers normally pass through [`object_store::ClientOptions`] — connect/request timeouts, +//! retries, proxy settings, `allow_http` — has no effect on COS URLs handled here. Properties +//! that are not recognized by this crate are dropped without a warning; callers that need strict +//! validation must pre-filter their property maps. use std::sync::Arc; @@ -16,6 +24,7 @@ use object_store::ObjectStore; use object_store_opendal::OpendalStore; use opendal::Operator; use opendal::services; +use tracing::warn; use url::Url; use vortex_utils::aliases::hash_map::HashMap; @@ -58,7 +67,63 @@ impl From for object_store::Error { /// Schemes handled by this crate. pub const COS_SCHEME: &str = "cos"; -/// Build an [`object_store::ObjectStore`] for a COS URL. +/// Strongly-typed configuration for building a [`OpendalStore`] against Tencent Cloud COS. +/// +/// The fields mirror the keyword arguments of the `CosStore` Python class. Building from a +/// `CosConfig` avoids the URL-round-trip that the legacy `make_opendal_store(url, properties)` +/// entry point used, and is the preferred way to construct a COS store. +#[derive(Debug, Clone, Default)] +pub struct CosConfig { + /// COS bucket name. May be overridden by `properties["bucket"]` when adapting a URL. + pub bucket: String, + /// COS endpoint, e.g. `https://cos.ap-guangzhou.myqcloud.com`. + pub endpoint: String, + /// Tencent Cloud secret id (mapped to `TENCENTCLOUD_SECRET_ID`). + pub secret_id: Option, + /// Tencent Cloud secret key (mapped to `TENCENTCLOUD_SECRET_KEY`). + pub secret_key: Option, + /// Optional root prefix applied to all operations. + pub root: Option, + /// When `true`, disable OpenDAL's automatic config loading (so only the explicit + /// configuration is used). + pub disable_config_load: bool, +} + +/// Build an [`object_store::ObjectStore`] for Tencent Cloud COS directly from a [`CosConfig`]. +/// +/// This is the preferred entry point for callers that have a strongly-typed configuration object +/// (such as the `CosStore` pyclass in `vortex-python`). It does not synthesize a URL and so is not +/// fragile against reordering of `bucket` vs URL host precedence. +pub fn make_cos_store(config: CosConfig) -> Result, OpenDALStoreError> { + if config.bucket.is_empty() { + return Err(OpenDALStoreError::MissingConfig("bucket")); + } + if config.endpoint.is_empty() { + return Err(OpenDALStoreError::MissingConfig("endpoint")); + } + + let mut builder = services::Cos::default() + .bucket(&config.bucket) + .endpoint(&config.endpoint); + + if let Some(root) = config.root.as_deref() { + builder = builder.root(root); + } + if let Some(secret_id) = config.secret_id.as_deref() { + builder = builder.secret_id(secret_id); + } + if let Some(secret_key) = config.secret_key.as_deref() { + builder = builder.secret_key(secret_key); + } + if config.disable_config_load { + builder = builder.disable_config_load(); + } + + let operator = build_operator(builder)?; + Ok(Arc::new(OpendalStore::new(operator))) +} + +/// 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 @@ -69,43 +134,70 @@ 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())), + if url.scheme() != COS_SCHEME { + return Err(OpenDALStoreError::UnsupportedScheme( + url.scheme().to_string(), + )); } + + // Translate the (URL, properties) pair into a strongly-typed `CosConfig` and delegate to + // `make_cos_store`. Bucket is taken from `properties["bucket"]` first (matching the historical + // precedence) and falls back to the URL host; endpoint is taken from `properties["endpoint"]` + // first and falls back to `COS_ENDPOINT`; credentials fall back to the variables that + // OpenDAL's COS builder reads. + let config = url_and_properties_to_cos_config(url, properties)?; + make_cos_store(config) } -fn make_cos_store( +fn url_and_properties_to_cos_config( url: &Url, properties: &HashMap, -) -> Result, OpenDALStoreError> { +) -> Result { + warn_on_unknown_properties(properties); + 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() + .or_else(|| std::env::var("COS_ENDPOINT").ok()) .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; - let mut builder = services::Cos::default().bucket(&bucket).endpoint(&endpoint); + Ok(CosConfig { + bucket, + endpoint, + secret_id: properties + .get("secret_id") + .cloned() + .or_else(|| std::env::var("TENCENTCLOUD_SECRET_ID").ok()), + secret_key: properties + .get("secret_key") + .cloned() + .or_else(|| std::env::var("TENCENTCLOUD_SECRET_KEY").ok()), + root: properties.get("root").cloned(), + disable_config_load: properties.get("disable_config_load").map(String::as_str) + == Some("true"), + }) +} - 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(); +fn warn_on_unknown_properties(properties: &HashMap) { + const KNOWN: &[&str] = &[ + "bucket", + "endpoint", + "secret_id", + "secret_key", + "root", + "disable_config_load", + ]; + for key in properties.keys() { + if !KNOWN.contains(&key.as_str()) { + warn!("ignoring unknown OpenDAL store property: {key}"); + } } - - let operator = build_operator(builder)?; - Ok(Arc::new(OpendalStore::new(operator))) } fn build_operator(builder: B) -> Result @@ -136,25 +228,75 @@ mod tests { fn cos_requires_endpoint() { let url = Url::parse("cos://my-bucket/path").unwrap(); let props = HashMap::new(); + // Clear any `COS_ENDPOINT` left over from the environment so the test is deterministic. + // SAFETY: this is a unit test, single-threaded. + let previous = std::env::var("COS_ENDPOINT").ok(); + unsafe { std::env::remove_var("COS_ENDPOINT") }; + let result = make_opendal_store(&url, &props); + if let Some(value) = previous { + unsafe { std::env::set_var("COS_ENDPOINT", value) }; + } assert!(matches!( - make_cos_store(&url, &props), + result, Err(OpenDALStoreError::MissingConfig("endpoint")) )); } + /// With explicit credentials and `disable_config_load=true` so that no environment variables + /// are consulted, the store should build successfully. The bucket is taken from the URL host. #[test] - fn cos_derives_bucket_from_host() { + fn cos_builds_with_explicit_config() { + let url = Url::parse("cos://my-bucket/path/to/dataset.vortex").unwrap(); + let mut props = HashMap::new(); + props.insert("endpoint".to_string(), "https://example.com".to_string()); + props.insert("secret_id".to_string(), "AKID".to_string()); + props.insert("secret_key".to_string(), "secret".to_string()); + props.insert("disable_config_load".to_string(), "true".to_string()); + + let store = make_opendal_store(&url, &props).expect("store should build"); + // Sanity: the returned store is a non-null `Arc`. + assert!(Arc::strong_count(&store) >= 1); + } + + /// When `properties` does not contain `endpoint`, the builder must fall back to the + /// `COS_ENDPOINT` environment variable instead of erroring out. + #[test] + fn cos_falls_back_to_cos_endpoint_env() { let url = Url::parse("cos://my-bucket/path").unwrap(); + // SAFETY: this is a unit test, single-threaded. + unsafe { std::env::set_var("COS_ENDPOINT", "https://example.com") }; let mut props = HashMap::new(); - props.insert( - "endpoint".to_string(), - "https://cos.ap-guangzhou.myqcloud.com".to_string(), + props.insert("secret_id".to_string(), "AKID".to_string()); + props.insert("secret_key".to_string(), "secret".to_string()); + props.insert("disable_config_load".to_string(), "true".to_string()); + + let result = make_opendal_store(&url, &props); + // SAFETY: this is a unit test, single-threaded. + unsafe { std::env::remove_var("COS_ENDPOINT") }; + assert!( + result.is_ok(), + "expected store to build with COS_ENDPOINT set" ); - // 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(_)) + } + + /// The strongly-typed `make_cos_store` entry point must reject an empty bucket or endpoint + /// with a `MissingConfig` error before consulting any environment or builder. + #[test] + fn cos_config_rejects_empty_fields() { + assert!(matches!( + make_cos_store(CosConfig { + bucket: String::new(), + ..CosConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("bucket")) + )); + assert!(matches!( + make_cos_store(CosConfig { + bucket: "b".to_string(), + endpoint: String::new(), + ..CosConfig::default() + }), + Err(OpenDALStoreError::MissingConfig("endpoint")) )); } } diff --git a/vortex-python/Cargo.toml b/vortex-python/Cargo.toml index 34a28b3b568..3447912861a 100644 --- a/vortex-python/Cargo.toml +++ b/vortex-python/Cargo.toml @@ -28,9 +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"] +# OpenDAL-backed object stores (Tencent COS) for `cos://` URLs. Pulls in `opendal`, so it is +# opt-in. +opendal = ["dep:vortex-object-store-opendal", "dep:vortex-utils"] [dependencies] arrow-array = { workspace = true } @@ -59,7 +59,7 @@ vortex-arrow = { workspace = true } vortex-object-store-opendal = { path = "../vortex-object-store-opendal", optional = true } vortex-python-abi = { path = "../vortex-python-abi" } vortex-tui = { workspace = true, optional = true } -vortex-utils = { workspace = true } +vortex-utils = { workspace = true, optional = true } [dev-dependencies] vortex-array = { path = "../vortex-array", features = ["_test-harness"] } diff --git a/vortex-python/python/vortex/_lib/store/__init__.pyi b/vortex-python/python/vortex/_lib/store/__init__.pyi index 2e6965872ff..95bc141a044 100644 --- a/vortex-python/python/vortex/_lib/store/__init__.pyi +++ b/vortex-python/python/vortex/_lib/store/__init__.pyi @@ -92,9 +92,6 @@ def from_url( # type: ignore[misc] # docstring in pyi file - `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 050d14766ff..c9127087bf0 100644 --- a/vortex-python/python/vortex/store/__init__.py +++ b/vortex-python/python/vortex/store/__init__.py @@ -23,15 +23,7 @@ from ._memory import MemoryStore from ._retry import BackoffConfig, RetryConfig -ObjectStore: TypeAlias = ( - AzureStore - | CosStore - | GCSStore - | HTTPStore - | S3Store - | LocalStore - | MemoryStore -) +ObjectStore: TypeAlias = AzureStore | CosStore | GCSStore | HTTPStore | S3Store | LocalStore | MemoryStore """All supported ObjectStore implementations.""" @@ -101,9 +93,6 @@ def from_url( # type: ignore[misc] # docstring in pyi file - ``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` diff --git a/vortex-python/python/vortex/store/_cos.py b/vortex-python/python/vortex/store/_cos.py index 8d69f5e4cb7..9a3699d7137 100644 --- a/vortex-python/python/vortex/store/_cos.py +++ b/vortex-python/python/vortex/store/_cos.py @@ -1,72 +1,32 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors -"""OpenDAL-backed object store for Tencent Cloud COS. +"""Tencent Cloud COS object store, backed by OpenDAL. -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. +This store is only available when Vortex is built with the ``opendal`` feature. +The class is re-exported from the native extension module; if the feature is +not enabled, instantiating :class:`CosStore` raises :class:`ImportError`. """ 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] +from typing import TYPE_CHECKING, Any - 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) +if TYPE_CHECKING: + from vortex._lib import CosStore - def apply(self) -> None: - """Export this store's configuration into the process environment. +try: + from vortex._lib import CosStore as CosStore # type: ignore[attr-defined, no-redef] +except ImportError: - 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: # type: ignore[no-redef] + """Placeholder; the real implementation requires the ``opendal`` feature.""" + def __init__(self, *args: Any, **kwargs: Any) -> None: + raise ImportError( + "CosStore requires Vortex to be built with the 'opendal' feature; " + "build with `maturin build --features opendal` or `maturin develop --features opendal`." + ) -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", - } +__all__ = ["CosStore"] diff --git a/vortex-python/src/io.rs b/vortex-python/src/io.rs index 1332a063f4c..a9dfa8c8f1a 100644 --- a/vortex-python/src/io.rs +++ b/vortex-python/src/io.rs @@ -193,7 +193,7 @@ impl<'py> FromPyObject<'_, 'py> for AnyVortexStore { /// path : str /// The file path. /// -/// 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 /// An optional object store configuration to use for writing the output. /// /// Examples @@ -335,7 +335,7 @@ impl PyVortexWriteOptions { /// path : str /// The file path. /// - /// 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 /// An optional object store configuration to use for writing the output. /// /// Examples diff --git a/vortex-python/src/object_store/registry.rs b/vortex-python/src/object_store/registry.rs index 17975b1e00e..240a73ced84 100644 --- a/vortex-python/src/object_store/registry.rs +++ b/vortex-python/src/object_store/registry.rs @@ -20,7 +20,7 @@ use object_store::registry::ObjectStoreRegistry; use parking_lot::RwLock; use url::Url; #[cfg(feature = "opendal")] -use vortex::utils::aliases::hash_map::HashMap as VortexHashMap; +use vortex_utils::aliases::hash_map::HashMap as VortexHashMap; #[derive(Debug, Default)] struct PathEntry { @@ -78,11 +78,25 @@ 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`). + // 1. Consult the user-registered map first. Every other scheme does this, so an explicit + // `Registry::register("cos://...", store)` should also win over the env-var / OpenDAL + // fallback below. This also means a previously-resolved OpenDAL store can be cached + // here for the lifetime of the process (the lookup at the bottom of this function + // inserts into the same map). + { + let map = self.map.read(); + + if let Some((store, depth)) = map.get(key).and_then(|entry| entry.lookup(to_resolve)) { + let path = path_suffix(to_resolve, depth)?; + return Ok((Arc::clone(store), path)); + } + } + + // 2. OpenDAL-backed schemes (Tencent COS) are not recognized by the `object_store` crate. + // 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") { + if to_resolve.scheme() == "cos" { let store = vortex_object_store_opendal::make_opendal_store(to_resolve, &VortexHashMap::new()) .map_err(|e| object_store::Error::Generic { @@ -95,16 +109,20 @@ impl ObjectStoreRegistry for Registry { source: Box::new(e), } })?; - return Ok((store, path)); - } - { - let map = self.map.read(); - - if let Some((store, depth)) = map.get(key).and_then(|entry| entry.lookup(to_resolve)) { - let path = path_suffix(to_resolve, depth)?; - return Ok((Arc::clone(store), path)); + // Cache the resolved store under the URL's authority so subsequent `resolve` calls + // for the same bucket/configuration share a single client. + let mut map = self.map.write(); + let mut entry = map.entry(key.to_string()).or_default(); + for segment in path_segments(to_resolve.path()) { + entry = entry.children.entry(segment.to_string()).or_default(); } + let stored = Arc::clone(match &entry.store { + None => entry.store.insert(Arc::clone(&store)), + Some(x) => x, // Racing creation - use existing + }); + + return Ok((stored, path)); } let normalized_env = std::env::vars().map(|(k, v)| (k.to_ascii_lowercase(), v)); diff --git a/vortex-python/src/opendal_store.rs b/vortex-python/src/opendal_store.rs index 34f572f50de..f0c8d868b79 100644 --- a/vortex-python/src/opendal_store.rs +++ b/vortex-python/src/opendal_store.rs @@ -15,28 +15,14 @@ use std::sync::Arc; use pyo3::exceptions::PyValueError; use pyo3::prelude::*; -use url::Url; -use vortex_utils::aliases::hash_map::HashMap; - -/// Build an `object_store::ObjectStore` 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(store) -} +use vortex_object_store_opendal::CosConfig; /// 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)] +#[pyclass(name = "CosStore", module = "vortex._lib", frozen, from_py_object)] +#[derive(Clone, Debug)] pub struct CosStore { store: Arc, } @@ -68,52 +54,62 @@ impl CosStore { 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)?; + let config = CosConfig { + bucket, + endpoint, + secret_id, + secret_key, + root, + disable_config_load, + }; + let store = vortex_object_store_opendal::make_cos_store(config) + .map_err(|e| PyValueError::new_err(e.to_string()))?; 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"))] +#[cfg(test)] mod tests { + use vortex_object_store_opendal::CosConfig; + use super::*; - /// A `CosStore` built in Rust must yield a non-null `Arc` that can be - /// handed to `read_url`/`write`. + /// `make_cos_store` (the strongly-typed entry point used by the `CosStore` pyclass) must + /// accept a fully-specified `CosConfig` and yield a usable `Arc`. + #[test] + fn cos_config_builds_object_store() { + let config = CosConfig { + bucket: "my-bucket".to_string(), + endpoint: "https://cos.ap-guangzhou.myqcloud.com".to_string(), + secret_id: Some("AKID".to_string()), + secret_key: Some("secret".to_string()), + root: Some("nested/prefix".to_string()), + disable_config_load: true, + }; + let store = + vortex_object_store_opendal::make_cos_store(config).expect("cos store should build"); + // Sanity-check the result is a non-null `Arc`. + assert!(Arc::strong_count(&store) >= 1); + } + + /// Constructing a `CosConfig` with an empty `bucket` must surface as a `MissingConfig` error + /// rather than silently building an invalid store. #[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`. - let arc: Arc = store.to_arc(); - assert!(Arc::strong_count(&arc) >= 1); + fn cos_config_rejects_empty_bucket() { + let result = vortex_object_store_opendal::make_cos_store(CosConfig { + bucket: String::new(), + endpoint: "https://cos.ap-guangzhou.myqcloud.com".to_string(), + ..CosConfig::default() + }); + assert!(matches!( + result, + Err(vortex_object_store_opendal::OpenDALStoreError::MissingConfig("bucket")) + )); } } From 35dd513e9633f6229a7e8588e8400109d38411b9 Mon Sep 17 00:00:00 2001 From: forwardxu Date: Sun, 26 Jul 2026 09:01:51 +0800 Subject: [PATCH 5/8] fix(opendal): cache cos:// per bucket, inject env lookup, declare CosStore stub - registry.rs: route cos:// through a unified build_store + cache_and_resolve pipeline, with the OpenDAL branch rooted at the bucket (depth 0) like s3:// so sibling objects share a client. Pull the per-key path-walk into an entry_at helper that register and cache_and_resolve both reuse, and add Arc::ptr_eq-based tests pinning the symmetry from both sides (per-key cache, per-bucket client, prefix registration at depth > 0). - vortex-object-store-opendal: make the env-var fallback injectable so the two cos tests no longer race against the global environment under cargo test's default multi-threaded runner (Rust 2024 makes std::env::set_var unsafe for exactly this reason). Merge the two flaky tests into a single endpoint-required + env-fallback pair driven by a fixed lookup map. - vortex/_lib/__init__.pyi: declare CosStore with its full constructor signature so basedpyright sees a real symbol (fixes 3 reportAttributeAccessIssue errors against _cos.py), and so .. autoclass:: in the docs page renders the real parameters instead of the runtime placeholder. - Drop now-unnecessary # pyright: ignore[reportMissingModuleSource] comments on the store/_*.py files that import vortex._lib. Signed-off-by: forwardxu --- vortex-object-store-opendal/src/lib.rs | 98 +++++---- vortex-python/python/vortex/_lib/__init__.pyi | 20 ++ vortex-python/python/vortex/store/__init__.py | 2 +- vortex-python/python/vortex/store/_aws.py | 2 +- vortex-python/python/vortex/store/_azure.py | 2 +- vortex-python/python/vortex/store/_cos.py | 31 +-- vortex-python/python/vortex/store/_gcs.py | 2 +- vortex-python/python/vortex/store/_http.py | 2 +- vortex-python/python/vortex/store/_local.py | 2 +- vortex-python/python/vortex/store/_memory.py | 2 +- vortex-python/src/object_store/registry.rs | 202 ++++++++++++------ 11 files changed, 240 insertions(+), 125 deletions(-) diff --git a/vortex-object-store-opendal/src/lib.rs b/vortex-object-store-opendal/src/lib.rs index 1aaf0dba1d6..92caee88092 100644 --- a/vortex-object-store-opendal/src/lib.rs +++ b/vortex-object-store-opendal/src/lib.rs @@ -145,14 +145,33 @@ pub fn make_opendal_store( // precedence) and falls back to the URL host; endpoint is taken from `properties["endpoint"]` // first and falls back to `COS_ENDPOINT`; credentials fall back to the variables that // OpenDAL's COS builder reads. - let config = url_and_properties_to_cos_config(url, properties)?; + let config = url_and_properties_to_cos_config(url, properties, env_var_lookup)?; make_cos_store(config) } -fn url_and_properties_to_cos_config( +/// Default environment-variable lookup: reads `key` from the process environment. +/// +/// The lookup is factored out so tests can pass a fixed map instead of mutating the global +/// environment, which is unsound when `cargo test` runs tests on multiple threads within one +/// process (the `unsafe std::env::set_var` block became `unsafe` in Rust 2024 for exactly this +/// reason). +fn env_var_lookup(key: &str) -> Option { + std::env::var(key).ok() +} + +/// Translate the (URL, properties) pair into a strongly-typed [`CosConfig`]. +/// +/// `env_lookup` is the source of truth for environment-variable fallbacks. The production +/// entry points pass [`env_var_lookup`]; tests pass a closure that returns from a fixed map, so +/// they do not race against the process environment. +fn url_and_properties_to_cos_config( url: &Url, properties: &HashMap, -) -> Result { + env_lookup: F, +) -> Result +where + F: Fn(&str) -> Option, +{ warn_on_unknown_properties(properties); let bucket = properties @@ -164,7 +183,7 @@ fn url_and_properties_to_cos_config( let endpoint = properties .get("endpoint") .cloned() - .or_else(|| std::env::var("COS_ENDPOINT").ok()) + .or_else(|| env_lookup("COS_ENDPOINT")) .ok_or(OpenDALStoreError::MissingConfig("endpoint"))?; Ok(CosConfig { @@ -173,11 +192,11 @@ fn url_and_properties_to_cos_config( secret_id: properties .get("secret_id") .cloned() - .or_else(|| std::env::var("TENCENTCLOUD_SECRET_ID").ok()), + .or_else(|| env_lookup("TENCENTCLOUD_SECRET_ID")), secret_key: properties .get("secret_key") .cloned() - .or_else(|| std::env::var("TENCENTCLOUD_SECRET_KEY").ok()), + .or_else(|| env_lookup("TENCENTCLOUD_SECRET_KEY")), root: properties.get("root").cloned(), disable_config_load: properties.get("disable_config_load").map(String::as_str) == Some("true"), @@ -224,61 +243,56 @@ mod tests { )); } + /// The endpoint is required. With a fixed env-lookup that returns `None`, the call must + /// fail with `MissingConfig("endpoint")` regardless of what the process environment + /// happens to contain. This makes the test deterministic under `cargo test`'s default + /// multi-threaded runner (and avoids the `unsafe std::env::set_var` that became unsound in + /// Rust 2024). #[test] fn cos_requires_endpoint() { let url = Url::parse("cos://my-bucket/path").unwrap(); let props = HashMap::new(); - // Clear any `COS_ENDPOINT` left over from the environment so the test is deterministic. - // SAFETY: this is a unit test, single-threaded. - let previous = std::env::var("COS_ENDPOINT").ok(); - unsafe { std::env::remove_var("COS_ENDPOINT") }; - let result = make_opendal_store(&url, &props); - if let Some(value) = previous { - unsafe { std::env::set_var("COS_ENDPOINT", value) }; - } + let result = url_and_properties_to_cos_config(&url, &props, |_| None).unwrap_err(); assert!(matches!( result, - Err(OpenDALStoreError::MissingConfig("endpoint")) + OpenDALStoreError::MissingConfig("endpoint") )); } - /// With explicit credentials and `disable_config_load=true` so that no environment variables - /// are consulted, the store should build successfully. The bucket is taken from the URL host. + /// When `properties` does not contain `endpoint`, the env-lookup should be consulted. + /// The fixed lookup returns `"https://example.com"`, so the build must succeed. + #[test] + fn cos_falls_back_to_cos_endpoint_env() { + let url = Url::parse("cos://my-bucket/path").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://example.com".to_string()), + _ => None, + }; + let props = HashMap::new(); + let config = url_and_properties_to_cos_config(&url, &props, env).expect("config"); + assert_eq!(config.endpoint, "https://example.com"); + } + + /// With a fixed env-lookup that returns explicit credentials, the strongly-typed + /// `make_cos_store` should build a store successfully. #[test] fn cos_builds_with_explicit_config() { let url = Url::parse("cos://my-bucket/path/to/dataset.vortex").unwrap(); + let env = |key: &str| match key { + "COS_ENDPOINT" => Some("https://example.com".to_string()), + "TENCENTCLOUD_SECRET_ID" => Some("AKID".to_string()), + "TENCENTCLOUD_SECRET_KEY" => Some("secret".to_string()), + _ => None, + }; let mut props = HashMap::new(); - props.insert("endpoint".to_string(), "https://example.com".to_string()); - props.insert("secret_id".to_string(), "AKID".to_string()); - props.insert("secret_key".to_string(), "secret".to_string()); props.insert("disable_config_load".to_string(), "true".to_string()); - let store = make_opendal_store(&url, &props).expect("store should build"); + let config = url_and_properties_to_cos_config(&url, &props, env).expect("config"); + let store = make_cos_store(config).expect("store should build"); // Sanity: the returned store is a non-null `Arc`. assert!(Arc::strong_count(&store) >= 1); } - /// When `properties` does not contain `endpoint`, the builder must fall back to the - /// `COS_ENDPOINT` environment variable instead of erroring out. - #[test] - fn cos_falls_back_to_cos_endpoint_env() { - let url = Url::parse("cos://my-bucket/path").unwrap(); - // SAFETY: this is a unit test, single-threaded. - unsafe { std::env::set_var("COS_ENDPOINT", "https://example.com") }; - let mut props = HashMap::new(); - props.insert("secret_id".to_string(), "AKID".to_string()); - props.insert("secret_key".to_string(), "secret".to_string()); - props.insert("disable_config_load".to_string(), "true".to_string()); - - let result = make_opendal_store(&url, &props); - // SAFETY: this is a unit test, single-threaded. - unsafe { std::env::remove_var("COS_ENDPOINT") }; - assert!( - result.is_ok(), - "expected store to build with COS_ENDPOINT set" - ); - } - /// The strongly-typed `make_cos_store` entry point must reject an empty bucket or endpoint /// with a `MissingConfig` error before consulting any environment or builder. #[test] diff --git a/vortex-python/python/vortex/_lib/__init__.pyi b/vortex-python/python/vortex/_lib/__init__.pyi index 79f62c4521f..09c9a520406 100644 --- a/vortex-python/python/vortex/_lib/__init__.pyi +++ b/vortex-python/python/vortex/_lib/__init__.pyi @@ -1,3 +1,23 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright the Vortex contributors +class CosStore: + """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)``. + + This class is only available when Vortex is built with the ``opendal`` feature. + """ + + def __init__( + self, + bucket: str, + endpoint: str, + *, + secret_id: str | None = None, + secret_key: str | None = None, + root: str | None = None, + disable_config_load: bool = False, + ) -> None: ... diff --git a/vortex-python/python/vortex/store/__init__.py b/vortex-python/python/vortex/store/__init__.py index c9127087bf0..ba67fd72f3f 100644 --- a/vortex-python/python/vortex/store/__init__.py +++ b/vortex-python/python/vortex/store/__init__.py @@ -4,7 +4,7 @@ from collections.abc import Callable from typing import TypeAlias, Unpack, overload -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store from ._aws import S3Config, S3Credential, S3CredentialProvider, S3Store from ._azure import ( AzureAccessKey, diff --git a/vortex-python/python/vortex/store/_aws.py b/vortex-python/python/vortex/store/_aws.py index 3f3079e9902..e3c02bf042e 100644 --- a/vortex-python/python/vortex/store/_aws.py +++ b/vortex-python/python/vortex/store/_aws.py @@ -7,7 +7,7 @@ from typing_extensions import override -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_azure.py b/vortex-python/python/vortex/store/_azure.py index 1b3ec1b4e52..4d54531f928 100644 --- a/vortex-python/python/vortex/store/_azure.py +++ b/vortex-python/python/vortex/store/_azure.py @@ -7,7 +7,7 @@ from typing_extensions import override -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_cos.py b/vortex-python/python/vortex/store/_cos.py index 9a3699d7137..666e0c38b39 100644 --- a/vortex-python/python/vortex/store/_cos.py +++ b/vortex-python/python/vortex/store/_cos.py @@ -13,20 +13,25 @@ from typing import TYPE_CHECKING, Any if TYPE_CHECKING: + # `vortex._lib.__init__.pyi` declares `CosStore` so type checkers see a real + # signature; the runtime branch below provides the same symbol when the + # `opendal` feature is enabled, and a placeholder that raises on + # instantiation when it is not. from vortex._lib import CosStore - -try: - from vortex._lib import CosStore as CosStore # type: ignore[attr-defined, no-redef] -except ImportError: - - class CosStore: # type: ignore[no-redef] - """Placeholder; the real implementation requires the ``opendal`` feature.""" - - def __init__(self, *args: Any, **kwargs: Any) -> None: - raise ImportError( - "CosStore requires Vortex to be built with the 'opendal' feature; " - "build with `maturin build --features opendal` or `maturin develop --features opendal`." - ) +else: + try: + from vortex._lib import CosStore as CosStore + except ImportError: + + class CosStore: + """Placeholder; the real implementation requires the ``opendal`` feature.""" + + def __init__(self, *args: Any, **kwargs: Any) -> None: + raise ImportError( + "CosStore requires Vortex to be built with the 'opendal' feature; " + + "build with `maturin build --features opendal` " + + "or `maturin develop --features opendal`." + ) __all__ = ["CosStore"] diff --git a/vortex-python/python/vortex/store/_gcs.py b/vortex-python/python/vortex/store/_gcs.py index e1512434093..9d00055ba82 100644 --- a/vortex-python/python/vortex/store/_gcs.py +++ b/vortex-python/python/vortex/store/_gcs.py @@ -7,7 +7,7 @@ from typing_extensions import override -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_http.py b/vortex-python/python/vortex/store/_http.py index 57d46bb5ebb..5da9c6b2cd6 100644 --- a/vortex-python/python/vortex/store/_http.py +++ b/vortex-python/python/vortex/store/_http.py @@ -5,7 +5,7 @@ from typing_extensions import override -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_local.py b/vortex-python/python/vortex/store/_local.py index 891c8e968b4..de9670c2d2c 100644 --- a/vortex-python/python/vortex/store/_local.py +++ b/vortex-python/python/vortex/store/_local.py @@ -6,7 +6,7 @@ from typing_extensions import override -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store class LocalStore(_store.LocalStore): diff --git a/vortex-python/python/vortex/store/_memory.py b/vortex-python/python/vortex/store/_memory.py index 763424732eb..7dbd48b3d6f 100644 --- a/vortex-python/python/vortex/store/_memory.py +++ b/vortex-python/python/vortex/store/_memory.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # SPDX-FileCopyrightText: Copyright (c) 2024 Development Seed -from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] +from .._lib import store as _store class MemoryStore(_store.MemoryStore): diff --git a/vortex-python/src/object_store/registry.rs b/vortex-python/src/object_store/registry.rs index 240a73ced84..144a0d0e536 100644 --- a/vortex-python/src/object_store/registry.rs +++ b/vortex-python/src/object_store/registry.rs @@ -66,89 +66,88 @@ pub(crate) struct Registry { impl ObjectStoreRegistry for Registry { fn register(&self, url: Url, store: Arc) -> Option> { let mut map = self.map.write(); - let key = url_key(&url); - let mut entry = map.entry(key.to_string()).or_default(); - - for segment in path_segments(url.path()) { - entry = entry.children.entry(segment.to_string()).or_default(); - } + let entry = entry_at( + &mut map, + url_key(&url), + url.path(), + num_segments(url.path()), + ); entry.store.replace(store) } fn resolve(&self, to_resolve: &Url) -> object_store::Result<(Arc, Path)> { let key = url_key(to_resolve); - // 1. Consult the user-registered map first. Every other scheme does this, so an explicit - // `Registry::register("cos://...", store)` should also win over the env-var / OpenDAL - // fallback below. This also means a previously-resolved OpenDAL store can be cached - // here for the lifetime of the process (the lookup at the bottom of this function - // inserts into the same map). + // 1. Look up the user-registered map first. Every other scheme does this, so an explicit + // `Registry::register("cos://...", store)` should also win over the build-and-cache + // fallback below. This also means a previously-resolved store (built by us for the + // same URL) is served from the cache here, without rebuilding the client. { let map = self.map.read(); - if let Some((store, depth)) = map.get(key).and_then(|entry| entry.lookup(to_resolve)) { let path = path_suffix(to_resolve, depth)?; return Ok((Arc::clone(store), path)); } } - // 2. OpenDAL-backed schemes (Tencent COS) are not recognized by the `object_store` crate. - // Resolve them via the optional `opendal` feature, relying on OpenDAL's environment- - // variable configuration (e.g. `TENCENTCLOUD_SECRET_ID`). - #[cfg(feature = "opendal")] - if 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), - } - })?; - - // Cache the resolved store under the URL's authority so subsequent `resolve` calls - // for the same bucket/configuration share a single client. - let mut map = self.map.write(); - let mut entry = map.entry(key.to_string()).or_default(); - for segment in path_segments(to_resolve.path()) { - entry = entry.children.entry(segment.to_string()).or_default(); - } - let stored = Arc::clone(match &entry.store { - None => entry.store.insert(Arc::clone(&store)), - Some(x) => x, // Racing creation - use existing - }); - - return Ok((stored, path)); - } + // 2. Build the store and the path *within* that store. `build_store` reports the path + // the way `parse_url_opts` does — i.e. the path is what the store will see as the + // key — which is what keeps every scheme on the one caching rule in + // `cache_and_resolve`. The depth is then just the difference in segment counts. + let (store, path) = build_store(to_resolve)?; + let depth = num_segments(to_resolve.path()) - num_segments(path.as_ref()); - let normalized_env = std::env::vars().map(|(k, v)| (k.to_ascii_lowercase(), v)); + // 3. Cache the store and return it alongside the path that lives inside the store. + self.cache_and_resolve(to_resolve, store, depth) + } +} - if let Ok((store, path)) = parse_url_opts(to_resolve, normalized_env) { - let depth = num_segments(to_resolve.path()) - num_segments(path.as_ref()); +impl Registry { + /// Caches `store` as the store serving the first `depth` path segments of `to_resolve`, and + /// returns it alongside the remaining path. + /// + /// If a racing `resolve` cached a store at the same position first, that store wins and the + /// one just built is dropped, so all callers of a given prefix share a single client. + fn cache_and_resolve( + &self, + to_resolve: &Url, + store: Arc, + depth: usize, + ) -> object_store::Result<(Arc, Path)> { + let path = path_suffix(to_resolve, depth)?; - let mut map = self.map.write(); - let mut entry = map.entry(key.to_string()).or_default(); - for segment in path_segments(to_resolve.path()).take(depth) { - entry = entry.children.entry(segment.to_string()).or_default(); - } - let store = Arc::clone(match &entry.store { - None => entry.store.insert(Arc::from(store)), - Some(x) => x, // Racing creation - use existing - }); + let mut map = self.map.write(); + let entry = entry_at(&mut map, url_key(to_resolve), to_resolve.path(), depth); + let stored = Arc::clone(match &entry.store { + None => entry.store.insert(store), + Some(existing) => existing, // Racing creation - use existing + }); - let path = path_suffix(to_resolve, depth)?; - return Ok((store, path)); - } + Ok((stored, path)) + } +} - Err(object_store::Error::Generic { - store: "ObjectStoreRegistry", - source: "URL could not be resolved".into(), - }) +/// Builds the [`ObjectStore`] that serves `to_resolve`, with the path within that store. +/// +/// This mirrors [`parse_url_opts`], extended with schemes the `object_store` crate does not +/// recognize. Reporting the path the way `parse_url_opts` does is what keeps every scheme on +/// the one caching rule in [`Registry::resolve`]: a scheme says where its store is rooted, and +/// the registry decides how to cache it. +fn build_store(to_resolve: &Url) -> object_store::Result<(Arc, Path)> { + // OpenDAL-backed schemes (Tencent COS) are not recognized by `object_store`, so build them + // from OpenDAL's own environment-variable configuration (e.g. `TENCENTCLOUD_SECRET_ID`). + // The operator is rooted at the bucket, which lives in the URL authority, so — exactly as + // for `s3://bucket/path` — the whole URL path is the object key. + #[cfg(feature = "opendal")] + if to_resolve.scheme() == vortex_object_store_opendal::COS_SCHEME { + let store = + vortex_object_store_opendal::make_opendal_store(to_resolve, &VortexHashMap::new())?; + return Ok((store, Path::from_url_path(to_resolve.path())?)); } + + let normalized_env = std::env::vars().map(|(k, v)| (k.to_ascii_lowercase(), v)); + let (store, path) = parse_url_opts(to_resolve, normalized_env)?; + Ok((Arc::from(store), path)) } /// Extracts the scheme and authority of a URL (components before the Path) @@ -181,10 +180,27 @@ fn path_suffix(url: &Url, depth: usize) -> Result { Ok(path) } +/// Walks to the [`PathEntry`] for `key` sitting `depth` segments into `path`, creating the +/// intermediate entries as needed. +fn entry_at<'a>( + map: &'a mut HashMap, + key: &str, + path: &str, + depth: usize, +) -> &'a mut PathEntry { + let mut current = map.entry(key.to_string()).or_default(); + for segment in path_segments(path).take(depth) { + current = current.children.entry(segment.to_string()).or_default(); + } + current +} + #[cfg(test)] mod tests { use std::fmt::Write; + use std::sync::Arc; + use object_store::ObjectStore; use object_store::registry::ObjectStoreRegistry; use url::Url; @@ -230,4 +246,64 @@ mod tests { assert!(debug_str.contains("us-east-3")); }); } + + /// Two resolves of the same URL must return the same `Arc` (the cached client) and the same + /// path, and two different keys must get distinct stores. This pins the symmetry the registry + /// promises: cache hit returns the same client and the same key. + #[test] + fn test_resolve_url_caches_per_key() { + with_var("AWS_REGION", "us-east-3", || { + let registry = Registry::default(); + let first = Url::parse("s3://my-bucket/first/second").unwrap(); + let second = Url::parse("s3://my-bucket/first/second").unwrap(); + + let (store_a, path_a) = registry.resolve(&first).unwrap(); + let (store_b, path_b) = registry.resolve(&second).unwrap(); + assert!(Arc::ptr_eq(&store_a, &store_b)); + assert_eq!(path_a, path_b); + + // A different bucket gets its own client. + let other = Url::parse("s3://other-bucket/first/second").unwrap(); + let (store_c, _) = registry.resolve(&other).unwrap(); + assert!(!Arc::ptr_eq(&store_a, &store_c)); + }); + } + + /// Two objects in the same bucket must share a cached client. This is the regression that + /// the bespoke "walk every segment then return the whole key" path caused: by walking all + /// segments it stored the store at full depth, so the next resolve saw the whole key + /// already consumed and returned an empty path. Pinning the path here guards against that. + #[test] + fn test_resolve_url_shared_client_same_bucket() { + with_var("AWS_REGION", "us-east-3", || { + let registry = Registry::default(); + let a = Url::parse("s3://my-bucket/path/to/data.vortex").unwrap(); + let b = Url::parse("s3://my-bucket/other/data.vortex").unwrap(); + + let (store_a, path_a) = registry.resolve(&a).unwrap(); + let (store_b, path_b) = registry.resolve(&b).unwrap(); + assert!(Arc::ptr_eq(&store_a, &store_b)); + assert_eq!(path_a.as_ref(), "path/to/data.vortex"); + assert_eq!(path_b.as_ref(), "other/data.vortex"); + }); + } + + /// Pinning the `entry_at` helper at depth > 0 (the path-prefix case) protects the shared + /// `entry_at` / `cache_and_resolve` helpers from regressing prefix registration. + #[test] + fn test_register_at_prefix_shares_store() { + let registry = Registry::default(); + let prefix = Url::parse("s3://my-bucket/prefix/").unwrap(); + let (inner, _) = registry.resolve(&prefix).unwrap(); + let cached: Arc = Arc::clone(&inner); + let replaced = registry.register(prefix.clone(), Arc::clone(&cached)); + // First registration at this prefix replaces nothing. + assert!(replaced.is_none()); + + // A resolve at a deeper key still walks through the registered prefix's store. + let deeper = Url::parse("s3://my-bucket/prefix/inner/key").unwrap(); + let (store_deeper, path_deeper) = registry.resolve(&deeper).unwrap(); + assert!(Arc::ptr_eq(&store_deeper, &cached)); + assert_eq!(path_deeper.as_ref(), "inner/key"); + } } From 05c7212fa8d711a2b83352efcff9468ef7e27cae Mon Sep 17 00:00:00 2001 From: forwardxu Date: Sun, 26 Jul 2026 09:10:01 +0800 Subject: [PATCH 6/8] docs(opendal): render CosStore signature by hand for default builds The previous `.. autoclass:: vortex.store.CosStore` directive rendered the `_cos.py` placeholder (`(*args, **kwargs) -> None`, "Placeholder; the real implementation requires the ``opendal`` feature") on default docs builds, because the docs CI builds vortex without the opendal feature, so autodoc sees only the placeholder. Even with `CosStore` declared in `vortex/_lib/__init__.pyi`, autodoc reads the runtime class, not the stub. Replace the autoclass with a hand-written `.. py:class::` directive that embeds the real signature and per-parameter descriptions. This is the "Failing that" branch in the upstream review: write the signature and parameters into the page by hand instead of relying on autoclass. Verified locally with sphinx-build 9.1.0: the rendered HTML now shows `vortex.store.CosStore(bucket, endpoint, *, secret_id=None, secret_key=None, root=None, disable_config_load=False)` regardless of which feature flags the build environment enables. Signed-off-by: forwardxu --- docs/api/python/store/opendal.rst | 24 ++++++++++++++++++++++-- 1 file changed, 22 insertions(+), 2 deletions(-) diff --git a/docs/api/python/store/opendal.rst b/docs/api/python/store/opendal.rst index 8ec0a3b8fe2..bac62f993be 100644 --- a/docs/api/python/store/opendal.rst +++ b/docs/api/python/store/opendal.rst @@ -8,8 +8,28 @@ Vortex can read from and write to Tencent Cloud COS through 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: +:class:`vortex.store.CosStore` +============================== + +.. py:class:: vortex.store.CosStore(bucket, endpoint, *, secret_id=None, secret_key=None, root=None, disable_config_load=False) + + A Tencent Cloud COS object store, backed by OpenDAL. Construct it with explicit + configuration and pass it to + :func:`vortex.io.read_url` / :func:`vortex.io.write` via the ``store=`` argument, + exactly like the built-in S3/Azure/GCS stores. + + The class is only available when Vortex is built with the ``opendal`` feature; on + a default build, instantiating it raises :class:`ImportError`. + + :param bucket: COS bucket name (e.g. ``"my-bucket"``). + :param endpoint: COS endpoint (e.g. ``"https://cos.ap-guangzhou.myqcloud.com"``). + :param secret_id: Optional Tencent Cloud secret id. Maps to the ``TENCENTCLOUD_SECRET_ID`` + environment variable when unset. + :param secret_key: Optional Tencent Cloud secret key. Maps to the + ``TENCENTCLOUD_SECRET_KEY`` environment variable when unset. + :param root: Optional key prefix applied to every operation. + :param disable_config_load: When ``True``, disable OpenDAL's automatic config loading + and rely only on the explicit configuration. Defaults to ``False``. Reading from COS ================ From bbcb83f3f15c76a996ce8618e04604ff446d628b Mon Sep 17 00:00:00 2001 From: forwardxu Date: Mon, 27 Jul 2026 08:44:34 +0800 Subject: [PATCH 7/8] fix(python): resolve pyright reportMissingModuleSource and Sphinx CosStore warnings - Add `# pyright: ignore[reportMissingModuleSource]` to the 7 store module files that import the dynamically-registered `vortex._lib.store` submodule. The submodule has no Python source (only .pyi stubs, which are excluded from basedpyright), so pyright could not resolve it. This was previously fixed by the same comments; commit b990c75d7 incorrectly removed them assuming the `_lib/__init__.pyi` CosStore declaration covered the store submodule too, but it does not. - Add `nitpick_ignore` entry for `vortex.store._cos.CosStore` in `docs/conf.py`. The `ObjectStore` type alias resolves to this private path, but the public class is documented via `.. py:class::` in `opendal.rst` which only adds the public path `vortex.store.CosStore` to the inventory. Signed-off-by: forwardxu --- docs/conf.py | 7 ++++++- vortex-python/python/vortex/store/__init__.py | 2 +- vortex-python/python/vortex/store/_aws.py | 2 +- vortex-python/python/vortex/store/_azure.py | 2 +- vortex-python/python/vortex/store/_gcs.py | 2 +- vortex-python/python/vortex/store/_http.py | 2 +- vortex-python/python/vortex/store/_local.py | 2 +- vortex-python/python/vortex/store/_memory.py | 2 +- 8 files changed, 13 insertions(+), 8 deletions(-) diff --git a/docs/conf.py b/docs/conf.py index a2cadbb255f..177f9a22fe4 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -54,7 +54,12 @@ git_root = Path(__file__).parent.parent nitpicky = True # ensures all :class:, :obj:, etc. links are valid -nitpick_ignore = [] +nitpick_ignore = [ + # `vortex.store.CosStore` is re-exported through the private `vortex.store._cos` module, + # and the `ObjectStore` type alias resolves to the private path. The public class is + # fully documented in `opendal.rst`; the private path is intentionally not. + ("py:class", "vortex.store._cos.CosStore"), +] doctest_global_setup = "import pyarrow; import vortex; import vortex as vx; import random; random.seed(a=0)" doctest_default_flags = ( diff --git a/vortex-python/python/vortex/store/__init__.py b/vortex-python/python/vortex/store/__init__.py index ba67fd72f3f..c9127087bf0 100644 --- a/vortex-python/python/vortex/store/__init__.py +++ b/vortex-python/python/vortex/store/__init__.py @@ -4,7 +4,7 @@ from collections.abc import Callable from typing import TypeAlias, Unpack, overload -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] from ._aws import S3Config, S3Credential, S3CredentialProvider, S3Store from ._azure import ( AzureAccessKey, diff --git a/vortex-python/python/vortex/store/_aws.py b/vortex-python/python/vortex/store/_aws.py index e3c02bf042e..3f3079e9902 100644 --- a/vortex-python/python/vortex/store/_aws.py +++ b/vortex-python/python/vortex/store/_aws.py @@ -7,7 +7,7 @@ from typing_extensions import override -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_azure.py b/vortex-python/python/vortex/store/_azure.py index 4d54531f928..1b3ec1b4e52 100644 --- a/vortex-python/python/vortex/store/_azure.py +++ b/vortex-python/python/vortex/store/_azure.py @@ -7,7 +7,7 @@ from typing_extensions import override -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_gcs.py b/vortex-python/python/vortex/store/_gcs.py index 9d00055ba82..e1512434093 100644 --- a/vortex-python/python/vortex/store/_gcs.py +++ b/vortex-python/python/vortex/store/_gcs.py @@ -7,7 +7,7 @@ from typing_extensions import override -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_http.py b/vortex-python/python/vortex/store/_http.py index 5da9c6b2cd6..57d46bb5ebb 100644 --- a/vortex-python/python/vortex/store/_http.py +++ b/vortex-python/python/vortex/store/_http.py @@ -5,7 +5,7 @@ from typing_extensions import override -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] from ._client import ClientConfig from ._retry import RetryConfig diff --git a/vortex-python/python/vortex/store/_local.py b/vortex-python/python/vortex/store/_local.py index de9670c2d2c..891c8e968b4 100644 --- a/vortex-python/python/vortex/store/_local.py +++ b/vortex-python/python/vortex/store/_local.py @@ -6,7 +6,7 @@ from typing_extensions import override -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] class LocalStore(_store.LocalStore): diff --git a/vortex-python/python/vortex/store/_memory.py b/vortex-python/python/vortex/store/_memory.py index 7dbd48b3d6f..763424732eb 100644 --- a/vortex-python/python/vortex/store/_memory.py +++ b/vortex-python/python/vortex/store/_memory.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: MIT # SPDX-FileCopyrightText: Copyright (c) 2024 Development Seed -from .._lib import store as _store +from .._lib import store as _store # pyright: ignore[reportMissingModuleSource] class MemoryStore(_store.MemoryStore): From 952fcb6e983653fa8517459c95343feb6e869ff5 Mon Sep 17 00:00:00 2001 From: forwardxu Date: Mon, 27 Jul 2026 10:46:57 +0800 Subject: [PATCH 8/8] fix: regenerate Cargo.lock after rebase onto b4f8d34b4 Update Cargo.lock to match the latest origin/develop (b4f8d34b4) which includes renovatebot dependency updates (reqsign 3.2.0, reqsign-file-read-tokio 3.0.3, reqsign-tencent-cos 3.0.3) and removes conflict markers left from the rebase. Signed-off-by: forwardxu --- Cargo.lock | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8ba8297a1d0..16b0a83f9ac 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -7327,14 +7327,13 @@ dependencies = [ [[package]] name = "reqsign-core" -version = "3.1.0" +version = "3.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "514a1e0b4aa288652a3fdbda4f0a610f379cdf5374e55a37c9edd03d57ed856b" +checksum = "7e38b44697c60a823705ccef85cb04d8e0527c9d16ed7c58bf1c6395bdd24ceb" dependencies = [ "anyhow", "base64", "bytes", - "form_urlencoded", "futures", "hex", "hmac", @@ -7343,15 +7342,15 @@ dependencies = [ "log", "percent-encoding", "sha1", - "sha2 0.11.0", + "sha2", "windows-sys 0.61.2", ] [[package]] name = "reqsign-file-read-tokio" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2b472a8d1f2e5a4be8ce13bb7bdf4b59e9bee613ce124aca23959ddb42176b39" +checksum = "688ff0ae421b8d4b92b53fdafaf53df2de28f428a9962edcf21702990b26f74b" dependencies = [ "anyhow", "reqsign-core", @@ -7360,9 +7359,9 @@ dependencies = [ [[package]] name = "reqsign-tencent-cos" -version = "3.0.2" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0f5c353c69bdbd6e3deb8b067a5547504cc155b7f024998fef005e37d65d646c" +checksum = "f6497dd9f6e3d1349b420521484099b284f95e8d3a65f088fccef42493a7b644" dependencies = [ "anyhow", "http", @@ -9991,14 +9990,8 @@ dependencies = [ "url", "vortex", "vortex-arrow", - "vortex-object-store-opendal", ->>>>>>> 64006a996 (feat(vortex-python): Add OpenDAL-backed CosStore to the Python object-store API) -======= "vortex-geo", "vortex-object-store-opendal", -======= - "vortex-object-store-opendal", ->>>>>>> 64006a996 (feat(vortex-python): Add OpenDAL-backed CosStore to the Python object-store API) "vortex-parquet-variant", ]