starred/windmill
1
0
Fork 0
mirror of https://github.com/windmill-labs/windmill.git synced 2026-09-21 08:02:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

chore: vendor the object_store fork and page filesystem levels natively

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011HeHeGLWDuu62fxjBZA15P
This commit is contained in:
Diego Imbert
2026-07-29 18:29:14 +02:00
co-authored by Claude Fable 5
parent af98815e5c
commit a3a7033c31
62 changed files with 29452 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/Cargo.lock
Generated
+7 -1
View File
@@ -8170,7 +8170,6 @@ dependencies = [
[[package]]
name = "object_store"
version = "0.12.0"
source = "git+https://github.com/diegoimbert/arrow-rs-object-store?rev=a75c080f2fe821c28754481d840cd0aeccb0f42e#a75c080f2fe821c28754481d840cd0aeccb0f42e"
dependencies = [
"async-trait",
"base64 0.22.1",
@@ -8178,29 +8177,36 @@ dependencies = [
"chrono",
"form_urlencoded",
"futures",
"getrandom 0.2.17",
"getrandom 0.3.4",
"http 1.4.2",
"http-body-util",
"httparse",
"humantime",
"hyper 1.11.0",
"hyper-util",
"itertools 0.14.0",
"md-5 0.10.6",
"nix 0.30.1",
"parking_lot",
"percent-encoding",
"quick-xml",
"rand 0.9.0",
"regex",
"reqwest 0.12.28",
"ring 0.17.14",
"rustls-pemfile 2.2.0",
"serde",
"serde_json",
"serde_urlencoded",
"tempfile",
"thiserror 2.0.19",
"tokio",
"tracing",
"url",
"walkdir",
"wasm-bindgen-futures",
"wasm-bindgen-test",
"web-time",
]
backend/Cargo.toml
+3 -2
View File
@@ -84,6 +84,7 @@ members = [
"./windmill-worker-volumes",
"./windmill-test-utils",
"./windmill-api-integration-tests",
"./vendor/object_store",
]
exclude = ["./windmill-duckdb-ffi-internal", "./parsers/windmill-parser-wasm"]
@@ -214,7 +215,7 @@ all_sqlx_features = ["all_languages", "enterprise", "enterprise_saml", "embeddin
"windmill-git-sync/all_sqlx_features"]
[patch.crates-io]
object_store = { git = "https://github.com/diegoimbert/arrow-rs-object-store", rev = "a75c080f2fe821c28754481d840cd0aeccb0f42e" }
object_store = { path = "vendor/object_store" }
# Use tiberius main branch for libgssapi 0.8.1 fix (https://github.com/prisma/tiberius/issues/343)
tiberius = { git = "https://github.com/prisma/tiberius", rev = "59db57960a14b422fb3a1309aa4aa47880896ff8" }
# Pin tokio-postgres / postgres-types / postgres-protocol to the
@@ -618,7 +619,7 @@ process-wrap = { version = "8.2.1", features = ["tokio1"] }
systemstat = "0.2.4"
datafusion = "47.0.0"
object_store = { git = "https://github.com/diegoimbert/arrow-rs-object-store", rev = "a75c080f2fe821c28754481d840cd0aeccb0f42e", features = ["aws", "azure", "gcp"] }
object_store = { path = "vendor/object_store", features = ["aws", "azure", "gcp"] }
openidconnect = { version = "4.0.0-rc.1" }
aws-config = "^1"
aws-sdk-bedrock = "1.129.0"
backend/vendor/object_store/Cargo.toml
Vendored
+103
View File
@@ -0,0 +1,103 @@
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements. See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership. The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
[package]
name = "object_store"
version = "0.12.0"
edition = "2021"
license = "MIT/Apache-2.0"
readme = "README.md"
description = "A generic object store interface for uniformly interacting with AWS S3, Google Cloud Storage, Azure Blob Storage and local files."
keywords = ["object", "storage", "cloud"]
repository = "https://github.com/apache/arrow-rs-object-store"
rust-version = "1.64.0"
[package.metadata.docs.rs]
all-features = true
[dependencies] # In alphabetical order
async-trait = "0.1.53"
bytes = "1.0"
chrono = { version = "0.4.34", default-features = false, features = ["clock"] }
futures = "0.3"
http = "1.2.0"
humantime = "2.1"
itertools = "0.14.0"
parking_lot = { version = "0.12" }
percent-encoding = "2.1"
thiserror = "2.0.2"
tracing = { version = "0.1" }
url = "2.2"
walkdir = { version = "2", optional = true }
# Cloud storage support
base64 = { version = "0.22", default-features = false, features = ["std"], optional = true }
form_urlencoded = { version = "1.2", optional = true }
http-body-util = { version = "0.1.2", optional = true }
httparse = { version = "1.8.0", default-features = false, features = ["std"], optional = true }
hyper = { version = "1.2", default-features = false, optional = true }
md-5 = { version = "0.10.6", default-features = false, optional = true }
quick-xml = { version = "0.37.0", features = ["serialize", "overlapped-lists"], optional = true }
rand = { version = "0.9", default-features = false, features = ["std", "std_rng", "thread_rng"], optional = true }
reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-native-roots", "http2"], optional = true }
ring = { version = "0.17", default-features = false, features = ["std"], optional = true }
rustls-pemfile = { version = "2.0", default-features = false, features = ["std"], optional = true }
serde = { version = "1.0", default-features = false, features = ["derive"], optional = true }
serde_json = { version = "1.0", default-features = false, features = ["std"], optional = true }
serde_urlencoded = { version = "0.7", optional = true }
tokio = { version = "1.29.0", features = ["sync", "macros", "rt", "time", "io-util"] }
[target.'cfg(target_family="unix")'.dev-dependencies]
nix = { version = "0.30.0", features = ["fs"] }
[target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]
web-time = { version = "1.1.0" }
wasm-bindgen-futures = "0.4.18"
[features]
default = ["fs"]
cloud = ["serde", "serde_json", "quick-xml", "hyper", "reqwest", "reqwest/stream", "chrono/serde", "base64", "rand", "ring", "http-body-util", "form_urlencoded", "serde_urlencoded"]
azure = ["cloud", "httparse"]
fs = ["walkdir"]
gcp = ["cloud", "rustls-pemfile"]
aws = ["cloud", "md-5"]
http = ["cloud"]
tls-webpki-roots = ["reqwest?/rustls-tls-webpki-roots"]
integration = ["rand"]
[dev-dependencies] # In alphabetical order
hyper = { version = "1.2", features = ["server"] }
hyper-util = "0.1"
rand = "0.9"
tempfile = "3.1.0"
regex = "1.11.1"
# The "gzip" feature for reqwest is enabled for an integration test.
reqwest = { version = "0.12", features = ["gzip"] }
[target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dev-dependencies]
wasm-bindgen-test = "*"
[dev-dependencies.getrandom_v03]
package = "getrandom"
version = "0.3"
features = ["wasm_js"]
[dev-dependencies.getrandom_v02]
package = "getrandom"
version = "0.2"
features = ["js"]
backend/vendor/object_store/LICENSE.txt
Vendored
+204
View File
@@ -0,0 +1,204 @@
Apache License
Version 2.0, January 2004
http://www.apache.org/licenses/
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
1. Definitions.
"License" shall mean the terms and conditions for use, reproduction,
and distribution as defined by Sections 1 through 9 of this document.
"Licensor" shall mean the copyright owner or entity authorized by
the copyright owner that is granting the License.
"Legal Entity" shall mean the union of the acting entity and all
other entities that control, are controlled by, or are under common
control with that entity. For the purposes of this definition,
"control" means (i) the power, direct or indirect, to cause the
direction or management of such entity, whether by contract or
otherwise, or (ii) ownership of fifty percent (50%) or more of the
outstanding shares, or (iii) beneficial ownership of such entity.
"You" (or "Your") shall mean an individual or Legal Entity
exercising permissions granted by this License.
"Source" form shall mean the preferred form for making modifications,
including but not limited to software source code, documentation
source, and configuration files.
"Object" form shall mean any form resulting from mechanical
transformation or translation of a Source form, including but
not limited to compiled object code, generated documentation,
and conversions to other media types.
"Work" shall mean the work of authorship, whether in Source or
Object form, made available under the License, as indicated by a
copyright notice that is included in or attached to the work
(an example is provided in the Appendix below).
"Derivative Works" shall mean any work, whether in Source or Object
form, that is based on (or derived from) the Work and for which the
editorial revisions, annotations, elaborations, or other modifications
represent, as a whole, an original work of authorship. For the purposes
of this License, Derivative Works shall not include works that remain
separable from, or merely link (or bind by name) to the interfaces of,
the Work and Derivative Works thereof.
"Contribution" shall mean any work of authorship, including
the original version of the Work and any modifications or additions
to that Work or Derivative Works thereof, that is intentionally
submitted to Licensor for inclusion in the Work by the copyright owner
or by an individual or Legal Entity authorized to submit on behalf of
the copyright owner. For the purposes of this definition, "submitted"
means any form of electronic, verbal, or written communication sent
to the Licensor or its representatives, including but not limited to
communication on electronic mailing lists, source code control systems,
and issue tracking systems that are managed by, or on behalf of, the
Licensor for the purpose of discussing and improving the Work, but
excluding communication that is conspicuously marked or otherwise
designated in writing by the copyright owner as "Not a Contribution."
"Contributor" shall mean Licensor and any individual or Legal Entity
on behalf of whom a Contribution has been received by Licensor and
subsequently incorporated within the Work.
2. Grant of Copyright License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
copyright license to reproduce, prepare Derivative Works of,
publicly display, publicly perform, sublicense, and distribute the
Work and such Derivative Works in Source or Object form.
3. Grant of Patent License. Subject to the terms and conditions of
this License, each Contributor hereby grants to You a perpetual,
worldwide, non-exclusive, no-charge, royalty-free, irrevocable
(except as stated in this section) patent license to make, have made,
use, offer to sell, sell, import, and otherwise transfer the Work,
where such license applies only to those patent claims licensable
by such Contributor that are necessarily infringed by their
Contribution(s) alone or by combination of their Contribution(s)
with the Work to which such Contribution(s) was submitted. If You
institute patent litigation against any entity (including a
cross-claim or counterclaim in a lawsuit) alleging that the Work
or a Contribution incorporated within the Work constitutes direct
or contributory patent infringement, then any patent licenses
granted to You under this License for that Work shall terminate
as of the date such litigation is filed.
4. Redistribution. You may reproduce and distribute copies of the
Work or Derivative Works thereof in any medium, with or without
modifications, and in Source or Object form, provided that You
meet the following conditions:
(a) You must give any other recipients of the Work or
Derivative Works a copy of this License; and
(b) You must cause any modified files to carry prominent notices
stating that You changed the files; and
(c) You must retain, in the Source form of any Derivative Works
that You distribute, all copyright, patent, trademark, and
attribution notices from the Source form of the Work,
excluding those notices that do not pertain to any part of
the Derivative Works; and
(d) If the Work includes a "NOTICE" text file as part of its
distribution, then any Derivative Works that You distribute must
include a readable copy of the attribution notices contained
within such NOTICE file, excluding those notices that do not
pertain to any part of the Derivative Works, in at least one
of the following places: within a NOTICE text file distributed
as part of the Derivative Works; within the Source form or
documentation, if provided along with the Derivative Works; or,
within a display generated by the Derivative Works, if and
wherever such third-party notices normally appear. The contents
of the NOTICE file are for informational purposes only and
do not modify the License. You may add Your own attribution
notices within Derivative Works that You distribute, alongside
or as an addendum to the NOTICE text from the Work, provided
that such additional attribution notices cannot be construed
as modifying the License.
You may add Your own copyright statement to Your modifications and
may provide additional or different license terms and conditions
for use, reproduction, or distribution of Your modifications, or
for any such Derivative Works as a whole, provided Your use,
reproduction, and distribution of the Work otherwise complies with
the conditions stated in this License.
5. Submission of Contributions. Unless You explicitly state otherwise,
any Contribution intentionally submitted for inclusion in the Work
by You to the Licensor shall be under the terms and conditions of
this License, without any additional terms or conditions.
Notwithstanding the above, nothing herein shall supersede or modify
the terms of any separate license agreement you may have executed
with Licensor regarding such Contributions.
6. Trademarks. This License does not grant permission to use the trade
names, trademarks, service marks, or product names of the Licensor,
except as required for reasonable and customary use in describing the
origin of the Work and reproducing the content of the NOTICE file.
7. Disclaimer of Warranty. Unless required by applicable law or
agreed to in writing, Licensor provides the Work (and each
Contributor provides its Contributions) on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
implied, including, without limitation, any warranties or conditions
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
PARTICULAR PURPOSE. You are solely responsible for determining the
appropriateness of using or redistributing the Work and assume any
risks associated with Your exercise of permissions under this License.
8. Limitation of Liability. In no event and under no legal theory,
whether in tort (including negligence), contract, or otherwise,
unless required by applicable law (such as deliberate and grossly
negligent acts) or agreed to in writing, shall any Contributor be
liable to You for damages, including any direct, indirect, special,
incidental, or consequential damages of any character arising as a
result of this License or out of the use or inability to use the
Work (including but not limited to damages for loss of goodwill,
work stoppage, computer failure or malfunction, or any and all
other commercial damages or losses), even if such Contributor
has been advised of the possibility of such damages.
9. Accepting Warranty or Additional Liability. While redistributing
the Work or Derivative Works thereof, You may choose to offer,
and charge a fee for, acceptance of support, warranty, indemnity,
or other liability obligations and/or rights consistent with this
License. However, in accepting such obligations, You may act only
on Your own behalf and on Your sole responsibility, not on behalf
of any other Contributor, and only if You agree to indemnify,
defend, and hold each Contributor harmless for any liability
incurred by, or claims asserted against, such Contributor by reason
of your accepting any such warranty or additional liability.
END OF TERMS AND CONDITIONS
APPENDIX: How to apply the Apache License to your work.
To apply the Apache License to your work, attach the following
boilerplate notice, with the fields enclosed by brackets "[]"
replaced with your own identifying information. (Don't include
the brackets!) The text should be enclosed in the appropriate
comment syntax for the file format. We also recommend that a
file or class name and description of purpose be included on the
same "printed page" as the copyright notice for easier
identification within third-party archives.
Copyright [yyyy] [name of copyright owner]
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
backend/vendor/object_store/NOTICE.txt
Vendored
+5
View File
@@ -0,0 +1,5 @@
Apache Arrow Object Store
Copyright 2020-2024 The Apache Software Foundation
This product includes software developed at
The Apache Software Foundation (http://www.apache.org/).
backend/vendor/object_store/VENDORED.md
Vendored
+21
View File
@@ -0,0 +1,21 @@
# Vendored object_store
Vendored copy of `apache/arrow-rs-object-store` at rev
`36752c975d4f29e20b57c91f81a10872dcd48ae7` (the rev this workspace previously
pinned via git), plus three Windmill commits — mirrored at
https://github.com/diegoimbert/arrow-rs-object-store/tree/windmill/list-delimited-page
(head `6a63a59ed5ee90c4e8e2bf502779437baf7aa806`):
- `fix: use runtime format widths accepted by recent rustc` (test-only)
- `feat: add ObjectStore::list_delimited_page for bounded, resumable delimited listing`
- `fix: saturate the page bound and disambiguate colliding entry names in the
default list_delimited_page`
- `perf: stat only the returned page in LocalFileSystem's delimited paging`
`list_delimited_page` is the reason for the fork: one bounded page of a
delimited listing with an opaque continuation token, served natively on
S3/GCS/Azure. It is a candidate for upstreaming; if it lands, this directory
goes away and the dependency returns to the upstream pin.
To update: rebase the commits above onto the new upstream rev, rerun the crate
test suite (`cargo test -p object_store --lib`), and replace this directory.
backend/vendor/object_store/src/attributes.rs
Vendored
+248
View File
@@ -0,0 +1,248 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::borrow::Cow;
use std::collections::HashMap;
use std::ops::Deref;
/// Additional object attribute types
#[non_exhaustive]
#[derive(Debug, Hash, Eq, PartialEq, Clone)]
pub enum Attribute {
/// Specifies how the object should be handled by a browser
///
/// See [Content-Disposition](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Disposition)
ContentDisposition,
/// Specifies the encodings applied to the object
///
/// See [Content-Encoding](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Encoding)
ContentEncoding,
/// Specifies the language of the object
///
/// See [Content-Language](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Language)
ContentLanguage,
/// Specifies the MIME type of the object
///
/// This takes precedence over any [ClientOptions](crate::ClientOptions) configuration
///
/// See [Content-Type](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Type)
ContentType,
/// Overrides cache control policy of the object
///
/// See [Cache-Control](https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Cache-Control)
CacheControl,
/// Specifies a user-defined metadata field for the object
///
/// The String is a user-defined key
Metadata(Cow<'static, str>),
}
/// The value of an [`Attribute`]
///
/// Provides efficient conversion from both static and owned strings
///
/// ```
/// # use object_store::AttributeValue;
/// // Can use static strings without needing an allocation
/// let value = AttributeValue::from("bar");
/// // Can also store owned strings
/// let value = AttributeValue::from("foo".to_string());
/// ```
#[derive(Debug, Hash, Eq, PartialEq, Clone)]
pub struct AttributeValue(Cow<'static, str>);
impl AsRef<str> for AttributeValue {
fn as_ref(&self) -> &str {
&self.0
}
}
impl From<&'static str> for AttributeValue {
fn from(value: &'static str) -> Self {
Self(Cow::Borrowed(value))
}
}
impl From<String> for AttributeValue {
fn from(value: String) -> Self {
Self(Cow::Owned(value))
}
}
impl Deref for AttributeValue {
type Target = str;
fn deref(&self) -> &Self::Target {
self.0.as_ref()
}
}
/// Additional attributes of an object
///
/// Attributes can be specified in [PutOptions](crate::PutOptions) and retrieved
/// from APIs returning [GetResult](crate::GetResult).
///
/// Unlike [`ObjectMeta`](crate::ObjectMeta), [`Attributes`] are not returned by
/// listing APIs
#[derive(Debug, Default, Eq, PartialEq, Clone)]
pub struct Attributes(HashMap<Attribute, AttributeValue>);
impl Attributes {
/// Create a new empty [`Attributes`]
pub fn new() -> Self {
Self::default()
}
/// Create a new [`Attributes`] with space for `capacity` [`Attribute`]
pub fn with_capacity(capacity: usize) -> Self {
Self(HashMap::with_capacity(capacity))
}
/// Insert a new [`Attribute`], [`AttributeValue`] pair
///
/// Returns the previous value for `key` if any
pub fn insert(&mut self, key: Attribute, value: AttributeValue) -> Option<AttributeValue> {
self.0.insert(key, value)
}
/// Returns the [`AttributeValue`] for `key` if any
pub fn get(&self, key: &Attribute) -> Option<&AttributeValue> {
self.0.get(key)
}
/// Removes the [`AttributeValue`] for `key` if any
pub fn remove(&mut self, key: &Attribute) -> Option<AttributeValue> {
self.0.remove(key)
}
/// Returns an [`AttributesIter`] over this
pub fn iter(&self) -> AttributesIter<'_> {
self.into_iter()
}
/// Returns the number of [`Attribute`] in this collection
#[inline]
pub fn len(&self) -> usize {
self.0.len()
}
/// Returns true if this contains no [`Attribute`]
#[inline]
pub fn is_empty(&self) -> bool {
self.0.is_empty()
}
}
impl<K, V> FromIterator<(K, V)> for Attributes
where
K: Into<Attribute>,
V: Into<AttributeValue>,
{
fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
Self(
iter.into_iter()
.map(|(k, v)| (k.into(), v.into()))
.collect(),
)
}
}
impl<'a> IntoIterator for &'a Attributes {
type Item = (&'a Attribute, &'a AttributeValue);
type IntoIter = AttributesIter<'a>;
fn into_iter(self) -> Self::IntoIter {
AttributesIter(self.0.iter())
}
}
/// Iterator over [`Attributes`]
#[derive(Debug)]
pub struct AttributesIter<'a>(std::collections::hash_map::Iter<'a, Attribute, AttributeValue>);
impl<'a> Iterator for AttributesIter<'a> {
type Item = (&'a Attribute, &'a AttributeValue);
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_attributes_basic() {
let mut attributes = Attributes::from_iter([
(Attribute::ContentDisposition, "inline"),
(Attribute::ContentEncoding, "gzip"),
(Attribute::ContentLanguage, "en-US"),
(Attribute::ContentType, "test"),
(Attribute::CacheControl, "control"),
(Attribute::Metadata("key1".into()), "value1"),
]);
assert!(!attributes.is_empty());
assert_eq!(attributes.len(), 6);
assert_eq!(
attributes.get(&Attribute::ContentType),
Some(&"test".into())
);
let metav = "control".into();
assert_eq!(attributes.get(&Attribute::CacheControl), Some(&metav));
assert_eq!(
attributes.insert(Attribute::CacheControl, "v1".into()),
Some(metav)
);
assert_eq!(attributes.len(), 6);
assert_eq!(
attributes.remove(&Attribute::CacheControl).unwrap(),
"v1".into()
);
assert_eq!(attributes.len(), 5);
let metav: AttributeValue = "v2".into();
attributes.insert(Attribute::CacheControl, metav.clone());
assert_eq!(attributes.get(&Attribute::CacheControl), Some(&metav));
assert_eq!(attributes.len(), 6);
assert_eq!(
attributes.get(&Attribute::ContentDisposition),
Some(&"inline".into())
);
assert_eq!(
attributes.get(&Attribute::ContentEncoding),
Some(&"gzip".into())
);
assert_eq!(
attributes.get(&Attribute::ContentLanguage),
Some(&"en-US".into())
);
assert_eq!(
attributes.get(&Attribute::Metadata("key1".into())),
Some(&"value1".into())
);
}
}
backend/vendor/object_store/src/aws/builder.rs
Vendored
+1614
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/aws/checksum.rs
Vendored
+63
View File
@@ -0,0 +1,63 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::config::Parse;
use std::str::FromStr;
#[allow(non_camel_case_types)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
/// Enum representing checksum algorithm supported by S3.
pub enum Checksum {
/// SHA-256 algorithm.
SHA256,
}
impl std::fmt::Display for Checksum {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match &self {
Self::SHA256 => write!(f, "sha256"),
}
}
}
impl FromStr for Checksum {
type Err = ();
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s.to_lowercase().as_str() {
"sha256" => Ok(Self::SHA256),
_ => Err(()),
}
}
}
impl TryFrom<&String> for Checksum {
type Error = ();
fn try_from(value: &String) -> Result<Self, Self::Error> {
value.parse()
}
}
impl Parse for Checksum {
fn parse(v: &str) -> crate::Result<Self> {
v.parse().map_err(|_| crate::Error::Generic {
store: "Config",
source: format!("\"{v}\" is not a valid checksum algorithm").into(),
})
}
}
backend/vendor/object_store/src/aws/client.rs
Vendored
+937
View File
@@ -0,0 +1,937 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::aws::builder::S3EncryptionHeaders;
use crate::aws::checksum::Checksum;
use crate::aws::credential::{AwsCredential, CredentialExt};
use crate::aws::{
AwsAuthorizer, AwsCredentialProvider, S3ConditionalPut, S3CopyIfNotExists, COPY_SOURCE_HEADER,
STORE, STRICT_PATH_ENCODE_SET, TAGS_HEADER,
};
use crate::client::builder::{HttpRequestBuilder, RequestBuilderError};
use crate::client::get::GetClient;
use crate::client::header::{get_etag, HeaderConfig};
use crate::client::header::{get_put_result, get_version};
use crate::client::list::ListClient;
use crate::client::retry::RetryExt;
use crate::client::s3::{
CompleteMultipartUpload, CompleteMultipartUploadResult, CopyPartResult,
InitiateMultipartUploadResult, ListResponse, PartMetadata,
};
use crate::client::{GetOptionsExt, HttpClient, HttpError, HttpResponse};
use crate::multipart::PartId;
use crate::path::DELIMITER;
use crate::{
Attribute, Attributes, ClientOptions, GetOptions, ListResult, MultipartId, Path,
PutMultipartOpts, PutPayload, PutResult, Result, RetryConfig, TagSet,
};
use async_trait::async_trait;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use bytes::{Buf, Bytes};
use http::header::{
CACHE_CONTROL, CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LANGUAGE, CONTENT_LENGTH,
CONTENT_TYPE,
};
use http::{HeaderMap, HeaderName, Method};
use itertools::Itertools;
use md5::{Digest, Md5};
use percent_encoding::{utf8_percent_encode, PercentEncode};
use quick_xml::events::{self as xml_events};
use ring::digest;
use ring::digest::Context;
use serde::{Deserialize, Serialize};
use std::sync::Arc;
const VERSION_HEADER: &str = "x-amz-version-id";
const SHA256_CHECKSUM: &str = "x-amz-checksum-sha256";
const USER_DEFINED_METADATA_HEADER_PREFIX: &str = "x-amz-meta-";
const ALGORITHM: &str = "x-amz-checksum-algorithm";
/// A specialized `Error` for object store-related errors
#[derive(Debug, thiserror::Error)]
pub(crate) enum Error {
#[error("Error performing DeleteObjects request: {}", source)]
DeleteObjectsRequest {
source: crate::client::retry::RetryError,
},
#[error(
"DeleteObjects request failed for key {}: {} (code: {})",
path,
message,
code
)]
DeleteFailed {
path: String,
code: String,
message: String,
},
#[error("Error getting DeleteObjects response body: {}", source)]
DeleteObjectsResponse { source: HttpError },
#[error("Got invalid DeleteObjects response: {}", source)]
InvalidDeleteObjectsResponse {
source: Box<dyn std::error::Error + Send + Sync + 'static>,
},
#[error("Error performing list request: {}", source)]
ListRequest {
source: crate::client::retry::RetryError,
},
#[error("Error getting list response body: {}", source)]
ListResponseBody { source: HttpError },
#[error("Error getting create multipart response body: {}", source)]
CreateMultipartResponseBody { source: HttpError },
#[error("Error performing complete multipart request: {}: {}", path, source)]
CompleteMultipartRequest {
source: crate::client::retry::RetryError,
path: String,
},
#[error("Error getting complete multipart response body: {}", source)]
CompleteMultipartResponseBody { source: HttpError },
#[error("Got invalid list response: {}", source)]
InvalidListResponse { source: quick_xml::de::DeError },
#[error("Got invalid multipart response: {}", source)]
InvalidMultipartResponse { source: quick_xml::de::DeError },
#[error("Unable to extract metadata from headers: {}", source)]
Metadata {
source: crate::client::header::Error,
},
}
impl From<Error> for crate::Error {
fn from(err: Error) -> Self {
match err {
Error::CompleteMultipartRequest { source, path } => source.error(STORE, path),
_ => Self::Generic {
store: STORE,
source: Box::new(err),
},
}
}
}
pub(crate) enum PutPartPayload<'a> {
Part(PutPayload),
Copy(&'a Path),
}
impl Default for PutPartPayload<'_> {
fn default() -> Self {
Self::Part(PutPayload::default())
}
}
pub(crate) enum CompleteMultipartMode {
Overwrite,
Create,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase", rename = "DeleteResult")]
struct BatchDeleteResponse {
#[serde(rename = "$value")]
content: Vec<DeleteObjectResult>,
}
#[derive(Deserialize)]
enum DeleteObjectResult {
#[allow(unused)]
Deleted(DeletedObject),
Error(DeleteError),
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase", rename = "Deleted")]
struct DeletedObject {
#[allow(dead_code)]
key: String,
}
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase", rename = "Error")]
struct DeleteError {
key: String,
code: String,
message: String,
}
impl From<DeleteError> for Error {
fn from(err: DeleteError) -> Self {
Self::DeleteFailed {
path: err.key,
code: err.code,
message: err.message,
}
}
}
#[derive(Debug)]
pub(crate) struct S3Config {
pub region: String,
pub endpoint: Option<String>,
pub bucket: String,
pub bucket_endpoint: String,
pub credentials: AwsCredentialProvider,
pub session_provider: Option<AwsCredentialProvider>,
pub retry_config: RetryConfig,
pub client_options: ClientOptions,
pub sign_payload: bool,
pub skip_signature: bool,
pub disable_tagging: bool,
pub checksum: Option<Checksum>,
pub copy_if_not_exists: Option<S3CopyIfNotExists>,
pub conditional_put: S3ConditionalPut,
pub request_payer: bool,
pub(super) encryption_headers: S3EncryptionHeaders,
}
impl S3Config {
pub(crate) fn path_url(&self, path: &Path) -> String {
format!("{}/{}", self.bucket_endpoint, encode_path(path))
}
async fn get_session_credential(&self) -> Result<SessionCredential<'_>> {
let credential = match self.skip_signature {
false => {
let provider = self.session_provider.as_ref().unwrap_or(&self.credentials);
Some(provider.get_credential().await?)
}
true => None,
};
Ok(SessionCredential {
credential,
session_token: self.session_provider.is_some(),
config: self,
})
}
pub(crate) async fn get_credential(&self) -> Result<Option<Arc<AwsCredential>>> {
Ok(match self.skip_signature {
false => Some(self.credentials.get_credential().await?),
true => None,
})
}
#[inline]
pub(crate) fn is_s3_express(&self) -> bool {
self.session_provider.is_some()
}
}
struct SessionCredential<'a> {
credential: Option<Arc<AwsCredential>>,
session_token: bool,
config: &'a S3Config,
}
impl SessionCredential<'_> {
fn authorizer(&self) -> Option<AwsAuthorizer<'_>> {
let mut authorizer =
AwsAuthorizer::new(self.credential.as_deref()?, "s3", &self.config.region)
.with_sign_payload(self.config.sign_payload)
.with_request_payer(self.config.request_payer);
if self.session_token {
let token = HeaderName::from_static("x-amz-s3session-token");
authorizer = authorizer.with_token_header(token)
}
Some(authorizer)
}
}
#[derive(Debug, thiserror::Error)]
pub enum RequestError {
#[error(transparent)]
Generic {
#[from]
source: crate::Error,
},
#[error("Retry")]
Retry {
source: crate::client::retry::RetryError,
path: String,
},
}
impl From<RequestError> for crate::Error {
fn from(value: RequestError) -> Self {
match value {
RequestError::Generic { source } => source,
RequestError::Retry { source, path } => source.error(STORE, path),
}
}
}
/// A builder for a request allowing customisation of the headers and query string
pub(crate) struct Request<'a> {
path: &'a Path,
config: &'a S3Config,
builder: HttpRequestBuilder,
payload_sha256: Option<digest::Digest>,
payload: Option<PutPayload>,
use_session_creds: bool,
idempotent: bool,
retry_on_conflict: bool,
retry_error_body: bool,
}
impl Request<'_> {
pub(crate) fn query<T: Serialize + ?Sized + Sync>(self, query: &T) -> Self {
let builder = self.builder.query(query);
Self { builder, ..self }
}
pub(crate) fn header<K>(self, k: K, v: &str) -> Self
where
K: TryInto<HeaderName>,
K::Error: Into<RequestBuilderError>,
{
let builder = self.builder.header(k, v);
Self { builder, ..self }
}
pub(crate) fn headers(self, headers: HeaderMap) -> Self {
let builder = self.builder.headers(headers);
Self { builder, ..self }
}
pub(crate) fn idempotent(self, idempotent: bool) -> Self {
Self { idempotent, ..self }
}
pub(crate) fn retry_on_conflict(self, retry_on_conflict: bool) -> Self {
Self {
retry_on_conflict,
..self
}
}
pub(crate) fn retry_error_body(self, retry_error_body: bool) -> Self {
Self {
retry_error_body,
..self
}
}
pub(crate) fn with_encryption_headers(self) -> Self {
let headers = self.config.encryption_headers.clone().into();
let builder = self.builder.headers(headers);
Self { builder, ..self }
}
pub(crate) fn with_session_creds(self, use_session_creds: bool) -> Self {
Self {
use_session_creds,
..self
}
}
pub(crate) fn with_tags(mut self, tags: TagSet) -> Self {
let tags = tags.encoded();
if !tags.is_empty() && !self.config.disable_tagging {
self.builder = self.builder.header(&TAGS_HEADER, tags);
}
self
}
pub(crate) fn with_attributes(self, attributes: Attributes) -> Self {
let mut has_content_type = false;
let mut builder = self.builder;
for (k, v) in &attributes {
builder = match k {
Attribute::CacheControl => builder.header(CACHE_CONTROL, v.as_ref()),
Attribute::ContentDisposition => builder.header(CONTENT_DISPOSITION, v.as_ref()),
Attribute::ContentEncoding => builder.header(CONTENT_ENCODING, v.as_ref()),
Attribute::ContentLanguage => builder.header(CONTENT_LANGUAGE, v.as_ref()),
Attribute::ContentType => {
has_content_type = true;
builder.header(CONTENT_TYPE, v.as_ref())
}
Attribute::Metadata(k_suffix) => builder.header(
&format!("{}{}", USER_DEFINED_METADATA_HEADER_PREFIX, k_suffix),
v.as_ref(),
),
};
}
if !has_content_type {
if let Some(value) = self.config.client_options.get_content_type(self.path) {
builder = builder.header(CONTENT_TYPE, value);
}
}
Self { builder, ..self }
}
pub(crate) fn with_extensions(self, extensions: ::http::Extensions) -> Self {
let builder = self.builder.extensions(extensions);
Self { builder, ..self }
}
pub(crate) fn with_payload(mut self, payload: PutPayload) -> Self {
if (!self.config.skip_signature && self.config.sign_payload)
|| self.config.checksum.is_some()
{
let mut sha256 = Context::new(&digest::SHA256);
payload.iter().for_each(|x| sha256.update(x));
let payload_sha256 = sha256.finish();
if let Some(Checksum::SHA256) = self.config.checksum {
self.builder = self
.builder
.header(SHA256_CHECKSUM, BASE64_STANDARD.encode(payload_sha256));
}
self.payload_sha256 = Some(payload_sha256);
}
let content_length = payload.content_length();
self.builder = self.builder.header(CONTENT_LENGTH, content_length);
self.payload = Some(payload);
self
}
pub(crate) async fn send(self) -> Result<HttpResponse, RequestError> {
let credential = match self.use_session_creds {
true => self.config.get_session_credential().await?,
false => SessionCredential {
credential: self.config.get_credential().await?,
session_token: false,
config: self.config,
},
};
let sha = self.payload_sha256.as_ref().map(|x| x.as_ref());
let path = self.path.as_ref();
self.builder
.with_aws_sigv4(credential.authorizer(), sha)
.retryable(&self.config.retry_config)
.retry_on_conflict(self.retry_on_conflict)
.idempotent(self.idempotent)
.retry_error_body(self.retry_error_body)
.payload(self.payload)
.send()
.await
.map_err(|source| {
let path = path.into();
RequestError::Retry { source, path }
})
}
pub(crate) async fn do_put(self) -> Result<PutResult> {
let response = self.send().await?;
Ok(get_put_result(response.headers(), VERSION_HEADER)
.map_err(|source| Error::Metadata { source })?)
}
}
#[derive(Debug)]
pub(crate) struct S3Client {
pub config: S3Config,
pub client: HttpClient,
}
impl S3Client {
pub(crate) fn new(config: S3Config, client: HttpClient) -> Self {
Self { config, client }
}
pub(crate) fn request<'a>(&'a self, method: Method, path: &'a Path) -> Request<'a> {
let url = self.config.path_url(path);
Request {
path,
builder: self.client.request(method, url),
payload: None,
payload_sha256: None,
config: &self.config,
use_session_creds: true,
idempotent: false,
retry_on_conflict: false,
retry_error_body: false,
}
}
/// Make an S3 Delete Objects request <https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html>
///
/// Produces a vector of results, one for each path in the input vector. If
/// the delete was successful, the path is returned in the `Ok` variant. If
/// there was an error for a certain path, the error will be returned in the
/// vector. If there was an issue with making the overall request, an error
/// will be returned at the top level.
pub(crate) async fn bulk_delete_request(&self, paths: Vec<Path>) -> Result<Vec<Result<Path>>> {
if paths.is_empty() {
return Ok(Vec::new());
}
let credential = self.config.get_session_credential().await?;
let url = format!("{}?delete", self.config.bucket_endpoint);
let mut buffer = Vec::new();
let mut writer = quick_xml::Writer::new(&mut buffer);
writer
.write_event(xml_events::Event::Start(
xml_events::BytesStart::new("Delete")
.with_attributes([("xmlns", "http://s3.amazonaws.com/doc/2006-03-01/")]),
))
.unwrap();
for path in &paths {
// <Object><Key>{path}</Key></Object>
writer
.write_event(xml_events::Event::Start(xml_events::BytesStart::new(
"Object",
)))
.unwrap();
writer
.write_event(xml_events::Event::Start(xml_events::BytesStart::new("Key")))
.unwrap();
writer
.write_event(xml_events::Event::Text(xml_events::BytesText::new(
path.as_ref(),
)))
.map_err(|err| crate::Error::Generic {
store: STORE,
source: Box::new(err),
})?;
writer
.write_event(xml_events::Event::End(xml_events::BytesEnd::new("Key")))
.unwrap();
writer
.write_event(xml_events::Event::End(xml_events::BytesEnd::new("Object")))
.unwrap();
}
writer
.write_event(xml_events::Event::End(xml_events::BytesEnd::new("Delete")))
.unwrap();
let body = Bytes::from(buffer);
let mut builder = self.client.request(Method::POST, url);
let digest = digest::digest(&digest::SHA256, &body);
builder = builder.header(SHA256_CHECKSUM, BASE64_STANDARD.encode(digest));
// S3 *requires* DeleteObjects to include a Content-MD5 header:
// https://docs.aws.amazon.com/AmazonS3/latest/API/API_DeleteObjects.html
// > "The Content-MD5 request header is required for all Multi-Object Delete requests"
// Some platforms, like MinIO, enforce this requirement and fail requests without the header.
let mut hasher = Md5::new();
hasher.update(&body);
builder = builder.header("Content-MD5", BASE64_STANDARD.encode(hasher.finalize()));
let response = builder
.header(CONTENT_TYPE, "application/xml")
.body(body)
.with_aws_sigv4(credential.authorizer(), Some(digest.as_ref()))
.send_retry(&self.config.retry_config)
.await
.map_err(|source| Error::DeleteObjectsRequest { source })?
.into_body()
.bytes()
.await
.map_err(|source| Error::DeleteObjectsResponse { source })?;
let response: BatchDeleteResponse =
quick_xml::de::from_reader(response.reader()).map_err(|err| {
Error::InvalidDeleteObjectsResponse {
source: Box::new(err),
}
})?;
// Assume all were ok, then fill in errors. This guarantees output order
// matches input order.
let mut results: Vec<Result<Path>> = paths.iter().cloned().map(Ok).collect();
for content in response.content.into_iter() {
if let DeleteObjectResult::Error(error) = content {
let path =
Path::parse(&error.key).map_err(|err| Error::InvalidDeleteObjectsResponse {
source: Box::new(err),
})?;
let i = paths.iter().find_position(|&p| p == &path).unwrap().0;
results[i] = Err(Error::from(error).into());
}
}
Ok(results)
}
/// Make an S3 Copy request <https://docs.aws.amazon.com/AmazonS3/latest/API/API_CopyObject.html>
pub(crate) fn copy_request<'a>(&'a self, from: &Path, to: &'a Path) -> Request<'a> {
let source = format!("{}/{}", self.config.bucket, encode_path(from));
let mut copy_source_encryption_headers = HeaderMap::new();
if let Some(customer_algorithm) = self
.config
.encryption_headers
.0
.get("x-amz-server-side-encryption-customer-algorithm")
{
copy_source_encryption_headers.insert(
"x-amz-copy-source-server-side-encryption-customer-algorithm",
customer_algorithm.clone(),
);
}
if let Some(customer_key) = self
.config
.encryption_headers
.0
.get("x-amz-server-side-encryption-customer-key")
{
copy_source_encryption_headers.insert(
"x-amz-copy-source-server-side-encryption-customer-key",
customer_key.clone(),
);
}
if let Some(customer_key_md5) = self
.config
.encryption_headers
.0
.get("x-amz-server-side-encryption-customer-key-MD5")
{
copy_source_encryption_headers.insert(
"x-amz-copy-source-server-side-encryption-customer-key-MD5",
customer_key_md5.clone(),
);
}
self.request(Method::PUT, to)
.idempotent(true)
.retry_error_body(true)
.header(&COPY_SOURCE_HEADER, &source)
.headers(self.config.encryption_headers.clone().into())
.headers(copy_source_encryption_headers)
.with_session_creds(false)
}
pub(crate) async fn create_multipart(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<MultipartId> {
let PutMultipartOpts {
tags,
attributes,
extensions,
} = opts;
let mut request = self.request(Method::POST, location);
if let Some(algorithm) = self.config.checksum {
match algorithm {
Checksum::SHA256 => {
request = request.header(ALGORITHM, "SHA256");
}
}
}
let response = request
.query(&[("uploads", "")])
.with_encryption_headers()
.with_attributes(attributes)
.with_tags(tags)
.with_extensions(extensions)
.idempotent(true)
.send()
.await?
.into_body()
.bytes()
.await
.map_err(|source| Error::CreateMultipartResponseBody { source })?;
let response: InitiateMultipartUploadResult = quick_xml::de::from_reader(response.reader())
.map_err(|source| Error::InvalidMultipartResponse { source })?;
Ok(response.upload_id)
}
pub(crate) async fn put_part(
&self,
path: &Path,
upload_id: &MultipartId,
part_idx: usize,
data: PutPartPayload<'_>,
) -> Result<PartId> {
let is_copy = matches!(data, PutPartPayload::Copy(_));
let part = (part_idx + 1).to_string();
let mut request = self
.request(Method::PUT, path)
.query(&[("partNumber", &part), ("uploadId", upload_id)])
.idempotent(true);
request = match data {
PutPartPayload::Part(payload) => request.with_payload(payload),
PutPartPayload::Copy(path) => request.header(
"x-amz-copy-source",
&format!("{}/{}", self.config.bucket, encode_path(path)),
),
};
if self
.config
.encryption_headers
.0
.contains_key("x-amz-server-side-encryption-customer-algorithm")
{
// If SSE-C is used, we must include the encryption headers in every upload request.
request = request.with_encryption_headers();
}
let (parts, body) = request.send().await?.into_parts();
let checksum_sha256 = parts
.headers
.get(SHA256_CHECKSUM)
.and_then(|v| v.to_str().ok())
.map(|v| v.to_string());
let e_tag = match is_copy {
false => get_etag(&parts.headers).map_err(|source| Error::Metadata { source })?,
true => {
let response = body
.bytes()
.await
.map_err(|source| Error::CreateMultipartResponseBody { source })?;
let response: CopyPartResult = quick_xml::de::from_reader(response.reader())
.map_err(|source| Error::InvalidMultipartResponse { source })?;
response.e_tag
}
};
let content_id = if self.config.checksum == Some(Checksum::SHA256) {
let meta = PartMetadata {
e_tag,
checksum_sha256,
};
quick_xml::se::to_string(&meta).unwrap()
} else {
e_tag
};
Ok(PartId { content_id })
}
pub(crate) async fn abort_multipart(&self, location: &Path, upload_id: &str) -> Result<()> {
self.request(Method::DELETE, location)
.query(&[("uploadId", upload_id)])
.with_encryption_headers()
.send()
.await?;
Ok(())
}
pub(crate) async fn complete_multipart(
&self,
location: &Path,
upload_id: &str,
parts: Vec<PartId>,
mode: CompleteMultipartMode,
) -> Result<PutResult> {
let parts = if parts.is_empty() {
// If no parts were uploaded, upload an empty part
// otherwise the completion request will fail
let part = self
.put_part(
location,
&upload_id.to_string(),
0,
PutPartPayload::default(),
)
.await?;
vec![part]
} else {
parts
};
let request = CompleteMultipartUpload::from(parts);
let body = quick_xml::se::to_string(&request).unwrap();
let credential = self.config.get_session_credential().await?;
let url = self.config.path_url(location);
let request = self
.client
.post(url)
.query(&[("uploadId", upload_id)])
.body(body)
.with_aws_sigv4(credential.authorizer(), None);
let request = match mode {
CompleteMultipartMode::Overwrite => request,
CompleteMultipartMode::Create => request.header("If-None-Match", "*"),
};
let response = request
.retryable(&self.config.retry_config)
.idempotent(true)
.retry_error_body(true)
.send()
.await
.map_err(|source| Error::CompleteMultipartRequest {
source,
path: location.as_ref().to_string(),
})?;
let version = get_version(response.headers(), VERSION_HEADER)
.map_err(|source| Error::Metadata { source })?;
let data = response
.into_body()
.bytes()
.await
.map_err(|source| Error::CompleteMultipartResponseBody { source })?;
let response: CompleteMultipartUploadResult = quick_xml::de::from_reader(data.reader())
.map_err(|source| Error::InvalidMultipartResponse { source })?;
Ok(PutResult {
e_tag: Some(response.e_tag),
version,
})
}
#[cfg(test)]
pub(crate) async fn get_object_tagging(&self, path: &Path) -> Result<HttpResponse> {
let credential = self.config.get_session_credential().await?;
let url = format!("{}?tagging", self.config.path_url(path));
let response = self
.client
.request(Method::GET, url)
.with_aws_sigv4(credential.authorizer(), None)
.send_retry(&self.config.retry_config)
.await
.map_err(|e| e.error(STORE, path.to_string()))?;
Ok(response)
}
}
#[async_trait]
impl GetClient for S3Client {
const STORE: &'static str = STORE;
const HEADER_CONFIG: HeaderConfig = HeaderConfig {
etag_required: false,
last_modified_required: false,
version_header: Some(VERSION_HEADER),
user_defined_metadata_prefix: Some(USER_DEFINED_METADATA_HEADER_PREFIX),
};
/// Make an S3 GET request <https://docs.aws.amazon.com/AmazonS3/latest/API/API_GetObject.html>
async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse> {
let credential = self.config.get_session_credential().await?;
let url = self.config.path_url(path);
let method = match options.head {
true => Method::HEAD,
false => Method::GET,
};
let mut builder = self.client.request(method, url);
if self
.config
.encryption_headers
.0
.contains_key("x-amz-server-side-encryption-customer-algorithm")
{
builder = builder.headers(self.config.encryption_headers.clone().into());
}
if let Some(v) = &options.version {
builder = builder.query(&[("versionId", v)])
}
let response = builder
.with_get_options(options)
.with_aws_sigv4(credential.authorizer(), None)
.send_retry(&self.config.retry_config)
.await
.map_err(|e| e.error(STORE, path.to_string()))?;
Ok(response)
}
}
#[async_trait]
impl ListClient for Arc<S3Client> {
/// Make an S3 List request <https://docs.aws.amazon.com/AmazonS3/latest/API/API_ListObjectsV2.html>
async fn list_request(
&self,
prefix: Option<&str>,
delimiter: bool,
token: Option<&str>,
offset: Option<&str>,
max_keys: Option<usize>,
) -> Result<(ListResult, Option<String>)> {
let credential = self.config.get_session_credential().await?;
let url = self.config.bucket_endpoint.clone();
let mut query = Vec::with_capacity(4);
if let Some(token) = token {
query.push(("continuation-token", token))
}
if delimiter {
query.push(("delimiter", DELIMITER))
}
query.push(("list-type", "2"));
if let Some(prefix) = prefix {
query.push(("prefix", prefix))
}
if let Some(offset) = offset {
query.push(("start-after", offset))
}
let max_keys = max_keys.map(|x| x.to_string());
if let Some(max_keys) = &max_keys {
query.push(("max-keys", max_keys))
}
let response = self
.client
.request(Method::GET, &url)
.query(&query)
.with_aws_sigv4(credential.authorizer(), None)
.send_retry(&self.config.retry_config)
.await
.map_err(|source| Error::ListRequest { source })?
.into_body()
.bytes()
.await
.map_err(|source| Error::ListResponseBody { source })?;
let mut response: ListResponse = quick_xml::de::from_reader(response.reader())
.map_err(|source| Error::InvalidListResponse { source })?;
let token = response.next_continuation_token.take();
Ok((response.try_into()?, token))
}
}
fn encode_path(path: &Path) -> PercentEncode<'_> {
utf8_percent_encode(path.as_ref(), &STRICT_PATH_ENCODE_SET)
}
backend/vendor/object_store/src/aws/credential.rs
Vendored
+1268
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/aws/dynamo.rs
Vendored
+594
View File
@@ -0,0 +1,594 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! A DynamoDB based lock system
use std::borrow::Cow;
use std::collections::HashMap;
use std::future::Future;
use std::time::{Duration, Instant};
use chrono::Utc;
use http::{Method, StatusCode};
use serde::ser::SerializeMap;
use serde::{Deserialize, Serialize, Serializer};
use crate::aws::client::S3Client;
use crate::aws::credential::CredentialExt;
use crate::aws::{AwsAuthorizer, AwsCredential};
use crate::client::get::GetClientExt;
use crate::client::retry::RetryExt;
use crate::client::retry::{RequestError, RetryError};
use crate::path::Path;
use crate::{Error, GetOptions, Result};
/// The exception returned by DynamoDB on conflict
const CONFLICT: &str = "ConditionalCheckFailedException";
const STORE: &str = "DynamoDB";
/// A DynamoDB-based commit protocol, used to provide conditional write support for S3
///
/// ## Limitations
///
/// Only conditional operations, e.g. `copy_if_not_exists` will be synchronized, and can
/// therefore race with non-conditional operations, e.g. `put`, `copy`, `delete`, or
/// conditional operations performed by writers not configured to synchronize with DynamoDB.
///
/// Workloads making use of this mechanism **must** ensure:
///
/// * Conditional and non-conditional operations are not performed on the same paths
/// * Conditional operations are only performed via similarly configured clients
///
/// Additionally as the locking mechanism relies on timeouts to detect stale locks,
/// performance will be poor for systems that frequently delete and then create
/// objects at the same path, instead being optimised for systems that primarily create
/// files with paths never used before, or perform conditional updates to existing files
///
/// ## Commit Protocol
///
/// The DynamoDB schema is as follows:
///
/// * A string partition key named `"path"`
/// * A string sort key named `"etag"`
/// * A numeric [TTL] attribute named `"ttl"`
/// * A numeric attribute named `"generation"`
/// * A numeric attribute named `"timeout"`
///
/// An appropriate DynamoDB table can be created with the CLI as follows:
///
/// ```bash
/// $ aws dynamodb create-table --table-name <TABLE_NAME> --key-schema AttributeName=path,KeyType=HASH AttributeName=etag,KeyType=RANGE --attribute-definitions AttributeName=path,AttributeType=S AttributeName=etag,AttributeType=S
/// $ aws dynamodb update-time-to-live --table-name <TABLE_NAME> --time-to-live-specification Enabled=true,AttributeName=ttl
/// ```
///
/// To perform a conditional operation on an object with a given `path` and `etag` (`*` if creating),
/// the commit protocol is as follows:
///
/// 1. Perform HEAD request on `path` and error on precondition mismatch
/// 2. Create record in DynamoDB with given `path` and `etag` with the configured timeout
/// 1. On Success: Perform operation with the configured timeout
/// 2. On Conflict:
/// 1. Periodically re-perform HEAD request on `path` and error on precondition mismatch
/// 2. If `timeout * max_skew_rate` passed, replace the record incrementing the `"generation"`
/// 1. On Success: GOTO 2.1
/// 2. On Conflict: GOTO 2.2
///
/// Provided no writer modifies an object with a given `path` and `etag` without first adding a
/// corresponding record to DynamoDB, we are guaranteed that only one writer will ever commit.
///
/// This is inspired by the [DynamoDB Lock Client] but simplified for the more limited
/// requirements of synchronizing object storage. The major changes are:
///
/// * Uses a monotonic generation count instead of a UUID rvn, as this is:
/// * Cheaper to generate, serialize and compare
/// * Cannot collide
/// * More human readable / interpretable
/// * Relies on [TTL] to eventually clean up old locks
///
/// It also draws inspiration from the DeltaLake [S3 Multi-Cluster] commit protocol, but
/// generalised to not make assumptions about the workload and not rely on first writing
/// to a temporary path.
///
/// [TTL]: https://docs.aws.amazon.com/amazondynamodb/latest/developerguide/howitworks-ttl.html
/// [DynamoDB Lock Client]: https://aws.amazon.com/blogs/database/building-distributed-locks-with-the-dynamodb-lock-client/
/// [S3 Multi-Cluster]: https://docs.google.com/document/d/1Gs4ZsTH19lMxth4BSdwlWjUNR-XhKHicDvBjd2RqNd8/edit#heading=h.mjjuxw9mcz9h
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct DynamoCommit {
table_name: String,
/// The number of milliseconds a lease is valid for
timeout: u64,
/// The maximum clock skew rate tolerated by the system
max_clock_skew_rate: u32,
/// The length of time a record will be retained in DynamoDB before being cleaned up
///
/// This is purely an optimisation to avoid indefinite growth of the DynamoDB table
/// and does not impact how long clients may wait to acquire a lock
ttl: Duration,
/// The backoff duration before retesting a condition
test_interval: Duration,
}
impl DynamoCommit {
/// Create a new [`DynamoCommit`] with a given table name
pub fn new(table_name: String) -> Self {
Self {
table_name,
timeout: 20_000,
max_clock_skew_rate: 3,
ttl: Duration::from_secs(60 * 60),
test_interval: Duration::from_millis(100),
}
}
/// Overrides the lock timeout.
///
/// A longer lock timeout reduces the probability of spurious commit failures and multi-writer
/// races, but will increase the time that writers must wait to reclaim a lock lost. The
/// default value of 20 seconds should be appropriate for must use-cases.
pub fn with_timeout(mut self, millis: u64) -> Self {
self.timeout = millis;
self
}
/// The maximum clock skew rate tolerated by the system.
///
/// An environment in which the clock on the fastest node ticks twice as fast as the slowest
/// node, would have a clock skew rate of 2. The default value of 3 should be appropriate
/// for most environments.
pub fn with_max_clock_skew_rate(mut self, rate: u32) -> Self {
self.max_clock_skew_rate = rate;
self
}
/// The length of time a record should be retained in DynamoDB before being cleaned up
///
/// This should be significantly larger than the configured lock timeout, with the default
/// value of 1 hour appropriate for most use-cases.
pub fn with_ttl(mut self, ttl: Duration) -> Self {
self.ttl = ttl;
self
}
/// Parse [`DynamoCommit`] from a string
pub(crate) fn from_str(value: &str) -> Option<Self> {
Some(match value.split_once(':') {
Some((table_name, timeout)) => {
Self::new(table_name.trim().to_string()).with_timeout(timeout.parse().ok()?)
}
None => Self::new(value.trim().to_string()),
})
}
/// Returns the name of the DynamoDB table.
pub(crate) fn table_name(&self) -> &str {
&self.table_name
}
pub(crate) async fn copy_if_not_exists(
&self,
client: &S3Client,
from: &Path,
to: &Path,
) -> Result<()> {
self.conditional_op(client, to, None, || async {
client.copy_request(from, to).send().await?;
Ok(())
})
.await
}
#[allow(clippy::future_not_send)] // Generics confound this lint
pub(crate) async fn conditional_op<F, Fut, T>(
&self,
client: &S3Client,
to: &Path,
etag: Option<&str>,
op: F,
) -> Result<T>
where
F: FnOnce() -> Fut,
Fut: Future<Output = Result<T, Error>>,
{
check_precondition(client, to, etag).await?;
let mut previous_lease = None;
loop {
let existing = previous_lease.as_ref();
match self.try_lock(client, to.as_ref(), etag, existing).await? {
TryLockResult::Ok(lease) => {
let expiry = lease.acquire + lease.timeout;
return match tokio::time::timeout_at(expiry.into(), op()).await {
Ok(Ok(v)) => Ok(v),
Ok(Err(e)) => Err(e),
Err(_) => Err(Error::Generic {
store: "DynamoDB",
source: format!(
"Failed to perform conditional operation in {} milliseconds",
self.timeout
)
.into(),
}),
};
}
TryLockResult::Conflict(conflict) => {
let mut interval = tokio::time::interval(self.test_interval);
let expiry = conflict.timeout * self.max_clock_skew_rate;
loop {
interval.tick().await;
check_precondition(client, to, etag).await?;
if conflict.acquire.elapsed() > expiry {
previous_lease = Some(conflict);
break;
}
}
}
}
}
}
/// Attempt to acquire a lock, reclaiming an existing lease if provided
async fn try_lock(
&self,
s3: &S3Client,
path: &str,
etag: Option<&str>,
existing: Option<&Lease>,
) -> Result<TryLockResult> {
let attributes;
let (next_gen, condition_expression, expression_attribute_values) = match existing {
None => (0_u64, "attribute_not_exists(#pk)", Map(&[])),
Some(existing) => {
attributes = [(":g", AttributeValue::Number(existing.generation))];
(
existing.generation.checked_add(1).unwrap(),
"attribute_exists(#pk) AND generation = :g",
Map(attributes.as_slice()),
)
}
};
let ttl = (Utc::now() + self.ttl).timestamp();
let items = [
("path", AttributeValue::from(path)),
("etag", AttributeValue::from(etag.unwrap_or("*"))),
("generation", AttributeValue::Number(next_gen)),
("timeout", AttributeValue::Number(self.timeout)),
("ttl", AttributeValue::Number(ttl as _)),
];
let names = [("#pk", "path")];
let req = PutItem {
table_name: &self.table_name,
condition_expression,
expression_attribute_values,
expression_attribute_names: Map(&names),
item: Map(&items),
return_values: None,
return_values_on_condition_check_failure: Some(ReturnValues::AllOld),
};
let credential = s3.config.get_credential().await?;
let acquire = Instant::now();
match self
.request(s3, credential.as_deref(), "DynamoDB_20120810.PutItem", req)
.await
{
Ok(_) => Ok(TryLockResult::Ok(Lease {
acquire,
generation: next_gen,
timeout: Duration::from_millis(self.timeout),
})),
Err(e) => match parse_error_response(&e) {
Some(e) if e.error.ends_with(CONFLICT) => match extract_lease(&e.item) {
Some(lease) => Ok(TryLockResult::Conflict(lease)),
None => Err(Error::Generic {
store: STORE,
source: "Failed to extract lease from conflict ReturnValuesOnConditionCheckFailure response".into()
}),
},
_ => Err(Error::Generic {
store: STORE,
source: Box::new(e),
}),
},
}
}
async fn request<R: Serialize + Send + Sync>(
&self,
s3: &S3Client,
cred: Option<&AwsCredential>,
target: &str,
req: R,
) -> Result<HttpResponse, RetryError> {
let region = &s3.config.region;
let authorizer = cred.map(|x| AwsAuthorizer::new(x, "dynamodb", region));
let builder = match &s3.config.endpoint {
Some(e) => s3.client.request(Method::POST, e),
None => {
let url = format!("https://dynamodb.{region}.amazonaws.com");
s3.client.request(Method::POST, url)
}
};
// TODO: Timeout
builder
.json(&req)
.header("X-Amz-Target", target)
.with_aws_sigv4(authorizer, None)
.send_retry(&s3.config.retry_config)
.await
}
}
#[derive(Debug)]
enum TryLockResult {
/// Successfully acquired a lease
Ok(Lease),
/// An existing lease was found
Conflict(Lease),
}
/// Validates that `path` has the given `etag` or doesn't exist if `None`
async fn check_precondition(client: &S3Client, path: &Path, etag: Option<&str>) -> Result<()> {
let options = GetOptions {
head: true,
..Default::default()
};
match etag {
Some(expected) => match client.get_opts(path, options).await {
Ok(r) => match r.meta.e_tag {
Some(actual) if expected == actual => Ok(()),
actual => Err(Error::Precondition {
path: path.to_string(),
source: format!("{} does not match {expected}", actual.unwrap_or_default())
.into(),
}),
},
Err(Error::NotFound { .. }) => Err(Error::Precondition {
path: path.to_string(),
source: format!("Object at location {path} not found").into(),
}),
Err(e) => Err(e),
},
None => match client.get_opts(path, options).await {
Ok(_) => Err(Error::AlreadyExists {
path: path.to_string(),
source: "Already Exists".to_string().into(),
}),
Err(Error::NotFound { .. }) => Ok(()),
Err(e) => Err(e),
},
}
}
/// Parses the error response if any
fn parse_error_response(e: &RetryError) -> Option<ErrorResponse<'_>> {
match e.inner() {
RequestError::Status {
status: StatusCode::BAD_REQUEST,
body: Some(b),
} => serde_json::from_str(b).ok(),
_ => None,
}
}
/// Extracts a lease from `item`, returning `None` on error
fn extract_lease(item: &HashMap<&str, AttributeValue<'_>>) -> Option<Lease> {
let generation = match item.get("generation") {
Some(AttributeValue::Number(generation)) => generation,
_ => return None,
};
let timeout = match item.get("timeout") {
Some(AttributeValue::Number(timeout)) => *timeout,
_ => return None,
};
Some(Lease {
acquire: Instant::now(),
generation: *generation,
timeout: Duration::from_millis(timeout),
})
}
/// A lock lease
#[derive(Debug, Clone)]
struct Lease {
acquire: Instant,
generation: u64,
timeout: Duration,
}
/// A DynamoDB [PutItem] payload
///
/// [PutItem]: https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_PutItem.html
#[derive(Serialize)]
#[serde(rename_all = "PascalCase")]
struct PutItem<'a> {
/// The table name
table_name: &'a str,
/// A condition that must be satisfied in order for a conditional PutItem operation to succeed.
condition_expression: &'a str,
/// One or more substitution tokens for attribute names in an expression
expression_attribute_names: Map<'a, &'a str, &'a str>,
/// One or more values that can be substituted in an expression
expression_attribute_values: Map<'a, &'a str, AttributeValue<'a>>,
/// A map of attribute name/value pairs, one for each attribute
item: Map<'a, &'a str, AttributeValue<'a>>,
/// Use ReturnValues if you want to get the item attributes as they appeared
/// before they were updated with the PutItem request.
#[serde(skip_serializing_if = "Option::is_none")]
return_values: Option<ReturnValues>,
/// An optional parameter that returns the item attributes for a PutItem operation
/// that failed a condition check.
#[serde(skip_serializing_if = "Option::is_none")]
return_values_on_condition_check_failure: Option<ReturnValues>,
}
#[derive(Deserialize)]
struct ErrorResponse<'a> {
#[serde(rename = "__type")]
error: &'a str,
#[serde(borrow, default, rename = "Item")]
item: HashMap<&'a str, AttributeValue<'a>>,
}
#[derive(Serialize)]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
enum ReturnValues {
AllOld,
}
/// A collection of key value pairs
///
/// This provides cheap, ordered serialization of maps
struct Map<'a, K, V>(&'a [(K, V)]);
impl<K: Serialize, V: Serialize> Serialize for Map<'_, K, V> {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
if self.0.is_empty() {
return serializer.serialize_none();
}
let mut map = serializer.serialize_map(Some(self.0.len()))?;
for (k, v) in self.0 {
map.serialize_entry(k, v)?
}
map.end()
}
}
/// A DynamoDB [AttributeValue]
///
/// [AttributeValue]: https://docs.aws.amazon.com/amazondynamodb/latest/APIReference/API_AttributeValue.html
#[derive(Debug, Serialize, Deserialize)]
enum AttributeValue<'a> {
#[serde(rename = "S")]
String(Cow<'a, str>),
#[serde(rename = "N", with = "number")]
Number(u64),
}
impl<'a> From<&'a str> for AttributeValue<'a> {
fn from(value: &'a str) -> Self {
Self::String(Cow::Borrowed(value))
}
}
/// Numbers are serialized as strings
mod number {
use serde::{Deserialize, Deserializer, Serializer};
pub(crate) fn serialize<S: Serializer>(v: &u64, s: S) -> Result<S::Ok, S::Error> {
s.serialize_str(&v.to_string())
}
pub(crate) fn deserialize<'de, D: Deserializer<'de>>(d: D) -> Result<u64, D::Error> {
let v: &str = Deserialize::deserialize(d)?;
v.parse().map_err(serde::de::Error::custom)
}
}
use crate::client::HttpResponse;
/// Re-export integration_test to be called by s3_test
#[cfg(test)]
pub(crate) use tests::integration_test;
#[cfg(test)]
mod tests {
use super::*;
use crate::aws::AmazonS3;
use crate::ObjectStore;
use rand::distr::Alphanumeric;
use rand::{rng, Rng};
#[test]
fn test_attribute_serde() {
let serde = serde_json::to_string(&AttributeValue::Number(23)).unwrap();
assert_eq!(serde, "{\"N\":\"23\"}");
let back: AttributeValue<'_> = serde_json::from_str(&serde).unwrap();
assert!(matches!(back, AttributeValue::Number(23)));
}
/// An integration test for DynamoDB
///
/// This is a function called by s3_test to avoid test concurrency issues
pub(crate) async fn integration_test(integration: &AmazonS3, d: &DynamoCommit) {
let client = integration.client.as_ref();
let src = Path::from("dynamo_path_src");
integration.put(&src, "asd".into()).await.unwrap();
let dst = Path::from("dynamo_path");
let _ = integration.delete(&dst).await; // Delete if present
// Create a lock if not already exists
let existing = match d.try_lock(client, dst.as_ref(), None, None).await.unwrap() {
TryLockResult::Conflict(l) => l,
TryLockResult::Ok(l) => l,
};
// Should not be able to acquire a lock again
let r = d.try_lock(client, dst.as_ref(), None, None).await;
assert!(matches!(r, Ok(TryLockResult::Conflict(_))));
// But should still be able to reclaim lock and perform copy
d.copy_if_not_exists(client, &src, &dst).await.unwrap();
match d.try_lock(client, dst.as_ref(), None, None).await.unwrap() {
TryLockResult::Conflict(new) => {
// Should have incremented generation to do so
assert_eq!(new.generation, existing.generation + 1);
}
_ => panic!("Should conflict"),
}
let rng = rng();
let etag = String::from_utf8(rng.sample_iter(Alphanumeric).take(32).collect()).unwrap();
let t = Some(etag.as_str());
let l = match d.try_lock(client, dst.as_ref(), t, None).await.unwrap() {
TryLockResult::Ok(l) => l,
_ => panic!("should not conflict"),
};
match d.try_lock(client, dst.as_ref(), t, None).await.unwrap() {
TryLockResult::Conflict(c) => assert_eq!(l.generation, c.generation),
_ => panic!("should conflict"),
}
match d.try_lock(client, dst.as_ref(), t, Some(&l)).await.unwrap() {
TryLockResult::Ok(new) => assert_eq!(new.generation, l.generation + 1),
_ => panic!("should not conflict"),
}
}
}
backend/vendor/object_store/src/aws/mod.rs
Vendored
+886
View File
@@ -0,0 +1,886 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An object store implementation for S3
//!
//! ## Multipart uploads
//!
//! Multipart uploads can be initiated with the [ObjectStore::put_multipart] method.
//!
//! If the writer fails for any reason, you may have parts uploaded to AWS but not
//! used that you will be charged for. [`MultipartUpload::abort`] may be invoked to drop
//! these unneeded parts, however, it is recommended that you consider implementing
//! [automatic cleanup] of unused parts that are older than some threshold.
//!
//! [automatic cleanup]: https://aws.amazon.com/blogs/aws/s3-lifecycle-management-update-support-for-multipart-uploads-and-delete-markers/
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::{StreamExt, TryStreamExt};
use reqwest::header::{HeaderName, IF_MATCH, IF_NONE_MATCH};
use reqwest::{Method, StatusCode};
use std::{sync::Arc, time::Duration};
use url::Url;
use crate::aws::client::{CompleteMultipartMode, PutPartPayload, RequestError, S3Client};
use crate::client::get::GetClientExt;
use crate::client::list::ListClientExt;
use crate::client::CredentialProvider;
use crate::multipart::{MultipartStore, PartId};
use crate::signer::Signer;
use crate::util::STRICT_ENCODE_SET;
use crate::{
Error, GetOptions, GetResult, ListPage, ListResult, MultipartId, MultipartUpload, ObjectMeta,
ObjectStore, Path, PutMode, PutMultipartOpts, PutOptions, PutPayload, PutResult, Result,
UploadPart,
};
static TAGS_HEADER: HeaderName = HeaderName::from_static("x-amz-tagging");
static COPY_SOURCE_HEADER: HeaderName = HeaderName::from_static("x-amz-copy-source");
mod builder;
mod checksum;
mod client;
mod credential;
mod dynamo;
mod precondition;
#[cfg(not(target_arch = "wasm32"))]
mod resolve;
pub use builder::{AmazonS3Builder, AmazonS3ConfigKey};
pub use checksum::Checksum;
pub use dynamo::DynamoCommit;
pub use precondition::{S3ConditionalPut, S3CopyIfNotExists};
#[cfg(not(target_arch = "wasm32"))]
pub use resolve::resolve_bucket_region;
/// This struct is used to maintain the URI path encoding
const STRICT_PATH_ENCODE_SET: percent_encoding::AsciiSet = STRICT_ENCODE_SET.remove(b'/');
const STORE: &str = "S3";
/// [`CredentialProvider`] for [`AmazonS3`]
pub type AwsCredentialProvider = Arc<dyn CredentialProvider<Credential = AwsCredential>>;
use crate::client::parts::Parts;
pub use credential::{AwsAuthorizer, AwsCredential};
/// Interface for [Amazon S3](https://aws.amazon.com/s3/).
#[derive(Debug, Clone)]
pub struct AmazonS3 {
client: Arc<S3Client>,
}
impl std::fmt::Display for AmazonS3 {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "AmazonS3({})", self.client.config.bucket)
}
}
impl AmazonS3 {
/// Returns the [`AwsCredentialProvider`] used by [`AmazonS3`]
pub fn credentials(&self) -> &AwsCredentialProvider {
&self.client.config.credentials
}
/// Create a full URL to the resource specified by `path` with this instance's configuration.
fn path_url(&self, path: &Path) -> String {
self.client.config.path_url(path)
}
}
#[async_trait]
impl Signer for AmazonS3 {
/// Create a URL containing the relevant [AWS SigV4] query parameters that authorize a request
/// via `method` to the resource at `path` valid for the duration specified in `expires_in`.
///
/// [AWS SigV4]: https://docs.aws.amazon.com/IAM/latest/UserGuide/create-signed-request.html
///
/// # Example
///
/// This example returns a URL that will enable a user to upload a file to
/// "some-folder/some-file.txt" in the next hour.
///
/// ```
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// # use object_store::{aws::AmazonS3Builder, path::Path, signer::Signer};
/// # use reqwest::Method;
/// # use std::time::Duration;
/// #
/// let region = "us-east-1";
/// let s3 = AmazonS3Builder::new()
/// .with_region(region)
/// .with_bucket_name("my-bucket")
/// .with_access_key_id("my-access-key-id")
/// .with_secret_access_key("my-secret-access-key")
/// .build()?;
///
/// let url = s3.signed_url(
/// Method::PUT,
/// &Path::from("some-folder/some-file.txt"),
/// Duration::from_secs(60 * 60)
/// ).await?;
/// # Ok(())
/// # }
/// ```
async fn signed_url(&self, method: Method, path: &Path, expires_in: Duration) -> Result<Url> {
let credential = self.credentials().get_credential().await?;
let authorizer = AwsAuthorizer::new(&credential, "s3", &self.client.config.region)
.with_request_payer(self.client.config.request_payer);
let path_url = self.path_url(path);
let mut url = path_url.parse().map_err(|e| Error::Generic {
store: STORE,
source: format!("Unable to parse url {path_url}: {e}").into(),
})?;
authorizer.sign(method, &mut url, expires_in);
Ok(url)
}
}
#[async_trait]
impl ObjectStore for AmazonS3 {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
let PutOptions {
mode,
tags,
attributes,
extensions,
} = opts;
let request = self
.client
.request(Method::PUT, location)
.with_payload(payload)
.with_attributes(attributes)
.with_tags(tags)
.with_extensions(extensions)
.with_encryption_headers();
match (mode, &self.client.config.conditional_put) {
(PutMode::Overwrite, _) => request.idempotent(true).do_put().await,
(PutMode::Create, S3ConditionalPut::Disabled) => Err(Error::NotImplemented),
(PutMode::Create, S3ConditionalPut::ETagMatch) => {
match request.header(&IF_NONE_MATCH, "*").do_put().await {
// Technically If-None-Match should return NotModified but some stores,
// such as R2, instead return PreconditionFailed
// https://developers.cloudflare.com/r2/api/s3/extensions/#conditional-operations-in-putobject
Err(e @ Error::NotModified { .. } | e @ Error::Precondition { .. }) => {
Err(Error::AlreadyExists {
path: location.to_string(),
source: Box::new(e),
})
}
r => r,
}
}
(PutMode::Create, S3ConditionalPut::Dynamo(d)) => {
d.conditional_op(&self.client, location, None, move || request.do_put())
.await
}
(PutMode::Update(v), put) => {
let etag = v.e_tag.ok_or_else(|| Error::Generic {
store: STORE,
source: "ETag required for conditional put".to_string().into(),
})?;
match put {
S3ConditionalPut::ETagMatch => {
match request
.header(&IF_MATCH, etag.as_str())
// Real S3 will occasionally report 409 Conflict
// if there are concurrent `If-Match` requests
// in flight, so we need to be prepared to retry
// 409 responses.
.retry_on_conflict(true)
.do_put()
.await
{
// Real S3 reports NotFound rather than PreconditionFailed when the
// object doesn't exist. Convert to PreconditionFailed for
// consistency with R2. This also matches what the HTTP spec
// says the behavior should be.
Err(Error::NotFound { path, source }) => {
Err(Error::Precondition { path, source })
}
r => r,
}
}
S3ConditionalPut::Dynamo(d) => {
d.conditional_op(&self.client, location, Some(&etag), move || {
request.do_put()
})
.await
}
S3ConditionalPut::Disabled => Err(Error::NotImplemented),
}
}
}
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
let upload_id = self.client.create_multipart(location, opts).await?;
Ok(Box::new(S3MultiPartUpload {
part_idx: 0,
state: Arc::new(UploadState {
client: Arc::clone(&self.client),
location: location.clone(),
upload_id: upload_id.clone(),
parts: Default::default(),
}),
}))
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
self.client.get_opts(location, options).await
}
async fn delete(&self, location: &Path) -> Result<()> {
self.client.request(Method::DELETE, location).send().await?;
Ok(())
}
fn delete_stream<'a>(
&'a self,
locations: BoxStream<'a, Result<Path>>,
) -> BoxStream<'a, Result<Path>> {
locations
.try_chunks(1_000)
.map(move |locations| async {
// Early return the error. We ignore the paths that have already been
// collected into the chunk.
let locations = locations.map_err(|e| e.1)?;
self.client
.bulk_delete_request(locations)
.await
.map(futures::stream::iter)
})
.buffered(20)
.try_flatten()
.boxed()
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
self.client.list(prefix)
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
if self.client.config.is_s3_express() {
let offset = offset.clone();
// S3 Express does not support start-after
return self
.client
.list(prefix)
.try_filter(move |f| futures::future::ready(f.location > offset))
.boxed();
}
self.client.list_with_offset(prefix, offset)
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
self.client.list_with_delimiter(prefix).await
}
async fn list_delimited_page(
&self,
prefix: Option<&Path>,
token: Option<&str>,
max_keys: Option<usize>,
) -> Result<ListPage> {
self.client.list_delimited_page(prefix, token, max_keys).await
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
self.client
.copy_request(from, to)
.idempotent(true)
.send()
.await?;
Ok(())
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
let (k, v, status) = match &self.client.config.copy_if_not_exists {
Some(S3CopyIfNotExists::Header(k, v)) => (k, v, StatusCode::PRECONDITION_FAILED),
Some(S3CopyIfNotExists::HeaderWithStatus(k, v, status)) => (k, v, *status),
Some(S3CopyIfNotExists::Multipart) => {
let upload_id = self
.client
.create_multipart(to, PutMultipartOpts::default())
.await?;
let res = async {
let part_id = self
.client
.put_part(to, &upload_id, 0, PutPartPayload::Copy(from))
.await?;
match self
.client
.complete_multipart(
to,
&upload_id,
vec![part_id],
CompleteMultipartMode::Create,
)
.await
{
Err(e @ Error::Precondition { .. }) => Err(Error::AlreadyExists {
path: to.to_string(),
source: Box::new(e),
}),
Ok(_) => Ok(()),
Err(e) => Err(e),
}
}
.await;
// If the multipart upload failed, make a best effort attempt to
// clean it up. It's the caller's responsibility to add a
// lifecycle rule if guaranteed cleanup is required, as we
// cannot protect against an ill-timed process crash.
if res.is_err() {
let _ = self.client.abort_multipart(to, &upload_id).await;
}
return res;
}
Some(S3CopyIfNotExists::Dynamo(lock)) => {
return lock.copy_if_not_exists(&self.client, from, to).await
}
None => {
return Err(Error::NotSupported {
source: "S3 does not support copy-if-not-exists".to_string().into(),
})
}
};
let req = self.client.copy_request(from, to);
match req.header(k, v).send().await {
Err(RequestError::Retry { source, path }) if source.status() == Some(status) => {
Err(Error::AlreadyExists {
source: Box::new(source),
path,
})
}
Err(e) => Err(e.into()),
Ok(_) => Ok(()),
}
}
}
#[derive(Debug)]
struct S3MultiPartUpload {
part_idx: usize,
state: Arc<UploadState>,
}
#[derive(Debug)]
struct UploadState {
parts: Parts,
location: Path,
upload_id: String,
client: Arc<S3Client>,
}
#[async_trait]
impl MultipartUpload for S3MultiPartUpload {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
let idx = self.part_idx;
self.part_idx += 1;
let state = Arc::clone(&self.state);
Box::pin(async move {
let part = state
.client
.put_part(
&state.location,
&state.upload_id,
idx,
PutPartPayload::Part(data),
)
.await?;
state.parts.put(idx, part);
Ok(())
})
}
async fn complete(&mut self) -> Result<PutResult> {
let parts = self.state.parts.finish(self.part_idx)?;
self.state
.client
.complete_multipart(
&self.state.location,
&self.state.upload_id,
parts,
CompleteMultipartMode::Overwrite,
)
.await
}
async fn abort(&mut self) -> Result<()> {
self.state
.client
.request(Method::DELETE, &self.state.location)
.query(&[("uploadId", &self.state.upload_id)])
.idempotent(true)
.send()
.await?;
Ok(())
}
}
#[async_trait]
impl MultipartStore for AmazonS3 {
async fn create_multipart(&self, path: &Path) -> Result<MultipartId> {
self.client
.create_multipart(path, PutMultipartOpts::default())
.await
}
async fn put_part(
&self,
path: &Path,
id: &MultipartId,
part_idx: usize,
data: PutPayload,
) -> Result<PartId> {
self.client
.put_part(path, id, part_idx, PutPartPayload::Part(data))
.await
}
async fn complete_multipart(
&self,
path: &Path,
id: &MultipartId,
parts: Vec<PartId>,
) -> Result<PutResult> {
self.client
.complete_multipart(path, id, parts, CompleteMultipartMode::Overwrite)
.await
}
async fn abort_multipart(&self, path: &Path, id: &MultipartId) -> Result<()> {
self.client
.request(Method::DELETE, path)
.query(&[("uploadId", id)])
.send()
.await?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::get::GetClient;
use crate::client::SpawnedReqwestConnector;
use crate::integration::*;
use crate::tests::*;
use crate::ClientOptions;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use http::HeaderMap;
const NON_EXISTENT_NAME: &str = "nonexistentname";
#[tokio::test]
async fn write_multipart_file_with_signature() {
maybe_skip_integration!();
let store = AmazonS3Builder::from_env()
.with_checksum_algorithm(Checksum::SHA256)
.build()
.unwrap();
let str = "test.bin";
let path = Path::parse(str).unwrap();
let opts = PutMultipartOpts::default();
let mut upload = store.put_multipart_opts(&path, opts).await.unwrap();
upload
.put_part(PutPayload::from(vec![0u8; 10_000_000]))
.await
.unwrap();
upload
.put_part(PutPayload::from(vec![0u8; 5_000_000]))
.await
.unwrap();
let res = upload.complete().await.unwrap();
assert!(res.e_tag.is_some(), "Should have valid etag");
store.delete(&path).await.unwrap();
}
#[tokio::test]
async fn write_multipart_file_with_signature_object_lock() {
maybe_skip_integration!();
let bucket = "test-object-lock";
let store = AmazonS3Builder::from_env()
.with_bucket_name(bucket)
.with_checksum_algorithm(Checksum::SHA256)
.build()
.unwrap();
let str = "test.bin";
let path = Path::parse(str).unwrap();
let opts = PutMultipartOpts::default();
let mut upload = store.put_multipart_opts(&path, opts).await.unwrap();
upload
.put_part(PutPayload::from(vec![0u8; 10_000_000]))
.await
.unwrap();
upload
.put_part(PutPayload::from(vec![0u8; 5_000_000]))
.await
.unwrap();
let res = upload.complete().await.unwrap();
assert!(res.e_tag.is_some(), "Should have valid etag");
store.delete(&path).await.unwrap();
}
#[tokio::test]
async fn s3_test() {
maybe_skip_integration!();
let config = AmazonS3Builder::from_env();
let integration = config.build().unwrap();
let config = &integration.client.config;
let test_not_exists = config.copy_if_not_exists.is_some();
let test_conditional_put = config.conditional_put != S3ConditionalPut::Disabled;
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
stream_get(&integration).await;
multipart(&integration, &integration).await;
multipart_race_condition(&integration, true).await;
multipart_out_of_order(&integration).await;
signing(&integration).await;
s3_encryption(&integration).await;
put_get_attributes(&integration).await;
// Object tagging is not supported by S3 Express One Zone
if config.session_provider.is_none() {
tagging(
Arc::new(AmazonS3 {
client: Arc::clone(&integration.client),
}),
!config.disable_tagging,
|p| {
let client = Arc::clone(&integration.client);
async move { client.get_object_tagging(&p).await }
},
)
.await;
}
if test_not_exists {
copy_if_not_exists(&integration).await;
}
if test_conditional_put {
put_opts(&integration, true).await;
}
// run integration test with unsigned payload enabled
let builder = AmazonS3Builder::from_env().with_unsigned_payload(true);
let integration = builder.build().unwrap();
put_get_delete_list(&integration).await;
// run integration test with checksum set to sha256
let builder = AmazonS3Builder::from_env().with_checksum_algorithm(Checksum::SHA256);
let integration = builder.build().unwrap();
put_get_delete_list(&integration).await;
match &integration.client.config.copy_if_not_exists {
Some(S3CopyIfNotExists::Dynamo(d)) => dynamo::integration_test(&integration, d).await,
_ => eprintln!("Skipping dynamo integration test - dynamo not configured"),
};
}
#[tokio::test]
async fn s3_test_get_nonexistent_location() {
maybe_skip_integration!();
let integration = AmazonS3Builder::from_env().build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = get_nonexistent_object(&integration, Some(location))
.await
.unwrap_err();
assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
}
#[tokio::test]
async fn s3_test_get_nonexistent_bucket() {
maybe_skip_integration!();
let config = AmazonS3Builder::from_env().with_bucket_name(NON_EXISTENT_NAME);
let integration = config.build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = integration.get(&location).await.unwrap_err();
assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
}
#[tokio::test]
async fn s3_test_put_nonexistent_bucket() {
maybe_skip_integration!();
let config = AmazonS3Builder::from_env().with_bucket_name(NON_EXISTENT_NAME);
let integration = config.build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let data = PutPayload::from("arbitrary data");
let err = integration.put(&location, data).await.unwrap_err();
assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
}
#[tokio::test]
async fn s3_test_delete_nonexistent_location() {
maybe_skip_integration!();
let integration = AmazonS3Builder::from_env().build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
integration.delete(&location).await.unwrap();
}
#[tokio::test]
async fn s3_test_delete_nonexistent_bucket() {
maybe_skip_integration!();
let config = AmazonS3Builder::from_env().with_bucket_name(NON_EXISTENT_NAME);
let integration = config.build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = integration.delete(&location).await.unwrap_err();
assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
}
#[tokio::test]
#[ignore = "Tests shouldn't call use remote services by default"]
async fn test_disable_creds() {
// https://registry.opendata.aws/daylight-osm/
let v1 = AmazonS3Builder::new()
.with_bucket_name("daylight-map-distribution")
.with_region("us-west-1")
.with_access_key_id("local")
.with_secret_access_key("development")
.build()
.unwrap();
let prefix = Path::from("release");
v1.list_with_delimiter(Some(&prefix)).await.unwrap_err();
let v2 = AmazonS3Builder::new()
.with_bucket_name("daylight-map-distribution")
.with_region("us-west-1")
.with_skip_signature(true)
.build()
.unwrap();
v2.list_with_delimiter(Some(&prefix)).await.unwrap();
}
async fn s3_encryption(store: &AmazonS3) {
maybe_skip_integration!();
let data = PutPayload::from(vec![3u8; 1024]);
let encryption_headers: HeaderMap = store.client.config.encryption_headers.clone().into();
let expected_encryption =
if let Some(encryption_type) = encryption_headers.get("x-amz-server-side-encryption") {
encryption_type
} else {
eprintln!("Skipping S3 encryption test - encryption not configured");
return;
};
let locations = [
Path::from("test-encryption-1"),
Path::from("test-encryption-2"),
Path::from("test-encryption-3"),
];
store.put(&locations[0], data.clone()).await.unwrap();
store.copy(&locations[0], &locations[1]).await.unwrap();
let mut upload = store.put_multipart(&locations[2]).await.unwrap();
upload.put_part(data.clone()).await.unwrap();
upload.complete().await.unwrap();
for location in &locations {
let res = store
.client
.get_request(location, GetOptions::default())
.await
.unwrap();
let headers = res.headers();
assert_eq!(
headers
.get("x-amz-server-side-encryption")
.expect("object is not encrypted"),
expected_encryption
);
store.delete(location).await.unwrap();
}
}
/// See CONTRIBUTING.md for the MinIO setup for this test.
#[tokio::test]
async fn test_s3_ssec_encryption_with_minio() {
if std::env::var("TEST_S3_SSEC_ENCRYPTION").is_err() {
eprintln!("Skipping S3 SSE-C encryption test");
return;
}
eprintln!("Running S3 SSE-C encryption test");
let customer_key = "1234567890abcdef1234567890abcdef";
let expected_md5 = "JMwgiexXqwuPqIPjYFmIZQ==";
let store = AmazonS3Builder::from_env()
.with_ssec_encryption(BASE64_STANDARD.encode(customer_key))
.with_client_options(ClientOptions::default().with_allow_invalid_certificates(true))
.build()
.unwrap();
let data = PutPayload::from(vec![3u8; 1024]);
let locations = [
Path::from("test-encryption-1"),
Path::from("test-encryption-2"),
Path::from("test-encryption-3"),
];
// Test put with sse-c.
store.put(&locations[0], data.clone()).await.unwrap();
// Test copy with sse-c.
store.copy(&locations[0], &locations[1]).await.unwrap();
// Test multipart upload with sse-c.
let mut upload = store.put_multipart(&locations[2]).await.unwrap();
upload.put_part(data.clone()).await.unwrap();
upload.complete().await.unwrap();
// Test get with sse-c.
for location in &locations {
let res = store
.client
.get_request(location, GetOptions::default())
.await
.unwrap();
let headers = res.headers();
assert_eq!(
headers
.get("x-amz-server-side-encryption-customer-algorithm")
.expect("object is not encrypted with SSE-C"),
"AES256"
);
assert_eq!(
headers
.get("x-amz-server-side-encryption-customer-key-MD5")
.expect("object is not encrypted with SSE-C"),
expected_md5
);
store.delete(location).await.unwrap();
}
}
/// Integration test that ensures I/O is done on an alternate threadpool
/// when using the `SpawnedReqwestConnector`.
#[test]
fn s3_alternate_threadpool_spawned_request_connector() {
maybe_skip_integration!();
let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>();
// Runtime with I/O enabled
let io_runtime = tokio::runtime::Builder::new_current_thread()
.enable_all() // <-- turns on IO
.build()
.unwrap();
// Runtime without I/O enabled
let non_io_runtime = tokio::runtime::Builder::new_current_thread()
// note: no call to enable_all
.build()
.unwrap();
// run the io runtime in a different thread
let io_handle = io_runtime.handle().clone();
let thread_handle = std::thread::spawn(move || {
io_runtime.block_on(async move {
shutdown_rx.await.unwrap();
});
});
let store = AmazonS3Builder::from_env()
// use different bucket to avoid collisions with other tests
.with_bucket_name("test-bucket-for-spawn")
.with_http_connector(SpawnedReqwestConnector::new(io_handle))
.build()
.unwrap();
// run a request on the non io runtime -- will fail if the connector
// does not spawn the request to the io runtime
non_io_runtime
.block_on(async move {
let path = Path::from("alternate_threadpool/test.txt");
store.delete(&path).await.ok(); // remove the file if it exists from prior runs
store.put(&path, "foo".into()).await?;
let res = store.get(&path).await?.bytes().await?;
assert_eq!(res.as_ref(), b"foo");
store.delete(&path).await?; // cleanup
Ok(()) as Result<()>
})
.expect("failed to run request on non io runtime");
// shutdown the io runtime and thread
shutdown_tx.send(()).ok();
thread_handle.join().expect("runtime thread panicked");
}
}
backend/vendor/object_store/src/aws/precondition.rs
Vendored
+278
View File
@@ -0,0 +1,278 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::aws::dynamo::DynamoCommit;
use crate::config::Parse;
use itertools::Itertools;
/// Configure how to provide [`ObjectStore::copy_if_not_exists`] for [`AmazonS3`].
///
/// [`ObjectStore::copy_if_not_exists`]: crate::ObjectStore::copy_if_not_exists
/// [`AmazonS3`]: super::AmazonS3
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum S3CopyIfNotExists {
/// Some S3-compatible stores, such as Cloudflare R2, support copy if not exists
/// semantics through custom headers.
///
/// If set, [`ObjectStore::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.
///
/// Encoded as `header:<HEADER_NAME>:<HEADER_VALUE>` ignoring whitespace
///
/// For example `header: cf-copy-destination-if-none-match: *`, would set
/// the header `cf-copy-destination-if-none-match` to `*`
///
/// [`ObjectStore::copy_if_not_exists`]: crate::ObjectStore::copy_if_not_exists
Header(String, String),
/// The same as [`S3CopyIfNotExists::Header`] but allows custom status code checking, for object stores that return values
/// other than 412.
///
/// Encoded as `header-with-status:<HEADER_NAME>:<HEADER_VALUE>:<STATUS>` ignoring whitespace
HeaderWithStatus(String, String, reqwest::StatusCode),
/// 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. See [the module
/// docs](super#multipart-uploads) for details.
///
/// Encoded as `multipart` ignoring whitespace.
Multipart,
/// The name of a DynamoDB table to use for coordination
///
/// Encoded as either `dynamo:<TABLE_NAME>` or `dynamo:<TABLE_NAME>:<TIMEOUT_MILLIS>`
/// ignoring whitespace. The default timeout is used if not specified
///
/// See [`DynamoCommit`] for more information
///
/// This will use the same region, credentials and endpoint as configured for S3
Dynamo(DynamoCommit),
}
impl std::fmt::Display for S3CopyIfNotExists {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Header(k, v) => write!(f, "header: {}: {}", k, v),
Self::HeaderWithStatus(k, v, code) => {
write!(f, "header-with-status: {k}: {v}: {}", code.as_u16())
}
Self::Multipart => f.write_str("multipart"),
Self::Dynamo(lock) => write!(f, "dynamo: {}", lock.table_name()),
}
}
}
impl S3CopyIfNotExists {
fn from_str(s: &str) -> Option<Self> {
if s.trim() == "multipart" {
return Some(Self::Multipart);
};
let (variant, value) = s.split_once(':')?;
match variant.trim() {
"header" => {
let (k, v) = value.split_once(':')?;
Some(Self::Header(k.trim().to_string(), v.trim().to_string()))
}
"header-with-status" => {
let (k, v, status) = value.split(':').collect_tuple()?;
let code = status.trim().parse().ok()?;
Some(Self::HeaderWithStatus(
k.trim().to_string(),
v.trim().to_string(),
code,
))
}
"dynamo" => Some(Self::Dynamo(DynamoCommit::from_str(value)?)),
_ => None,
}
}
}
impl Parse for S3CopyIfNotExists {
fn parse(v: &str) -> crate::Result<Self> {
Self::from_str(v).ok_or_else(|| crate::Error::Generic {
store: "Config",
source: format!("Failed to parse \"{v}\" as S3CopyIfNotExists").into(),
})
}
}
/// Configure how to provide conditional put support for [`AmazonS3`].
///
/// [`AmazonS3`]: super::AmazonS3
#[derive(Debug, Clone, Eq, PartialEq, Default)]
#[allow(missing_copy_implementations)]
#[non_exhaustive]
pub enum S3ConditionalPut {
/// Some S3-compatible stores, such as Cloudflare R2 and minio support conditional
/// put using the standard [HTTP precondition] headers If-Match and If-None-Match
///
/// Encoded as `etag` ignoring whitespace
///
/// [HTTP precondition]: https://datatracker.ietf.org/doc/html/rfc9110#name-preconditions
#[default]
ETagMatch,
/// The name of a DynamoDB table to use for coordination
///
/// Encoded as either `dynamo:<TABLE_NAME>` or `dynamo:<TABLE_NAME>:<TIMEOUT_MILLIS>`
/// ignoring whitespace. The default timeout is used if not specified
///
/// See [`DynamoCommit`] for more information
///
/// This will use the same region, credentials and endpoint as configured for S3
Dynamo(DynamoCommit),
/// Disable `conditional put`
Disabled,
}
impl std::fmt::Display for S3ConditionalPut {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::ETagMatch => write!(f, "etag"),
Self::Dynamo(lock) => write!(f, "dynamo: {}", lock.table_name()),
Self::Disabled => write!(f, "disabled"),
}
}
}
impl S3ConditionalPut {
fn from_str(s: &str) -> Option<Self> {
match s.trim() {
"etag" => Some(Self::ETagMatch),
"disabled" => Some(Self::Disabled),
trimmed => match trimmed.split_once(':')? {
("dynamo", s) => Some(Self::Dynamo(DynamoCommit::from_str(s)?)),
_ => None,
},
}
}
}
impl Parse for S3ConditionalPut {
fn parse(v: &str) -> crate::Result<Self> {
Self::from_str(v).ok_or_else(|| crate::Error::Generic {
store: "Config",
source: format!("Failed to parse \"{v}\" as S3PutConditional").into(),
})
}
}
#[cfg(test)]
mod tests {
use super::S3CopyIfNotExists;
use crate::aws::{DynamoCommit, S3ConditionalPut};
#[test]
fn parse_s3_copy_if_not_exists_header() {
let input = "header: cf-copy-destination-if-none-match: *";
let expected = Some(S3CopyIfNotExists::Header(
"cf-copy-destination-if-none-match".to_owned(),
"*".to_owned(),
));
assert_eq!(expected, S3CopyIfNotExists::from_str(input));
}
#[test]
fn parse_s3_copy_if_not_exists_header_with_status() {
let input = "header-with-status:key:value:403";
let expected = Some(S3CopyIfNotExists::HeaderWithStatus(
"key".to_owned(),
"value".to_owned(),
reqwest::StatusCode::FORBIDDEN,
));
assert_eq!(expected, S3CopyIfNotExists::from_str(input));
}
#[test]
fn parse_s3_copy_if_not_exists_dynamo() {
let input = "dynamo: table:100";
let expected = Some(S3CopyIfNotExists::Dynamo(
DynamoCommit::new("table".into()).with_timeout(100),
));
assert_eq!(expected, S3CopyIfNotExists::from_str(input));
}
#[test]
fn parse_s3_condition_put_dynamo() {
let input = "dynamo: table:1300";
let expected = Some(S3ConditionalPut::Dynamo(
DynamoCommit::new("table".into()).with_timeout(1300),
));
assert_eq!(expected, S3ConditionalPut::from_str(input));
}
#[test]
fn parse_s3_copy_if_not_exists_header_whitespace_invariant() {
let expected = Some(S3CopyIfNotExists::Header(
"cf-copy-destination-if-none-match".to_owned(),
"*".to_owned(),
));
const INPUTS: &[&str] = &[
"header:cf-copy-destination-if-none-match:*",
"header: cf-copy-destination-if-none-match:*",
"header: cf-copy-destination-if-none-match: *",
"header : cf-copy-destination-if-none-match: *",
"header : cf-copy-destination-if-none-match : *",
"header : cf-copy-destination-if-none-match : * ",
];
for input in INPUTS {
assert_eq!(expected, S3CopyIfNotExists::from_str(input));
}
}
#[test]
fn parse_s3_copy_if_not_exists_header_with_status_whitespace_invariant() {
let expected = Some(S3CopyIfNotExists::HeaderWithStatus(
"key".to_owned(),
"value".to_owned(),
reqwest::StatusCode::FORBIDDEN,
));
const INPUTS: &[&str] = &[
"header-with-status:key:value:403",
"header-with-status: key:value:403",
"header-with-status: key: value:403",
"header-with-status: key: value: 403",
"header-with-status : key: value: 403",
"header-with-status : key : value: 403",
"header-with-status : key : value : 403",
"header-with-status : key : value : 403 ",
];
for input in INPUTS {
assert_eq!(expected, S3CopyIfNotExists::from_str(input));
}
}
}
backend/vendor/object_store/src/aws/resolve.rs
Vendored
+89
View File
@@ -0,0 +1,89 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::aws::STORE;
use crate::{ClientOptions, Result};
/// A specialized `Error` for object store-related errors
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("Bucket '{}' not found", bucket)]
BucketNotFound { bucket: String },
#[error("Failed to resolve region for bucket '{}'", bucket)]
ResolveRegion {
bucket: String,
source: reqwest::Error,
},
#[error("Failed to parse the region for bucket '{}'", bucket)]
RegionParse { bucket: String },
}
impl From<Error> for crate::Error {
fn from(source: Error) -> Self {
Self::Generic {
store: STORE,
source: Box::new(source),
}
}
}
/// Get the bucket region using the [HeadBucket API]. This will fail if the bucket does not exist.
///
/// [HeadBucket API]: https://docs.aws.amazon.com/AmazonS3/latest/API/API_HeadBucket.html
pub async fn resolve_bucket_region(bucket: &str, client_options: &ClientOptions) -> Result<String> {
use reqwest::StatusCode;
let endpoint = format!("https://{}.s3.amazonaws.com", bucket);
let client = client_options.client()?;
let response = client.head(&endpoint).send().await.map_err(|source| {
let bucket = bucket.into();
Error::ResolveRegion { bucket, source }
})?;
if response.status() == StatusCode::NOT_FOUND {
let bucket = bucket.into();
return Err(Error::BucketNotFound { bucket }.into());
}
let region = response
.headers()
.get("x-amz-bucket-region")
.and_then(|x| x.to_str().ok())
.ok_or_else(|| Error::RegionParse {
bucket: bucket.into(),
})?;
Ok(region.to_string())
}
#[cfg(test)]
mod tests {
use super::*;
#[tokio::test]
async fn test_bucket_does_not_exist() {
let bucket = "please-dont-exist";
let result = resolve_bucket_region(bucket, &ClientOptions::new()).await;
assert!(result.is_err());
}
}
backend/vendor/object_store/src/azure/builder.rs
Vendored
+1242
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/azure/client.rs
Vendored
+1536
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/azure/credential.rs
Vendored
+1220
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/azure/mod.rs
Vendored
+405
View File
@@ -0,0 +1,405 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An object store implementation for Azure blob storage
//!
//! ## Streaming uploads
//!
//! [ObjectStore::put_multipart] will upload data in blocks and write a blob from those blocks.
//!
//! Unused blocks will automatically be dropped after 7 days.
use crate::{
multipart::{MultipartStore, PartId},
path::Path,
signer::Signer,
GetOptions, GetResult, ListPage, ListResult, MultipartId, MultipartUpload, ObjectMeta, ObjectStore,
PutMultipartOpts, PutOptions, PutPayload, PutResult, Result, UploadPart,
};
use async_trait::async_trait;
use futures::stream::{BoxStream, StreamExt, TryStreamExt};
use reqwest::Method;
use std::fmt::Debug;
use std::sync::Arc;
use std::time::Duration;
use url::Url;
use crate::client::get::GetClientExt;
use crate::client::list::ListClientExt;
use crate::client::CredentialProvider;
pub use credential::{authority_hosts, AzureAccessKey, AzureAuthorizer};
mod builder;
mod client;
mod credential;
/// [`CredentialProvider`] for [`MicrosoftAzure`]
pub type AzureCredentialProvider = Arc<dyn CredentialProvider<Credential = AzureCredential>>;
use crate::azure::client::AzureClient;
use crate::client::parts::Parts;
pub use builder::{AzureConfigKey, MicrosoftAzureBuilder};
pub use credential::AzureCredential;
const STORE: &str = "MicrosoftAzure";
/// Interface for [Microsoft Azure Blob Storage](https://azure.microsoft.com/en-us/services/storage/blobs/).
#[derive(Debug)]
pub struct MicrosoftAzure {
client: Arc<AzureClient>,
}
impl MicrosoftAzure {
/// Returns the [`AzureCredentialProvider`] used by [`MicrosoftAzure`]
pub fn credentials(&self) -> &AzureCredentialProvider {
&self.client.config().credentials
}
/// Create a full URL to the resource specified by `path` with this instance's configuration.
fn path_url(&self, path: &Path) -> Url {
self.client.config().path_url(path)
}
}
impl std::fmt::Display for MicrosoftAzure {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"MicrosoftAzure {{ account: {}, container: {} }}",
self.client.config().account,
self.client.config().container
)
}
}
#[async_trait]
impl ObjectStore for MicrosoftAzure {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
self.client.put_blob(location, payload, opts).await
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
Ok(Box::new(AzureMultiPartUpload {
part_idx: 0,
opts,
state: Arc::new(UploadState {
client: Arc::clone(&self.client),
location: location.clone(),
parts: Default::default(),
}),
}))
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
self.client.get_opts(location, options).await
}
async fn delete(&self, location: &Path) -> Result<()> {
self.client.delete_request(location, &()).await
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
self.client.list(prefix)
}
fn delete_stream<'a>(
&'a self,
locations: BoxStream<'a, Result<Path>>,
) -> BoxStream<'a, Result<Path>> {
locations
.try_chunks(256)
.map(move |locations| async {
// Early return the error. We ignore the paths that have already been
// collected into the chunk.
let locations = locations.map_err(|e| e.1)?;
self.client
.bulk_delete_request(locations)
.await
.map(futures::stream::iter)
})
.buffered(20)
.try_flatten()
.boxed()
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
self.client.list_with_delimiter(prefix).await
}
async fn list_delimited_page(
&self,
prefix: Option<&Path>,
token: Option<&str>,
max_keys: Option<usize>,
) -> Result<ListPage> {
self.client.list_delimited_page(prefix, token, max_keys).await
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
self.client.copy_request(from, to, true).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
self.client.copy_request(from, to, false).await
}
}
#[async_trait]
impl Signer for MicrosoftAzure {
/// Create a URL containing the relevant [Service SAS] query parameters that authorize a request
/// via `method` to the resource at `path` valid for the duration specified in `expires_in`.
///
/// [Service SAS]: https://learn.microsoft.com/en-us/rest/api/storageservices/create-service-sas
///
/// # Example
///
/// This example returns a URL that will enable a user to upload a file to
/// "some-folder/some-file.txt" in the next hour.
///
/// ```
/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
/// # use object_store::{azure::MicrosoftAzureBuilder, path::Path, signer::Signer};
/// # use reqwest::Method;
/// # use std::time::Duration;
/// #
/// let azure = MicrosoftAzureBuilder::new()
/// .with_account("my-account")
/// .with_access_key("my-access-key")
/// .with_container_name("my-container")
/// .build()?;
///
/// let url = azure.signed_url(
/// Method::PUT,
/// &Path::from("some-folder/some-file.txt"),
/// Duration::from_secs(60 * 60)
/// ).await?;
/// # Ok(())
/// # }
/// ```
async fn signed_url(&self, method: Method, path: &Path, expires_in: Duration) -> Result<Url> {
let mut url = self.path_url(path);
let signer = self.client.signer(expires_in).await?;
signer.sign(&method, &mut url)?;
Ok(url)
}
async fn signed_urls(
&self,
method: Method,
paths: &[Path],
expires_in: Duration,
) -> Result<Vec<Url>> {
let mut urls = Vec::with_capacity(paths.len());
let signer = self.client.signer(expires_in).await?;
for path in paths {
let mut url = self.path_url(path);
signer.sign(&method, &mut url)?;
urls.push(url);
}
Ok(urls)
}
}
/// Relevant docs: <https://azure.github.io/Storage/docs/application-and-user-data/basics/azure-blob-storage-upload-apis/>
/// In Azure Blob Store, parts are "blocks"
/// put_multipart_part -> PUT block
/// complete -> PUT block list
/// abort -> No equivalent; blocks are simply dropped after 7 days
#[derive(Debug)]
struct AzureMultiPartUpload {
part_idx: usize,
state: Arc<UploadState>,
opts: PutMultipartOpts,
}
#[derive(Debug)]
struct UploadState {
location: Path,
parts: Parts,
client: Arc<AzureClient>,
}
#[async_trait]
impl MultipartUpload for AzureMultiPartUpload {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
let idx = self.part_idx;
self.part_idx += 1;
let state = Arc::clone(&self.state);
Box::pin(async move {
let part = state.client.put_block(&state.location, idx, data).await?;
state.parts.put(idx, part);
Ok(())
})
}
async fn complete(&mut self) -> Result<PutResult> {
let parts = self.state.parts.finish(self.part_idx)?;
self.state
.client
.put_block_list(&self.state.location, parts, std::mem::take(&mut self.opts))
.await
}
async fn abort(&mut self) -> Result<()> {
// Nothing to do
Ok(())
}
}
#[async_trait]
impl MultipartStore for MicrosoftAzure {
async fn create_multipart(&self, _: &Path) -> Result<MultipartId> {
Ok(String::new())
}
async fn put_part(
&self,
path: &Path,
_: &MultipartId,
part_idx: usize,
data: PutPayload,
) -> Result<PartId> {
self.client.put_block(path, part_idx, data).await
}
async fn complete_multipart(
&self,
path: &Path,
_: &MultipartId,
parts: Vec<PartId>,
) -> Result<PutResult> {
self.client
.put_block_list(path, parts, Default::default())
.await
}
async fn abort_multipart(&self, _: &Path, _: &MultipartId) -> Result<()> {
// There is no way to drop blocks that have been uploaded. Instead, they simply
// expire in 7 days.
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::integration::*;
use crate::tests::*;
use bytes::Bytes;
#[tokio::test]
async fn azure_blob_test() {
maybe_skip_integration!();
let integration = MicrosoftAzureBuilder::from_env().build().unwrap();
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
stream_get(&integration).await;
put_opts(&integration, true).await;
multipart(&integration, &integration).await;
multipart_race_condition(&integration, false).await;
multipart_out_of_order(&integration).await;
signing(&integration).await;
let validate = !integration.client.config().disable_tagging;
tagging(
Arc::new(MicrosoftAzure {
client: Arc::clone(&integration.client),
}),
validate,
|p| {
let client = Arc::clone(&integration.client);
async move { client.get_blob_tagging(&p).await }
},
)
.await;
// Azurite doesn't support attributes properly
if !integration.client.config().is_emulator {
put_get_attributes(&integration).await;
}
}
#[ignore = "Used for manual testing against a real storage account."]
#[tokio::test]
async fn test_user_delegation_key() {
let account = std::env::var("AZURE_ACCOUNT_NAME").unwrap();
let container = std::env::var("AZURE_CONTAINER_NAME").unwrap();
let client_id = std::env::var("AZURE_CLIENT_ID").unwrap();
let client_secret = std::env::var("AZURE_CLIENT_SECRET").unwrap();
let tenant_id = std::env::var("AZURE_TENANT_ID").unwrap();
let integration = MicrosoftAzureBuilder::new()
.with_account(account)
.with_container_name(container)
.with_client_id(client_id)
.with_client_secret(client_secret)
.with_tenant_id(&tenant_id)
.build()
.unwrap();
let data = Bytes::from("hello world");
let path = Path::from("file.txt");
integration.put(&path, data.clone().into()).await.unwrap();
let signed = integration
.signed_url(Method::GET, &path, Duration::from_secs(60))
.await
.unwrap();
let resp = reqwest::get(signed).await.unwrap();
let loaded = resp.bytes().await.unwrap();
assert_eq!(data, loaded);
}
#[test]
fn azure_test_config_get_value() {
let azure_client_id = "object_store:fake_access_key_id".to_string();
let azure_storage_account_name = "object_store:fake_secret_key".to_string();
let azure_storage_token = "object_store:fake_default_region".to_string();
let builder = MicrosoftAzureBuilder::new()
.with_config(AzureConfigKey::ClientId, &azure_client_id)
.with_config(AzureConfigKey::AccountName, &azure_storage_account_name)
.with_config(AzureConfigKey::Token, &azure_storage_token);
assert_eq!(
builder.get_config_value(&AzureConfigKey::ClientId).unwrap(),
azure_client_id
);
assert_eq!(
builder
.get_config_value(&AzureConfigKey::AccountName)
.unwrap(),
azure_storage_account_name
);
assert_eq!(
builder.get_config_value(&AzureConfigKey::Token).unwrap(),
azure_storage_token
);
}
}
backend/vendor/object_store/src/buffered.rs
Vendored
+679
View File
@@ -0,0 +1,679 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Utilities for performing tokio-style buffered IO
use crate::path::Path;
use crate::{
Attributes, ObjectMeta, ObjectStore, PutMultipartOpts, PutOptions, PutPayloadMut, TagSet,
WriteMultipart,
};
use bytes::Bytes;
use futures::future::{BoxFuture, FutureExt};
use futures::ready;
use std::cmp::Ordering;
use std::io::{Error, ErrorKind, SeekFrom};
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::io::{AsyncBufRead, AsyncRead, AsyncSeek, AsyncWrite, ReadBuf};
/// The default buffer size used by [`BufReader`]
pub const DEFAULT_BUFFER_SIZE: usize = 1024 * 1024;
/// An async-buffered reader compatible with the tokio IO traits
///
/// Internally this maintains a buffer of the requested size, and uses [`ObjectStore::get_range`]
/// to populate its internal buffer once depleted. This buffer is cleared on seek.
///
/// Whilst simple, this interface will typically be outperformed by the native [`ObjectStore`]
/// methods that better map to the network APIs. This is because most object stores have
/// very [high first-byte latencies], on the order of 100-200ms, and so avoiding unnecessary
/// round-trips is critical to throughput.
///
/// Systems looking to sequentially scan a file should instead consider using [`ObjectStore::get`],
/// or [`ObjectStore::get_opts`], or [`ObjectStore::get_range`] to read a particular range.
///
/// Systems looking to read multiple ranges of a file should instead consider using
/// [`ObjectStore::get_ranges`], which will optimise the vectored IO.
///
/// [high first-byte latencies]: https://docs.aws.amazon.com/AmazonS3/latest/userguide/optimizing-performance.html
pub struct BufReader {
/// The object store to fetch data from
store: Arc<dyn ObjectStore>,
/// The size of the object
size: u64,
/// The path to the object
path: Path,
/// The current position in the object
cursor: u64,
/// The number of bytes to read in a single request
capacity: usize,
/// The buffered data if any
buffer: Buffer,
}
impl std::fmt::Debug for BufReader {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BufReader")
.field("path", &self.path)
.field("size", &self.size)
.field("capacity", &self.capacity)
.finish()
}
}
enum Buffer {
Empty,
Pending(BoxFuture<'static, std::io::Result<Bytes>>),
Ready(Bytes),
}
impl BufReader {
/// Create a new [`BufReader`] from the provided [`ObjectMeta`] and [`ObjectStore`]
pub fn new(store: Arc<dyn ObjectStore>, meta: &ObjectMeta) -> Self {
Self::with_capacity(store, meta, DEFAULT_BUFFER_SIZE)
}
/// Create a new [`BufReader`] from the provided [`ObjectMeta`], [`ObjectStore`], and `capacity`
pub fn with_capacity(store: Arc<dyn ObjectStore>, meta: &ObjectMeta, capacity: usize) -> Self {
Self {
path: meta.location.clone(),
size: meta.size as _,
store,
capacity,
cursor: 0,
buffer: Buffer::Empty,
}
}
fn poll_fill_buf_impl(
&mut self,
cx: &mut Context<'_>,
amnt: usize,
) -> Poll<std::io::Result<&[u8]>> {
let buf = &mut self.buffer;
loop {
match buf {
Buffer::Empty => {
let store = Arc::clone(&self.store);
let path = self.path.clone();
let start = self.cursor.min(self.size) as _;
let end = self.cursor.saturating_add(amnt as u64).min(self.size) as _;
if start == end {
return Poll::Ready(Ok(&[]));
}
*buf = Buffer::Pending(Box::pin(async move {
Ok(store.get_range(&path, start..end).await?)
}))
}
Buffer::Pending(fut) => match ready!(fut.poll_unpin(cx)) {
Ok(b) => *buf = Buffer::Ready(b),
Err(e) => return Poll::Ready(Err(e)),
},
Buffer::Ready(r) => return Poll::Ready(Ok(r)),
}
}
}
}
impl AsyncSeek for BufReader {
fn start_seek(mut self: Pin<&mut Self>, position: SeekFrom) -> std::io::Result<()> {
self.cursor = match position {
SeekFrom::Start(offset) => offset,
SeekFrom::End(offset) => checked_add_signed(self.size, offset).ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!(
"Seeking {offset} from end of {} byte file would result in overflow",
self.size
),
)
})?,
SeekFrom::Current(offset) => {
checked_add_signed(self.cursor, offset).ok_or_else(|| {
Error::new(
ErrorKind::InvalidInput,
format!(
"Seeking {offset} from current offset of {} would result in overflow",
self.cursor
),
)
})?
}
};
self.buffer = Buffer::Empty;
Ok(())
}
fn poll_complete(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<std::io::Result<u64>> {
Poll::Ready(Ok(self.cursor))
}
}
impl AsyncRead for BufReader {
fn poll_read(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
out: &mut ReadBuf<'_>,
) -> Poll<std::io::Result<()>> {
// Read the maximum of the internal buffer and `out`
let to_read = out.remaining().max(self.capacity);
let r = match ready!(self.poll_fill_buf_impl(cx, to_read)) {
Ok(buf) => {
let to_consume = out.remaining().min(buf.len());
out.put_slice(&buf[..to_consume]);
self.consume(to_consume);
Ok(())
}
Err(e) => Err(e),
};
Poll::Ready(r)
}
}
impl AsyncBufRead for BufReader {
fn poll_fill_buf(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<&[u8]>> {
let capacity = self.capacity;
self.get_mut().poll_fill_buf_impl(cx, capacity)
}
fn consume(mut self: Pin<&mut Self>, amt: usize) {
match &mut self.buffer {
Buffer::Empty => assert_eq!(amt, 0, "cannot consume from empty buffer"),
Buffer::Ready(b) => match b.len().cmp(&amt) {
Ordering::Less => panic!("{amt} exceeds buffer sized of {}", b.len()),
Ordering::Greater => *b = b.slice(amt..),
Ordering::Equal => self.buffer = Buffer::Empty,
},
Buffer::Pending(_) => panic!("cannot consume from pending buffer"),
}
self.cursor += amt as u64;
}
}
/// An async buffered writer compatible with the tokio IO traits
///
/// This writer adaptively uses [`ObjectStore::put`] or
/// [`ObjectStore::put_multipart`] depending on the amount of data that has
/// been written.
///
/// Up to `capacity` bytes will be buffered in memory, and flushed on shutdown
/// using [`ObjectStore::put`]. If `capacity` is exceeded, data will instead be
/// streamed using [`ObjectStore::put_multipart`]
pub struct BufWriter {
capacity: usize,
max_concurrency: usize,
attributes: Option<Attributes>,
tags: Option<TagSet>,
extensions: Option<::http::Extensions>,
state: BufWriterState,
store: Arc<dyn ObjectStore>,
}
impl std::fmt::Debug for BufWriter {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("BufWriter")
.field("capacity", &self.capacity)
.finish()
}
}
enum BufWriterState {
/// Buffer up to capacity bytes
Buffer(Path, PutPayloadMut),
/// [`ObjectStore::put_multipart`]
Prepare(BoxFuture<'static, crate::Result<WriteMultipart>>),
/// Write to a multipart upload
Write(Option<WriteMultipart>),
/// [`ObjectStore::put`]
Flush(BoxFuture<'static, crate::Result<()>>),
}
impl BufWriter {
/// Create a new [`BufWriter`] from the provided [`ObjectStore`] and [`Path`]
pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
Self::with_capacity(store, path, 10 * 1024 * 1024)
}
/// Create a new [`BufWriter`] from the provided [`ObjectStore`], [`Path`] and `capacity`
pub fn with_capacity(store: Arc<dyn ObjectStore>, path: Path, capacity: usize) -> Self {
Self {
capacity,
store,
max_concurrency: 8,
attributes: None,
tags: None,
extensions: None,
state: BufWriterState::Buffer(path, PutPayloadMut::new()),
}
}
/// Override the maximum number of in-flight requests for this writer
///
/// Defaults to 8
pub fn with_max_concurrency(self, max_concurrency: usize) -> Self {
Self {
max_concurrency,
..self
}
}
/// Set the attributes of the uploaded object
pub fn with_attributes(self, attributes: Attributes) -> Self {
Self {
attributes: Some(attributes),
..self
}
}
/// Set the tags of the uploaded object
pub fn with_tags(self, tags: TagSet) -> Self {
Self {
tags: Some(tags),
..self
}
}
/// Set the extensions of the uploaded object
///
/// Implementation-specific extensions. Intended for use by [`ObjectStore`] implementations
/// that need to pass context-specific information (like tracing spans) via trait methods.
///
/// These extensions are ignored entirely by backends offered through this crate.
pub fn with_extensions(self, extensions: ::http::Extensions) -> Self {
Self {
extensions: Some(extensions),
..self
}
}
/// Write data to the writer in [`Bytes`].
///
/// Unlike [`AsyncWrite::poll_write`], `put` can write data without extra copying.
///
/// This API is recommended while the data source generates [`Bytes`].
pub async fn put(&mut self, bytes: Bytes) -> crate::Result<()> {
loop {
return match &mut self.state {
BufWriterState::Write(Some(write)) => {
write.wait_for_capacity(self.max_concurrency).await?;
write.put(bytes);
Ok(())
}
BufWriterState::Write(None) | BufWriterState::Flush(_) => {
panic!("Already shut down")
}
// NOTE
//
// This case should never happen in practice, but rust async API does
// make it possible for users to call `put` before `poll_write` returns `Ready`.
//
// We allow such usage by `await` the future and continue the loop.
BufWriterState::Prepare(f) => {
self.state = BufWriterState::Write(f.await?.into());
continue;
}
BufWriterState::Buffer(path, b) => {
if b.content_length().saturating_add(bytes.len()) < self.capacity {
b.push(bytes);
Ok(())
} else {
let buffer = std::mem::take(b);
let path = std::mem::take(path);
let opts = PutMultipartOpts {
attributes: self.attributes.take().unwrap_or_default(),
tags: self.tags.take().unwrap_or_default(),
extensions: self.extensions.take().unwrap_or_default(),
};
let upload = self.store.put_multipart_opts(&path, opts).await?;
let mut chunked =
WriteMultipart::new_with_chunk_size(upload, self.capacity);
for chunk in buffer.freeze() {
chunked.put(chunk);
}
chunked.put(bytes);
self.state = BufWriterState::Write(Some(chunked));
Ok(())
}
}
};
}
}
/// Abort this writer, cleaning up any partially uploaded state
///
/// # Panic
///
/// Panics if this writer has already been shutdown or aborted
pub async fn abort(&mut self) -> crate::Result<()> {
match &mut self.state {
BufWriterState::Buffer(_, _) | BufWriterState::Prepare(_) => Ok(()),
BufWriterState::Flush(_) => panic!("Already shut down"),
BufWriterState::Write(x) => x.take().unwrap().abort().await,
}
}
}
impl AsyncWrite for BufWriter {
fn poll_write(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
buf: &[u8],
) -> Poll<Result<usize, Error>> {
let cap = self.capacity;
let max_concurrency = self.max_concurrency;
loop {
return match &mut self.state {
BufWriterState::Write(Some(write)) => {
ready!(write.poll_for_capacity(cx, max_concurrency))?;
write.write(buf);
Poll::Ready(Ok(buf.len()))
}
BufWriterState::Write(None) | BufWriterState::Flush(_) => {
panic!("Already shut down")
}
BufWriterState::Prepare(f) => {
self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
continue;
}
BufWriterState::Buffer(path, b) => {
if b.content_length().saturating_add(buf.len()) >= cap {
let buffer = std::mem::take(b);
let path = std::mem::take(path);
let opts = PutMultipartOpts {
attributes: self.attributes.take().unwrap_or_default(),
tags: self.tags.take().unwrap_or_default(),
extensions: self.extensions.take().unwrap_or_default(),
};
let store = Arc::clone(&self.store);
self.state = BufWriterState::Prepare(Box::pin(async move {
let upload = store.put_multipart_opts(&path, opts).await?;
let mut chunked = WriteMultipart::new_with_chunk_size(upload, cap);
for chunk in buffer.freeze() {
chunked.put(chunk);
}
Ok(chunked)
}));
continue;
}
b.extend_from_slice(buf);
Poll::Ready(Ok(buf.len()))
}
};
}
}
fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
loop {
return match &mut self.state {
BufWriterState::Write(_) | BufWriterState::Buffer(_, _) => Poll::Ready(Ok(())),
BufWriterState::Flush(_) => panic!("Already shut down"),
BufWriterState::Prepare(f) => {
self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
continue;
}
};
}
}
fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
loop {
match &mut self.state {
BufWriterState::Prepare(f) => {
self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
}
BufWriterState::Buffer(p, b) => {
let buf = std::mem::take(b);
let path = std::mem::take(p);
let opts = PutOptions {
attributes: self.attributes.take().unwrap_or_default(),
tags: self.tags.take().unwrap_or_default(),
..Default::default()
};
let store = Arc::clone(&self.store);
self.state = BufWriterState::Flush(Box::pin(async move {
store.put_opts(&path, buf.into(), opts).await?;
Ok(())
}));
}
BufWriterState::Flush(f) => return f.poll_unpin(cx).map_err(std::io::Error::from),
BufWriterState::Write(x) => {
let upload = x.take().ok_or_else(|| {
std::io::Error::new(
ErrorKind::InvalidInput,
"Cannot shutdown a writer that has already been shut down",
)
})?;
self.state = BufWriterState::Flush(
async move {
upload.finish().await?;
Ok(())
}
.boxed(),
)
}
}
}
}
}
/// Port of standardised function as requires Rust 1.66
///
/// <https://github.com/rust-lang/rust/pull/87601/files#diff-b9390ee807a1dae3c3128dce36df56748ad8d23c6e361c0ebba4d744bf6efdb9R1533>
#[inline]
fn checked_add_signed(a: u64, rhs: i64) -> Option<u64> {
let (res, overflowed) = a.overflowing_add(rhs as _);
let overflow = overflowed ^ (rhs < 0);
(!overflow).then_some(res)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::memory::InMemory;
use crate::path::Path;
use crate::{Attribute, GetOptions};
use itertools::Itertools;
use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncSeekExt, AsyncWriteExt};
#[tokio::test]
async fn test_buf_reader() {
let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
let existent = Path::from("exists.txt");
const BYTES: usize = 4096;
let data: Bytes = b"12345678".iter().cycle().copied().take(BYTES).collect();
store.put(&existent, data.clone().into()).await.unwrap();
let meta = store.head(&existent).await.unwrap();
let mut reader = BufReader::new(Arc::clone(&store), &meta);
let mut out = Vec::with_capacity(BYTES);
let read = reader.read_to_end(&mut out).await.unwrap();
assert_eq!(read, BYTES);
assert_eq!(&out, &data);
let err = reader.seek(SeekFrom::Current(i64::MIN)).await.unwrap_err();
assert_eq!(
err.to_string(),
"Seeking -9223372036854775808 from current offset of 4096 would result in overflow"
);
reader.rewind().await.unwrap();
let err = reader.seek(SeekFrom::Current(-1)).await.unwrap_err();
assert_eq!(
err.to_string(),
"Seeking -1 from current offset of 0 would result in overflow"
);
// Seeking beyond the bounds of the file is permitted but should return no data
reader.seek(SeekFrom::Start(u64::MAX)).await.unwrap();
let buf = reader.fill_buf().await.unwrap();
assert!(buf.is_empty());
let err = reader.seek(SeekFrom::Current(1)).await.unwrap_err();
assert_eq!(
err.to_string(),
"Seeking 1 from current offset of 18446744073709551615 would result in overflow"
);
for capacity in [200, 1024, 4096, DEFAULT_BUFFER_SIZE] {
let store = Arc::clone(&store);
let mut reader = BufReader::with_capacity(store, &meta, capacity);
let mut bytes_read = 0;
loop {
let buf = reader.fill_buf().await.unwrap();
if buf.is_empty() {
assert_eq!(bytes_read, BYTES);
break;
}
assert!(buf.starts_with(b"12345678"));
bytes_read += 8;
reader.consume(8);
}
let mut buf = Vec::with_capacity(76);
reader.seek(SeekFrom::Current(-76)).await.unwrap();
reader.read_to_end(&mut buf).await.unwrap();
assert_eq!(&buf, &data[BYTES - 76..]);
reader.rewind().await.unwrap();
let buffer = reader.fill_buf().await.unwrap();
assert_eq!(buffer, &data[..capacity.min(BYTES)]);
reader.seek(SeekFrom::Start(325)).await.unwrap();
let buffer = reader.fill_buf().await.unwrap();
assert_eq!(buffer, &data[325..(325 + capacity).min(BYTES)]);
reader.seek(SeekFrom::End(0)).await.unwrap();
let buffer = reader.fill_buf().await.unwrap();
assert!(buffer.is_empty());
}
}
// Note: `BufWriter::with_tags` functionality is tested in `crate::tests::tagging`
#[tokio::test]
async fn test_buf_writer() {
let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
let path = Path::from("file.txt");
let attributes = Attributes::from_iter([
(Attribute::ContentType, "text/html"),
(Attribute::CacheControl, "max-age=604800"),
]);
// Test put
let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30)
.with_attributes(attributes.clone());
writer.write_all(&[0; 20]).await.unwrap();
writer.flush().await.unwrap();
writer.write_all(&[0; 5]).await.unwrap();
writer.shutdown().await.unwrap();
let response = store
.get_opts(
&path,
GetOptions {
head: true,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(response.meta.size, 25);
assert_eq!(response.attributes, attributes);
// Test multipart
let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30)
.with_attributes(attributes.clone());
writer.write_all(&[0; 20]).await.unwrap();
writer.flush().await.unwrap();
writer.write_all(&[0; 20]).await.unwrap();
writer.shutdown().await.unwrap();
let response = store
.get_opts(
&path,
GetOptions {
head: true,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(response.meta.size, 40);
assert_eq!(response.attributes, attributes);
}
#[tokio::test]
async fn test_buf_writer_with_put() {
let store = Arc::new(InMemory::new()) as Arc<dyn ObjectStore>;
let path = Path::from("file.txt");
// Test put
let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30);
writer
.put(Bytes::from((0..20).collect_vec()))
.await
.unwrap();
writer
.put(Bytes::from((20..25).collect_vec()))
.await
.unwrap();
writer.shutdown().await.unwrap();
let response = store
.get_opts(
&path,
GetOptions {
head: true,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(response.meta.size, 25);
assert_eq!(response.bytes().await.unwrap(), (0..25).collect_vec());
// Test multipart
let mut writer = BufWriter::with_capacity(Arc::clone(&store), path.clone(), 30);
writer
.put(Bytes::from((0..20).collect_vec()))
.await
.unwrap();
writer
.put(Bytes::from((20..40).collect_vec()))
.await
.unwrap();
writer.shutdown().await.unwrap();
let response = store
.get_opts(
&path,
GetOptions {
head: true,
..Default::default()
},
)
.await
.unwrap();
assert_eq!(response.meta.size, 40);
assert_eq!(response.bytes().await.unwrap(), (0..40).collect_vec());
}
}
backend/vendor/object_store/src/chunked.rs
Vendored
+236
View File
@@ -0,0 +1,236 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! A [`ChunkedStore`] that can be used to test streaming behaviour
use std::fmt::{Debug, Display, Formatter};
use std::ops::Range;
use std::sync::Arc;
use async_trait::async_trait;
use bytes::{BufMut, Bytes, BytesMut};
use futures::stream::BoxStream;
use futures::StreamExt;
use crate::path::Path;
use crate::{
GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
PutMultipartOpts, PutOptions, PutResult,
};
use crate::{PutPayload, Result};
/// Wraps a [`ObjectStore`] and makes its get response return chunks
/// in a controllable manner.
///
/// A `ChunkedStore` makes the memory consumption and performance of
/// the wrapped [`ObjectStore`] worse. It is intended for use within
/// tests, to control the chunks in the produced output streams. For
/// example, it is used to verify the delimiting logic in
/// newline_delimited_stream.
#[derive(Debug)]
pub struct ChunkedStore {
inner: Arc<dyn ObjectStore>,
chunk_size: usize, // chunks are in memory, so we use usize not u64
}
impl ChunkedStore {
/// Creates a new [`ChunkedStore`] with the specified chunk_size
pub fn new(inner: Arc<dyn ObjectStore>, chunk_size: usize) -> Self {
Self { inner, chunk_size }
}
}
impl Display for ChunkedStore {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
write!(f, "ChunkedStore({})", self.inner)
}
}
#[async_trait]
impl ObjectStore for ChunkedStore {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart(&self, location: &Path) -> Result<Box<dyn MultipartUpload>> {
self.inner.put_multipart(location).await
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
self.inner.put_multipart_opts(location, opts).await
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
let r = self.inner.get_opts(location, options).await?;
let stream = match r.payload {
#[cfg(all(feature = "fs", not(target_arch = "wasm32")))]
GetResultPayload::File(file, path) => {
crate::local::chunked_stream(file, path, r.range.clone(), self.chunk_size)
}
GetResultPayload::Stream(stream) => {
let buffer = BytesMut::new();
futures::stream::unfold(
(stream, buffer, false, self.chunk_size),
|(mut stream, mut buffer, mut exhausted, chunk_size)| async move {
// Keep accumulating bytes until we reach capacity as long as
// the stream can provide them:
if exhausted {
return None;
}
while buffer.len() < chunk_size {
match stream.next().await {
None => {
exhausted = true;
let slice = buffer.split_off(0).freeze();
return Some((
Ok(slice),
(stream, buffer, exhausted, chunk_size),
));
}
Some(Ok(bytes)) => {
buffer.put(bytes);
}
Some(Err(e)) => {
return Some((
Err(crate::Error::Generic {
store: "ChunkedStore",
source: Box::new(e),
}),
(stream, buffer, exhausted, chunk_size),
))
}
};
}
// Return the chunked values as the next value in the stream
let slice = buffer.split_to(chunk_size).freeze();
Some((Ok(slice), (stream, buffer, exhausted, chunk_size)))
},
)
.boxed()
}
};
Ok(GetResult {
payload: GetResultPayload::Stream(stream),
..r
})
}
async fn get_range(&self, location: &Path, range: Range<u64>) -> Result<Bytes> {
self.inner.get_range(location, range).await
}
async fn head(&self, location: &Path) -> Result<ObjectMeta> {
self.inner.head(location).await
}
async fn delete(&self, location: &Path) -> Result<()> {
self.inner.delete(location).await
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
self.inner.list(prefix)
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
self.inner.list_with_offset(prefix, offset)
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
self.inner.list_with_delimiter(prefix).await
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
self.inner.copy(from, to).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
self.inner.copy_if_not_exists(from, to).await
}
}
#[cfg(test)]
mod tests {
use futures::StreamExt;
#[cfg(feature = "fs")]
use crate::integration::*;
#[cfg(feature = "fs")]
use crate::local::LocalFileSystem;
use crate::memory::InMemory;
use crate::path::Path;
use super::*;
#[tokio::test]
async fn test_chunked_basic() {
let location = Path::parse("test").unwrap();
let store: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
store.put(&location, vec![0; 1001].into()).await.unwrap();
for chunk_size in [10, 20, 31] {
let store = ChunkedStore::new(Arc::clone(&store), chunk_size);
let mut s = match store.get(&location).await.unwrap().payload {
GetResultPayload::Stream(s) => s,
_ => unreachable!(),
};
let mut remaining = 1001;
while let Some(next) = s.next().await {
let size = next.unwrap().len() as u64;
let expected = remaining.min(chunk_size as u64);
assert_eq!(size, expected);
remaining -= expected;
}
assert_eq!(remaining, 0);
}
}
#[cfg(feature = "fs")]
#[tokio::test]
async fn test_chunked() {
let temporary = tempfile::tempdir().unwrap();
let integrations: &[Arc<dyn ObjectStore>] = &[
Arc::new(InMemory::new()),
Arc::new(LocalFileSystem::new_with_prefix(temporary.path()).unwrap()),
];
for integration in integrations {
let integration = ChunkedStore::new(Arc::clone(integration), 100);
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
stream_get(&integration).await;
}
}
}
backend/vendor/object_store/src/client/backoff.rs
Vendored
+157
View File
@@ -0,0 +1,157 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use rand::{prelude::*, rng};
use std::time::Duration;
/// Exponential backoff with decorrelated jitter algorithm
///
/// The first backoff will always be `init_backoff`.
///
/// Subsequent backoffs will pick a random value between `init_backoff` and
/// `base * previous` where `previous` is the duration of the previous backoff
///
/// See <https://aws.amazon.com/blogs/architecture/exponential-backoff-and-jitter/>
#[allow(missing_copy_implementations)]
#[derive(Debug, Clone)]
pub struct BackoffConfig {
/// The initial backoff duration
pub init_backoff: Duration,
/// The maximum backoff duration
pub max_backoff: Duration,
/// The multiplier to use for the next backoff duration
pub base: f64,
}
impl Default for BackoffConfig {
fn default() -> Self {
Self {
init_backoff: Duration::from_millis(100),
max_backoff: Duration::from_secs(15),
base: 2.,
}
}
}
/// [`Backoff`] can be created from a [`BackoffConfig`]
///
/// Consecutive calls to [`Backoff::next`] will return the next backoff interval
///
pub(crate) struct Backoff {
init_backoff: f64,
next_backoff_secs: f64,
max_backoff_secs: f64,
base: f64,
rng: Option<Box<dyn RngCore + Sync + Send>>,
}
impl std::fmt::Debug for Backoff {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Backoff")
.field("init_backoff", &self.init_backoff)
.field("next_backoff_secs", &self.next_backoff_secs)
.field("max_backoff_secs", &self.max_backoff_secs)
.field("base", &self.base)
.finish()
}
}
impl Backoff {
/// Create a new [`Backoff`] from the provided [`BackoffConfig`]
pub(crate) fn new(config: &BackoffConfig) -> Self {
Self::new_with_rng(config, None)
}
/// Creates a new `Backoff` with the optional `rng`
///
/// Used [`rand::rng()`] if no rng provided
pub(crate) fn new_with_rng(
config: &BackoffConfig,
rng: Option<Box<dyn RngCore + Sync + Send>>,
) -> Self {
let init_backoff = config.init_backoff.as_secs_f64();
Self {
init_backoff,
next_backoff_secs: init_backoff,
max_backoff_secs: config.max_backoff.as_secs_f64(),
base: config.base,
rng,
}
}
/// Returns the next backoff duration to wait for
pub(crate) fn next(&mut self) -> Duration {
let range = self.init_backoff..(self.next_backoff_secs * self.base);
let rand_backoff = match self.rng.as_mut() {
Some(rng) => rng.random_range(range),
None => rng().random_range(range),
};
let next_backoff = self.max_backoff_secs.min(rand_backoff);
Duration::from_secs_f64(std::mem::replace(&mut self.next_backoff_secs, next_backoff))
}
}
#[cfg(test)]
mod tests {
use super::*;
use rand::rngs::mock::StepRng;
#[test]
fn test_backoff() {
let init_backoff_secs = 1.;
let max_backoff_secs = 500.;
let base = 3.;
let config = BackoffConfig {
init_backoff: Duration::from_secs_f64(init_backoff_secs),
max_backoff: Duration::from_secs_f64(max_backoff_secs),
base,
};
let assert_fuzzy_eq = |a: f64, b: f64| assert!((b - a).abs() < 0.0001, "{a} != {b}");
// Create a static rng that takes the minimum of the range
let rng = Box::new(StepRng::new(0, 0));
let mut backoff = Backoff::new_with_rng(&config, Some(rng));
for _ in 0..20 {
assert_eq!(backoff.next().as_secs_f64(), init_backoff_secs);
}
// Create a static rng that takes the maximum of the range
let rng = Box::new(StepRng::new(u64::MAX, 0));
let mut backoff = Backoff::new_with_rng(&config, Some(rng));
for i in 0..20 {
let value = (base.powi(i) * init_backoff_secs).min(max_backoff_secs);
assert_fuzzy_eq(backoff.next().as_secs_f64(), value);
}
// Create a static rng that takes the mid point of the range
let rng = Box::new(StepRng::new(u64::MAX / 2, 0));
let mut backoff = Backoff::new_with_rng(&config, Some(rng));
let mut value = init_backoff_secs;
for _ in 0..20 {
assert_fuzzy_eq(backoff.next().as_secs_f64(), value);
value =
(init_backoff_secs + (value * base - init_backoff_secs) / 2.).min(max_backoff_secs);
}
}
}
backend/vendor/object_store/src/client/builder.rs
Vendored
+329
View File
@@ -0,0 +1,329 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::{HttpClient, HttpError, HttpErrorKind, HttpRequest, HttpRequestBody};
use http::header::{InvalidHeaderName, InvalidHeaderValue};
use http::uri::InvalidUri;
use http::{HeaderName, HeaderValue, Method, Uri};
#[derive(Debug, thiserror::Error)]
pub(crate) enum RequestBuilderError {
#[error("Invalid URI")]
InvalidUri(#[from] InvalidUri),
#[error("Invalid Header Value")]
InvalidHeaderValue(#[from] InvalidHeaderValue),
#[error("Invalid Header Name")]
InvalidHeaderName(#[from] InvalidHeaderName),
#[error("JSON serialization error")]
SerdeJson(#[from] serde_json::Error),
#[error("URL serialization error")]
SerdeUrl(#[from] serde_urlencoded::ser::Error),
}
impl From<RequestBuilderError> for HttpError {
fn from(value: RequestBuilderError) -> Self {
Self::new(HttpErrorKind::Request, value)
}
}
impl From<std::convert::Infallible> for RequestBuilderError {
fn from(value: std::convert::Infallible) -> Self {
match value {}
}
}
pub(crate) struct HttpRequestBuilder {
client: HttpClient,
request: Result<HttpRequest, RequestBuilderError>,
}
impl HttpRequestBuilder {
pub(crate) fn new(client: HttpClient) -> Self {
Self {
client,
request: Ok(HttpRequest::new(HttpRequestBody::empty())),
}
}
#[cfg(any(feature = "aws", feature = "azure"))]
pub(crate) fn from_parts(client: HttpClient, request: HttpRequest) -> Self {
Self {
client,
request: Ok(request),
}
}
pub(crate) fn method(mut self, method: Method) -> Self {
if let Ok(r) = &mut self.request {
*r.method_mut() = method;
}
self
}
pub(crate) fn uri<U>(mut self, url: U) -> Self
where
U: TryInto<Uri>,
U::Error: Into<RequestBuilderError>,
{
match (url.try_into(), &mut self.request) {
(Ok(uri), Ok(r)) => *r.uri_mut() = uri,
(Err(e), Ok(_)) => self.request = Err(e.into()),
(_, Err(_)) => {}
}
self
}
pub(crate) fn extensions(mut self, extensions: ::http::Extensions) -> Self {
if let Ok(r) = &mut self.request {
*r.extensions_mut() = extensions;
}
self
}
pub(crate) fn header<K, V>(mut self, name: K, value: V) -> Self
where
K: TryInto<HeaderName>,
K::Error: Into<RequestBuilderError>,
V: TryInto<HeaderValue>,
V::Error: Into<RequestBuilderError>,
{
match (name.try_into(), value.try_into(), &mut self.request) {
(Ok(name), Ok(value), Ok(r)) => {
r.headers_mut().insert(name, value);
}
(Err(e), _, Ok(_)) => self.request = Err(e.into()),
(_, Err(e), Ok(_)) => self.request = Err(e.into()),
(_, _, Err(_)) => {}
}
self
}
#[cfg(feature = "aws")]
pub(crate) fn headers(mut self, headers: http::HeaderMap) -> Self {
use http::header::{Entry, OccupiedEntry};
if let Ok(ref mut req) = self.request {
// IntoIter of HeaderMap yields (Option<HeaderName>, HeaderValue).
// The first time a name is yielded, it will be Some(name), and if
// there are more values with the same name, the next yield will be
// None.
let mut prev_entry: Option<OccupiedEntry<'_, _>> = None;
for (key, value) in headers {
match key {
Some(key) => match req.headers_mut().entry(key) {
Entry::Occupied(mut e) => {
e.insert(value);
prev_entry = Some(e);
}
Entry::Vacant(e) => {
let e = e.insert_entry(value);
prev_entry = Some(e);
}
},
None => match prev_entry {
Some(ref mut entry) => {
entry.append(value);
}
None => unreachable!("HeaderMap::into_iter yielded None first"),
},
}
}
}
self
}
#[cfg(feature = "gcp")]
pub(crate) fn bearer_auth(mut self, token: &str) -> Self {
let value = HeaderValue::try_from(format!("Bearer {}", token));
match (value, &mut self.request) {
(Ok(mut v), Ok(r)) => {
v.set_sensitive(true);
r.headers_mut().insert(http::header::AUTHORIZATION, v);
}
(Err(e), Ok(_)) => self.request = Err(e.into()),
(_, Err(_)) => {}
}
self
}
#[cfg(any(feature = "aws", feature = "gcp"))]
pub(crate) fn json<S: serde::Serialize>(mut self, s: S) -> Self {
match (serde_json::to_vec(&s), &mut self.request) {
(Ok(json), Ok(request)) => {
*request.body_mut() = json.into();
}
(Err(e), Ok(_)) => self.request = Err(e.into()),
(_, Err(_)) => {}
}
self
}
#[cfg(any(test, feature = "aws", feature = "gcp", feature = "azure"))]
pub(crate) fn query<T: serde::Serialize + ?Sized>(mut self, query: &T) -> Self {
let mut error = None;
if let Ok(ref mut req) = self.request {
let mut out = format!("{}?", req.uri().path());
let start_position = out.len();
let mut encoder = form_urlencoded::Serializer::for_suffix(&mut out, start_position);
let serializer = serde_urlencoded::Serializer::new(&mut encoder);
if let Err(err) = query.serialize(serializer) {
error = Some(err.into());
}
match http::uri::PathAndQuery::from_maybe_shared(out) {
Ok(p) => {
let mut parts = req.uri().clone().into_parts();
parts.path_and_query = Some(p);
*req.uri_mut() = Uri::from_parts(parts).unwrap();
}
Err(err) => error = Some(err.into()),
}
}
if let Some(err) = error {
self.request = Err(err);
}
self
}
#[cfg(any(feature = "gcp", feature = "azure"))]
pub(crate) fn form<T: serde::Serialize>(mut self, form: T) -> Self {
let mut error = None;
if let Ok(ref mut req) = self.request {
match serde_urlencoded::to_string(form) {
Ok(body) => {
req.headers_mut().insert(
http::header::CONTENT_TYPE,
HeaderValue::from_static("application/x-www-form-urlencoded"),
);
*req.body_mut() = body.into();
}
Err(err) => error = Some(err.into()),
}
}
if let Some(err) = error {
self.request = Err(err);
}
self
}
#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
pub(crate) fn body(mut self, b: impl Into<HttpRequestBody>) -> Self {
if let Ok(r) = &mut self.request {
*r.body_mut() = b.into();
}
self
}
pub(crate) fn into_parts(self) -> (HttpClient, Result<HttpRequest, RequestBuilderError>) {
(self.client, self.request)
}
}
#[cfg(any(test, feature = "azure"))]
pub(crate) fn add_query_pairs<I, K, V>(uri: &mut Uri, query_pairs: I)
where
I: IntoIterator,
I::Item: std::borrow::Borrow<(K, V)>,
K: AsRef<str>,
V: AsRef<str>,
{
let mut parts = uri.clone().into_parts();
let mut out = match parts.path_and_query {
Some(p) => match p.query() {
Some(query) => format!("{}?{}", p.path(), query),
None => format!("{}?", p.path()),
},
None => "/?".to_string(),
};
let mut serializer = if out.ends_with('?') {
let start_position = out.len();
form_urlencoded::Serializer::for_suffix(&mut out, start_position)
} else {
form_urlencoded::Serializer::new(&mut out)
};
serializer.extend_pairs(query_pairs);
parts.path_and_query = Some(out.try_into().unwrap());
*uri = Uri::from_parts(parts).unwrap();
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_add_query_pairs() {
let mut uri = Uri::from_static("https://foo@example.com/bananas");
add_query_pairs(&mut uri, [("foo", "1")]);
assert_eq!(uri.to_string(), "https://foo@example.com/bananas?foo=1");
add_query_pairs(&mut uri, [("bingo", "foo"), ("auth", "test")]);
assert_eq!(
uri.to_string(),
"https://foo@example.com/bananas?foo=1&bingo=foo&auth=test"
);
add_query_pairs(&mut uri, [("t1", "funky shenanigans"), ("a", "😀")]);
assert_eq!(
uri.to_string(),
"https://foo@example.com/bananas?foo=1&bingo=foo&auth=test&t1=funky+shenanigans&a=%F0%9F%98%80"
);
}
#[test]
fn test_add_query_pairs_no_path() {
let mut uri = Uri::from_static("https://foo@example.com");
add_query_pairs(&mut uri, [("foo", "1")]);
assert_eq!(uri.to_string(), "https://foo@example.com/?foo=1");
}
#[test]
fn test_request_builder_query() {
let client = HttpClient::new(reqwest::Client::new());
assert_request_uri(
HttpRequestBuilder::new(client.clone()).uri("http://example.com/bananas"),
"http://example.com/bananas",
);
assert_request_uri(
HttpRequestBuilder::new(client.clone())
.uri("http://example.com/bananas")
.query(&[("foo", "1")]),
"http://example.com/bananas?foo=1",
);
assert_request_uri(
HttpRequestBuilder::new(client.clone())
.uri("http://example.com")
.query(&[("foo", "1")]),
"http://example.com/?foo=1",
);
}
fn assert_request_uri(builder: HttpRequestBuilder, expected: &str) {
assert_eq!(builder.into_parts().1.unwrap().uri().to_string(), expected)
}
}
backend/vendor/object_store/src/client/dns.rs
Vendored
+50
View File
@@ -0,0 +1,50 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::net::ToSocketAddrs;
use rand::prelude::SliceRandom;
use reqwest::dns::{Addrs, Name, Resolve, Resolving};
use tokio::task::JoinSet;
type DynErr = Box<dyn std::error::Error + Send + Sync>;
#[derive(Debug)]
pub(crate) struct ShuffleResolver;
impl Resolve for ShuffleResolver {
fn resolve(&self, name: Name) -> Resolving {
Box::pin(async move {
// use `JoinSet` to propagate cancelation
let mut tasks = JoinSet::new();
tasks.spawn_blocking(move || {
let it = (name.as_str(), 0).to_socket_addrs()?;
let mut addrs = it.collect::<Vec<_>>();
addrs.shuffle(&mut rand::rng());
Ok(Box::new(addrs.into_iter()) as Addrs)
});
tasks
.join_next()
.await
.expect("spawned on task")
.map_err(|err| Box::new(err) as DynErr)?
})
}
}
backend/vendor/object_store/src/client/get.rs
Vendored
+429
View File
@@ -0,0 +1,429 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::ops::Range;
use crate::client::header::{header_meta, HeaderConfig};
use crate::client::HttpResponse;
use crate::path::Path;
use crate::{Attribute, Attributes, GetOptions, GetRange, GetResult, GetResultPayload, Result};
use async_trait::async_trait;
use futures::{StreamExt, TryStreamExt};
use http::header::{
CACHE_CONTROL, CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LANGUAGE, CONTENT_RANGE,
CONTENT_TYPE,
};
use http::StatusCode;
use reqwest::header::ToStrError;
/// A client that can perform a get request
#[async_trait]
pub(crate) trait GetClient: Send + Sync + 'static {
const STORE: &'static str;
/// Configure the [`HeaderConfig`] for this client
const HEADER_CONFIG: HeaderConfig;
async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse>;
}
/// Extension trait for [`GetClient`] that adds common retrieval functionality
#[async_trait]
pub(crate) trait GetClientExt {
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult>;
}
#[async_trait]
impl<T: GetClient> GetClientExt for T {
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
let range = options.range.clone();
if let Some(r) = range.as_ref() {
r.is_valid().map_err(|e| crate::Error::Generic {
store: T::STORE,
source: Box::new(e),
})?;
}
let response = self.get_request(location, options).await?;
get_result::<T>(location, range, response).map_err(|e| crate::Error::Generic {
store: T::STORE,
source: Box::new(e),
})
}
}
struct ContentRange {
/// The range of the object returned
range: Range<u64>,
/// The total size of the object being requested
size: u64,
}
impl ContentRange {
/// Parse a content range of the form `bytes <range-start>-<range-end>/<size>`
///
/// <https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Range>
fn from_str(s: &str) -> Option<Self> {
let rem = s.trim().strip_prefix("bytes ")?;
let (range, size) = rem.split_once('/')?;
let size = size.parse().ok()?;
let (start_s, end_s) = range.split_once('-')?;
let start = start_s.parse().ok()?;
let end: u64 = end_s.parse().ok()?;
Some(Self {
size,
range: start..end + 1,
})
}
}
/// A specialized `Error` for get-related errors
#[derive(Debug, thiserror::Error)]
enum GetResultError {
#[error(transparent)]
Header {
#[from]
source: crate::client::header::Error,
},
#[error(transparent)]
InvalidRangeRequest {
#[from]
source: crate::util::InvalidGetRange,
},
#[error("Received non-partial response when range requested")]
NotPartial,
#[error("Content-Range header not present in partial response")]
NoContentRange,
#[error("Failed to parse value for CONTENT_RANGE header: \"{value}\"")]
ParseContentRange { value: String },
#[error("Content-Range header contained non UTF-8 characters")]
InvalidContentRange { source: ToStrError },
#[error("Cache-Control header contained non UTF-8 characters")]
InvalidCacheControl { source: ToStrError },
#[error("Content-Disposition header contained non UTF-8 characters")]
InvalidContentDisposition { source: ToStrError },
#[error("Content-Encoding header contained non UTF-8 characters")]
InvalidContentEncoding { source: ToStrError },
#[error("Content-Language header contained non UTF-8 characters")]
InvalidContentLanguage { source: ToStrError },
#[error("Content-Type header contained non UTF-8 characters")]
InvalidContentType { source: ToStrError },
#[error("Metadata value for \"{key:?}\" contained non UTF-8 characters")]
InvalidMetadata { key: String },
#[error("Requested {expected:?}, got {actual:?}")]
UnexpectedRange {
expected: Range<u64>,
actual: Range<u64>,
},
}
fn get_result<T: GetClient>(
location: &Path,
range: Option<GetRange>,
response: HttpResponse,
) -> Result<GetResult, GetResultError> {
let mut meta = header_meta(location, response.headers(), T::HEADER_CONFIG)?;
// ensure that we receive the range we asked for
let range = if let Some(expected) = range {
if response.status() != StatusCode::PARTIAL_CONTENT {
return Err(GetResultError::NotPartial);
}
let val = response
.headers()
.get(CONTENT_RANGE)
.ok_or(GetResultError::NoContentRange)?;
let value = val
.to_str()
.map_err(|source| GetResultError::InvalidContentRange { source })?;
let value = ContentRange::from_str(value).ok_or_else(|| {
let value = value.into();
GetResultError::ParseContentRange { value }
})?;
let actual = value.range;
// Update size to reflect full size of object (#5272)
meta.size = value.size;
let expected = expected.as_range(meta.size)?;
if actual != expected {
return Err(GetResultError::UnexpectedRange { expected, actual });
}
actual
} else {
0..meta.size
};
macro_rules! parse_attributes {
($headers:expr, $(($header:expr, $attr:expr, $map_err:expr)),*) => {{
let mut attributes = Attributes::new();
$(
if let Some(x) = $headers.get($header) {
let x = x.to_str().map_err($map_err)?;
attributes.insert($attr, x.to_string().into());
}
)*
attributes
}}
}
let mut attributes = parse_attributes!(
response.headers(),
(CACHE_CONTROL, Attribute::CacheControl, |source| {
GetResultError::InvalidCacheControl { source }
}),
(
CONTENT_DISPOSITION,
Attribute::ContentDisposition,
|source| GetResultError::InvalidContentDisposition { source }
),
(CONTENT_ENCODING, Attribute::ContentEncoding, |source| {
GetResultError::InvalidContentEncoding { source }
}),
(CONTENT_LANGUAGE, Attribute::ContentLanguage, |source| {
GetResultError::InvalidContentLanguage { source }
}),
(CONTENT_TYPE, Attribute::ContentType, |source| {
GetResultError::InvalidContentType { source }
})
);
// Add attributes that match the user-defined metadata prefix (e.g. x-amz-meta-)
if let Some(prefix) = T::HEADER_CONFIG.user_defined_metadata_prefix {
for (key, val) in response.headers() {
if let Some(suffix) = key.as_str().strip_prefix(prefix) {
if let Ok(val_str) = val.to_str() {
attributes.insert(
Attribute::Metadata(suffix.to_string().into()),
val_str.to_string().into(),
);
} else {
return Err(GetResultError::InvalidMetadata {
key: key.to_string(),
});
}
}
}
}
let stream = response
.into_body()
.bytes_stream()
.map_err(|source| crate::Error::Generic {
store: T::STORE,
source: Box::new(source),
})
.boxed();
Ok(GetResult {
range,
meta,
attributes,
payload: GetResultPayload::Stream(stream),
})
}
#[cfg(test)]
mod tests {
use super::*;
use http::header::*;
struct TestClient {}
#[async_trait]
impl GetClient for TestClient {
const STORE: &'static str = "TEST";
const HEADER_CONFIG: HeaderConfig = HeaderConfig {
etag_required: false,
last_modified_required: false,
version_header: None,
user_defined_metadata_prefix: Some("x-test-meta-"),
};
async fn get_request(&self, _: &Path, _: GetOptions) -> Result<HttpResponse> {
unimplemented!()
}
}
fn make_response(
object_size: usize,
range: Option<Range<usize>>,
status: StatusCode,
content_range: Option<&str>,
headers: Option<Vec<(&str, &str)>>,
) -> HttpResponse {
let mut builder = http::Response::builder();
if let Some(range) = content_range {
builder = builder.header(CONTENT_RANGE, range);
}
let body = match range {
Some(range) => vec![0_u8; range.end - range.start],
None => vec![0_u8; object_size],
};
if let Some(headers) = headers {
for (key, value) in headers {
builder = builder.header(key, value);
}
}
builder
.status(status)
.header(CONTENT_LENGTH, object_size)
.body(body.into())
.unwrap()
}
#[tokio::test]
async fn test_get_result() {
let path = Path::from("test");
let resp = make_response(12, None, StatusCode::OK, None, None);
let res = get_result::<TestClient>(&path, None, resp).unwrap();
assert_eq!(res.meta.size, 12);
assert_eq!(res.range, 0..12);
let bytes = res.bytes().await.unwrap();
assert_eq!(bytes.len(), 12);
let get_range = GetRange::from(2..3);
let resp = make_response(
12,
Some(2..3),
StatusCode::PARTIAL_CONTENT,
Some("bytes 2-2/12"),
None,
);
let res = get_result::<TestClient>(&path, Some(get_range.clone()), resp).unwrap();
assert_eq!(res.meta.size, 12);
assert_eq!(res.range, 2..3);
let bytes = res.bytes().await.unwrap();
assert_eq!(bytes.len(), 1);
let resp = make_response(12, Some(2..3), StatusCode::OK, None, None);
let err = get_result::<TestClient>(&path, Some(get_range.clone()), resp).unwrap_err();
assert_eq!(
err.to_string(),
"Received non-partial response when range requested"
);
let resp = make_response(
12,
Some(2..3),
StatusCode::PARTIAL_CONTENT,
Some("bytes 2-3/12"),
None,
);
let err = get_result::<TestClient>(&path, Some(get_range.clone()), resp).unwrap_err();
assert_eq!(err.to_string(), "Requested 2..3, got 2..4");
let resp = make_response(
12,
Some(2..3),
StatusCode::PARTIAL_CONTENT,
Some("bytes 2-2/*"),
None,
);
let err = get_result::<TestClient>(&path, Some(get_range.clone()), resp).unwrap_err();
assert_eq!(
err.to_string(),
"Failed to parse value for CONTENT_RANGE header: \"bytes 2-2/*\""
);
let resp = make_response(12, Some(2..3), StatusCode::PARTIAL_CONTENT, None, None);
let err = get_result::<TestClient>(&path, Some(get_range.clone()), resp).unwrap_err();
assert_eq!(
err.to_string(),
"Content-Range header not present in partial response"
);
let resp = make_response(
2,
Some(2..3),
StatusCode::PARTIAL_CONTENT,
Some("bytes 2-3/2"),
None,
);
let err = get_result::<TestClient>(&path, Some(get_range.clone()), resp).unwrap_err();
assert_eq!(
err.to_string(),
"Wanted range starting at 2, but object was only 2 bytes long"
);
let resp = make_response(
6,
Some(2..6),
StatusCode::PARTIAL_CONTENT,
Some("bytes 2-5/6"),
None,
);
let res = get_result::<TestClient>(&path, Some(GetRange::Suffix(4)), resp).unwrap();
assert_eq!(res.meta.size, 6);
assert_eq!(res.range, 2..6);
let bytes = res.bytes().await.unwrap();
assert_eq!(bytes.len(), 4);
let resp = make_response(
6,
Some(2..6),
StatusCode::PARTIAL_CONTENT,
Some("bytes 2-3/6"),
None,
);
let err = get_result::<TestClient>(&path, Some(GetRange::Suffix(4)), resp).unwrap_err();
assert_eq!(err.to_string(), "Requested 2..6, got 2..4");
let resp = make_response(
12,
None,
StatusCode::OK,
None,
Some(vec![("x-test-meta-foo", "bar")]),
);
let res = get_result::<TestClient>(&path, None, resp).unwrap();
assert_eq!(res.meta.size, 12);
assert_eq!(res.range, 0..12);
assert_eq!(
res.attributes.get(&Attribute::Metadata("foo".into())),
Some(&"bar".into())
);
let bytes = res.bytes().await.unwrap();
assert_eq!(bytes.len(), 12);
}
}
backend/vendor/object_store/src/client/header.rs
Vendored
+166
View File
@@ -0,0 +1,166 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Logic for extracting ObjectMeta from headers used by AWS, GCP and Azure
use crate::path::Path;
use crate::ObjectMeta;
use chrono::{DateTime, TimeZone, Utc};
use http::header::{CONTENT_LENGTH, ETAG, LAST_MODIFIED};
use http::HeaderMap;
#[derive(Debug, Copy, Clone)]
/// Configuration for header extraction
pub(crate) struct HeaderConfig {
/// Whether to require an ETag header when extracting [`ObjectMeta`] from headers.
///
/// Defaults to `true`
pub etag_required: bool,
/// Whether to require a Last-Modified header when extracting [`ObjectMeta`] from headers.
///
/// Defaults to `true`
pub last_modified_required: bool,
/// The version header name if any
pub version_header: Option<&'static str>,
/// The user defined metadata prefix if any
pub user_defined_metadata_prefix: Option<&'static str>,
}
#[derive(Debug, thiserror::Error)]
pub(crate) enum Error {
#[error("ETag Header missing from response")]
MissingEtag,
#[error("Received header containing non-ASCII data")]
BadHeader { source: reqwest::header::ToStrError },
#[error("Last-Modified Header missing from response")]
MissingLastModified,
#[error("Content-Length Header missing from response")]
MissingContentLength,
#[error("Invalid last modified '{}': {}", last_modified, source)]
InvalidLastModified {
last_modified: String,
source: chrono::ParseError,
},
#[error("Invalid content length '{}': {}", content_length, source)]
InvalidContentLength {
content_length: String,
source: std::num::ParseIntError,
},
}
/// Extracts a PutResult from the provided [`HeaderMap`]
#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
pub(crate) fn get_put_result(
headers: &HeaderMap,
version: &str,
) -> Result<crate::PutResult, Error> {
let e_tag = Some(get_etag(headers)?);
let version = get_version(headers, version)?;
Ok(crate::PutResult { e_tag, version })
}
/// Extracts a optional version from the provided [`HeaderMap`]
#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
pub(crate) fn get_version(headers: &HeaderMap, version: &str) -> Result<Option<String>, Error> {
Ok(match headers.get(version) {
Some(x) => Some(
x.to_str()
.map_err(|source| Error::BadHeader { source })?
.to_string(),
),
None => None,
})
}
/// Extracts an etag from the provided [`HeaderMap`]
pub(crate) fn get_etag(headers: &HeaderMap) -> Result<String, Error> {
let e_tag = headers.get(ETAG).ok_or(Error::MissingEtag)?;
Ok(e_tag
.to_str()
.map_err(|source| Error::BadHeader { source })?
.to_string())
}
/// Extracts [`ObjectMeta`] from the provided [`HeaderMap`]
pub(crate) fn header_meta(
location: &Path,
headers: &HeaderMap,
cfg: HeaderConfig,
) -> Result<ObjectMeta, Error> {
let last_modified = match headers.get(LAST_MODIFIED) {
Some(last_modified) => {
let last_modified = last_modified
.to_str()
.map_err(|source| Error::BadHeader { source })?;
DateTime::parse_from_rfc2822(last_modified)
.map_err(|source| Error::InvalidLastModified {
last_modified: last_modified.into(),
source,
})?
.with_timezone(&Utc)
}
None if cfg.last_modified_required => return Err(Error::MissingLastModified),
None => Utc.timestamp_nanos(0),
};
let e_tag = match get_etag(headers) {
Ok(e_tag) => Some(e_tag),
Err(Error::MissingEtag) if !cfg.etag_required => None,
Err(e) => return Err(e),
};
let content_length = headers
.get(CONTENT_LENGTH)
.ok_or(Error::MissingContentLength)?;
let content_length = content_length
.to_str()
.map_err(|source| Error::BadHeader { source })?;
let size = content_length
.parse()
.map_err(|source| Error::InvalidContentLength {
content_length: content_length.into(),
source,
})?;
let version = match cfg.version_header.and_then(|h| headers.get(h)) {
Some(v) => Some(
v.to_str()
.map_err(|source| Error::BadHeader { source })?
.to_string(),
),
None => None,
};
Ok(ObjectMeta {
location: location.clone(),
last_modified,
version,
size,
e_tag,
})
}
backend/vendor/object_store/src/client/http/body.rs
Vendored
+234
View File
@@ -0,0 +1,234 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::{HttpError, HttpErrorKind};
use crate::{collect_bytes, PutPayload};
use bytes::Bytes;
use futures::stream::BoxStream;
use futures::StreamExt;
use http_body_util::combinators::BoxBody;
use http_body_util::{BodyExt, Full};
use hyper::body::{Body, Frame, SizeHint};
use std::pin::Pin;
use std::task::{Context, Poll};
/// An HTTP Request
pub type HttpRequest = http::Request<HttpRequestBody>;
/// The [`Body`] of an [`HttpRequest`]
#[derive(Debug, Clone)]
pub struct HttpRequestBody(Inner);
impl HttpRequestBody {
/// An empty [`HttpRequestBody`]
pub fn empty() -> Self {
Self(Inner::Bytes(Bytes::new()))
}
#[cfg(not(target_arch = "wasm32"))]
pub(crate) fn into_reqwest(self) -> reqwest::Body {
match self.0 {
Inner::Bytes(b) => b.into(),
Inner::PutPayload(_, payload) => reqwest::Body::wrap_stream(futures::stream::iter(
payload.into_iter().map(Ok::<_, HttpError>),
)),
}
}
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
pub(crate) fn into_reqwest(self) -> reqwest::Body {
match self.0 {
Inner::Bytes(b) => b.into(),
Inner::PutPayload(_, payload) => Bytes::from(payload).into(),
}
}
/// Returns true if this body is empty
pub fn is_empty(&self) -> bool {
match &self.0 {
Inner::Bytes(x) => x.is_empty(),
Inner::PutPayload(_, x) => x.iter().any(|x| !x.is_empty()),
}
}
/// Returns the total length of the [`Bytes`] in this body
pub fn content_length(&self) -> usize {
match &self.0 {
Inner::Bytes(x) => x.len(),
Inner::PutPayload(_, x) => x.content_length(),
}
}
/// If this body consists of a single contiguous [`Bytes`], returns it
pub fn as_bytes(&self) -> Option<&Bytes> {
match &self.0 {
Inner::Bytes(x) => Some(x),
_ => None,
}
}
}
impl From<Bytes> for HttpRequestBody {
fn from(value: Bytes) -> Self {
Self(Inner::Bytes(value))
}
}
impl From<Vec<u8>> for HttpRequestBody {
fn from(value: Vec<u8>) -> Self {
Self(Inner::Bytes(value.into()))
}
}
impl From<String> for HttpRequestBody {
fn from(value: String) -> Self {
Self(Inner::Bytes(value.into()))
}
}
impl From<PutPayload> for HttpRequestBody {
fn from(value: PutPayload) -> Self {
Self(Inner::PutPayload(0, value))
}
}
#[derive(Debug, Clone)]
enum Inner {
Bytes(Bytes),
PutPayload(usize, PutPayload),
}
impl Body for HttpRequestBody {
type Data = Bytes;
type Error = HttpError;
fn poll_frame(
mut self: Pin<&mut Self>,
_cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
Poll::Ready(match &mut self.0 {
Inner::Bytes(bytes) => {
let out = bytes.split_off(0);
if out.is_empty() {
None
} else {
Some(Ok(Frame::data(out)))
}
}
Inner::PutPayload(offset, payload) => {
let slice = payload.as_ref();
if *offset == slice.len() {
None
} else {
Some(Ok(Frame::data(
slice[std::mem::replace(offset, *offset + 1)].clone(),
)))
}
}
})
}
fn is_end_stream(&self) -> bool {
match self.0 {
Inner::Bytes(ref bytes) => bytes.is_empty(),
Inner::PutPayload(offset, ref body) => offset == body.as_ref().len(),
}
}
fn size_hint(&self) -> SizeHint {
match self.0 {
Inner::Bytes(ref bytes) => SizeHint::with_exact(bytes.len() as u64),
Inner::PutPayload(offset, ref payload) => {
let iter = payload.as_ref().iter().skip(offset);
SizeHint::with_exact(iter.map(|x| x.len() as u64).sum())
}
}
}
}
/// An HTTP response
pub type HttpResponse = http::Response<HttpResponseBody>;
/// The body of an [`HttpResponse`]
#[derive(Debug)]
pub struct HttpResponseBody(BoxBody<Bytes, HttpError>);
impl HttpResponseBody {
/// Create an [`HttpResponseBody`] from the provided [`Body`]
///
/// Note: [`BodyExt::map_err`] can be used to alter error variants
pub fn new<B>(body: B) -> Self
where
B: Body<Data = Bytes, Error = HttpError> + Send + Sync + 'static,
{
Self(BoxBody::new(body))
}
/// Collects this response into a [`Bytes`]
pub async fn bytes(self) -> Result<Bytes, HttpError> {
let size_hint = self.0.size_hint().lower();
let s = self.0.into_data_stream();
collect_bytes(s, Some(size_hint)).await
}
/// Returns a stream of this response data
pub fn bytes_stream(self) -> BoxStream<'static, Result<Bytes, HttpError>> {
self.0.into_data_stream().boxed()
}
/// Returns the response as a [`String`]
pub(crate) async fn text(self) -> Result<String, HttpError> {
let b = self.bytes().await?;
String::from_utf8(b.into()).map_err(|e| HttpError::new(HttpErrorKind::Decode, e))
}
#[cfg(any(feature = "aws", feature = "gcp", feature = "azure"))]
pub(crate) async fn json<B: serde::de::DeserializeOwned>(self) -> Result<B, HttpError> {
let b = self.bytes().await?;
serde_json::from_slice(&b).map_err(|e| HttpError::new(HttpErrorKind::Decode, e))
}
}
impl Body for HttpResponseBody {
type Data = Bytes;
type Error = HttpError;
fn poll_frame(
mut self: Pin<&mut Self>,
cx: &mut Context<'_>,
) -> Poll<Option<Result<Frame<Self::Data>, Self::Error>>> {
Pin::new(&mut self.0).poll_frame(cx)
}
}
impl From<Bytes> for HttpResponseBody {
fn from(value: Bytes) -> Self {
Self::new(Full::new(value).map_err(|e| match e {}))
}
}
impl From<Vec<u8>> for HttpResponseBody {
fn from(value: Vec<u8>) -> Self {
Bytes::from(value).into()
}
}
impl From<String> for HttpResponseBody {
fn from(value: String) -> Self {
Bytes::from(value).into()
}
}
backend/vendor/object_store/src/client/http/connection.rs
Vendored
+384
View File
@@ -0,0 +1,384 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::builder::{HttpRequestBuilder, RequestBuilderError};
use crate::client::{HttpRequest, HttpResponse, HttpResponseBody};
use crate::ClientOptions;
use async_trait::async_trait;
use http::{Method, Uri};
use http_body_util::BodyExt;
use std::error::Error;
use std::sync::Arc;
use tokio::runtime::Handle;
/// An HTTP protocol error
///
/// Clients should return this when an HTTP request fails to be completed, e.g. because
/// of a connection issue. This does **not** include HTTP requests that are return
/// non 2xx Status Codes, as these should instead be returned as an [`HttpResponse`]
/// with the appropriate status code set.
#[derive(Debug, thiserror::Error)]
#[error("HTTP error: {source}")]
pub struct HttpError {
kind: HttpErrorKind,
#[source]
source: Box<dyn Error + Send + Sync>,
}
/// Identifies the kind of [`HttpError`]
///
/// This is used, among other things, to determine if a request can be retried
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum HttpErrorKind {
/// An error occurred whilst connecting to the remote
///
/// Will be automatically retried
Connect,
/// An error occurred whilst making the request
///
/// Will be automatically retried
Request,
/// Request timed out
///
/// Will be automatically retried if the request is idempotent
Timeout,
/// The request was aborted
///
/// Will be automatically retried if the request is idempotent
Interrupted,
/// An error occurred whilst decoding the response
///
/// Will not be automatically retried
Decode,
/// An unknown error occurred
///
/// Will not be automatically retried
Unknown,
}
impl HttpError {
/// Create a new [`HttpError`] with the optional status code
pub fn new<E>(kind: HttpErrorKind, e: E) -> Self
where
E: Error + Send + Sync + 'static,
{
Self {
kind,
source: Box::new(e),
}
}
pub(crate) fn reqwest(e: reqwest::Error) -> Self {
#[cfg(not(target_arch = "wasm32"))]
let is_connect = || e.is_connect();
#[cfg(target_arch = "wasm32")]
let is_connect = || false;
let mut kind = if e.is_timeout() {
HttpErrorKind::Timeout
} else if is_connect() {
HttpErrorKind::Connect
} else if e.is_decode() {
HttpErrorKind::Decode
} else {
HttpErrorKind::Unknown
};
// Reqwest error variants aren't great, attempt to refine them
let mut source = e.source();
while let Some(e) = source {
if let Some(e) = e.downcast_ref::<hyper::Error>() {
if e.is_closed() || e.is_incomplete_message() || e.is_body_write_aborted() {
kind = HttpErrorKind::Request;
} else if e.is_timeout() {
kind = HttpErrorKind::Timeout;
}
break;
}
if let Some(e) = e.downcast_ref::<std::io::Error>() {
match e.kind() {
std::io::ErrorKind::TimedOut => kind = HttpErrorKind::Timeout,
std::io::ErrorKind::ConnectionAborted
| std::io::ErrorKind::BrokenPipe
| std::io::ErrorKind::UnexpectedEof => kind = HttpErrorKind::Interrupted,
_ => {}
}
break;
}
source = e.source();
}
Self {
kind,
// We strip URL as it will be included by RetryError if not sensitive
source: Box::new(e.without_url()),
}
}
/// Returns the [`HttpErrorKind`]
pub fn kind(&self) -> HttpErrorKind {
self.kind
}
}
/// An asynchronous function from a [`HttpRequest`] to a [`HttpResponse`].
#[async_trait]
pub trait HttpService: std::fmt::Debug + Send + Sync + 'static {
/// Perform [`HttpRequest`] returning [`HttpResponse`]
async fn call(&self, req: HttpRequest) -> Result<HttpResponse, HttpError>;
}
/// An HTTP client
#[derive(Debug, Clone)]
pub struct HttpClient(Arc<dyn HttpService>);
impl HttpClient {
/// Create a new [`HttpClient`] from an [`HttpService`]
pub fn new(service: impl HttpService + 'static) -> Self {
Self(Arc::new(service))
}
/// Performs [`HttpRequest`] using this client
pub async fn execute(&self, request: HttpRequest) -> Result<HttpResponse, HttpError> {
self.0.call(request).await
}
#[allow(unused)]
pub(crate) fn get<U>(&self, url: U) -> HttpRequestBuilder
where
U: TryInto<Uri>,
U::Error: Into<RequestBuilderError>,
{
self.request(Method::GET, url)
}
#[allow(unused)]
pub(crate) fn post<U>(&self, url: U) -> HttpRequestBuilder
where
U: TryInto<Uri>,
U::Error: Into<RequestBuilderError>,
{
self.request(Method::POST, url)
}
#[allow(unused)]
pub(crate) fn put<U>(&self, url: U) -> HttpRequestBuilder
where
U: TryInto<Uri>,
U::Error: Into<RequestBuilderError>,
{
self.request(Method::PUT, url)
}
#[allow(unused)]
pub(crate) fn delete<U>(&self, url: U) -> HttpRequestBuilder
where
U: TryInto<Uri>,
U::Error: Into<RequestBuilderError>,
{
self.request(Method::DELETE, url)
}
pub(crate) fn request<U>(&self, method: Method, url: U) -> HttpRequestBuilder
where
U: TryInto<Uri>,
U::Error: Into<RequestBuilderError>,
{
HttpRequestBuilder::new(self.clone())
.uri(url)
.method(method)
}
}
#[async_trait]
#[cfg(not(target_arch = "wasm32"))]
impl HttpService for reqwest::Client {
async fn call(&self, req: HttpRequest) -> Result<HttpResponse, HttpError> {
let (parts, body) = req.into_parts();
let url = parts.uri.to_string().parse().unwrap();
let mut req = reqwest::Request::new(parts.method, url);
*req.headers_mut() = parts.headers;
*req.body_mut() = Some(body.into_reqwest());
let r = self.execute(req).await.map_err(HttpError::reqwest)?;
let res: http::Response<reqwest::Body> = r.into();
let (parts, body) = res.into_parts();
let body = HttpResponseBody::new(body.map_err(HttpError::reqwest));
Ok(HttpResponse::from_parts(parts, body))
}
}
#[async_trait]
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
impl HttpService for reqwest::Client {
async fn call(&self, req: HttpRequest) -> Result<HttpResponse, HttpError> {
use futures::{
channel::{mpsc, oneshot},
SinkExt, StreamExt, TryStreamExt,
};
use http_body_util::{Empty, StreamBody};
use wasm_bindgen_futures::spawn_local;
let (parts, body) = req.into_parts();
let url = parts.uri.to_string().parse().unwrap();
let mut req = reqwest::Request::new(parts.method, url);
*req.headers_mut() = parts.headers;
*req.body_mut() = Some(body.into_reqwest());
let (mut tx, rx) = mpsc::channel(1);
let (tx_parts, rx_parts) = oneshot::channel();
let res_fut = self.execute(req);
spawn_local(async move {
match res_fut.await.map_err(HttpError::reqwest) {
Err(err) => {
let _ = tx_parts.send(Err(err));
drop(tx);
}
Ok(res) => {
let (mut parts, _) = http::Response::new(Empty::<()>::new()).into_parts();
parts.headers = res.headers().clone();
parts.status = res.status();
let _ = tx_parts.send(Ok(parts));
let mut stream = res.bytes_stream().map_err(HttpError::reqwest);
while let Some(chunk) = stream.next().await {
if let Err(_e) = tx.send(chunk).await {
// Disconnected due to a transitive drop of the receiver
break;
}
}
}
}
});
let parts = rx_parts.await.unwrap()?;
let safe_stream = rx.map(|chunk| {
let frame = hyper::body::Frame::data(chunk?);
Ok(frame)
});
let body = HttpResponseBody::new(StreamBody::new(safe_stream));
Ok(HttpResponse::from_parts(parts, body))
}
}
/// A factory for [`HttpClient`]
pub trait HttpConnector: std::fmt::Debug + Send + Sync + 'static {
/// Create a new [`HttpClient`] with the provided [`ClientOptions`]
fn connect(&self, options: &ClientOptions) -> crate::Result<HttpClient>;
}
/// [`HttpConnector`] using [`reqwest::Client`]
#[derive(Debug, Default)]
#[allow(missing_copy_implementations)]
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
pub struct ReqwestConnector {}
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
impl HttpConnector for ReqwestConnector {
fn connect(&self, options: &ClientOptions) -> crate::Result<HttpClient> {
let client = options.client()?;
Ok(HttpClient::new(client))
}
}
/// [`reqwest::Client`] connector that performs all I/O on the provided tokio
/// [`Runtime`] (thread pool).
///
/// This adapter is most useful when you wish to segregate I/O from CPU bound
/// work that may be happening on the [`Runtime`].
///
/// [`Runtime`]: tokio::runtime::Runtime
///
/// # Example: Spawning requests on separate runtime
///
/// ```
/// # use std::sync::Arc;
/// # use tokio::runtime::Runtime;
/// # use object_store::azure::MicrosoftAzureBuilder;
/// # use object_store::client::SpawnedReqwestConnector;
/// # use object_store::ObjectStore;
/// # fn get_io_runtime() -> Runtime {
/// # tokio::runtime::Builder::new_current_thread().build().unwrap()
/// # }
/// # fn main() -> Result<(), object_store::Error> {
/// // create a tokio runtime for I/O.
/// let io_runtime: Runtime = get_io_runtime();
/// // configure a store using the runtime.
/// let handle = io_runtime.handle().clone(); // get a handle to the same runtime
/// let store: Arc<dyn ObjectStore> = Arc::new(
/// MicrosoftAzureBuilder::new()
/// .with_http_connector(SpawnedReqwestConnector::new(handle))
/// .with_container_name("my_container")
/// .with_account("my_account")
/// .build()?
/// );
/// // any requests made using store will be spawned on the io_runtime
/// # Ok(())
/// # }
/// ```
#[derive(Debug)]
#[allow(missing_copy_implementations)]
#[cfg(not(target_arch = "wasm32"))]
pub struct SpawnedReqwestConnector {
runtime: Handle,
}
#[cfg(not(target_arch = "wasm32"))]
impl SpawnedReqwestConnector {
/// Create a new [`SpawnedReqwestConnector`] with the provided [`Handle`] to
/// a tokio [`Runtime`]
///
/// [`Runtime`]: tokio::runtime::Runtime
pub fn new(runtime: Handle) -> Self {
Self { runtime }
}
}
#[cfg(not(target_arch = "wasm32"))]
impl HttpConnector for SpawnedReqwestConnector {
fn connect(&self, options: &ClientOptions) -> crate::Result<HttpClient> {
let spawn_service = super::SpawnService::new(options.client()?, self.runtime.clone());
Ok(HttpClient::new(spawn_service))
}
}
#[cfg(all(target_arch = "wasm32", target_os = "wasi"))]
pub(crate) fn http_connector(
custom: Option<Arc<dyn HttpConnector>>,
) -> crate::Result<Arc<dyn HttpConnector>> {
match custom {
Some(x) => Ok(x),
None => Err(crate::Error::NotSupported {
source: "WASI architectures must provide an HTTPConnector"
.to_string()
.into(),
}),
}
}
#[cfg(not(all(target_arch = "wasm32", target_os = "wasi")))]
pub(crate) fn http_connector(
custom: Option<Arc<dyn HttpConnector>>,
) -> crate::Result<Arc<dyn HttpConnector>> {
match custom {
Some(x) => Ok(x),
None => Ok(Arc::new(ReqwestConnector {})),
}
}
backend/vendor/object_store/src/client/http/mod.rs
Vendored
+27
View File
@@ -0,0 +1,27 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! HTTP client abstraction
mod body;
pub use body::*;
mod connection;
pub use connection::*;
mod spawn;
pub use spawn::*;
backend/vendor/object_store/src/client/http/spawn.rs
Vendored
+167
View File
@@ -0,0 +1,167 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::{
HttpError, HttpErrorKind, HttpRequest, HttpResponse, HttpResponseBody, HttpService,
};
use async_trait::async_trait;
use bytes::Bytes;
use http::Response;
use http_body_util::BodyExt;
use hyper::body::{Body, Frame};
use std::pin::Pin;
use std::task::{Context, Poll};
use thiserror::Error;
use tokio::runtime::Handle;
use tokio::task::JoinHandle;
/// Spawn error
#[derive(Debug, Error)]
#[error("SpawnError")]
struct SpawnError {}
impl From<SpawnError> for HttpError {
fn from(value: SpawnError) -> Self {
Self::new(HttpErrorKind::Interrupted, value)
}
}
/// Wraps a provided [`HttpService`] and runs it on a separate tokio runtime
///
/// See example on [`SpawnedReqwestConnector`]
///
/// [`SpawnedReqwestConnector`]: crate::client::http::SpawnedReqwestConnector
#[derive(Debug)]
pub struct SpawnService<T: HttpService + Clone> {
inner: T,
runtime: Handle,
}
impl<T: HttpService + Clone> SpawnService<T> {
/// Creates a new [`SpawnService`] from the provided
pub fn new(inner: T, runtime: Handle) -> Self {
Self { inner, runtime }
}
}
#[async_trait]
impl<T: HttpService + Clone> HttpService for SpawnService<T> {
async fn call(&self, req: HttpRequest) -> Result<HttpResponse, HttpError> {
let inner = self.inner.clone();
let (send, recv) = tokio::sync::oneshot::channel();
// We use an unbounded channel to prevent backpressure across the runtime boundary
// which could in turn starve the underlying IO operations
let (sender, receiver) = tokio::sync::mpsc::unbounded_channel();
let handle = SpawnHandle(self.runtime.spawn(async move {
let r = match HttpService::call(&inner, req).await {
Ok(resp) => resp,
Err(e) => {
let _ = send.send(Err(e));
return;
}
};
let (parts, mut body) = r.into_parts();
if send.send(Ok(parts)).is_err() {
return;
}
while let Some(x) = body.frame().await {
sender.send(x).unwrap();
}
}));
let parts = recv.await.map_err(|_| SpawnError {})??;
Ok(Response::from_parts(
parts,
HttpResponseBody::new(SpawnBody {
stream: receiver,
_worker: handle,
}),
))
}
}
/// A wrapper around a [`JoinHandle`] that aborts on drop
struct SpawnHandle(JoinHandle<()>);
impl Drop for SpawnHandle {
fn drop(&mut self) {
self.0.abort();
}
}
type StreamItem = Result<Frame<Bytes>, HttpError>;
struct SpawnBody {
stream: tokio::sync::mpsc::UnboundedReceiver<StreamItem>,
_worker: SpawnHandle,
}
impl Body for SpawnBody {
type Data = Bytes;
type Error = HttpError;
fn poll_frame(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<StreamItem>> {
self.stream.poll_recv(cx)
}
}
#[cfg(not(target_arch = "wasm32"))]
#[cfg(test)]
mod tests {
use super::*;
use crate::client::mock_server::MockServer;
use crate::client::retry::RetryExt;
use crate::client::HttpClient;
use crate::RetryConfig;
async fn test_client(client: HttpClient) {
let (send, recv) = tokio::sync::oneshot::channel();
let mock = MockServer::new().await;
mock.push(Response::new("BANANAS".to_string()));
let url = mock.url().to_string();
let thread = std::thread::spawn(|| {
futures::executor::block_on(async move {
let retry = RetryConfig::default();
let ret = client.get(url).send_retry(&retry).await.unwrap();
let payload = ret.into_body().bytes().await.unwrap();
assert_eq!(payload.as_ref(), b"BANANAS");
let _ = send.send(());
})
});
recv.await.unwrap();
thread.join().unwrap();
}
#[tokio::test]
async fn test_spawn() {
let client = HttpClient::new(SpawnService::new(reqwest::Client::new(), Handle::current()));
test_client(client).await;
}
#[tokio::test]
#[should_panic]
async fn test_no_spawn() {
let client = HttpClient::new(reqwest::Client::new());
test_client(client).await;
}
}
backend/vendor/object_store/src/client/list.rs
Vendored
+157
View File
@@ -0,0 +1,157 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::pagination::stream_paginated;
use crate::path::Path;
use crate::Result;
use crate::{ListPage, ListResult, ObjectMeta};
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::{StreamExt, TryStreamExt};
use std::collections::BTreeSet;
/// A client that can perform paginated list requests
#[async_trait]
pub(crate) trait ListClient: Send + Sync + 'static {
async fn list_request(
&self,
prefix: Option<&str>,
delimiter: bool,
token: Option<&str>,
offset: Option<&str>,
max_keys: Option<usize>,
) -> Result<(ListResult, Option<String>)>;
}
/// Extension trait for [`ListClient`] that adds common listing functionality
#[async_trait]
pub(crate) trait ListClientExt {
fn list_paginated(
&self,
prefix: Option<&Path>,
delimiter: bool,
offset: Option<&Path>,
) -> BoxStream<'static, Result<ListResult>>;
/// A single delimited page: one backend request, resumable via the
/// returned continuation token.
async fn list_delimited_page(
&self,
prefix: Option<&Path>,
token: Option<&str>,
max_keys: Option<usize>,
) -> Result<ListPage>;
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>>;
#[allow(unused)]
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>>;
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult>;
}
#[async_trait]
impl<T: ListClient + Clone> ListClientExt for T {
fn list_paginated(
&self,
prefix: Option<&Path>,
delimiter: bool,
offset: Option<&Path>,
) -> BoxStream<'static, Result<ListResult>> {
let offset = offset.map(|x| x.to_string());
let prefix = prefix
.filter(|x| !x.as_ref().is_empty())
.map(|p| format!("{}{}", p.as_ref(), crate::path::DELIMITER));
stream_paginated(
self.clone(),
(prefix, offset),
move |client, (prefix, offset), token| async move {
let (r, next_token) = client
.list_request(
prefix.as_deref(),
delimiter,
token.as_deref(),
offset.as_deref(),
None,
)
.await?;
Ok((r, (prefix, offset), next_token))
},
)
.boxed()
}
async fn list_delimited_page(
&self,
prefix: Option<&Path>,
token: Option<&str>,
max_keys: Option<usize>,
) -> Result<ListPage> {
let prefix = prefix
.filter(|x| !x.as_ref().is_empty())
.map(|p| format!("{}{}", p.as_ref(), crate::path::DELIMITER));
let (r, next_token) = self
.list_request(prefix.as_deref(), true, token, None, max_keys)
.await?;
Ok(ListPage {
common_prefixes: r.common_prefixes,
objects: r.objects,
next_token,
})
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
self.list_paginated(prefix, false, None)
.map_ok(|r| futures::stream::iter(r.objects.into_iter().map(Ok)))
.try_flatten()
.boxed()
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
self.list_paginated(prefix, false, Some(offset))
.map_ok(|r| futures::stream::iter(r.objects.into_iter().map(Ok)))
.try_flatten()
.boxed()
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
let mut stream = self.list_paginated(prefix, true, None);
let mut common_prefixes = BTreeSet::new();
let mut objects = Vec::new();
while let Some(result) = stream.next().await {
let response = result?;
common_prefixes.extend(response.common_prefixes.into_iter());
objects.extend(response.objects.into_iter());
}
Ok(ListResult {
common_prefixes: common_prefixes.into_iter().collect(),
objects,
})
}
}
backend/vendor/object_store/src/client/mock_server.rs
Vendored
+131
View File
@@ -0,0 +1,131 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use futures::future::BoxFuture;
use futures::FutureExt;
use hyper::body::Incoming;
use hyper::server::conn::http1;
use hyper::service::service_fn;
use hyper::{Request, Response};
use hyper_util::rt::TokioIo;
use parking_lot::Mutex;
use std::collections::VecDeque;
use std::convert::Infallible;
use std::future::Future;
use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net::TcpListener;
use tokio::sync::oneshot;
use tokio::task::{JoinHandle, JoinSet};
pub(crate) type ResponseFn =
Box<dyn FnOnce(Request<Incoming>) -> BoxFuture<'static, Response<String>> + Send>;
/// A mock server
pub(crate) struct MockServer {
responses: Arc<Mutex<VecDeque<ResponseFn>>>,
shutdown: oneshot::Sender<()>,
handle: JoinHandle<()>,
url: String,
}
impl MockServer {
pub(crate) async fn new() -> Self {
let responses: Arc<Mutex<VecDeque<ResponseFn>>> =
Arc::new(Mutex::new(VecDeque::with_capacity(10)));
let addr = SocketAddr::from(([127, 0, 0, 1], 0));
let listener = TcpListener::bind(addr).await.unwrap();
let (shutdown, mut rx) = oneshot::channel::<()>();
let url = format!("http://{}", listener.local_addr().unwrap());
let r = Arc::clone(&responses);
let handle = tokio::spawn(async move {
let mut set = JoinSet::new();
loop {
let (stream, _) = tokio::select! {
conn = listener.accept() => conn.unwrap(),
_ = &mut rx => break,
};
let r = Arc::clone(&r);
set.spawn(async move {
let _ = http1::Builder::new()
.serve_connection(
TokioIo::new(stream),
service_fn(move |req| {
let r = Arc::clone(&r);
let next = r.lock().pop_front();
async move {
Ok::<_, Infallible>(match next {
Some(r) => r(req).await,
None => Response::new("Hello World".to_string()),
})
}
}),
)
.await;
});
}
set.abort_all();
});
Self {
responses,
shutdown,
handle,
url,
}
}
/// The url of the mock server
pub(crate) fn url(&self) -> &str {
&self.url
}
/// Add a response
pub(crate) fn push(&self, response: Response<String>) {
self.push_fn(|_| response)
}
/// Add a response function
pub(crate) fn push_fn<F>(&self, f: F)
where
F: FnOnce(Request<Incoming>) -> Response<String> + Send + 'static,
{
let f = Box::new(|req| async move { f(req) }.boxed());
self.responses.lock().push_back(f)
}
pub(crate) fn push_async_fn<F, Fut>(&self, f: F)
where
F: FnOnce(Request<Incoming>) -> Fut + Send + 'static,
Fut: Future<Output = Response<String>> + Send + 'static,
{
self.responses.lock().push_back(Box::new(|r| f(r).boxed()))
}
/// Shutdown the mock server
pub(crate) async fn shutdown(self) {
let _ = self.shutdown.send(());
self.handle.await.unwrap()
}
}
backend/vendor/object_store/src/client/mod.rs
Vendored
+1024
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/client/pagination.rs
Vendored
+80
View File
@@ -0,0 +1,80 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::Result;
use futures::Stream;
use std::future::Future;
/// Takes a paginated operation `op` that when called with:
///
/// - A state `S`
/// - An optional next token `Option<String>`
///
/// Returns
///
/// - A response value `T`
/// - The next state `S`
/// - The next continuation token `Option<String>`
///
/// And converts it into a `Stream<Result<T>>` which will first call `op(state, None)`, and yield
/// the returned response `T`. If the returned continuation token was `None` the stream will then
/// finish, otherwise it will continue to call `op(state, token)` with the values returned by the
/// previous call to `op`, until a continuation token of `None` is returned
///
pub(crate) fn stream_paginated<F, Fut, S, T, C>(
client: C,
state: S,
op: F,
) -> impl Stream<Item = Result<T>>
where
C: Clone,
F: Fn(C, S, Option<String>) -> Fut + Copy,
Fut: Future<Output = Result<(T, S, Option<String>)>>,
{
enum PaginationState<T> {
Start(T),
HasMore(T, String),
Done,
}
futures::stream::unfold(PaginationState::Start(state), move |state| {
let client = client.clone();
async move {
let (s, page_token) = match state {
PaginationState::Start(s) => (s, None),
PaginationState::HasMore(s, page_token) if !page_token.is_empty() => {
(s, Some(page_token))
}
_ => {
return None;
}
};
let (resp, s, continuation) = match op(client, s, page_token).await {
Ok(resp) => resp,
Err(e) => return Some((Err(e), PaginationState::Done)),
};
let next_state = match continuation {
Some(token) => PaginationState::HasMore(s, token),
None => PaginationState::Done,
};
Some((Ok(resp), next_state))
}
})
}
backend/vendor/object_store/src/client/parts.rs
Vendored
+48
View File
@@ -0,0 +1,48 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::multipart::PartId;
use parking_lot::Mutex;
/// An interior mutable collection of upload parts and their corresponding part index
#[derive(Debug, Default)]
pub(crate) struct Parts(Mutex<Vec<(usize, PartId)>>);
impl Parts {
/// Record the [`PartId`] for a given index
///
/// Note: calling this method multiple times with the same `part_idx`
/// will result in multiple [`PartId`] in the final output
pub(crate) fn put(&self, part_idx: usize, id: PartId) {
self.0.lock().push((part_idx, id))
}
/// Produce the final list of [`PartId`] ordered by `part_idx`
///
/// `expected` is the number of parts expected in the final result
pub(crate) fn finish(&self, expected: usize) -> crate::Result<Vec<PartId>> {
let mut parts = self.0.lock();
if parts.len() != expected {
return Err(crate::Error::Generic {
store: "Parts",
source: "Missing part".to_string().into(),
});
}
parts.sort_unstable_by_key(|(idx, _)| *idx);
Ok(parts.drain(..).map(|(_, v)| v).collect())
}
}
backend/vendor/object_store/src/client/retry.rs
Vendored
+757
View File
@@ -0,0 +1,757 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! A shared HTTP client implementation incorporating retries
use crate::client::backoff::{Backoff, BackoffConfig};
use crate::client::builder::HttpRequestBuilder;
use crate::client::{HttpClient, HttpError, HttpErrorKind, HttpRequest, HttpResponse};
use crate::PutPayload;
use futures::future::BoxFuture;
use http::{Method, Uri};
use reqwest::header::LOCATION;
use reqwest::StatusCode;
#[cfg(not(all(target_arch = "wasm32", target_os = "unknown")))]
use std::time::{Duration, Instant};
use tracing::info;
#[cfg(all(target_arch = "wasm32", target_os = "unknown"))]
use web_time::{Duration, Instant};
/// Retry request error
#[derive(Debug, thiserror::Error)]
pub struct RetryError {
method: Method,
uri: Option<Uri>,
retries: usize,
max_retries: usize,
elapsed: Duration,
retry_timeout: Duration,
inner: RequestError,
}
impl std::fmt::Display for RetryError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "Error performing {} ", self.method)?;
match &self.uri {
Some(uri) => write!(f, "{uri} ")?,
None => write!(f, "REDACTED ")?,
}
write!(f, "in {:?}", self.elapsed)?;
if self.retries != 0 {
write!(
f,
", after {} retries, max_retries: {}, retry_timeout: {:?} ",
self.retries, self.max_retries, self.retry_timeout
)?;
}
write!(f, " - {}", self.inner)
}
}
/// Context of the retry loop
struct RetryContext {
method: Method,
uri: Option<Uri>,
retries: usize,
max_retries: usize,
start: Instant,
retry_timeout: Duration,
}
impl RetryContext {
fn err(self, error: RequestError) -> RetryError {
RetryError {
uri: self.uri,
method: self.method,
retries: self.retries,
max_retries: self.max_retries,
elapsed: self.start.elapsed(),
retry_timeout: self.retry_timeout,
inner: error,
}
}
fn exhausted(&self) -> bool {
self.retries == self.max_retries || self.start.elapsed() > self.retry_timeout
}
}
/// The reason a request failed
#[derive(Debug, thiserror::Error)]
pub enum RequestError {
#[error("Received redirect without LOCATION, this normally indicates an incorrectly configured region"
)]
BareRedirect,
#[error("Server returned non-2xx status code: {status}: {}", body.as_deref().unwrap_or(""))]
Status {
status: StatusCode,
body: Option<String>,
},
#[error("Server returned error response: {body}")]
Response { status: StatusCode, body: String },
#[error(transparent)]
Http(#[from] HttpError),
}
impl RetryError {
/// Returns the underlying [`RequestError`]
pub fn inner(&self) -> &RequestError {
&self.inner
}
/// Returns the status code associated with this error if any
pub fn status(&self) -> Option<StatusCode> {
match &self.inner {
RequestError::Status { status, .. } | RequestError::Response { status, .. } => {
Some(*status)
}
RequestError::BareRedirect | RequestError::Http(_) => None,
}
}
/// Returns the error body if any
pub fn body(&self) -> Option<&str> {
match &self.inner {
RequestError::Status { body, .. } => body.as_deref(),
RequestError::Response { body, .. } => Some(body),
RequestError::BareRedirect | RequestError::Http(_) => None,
}
}
pub fn error(self, store: &'static str, path: String) -> crate::Error {
match self.status() {
Some(StatusCode::NOT_FOUND) => crate::Error::NotFound {
path,
source: Box::new(self),
},
Some(StatusCode::NOT_MODIFIED) => crate::Error::NotModified {
path,
source: Box::new(self),
},
Some(StatusCode::PRECONDITION_FAILED) => crate::Error::Precondition {
path,
source: Box::new(self),
},
Some(StatusCode::CONFLICT) => crate::Error::AlreadyExists {
path,
source: Box::new(self),
},
Some(StatusCode::FORBIDDEN) => crate::Error::PermissionDenied {
path,
source: Box::new(self),
},
Some(StatusCode::UNAUTHORIZED) => crate::Error::Unauthenticated {
path,
source: Box::new(self),
},
_ => crate::Error::Generic {
store,
source: Box::new(self),
},
}
}
}
impl From<RetryError> for std::io::Error {
fn from(err: RetryError) -> Self {
use std::io::ErrorKind;
let kind = match err.status() {
Some(StatusCode::NOT_FOUND) => ErrorKind::NotFound,
Some(StatusCode::BAD_REQUEST) => ErrorKind::InvalidInput,
Some(StatusCode::UNAUTHORIZED) | Some(StatusCode::FORBIDDEN) => {
ErrorKind::PermissionDenied
}
_ => match &err.inner {
RequestError::Http(h) => match h.kind() {
HttpErrorKind::Timeout => ErrorKind::TimedOut,
HttpErrorKind::Connect => ErrorKind::NotConnected,
_ => ErrorKind::Other,
},
_ => ErrorKind::Other,
},
};
Self::new(kind, err)
}
}
pub(crate) type Result<T, E = RetryError> = std::result::Result<T, E>;
/// 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`] for more information
///
/// [safe]: https://datatracker.ietf.org/doc/html/rfc7231#section-4.2.1
#[derive(Debug, Clone)]
pub struct RetryConfig {
/// The backoff configuration
pub backoff: BackoffConfig,
/// The maximum number of times to retry a request
///
/// Set to 0 to disable retries
pub max_retries: usize,
/// 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
pub retry_timeout: Duration,
}
impl Default for RetryConfig {
fn default() -> Self {
Self {
backoff: Default::default(),
max_retries: 10,
retry_timeout: Duration::from_secs(3 * 60),
}
}
}
fn body_contains_error(response_body: &str) -> bool {
response_body.contains("InternalError") || response_body.contains("SlowDown")
}
pub(crate) struct RetryableRequest {
client: HttpClient,
request: HttpRequest,
max_retries: usize,
retry_timeout: Duration,
backoff: Backoff,
sensitive: bool,
idempotent: Option<bool>,
retry_on_conflict: bool,
payload: Option<PutPayload>,
retry_error_body: bool,
}
impl RetryableRequest {
/// Set whether this request is idempotent
///
/// An idempotent request will be retried on timeout even if the request
/// method is not [safe](https://datatracker.ietf.org/doc/html/rfc7231#section-4.2.1)
pub(crate) fn idempotent(self, idempotent: bool) -> Self {
Self {
idempotent: Some(idempotent),
..self
}
}
/// Set whether this request should be retried on a 409 Conflict response.
#[cfg(feature = "aws")]
pub(crate) fn retry_on_conflict(self, retry_on_conflict: bool) -> Self {
Self {
retry_on_conflict,
..self
}
}
/// Set whether this request contains sensitive data
///
/// This will avoid printing out the URL in error messages
#[allow(unused)]
pub(crate) fn sensitive(self, sensitive: bool) -> Self {
Self { sensitive, ..self }
}
/// Provide a [`PutPayload`]
pub(crate) fn payload(self, payload: Option<PutPayload>) -> Self {
Self { payload, ..self }
}
#[allow(unused)]
pub(crate) fn retry_error_body(self, retry_error_body: bool) -> Self {
Self {
retry_error_body,
..self
}
}
pub(crate) async fn send(self) -> Result<HttpResponse> {
let mut ctx = RetryContext {
retries: 0,
uri: (!self.sensitive).then(|| self.request.uri().clone()),
method: self.request.method().clone(),
max_retries: self.max_retries,
start: Instant::now(),
retry_timeout: self.retry_timeout,
};
let mut backoff = self.backoff;
let is_idempotent = self
.idempotent
.unwrap_or_else(|| self.request.method().is_safe());
loop {
let mut request = self.request.clone();
if let Some(payload) = &self.payload {
*request.body_mut() = payload.clone().into();
}
match self.client.execute(request).await {
Ok(r) => {
let status = r.status();
if status.is_success() {
// For certain S3 requests, 200 response may contain `InternalError` or
// `SlowDown` in the message. These responses should be handled similarly
// to r5xx errors.
// More info here: https://repost.aws/knowledge-center/s3-resolve-200-internalerror
if !self.retry_error_body {
return Ok(r);
}
let (parts, body) = r.into_parts();
let body = match body.text().await {
Ok(body) => body,
Err(e) => return Err(ctx.err(RequestError::Http(e))),
};
if !body_contains_error(&body) {
// Success response and no error, clone and return response
return Ok(HttpResponse::from_parts(parts, body.into()));
} else {
// Retry as if this was a 5xx response
if ctx.exhausted() {
return Err(ctx.err(RequestError::Response { body, status }));
}
let sleep = backoff.next();
ctx.retries += 1;
info!(
"Encountered a response status of {} but body contains Error, backing off for {} seconds, retry {} of {}",
status,
sleep.as_secs_f32(),
ctx.retries,
ctx.max_retries,
);
tokio::time::sleep(sleep).await;
}
} else if status == StatusCode::NOT_MODIFIED {
return Err(ctx.err(RequestError::Status { status, body: None }));
} else if status.is_redirection() {
let is_bare_redirect = !r.headers().contains_key(LOCATION);
return match is_bare_redirect {
true => Err(ctx.err(RequestError::BareRedirect)),
false => Err(ctx.err(RequestError::Status {
body: None,
status: r.status(),
})),
};
} else {
let status = r.status();
if ctx.exhausted()
|| !(status.is_server_error()
|| (self.retry_on_conflict && status == StatusCode::CONFLICT))
{
let source = match status.is_client_error() {
true => match r.into_body().text().await {
Ok(body) => RequestError::Status {
status,
body: Some(body),
},
Err(e) => RequestError::Http(e),
},
false => RequestError::Status { status, body: None },
};
return Err(ctx.err(source));
};
let sleep = backoff.next();
ctx.retries += 1;
info!(
"Encountered server error, backing off for {} seconds, retry {} of {}",
sleep.as_secs_f32(),
ctx.retries,
ctx.max_retries,
);
tokio::time::sleep(sleep).await;
}
}
Err(e) => {
// let e = sanitize_err(e);
let do_retry = match e.kind() {
HttpErrorKind::Connect | HttpErrorKind::Request => true, // Request not sent, can retry
HttpErrorKind::Timeout | HttpErrorKind::Interrupted => is_idempotent,
HttpErrorKind::Unknown | HttpErrorKind::Decode => false,
};
if ctx.retries == ctx.max_retries
|| ctx.start.elapsed() > ctx.retry_timeout
|| !do_retry
{
return Err(ctx.err(RequestError::Http(e)));
}
let sleep = backoff.next();
ctx.retries += 1;
info!(
"Encountered transport error backing off for {} seconds, retry {} of {}: {}",
sleep.as_secs_f32(),
ctx.retries,
ctx.max_retries,
e,
);
tokio::time::sleep(sleep).await;
}
}
}
}
}
pub(crate) trait RetryExt {
/// Return a [`RetryableRequest`]
fn retryable(self, config: &RetryConfig) -> RetryableRequest;
/// Dispatch a request with the given retry configuration
///
/// # Panic
///
/// This will panic if the request body is a stream
fn send_retry(self, config: &RetryConfig) -> BoxFuture<'static, Result<HttpResponse>>;
}
impl RetryExt for HttpRequestBuilder {
fn retryable(self, config: &RetryConfig) -> RetryableRequest {
let (client, request) = self.into_parts();
let request = request.expect("request must be valid");
RetryableRequest {
client,
request,
max_retries: config.max_retries,
retry_timeout: config.retry_timeout,
backoff: Backoff::new(&config.backoff),
idempotent: None,
payload: None,
sensitive: false,
retry_on_conflict: false,
retry_error_body: false,
}
}
fn send_retry(self, config: &RetryConfig) -> BoxFuture<'static, Result<HttpResponse>> {
let request = self.retryable(config);
Box::pin(async move { request.send().await })
}
}
#[cfg(not(target_arch = "wasm32"))]
#[cfg(test)]
mod tests {
use crate::client::mock_server::MockServer;
use crate::client::retry::{body_contains_error, RequestError, RetryExt};
use crate::client::HttpClient;
use crate::RetryConfig;
use hyper::header::LOCATION;
use hyper::Response;
use reqwest::{Client, Method, StatusCode};
use std::time::Duration;
#[test]
fn test_body_contains_error() {
// Example error message provided by https://repost.aws/knowledge-center/s3-resolve-200-internalerror
let error_response = "AmazonS3Exception: We encountered an internal error. Please try again. (Service: Amazon S3; Status Code: 200; Error Code: InternalError; Request ID: 0EXAMPLE9AAEB265)";
assert!(body_contains_error(error_response));
let error_response_2 = "<?xml version=\"1.0\" encoding=\"UTF-8\"?><Error><Code>SlowDown</Code><Message>Please reduce your request rate.</Message><RequestId>123</RequestId><HostId>456</HostId></Error>";
assert!(body_contains_error(error_response_2));
// Example success response from https://docs.aws.amazon.com/AmazonS3/latest/API/API_CopyObject.html
let success_response = "<CopyObjectResult><LastModified>2009-10-12T17:50:30.000Z</LastModified><ETag>\"9b2cf535f27731c974343645a3985328\"</ETag></CopyObjectResult>";
assert!(!body_contains_error(success_response));
}
#[tokio::test]
async fn test_retry() {
let mock = MockServer::new().await;
let retry = RetryConfig {
backoff: Default::default(),
max_retries: 2,
retry_timeout: Duration::from_secs(1000),
};
let client = HttpClient::new(
Client::builder()
.timeout(Duration::from_millis(100))
.build()
.unwrap(),
);
let do_request = || client.request(Method::GET, mock.url()).send_retry(&retry);
// Simple request should work
let r = do_request().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
// Returns client errors immediately with status message
mock.push(
Response::builder()
.status(StatusCode::BAD_REQUEST)
.body("cupcakes".to_string())
.unwrap(),
);
let e = do_request().await.unwrap_err();
assert_eq!(e.status().unwrap(), StatusCode::BAD_REQUEST);
assert_eq!(e.body(), Some("cupcakes"));
assert_eq!(
e.inner().to_string(),
"Server returned non-2xx status code: 400 Bad Request: cupcakes"
);
// Handles client errors with no payload
mock.push(
Response::builder()
.status(StatusCode::BAD_REQUEST)
.body("NAUGHTY NAUGHTY".to_string())
.unwrap(),
);
let e = do_request().await.unwrap_err();
assert_eq!(e.status().unwrap(), StatusCode::BAD_REQUEST);
assert_eq!(e.body(), Some("NAUGHTY NAUGHTY"));
assert_eq!(
e.inner().to_string(),
"Server returned non-2xx status code: 400 Bad Request: NAUGHTY NAUGHTY"
);
// Should retry server error request
mock.push(
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.body(String::new())
.unwrap(),
);
let r = do_request().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
// Accepts 204 status code
mock.push(
Response::builder()
.status(StatusCode::NO_CONTENT)
.body(String::new())
.unwrap(),
);
let r = do_request().await.unwrap();
assert_eq!(r.status(), StatusCode::NO_CONTENT);
// Follows 402 redirects
mock.push(
Response::builder()
.status(StatusCode::FOUND)
.header(LOCATION, "/foo")
.body(String::new())
.unwrap(),
);
let r = do_request().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
// Follows 401 redirects
mock.push(
Response::builder()
.status(StatusCode::FOUND)
.header(LOCATION, "/bar")
.body(String::new())
.unwrap(),
);
let r = do_request().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
// Handles redirect loop
for _ in 0..10 {
mock.push(
Response::builder()
.status(StatusCode::FOUND)
.header(LOCATION, "/bar")
.body(String::new())
.unwrap(),
);
}
let e = do_request().await.unwrap_err().to_string();
assert!(e.contains("error following redirect"), "{}", e);
// Handles redirect missing location
mock.push(
Response::builder()
.status(StatusCode::FOUND)
.body(String::new())
.unwrap(),
);
let e = do_request().await.unwrap_err();
assert!(matches!(e.inner, RequestError::BareRedirect));
assert_eq!(e.inner().to_string(), "Received redirect without LOCATION, this normally indicates an incorrectly configured region");
// Gives up after the retrying the specified number of times
for _ in 0..=retry.max_retries {
mock.push(
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.body("ignored".to_string())
.unwrap(),
);
}
let e = do_request().await.unwrap_err().to_string();
assert!(
e.contains(" after 2 retries, max_retries: 2, retry_timeout: 1000s - Server returned non-2xx status code: 502 Bad Gateway"),
"{e}"
);
// Panic results in an incomplete message error in the client
mock.push_fn(|_| panic!());
let r = do_request().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
// Gives up after retrying multiple panics
for _ in 0..=retry.max_retries {
mock.push_fn(|_| panic!());
}
let e = do_request().await.unwrap_err().to_string();
assert!(
e.contains("after 2 retries, max_retries: 2, retry_timeout: 1000s - HTTP error: error sending request"),
"{e}"
);
// Retries on client timeout
mock.push_async_fn(|_| async move {
tokio::time::sleep(Duration::from_secs(10)).await;
panic!()
});
do_request().await.unwrap();
// Does not retry PUT request
mock.push_async_fn(|_| async move {
tokio::time::sleep(Duration::from_secs(10)).await;
panic!()
});
let res = client.request(Method::PUT, mock.url()).send_retry(&retry);
let e = res.await.unwrap_err().to_string();
assert!(
!e.contains("retries") && e.contains("error sending request"),
"{e}"
);
let url = format!("{}/SENSITIVE", mock.url());
for _ in 0..=retry.max_retries {
mock.push(
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.body("ignored".to_string())
.unwrap(),
);
}
let res = client.request(Method::GET, url).send_retry(&retry).await;
let err = res.unwrap_err().to_string();
assert!(err.contains("SENSITIVE"), "{err}");
let url = format!("{}/SENSITIVE", mock.url());
for _ in 0..=retry.max_retries {
mock.push(
Response::builder()
.status(StatusCode::BAD_GATEWAY)
.body("ignored".to_string())
.unwrap(),
);
}
// Sensitive requests should strip URL from error
let req = client
.request(Method::GET, &url)
.retryable(&retry)
.sensitive(true);
let err = req.send().await.unwrap_err().to_string();
assert!(!err.contains("SENSITIVE"), "{err}");
for _ in 0..=retry.max_retries {
mock.push_fn(|_| panic!());
}
let req = client
.request(Method::GET, &url)
.retryable(&retry)
.sensitive(true);
let err = req.send().await.unwrap_err().to_string();
assert!(!err.contains("SENSITIVE"), "{err}");
// Success response with error in body is retried
mock.push(
Response::builder()
.status(StatusCode::OK)
.body("InternalError".to_string())
.unwrap(),
);
let req = client
.request(Method::PUT, &url)
.retryable(&retry)
.idempotent(true)
.retry_error_body(true);
let r = req.send().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
// Response with InternalError should have been retried
let b = r.into_body().text().await.unwrap();
assert!(!b.contains("InternalError"));
// Should not retry success response with no error in body
mock.push(
Response::builder()
.status(StatusCode::OK)
.body("success".to_string())
.unwrap(),
);
let req = client
.request(Method::PUT, &url)
.retryable(&retry)
.idempotent(true)
.retry_error_body(true);
let r = req.send().await.unwrap();
assert_eq!(r.status(), StatusCode::OK);
let b = r.into_body().text().await.unwrap();
assert!(b.contains("success"));
// Shutdown
mock.shutdown().await
}
}
backend/vendor/object_store/src/client/s3.rs
Vendored
+157
View File
@@ -0,0 +1,157 @@
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! The list and multipart API used by both GCS and S3
use crate::multipart::PartId;
use crate::path::Path;
use crate::{ListResult, ObjectMeta, Result};
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct ListResponse {
#[serde(default)]
pub contents: Vec<ListContents>,
#[serde(default)]
pub common_prefixes: Vec<ListPrefix>,
#[serde(default)]
pub next_continuation_token: Option<String>,
}
impl TryFrom<ListResponse> for ListResult {
type Error = crate::Error;
fn try_from(value: ListResponse) -> Result<Self> {
let common_prefixes = value
.common_prefixes
.into_iter()
.map(|x| Ok(Path::parse(x.prefix)?))
.collect::<Result<_>>()?;
let objects = value
.contents
.into_iter()
.map(TryFrom::try_from)
.collect::<Result<_>>()?;
Ok(Self {
common_prefixes,
objects,
})
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct ListPrefix {
pub prefix: String,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub struct ListContents {
pub key: String,
pub size: u64,
pub last_modified: DateTime<Utc>,
#[serde(rename = "ETag")]
pub e_tag: Option<String>,
}
impl TryFrom<ListContents> for ObjectMeta {
type Error = crate::Error;
fn try_from(value: ListContents) -> Result<Self> {
Ok(Self {
location: Path::parse(value.key)?,
last_modified: value.last_modified,
size: value.size,
e_tag: value.e_tag,
version: None,
})
}
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(crate) struct InitiateMultipartUploadResult {
pub upload_id: String,
}
#[cfg(feature = "aws")]
#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(crate) struct CopyPartResult {
#[serde(rename = "ETag")]
pub e_tag: String,
}
#[derive(Debug, Serialize)]
#[serde(rename_all = "PascalCase")]
pub(crate) struct CompleteMultipartUpload {
pub part: Vec<MultipartPart>,
}
#[derive(Serialize, Deserialize)]
pub(crate) struct PartMetadata {
pub e_tag: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub checksum_sha256: Option<String>,
}
impl From<Vec<PartId>> for CompleteMultipartUpload {
fn from(value: Vec<PartId>) -> Self {
let part = value
.into_iter()
.enumerate()
.map(|(part_idx, part)| {
let md = match quick_xml::de::from_str::<PartMetadata>(&part.content_id) {
Ok(md) => md,
// fallback to old way
Err(_) => PartMetadata {
e_tag: part.content_id.clone(),
checksum_sha256: None,
},
};
MultipartPart {
e_tag: md.e_tag,
part_number: part_idx + 1,
checksum_sha256: md.checksum_sha256,
}
})
.collect();
Self { part }
}
}
#[derive(Debug, Serialize)]
pub(crate) struct MultipartPart {
#[serde(rename = "ETag")]
pub e_tag: String,
#[serde(rename = "PartNumber")]
pub part_number: usize,
#[serde(rename = "ChecksumSHA256")]
#[serde(skip_serializing_if = "Option::is_none")]
pub checksum_sha256: Option<String>,
}
#[derive(Debug, Deserialize)]
#[serde(rename_all = "PascalCase")]
pub(crate) struct CompleteMultipartUploadResult {
#[serde(rename = "ETag")]
pub e_tag: String,
}
backend/vendor/object_store/src/client/token.rs
Vendored
+155
View File
@@ -0,0 +1,155 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::future::Future;
use std::time::{Duration, Instant};
use tokio::sync::Mutex;
/// A temporary authentication token with an associated expiry
#[derive(Debug, Clone)]
pub(crate) struct TemporaryToken<T> {
/// 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<Instant>,
}
/// Provides [`TokenCache::get_or_insert_with`] which can be used to cache a
/// [`TemporaryToken`] based on its expiry
#[derive(Debug)]
pub(crate) struct TokenCache<T> {
cache: Mutex<Option<(TemporaryToken<T>, Instant)>>,
min_ttl: Duration,
fetch_backoff: Duration,
}
impl<T> Default for TokenCache<T> {
fn default() -> Self {
Self {
cache: Default::default(),
min_ttl: Duration::from_secs(300),
// How long to wait before re-attempting a token fetch after receiving one that
// is still within the min-ttl
fetch_backoff: Duration::from_millis(100),
}
}
}
impl<T: Clone + Send> TokenCache<T> {
/// Override the minimum remaining TTL for a cached token to be used
#[cfg(any(feature = "aws", feature = "gcp"))]
pub(crate) fn with_min_ttl(self, min_ttl: Duration) -> Self {
Self { min_ttl, ..self }
}
pub(crate) async fn get_or_insert_with<F, Fut, E>(&self, f: F) -> Result<T, E>
where
F: FnOnce() -> Fut + Send,
Fut: Future<Output = Result<TemporaryToken<T>, E>> + Send,
{
let now = Instant::now();
let mut locked = self.cache.lock().await;
if let Some((cached, fetched_at)) = locked.as_ref() {
match cached.expiry {
Some(ttl) => {
if ttl.checked_duration_since(now).unwrap_or_default() > 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
(fetched_at.elapsed() < self.fetch_backoff && ttl.checked_duration_since(now).is_some())
{
return Ok(cached.token.clone());
}
}
None => return Ok(cached.token.clone()),
}
}
let cached = f().await?;
let token = cached.token.clone();
*locked = Some((cached, Instant::now()));
Ok(token)
}
}
#[cfg(test)]
mod test {
use crate::client::token::{TemporaryToken, TokenCache};
use std::sync::atomic::{AtomicU32, Ordering};
use std::time::{Duration, Instant};
// Helper function to create a token with a specific expiry duration from now
fn create_token(expiry_duration: Option<Duration>) -> TemporaryToken<String> {
TemporaryToken {
token: "test_token".to_string(),
expiry: expiry_duration.map(|d| Instant::now() + d),
}
}
#[tokio::test]
async fn test_expired_token_is_refreshed() {
let cache = TokenCache::default();
static COUNTER: AtomicU32 = AtomicU32::new(0);
async fn get_token() -> Result<TemporaryToken<String>, String> {
COUNTER.fetch_add(1, Ordering::SeqCst);
Ok::<_, String>(create_token(Some(Duration::from_secs(0))))
}
// Should fetch initial token
let _ = cache.get_or_insert_with(get_token).await.unwrap();
assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
tokio::time::sleep(Duration::from_millis(2)).await;
// Token is expired, so should fetch again
let _ = cache.get_or_insert_with(get_token).await.unwrap();
assert_eq!(COUNTER.load(Ordering::SeqCst), 2);
}
#[tokio::test]
async fn test_min_ttl_causes_refresh() {
let cache = TokenCache {
cache: Default::default(),
min_ttl: Duration::from_secs(1),
fetch_backoff: Duration::from_millis(1),
};
static COUNTER: AtomicU32 = AtomicU32::new(0);
async fn get_token() -> Result<TemporaryToken<String>, String> {
COUNTER.fetch_add(1, Ordering::SeqCst);
Ok::<_, String>(create_token(Some(Duration::from_millis(100))))
}
// Initial fetch
let _ = cache.get_or_insert_with(get_token).await.unwrap();
assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
// Should not fetch again since not expired and within fetch_backoff
let _ = cache.get_or_insert_with(get_token).await.unwrap();
assert_eq!(COUNTER.load(Ordering::SeqCst), 1);
tokio::time::sleep(Duration::from_millis(2)).await;
// Should fetch, since we've passed fetch_backoff
let _ = cache.get_or_insert_with(get_token).await.unwrap();
assert_eq!(COUNTER.load(Ordering::SeqCst), 2);
}
}
backend/vendor/object_store/src/config.rs
Vendored
+143
View File
@@ -0,0 +1,143 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::fmt::{Debug, Display, Formatter};
use std::str::FromStr;
use std::time::Duration;
use humantime::{format_duration, parse_duration};
use reqwest::header::HeaderValue;
use crate::{Error, Result};
/// Provides deferred parsing of a value
///
/// This allows builders to defer fallibility to build
#[derive(Debug, Clone)]
pub(crate) enum ConfigValue<T> {
Parsed(T),
Deferred(String),
}
impl<T: Display> Display for ConfigValue<T> {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::Parsed(v) => write!(f, "{v}"),
Self::Deferred(v) => write!(f, "{v}"),
}
}
}
impl<T> From<T> for ConfigValue<T> {
fn from(value: T) -> Self {
Self::Parsed(value)
}
}
impl<T: Parse + Clone> ConfigValue<T> {
pub(crate) fn parse(&mut self, v: impl Into<String>) {
*self = Self::Deferred(v.into())
}
pub(crate) fn get(&self) -> Result<T> {
match self {
Self::Parsed(v) => Ok(v.clone()),
Self::Deferred(v) => T::parse(v),
}
}
}
impl<T: Default> Default for ConfigValue<T> {
fn default() -> Self {
Self::Parsed(T::default())
}
}
/// A value that can be stored in [`ConfigValue`]
pub(crate) trait Parse: Sized {
fn parse(v: &str) -> Result<Self>;
}
impl Parse for bool {
fn parse(v: &str) -> Result<Self> {
let lower = v.to_ascii_lowercase();
match lower.as_str() {
"1" | "true" | "on" | "yes" | "y" => Ok(true),
"0" | "false" | "off" | "no" | "n" => Ok(false),
_ => Err(Error::Generic {
store: "Config",
source: format!("failed to parse \"{v}\" as boolean").into(),
}),
}
}
}
impl Parse for Duration {
fn parse(v: &str) -> Result<Self> {
parse_duration(v).map_err(|_| Error::Generic {
store: "Config",
source: format!("failed to parse \"{v}\" as Duration").into(),
})
}
}
impl Parse for usize {
fn parse(v: &str) -> Result<Self> {
Self::from_str(v).map_err(|_| Error::Generic {
store: "Config",
source: format!("failed to parse \"{v}\" as usize").into(),
})
}
}
impl Parse for u32 {
fn parse(v: &str) -> Result<Self> {
Self::from_str(v).map_err(|_| Error::Generic {
store: "Config",
source: format!("failed to parse \"{v}\" as u32").into(),
})
}
}
impl Parse for HeaderValue {
fn parse(v: &str) -> Result<Self> {
Self::from_str(v).map_err(|_| Error::Generic {
store: "Config",
source: format!("failed to parse \"{v}\" as HeaderValue").into(),
})
}
}
pub(crate) fn fmt_duration(duration: &ConfigValue<Duration>) -> String {
match duration {
ConfigValue::Parsed(v) => format_duration(*v).to_string(),
ConfigValue::Deferred(v) => v.clone(),
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::time::Duration;
#[test]
fn test_parse_duration() {
let duration = Duration::from_secs(60);
assert_eq!(Duration::parse("60 seconds").unwrap(), duration);
assert_eq!(Duration::parse("60 s").unwrap(), duration);
assert_eq!(Duration::parse("60s").unwrap(), duration)
}
}
backend/vendor/object_store/src/delimited.rs
Vendored
+272
View File
@@ -0,0 +1,272 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Utility for streaming newline delimited files from object storage
use std::collections::VecDeque;
use bytes::Bytes;
use futures::{Stream, StreamExt};
use super::Result;
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("encountered unterminated string")]
UnterminatedString,
#[error("encountered trailing escape character")]
TrailingEscape,
}
impl From<Error> for super::Error {
fn from(err: Error) -> Self {
Self::Generic {
store: "LineDelimiter",
source: Box::new(err),
}
}
}
/// The ASCII encoding of `"`
const QUOTE: u8 = b'"';
/// The ASCII encoding of `\n`
const NEWLINE: u8 = b'\n';
/// The ASCII encoding of `\`
const ESCAPE: u8 = b'\\';
/// [`LineDelimiter`] is provided with a stream of [`Bytes`] and returns an iterator
/// of [`Bytes`] containing a whole number of new line delimited records
#[derive(Debug, Default)]
struct LineDelimiter {
/// Complete chunks of [`Bytes`]
complete: VecDeque<Bytes>,
/// Remainder bytes that form the next record
remainder: Vec<u8>,
/// True if the last character was the escape character
is_escape: bool,
/// True if currently processing a quoted string
is_quote: bool,
}
impl LineDelimiter {
/// Creates a new [`LineDelimiter`] with the provided delimiter
fn new() -> Self {
Self::default()
}
/// Adds the next set of [`Bytes`]
fn push(&mut self, val: impl Into<Bytes>) {
let val: Bytes = val.into();
let is_escape = &mut self.is_escape;
let is_quote = &mut self.is_quote;
let mut record_ends = val.iter().enumerate().filter_map(|(idx, v)| {
if *is_escape {
*is_escape = false;
None
} else if *v == ESCAPE {
*is_escape = true;
None
} else if *v == QUOTE {
*is_quote = !*is_quote;
None
} else if *is_quote {
None
} else {
(*v == NEWLINE).then_some(idx + 1)
}
});
let start_offset = match self.remainder.is_empty() {
true => 0,
false => match record_ends.next() {
Some(idx) => {
self.remainder.extend_from_slice(&val[0..idx]);
self.complete
.push_back(Bytes::from(std::mem::take(&mut self.remainder)));
idx
}
None => {
self.remainder.extend_from_slice(&val);
return;
}
},
};
let end_offset = record_ends.next_back().unwrap_or(start_offset);
if start_offset != end_offset {
self.complete.push_back(val.slice(start_offset..end_offset));
}
if end_offset != val.len() {
self.remainder.extend_from_slice(&val[end_offset..])
}
}
/// Marks the end of the stream, delimiting any remaining bytes
///
/// Returns `true` if there is no remaining data to be read
fn finish(&mut self) -> Result<bool> {
if !self.remainder.is_empty() {
if self.is_quote {
Err(Error::UnterminatedString)?;
}
if self.is_escape {
Err(Error::TrailingEscape)?;
}
self.complete
.push_back(Bytes::from(std::mem::take(&mut self.remainder)))
}
Ok(self.complete.is_empty())
}
}
impl Iterator for LineDelimiter {
type Item = Bytes;
fn next(&mut self) -> Option<Self::Item> {
self.complete.pop_front()
}
}
/// Given a [`Stream`] of [`Bytes`] returns a [`Stream`] where each
/// yielded [`Bytes`] contains a whole number of new line delimited records
/// accounting for `\` style escapes and `"` quotes
pub fn newline_delimited_stream<S>(s: S) -> impl Stream<Item = Result<Bytes>>
where
S: Stream<Item = Result<Bytes>> + Unpin,
{
let delimiter = LineDelimiter::new();
futures::stream::unfold(
(s, delimiter, false),
|(mut s, mut delimiter, mut exhausted)| async move {
loop {
if let Some(next) = delimiter.next() {
return Some((Ok(next), (s, delimiter, exhausted)));
} else if exhausted {
return None;
}
match s.next().await {
Some(Ok(bytes)) => delimiter.push(bytes),
Some(Err(e)) => return Some((Err(e), (s, delimiter, exhausted))),
None => {
exhausted = true;
match delimiter.finish() {
Ok(true) => return None,
Ok(false) => continue,
Err(e) => return Some((Err(e), (s, delimiter, exhausted))),
}
}
}
}
},
)
}
#[cfg(test)]
mod tests {
use futures::stream::{BoxStream, TryStreamExt};
use super::*;
#[test]
fn test_delimiter() {
let mut delimiter = LineDelimiter::new();
delimiter.push("hello\nworld");
delimiter.push("\n\n");
assert_eq!(delimiter.next().unwrap(), Bytes::from("hello\n"));
assert_eq!(delimiter.next().unwrap(), Bytes::from("world\n"));
assert_eq!(delimiter.next().unwrap(), Bytes::from("\n"));
assert!(delimiter.next().is_none());
}
#[test]
fn test_delimiter_escaped() {
let mut delimiter = LineDelimiter::new();
delimiter.push("");
delimiter.push("fo\\\n\"foo");
delimiter.push("bo\n\"bar\n");
delimiter.push("\"he");
delimiter.push("llo\"\n");
assert_eq!(
delimiter.next().unwrap(),
Bytes::from("fo\\\n\"foobo\n\"bar\n")
);
assert_eq!(delimiter.next().unwrap(), Bytes::from("\"hello\"\n"));
assert!(delimiter.next().is_none());
// Verify can push further data
delimiter.push("\"foo\nbar\",\"fiz\\\"inner\\\"\"\nhello");
assert!(!delimiter.finish().unwrap());
assert_eq!(
delimiter.next().unwrap(),
Bytes::from("\"foo\nbar\",\"fiz\\\"inner\\\"\"\n")
);
assert_eq!(delimiter.next().unwrap(), Bytes::from("hello"));
assert!(delimiter.finish().unwrap());
assert!(delimiter.next().is_none());
}
#[tokio::test]
async fn test_delimiter_stream() {
let input = vec!["hello\nworld\nbin", "go\ncup", "cakes"];
let input_stream = futures::stream::iter(input.into_iter().map(|s| Ok(Bytes::from(s))));
let stream = newline_delimited_stream(input_stream);
let results: Vec<_> = stream.try_collect().await.unwrap();
assert_eq!(
results,
vec![
Bytes::from("hello\nworld\n"),
Bytes::from("bingo\n"),
Bytes::from("cupcakes")
]
)
}
#[tokio::test]
async fn test_delimiter_unfold_stream() {
let input_stream: BoxStream<'static, Result<Bytes>> = futures::stream::unfold(
VecDeque::from(["hello\nworld\nbin", "go\ncup", "cakes"]),
|mut input| async move {
if !input.is_empty() {
Some((Ok(Bytes::from(input.pop_front().unwrap())), input))
} else {
None
}
},
)
.boxed();
let stream = newline_delimited_stream(input_stream);
let results: Vec<_> = stream.try_collect().await.unwrap();
assert_eq!(
results,
vec![
Bytes::from("hello\nworld\n"),
Bytes::from("bingo\n"),
Bytes::from("cupcakes")
]
)
}
}
backend/vendor/object_store/src/gcp/builder.rs
Vendored
+738
View File
@@ -0,0 +1,738 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::{http_connector, HttpConnector, TokenCredentialProvider};
use crate::config::ConfigValue;
use crate::gcp::client::{GoogleCloudStorageClient, GoogleCloudStorageConfig};
use crate::gcp::credential::{
ApplicationDefaultCredentials, InstanceCredentialProvider, ServiceAccountCredentials,
DEFAULT_GCS_BASE_URL,
};
use crate::gcp::{
credential, GcpCredential, GcpCredentialProvider, GcpSigningCredential,
GcpSigningCredentialProvider, GoogleCloudStorage, STORE,
};
use crate::{ClientConfigKey, ClientOptions, Result, RetryConfig, StaticCredentialProvider};
use serde::{Deserialize, Serialize};
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;
use url::Url;
use super::credential::{AuthorizedUserSigningCredentials, InstanceSigningCredentialProvider};
const TOKEN_MIN_TTL: Duration = Duration::from_secs(4 * 60);
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("Missing bucket name")]
MissingBucketName {},
#[error("One of service account path or service account key may be provided.")]
ServiceAccountPathAndKeyProvided,
#[error("Unable parse source url. Url: {}, Error: {}", url, source)]
UnableToParseUrl {
source: url::ParseError,
url: String,
},
#[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 },
#[error("Configuration key: '{}' is not known.", key)]
UnknownConfigurationKey { key: String },
#[error("GCP credential error: {}", source)]
Credential { source: credential::Error },
}
impl From<Error> for crate::Error {
fn from(err: Error) -> Self {
match err {
Error::UnknownConfigurationKey { key } => {
Self::UnknownConfigurationKey { store: STORE, key }
}
_ => Self::Generic {
store: STORE,
source: Box::new(err),
},
}
}
}
/// Configure a connection to Google Cloud Storage.
///
/// 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).
///
/// # Example
/// ```
/// # let BUCKET_NAME = "foo";
/// # use object_store::gcp::GoogleCloudStorageBuilder;
/// let gcs = GoogleCloudStorageBuilder::from_env().with_bucket_name(BUCKET_NAME).build();
/// ```
#[derive(Debug, Clone)]
pub struct GoogleCloudStorageBuilder {
/// Bucket name
bucket_name: Option<String>,
/// Url
url: Option<String>,
/// Path to the service account file
service_account_path: Option<String>,
/// The serialized service account key
service_account_key: Option<String>,
/// Path to the application credentials file.
application_credentials_path: Option<String>,
/// Retry config
retry_config: RetryConfig,
/// Client options
client_options: ClientOptions,
/// Credentials
credentials: Option<GcpCredentialProvider>,
/// Skip signing requests
skip_signature: ConfigValue<bool>,
/// Credentials for sign url
signing_credentials: Option<GcpSigningCredentialProvider>,
/// The [`HttpConnector`] to use
http_connector: Option<Arc<dyn HttpConnector>>,
}
/// Configuration keys for [`GoogleCloudStorageBuilder`]
///
/// Configuration via keys can be done via [`GoogleCloudStorageBuilder::with_config`]
///
/// # Example
/// ```
/// # use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey};
/// let builder = GoogleCloudStorageBuilder::new()
/// .with_config("google_service_account".parse().unwrap(), "my-service-account")
/// .with_config(GoogleConfigKey::Bucket, "my-bucket");
/// ```
#[derive(PartialEq, Eq, Hash, Clone, Debug, Copy, Serialize, Deserialize)]
#[non_exhaustive]
pub enum GoogleConfigKey {
/// Path to the service account file
///
/// Supported keys:
/// - `google_service_account`
/// - `service_account`
/// - `google_service_account_path`
/// - `service_account_path`
ServiceAccount,
/// The serialized service account key.
///
/// Supported keys:
/// - `google_service_account_key`
/// - `service_account_key`
ServiceAccountKey,
/// Bucket name
///
/// See [`GoogleCloudStorageBuilder::with_bucket_name`] for details.
///
/// Supported keys:
/// - `google_bucket`
/// - `google_bucket_name`
/// - `bucket`
/// - `bucket_name`
Bucket,
/// Application credentials path
///
/// See [`GoogleCloudStorageBuilder::with_application_credentials`].
ApplicationCredentials,
/// Skip signing request
SkipSignature,
/// Client options
Client(ClientConfigKey),
}
impl AsRef<str> for GoogleConfigKey {
fn as_ref(&self) -> &str {
match self {
Self::ServiceAccount => "google_service_account",
Self::ServiceAccountKey => "google_service_account_key",
Self::Bucket => "google_bucket",
Self::ApplicationCredentials => "google_application_credentials",
Self::SkipSignature => "google_skip_signature",
Self::Client(key) => key.as_ref(),
}
}
}
impl FromStr for GoogleConfigKey {
type Err = crate::Error;
fn from_str(s: &str) -> Result<Self, Self::Err> {
match s {
"google_service_account"
| "service_account"
| "google_service_account_path"
| "service_account_path" => Ok(Self::ServiceAccount),
"google_service_account_key" | "service_account_key" => Ok(Self::ServiceAccountKey),
"google_bucket" | "google_bucket_name" | "bucket" | "bucket_name" => Ok(Self::Bucket),
"google_application_credentials" => Ok(Self::ApplicationCredentials),
"google_skip_signature" | "skip_signature" => Ok(Self::SkipSignature),
_ => match s.strip_prefix("google_").unwrap_or(s).parse() {
Ok(key) => Ok(Self::Client(key)),
Err(_) => Err(Error::UnknownConfigurationKey { key: s.into() }.into()),
},
}
}
}
impl Default for GoogleCloudStorageBuilder {
fn default() -> Self {
Self {
bucket_name: None,
service_account_path: None,
service_account_key: None,
application_credentials_path: None,
retry_config: Default::default(),
client_options: ClientOptions::new().with_allow_http(true),
url: None,
credentials: None,
skip_signature: Default::default(),
signing_credentials: None,
http_connector: None,
}
}
}
impl GoogleCloudStorageBuilder {
/// Create a new [`GoogleCloudStorageBuilder`] with default values.
pub fn new() -> Self {
Default::default()
}
/// Create an instance of [`GoogleCloudStorageBuilder`] with values pre-populated from environment variables.
///
/// Variables extracted from environment:
/// * GOOGLE_SERVICE_ACCOUNT: location of service account file
/// * GOOGLE_SERVICE_ACCOUNT_PATH: (alias) location of service account file
/// * SERVICE_ACCOUNT: (alias) location of service account file
/// * GOOGLE_SERVICE_ACCOUNT_KEY: JSON serialized service account key
/// * GOOGLE_BUCKET: bucket name
/// * GOOGLE_BUCKET_NAME: (alias) bucket name
///
/// # Example
/// ```
/// use object_store::gcp::GoogleCloudStorageBuilder;
///
/// let gcs = GoogleCloudStorageBuilder::from_env()
/// .with_bucket_name("foo")
/// .build();
/// ```
pub fn from_env() -> Self {
let mut builder = Self::default();
if let Ok(service_account_path) = std::env::var("SERVICE_ACCOUNT") {
builder.service_account_path = Some(service_account_path);
}
for (os_key, os_value) in std::env::vars_os() {
if let (Some(key), Some(value)) = (os_key.to_str(), os_value.to_str()) {
if key.starts_with("GOOGLE_") {
if let Ok(config_key) = key.to_ascii_lowercase().parse() {
builder = builder.with_config(config_key, value);
}
}
}
}
builder
}
/// Parse available connection info form a well-known storage URL.
///
/// The supported url schemes are:
///
/// - `gs://<bucket>/<path>`
///
/// Note: Settings derived from the URL will override any others set on this builder
///
/// # Example
/// ```
/// use object_store::gcp::GoogleCloudStorageBuilder;
///
/// let gcs = GoogleCloudStorageBuilder::from_env()
/// .with_url("gs://bucket/path")
/// .build();
/// ```
pub fn with_url(mut self, url: impl Into<String>) -> Self {
self.url = Some(url.into());
self
}
/// Set an option on the builder via a key - value pair.
pub fn with_config(mut self, key: GoogleConfigKey, value: impl Into<String>) -> Self {
match key {
GoogleConfigKey::ServiceAccount => self.service_account_path = Some(value.into()),
GoogleConfigKey::ServiceAccountKey => self.service_account_key = Some(value.into()),
GoogleConfigKey::Bucket => self.bucket_name = Some(value.into()),
GoogleConfigKey::ApplicationCredentials => {
self.application_credentials_path = Some(value.into())
}
GoogleConfigKey::SkipSignature => self.skip_signature.parse(value),
GoogleConfigKey::Client(key) => {
self.client_options = self.client_options.with_config(key, value)
}
};
self
}
/// Get config value via a [`GoogleConfigKey`].
///
/// # Example
/// ```
/// use object_store::gcp::{GoogleCloudStorageBuilder, GoogleConfigKey};
///
/// let builder = GoogleCloudStorageBuilder::from_env()
/// .with_service_account_key("foo");
/// let service_account_key = builder.get_config_value(&GoogleConfigKey::ServiceAccountKey).unwrap_or_default();
/// assert_eq!("foo", &service_account_key);
/// ```
pub fn get_config_value(&self, key: &GoogleConfigKey) -> Option<String> {
match key {
GoogleConfigKey::ServiceAccount => self.service_account_path.clone(),
GoogleConfigKey::ServiceAccountKey => self.service_account_key.clone(),
GoogleConfigKey::Bucket => self.bucket_name.clone(),
GoogleConfigKey::ApplicationCredentials => self.application_credentials_path.clone(),
GoogleConfigKey::SkipSignature => Some(self.skip_signature.to_string()),
GoogleConfigKey::Client(key) => self.client_options.get_config_value(key),
}
}
/// Sets properties on this builder based on a URL
///
/// This is a separate member function to allow fallible computation to
/// be deferred until [`Self::build`] which in turn allows deriving [`Clone`]
fn parse_url(&mut self, url: &str) -> Result<()> {
let parsed = Url::parse(url).map_err(|source| Error::UnableToParseUrl {
source,
url: url.to_string(),
})?;
let host = parsed.host_str().ok_or_else(|| Error::UrlNotRecognised {
url: url.to_string(),
})?;
match parsed.scheme() {
"gs" => self.bucket_name = Some(host.to_string()),
scheme => {
let scheme = scheme.to_string();
return Err(Error::UnknownUrlScheme { scheme }.into());
}
}
Ok(())
}
/// Set the bucket name (required)
pub fn with_bucket_name(mut self, bucket_name: impl Into<String>) -> Self {
self.bucket_name = Some(bucket_name.into());
self
}
/// Set the path to the service account file.
///
/// This or [`GoogleCloudStorageBuilder::with_service_account_key`] must be
/// set.
///
/// Example `"/tmp/gcs.json"`.
///
/// Example contents of `gcs.json`:
///
/// ```json
/// {
/// "gcs_base_url": "https://localhost:4443",
/// "disable_oauth": true,
/// "client_email": "",
/// "private_key": ""
/// }
/// ```
pub fn with_service_account_path(mut self, service_account_path: impl Into<String>) -> Self {
self.service_account_path = Some(service_account_path.into());
self
}
/// Set the service account key. The service account must be in the JSON
/// format.
///
/// This or [`GoogleCloudStorageBuilder::with_service_account_path`] must be
/// set.
pub fn with_service_account_key(mut self, service_account: impl Into<String>) -> Self {
self.service_account_key = Some(service_account.into());
self
}
/// Set the path to the application credentials file.
///
/// <https://cloud.google.com/docs/authentication/provide-credentials-adc>
pub fn with_application_credentials(
mut self,
application_credentials_path: impl Into<String>,
) -> Self {
self.application_credentials_path = Some(application_credentials_path.into());
self
}
/// If enabled, [`GoogleCloudStorage`] will not fetch credentials and will not sign requests.
///
/// This can be useful when interacting with public GCS buckets that deny authorized requests.
pub fn with_skip_signature(mut self, skip_signature: bool) -> Self {
self.skip_signature = skip_signature.into();
self
}
/// Set the credential provider overriding any other options
pub fn with_credentials(mut self, credentials: GcpCredentialProvider) -> Self {
self.credentials = Some(credentials);
self
}
/// Set the retry configuration
pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
self.retry_config = retry_config;
self
}
/// Set the proxy_url to be used by the underlying client
pub fn with_proxy_url(mut self, proxy_url: impl Into<String>) -> Self {
self.client_options = self.client_options.with_proxy_url(proxy_url);
self
}
/// Set a trusted proxy CA certificate
pub fn with_proxy_ca_certificate(mut self, proxy_ca_certificate: impl Into<String>) -> Self {
self.client_options = self
.client_options
.with_proxy_ca_certificate(proxy_ca_certificate);
self
}
/// Set a list of hosts to exclude from proxy connections
pub fn with_proxy_excludes(mut self, proxy_excludes: impl Into<String>) -> Self {
self.client_options = self.client_options.with_proxy_excludes(proxy_excludes);
self
}
/// Sets the client options, overriding any already set
pub fn with_client_options(mut self, options: ClientOptions) -> Self {
self.client_options = options;
self
}
/// The [`HttpConnector`] to use
///
/// On non-WASM32 platforms uses [`reqwest`] by default, on WASM32 platforms must be provided
pub fn with_http_connector<C: HttpConnector>(mut self, connector: C) -> Self {
self.http_connector = Some(Arc::new(connector));
self
}
/// Configure a connection to Google Cloud Storage, returning a
/// new [`GoogleCloudStorage`] and consuming `self`
pub fn build(mut self) -> Result<GoogleCloudStorage> {
if let Some(url) = self.url.take() {
self.parse_url(&url)?;
}
let bucket_name = self.bucket_name.ok_or(Error::MissingBucketName {})?;
let http = http_connector(self.http_connector)?;
// First try to initialize from the service account information.
let service_account_credentials =
match (self.service_account_path, self.service_account_key) {
(Some(path), None) => Some(
ServiceAccountCredentials::from_file(path)
.map_err(|source| Error::Credential { source })?,
),
(None, Some(key)) => Some(
ServiceAccountCredentials::from_key(&key)
.map_err(|source| Error::Credential { source })?,
),
(None, None) => None,
(Some(_), Some(_)) => return Err(Error::ServiceAccountPathAndKeyProvided.into()),
};
// Then try to initialize from the application credentials file, or the environment.
let application_default_credentials =
ApplicationDefaultCredentials::read(self.application_credentials_path.as_deref())?;
let disable_oauth = service_account_credentials
.as_ref()
.map(|c| c.disable_oauth)
.unwrap_or(false);
let gcs_base_url: String = service_account_credentials
.as_ref()
.and_then(|c| c.gcs_base_url.clone())
.unwrap_or_else(|| DEFAULT_GCS_BASE_URL.to_string());
let credentials = if let Some(credentials) = self.credentials {
credentials
} else if disable_oauth {
Arc::new(StaticCredentialProvider::new(GcpCredential {
bearer: "".to_string(),
})) as _
} else if let Some(credentials) = service_account_credentials.clone() {
Arc::new(TokenCredentialProvider::new(
credentials.token_provider()?,
http.connect(&self.client_options)?,
self.retry_config.clone(),
)) as _
} else if let Some(credentials) = application_default_credentials.clone() {
match credentials {
ApplicationDefaultCredentials::AuthorizedUser(token) => Arc::new(
TokenCredentialProvider::new(
token,
http.connect(&self.client_options)?,
self.retry_config.clone(),
)
.with_min_ttl(TOKEN_MIN_TTL),
) as _,
ApplicationDefaultCredentials::ServiceAccount(token) => {
Arc::new(TokenCredentialProvider::new(
token.token_provider()?,
http.connect(&self.client_options)?,
self.retry_config.clone(),
)) as _
}
}
} else {
Arc::new(
TokenCredentialProvider::new(
InstanceCredentialProvider::default(),
http.connect(&self.client_options.metadata_options())?,
self.retry_config.clone(),
)
.with_min_ttl(TOKEN_MIN_TTL),
) as _
};
let signing_credentials = if let Some(signing_credentials) = self.signing_credentials {
signing_credentials
} else if disable_oauth {
Arc::new(StaticCredentialProvider::new(GcpSigningCredential {
email: "".to_string(),
private_key: None,
})) as _
} else if let Some(credentials) = service_account_credentials.clone() {
credentials.signing_credentials()?
} else if let Some(credentials) = application_default_credentials.clone() {
match credentials {
ApplicationDefaultCredentials::AuthorizedUser(token) => {
Arc::new(TokenCredentialProvider::new(
AuthorizedUserSigningCredentials::from(token)?,
http.connect(&self.client_options)?,
self.retry_config.clone(),
)) as _
}
ApplicationDefaultCredentials::ServiceAccount(token) => {
token.signing_credentials()?
}
}
} else {
Arc::new(TokenCredentialProvider::new(
InstanceSigningCredentialProvider::default(),
http.connect(&self.client_options.metadata_options())?,
self.retry_config.clone(),
)) as _
};
let config = GoogleCloudStorageConfig {
base_url: gcs_base_url,
credentials,
signing_credentials,
bucket_name,
retry_config: self.retry_config,
client_options: self.client_options,
skip_signature: self.skip_signature.get()?,
};
let http_client = http.connect(&config.client_options)?;
Ok(GoogleCloudStorage {
client: Arc::new(GoogleCloudStorageClient::new(config, http_client)?),
})
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::collections::HashMap;
use std::io::Write;
use tempfile::NamedTempFile;
const FAKE_KEY: &str = r#"{"private_key": "private_key", "private_key_id": "private_key_id", "client_email":"client_email", "disable_oauth":true}"#;
#[test]
fn gcs_test_service_account_key_and_path() {
let mut tfile = NamedTempFile::new().unwrap();
write!(tfile, "{FAKE_KEY}").unwrap();
let _ = GoogleCloudStorageBuilder::new()
.with_service_account_key(FAKE_KEY)
.with_service_account_path(tfile.path().to_str().unwrap())
.with_bucket_name("foo")
.build()
.unwrap_err();
}
#[test]
fn gcs_test_config_from_map() {
let google_service_account = "object_store:fake_service_account".to_string();
let google_bucket_name = "object_store:fake_bucket".to_string();
let options = HashMap::from([
("google_service_account", google_service_account.clone()),
("google_bucket_name", google_bucket_name.clone()),
]);
let builder = options
.iter()
.fold(GoogleCloudStorageBuilder::new(), |builder, (key, value)| {
builder.with_config(key.parse().unwrap(), value)
});
assert_eq!(
builder.service_account_path.unwrap(),
google_service_account.as_str()
);
assert_eq!(builder.bucket_name.unwrap(), google_bucket_name.as_str());
}
#[test]
fn gcs_test_config_aliases() {
// Service account path
for alias in [
"google_service_account",
"service_account",
"google_service_account_path",
"service_account_path",
] {
let builder = GoogleCloudStorageBuilder::new()
.with_config(alias.parse().unwrap(), "/fake/path.json");
assert_eq!("/fake/path.json", builder.service_account_path.unwrap());
}
// Service account key
for alias in ["google_service_account_key", "service_account_key"] {
let builder =
GoogleCloudStorageBuilder::new().with_config(alias.parse().unwrap(), FAKE_KEY);
assert_eq!(FAKE_KEY, builder.service_account_key.unwrap());
}
// Bucket name
for alias in [
"google_bucket",
"google_bucket_name",
"bucket",
"bucket_name",
] {
let builder =
GoogleCloudStorageBuilder::new().with_config(alias.parse().unwrap(), "fake_bucket");
assert_eq!("fake_bucket", builder.bucket_name.unwrap());
}
}
#[tokio::test]
async fn gcs_test_proxy_url() {
let mut tfile = NamedTempFile::new().unwrap();
write!(tfile, "{FAKE_KEY}").unwrap();
let service_account_path = tfile.path();
let gcs = GoogleCloudStorageBuilder::new()
.with_service_account_path(service_account_path.to_str().unwrap())
.with_bucket_name("foo")
.with_proxy_url("https://example.com")
.build();
assert!(gcs.is_ok());
let err = GoogleCloudStorageBuilder::new()
.with_service_account_path(service_account_path.to_str().unwrap())
.with_bucket_name("foo")
.with_proxy_url("asdf://example.com")
.build()
.unwrap_err()
.to_string();
assert_eq!("Generic HTTP client error: builder error", err);
}
#[test]
fn gcs_test_urls() {
let mut builder = GoogleCloudStorageBuilder::new();
builder.parse_url("gs://bucket/path").unwrap();
assert_eq!(builder.bucket_name.as_deref(), Some("bucket"));
builder.parse_url("gs://bucket.mydomain/path").unwrap();
assert_eq!(builder.bucket_name.as_deref(), Some("bucket.mydomain"));
builder.parse_url("mailto://bucket/path").unwrap_err();
}
#[test]
fn gcs_test_service_account_key_only() {
let _ = GoogleCloudStorageBuilder::new()
.with_service_account_key(FAKE_KEY)
.with_bucket_name("foo")
.build()
.unwrap();
}
#[test]
fn gcs_test_config_get_value() {
let google_service_account = "object_store:fake_service_account".to_string();
let google_bucket_name = "object_store:fake_bucket".to_string();
let builder = GoogleCloudStorageBuilder::new()
.with_config(GoogleConfigKey::ServiceAccount, &google_service_account)
.with_config(GoogleConfigKey::Bucket, &google_bucket_name);
assert_eq!(
builder
.get_config_value(&GoogleConfigKey::ServiceAccount)
.unwrap(),
google_service_account
);
assert_eq!(
builder.get_config_value(&GoogleConfigKey::Bucket).unwrap(),
google_bucket_name
);
}
#[test]
fn gcp_test_client_opts() {
let key = "GOOGLE_PROXY_URL";
if let Ok(config_key) = key.to_ascii_lowercase().parse() {
assert_eq!(
GoogleConfigKey::Client(ClientConfigKey::ProxyUrl),
config_key
);
} else {
panic!("{} not propagated as ClientConfigKey", key);
}
}
}
backend/vendor/object_store/src/gcp/client.rs
Vendored
+707
View File
@@ -0,0 +1,707 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::builder::HttpRequestBuilder;
use crate::client::get::GetClient;
use crate::client::header::{get_put_result, get_version, HeaderConfig};
use crate::client::list::ListClient;
use crate::client::retry::RetryExt;
use crate::client::s3::{
CompleteMultipartUpload, CompleteMultipartUploadResult, InitiateMultipartUploadResult,
ListResponse,
};
use crate::client::{GetOptionsExt, HttpClient, HttpError, HttpResponse};
use crate::gcp::credential::CredentialExt;
use crate::gcp::{GcpCredential, GcpCredentialProvider, GcpSigningCredentialProvider, STORE};
use crate::multipart::PartId;
use crate::path::{Path, DELIMITER};
use crate::util::hex_encode;
use crate::{
Attribute, Attributes, ClientOptions, GetOptions, ListResult, MultipartId, PutMode,
PutMultipartOpts, PutOptions, PutPayload, PutResult, Result, RetryConfig,
};
use async_trait::async_trait;
use base64::prelude::BASE64_STANDARD;
use base64::Engine;
use bytes::Buf;
use http::header::{
CACHE_CONTROL, CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LANGUAGE, CONTENT_LENGTH,
CONTENT_TYPE,
};
use http::{HeaderName, Method, StatusCode};
use percent_encoding::{percent_encode, utf8_percent_encode, NON_ALPHANUMERIC};
use serde::{Deserialize, Serialize};
use std::sync::Arc;
const VERSION_HEADER: &str = "x-goog-generation";
const DEFAULT_CONTENT_TYPE: &str = "application/octet-stream";
const USER_DEFINED_METADATA_HEADER_PREFIX: &str = "x-goog-meta-";
static VERSION_MATCH: HeaderName = HeaderName::from_static("x-goog-if-generation-match");
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("Error performing list request: {}", source)]
ListRequest {
source: crate::client::retry::RetryError,
},
#[error("Error getting list response body: {}", source)]
ListResponseBody { source: HttpError },
#[error("Got invalid list response: {}", source)]
InvalidListResponse { source: quick_xml::de::DeError },
#[error("Error performing get request {}: {}", path, source)]
GetRequest {
source: crate::client::retry::RetryError,
path: String,
},
#[error("Error performing request {}: {}", path, source)]
Request {
source: crate::client::retry::RetryError,
path: String,
},
#[error("Error getting put response body: {}", source)]
PutResponseBody { source: HttpError },
#[error("Got invalid put request: {}", source)]
InvalidPutRequest { source: quick_xml::se::SeError },
#[error("Got invalid put response: {}", source)]
InvalidPutResponse { source: quick_xml::de::DeError },
#[error("Unable to extract metadata from headers: {}", source)]
Metadata {
source: crate::client::header::Error,
},
#[error("Version required for conditional update")]
MissingVersion,
#[error("Error performing complete multipart request: {}", source)]
CompleteMultipartRequest {
source: crate::client::retry::RetryError,
},
#[error("Error getting complete multipart response body: {}", source)]
CompleteMultipartResponseBody { source: HttpError },
#[error("Got invalid multipart response: {}", source)]
InvalidMultipartResponse { source: quick_xml::de::DeError },
#[error("Error signing blob: {}", source)]
SignBlobRequest {
source: crate::client::retry::RetryError,
},
#[error("Got invalid signing blob response: {}", source)]
InvalidSignBlobResponse { source: HttpError },
#[error("Got invalid signing blob signature: {}", source)]
InvalidSignBlobSignature { source: base64::DecodeError },
}
impl From<Error> for crate::Error {
fn from(err: Error) -> Self {
match err {
Error::GetRequest { source, path } | Error::Request { source, path } => {
source.error(STORE, path)
}
_ => Self::Generic {
store: STORE,
source: Box::new(err),
},
}
}
}
#[derive(Debug)]
pub(crate) struct GoogleCloudStorageConfig {
pub base_url: String,
pub credentials: GcpCredentialProvider,
pub signing_credentials: GcpSigningCredentialProvider,
pub bucket_name: String,
pub retry_config: RetryConfig,
pub client_options: ClientOptions,
pub skip_signature: bool,
}
impl GoogleCloudStorageConfig {
pub(crate) fn path_url(&self, path: &Path) -> String {
format!("{}/{}/{}", self.base_url, self.bucket_name, path)
}
pub(crate) async fn get_credential(&self) -> Result<Option<Arc<GcpCredential>>> {
Ok(match self.skip_signature {
false => Some(self.credentials.get_credential().await?),
true => None,
})
}
}
/// A builder for a put request allowing customisation of the headers and query string
pub(crate) struct Request<'a> {
path: &'a Path,
config: &'a GoogleCloudStorageConfig,
payload: Option<PutPayload>,
builder: HttpRequestBuilder,
idempotent: bool,
}
impl Request<'_> {
fn header(self, k: &HeaderName, v: &str) -> Self {
let builder = self.builder.header(k, v);
Self { builder, ..self }
}
fn query<T: Serialize + ?Sized + Sync>(self, query: &T) -> Self {
let builder = self.builder.query(query);
Self { builder, ..self }
}
fn idempotent(mut self, idempotent: bool) -> Self {
self.idempotent = idempotent;
self
}
fn with_attributes(self, attributes: Attributes) -> Self {
let mut builder = self.builder;
let mut has_content_type = false;
for (k, v) in &attributes {
builder = match k {
Attribute::CacheControl => builder.header(CACHE_CONTROL, v.as_ref()),
Attribute::ContentDisposition => builder.header(CONTENT_DISPOSITION, v.as_ref()),
Attribute::ContentEncoding => builder.header(CONTENT_ENCODING, v.as_ref()),
Attribute::ContentLanguage => builder.header(CONTENT_LANGUAGE, v.as_ref()),
Attribute::ContentType => {
has_content_type = true;
builder.header(CONTENT_TYPE, v.as_ref())
}
Attribute::Metadata(k_suffix) => builder.header(
&format!("{}{}", USER_DEFINED_METADATA_HEADER_PREFIX, k_suffix),
v.as_ref(),
),
};
}
if !has_content_type {
let value = self.config.client_options.get_content_type(self.path);
builder = builder.header(CONTENT_TYPE, value.unwrap_or(DEFAULT_CONTENT_TYPE))
}
Self { builder, ..self }
}
fn with_payload(self, payload: PutPayload) -> Self {
let content_length = payload.content_length();
Self {
builder: self.builder.header(CONTENT_LENGTH, content_length),
payload: Some(payload),
..self
}
}
fn with_extensions(self, extensions: ::http::Extensions) -> Self {
let builder = self.builder.extensions(extensions);
Self { builder, ..self }
}
async fn send(self) -> Result<HttpResponse> {
let credential = self.config.credentials.get_credential().await?;
let resp = self
.builder
.bearer_auth(&credential.bearer)
.retryable(&self.config.retry_config)
.idempotent(self.idempotent)
.payload(self.payload)
.send()
.await
.map_err(|source| {
let path = self.path.as_ref().into();
Error::Request { source, path }
})?;
Ok(resp)
}
async fn do_put(self) -> Result<PutResult> {
let response = self.send().await?;
Ok(get_put_result(response.headers(), VERSION_HEADER)
.map_err(|source| Error::Metadata { source })?)
}
}
/// Sign Blob Request Body
#[derive(Debug, Serialize)]
struct SignBlobBody {
/// The payload to sign
payload: String,
}
/// Sign Blob Response
#[derive(Deserialize)]
#[serde(rename_all = "camelCase")]
struct SignBlobResponse {
/// The signature for the payload
signed_blob: String,
}
#[derive(Debug)]
pub(crate) struct GoogleCloudStorageClient {
config: GoogleCloudStorageConfig,
client: HttpClient,
bucket_name_encoded: String,
// TODO: Hook this up in tests
max_list_results: Option<String>,
}
impl GoogleCloudStorageClient {
pub(crate) fn new(config: GoogleCloudStorageConfig, client: HttpClient) -> Result<Self> {
let bucket_name_encoded =
percent_encode(config.bucket_name.as_bytes(), NON_ALPHANUMERIC).to_string();
Ok(Self {
config,
client,
bucket_name_encoded,
max_list_results: None,
})
}
pub(crate) fn config(&self) -> &GoogleCloudStorageConfig {
&self.config
}
async fn get_credential(&self) -> Result<Option<Arc<GcpCredential>>> {
self.config.get_credential().await
}
/// Create a signature from a string-to-sign using Google Cloud signBlob method.
/// form like:
/// ```plaintext
/// curl -X POST --data-binary @JSON_FILE_NAME \
/// -H "Authorization: Bearer OAUTH2_TOKEN" \
/// -H "Content-Type: application/json" \
/// "https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/SERVICE_ACCOUNT_EMAIL:signBlob"
/// ```
///
/// 'JSON_FILE_NAME' is a file containing the following JSON object:
/// ```plaintext
/// {
/// "payload": "REQUEST_INFORMATION"
/// }
/// ```
pub(crate) async fn sign_blob(
&self,
string_to_sign: &str,
client_email: &str,
) -> Result<String> {
let credential = self.get_credential().await?;
let body = SignBlobBody {
payload: BASE64_STANDARD.encode(string_to_sign),
};
let url = format!(
"https://iamcredentials.googleapis.com/v1/projects/-/serviceAccounts/{}:signBlob",
client_email
);
let response = self
.client
.post(&url)
.with_bearer_auth(credential.as_deref())
.json(&body)
.retryable(&self.config.retry_config)
.idempotent(true)
.send()
.await
.map_err(|source| Error::SignBlobRequest { source })?
.into_body()
.json::<SignBlobResponse>()
.await
.map_err(|source| Error::InvalidSignBlobResponse { source })?;
let signed_blob = BASE64_STANDARD
.decode(response.signed_blob)
.map_err(|source| Error::InvalidSignBlobSignature { source })?;
Ok(hex_encode(&signed_blob))
}
pub(crate) fn object_url(&self, path: &Path) -> String {
let encoded = utf8_percent_encode(path.as_ref(), NON_ALPHANUMERIC);
format!(
"{}/{}/{}",
self.config.base_url, self.bucket_name_encoded, encoded
)
}
/// Perform a put request <https://cloud.google.com/storage/docs/xml-api/put-object-upload>
///
/// Returns the new ETag
pub(crate) fn request<'a>(&'a self, method: Method, path: &'a Path) -> Request<'a> {
let builder = self.client.request(method, self.object_url(path));
Request {
path,
builder,
payload: None,
config: &self.config,
idempotent: false,
}
}
pub(crate) async fn put(
&self,
path: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
let PutOptions {
mode,
// not supported by GCP
tags: _,
attributes,
extensions,
} = opts;
let builder = self
.request(Method::PUT, path)
.with_payload(payload)
.with_attributes(attributes)
.with_extensions(extensions);
let builder = match &mode {
PutMode::Overwrite => builder.idempotent(true),
PutMode::Create => builder.header(&VERSION_MATCH, "0"),
PutMode::Update(v) => {
let etag = v.version.as_ref().ok_or(Error::MissingVersion)?;
builder.header(&VERSION_MATCH, etag)
}
};
match (mode, builder.do_put().await) {
(PutMode::Create, Err(crate::Error::Precondition { path, source })) => {
Err(crate::Error::AlreadyExists { path, source })
}
(_, r) => r,
}
}
/// Perform a put part request <https://cloud.google.com/storage/docs/xml-api/put-object-multipart>
///
/// Returns the new [`PartId`]
pub(crate) async fn put_part(
&self,
path: &Path,
upload_id: &MultipartId,
part_idx: usize,
data: PutPayload,
) -> Result<PartId> {
let query = &[
("partNumber", &format!("{}", part_idx + 1)),
("uploadId", upload_id),
];
let result = self
.request(Method::PUT, path)
.with_payload(data)
.query(query)
.idempotent(true)
.do_put()
.await?;
Ok(PartId {
content_id: result.e_tag.unwrap(),
})
}
/// Initiate a multipart upload <https://cloud.google.com/storage/docs/xml-api/post-object-multipart>
pub(crate) async fn multipart_initiate(
&self,
path: &Path,
opts: PutMultipartOpts,
) -> Result<MultipartId> {
let PutMultipartOpts {
// not supported by GCP
tags: _,
attributes,
extensions,
} = opts;
let response = self
.request(Method::POST, path)
.with_attributes(attributes)
.with_extensions(extensions)
.header(&CONTENT_LENGTH, "0")
.query(&[("uploads", "")])
.send()
.await?;
let data = response
.into_body()
.bytes()
.await
.map_err(|source| Error::PutResponseBody { source })?;
let result: InitiateMultipartUploadResult =
quick_xml::de::from_reader(data.as_ref().reader())
.map_err(|source| Error::InvalidPutResponse { source })?;
Ok(result.upload_id)
}
/// Cleanup unused parts <https://cloud.google.com/storage/docs/xml-api/delete-multipart>
pub(crate) async fn multipart_cleanup(
&self,
path: &Path,
multipart_id: &MultipartId,
) -> Result<()> {
let credential = self.get_credential().await?;
let url = self.object_url(path);
self.client
.request(Method::DELETE, &url)
.with_bearer_auth(credential.as_deref())
.header(CONTENT_TYPE, "application/octet-stream")
.header(CONTENT_LENGTH, "0")
.query(&[("uploadId", multipart_id)])
.send_retry(&self.config.retry_config)
.await
.map_err(|source| {
let path = path.as_ref().into();
Error::Request { source, path }
})?;
Ok(())
}
pub(crate) async fn multipart_complete(
&self,
path: &Path,
multipart_id: &MultipartId,
completed_parts: Vec<PartId>,
) -> Result<PutResult> {
if completed_parts.is_empty() {
// GCS doesn't allow empty multipart uploads, so fallback to regular upload.
self.multipart_cleanup(path, multipart_id).await?;
let result = self
.put(path, PutPayload::new(), Default::default())
.await?;
return Ok(result);
}
let upload_id = multipart_id.clone();
let url = self.object_url(path);
let upload_info = CompleteMultipartUpload::from(completed_parts);
let credential = self.get_credential().await?;
let data = quick_xml::se::to_string(&upload_info)
.map_err(|source| Error::InvalidPutRequest { source })?
// We cannot disable the escaping that transforms "/" to "&quote;" :(
// https://github.com/tafia/quick-xml/issues/362
// https://github.com/tafia/quick-xml/issues/350
.replace("&quot;", "\"");
let response = self
.client
.request(Method::POST, &url)
.with_bearer_auth(credential.as_deref())
.query(&[("uploadId", upload_id)])
.body(data)
.retryable(&self.config.retry_config)
.idempotent(true)
.send()
.await
.map_err(|source| Error::CompleteMultipartRequest { source })?;
let version = get_version(response.headers(), VERSION_HEADER)
.map_err(|source| Error::Metadata { source })?;
let data = response
.into_body()
.bytes()
.await
.map_err(|source| Error::CompleteMultipartResponseBody { source })?;
let response: CompleteMultipartUploadResult = quick_xml::de::from_reader(data.reader())
.map_err(|source| Error::InvalidMultipartResponse { source })?;
Ok(PutResult {
e_tag: Some(response.e_tag),
version,
})
}
/// Perform a delete request <https://cloud.google.com/storage/docs/xml-api/delete-object>
pub(crate) async fn delete_request(&self, path: &Path) -> Result<()> {
self.request(Method::DELETE, path).send().await?;
Ok(())
}
/// Perform a copy request <https://cloud.google.com/storage/docs/xml-api/put-object-copy>
pub(crate) async fn copy_request(
&self,
from: &Path,
to: &Path,
if_not_exists: bool,
) -> Result<()> {
let credential = self.get_credential().await?;
let url = self.object_url(to);
let from = utf8_percent_encode(from.as_ref(), NON_ALPHANUMERIC);
let source = format!("{}/{}", self.bucket_name_encoded, from);
let mut builder = self
.client
.request(Method::PUT, url)
.header("x-goog-copy-source", source);
if if_not_exists {
builder = builder.header(&VERSION_MATCH, 0);
}
builder
.with_bearer_auth(credential.as_deref())
// Needed if reqwest is compiled with native-tls instead of rustls-tls
// See https://github.com/apache/arrow-rs/pull/3921
.header(CONTENT_LENGTH, 0)
.retryable(&self.config.retry_config)
.idempotent(!if_not_exists)
.send()
.await
.map_err(|err| match err.status() {
Some(StatusCode::PRECONDITION_FAILED) => crate::Error::AlreadyExists {
source: Box::new(err),
path: to.to_string(),
},
_ => err.error(STORE, from.to_string()),
})?;
Ok(())
}
}
#[async_trait]
impl GetClient for GoogleCloudStorageClient {
const STORE: &'static str = STORE;
const HEADER_CONFIG: HeaderConfig = HeaderConfig {
etag_required: true,
last_modified_required: true,
version_header: Some(VERSION_HEADER),
user_defined_metadata_prefix: Some(USER_DEFINED_METADATA_HEADER_PREFIX),
};
/// Perform a get request <https://cloud.google.com/storage/docs/xml-api/get-object-download>
async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse> {
let credential = self.get_credential().await?;
let url = self.object_url(path);
let method = match options.head {
true => Method::HEAD,
false => Method::GET,
};
let mut request = self.client.request(method, url);
if let Some(version) = &options.version {
request = request.query(&[("generation", version)]);
}
let response = request
.with_bearer_auth(credential.as_deref())
.with_get_options(options)
.send_retry(&self.config.retry_config)
.await
.map_err(|source| {
let path = path.as_ref().into();
Error::GetRequest { source, path }
})?;
Ok(response)
}
}
#[async_trait]
impl ListClient for Arc<GoogleCloudStorageClient> {
/// Perform a list request <https://cloud.google.com/storage/docs/xml-api/get-bucket-list>
async fn list_request(
&self,
prefix: Option<&str>,
delimiter: bool,
page_token: Option<&str>,
offset: Option<&str>,
max_keys: Option<usize>,
) -> Result<(ListResult, Option<String>)> {
let credential = self.get_credential().await?;
let url = format!("{}/{}", self.config.base_url, self.bucket_name_encoded);
let mut query = Vec::with_capacity(5);
query.push(("list-type", "2"));
if delimiter {
query.push(("delimiter", DELIMITER))
}
if let Some(prefix) = &prefix {
query.push(("prefix", prefix))
}
if let Some(page_token) = page_token {
query.push(("continuation-token", page_token))
}
let max_keys = max_keys.map(|x| x.to_string());
if let Some(max_keys) = &max_keys {
query.push(("max-keys", max_keys))
} else if let Some(max_results) = &self.max_list_results {
query.push(("max-keys", max_results))
}
if let Some(offset) = offset {
query.push(("start-after", offset))
}
let response = self
.client
.request(Method::GET, url)
.query(&query)
.with_bearer_auth(credential.as_deref())
.send_retry(&self.config.retry_config)
.await
.map_err(|source| Error::ListRequest { source })?
.into_body()
.bytes()
.await
.map_err(|source| Error::ListResponseBody { source })?;
let mut response: ListResponse = quick_xml::de::from_reader(response.reader())
.map_err(|source| Error::InvalidListResponse { source })?;
let token = response.next_continuation_token.take();
Ok((response.try_into()?, token))
}
}
backend/vendor/object_store/src/gcp/credential.rs
Vendored
+960
View File
@@ -0,0 +1,960 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use super::client::GoogleCloudStorageClient;
use crate::client::builder::HttpRequestBuilder;
use crate::client::retry::RetryExt;
use crate::client::token::TemporaryToken;
use crate::client::{HttpClient, HttpError, TokenProvider};
use crate::gcp::{GcpSigningCredentialProvider, STORE};
use crate::util::{hex_digest, hex_encode, STRICT_ENCODE_SET};
use crate::{RetryConfig, StaticCredentialProvider};
use async_trait::async_trait;
use base64::prelude::BASE64_URL_SAFE_NO_PAD;
use base64::Engine;
use chrono::{DateTime, Utc};
use futures::TryFutureExt;
use http::{HeaderMap, Method};
use itertools::Itertools;
use percent_encoding::utf8_percent_encode;
use ring::signature::RsaKeyPair;
use serde::Deserialize;
use std::collections::BTreeMap;
use std::env;
use std::fs::File;
use std::io::BufReader;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant};
use tracing::info;
use url::Url;
pub(crate) const DEFAULT_SCOPE: &str = "https://www.googleapis.com/auth/cloud-platform";
pub(crate) const DEFAULT_GCS_BASE_URL: &str = "https://storage.googleapis.com";
const DEFAULT_GCS_PLAYLOAD_STRING: &str = "UNSIGNED-PAYLOAD";
const DEFAULT_GCS_SIGN_BLOB_HOST: &str = "storage.googleapis.com";
const DEFAULT_METADATA_HOST: &str = "metadata.google.internal";
const DEFAULT_METADATA_IP: &str = "169.254.169.254";
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("Unable to open service account file from {}: {}", path.display(), source)]
OpenCredentials {
source: std::io::Error,
path: PathBuf,
},
#[error("Unable to decode service account file: {}", source)]
DecodeCredentials { source: serde_json::Error },
#[error("No RSA key found in pem file")]
MissingKey,
#[error("Invalid RSA key: {}", source)]
InvalidKey {
#[from]
source: ring::error::KeyRejected,
},
#[error("Error signing: {}", source)]
Sign { source: ring::error::Unspecified },
#[error("Error encoding jwt payload: {}", source)]
Encode { source: serde_json::Error },
#[error("Unsupported key encoding: {}", encoding)]
UnsupportedKey { encoding: String },
#[error("Error performing token request: {}", source)]
TokenRequest {
source: crate::client::retry::RetryError,
},
#[error("Error getting token response body: {}", source)]
TokenResponseBody { source: HttpError },
}
impl From<Error> for crate::Error {
fn from(value: Error) -> Self {
Self::Generic {
store: STORE,
source: Box::new(value),
}
}
}
/// A Google Cloud Storage Credential for signing
#[derive(Debug)]
pub struct GcpSigningCredential {
/// The email of the service account
pub email: String,
/// An optional RSA private key
///
/// If provided this will be used to sign the URL, otherwise a call will be made to
/// [`iam.serviceAccounts.signBlob`]. This allows supporting credential sources
/// that don't expose the service account private key, e.g. [IMDS].
///
/// [IMDS]: https://cloud.google.com/docs/authentication/get-id-token#metadata-server
/// [`iam.serviceAccounts.signBlob`]: https://cloud.google.com/storage/docs/authentication/creating-signatures
pub private_key: Option<ServiceAccountKey>,
}
/// A private RSA key for a service account
#[derive(Debug)]
pub struct ServiceAccountKey(RsaKeyPair);
impl ServiceAccountKey {
/// Parses a pem-encoded RSA key
pub fn from_pem(encoded: &[u8]) -> Result<Self> {
use rustls_pemfile::Item;
use std::io::Cursor;
let mut cursor = Cursor::new(encoded);
let mut reader = BufReader::new(&mut cursor);
// Reading from string is infallible
match rustls_pemfile::read_one(&mut reader).unwrap() {
Some(Item::Pkcs8Key(key)) => Self::from_pkcs8(key.secret_pkcs8_der()),
Some(Item::Pkcs1Key(key)) => Self::from_der(key.secret_pkcs1_der()),
_ => Err(Error::MissingKey),
}
}
/// Parses an unencrypted PKCS#8-encoded RSA private key.
pub fn from_pkcs8(key: &[u8]) -> Result<Self> {
Ok(Self(RsaKeyPair::from_pkcs8(key)?))
}
/// Parses an unencrypted PKCS#8-encoded RSA private key.
pub fn from_der(key: &[u8]) -> Result<Self> {
Ok(Self(RsaKeyPair::from_der(key)?))
}
fn sign(&self, string_to_sign: &str) -> Result<String> {
let mut signature = vec![0; self.0.public().modulus_len()];
self.0
.sign(
&ring::signature::RSA_PKCS1_SHA256,
&ring::rand::SystemRandom::new(),
string_to_sign.as_bytes(),
&mut signature,
)
.map_err(|source| Error::Sign { source })?;
Ok(hex_encode(&signature))
}
}
/// A Google Cloud Storage Credential
#[derive(Debug, Eq, PartialEq)]
pub struct GcpCredential {
/// An HTTP bearer token
pub bearer: String,
}
pub(crate) type Result<T, E = Error> = std::result::Result<T, E>;
#[derive(Debug, Default, serde::Serialize)]
pub(crate) struct JwtHeader<'a> {
/// The type of JWS: it can only be "JWT" here
///
/// Defined in [RFC7515#4.1.9](https://tools.ietf.org/html/rfc7515#section-4.1.9).
#[serde(skip_serializing_if = "Option::is_none")]
pub typ: Option<&'a str>,
/// The algorithm used
///
/// Defined in [RFC7515#4.1.1](https://tools.ietf.org/html/rfc7515#section-4.1.1).
pub alg: &'a str,
/// Content type
///
/// Defined in [RFC7519#5.2](https://tools.ietf.org/html/rfc7519#section-5.2).
#[serde(skip_serializing_if = "Option::is_none")]
pub cty: Option<&'a str>,
/// JSON Key URL
///
/// Defined in [RFC7515#4.1.2](https://tools.ietf.org/html/rfc7515#section-4.1.2).
#[serde(skip_serializing_if = "Option::is_none")]
pub jku: Option<&'a str>,
/// Key ID
///
/// Defined in [RFC7515#4.1.4](https://tools.ietf.org/html/rfc7515#section-4.1.4).
#[serde(skip_serializing_if = "Option::is_none")]
pub kid: Option<&'a str>,
/// X.509 URL
///
/// Defined in [RFC7515#4.1.5](https://tools.ietf.org/html/rfc7515#section-4.1.5).
#[serde(skip_serializing_if = "Option::is_none")]
pub x5u: Option<&'a str>,
/// X.509 certificate thumbprint
///
/// Defined in [RFC7515#4.1.7](https://tools.ietf.org/html/rfc7515#section-4.1.7).
#[serde(skip_serializing_if = "Option::is_none")]
pub x5t: Option<&'a str>,
}
#[derive(serde::Serialize)]
struct TokenClaims<'a> {
iss: &'a str,
sub: &'a str,
scope: &'a str,
exp: u64,
iat: u64,
}
#[derive(serde::Deserialize, Debug)]
struct TokenResponse {
access_token: String,
expires_in: u64,
id_token: Option<String>,
}
/// Self-signed JWT (JSON Web Token).
///
/// # References
/// - <https://google.aip.dev/auth/4111>
#[derive(Debug)]
pub(crate) struct SelfSignedJwt {
issuer: String,
scope: String,
private_key: ServiceAccountKey,
key_id: String,
}
impl SelfSignedJwt {
/// Create a new [`SelfSignedJwt`]
pub(crate) fn new(
key_id: String,
issuer: String,
private_key: ServiceAccountKey,
scope: String,
) -> Result<Self> {
Ok(Self {
issuer,
scope,
private_key,
key_id,
})
}
}
#[async_trait]
impl TokenProvider for SelfSignedJwt {
type Credential = GcpCredential;
/// Fetch a fresh token
async fn fetch_token(
&self,
_client: &HttpClient,
_retry: &RetryConfig,
) -> crate::Result<TemporaryToken<Arc<GcpCredential>>> {
let now = seconds_since_epoch();
let exp = now + 3600;
let claims = TokenClaims {
iss: &self.issuer,
sub: &self.issuer,
scope: &self.scope,
iat: now,
exp,
};
let jwt_header = b64_encode_obj(&JwtHeader {
alg: "RS256",
typ: Some("JWT"),
kid: Some(&self.key_id),
..Default::default()
})?;
let claim_str = b64_encode_obj(&claims)?;
let message = [jwt_header.as_ref(), claim_str.as_ref()].join(".");
let mut sig_bytes = vec![0; self.private_key.0.public().modulus_len()];
self.private_key
.0
.sign(
&ring::signature::RSA_PKCS1_SHA256,
&ring::rand::SystemRandom::new(),
message.as_bytes(),
&mut sig_bytes,
)
.map_err(|source| Error::Sign { source })?;
let signature = BASE64_URL_SAFE_NO_PAD.encode(sig_bytes);
let bearer = [message, signature].join(".");
Ok(TemporaryToken {
token: Arc::new(GcpCredential { bearer }),
expiry: Some(Instant::now() + Duration::from_secs(3600)),
})
}
}
fn read_credentials_file<T>(service_account_path: impl AsRef<std::path::Path>) -> Result<T>
where
T: serde::de::DeserializeOwned,
{
let file = File::open(&service_account_path).map_err(|source| {
let path = service_account_path.as_ref().to_owned();
Error::OpenCredentials { source, path }
})?;
let reader = BufReader::new(file);
serde_json::from_reader(reader).map_err(|source| Error::DecodeCredentials { source })
}
/// A deserialized `service-account-********.json`-file.
#[derive(serde::Deserialize, Debug, Clone)]
pub(crate) struct ServiceAccountCredentials {
/// The private key in RSA format.
pub private_key: String,
/// The private key ID
pub private_key_id: String,
/// The email address associated with the service account.
pub client_email: String,
/// Base URL for GCS
#[serde(default)]
pub gcs_base_url: Option<String>,
/// Disable oauth and use empty tokens.
#[serde(default)]
pub disable_oauth: bool,
}
impl ServiceAccountCredentials {
/// Create a new [`ServiceAccountCredentials`] from a file.
pub(crate) fn from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
read_credentials_file(path)
}
/// Create a new [`ServiceAccountCredentials`] from a string.
pub(crate) fn from_key(key: &str) -> Result<Self> {
serde_json::from_str(key).map_err(|source| Error::DecodeCredentials { source })
}
/// Create a [`SelfSignedJwt`] from this credentials struct.
///
/// We use a scope of [`DEFAULT_SCOPE`] as opposed to an audience
/// as GCS appears to not support audience
///
/// # References
/// - <https://stackoverflow.com/questions/63222450/service-account-authorization-without-oauth-can-we-get-file-from-google-cloud/71834557#71834557>
/// - <https://www.codejam.info/2022/05/google-cloud-service-account-authorization-without-oauth.html>
pub(crate) fn token_provider(self) -> crate::Result<SelfSignedJwt> {
Ok(SelfSignedJwt::new(
self.private_key_id,
self.client_email,
ServiceAccountKey::from_pem(self.private_key.as_bytes())?,
DEFAULT_SCOPE.to_string(),
)?)
}
pub(crate) fn signing_credentials(self) -> crate::Result<GcpSigningCredentialProvider> {
Ok(Arc::new(StaticCredentialProvider::new(
GcpSigningCredential {
email: self.client_email,
private_key: Some(ServiceAccountKey::from_pem(self.private_key.as_bytes())?),
},
)))
}
}
/// Returns the number of seconds since unix epoch
fn seconds_since_epoch() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::SystemTime::UNIX_EPOCH)
.unwrap()
.as_secs()
}
fn b64_encode_obj<T: serde::Serialize>(obj: &T) -> Result<String> {
let string = serde_json::to_string(obj).map_err(|source| Error::Encode { source })?;
Ok(BASE64_URL_SAFE_NO_PAD.encode(string))
}
/// A provider that uses the Google Cloud Platform metadata server to fetch a token.
///
/// <https://cloud.google.com/docs/authentication/get-id-token#metadata-server>
#[derive(Debug, Default)]
pub(crate) struct InstanceCredentialProvider {}
/// Make a request to the metadata server to fetch a token, using a a given hostname.
async fn make_metadata_request(
client: &HttpClient,
hostname: &str,
retry: &RetryConfig,
) -> crate::Result<TokenResponse> {
let url =
format!("http://{hostname}/computeMetadata/v1/instance/service-accounts/default/token");
let response: TokenResponse = client
.get(url)
.header("Metadata-Flavor", "Google")
.query(&[("audience", "https://www.googleapis.com/oauth2/v4/token")])
.send_retry(retry)
.await
.map_err(|source| Error::TokenRequest { source })?
.into_body()
.json()
.await
.map_err(|source| Error::TokenResponseBody { source })?;
Ok(response)
}
#[async_trait]
impl TokenProvider for InstanceCredentialProvider {
type Credential = GcpCredential;
/// Fetch a token from the metadata server.
/// Since the connection is local we need to enable http access and don't actually use the client object passed in.
/// Respects the `GCE_METADATA_HOST`, `GCE_METADATA_ROOT`, and `GCE_METADATA_IP`
/// environment variables.
///
/// References: <https://googleapis.dev/python/google-auth/latest/reference/google.auth.environment_vars.html>
async fn fetch_token(
&self,
client: &HttpClient,
retry: &RetryConfig,
) -> crate::Result<TemporaryToken<Arc<GcpCredential>>> {
let metadata_host = if let Ok(host) = env::var("GCE_METADATA_HOST") {
host
} else if let Ok(host) = env::var("GCE_METADATA_ROOT") {
host
} else {
DEFAULT_METADATA_HOST.to_string()
};
let metadata_ip = if let Ok(ip) = env::var("GCE_METADATA_IP") {
ip
} else {
DEFAULT_METADATA_IP.to_string()
};
info!("fetching token from metadata server");
let response = make_metadata_request(client, &metadata_host, retry)
.or_else(|_| make_metadata_request(client, &metadata_ip, retry))
.await?;
let token = TemporaryToken {
token: Arc::new(GcpCredential {
bearer: response.access_token,
}),
expiry: Some(Instant::now() + Duration::from_secs(response.expires_in)),
};
Ok(token)
}
}
/// Make a request to the metadata server to fetch the client email, using a given hostname.
async fn make_metadata_request_for_email(
client: &HttpClient,
hostname: &str,
retry: &RetryConfig,
) -> crate::Result<String> {
let url =
format!("http://{hostname}/computeMetadata/v1/instance/service-accounts/default/email",);
let response = client
.get(url)
.header("Metadata-Flavor", "Google")
.send_retry(retry)
.await
.map_err(|source| Error::TokenRequest { source })?
.into_body()
.text()
.await
.map_err(|source| Error::TokenResponseBody { source })?;
Ok(response)
}
/// A provider that uses the Google Cloud Platform metadata server to fetch a email for signing.
///
/// <https://cloud.google.com/appengine/docs/legacy/standard/java/accessing-instance-metadata>
#[derive(Debug, Default)]
pub(crate) struct InstanceSigningCredentialProvider {}
#[async_trait]
impl TokenProvider for InstanceSigningCredentialProvider {
type Credential = GcpSigningCredential;
/// Fetch a token from the metadata server.
/// Since the connection is local we need to enable http access and don't actually use the client object passed in.
/// Respects the `GCE_METADATA_HOST`, `GCE_METADATA_ROOT`, and `GCE_METADATA_IP`
/// environment variables.
///
/// References: <https://googleapis.dev/python/google-auth/latest/reference/google.auth.environment_vars.html>
async fn fetch_token(
&self,
client: &HttpClient,
retry: &RetryConfig,
) -> crate::Result<TemporaryToken<Arc<GcpSigningCredential>>> {
let metadata_host = if let Ok(host) = env::var("GCE_METADATA_HOST") {
host
} else if let Ok(host) = env::var("GCE_METADATA_ROOT") {
host
} else {
DEFAULT_METADATA_HOST.to_string()
};
let metadata_ip = if let Ok(ip) = env::var("GCE_METADATA_IP") {
ip
} else {
DEFAULT_METADATA_IP.to_string()
};
info!("fetching token from metadata server");
let email = make_metadata_request_for_email(client, &metadata_host, retry)
.or_else(|_| make_metadata_request_for_email(client, &metadata_ip, retry))
.await?;
let token = TemporaryToken {
token: Arc::new(GcpSigningCredential {
email,
private_key: None,
}),
expiry: None,
};
Ok(token)
}
}
/// A deserialized `application_default_credentials.json`-file.
///
/// # References
/// - <https://cloud.google.com/docs/authentication/application-default-credentials#personal>
/// - <https://google.aip.dev/auth/4110>
#[derive(serde::Deserialize, Clone)]
#[serde(tag = "type")]
pub(crate) enum ApplicationDefaultCredentials {
/// Service Account.
///
/// # References
/// - <https://google.aip.dev/auth/4112>
#[serde(rename = "service_account")]
ServiceAccount(ServiceAccountCredentials),
/// Authorized user via "gcloud CLI Integration".
///
/// # References
/// - <https://google.aip.dev/auth/4113>
#[serde(rename = "authorized_user")]
AuthorizedUser(AuthorizedUserCredentials),
}
impl ApplicationDefaultCredentials {
const CREDENTIALS_PATH: &'static str = if cfg!(windows) {
"gcloud/application_default_credentials.json"
} else {
".config/gcloud/application_default_credentials.json"
};
// Create a new application default credential in the following situations:
// 1. a file is passed in and the type matches.
// 2. without argument if the well-known configuration file is present.
pub(crate) fn read(path: Option<&str>) -> Result<Option<Self>, Error> {
if let Some(path) = path {
return read_credentials_file::<Self>(path).map(Some);
}
let home_var = if cfg!(windows) { "APPDATA" } else { "HOME" };
if let Some(home) = env::var_os(home_var) {
let path = Path::new(&home).join(Self::CREDENTIALS_PATH);
// It's expected for this file to not exist unless it has been explicitly configured by the user.
if path.exists() {
return read_credentials_file::<Self>(path).map(Some);
}
}
Ok(None)
}
}
const DEFAULT_TOKEN_GCP_URI: &str = "https://accounts.google.com/o/oauth2/token";
/// <https://google.aip.dev/auth/4113>
#[derive(Debug, Deserialize, Clone)]
pub(crate) struct AuthorizedUserCredentials {
client_id: String,
client_secret: String,
refresh_token: String,
}
#[derive(Debug, Deserialize)]
pub(crate) struct AuthorizedUserSigningCredentials {
credential: AuthorizedUserCredentials,
}
///<https://oauth2.googleapis.com/tokeninfo?access_token=ACCESS_TOKEN>
#[derive(Debug, Deserialize)]
struct EmailResponse {
email: String,
}
#[derive(Debug, Deserialize)]
struct IdTokenClaims {
email: String,
}
async fn get_token_response(
client_id: &str,
client_secret: &str,
refresh_token: &str,
client: &HttpClient,
retry: &RetryConfig,
) -> Result<TokenResponse> {
client
.post(DEFAULT_TOKEN_GCP_URI)
.form([
("grant_type", "refresh_token"),
("client_id", client_id),
("client_secret", client_secret),
("refresh_token", refresh_token),
])
.retryable(retry)
.idempotent(true)
.send()
.await
.map_err(|source| Error::TokenRequest { source })?
.into_body()
.json::<TokenResponse>()
.await
.map_err(|source| Error::TokenResponseBody { source })
}
impl AuthorizedUserSigningCredentials {
pub(crate) fn from(credential: AuthorizedUserCredentials) -> crate::Result<Self> {
Ok(Self { credential })
}
async fn client_email(
&self,
client: &HttpClient,
retry: &RetryConfig,
) -> crate::Result<String> {
let response = get_token_response(
&self.credential.client_id,
&self.credential.client_secret,
&self.credential.refresh_token,
client,
retry,
)
.await?;
// Extract email from id_token if available
if let Some(id_token) = response.id_token {
// Split the JWT string by dots to get the payload section
let parts: Vec<&str> = id_token.split('.').collect();
if parts.len() == 3 {
// Decode the base64-encoded payload (middle part)
if let Ok(payload) = BASE64_URL_SAFE_NO_PAD.decode(parts[1]) {
// Parse the payload as JSON and extract the email
if let Ok(claims) = serde_json::from_slice::<IdTokenClaims>(&payload) {
return Ok(claims.email);
}
}
// If any of the parsing steps fail, fallback to other method
}
}
// Fallback to the original method if id_token is not available or invalid
let response = client
.get("https://oauth2.googleapis.com/tokeninfo")
.query(&[("access_token", response.access_token)])
.send_retry(retry)
.await
.map_err(|source| Error::TokenRequest { source })?
.into_body()
.json::<EmailResponse>()
.await
.map_err(|source: HttpError| Error::TokenResponseBody { source })?;
Ok(response.email)
}
}
#[async_trait]
impl TokenProvider for AuthorizedUserSigningCredentials {
type Credential = GcpSigningCredential;
async fn fetch_token(
&self,
client: &HttpClient,
retry: &RetryConfig,
) -> crate::Result<TemporaryToken<Arc<GcpSigningCredential>>> {
let email = self.client_email(client, retry).await?;
Ok(TemporaryToken {
token: Arc::new(GcpSigningCredential {
email,
private_key: None,
}),
expiry: None,
})
}
}
#[async_trait]
impl TokenProvider for AuthorizedUserCredentials {
type Credential = GcpCredential;
async fn fetch_token(
&self,
client: &HttpClient,
retry: &RetryConfig,
) -> crate::Result<TemporaryToken<Arc<GcpCredential>>> {
let response = get_token_response(
&self.client_id,
&self.client_secret,
&self.refresh_token,
client,
retry,
)
.await?;
Ok(TemporaryToken {
token: Arc::new(GcpCredential {
bearer: response.access_token,
}),
expiry: Some(Instant::now() + Duration::from_secs(response.expires_in)),
})
}
}
/// Trim whitespace from header values
fn trim_header_value(value: &str) -> String {
let mut ret = value.to_string();
ret.retain(|c| !c.is_whitespace());
ret
}
/// A Google Cloud Storage Authorizer for generating signed URL using [Google SigV4]
///
/// [Google SigV4]: https://cloud.google.com/storage/docs/access-control/signed-urls
#[derive(Debug)]
pub(crate) struct GCSAuthorizer {
date: Option<DateTime<Utc>>,
credential: Arc<GcpSigningCredential>,
}
impl GCSAuthorizer {
/// Create a new [`GCSAuthorizer`]
pub(crate) fn new(credential: Arc<GcpSigningCredential>) -> Self {
Self {
date: None,
credential,
}
}
pub(crate) async fn sign(
&self,
method: Method,
url: &mut Url,
expires_in: Duration,
client: &GoogleCloudStorageClient,
) -> crate::Result<()> {
let email = &self.credential.email;
let date = self.date.unwrap_or_else(Utc::now);
let scope = self.scope(date);
let credential_with_scope = format!("{}/{}", email, scope);
let mut headers = HeaderMap::new();
headers.insert("host", DEFAULT_GCS_SIGN_BLOB_HOST.parse().unwrap());
let (_, signed_headers) = Self::canonicalize_headers(&headers);
url.query_pairs_mut()
.append_pair("X-Goog-Algorithm", "GOOG4-RSA-SHA256")
.append_pair("X-Goog-Credential", &credential_with_scope)
.append_pair("X-Goog-Date", &date.format("%Y%m%dT%H%M%SZ").to_string())
.append_pair("X-Goog-Expires", &expires_in.as_secs().to_string())
.append_pair("X-Goog-SignedHeaders", &signed_headers);
let string_to_sign = self.string_to_sign(date, &method, url, &headers);
let signature = match &self.credential.private_key {
Some(key) => key.sign(&string_to_sign)?,
None => client.sign_blob(&string_to_sign, email).await?,
};
url.query_pairs_mut()
.append_pair("X-Goog-Signature", &signature);
Ok(())
}
/// Get scope for the request
///
/// <https://cloud.google.com/storage/docs/authentication/signatures#credential-scope>
fn scope(&self, date: DateTime<Utc>) -> String {
format!("{}/auto/storage/goog4_request", date.format("%Y%m%d"),)
}
/// Canonicalizes query parameters into the GCP canonical form
/// form like:
///```plaintext
///HTTP_VERB
///PATH_TO_RESOURCE
///CANONICAL_QUERY_STRING
///CANONICAL_HEADERS
///
///SIGNED_HEADERS
///PAYLOAD
///```
///
/// <https://cloud.google.com/storage/docs/authentication/canonical-requests>
fn canonicalize_request(url: &Url, method: &Method, headers: &HeaderMap) -> String {
let verb = method.as_str();
let path = url.path();
let query = Self::canonicalize_query(url);
let (canonical_headers, signed_headers) = Self::canonicalize_headers(headers);
format!(
"{}\n{}\n{}\n{}\n\n{}\n{}",
verb, path, query, canonical_headers, signed_headers, DEFAULT_GCS_PLAYLOAD_STRING
)
}
/// Canonicalizes query parameters into the GCP canonical form
/// form like `max-keys=2&prefix=object`
///
/// <https://cloud.google.com/storage/docs/authentication/canonical-requests#about-query-strings>
fn canonicalize_query(url: &Url) -> String {
url.query_pairs()
.sorted_unstable_by(|a, b| a.0.cmp(&b.0))
.map(|(k, v)| {
format!(
"{}={}",
utf8_percent_encode(k.as_ref(), &STRICT_ENCODE_SET),
utf8_percent_encode(v.as_ref(), &STRICT_ENCODE_SET)
)
})
.join("&")
}
/// Canonicalizes header into the GCP canonical form
///
/// <https://cloud.google.com/storage/docs/authentication/canonical-requests#about-headers>
fn canonicalize_headers(header_map: &HeaderMap) -> (String, String) {
//FIXME add error handling for invalid header values
let mut headers = BTreeMap::<String, Vec<&str>>::new();
for (k, v) in header_map {
headers
.entry(k.as_str().to_lowercase())
.or_default()
.push(std::str::from_utf8(v.as_bytes()).unwrap());
}
let canonicalize_headers = headers
.iter()
.map(|(k, v)| {
format!(
"{}:{}",
k.trim(),
v.iter().map(|v| trim_header_value(v)).join(",")
)
})
.join("\n");
let signed_headers = headers.keys().join(";");
(canonicalize_headers, signed_headers)
}
///construct the string to sign
///form like:
///```plaintext
///SIGNING_ALGORITHM
///ACTIVE_DATETIME
///CREDENTIAL_SCOPE
///HASHED_CANONICAL_REQUEST
///```
///`ACTIVE_DATETIME` format:`YYYYMMDD'T'HHMMSS'Z'`
/// <https://cloud.google.com/storage/docs/authentication/signatures#string-to-sign>
pub(crate) fn string_to_sign(
&self,
date: DateTime<Utc>,
request_method: &Method,
url: &Url,
headers: &HeaderMap,
) -> String {
let canonical_request = Self::canonicalize_request(url, request_method, headers);
let hashed_canonical_req = hex_digest(canonical_request.as_bytes());
let scope = self.scope(date);
format!(
"{}\n{}\n{}\n{}",
"GOOG4-RSA-SHA256",
date.format("%Y%m%dT%H%M%SZ"),
scope,
hashed_canonical_req
)
}
}
pub(crate) trait CredentialExt {
/// Apply bearer authentication to the request if the credential is not None
fn with_bearer_auth(self, credential: Option<&GcpCredential>) -> Self;
}
impl CredentialExt for HttpRequestBuilder {
fn with_bearer_auth(self, credential: Option<&GcpCredential>) -> Self {
match credential {
Some(credential) => {
if credential.bearer.is_empty() {
self
} else {
self.bearer_auth(&credential.bearer)
}
}
None => self,
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_canonicalize_headers() {
let mut input_header = HeaderMap::new();
input_header.insert("content-type", "text/plain".parse().unwrap());
input_header.insert("host", "storage.googleapis.com".parse().unwrap());
input_header.insert("x-goog-meta-reviewer", "jane".parse().unwrap());
input_header.append("x-goog-meta-reviewer", "john".parse().unwrap());
assert_eq!(
GCSAuthorizer::canonicalize_headers(&input_header),
(
"content-type:text/plain
host:storage.googleapis.com
x-goog-meta-reviewer:jane,john"
.into(),
"content-type;host;x-goog-meta-reviewer".to_string()
)
);
}
#[test]
fn test_canonicalize_query() {
let mut url = Url::parse("https://storage.googleapis.com/bucket/object").unwrap();
url.query_pairs_mut()
.append_pair("max-keys", "2")
.append_pair("prefix", "object");
assert_eq!(
GCSAuthorizer::canonicalize_query(&url),
"max-keys=2&prefix=object".to_string()
);
}
}
backend/vendor/object_store/src/gcp/mod.rs
Vendored
+431
View File
@@ -0,0 +1,431 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An object store implementation for Google Cloud Storage
//!
//! ## Multipart uploads
//!
//! [Multipart uploads](https://cloud.google.com/storage/docs/multipart-uploads)
//! can be initiated with the [ObjectStore::put_multipart] method. If neither
//! [`MultipartUpload::complete`] nor [`MultipartUpload::abort`] is invoked, you may
//! have parts uploaded to GCS but not used, that you will be charged for. It is recommended
//! you configure a [lifecycle rule] to abort incomplete multipart uploads after a certain
//! period of time to avoid being charged for storing partial uploads.
//!
//! ## Using HTTP/2
//!
//! Google Cloud Storage supports both HTTP/2 and HTTP/1. HTTP/1 is used by default
//! because it allows much higher throughput in our benchmarks (see
//! [#5194](https://github.com/apache/arrow-rs/issues/5194)). HTTP/2 can be
//! enabled by setting [crate::ClientConfigKey::Http1Only] to false.
//!
//! [lifecycle rule]: https://cloud.google.com/storage/docs/lifecycle#abort-mpu
use std::sync::Arc;
use std::time::Duration;
use crate::client::CredentialProvider;
use crate::gcp::credential::GCSAuthorizer;
use crate::signer::Signer;
use crate::{
multipart::PartId, path::Path, GetOptions, GetResult, ListPage, ListResult, MultipartId, MultipartUpload,
ObjectMeta, ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, Result,
UploadPart,
};
use async_trait::async_trait;
use client::GoogleCloudStorageClient;
use futures::stream::BoxStream;
use http::Method;
use url::Url;
use crate::client::get::GetClientExt;
use crate::client::list::ListClientExt;
use crate::client::parts::Parts;
use crate::multipart::MultipartStore;
pub use builder::{GoogleCloudStorageBuilder, GoogleConfigKey};
pub use credential::{GcpCredential, GcpSigningCredential, ServiceAccountKey};
mod builder;
mod client;
mod credential;
const STORE: &str = "GCS";
/// [`CredentialProvider`] for [`GoogleCloudStorage`]
pub type GcpCredentialProvider = Arc<dyn CredentialProvider<Credential = GcpCredential>>;
/// [`GcpSigningCredential`] for [`GoogleCloudStorage`]
pub type GcpSigningCredentialProvider =
Arc<dyn CredentialProvider<Credential = GcpSigningCredential>>;
/// Interface for [Google Cloud Storage](https://cloud.google.com/storage/).
#[derive(Debug, Clone)]
pub struct GoogleCloudStorage {
client: Arc<GoogleCloudStorageClient>,
}
impl std::fmt::Display for GoogleCloudStorage {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"GoogleCloudStorage({})",
self.client.config().bucket_name
)
}
}
impl GoogleCloudStorage {
/// Returns the [`GcpCredentialProvider`] used by [`GoogleCloudStorage`]
pub fn credentials(&self) -> &GcpCredentialProvider {
&self.client.config().credentials
}
/// Returns the [`GcpSigningCredentialProvider`] used by [`GoogleCloudStorage`]
pub fn signing_credentials(&self) -> &GcpSigningCredentialProvider {
&self.client.config().signing_credentials
}
}
#[derive(Debug)]
struct GCSMultipartUpload {
state: Arc<UploadState>,
part_idx: usize,
}
#[derive(Debug)]
struct UploadState {
client: Arc<GoogleCloudStorageClient>,
path: Path,
multipart_id: MultipartId,
parts: Parts,
}
#[async_trait]
impl MultipartUpload for GCSMultipartUpload {
fn put_part(&mut self, payload: PutPayload) -> UploadPart {
let idx = self.part_idx;
self.part_idx += 1;
let state = Arc::clone(&self.state);
Box::pin(async move {
let part = state
.client
.put_part(&state.path, &state.multipart_id, idx, payload)
.await?;
state.parts.put(idx, part);
Ok(())
})
}
async fn complete(&mut self) -> Result<PutResult> {
let parts = self.state.parts.finish(self.part_idx)?;
self.state
.client
.multipart_complete(&self.state.path, &self.state.multipart_id, parts)
.await
}
async fn abort(&mut self) -> Result<()> {
self.state
.client
.multipart_cleanup(&self.state.path, &self.state.multipart_id)
.await
}
}
#[async_trait]
impl ObjectStore for GoogleCloudStorage {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
self.client.put(location, payload, opts).await
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
let upload_id = self.client.multipart_initiate(location, opts).await?;
Ok(Box::new(GCSMultipartUpload {
part_idx: 0,
state: Arc::new(UploadState {
client: Arc::clone(&self.client),
path: location.clone(),
multipart_id: upload_id.clone(),
parts: Default::default(),
}),
}))
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
self.client.get_opts(location, options).await
}
async fn delete(&self, location: &Path) -> Result<()> {
self.client.delete_request(location).await
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
self.client.list(prefix)
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
self.client.list_with_offset(prefix, offset)
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
self.client.list_with_delimiter(prefix).await
}
async fn list_delimited_page(
&self,
prefix: Option<&Path>,
token: Option<&str>,
max_keys: Option<usize>,
) -> Result<ListPage> {
self.client.list_delimited_page(prefix, token, max_keys).await
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
self.client.copy_request(from, to, false).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
self.client.copy_request(from, to, true).await
}
}
#[async_trait]
impl MultipartStore for GoogleCloudStorage {
async fn create_multipart(&self, path: &Path) -> Result<MultipartId> {
self.client
.multipart_initiate(path, PutMultipartOpts::default())
.await
}
async fn put_part(
&self,
path: &Path,
id: &MultipartId,
part_idx: usize,
payload: PutPayload,
) -> Result<PartId> {
self.client.put_part(path, id, part_idx, payload).await
}
async fn complete_multipart(
&self,
path: &Path,
id: &MultipartId,
parts: Vec<PartId>,
) -> Result<PutResult> {
self.client.multipart_complete(path, id, parts).await
}
async fn abort_multipart(&self, path: &Path, id: &MultipartId) -> Result<()> {
self.client.multipart_cleanup(path, id).await
}
}
#[async_trait]
impl Signer for GoogleCloudStorage {
async fn signed_url(&self, method: Method, path: &Path, expires_in: Duration) -> Result<Url> {
if expires_in.as_secs() > 604800 {
return Err(crate::Error::Generic {
store: STORE,
source: "Expiration Time can't be longer than 604800 seconds (7 days).".into(),
});
}
let config = self.client.config();
let path_url = config.path_url(path);
let mut url = Url::parse(&path_url).map_err(|e| crate::Error::Generic {
store: STORE,
source: format!("Unable to parse url {path_url}: {e}").into(),
})?;
let signing_credentials = self.signing_credentials().get_credential().await?;
let authorizer = GCSAuthorizer::new(signing_credentials);
authorizer
.sign(method, &mut url, expires_in, &self.client)
.await?;
Ok(url)
}
}
#[cfg(test)]
mod test {
use credential::DEFAULT_GCS_BASE_URL;
use crate::integration::*;
use crate::tests::*;
use super::*;
const NON_EXISTENT_NAME: &str = "nonexistentname";
#[tokio::test]
async fn gcs_test() {
maybe_skip_integration!();
let integration = GoogleCloudStorageBuilder::from_env().build().unwrap();
put_get_delete_list(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
if integration.client.config().base_url == DEFAULT_GCS_BASE_URL {
// Fake GCS server doesn't currently honor ifGenerationMatch
// https://github.com/fsouza/fake-gcs-server/issues/994
copy_if_not_exists(&integration).await;
// Fake GCS server does not yet implement XML Multipart uploads
// https://github.com/fsouza/fake-gcs-server/issues/852
stream_get(&integration).await;
multipart(&integration, &integration).await;
multipart_race_condition(&integration, true).await;
multipart_out_of_order(&integration).await;
// Fake GCS server doesn't currently honor preconditions
get_opts(&integration).await;
put_opts(&integration, true).await;
// Fake GCS server doesn't currently support attributes
put_get_attributes(&integration).await;
}
}
#[tokio::test]
#[ignore]
async fn gcs_test_sign() {
maybe_skip_integration!();
let integration = GoogleCloudStorageBuilder::from_env().build().unwrap();
let client = reqwest::Client::new();
let path = Path::from("test_sign");
let url = integration
.signed_url(Method::PUT, &path, Duration::from_secs(3600))
.await
.unwrap();
println!("PUT {url}");
let resp = client.put(url).body("data").send().await.unwrap();
resp.error_for_status().unwrap();
let url = integration
.signed_url(Method::GET, &path, Duration::from_secs(3600))
.await
.unwrap();
println!("GET {url}");
let resp = client.get(url).send().await.unwrap();
let resp = resp.error_for_status().unwrap();
let data = resp.bytes().await.unwrap();
assert_eq!(data.as_ref(), b"data");
}
#[tokio::test]
async fn gcs_test_get_nonexistent_location() {
maybe_skip_integration!();
let integration = GoogleCloudStorageBuilder::from_env().build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = integration.get(&location).await.unwrap_err();
assert!(
matches!(err, crate::Error::NotFound { .. }),
"unexpected error type: {err}"
);
}
#[tokio::test]
async fn gcs_test_get_nonexistent_bucket() {
maybe_skip_integration!();
let config = GoogleCloudStorageBuilder::from_env();
let integration = config.with_bucket_name(NON_EXISTENT_NAME).build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = get_nonexistent_object(&integration, Some(location))
.await
.unwrap_err();
assert!(
matches!(err, crate::Error::NotFound { .. }),
"unexpected error type: {err}"
);
}
#[tokio::test]
async fn gcs_test_delete_nonexistent_location() {
maybe_skip_integration!();
let integration = GoogleCloudStorageBuilder::from_env().build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = integration.delete(&location).await.unwrap_err();
assert!(
matches!(err, crate::Error::NotFound { .. }),
"unexpected error type: {err}"
);
}
#[tokio::test]
async fn gcs_test_delete_nonexistent_bucket() {
maybe_skip_integration!();
let config = GoogleCloudStorageBuilder::from_env();
let integration = config.with_bucket_name(NON_EXISTENT_NAME).build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let err = integration.delete(&location).await.unwrap_err();
assert!(
matches!(err, crate::Error::NotFound { .. }),
"unexpected error type: {err}"
);
}
#[tokio::test]
async fn gcs_test_put_nonexistent_bucket() {
maybe_skip_integration!();
let config = GoogleCloudStorageBuilder::from_env();
let integration = config.with_bucket_name(NON_EXISTENT_NAME).build().unwrap();
let location = Path::from_iter([NON_EXISTENT_NAME]);
let data = PutPayload::from("arbitrary data");
let err = integration
.put(&location, data)
.await
.unwrap_err()
.to_string();
assert!(
err.contains("Server returned non-2xx status code: 404 Not Found"),
"{}",
err
)
}
}
backend/vendor/object_store/src/http/client.rs
Vendored
+478
View File
@@ -0,0 +1,478 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use crate::client::get::GetClient;
use crate::client::header::HeaderConfig;
use crate::client::retry::{self, RetryConfig, RetryExt};
use crate::client::{GetOptionsExt, HttpClient, HttpError, HttpResponse};
use crate::path::{Path, DELIMITER};
use crate::util::deserialize_rfc1123;
use crate::{Attribute, Attributes, ClientOptions, GetOptions, ObjectMeta, PutPayload, Result};
use async_trait::async_trait;
use bytes::Buf;
use chrono::{DateTime, Utc};
use http::header::{
CACHE_CONTROL, CONTENT_DISPOSITION, CONTENT_ENCODING, CONTENT_LANGUAGE, CONTENT_LENGTH,
CONTENT_TYPE,
};
use percent_encoding::percent_decode_str;
use reqwest::{Method, StatusCode};
use serde::Deserialize;
use url::Url;
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("Request error: {}", source)]
Request { source: retry::RetryError },
#[error("Request error: {}", source)]
Reqwest { source: HttpError },
#[error("Range request not supported by {}", href)]
RangeNotSupported { href: String },
#[error("Error decoding PROPFIND response: {}", source)]
InvalidPropFind { source: quick_xml::de::DeError },
#[error("Missing content size for {}", href)]
MissingSize { href: String },
#[error("Error getting properties of \"{}\" got \"{}\"", href, status)]
PropStatus { href: String, status: String },
#[error("Failed to parse href \"{}\": {}", href, source)]
InvalidHref {
href: String,
source: url::ParseError,
},
#[error("Path \"{}\" contained non-unicode characters: {}", path, source)]
NonUnicode {
path: String,
source: std::str::Utf8Error,
},
#[error("Encountered invalid path \"{}\": {}", path, source)]
InvalidPath {
path: String,
source: crate::path::Error,
},
}
impl From<Error> for crate::Error {
fn from(err: Error) -> Self {
Self::Generic {
store: "HTTP",
source: Box::new(err),
}
}
}
/// Internal client for HttpStore
#[derive(Debug)]
pub(crate) struct Client {
url: Url,
client: HttpClient,
retry_config: RetryConfig,
client_options: ClientOptions,
}
impl Client {
pub(crate) fn new(
url: Url,
client: HttpClient,
client_options: ClientOptions,
retry_config: RetryConfig,
) -> Self {
Self {
url,
retry_config,
client_options,
client,
}
}
pub(crate) fn base_url(&self) -> &Url {
&self.url
}
fn path_url(&self, location: &Path) -> String {
let mut url = self.url.clone();
url.path_segments_mut().unwrap().extend(location.parts());
url.to_string()
}
/// Create a directory with `path` using MKCOL
async fn make_directory(&self, path: &str) -> Result<(), Error> {
let method = Method::from_bytes(b"MKCOL").unwrap();
let mut url = self.url.clone();
url.path_segments_mut()
.unwrap()
.extend(path.split(DELIMITER));
self.client
.request(method, String::from(url))
.send_retry(&self.retry_config)
.await
.map_err(|source| Error::Request { source })?;
Ok(())
}
/// Recursively create parent directories
async fn create_parent_directories(&self, location: &Path) -> Result<()> {
let mut stack = vec![];
// Walk backwards until a request succeeds
let mut last_prefix = location.as_ref();
while let Some((prefix, _)) = last_prefix.rsplit_once(DELIMITER) {
last_prefix = prefix;
match self.make_directory(prefix).await {
Ok(_) => break,
Err(Error::Request { source })
if matches!(source.status(), Some(StatusCode::CONFLICT)) =>
{
// Need to create parent
stack.push(prefix)
}
Err(e) => return Err(e.into()),
}
}
// Retry the failed requests, which should now succeed
for prefix in stack.into_iter().rev() {
self.make_directory(prefix).await?;
}
Ok(())
}
pub(crate) async fn put(
&self,
location: &Path,
payload: PutPayload,
attributes: Attributes,
) -> Result<HttpResponse> {
let mut retry = false;
loop {
let url = self.path_url(location);
let mut builder = self.client.put(url);
let mut has_content_type = false;
for (k, v) in &attributes {
builder = match k {
Attribute::CacheControl => builder.header(CACHE_CONTROL, v.as_ref()),
Attribute::ContentDisposition => {
builder.header(CONTENT_DISPOSITION, v.as_ref())
}
Attribute::ContentEncoding => builder.header(CONTENT_ENCODING, v.as_ref()),
Attribute::ContentLanguage => builder.header(CONTENT_LANGUAGE, v.as_ref()),
Attribute::ContentType => {
has_content_type = true;
builder.header(CONTENT_TYPE, v.as_ref())
}
// Ignore metadata attributes
Attribute::Metadata(_) => builder,
};
}
if !has_content_type {
if let Some(value) = self.client_options.get_content_type(location) {
builder = builder.header(CONTENT_TYPE, value);
}
}
let resp = builder
.header(CONTENT_LENGTH, payload.content_length())
.retryable(&self.retry_config)
.idempotent(true)
.payload(Some(payload.clone()))
.send()
.await;
match resp {
Ok(response) => return Ok(response),
Err(source) => match source.status() {
// Some implementations return 404 instead of 409
Some(StatusCode::CONFLICT | StatusCode::NOT_FOUND) if !retry => {
retry = true;
self.create_parent_directories(location).await?
}
_ => return Err(Error::Request { source }.into()),
},
}
}
}
pub(crate) async fn list(&self, location: Option<&Path>, depth: &str) -> Result<MultiStatus> {
let url = location
.map(|path| self.path_url(path))
.unwrap_or_else(|| self.url.to_string());
let method = Method::from_bytes(b"PROPFIND").unwrap();
let result = self
.client
.request(method, url)
.header("Depth", depth)
.retryable(&self.retry_config)
.idempotent(true)
.send()
.await;
let response = match result {
Ok(result) => result
.into_body()
.bytes()
.await
.map_err(|source| Error::Reqwest { source })?,
Err(e) if matches!(e.status(), Some(StatusCode::NOT_FOUND)) => {
return match depth {
"0" => {
let path = location.map(|x| x.as_ref()).unwrap_or("");
Err(crate::Error::NotFound {
path: path.to_string(),
source: Box::new(e),
})
}
_ => {
// If prefix not found, return empty result set
Ok(Default::default())
}
};
}
Err(source) => return Err(Error::Request { source }.into()),
};
let status = quick_xml::de::from_reader(response.reader())
.map_err(|source| Error::InvalidPropFind { source })?;
Ok(status)
}
pub(crate) async fn delete(&self, path: &Path) -> Result<()> {
let url = self.path_url(path);
self.client
.delete(url)
.send_retry(&self.retry_config)
.await
.map_err(|source| match source.status() {
Some(StatusCode::NOT_FOUND) => crate::Error::NotFound {
source: Box::new(source),
path: path.to_string(),
},
_ => Error::Request { source }.into(),
})?;
Ok(())
}
pub(crate) async fn copy(&self, from: &Path, to: &Path, overwrite: bool) -> Result<()> {
let mut retry = false;
loop {
let method = Method::from_bytes(b"COPY").unwrap();
let mut builder = self
.client
.request(method, self.path_url(from))
.header("Destination", self.path_url(to).as_str());
if !overwrite {
// While the Overwrite header appears to duplicate
// the functionality of the If-Match: * header of HTTP/1.1, If-Match
// applies only to the Request-URI, and not to the Destination of a COPY
// or MOVE.
builder = builder.header("Overwrite", "F");
}
return match builder.send_retry(&self.retry_config).await {
Ok(_) => Ok(()),
Err(source) => Err(match source.status() {
Some(StatusCode::PRECONDITION_FAILED) if !overwrite => {
crate::Error::AlreadyExists {
path: to.to_string(),
source: Box::new(source),
}
}
// Some implementations return 404 instead of 409
Some(StatusCode::CONFLICT | StatusCode::NOT_FOUND) if !retry => {
retry = true;
self.create_parent_directories(to).await?;
continue;
}
_ => Error::Request { source }.into(),
}),
};
}
}
}
#[async_trait]
impl GetClient for Client {
const STORE: &'static str = "HTTP";
/// Override the [`HeaderConfig`] to be less strict to support a
/// broader range of HTTP servers (#4831)
const HEADER_CONFIG: HeaderConfig = HeaderConfig {
etag_required: false,
last_modified_required: false,
version_header: None,
user_defined_metadata_prefix: None,
};
async fn get_request(&self, path: &Path, options: GetOptions) -> Result<HttpResponse> {
let url = self.path_url(path);
let method = match options.head {
true => Method::HEAD,
false => Method::GET,
};
let has_range = options.range.is_some();
let builder = self.client.request(method, url);
let res = builder
.with_get_options(options)
.send_retry(&self.retry_config)
.await
.map_err(|source| match source.status() {
// Some stores return METHOD_NOT_ALLOWED for get on directories
Some(StatusCode::NOT_FOUND | StatusCode::METHOD_NOT_ALLOWED) => {
crate::Error::NotFound {
source: Box::new(source),
path: path.to_string(),
}
}
_ => Error::Request { source }.into(),
})?;
// We expect a 206 Partial Content response if a range was requested
// a 200 OK response would indicate the server did not fulfill the request
if has_range && res.status() != StatusCode::PARTIAL_CONTENT {
return Err(crate::Error::NotSupported {
source: Box::new(Error::RangeNotSupported {
href: path.to_string(),
}),
});
}
Ok(res)
}
}
/// The response returned by a PROPFIND request, i.e. list
#[derive(Deserialize, Default)]
pub(crate) struct MultiStatus {
pub response: Vec<MultiStatusResponse>,
}
#[derive(Deserialize)]
pub(crate) struct MultiStatusResponse {
href: String,
#[serde(rename = "propstat")]
prop_stat: PropStat,
}
impl MultiStatusResponse {
/// Returns an error if this response is not OK
pub(crate) fn check_ok(&self) -> Result<()> {
match self.prop_stat.status.contains("200 OK") {
true => Ok(()),
false => Err(Error::PropStatus {
href: self.href.clone(),
status: self.prop_stat.status.clone(),
}
.into()),
}
}
/// Returns the resolved path of this element relative to `base_url`
pub(crate) fn path(&self, base_url: &Url) -> Result<Path> {
let url = Url::options()
.base_url(Some(base_url))
.parse(&self.href)
.map_err(|source| Error::InvalidHref {
href: self.href.clone(),
source,
})?;
// Reverse any percent encoding
let path = percent_decode_str(url.path())
.decode_utf8()
.map_err(|source| Error::NonUnicode {
path: url.path().into(),
source,
})?;
Ok(Path::parse(path.as_ref()).map_err(|source| {
let path = path.into();
Error::InvalidPath { path, source }
})?)
}
fn size(&self) -> Result<u64> {
let size = self
.prop_stat
.prop
.content_length
.ok_or_else(|| Error::MissingSize {
href: self.href.clone(),
})?;
Ok(size)
}
/// Returns this objects metadata as [`ObjectMeta`]
pub(crate) fn object_meta(&self, base_url: &Url) -> Result<ObjectMeta> {
let last_modified = self.prop_stat.prop.last_modified;
Ok(ObjectMeta {
location: self.path(base_url)?,
last_modified,
size: self.size()?,
e_tag: self.prop_stat.prop.e_tag.clone(),
version: None,
})
}
/// Returns true if this is a directory / collection
pub(crate) fn is_dir(&self) -> bool {
self.prop_stat.prop.resource_type.collection.is_some()
}
}
#[derive(Deserialize)]
pub(crate) struct PropStat {
prop: Prop,
status: String,
}
#[derive(Deserialize)]
pub(crate) struct Prop {
#[serde(deserialize_with = "deserialize_rfc1123", rename = "getlastmodified")]
last_modified: DateTime<Utc>,
#[serde(rename = "getcontentlength")]
content_length: Option<u64>,
#[serde(rename = "resourcetype")]
resource_type: ResourceType,
#[serde(rename = "getetag")]
e_tag: Option<String>,
}
#[derive(Deserialize)]
pub(crate) struct ResourceType {
collection: Option<()>,
}
backend/vendor/object_store/src/http/mod.rs
Vendored
+290
View File
@@ -0,0 +1,290 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An object store implementation for generic HTTP servers
//!
//! This follows [rfc2518] commonly known as [WebDAV]
//!
//! Basic get support will work out of the box with most HTTP servers,
//! even those that don't explicitly support [rfc2518]
//!
//! Other operations such as list, delete, copy, etc... will likely
//! require server-side configuration. A list of HTTP servers with support
//! can be found [here](https://wiki.archlinux.org/title/WebDAV#Server)
//!
//! Multipart uploads are not currently supported
//!
//! [rfc2518]: https://datatracker.ietf.org/doc/html/rfc2518
//! [WebDAV]: https://en.wikipedia.org/wiki/WebDAV
use std::sync::Arc;
use async_trait::async_trait;
use futures::stream::BoxStream;
use futures::{StreamExt, TryStreamExt};
use itertools::Itertools;
use url::Url;
use crate::client::get::GetClientExt;
use crate::client::header::get_etag;
use crate::client::{http_connector, HttpConnector};
use crate::http::client::Client;
use crate::path::Path;
use crate::{
ClientConfigKey, ClientOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta,
ObjectStore, PutMode, PutMultipartOpts, PutOptions, PutPayload, PutResult, Result, RetryConfig,
};
mod client;
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("Must specify a URL")]
MissingUrl,
#[error("Unable parse source url. Url: {}, Error: {}", url, source)]
UnableToParseUrl {
source: url::ParseError,
url: String,
},
#[error("Unable to extract metadata from headers: {}", source)]
Metadata {
source: crate::client::header::Error,
},
}
impl From<Error> for crate::Error {
fn from(err: Error) -> Self {
Self::Generic {
store: "HTTP",
source: Box::new(err),
}
}
}
/// An [`ObjectStore`] implementation for generic HTTP servers
///
/// See [`crate::http`] for more information
#[derive(Debug)]
pub struct HttpStore {
client: Arc<Client>,
}
impl std::fmt::Display for HttpStore {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "HttpStore")
}
}
#[async_trait]
impl ObjectStore for HttpStore {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
if opts.mode != PutMode::Overwrite {
// TODO: Add support for If header - https://datatracker.ietf.org/doc/html/rfc2518#section-9.4
return Err(crate::Error::NotImplemented);
}
let response = self.client.put(location, payload, opts.attributes).await?;
let e_tag = match get_etag(response.headers()) {
Ok(e_tag) => Some(e_tag),
Err(crate::client::header::Error::MissingEtag) => None,
Err(source) => return Err(Error::Metadata { source }.into()),
};
Ok(PutResult {
e_tag,
version: None,
})
}
async fn put_multipart_opts(
&self,
_location: &Path,
_opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
Err(crate::Error::NotImplemented)
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
self.client.get_opts(location, options).await
}
async fn delete(&self, location: &Path) -> Result<()> {
self.client.delete(location).await
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
let prefix_len = prefix.map(|p| p.as_ref().len()).unwrap_or_default();
let prefix = prefix.cloned();
let client = Arc::clone(&self.client);
futures::stream::once(async move {
let status = client.list(prefix.as_ref(), "infinity").await?;
let iter = status
.response
.into_iter()
.filter(|r| !r.is_dir())
.map(move |response| {
response.check_ok()?;
response.object_meta(client.base_url())
})
// Filter out exact prefix matches
.filter_ok(move |r| r.location.as_ref().len() > prefix_len);
Ok::<_, crate::Error>(futures::stream::iter(iter))
})
.try_flatten()
.boxed()
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
let status = self.client.list(prefix, "1").await?;
let prefix_len = prefix.map(|p| p.as_ref().len()).unwrap_or(0);
let mut objects: Vec<ObjectMeta> = Vec::with_capacity(status.response.len());
let mut common_prefixes = Vec::with_capacity(status.response.len());
for response in status.response {
response.check_ok()?;
match response.is_dir() {
false => {
let meta = response.object_meta(self.client.base_url())?;
// Filter out exact prefix matches
if meta.location.as_ref().len() > prefix_len {
objects.push(meta);
}
}
true => {
let path = response.path(self.client.base_url())?;
// Exclude the current object
if path.as_ref().len() > prefix_len {
common_prefixes.push(path);
}
}
}
}
Ok(ListResult {
common_prefixes,
objects,
})
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
self.client.copy(from, to, true).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
self.client.copy(from, to, false).await
}
}
/// Configure a connection to a generic HTTP server
#[derive(Debug, Default, Clone)]
pub struct HttpBuilder {
url: Option<String>,
client_options: ClientOptions,
retry_config: RetryConfig,
http_connector: Option<Arc<dyn HttpConnector>>,
}
impl HttpBuilder {
/// Create a new [`HttpBuilder`] with default values.
pub fn new() -> Self {
Default::default()
}
/// Set the URL
pub fn with_url(mut self, url: impl Into<String>) -> Self {
self.url = Some(url.into());
self
}
/// Set the retry configuration
pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
self.retry_config = retry_config;
self
}
/// Set individual client configuration without overriding the entire config
pub fn with_config(mut self, key: ClientConfigKey, value: impl Into<String>) -> Self {
self.client_options = self.client_options.with_config(key, value);
self
}
/// Sets the client options, overriding any already set
pub fn with_client_options(mut self, options: ClientOptions) -> Self {
self.client_options = options;
self
}
/// The [`HttpConnector`] to use
///
/// On non-WASM32 platforms uses [`reqwest`] by default, on WASM32 platforms must be provided
pub fn with_http_connector<C: HttpConnector>(mut self, connector: C) -> Self {
self.http_connector = Some(Arc::new(connector));
self
}
/// Build an [`HttpStore`] with the configured options
pub fn build(self) -> Result<HttpStore> {
let url = self.url.ok_or(Error::MissingUrl)?;
let parsed = Url::parse(&url).map_err(|source| Error::UnableToParseUrl { url, source })?;
let client = http_connector(self.http_connector)?.connect(&self.client_options)?;
Ok(HttpStore {
client: Arc::new(Client::new(
parsed,
client,
self.client_options,
self.retry_config,
)),
})
}
}
#[cfg(test)]
mod tests {
use crate::integration::*;
use crate::tests::*;
use super::*;
#[tokio::test]
async fn http_test() {
maybe_skip_integration!();
let url = std::env::var("HTTP_URL").expect("HTTP_URL must be set");
let options = ClientOptions::new().with_allow_http(true);
let integration = HttpBuilder::new()
.with_url(url)
.with_client_options(options)
.build()
.unwrap();
put_get_delete_list(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
}
}
backend/vendor/object_store/src/integration.rs
Vendored
+1231
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/lib.rs
Vendored
+1830
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/limit.rs
Vendored
+320
View File
@@ -0,0 +1,320 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An object store that limits the maximum concurrency of the wrapped implementation
use crate::{
BoxStream, GetOptions, GetResult, GetResultPayload, ListResult, MultipartUpload, ObjectMeta,
ObjectStore, Path, PutMultipartOpts, PutOptions, PutPayload, PutResult, Result, StreamExt,
UploadPart,
};
use async_trait::async_trait;
use bytes::Bytes;
use futures::{FutureExt, Stream};
use std::ops::Range;
use std::pin::Pin;
use std::sync::Arc;
use std::task::{Context, Poll};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
/// Store wrapper that wraps an inner store and limits the maximum number of concurrent
/// object store operations. Where each call to an [`ObjectStore`] member function is
/// considered a single operation, even if it may result in more than one network call
///
/// ```
/// # use object_store::memory::InMemory;
/// # use object_store::limit::LimitStore;
///
/// // Create an in-memory `ObjectStore` limited to 20 concurrent requests
/// let store = LimitStore::new(InMemory::new(), 20);
/// ```
///
#[derive(Debug)]
pub struct LimitStore<T: ObjectStore> {
inner: Arc<T>,
max_requests: usize,
semaphore: Arc<Semaphore>,
}
impl<T: ObjectStore> LimitStore<T> {
/// Create new limit store that will limit the maximum
/// number of outstanding concurrent requests to
/// `max_requests`
pub fn new(inner: T, max_requests: usize) -> Self {
Self {
inner: Arc::new(inner),
max_requests,
semaphore: Arc::new(Semaphore::new(max_requests)),
}
}
}
impl<T: ObjectStore> std::fmt::Display for LimitStore<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "LimitStore({}, {})", self.max_requests, self.inner)
}
}
#[async_trait]
impl<T: ObjectStore> ObjectStore for LimitStore<T> {
async fn put(&self, location: &Path, payload: PutPayload) -> Result<PutResult> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.put(location, payload).await
}
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart(&self, location: &Path) -> Result<Box<dyn MultipartUpload>> {
let upload = self.inner.put_multipart(location).await?;
Ok(Box::new(LimitUpload {
semaphore: Arc::clone(&self.semaphore),
upload,
}))
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
let upload = self.inner.put_multipart_opts(location, opts).await?;
Ok(Box::new(LimitUpload {
semaphore: Arc::clone(&self.semaphore),
upload,
}))
}
async fn get(&self, location: &Path) -> Result<GetResult> {
let permit = Arc::clone(&self.semaphore).acquire_owned().await.unwrap();
let r = self.inner.get(location).await?;
Ok(permit_get_result(r, permit))
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
let permit = Arc::clone(&self.semaphore).acquire_owned().await.unwrap();
let r = self.inner.get_opts(location, options).await?;
Ok(permit_get_result(r, permit))
}
async fn get_range(&self, location: &Path, range: Range<u64>) -> Result<Bytes> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.get_range(location, range).await
}
async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.get_ranges(location, ranges).await
}
async fn head(&self, location: &Path) -> Result<ObjectMeta> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.head(location).await
}
async fn delete(&self, location: &Path) -> Result<()> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.delete(location).await
}
fn delete_stream<'a>(
&'a self,
locations: BoxStream<'a, Result<Path>>,
) -> BoxStream<'a, Result<Path>> {
self.inner.delete_stream(locations)
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
let prefix = prefix.cloned();
let inner = Arc::clone(&self.inner);
let fut = Arc::clone(&self.semaphore)
.acquire_owned()
.map(move |permit| {
let s = inner.list(prefix.as_ref());
PermitWrapper::new(s, permit.unwrap())
});
fut.into_stream().flatten().boxed()
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
let prefix = prefix.cloned();
let offset = offset.clone();
let inner = Arc::clone(&self.inner);
let fut = Arc::clone(&self.semaphore)
.acquire_owned()
.map(move |permit| {
let s = inner.list_with_offset(prefix.as_ref(), &offset);
PermitWrapper::new(s, permit.unwrap())
});
fut.into_stream().flatten().boxed()
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.list_with_delimiter(prefix).await
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.copy(from, to).await
}
async fn rename(&self, from: &Path, to: &Path) -> Result<()> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.rename(from, to).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.copy_if_not_exists(from, to).await
}
async fn rename_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
let _permit = self.semaphore.acquire().await.unwrap();
self.inner.rename_if_not_exists(from, to).await
}
}
fn permit_get_result(r: GetResult, permit: OwnedSemaphorePermit) -> GetResult {
let payload = match r.payload {
#[cfg(all(feature = "fs", not(target_arch = "wasm32")))]
v @ GetResultPayload::File(_, _) => v,
GetResultPayload::Stream(s) => {
GetResultPayload::Stream(PermitWrapper::new(s, permit).boxed())
}
};
GetResult { payload, ..r }
}
/// Combines an [`OwnedSemaphorePermit`] with some other type
struct PermitWrapper<T> {
inner: T,
#[allow(dead_code)]
permit: OwnedSemaphorePermit,
}
impl<T> PermitWrapper<T> {
fn new(inner: T, permit: OwnedSemaphorePermit) -> Self {
Self { inner, permit }
}
}
impl<T: Stream + Unpin> Stream for PermitWrapper<T> {
type Item = T::Item;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
Pin::new(&mut self.inner).poll_next(cx)
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.inner.size_hint()
}
}
/// An [`MultipartUpload`] wrapper that limits the maximum number of concurrent requests
#[derive(Debug)]
pub struct LimitUpload {
upload: Box<dyn MultipartUpload>,
semaphore: Arc<Semaphore>,
}
impl LimitUpload {
/// Create a new [`LimitUpload`] limiting `upload` to `max_concurrency` concurrent requests
pub fn new(upload: Box<dyn MultipartUpload>, max_concurrency: usize) -> Self {
Self {
upload,
semaphore: Arc::new(Semaphore::new(max_concurrency)),
}
}
}
#[async_trait]
impl MultipartUpload for LimitUpload {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
let upload = self.upload.put_part(data);
let s = Arc::clone(&self.semaphore);
Box::pin(async move {
let _permit = s.acquire().await.unwrap();
upload.await
})
}
async fn complete(&mut self) -> Result<PutResult> {
let _permit = self.semaphore.acquire().await.unwrap();
self.upload.complete().await
}
async fn abort(&mut self) -> Result<()> {
let _permit = self.semaphore.acquire().await.unwrap();
self.upload.abort().await
}
}
#[cfg(test)]
mod tests {
use crate::integration::*;
use crate::limit::LimitStore;
use crate::memory::InMemory;
use crate::ObjectStore;
use futures::stream::StreamExt;
use std::pin::Pin;
use std::time::Duration;
use tokio::time::timeout;
#[tokio::test]
async fn limit_test() {
let max_requests = 10;
let memory = InMemory::new();
let integration = LimitStore::new(memory, max_requests);
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
stream_get(&integration).await;
let mut streams = Vec::with_capacity(max_requests);
for _ in 0..max_requests {
let mut stream = integration.list(None).peekable();
Pin::new(&mut stream).peek().await; // Ensure semaphore is acquired
streams.push(stream);
}
let t = Duration::from_millis(20);
// Expect to not be able to make another request
let fut = integration.list(None).collect::<Vec<_>>();
assert!(timeout(t, fut).await.is_err());
// Drop one of the streams
streams.pop();
// Can now make another request
integration.list(None).collect::<Vec<_>>().await;
}
}
backend/vendor/object_store/src/local.rs
Vendored
+1870
View File
File diff suppressed because it is too large Load Diff
backend/vendor/object_store/src/memory.rs
Vendored
+630
View File
@@ -0,0 +1,630 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An in-memory object store implementation
use std::collections::{BTreeMap, BTreeSet, HashMap};
use std::ops::Range;
use std::sync::Arc;
use async_trait::async_trait;
use bytes::Bytes;
use chrono::{DateTime, Utc};
use futures::{stream::BoxStream, StreamExt};
use parking_lot::RwLock;
use crate::multipart::{MultipartStore, PartId};
use crate::util::InvalidGetRange;
use crate::{
path::Path, Attributes, GetRange, GetResult, GetResultPayload, ListResult, MultipartId,
MultipartUpload, ObjectMeta, ObjectStore, PutMode, PutMultipartOpts, PutOptions, PutResult,
Result, UpdateVersion, UploadPart,
};
use crate::{GetOptions, PutPayload};
/// A specialized `Error` for in-memory object store-related errors
#[derive(Debug, thiserror::Error)]
enum Error {
#[error("No data in memory found. Location: {path}")]
NoDataInMemory { path: String },
#[error("Invalid range: {source}")]
Range { source: InvalidGetRange },
#[error("Object already exists at that location: {path}")]
AlreadyExists { path: String },
#[error("ETag required for conditional update")]
MissingETag,
#[error("MultipartUpload not found: {id}")]
UploadNotFound { id: String },
#[error("Missing part at index: {part}")]
MissingPart { part: usize },
}
impl From<Error> for super::Error {
fn from(source: Error) -> Self {
match source {
Error::NoDataInMemory { ref path } => Self::NotFound {
path: path.into(),
source: source.into(),
},
Error::AlreadyExists { ref path } => Self::AlreadyExists {
path: path.into(),
source: source.into(),
},
_ => Self::Generic {
store: "InMemory",
source: Box::new(source),
},
}
}
}
/// In-memory storage suitable for testing or for opting out of using a cloud
/// storage provider.
#[derive(Debug, Default)]
pub struct InMemory {
storage: SharedStorage,
}
#[derive(Debug, Clone)]
struct Entry {
data: Bytes,
last_modified: DateTime<Utc>,
attributes: Attributes,
e_tag: usize,
}
impl Entry {
fn new(
data: Bytes,
last_modified: DateTime<Utc>,
e_tag: usize,
attributes: Attributes,
) -> Self {
Self {
data,
last_modified,
e_tag,
attributes,
}
}
}
#[derive(Debug, Default, Clone)]
struct Storage {
next_etag: usize,
map: BTreeMap<Path, Entry>,
uploads: HashMap<usize, PartStorage>,
}
#[derive(Debug, Default, Clone)]
struct PartStorage {
parts: Vec<Option<Bytes>>,
}
type SharedStorage = Arc<RwLock<Storage>>;
impl Storage {
fn insert(&mut self, location: &Path, bytes: Bytes, attributes: Attributes) -> usize {
let etag = self.next_etag;
self.next_etag += 1;
let entry = Entry::new(bytes, Utc::now(), etag, attributes);
self.overwrite(location, entry);
etag
}
fn overwrite(&mut self, location: &Path, entry: Entry) {
self.map.insert(location.clone(), entry);
}
fn create(&mut self, location: &Path, entry: Entry) -> Result<()> {
use std::collections::btree_map;
match self.map.entry(location.clone()) {
btree_map::Entry::Occupied(_) => Err(Error::AlreadyExists {
path: location.to_string(),
}
.into()),
btree_map::Entry::Vacant(v) => {
v.insert(entry);
Ok(())
}
}
}
fn update(&mut self, location: &Path, v: UpdateVersion, entry: Entry) -> Result<()> {
match self.map.get_mut(location) {
// Return Precondition instead of NotFound for consistency with stores
None => Err(crate::Error::Precondition {
path: location.to_string(),
source: format!("Object at location {location} not found").into(),
}),
Some(e) => {
let existing = e.e_tag.to_string();
let expected = v.e_tag.ok_or(Error::MissingETag)?;
if existing == expected {
*e = entry;
Ok(())
} else {
Err(crate::Error::Precondition {
path: location.to_string(),
source: format!("{existing} does not match {expected}").into(),
})
}
}
}
}
fn upload_mut(&mut self, id: &MultipartId) -> Result<&mut PartStorage> {
let parts = id
.parse()
.ok()
.and_then(|x| self.uploads.get_mut(&x))
.ok_or_else(|| Error::UploadNotFound { id: id.into() })?;
Ok(parts)
}
fn remove_upload(&mut self, id: &MultipartId) -> Result<PartStorage> {
let parts = id
.parse()
.ok()
.and_then(|x| self.uploads.remove(&x))
.ok_or_else(|| Error::UploadNotFound { id: id.into() })?;
Ok(parts)
}
}
impl std::fmt::Display for InMemory {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "InMemory")
}
}
#[async_trait]
impl ObjectStore for InMemory {
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
let mut storage = self.storage.write();
let etag = storage.next_etag;
let entry = Entry::new(payload.into(), Utc::now(), etag, opts.attributes);
match opts.mode {
PutMode::Overwrite => storage.overwrite(location, entry),
PutMode::Create => storage.create(location, entry)?,
PutMode::Update(v) => storage.update(location, v, entry)?,
}
storage.next_etag += 1;
Ok(PutResult {
e_tag: Some(etag.to_string()),
version: None,
})
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
Ok(Box::new(InMemoryUpload {
location: location.clone(),
attributes: opts.attributes,
parts: vec![],
storage: Arc::clone(&self.storage),
}))
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
let entry = self.entry(location)?;
let e_tag = entry.e_tag.to_string();
let meta = ObjectMeta {
location: location.clone(),
last_modified: entry.last_modified,
size: entry.data.len() as u64,
e_tag: Some(e_tag),
version: None,
};
options.check_preconditions(&meta)?;
let (range, data) = match options.range {
Some(range) => {
let r = range
.as_range(entry.data.len() as u64)
.map_err(|source| Error::Range { source })?;
(
r.clone(),
entry.data.slice(r.start as usize..r.end as usize),
)
}
None => (0..entry.data.len() as u64, entry.data),
};
let stream = futures::stream::once(futures::future::ready(Ok(data)));
Ok(GetResult {
payload: GetResultPayload::Stream(stream.boxed()),
attributes: entry.attributes,
meta,
range,
})
}
async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
let entry = self.entry(location)?;
ranges
.iter()
.map(|range| {
let r = GetRange::Bounded(range.clone())
.as_range(entry.data.len() as u64)
.map_err(|source| Error::Range { source })?;
let r_end = usize::try_from(r.end).map_err(|_e| Error::Range {
source: InvalidGetRange::TooLarge {
requested: r.end,
max: usize::MAX as u64,
},
})?;
let r_start = usize::try_from(r.start).map_err(|_e| Error::Range {
source: InvalidGetRange::TooLarge {
requested: r.start,
max: usize::MAX as u64,
},
})?;
Ok(entry.data.slice(r_start..r_end))
})
.collect()
}
async fn head(&self, location: &Path) -> Result<ObjectMeta> {
let entry = self.entry(location)?;
Ok(ObjectMeta {
location: location.clone(),
last_modified: entry.last_modified,
size: entry.data.len() as u64,
e_tag: Some(entry.e_tag.to_string()),
version: None,
})
}
async fn delete(&self, location: &Path) -> Result<()> {
self.storage.write().map.remove(location);
Ok(())
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
let root = Path::default();
let prefix = prefix.unwrap_or(&root);
let storage = self.storage.read();
let values: Vec<_> = storage
.map
.range((prefix)..)
.take_while(|(key, _)| key.as_ref().starts_with(prefix.as_ref()))
.filter(|(key, _)| {
// Don't return for exact prefix match
key.prefix_match(prefix)
.map(|mut x| x.next().is_some())
.unwrap_or(false)
})
.map(|(key, value)| {
Ok(ObjectMeta {
location: key.clone(),
last_modified: value.last_modified,
size: value.data.len() as u64,
e_tag: Some(value.e_tag.to_string()),
version: None,
})
})
.collect();
futures::stream::iter(values).boxed()
}
/// The memory implementation returns all results, as opposed to the cloud
/// versions which limit their results to 1k or more because of API
/// limitations.
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
let root = Path::default();
let prefix = prefix.unwrap_or(&root);
let mut common_prefixes = BTreeSet::new();
// Only objects in this base level should be returned in the
// response. Otherwise, we just collect the common prefixes.
let mut objects = vec![];
for (k, v) in self.storage.read().map.range((prefix)..) {
if !k.as_ref().starts_with(prefix.as_ref()) {
break;
}
let mut parts = match k.prefix_match(prefix) {
Some(parts) => parts,
None => continue,
};
// Pop first element
let common_prefix = match parts.next() {
Some(p) => p,
// Should only return children of the prefix
None => continue,
};
if parts.next().is_some() {
common_prefixes.insert(prefix.child(common_prefix));
} else {
let object = ObjectMeta {
location: k.clone(),
last_modified: v.last_modified,
size: v.data.len() as u64,
e_tag: Some(v.e_tag.to_string()),
version: None,
};
objects.push(object);
}
}
Ok(ListResult {
objects,
common_prefixes: common_prefixes.into_iter().collect(),
})
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
let entry = self.entry(from)?;
self.storage
.write()
.insert(to, entry.data, entry.attributes);
Ok(())
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
let entry = self.entry(from)?;
let mut storage = self.storage.write();
if storage.map.contains_key(to) {
return Err(Error::AlreadyExists {
path: to.to_string(),
}
.into());
}
storage.insert(to, entry.data, entry.attributes);
Ok(())
}
}
#[async_trait]
impl MultipartStore for InMemory {
async fn create_multipart(&self, _path: &Path) -> Result<MultipartId> {
let mut storage = self.storage.write();
let etag = storage.next_etag;
storage.next_etag += 1;
storage.uploads.insert(etag, Default::default());
Ok(etag.to_string())
}
async fn put_part(
&self,
_path: &Path,
id: &MultipartId,
part_idx: usize,
payload: PutPayload,
) -> Result<PartId> {
let mut storage = self.storage.write();
let upload = storage.upload_mut(id)?;
if part_idx <= upload.parts.len() {
upload.parts.resize(part_idx + 1, None);
}
upload.parts[part_idx] = Some(payload.into());
Ok(PartId {
content_id: Default::default(),
})
}
async fn complete_multipart(
&self,
path: &Path,
id: &MultipartId,
_parts: Vec<PartId>,
) -> Result<PutResult> {
let mut storage = self.storage.write();
let upload = storage.remove_upload(id)?;
let mut cap = 0;
for (part, x) in upload.parts.iter().enumerate() {
cap += x.as_ref().ok_or(Error::MissingPart { part })?.len();
}
let mut buf = Vec::with_capacity(cap);
for x in &upload.parts {
buf.extend_from_slice(x.as_ref().unwrap())
}
let etag = storage.insert(path, buf.into(), Default::default());
Ok(PutResult {
e_tag: Some(etag.to_string()),
version: None,
})
}
async fn abort_multipart(&self, _path: &Path, id: &MultipartId) -> Result<()> {
self.storage.write().remove_upload(id)?;
Ok(())
}
}
impl InMemory {
/// Create new in-memory storage.
pub fn new() -> Self {
Self::default()
}
/// Creates a fork of the store, with the current content copied into the
/// new store.
pub fn fork(&self) -> Self {
let storage = self.storage.read();
let storage = Arc::new(RwLock::new(storage.clone()));
Self { storage }
}
fn entry(&self, location: &Path) -> Result<Entry> {
let storage = self.storage.read();
let value = storage
.map
.get(location)
.cloned()
.ok_or_else(|| Error::NoDataInMemory {
path: location.to_string(),
})?;
Ok(value)
}
}
#[derive(Debug)]
struct InMemoryUpload {
location: Path,
attributes: Attributes,
parts: Vec<PutPayload>,
storage: Arc<RwLock<Storage>>,
}
#[async_trait]
impl MultipartUpload for InMemoryUpload {
fn put_part(&mut self, payload: PutPayload) -> UploadPart {
self.parts.push(payload);
Box::pin(futures::future::ready(Ok(())))
}
async fn complete(&mut self) -> Result<PutResult> {
let cap = self.parts.iter().map(|x| x.content_length()).sum();
let mut buf = Vec::with_capacity(cap);
let parts = self.parts.iter().flatten();
parts.for_each(|x| buf.extend_from_slice(x));
let etag = self.storage.write().insert(
&self.location,
buf.into(),
std::mem::take(&mut self.attributes),
);
Ok(PutResult {
e_tag: Some(etag.to_string()),
version: None,
})
}
async fn abort(&mut self) -> Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::integration::*;
use super::*;
#[tokio::test]
async fn in_memory_test() {
let integration = InMemory::new();
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
stream_get(&integration).await;
put_opts(&integration, true).await;
multipart(&integration, &integration).await;
put_get_attributes(&integration).await;
}
#[tokio::test]
async fn box_test() {
let integration: Box<dyn ObjectStore> = Box::new(InMemory::new());
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
stream_get(&integration).await;
}
#[tokio::test]
async fn arc_test() {
let integration: Arc<dyn ObjectStore> = Arc::new(InMemory::new());
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
stream_get(&integration).await;
}
#[tokio::test]
async fn unknown_length() {
let integration = InMemory::new();
let location = Path::from("some_file");
let data = Bytes::from("arbitrary data");
integration
.put(&location, data.clone().into())
.await
.unwrap();
let read_data = integration
.get(&location)
.await
.unwrap()
.bytes()
.await
.unwrap();
assert_eq!(&*read_data, data);
}
const NON_EXISTENT_NAME: &str = "nonexistentname";
#[tokio::test]
async fn nonexistent_location() {
let integration = InMemory::new();
let location = Path::from(NON_EXISTENT_NAME);
let err = get_nonexistent_object(&integration, Some(location))
.await
.unwrap_err();
if let crate::Error::NotFound { path, source } = err {
let source_variant = source.downcast_ref::<Error>();
assert!(
matches!(source_variant, Some(Error::NoDataInMemory { .. }),),
"got: {source_variant:?}"
);
assert_eq!(path, NON_EXISTENT_NAME);
} else {
panic!("unexpected error type: {err:?}");
}
}
}
backend/vendor/object_store/src/multipart.rs
Vendored
+84
View File
@@ -0,0 +1,84 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Cloud Multipart Upload
//!
//! This crate provides an asynchronous interface for multipart file uploads to
//! cloud storage services. It's designed to offer efficient, non-blocking operations,
//! especially useful when dealing with large files or high-throughput systems.
use async_trait::async_trait;
use crate::path::Path;
use crate::{MultipartId, PutPayload, PutResult, Result};
/// Represents a part of a file that has been successfully uploaded in a multipart upload process.
#[derive(Debug, Clone)]
pub struct PartId {
/// Id of this part
pub content_id: String,
}
/// A low-level interface for interacting with multipart upload APIs
///
/// Most use-cases should prefer [`ObjectStore::put_multipart`] as this is supported by more
/// backends, including [`LocalFileSystem`], and automatically handles uploading fixed
/// size parts of sufficient size in parallel
///
/// [`ObjectStore::put_multipart`]: crate::ObjectStore::put_multipart
/// [`LocalFileSystem`]: crate::local::LocalFileSystem
#[async_trait]
pub trait MultipartStore: Send + Sync + 'static {
/// Creates a new multipart upload, returning the [`MultipartId`]
async fn create_multipart(&self, path: &Path) -> Result<MultipartId>;
/// Uploads a new part with index `part_idx`
///
/// `part_idx` should be an integer in the range `0..N` where `N` is the number of
/// parts in the upload. Parts may be uploaded concurrently and in any order.
///
/// Most stores require that all parts excluding the last are at least 5 MiB, and some
/// further require that all parts excluding the last be the same size, e.g. [R2].
/// [`WriteMultipart`] performs writes in fixed size blocks of 5 MiB, and clients wanting
/// to maximise compatibility should look to do likewise.
///
/// [R2]: https://developers.cloudflare.com/r2/objects/multipart-objects/#limitations
/// [`WriteMultipart`]: crate::upload::WriteMultipart
async fn put_part(
&self,
path: &Path,
id: &MultipartId,
part_idx: usize,
data: PutPayload,
) -> Result<PartId>;
/// Completes a multipart upload
///
/// The `i`'th value of `parts` must be a [`PartId`] returned by a call to [`Self::put_part`]
/// with a `part_idx` of `i`, and the same `path` and `id` as provided to this method. Calling
/// this method with out of sequence or repeated [`PartId`], or [`PartId`] returned for other
/// values of `path` or `id`, will result in implementation-defined behaviour
async fn complete_multipart(
&self,
path: &Path,
id: &MultipartId,
parts: Vec<PartId>,
) -> Result<PutResult>;
/// Aborts a multipart upload
async fn abort_multipart(&self, path: &Path, id: &MultipartId) -> Result<()>;
}
backend/vendor/object_store/src/parse.rs
Vendored
+373
View File
@@ -0,0 +1,373 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
#[cfg(all(feature = "fs", not(target_arch = "wasm32")))]
use crate::local::LocalFileSystem;
use crate::memory::InMemory;
use crate::path::Path;
use crate::ObjectStore;
use url::Url;
#[derive(Debug, thiserror::Error)]
pub enum Error {
#[error("Unable to recognise URL \"{}\"", url)]
Unrecognised { url: Url },
#[error(transparent)]
Path {
#[from]
source: crate::path::Error,
},
}
impl From<Error> for super::Error {
fn from(e: Error) -> Self {
Self::Generic {
store: "URL",
source: Box::new(e),
}
}
}
/// Recognizes various URL formats, identifying the relevant [`ObjectStore`]
///
/// See [`ObjectStoreScheme::parse`] for more details
///
/// # Supported formats:
/// - `file:///path/to/my/file` -> [`LocalFileSystem`]
/// - `memory:///` -> [`InMemory`]
/// - `s3://bucket/path` -> [`AmazonS3`](crate::aws::AmazonS3) (also supports `s3a`)
/// - `gs://bucket/path` -> [`GoogleCloudStorage`](crate::gcp::GoogleCloudStorage)
/// - `az://account/container/path` -> [`MicrosoftAzure`](crate::azure::MicrosoftAzure) (also supports `adl`, `azure`, `abfs`, `abfss`)
/// - `http://mydomain/path` -> [`HttpStore`](crate::http::HttpStore)
/// - `https://mydomain/path` -> [`HttpStore`](crate::http::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` -> [`MicrosoftAzure`](crate::azure::MicrosoftAzure)
/// - `amazonaws.com` -> [`AmazonS3`](crate::aws::AmazonS3)
/// - `r2.cloudflarestorage.com` -> [`AmazonS3`](crate::aws::AmazonS3)
///
#[non_exhaustive] // permit new variants
#[derive(Debug, Eq, PartialEq, Clone)]
pub enum ObjectStoreScheme {
/// Url corresponding to [`LocalFileSystem`]
Local,
/// Url corresponding to [`InMemory`]
Memory,
/// Url corresponding to [`AmazonS3`](crate::aws::AmazonS3)
AmazonS3,
/// Url corresponding to [`GoogleCloudStorage`](crate::gcp::GoogleCloudStorage)
GoogleCloudStorage,
/// Url corresponding to [`MicrosoftAzure`](crate::azure::MicrosoftAzure)
MicrosoftAzure,
/// Url corresponding to [`HttpStore`](crate::http::HttpStore)
Http,
}
impl ObjectStoreScheme {
/// Create an [`ObjectStoreScheme`] from the provided [`Url`]
///
/// Returns the [`ObjectStoreScheme`] and the remaining [`Path`]
///
/// # Example
/// ```
/// # use url::Url;
/// # use object_store::ObjectStoreScheme;
/// let url: Url = "file:///path/to/my/file".parse().unwrap();
/// let (scheme, path) = ObjectStoreScheme::parse(&url).unwrap();
/// assert_eq!(scheme, ObjectStoreScheme::Local);
/// assert_eq!(path.as_ref(), "path/to/my/file");
///
/// let url: Url = "https://blob.core.windows.net/path/to/my/file".parse().unwrap();
/// let (scheme, path) = ObjectStoreScheme::parse(&url).unwrap();
/// assert_eq!(scheme, ObjectStoreScheme::MicrosoftAzure);
/// assert_eq!(path.as_ref(), "path/to/my/file");
///
/// let url: Url = "https://example.com/path/to/my/file".parse().unwrap();
/// let (scheme, path) = ObjectStoreScheme::parse(&url).unwrap();
/// assert_eq!(scheme, ObjectStoreScheme::Http);
/// assert_eq!(path.as_ref(), "path/to/my/file");
/// ```
pub fn parse(url: &Url) -> Result<(Self, Path), Error> {
let strip_bucket = || Some(url.path().strip_prefix('/')?.split_once('/')?.1);
let (scheme, path) = match (url.scheme(), url.host_str()) {
("file", None) => (Self::Local, url.path()),
("memory", None) => (Self::Memory, url.path()),
("s3" | "s3a", Some(_)) => (Self::AmazonS3, url.path()),
("gs", Some(_)) => (Self::GoogleCloudStorage, url.path()),
("az" | "adl" | "azure" | "abfs" | "abfss", Some(_)) => {
(Self::MicrosoftAzure, url.path())
}
("http", Some(_)) => (Self::Http, url.path()),
("https", Some(host)) => {
if host.ends_with("dfs.core.windows.net")
|| host.ends_with("blob.core.windows.net")
|| host.ends_with("dfs.fabric.microsoft.com")
|| host.ends_with("blob.fabric.microsoft.com")
{
(Self::MicrosoftAzure, url.path())
} else if host.ends_with("amazonaws.com") {
match host.starts_with("s3") {
true => (Self::AmazonS3, strip_bucket().unwrap_or_default()),
false => (Self::AmazonS3, url.path()),
}
} else if host.ends_with("r2.cloudflarestorage.com") {
(Self::AmazonS3, strip_bucket().unwrap_or_default())
} else {
(Self::Http, url.path())
}
}
_ => return Err(Error::Unrecognised { url: url.clone() }),
};
Ok((scheme, Path::from_url_path(path)?))
}
}
#[cfg(feature = "cloud")]
macro_rules! builder_opts {
($builder:ty, $url:expr, $options:expr) => {{
let builder = $options.into_iter().fold(
<$builder>::new().with_url($url.to_string()),
|builder, (key, value)| match key.as_ref().parse() {
Ok(k) => builder.with_config(k, value),
Err(_) => builder,
},
);
Box::new(builder.build()?) as _
}};
}
/// Create an [`ObjectStore`] based on the provided `url`
///
/// Returns
/// - An [`ObjectStore`] of the corresponding type
/// - The [`Path`] into the [`ObjectStore`] of the addressed resource
pub fn parse_url(url: &Url) -> Result<(Box<dyn ObjectStore>, Path), super::Error> {
parse_url_opts(url, std::iter::empty::<(&str, &str)>())
}
/// Create an [`ObjectStore`] based on the provided `url` and options
///
/// Returns
/// - An [`ObjectStore`] of the corresponding type
/// - The [`Path`] into the [`ObjectStore`] of the addressed resource
pub fn parse_url_opts<I, K, V>(
url: &Url,
options: I,
) -> Result<(Box<dyn ObjectStore>, Path), super::Error>
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<str>,
V: Into<String>,
{
let _options = options;
let (scheme, path) = ObjectStoreScheme::parse(url)?;
let path = Path::parse(path)?;
let store = match scheme {
#[cfg(all(feature = "fs", not(target_arch = "wasm32")))]
ObjectStoreScheme::Local => Box::new(LocalFileSystem::new()) as _,
ObjectStoreScheme::Memory => Box::new(InMemory::new()) as _,
#[cfg(feature = "aws")]
ObjectStoreScheme::AmazonS3 => {
builder_opts!(crate::aws::AmazonS3Builder, url, _options)
}
#[cfg(feature = "gcp")]
ObjectStoreScheme::GoogleCloudStorage => {
builder_opts!(crate::gcp::GoogleCloudStorageBuilder, url, _options)
}
#[cfg(feature = "azure")]
ObjectStoreScheme::MicrosoftAzure => {
builder_opts!(crate::azure::MicrosoftAzureBuilder, url, _options)
}
#[cfg(feature = "http")]
ObjectStoreScheme::Http => {
let url = &url[..url::Position::BeforePath];
builder_opts!(crate::http::HttpBuilder, url, _options)
}
#[cfg(not(all(
feature = "aws",
feature = "azure",
feature = "gcp",
feature = "http",
not(target_arch = "wasm32")
)))]
s => {
return Err(super::Error::Generic {
store: "parse_url",
source: format!("feature for {s:?} not enabled").into(),
})
}
};
Ok((store, path))
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_parse() {
let cases = [
("file:/path", (ObjectStoreScheme::Local, "path")),
("file:///path", (ObjectStoreScheme::Local, "path")),
("memory:/path", (ObjectStoreScheme::Memory, "path")),
("memory:///", (ObjectStoreScheme::Memory, "")),
("s3://bucket/path", (ObjectStoreScheme::AmazonS3, "path")),
("s3a://bucket/path", (ObjectStoreScheme::AmazonS3, "path")),
(
"https://s3.region.amazonaws.com/bucket",
(ObjectStoreScheme::AmazonS3, ""),
),
(
"https://s3.region.amazonaws.com/bucket/path",
(ObjectStoreScheme::AmazonS3, "path"),
),
(
"https://bucket.s3.region.amazonaws.com",
(ObjectStoreScheme::AmazonS3, ""),
),
(
"https://ACCOUNT_ID.r2.cloudflarestorage.com/bucket",
(ObjectStoreScheme::AmazonS3, ""),
),
(
"https://ACCOUNT_ID.r2.cloudflarestorage.com/bucket/path",
(ObjectStoreScheme::AmazonS3, "path"),
),
(
"abfs://container/path",
(ObjectStoreScheme::MicrosoftAzure, "path"),
),
(
"abfs://file_system@account_name.dfs.core.windows.net/path",
(ObjectStoreScheme::MicrosoftAzure, "path"),
),
(
"abfss://file_system@account_name.dfs.core.windows.net/path",
(ObjectStoreScheme::MicrosoftAzure, "path"),
),
(
"https://account.dfs.core.windows.net",
(ObjectStoreScheme::MicrosoftAzure, ""),
),
(
"https://account.blob.core.windows.net",
(ObjectStoreScheme::MicrosoftAzure, ""),
),
(
"gs://bucket/path",
(ObjectStoreScheme::GoogleCloudStorage, "path"),
),
(
"gs://test.example.com/path",
(ObjectStoreScheme::GoogleCloudStorage, "path"),
),
("http://mydomain/path", (ObjectStoreScheme::Http, "path")),
("https://mydomain/path", (ObjectStoreScheme::Http, "path")),
(
"s3://bucket/foo%20bar",
(ObjectStoreScheme::AmazonS3, "foo bar"),
),
(
"https://foo/bar%20baz",
(ObjectStoreScheme::Http, "bar baz"),
),
(
"file:///bar%252Efoo",
(ObjectStoreScheme::Local, "bar%2Efoo"),
),
(
"abfss://file_system@account.dfs.fabric.microsoft.com/",
(ObjectStoreScheme::MicrosoftAzure, ""),
),
(
"abfss://file_system@account.dfs.fabric.microsoft.com/",
(ObjectStoreScheme::MicrosoftAzure, ""),
),
(
"https://account.dfs.fabric.microsoft.com/",
(ObjectStoreScheme::MicrosoftAzure, ""),
),
(
"https://account.dfs.fabric.microsoft.com/container",
(ObjectStoreScheme::MicrosoftAzure, "container"),
),
(
"https://account.blob.fabric.microsoft.com/",
(ObjectStoreScheme::MicrosoftAzure, ""),
),
(
"https://account.blob.fabric.microsoft.com/container",
(ObjectStoreScheme::MicrosoftAzure, "container"),
),
];
for (s, (expected_scheme, expected_path)) in cases {
let url = Url::parse(s).unwrap();
let (scheme, path) = ObjectStoreScheme::parse(&url).unwrap();
assert_eq!(scheme, expected_scheme, "{s}");
assert_eq!(path, Path::parse(expected_path).unwrap(), "{s}");
}
let neg_cases = [
"unix:/run/foo.socket",
"file://remote/path",
"memory://remote/",
];
for s in neg_cases {
let url = Url::parse(s).unwrap();
assert!(ObjectStoreScheme::parse(&url).is_err());
}
}
#[test]
fn test_url_spaces() {
let url = Url::parse("file:///my file with spaces").unwrap();
assert_eq!(url.path(), "/my%20file%20with%20spaces");
let (_, path) = parse_url(&url).unwrap();
assert_eq!(path.as_ref(), "my file with spaces");
}
#[tokio::test]
#[cfg(all(feature = "http", not(target_arch = "wasm32")))]
async fn test_url_http() {
use crate::client::mock_server::MockServer;
use http::{header::USER_AGENT, Response};
let server = MockServer::new().await;
server.push_fn(|r| {
assert_eq!(r.uri().path(), "/foo/bar");
assert_eq!(r.headers().get(USER_AGENT).unwrap(), "test_url");
Response::new(String::new())
});
let test = format!("{}/foo/bar", server.url());
let opts = [("user_agent", "test_url"), ("allow_http", "true")];
let url = test.parse().unwrap();
let (store, path) = parse_url_opts(&url, opts).unwrap();
assert_eq!(path.as_ref(), "foo/bar");
store.get(&path).await.unwrap();
server.shutdown().await;
}
}
backend/vendor/object_store/src/path/mod.rs
Vendored
+614
View File
@@ -0,0 +1,614 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Path abstraction for Object Storage
use itertools::Itertools;
use percent_encoding::percent_decode;
use std::fmt::Formatter;
#[cfg(not(target_arch = "wasm32"))]
use url::Url;
/// The delimiter to separate object namespaces, creating a directory structure.
pub const DELIMITER: &str = "/";
/// The path delimiter as a single byte
pub const DELIMITER_BYTE: u8 = DELIMITER.as_bytes()[0];
mod parts;
pub use parts::{InvalidPart, PathPart};
/// Error returned by [`Path::parse`]
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum Error {
/// Error when there's an empty segment between two slashes `/` in the path
#[error("Path \"{}\" contained empty path segment", path)]
EmptySegment {
/// The source path
path: String,
},
/// Error when an invalid segment is encountered in the given path
#[error("Error parsing Path \"{}\": {}", path, source)]
BadSegment {
/// The source path
path: String,
/// The part containing the error
source: InvalidPart,
},
/// Error when path cannot be canonicalized
#[error("Failed to canonicalize path \"{}\": {}", path.display(), source)]
Canonicalize {
/// The source path
path: std::path::PathBuf,
/// The underlying error
source: std::io::Error,
},
/// Error when the path is not a valid URL
#[error("Unable to convert path \"{}\" to URL", path.display())]
InvalidPath {
/// The source path
path: std::path::PathBuf,
},
/// Error when a path contains non-unicode characters
#[error("Path \"{}\" contained non-unicode characters: {}", path, source)]
NonUnicode {
/// The source path
path: String,
/// The underlying `UTF8Error`
source: std::str::Utf8Error,
},
/// Error when the a path doesn't start with given prefix
#[error("Path {} does not start with prefix {}", path, prefix)]
PrefixMismatch {
/// The source path
path: String,
/// The mismatched prefix
prefix: String,
},
}
/// A parsed path representation that can be safely written to object storage
///
/// A [`Path`] maintains the following invariants:
///
/// * Paths are delimited by `/`
/// * Paths do not contain leading or trailing `/`
/// * Paths do not contain relative path segments, i.e. `.` or `..`
/// * Paths do not contain empty path segments
/// * Paths do not contain any ASCII control characters
///
/// There are no enforced restrictions on path length, however, it should be noted that most
/// object stores do not permit paths longer than 1024 bytes, and many filesystems do not
/// support path segments longer than 255 bytes.
///
/// # Encode
///
/// In theory object stores support any UTF-8 character sequence, however, certain character
/// sequences cause compatibility problems with some applications and protocols. Additionally
/// some filesystems may impose character restrictions, see [`LocalFileSystem`]. As such the
/// naming guidelines for [S3], [GCS] and [Azure Blob Storage] all recommend sticking to a
/// limited character subset.
///
/// [S3]: https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html
/// [GCS]: https://cloud.google.com/storage/docs/naming-objects
/// [Azure Blob Storage]: https://docs.microsoft.com/en-us/rest/api/storageservices/Naming-and-Referencing-Containers--Blobs--and-Metadata#blob-names
///
/// A string containing potentially problematic path segments can therefore be encoded to a [`Path`]
/// using [`Path::from`] or [`Path::from_iter`]. This will percent encode any problematic
/// segments according to [RFC 1738].
///
/// ```
/// # use object_store::path::Path;
/// assert_eq!(Path::from("foo/bar").as_ref(), "foo/bar");
/// assert_eq!(Path::from("foo//bar").as_ref(), "foo/bar");
/// assert_eq!(Path::from("foo/../bar").as_ref(), "foo/%2E%2E/bar");
/// assert_eq!(Path::from("/").as_ref(), "");
/// assert_eq!(Path::from_iter(["foo", "foo/bar"]).as_ref(), "foo/foo%2Fbar");
/// ```
///
/// Note: if provided with an already percent encoded string, this will encode it again
///
/// ```
/// # use object_store::path::Path;
/// assert_eq!(Path::from("foo/foo%2Fbar").as_ref(), "foo/foo%252Fbar");
/// ```
///
/// # Parse
///
/// Alternatively a [`Path`] can be parsed from an existing string, returning an
/// error if it is invalid. Unlike the encoding methods above, this will permit
/// arbitrary unicode, including percent encoded sequences.
///
/// ```
/// # use object_store::path::Path;
/// assert_eq!(Path::parse("/foo/foo%2Fbar").unwrap().as_ref(), "foo/foo%2Fbar");
/// Path::parse("..").unwrap_err(); // Relative path segments are disallowed
/// Path::parse("/foo//").unwrap_err(); // Empty path segments are disallowed
/// Path::parse("\x00").unwrap_err(); // ASCII control characters are disallowed
/// ```
///
/// [RFC 1738]: https://www.ietf.org/rfc/rfc1738.txt
/// [`LocalFileSystem`]: crate::local::LocalFileSystem
#[derive(Debug, Clone, Default, PartialEq, Eq, Hash, Ord, PartialOrd)]
pub struct Path {
/// The raw path with no leading or trailing delimiters
raw: String,
}
impl Path {
/// Parse a string as a [`Path`], returning a [`Error`] if invalid,
/// as defined on the docstring for [`Path`]
///
/// Note: this will strip any leading `/` or trailing `/`
pub fn parse(path: impl AsRef<str>) -> Result<Self, Error> {
let path = path.as_ref();
let stripped = path.strip_prefix(DELIMITER).unwrap_or(path);
if stripped.is_empty() {
return Ok(Default::default());
}
let stripped = stripped.strip_suffix(DELIMITER).unwrap_or(stripped);
for segment in stripped.split(DELIMITER) {
if segment.is_empty() {
return Err(Error::EmptySegment { path: path.into() });
}
PathPart::parse(segment).map_err(|source| {
let path = path.into();
Error::BadSegment { source, path }
})?;
}
Ok(Self {
raw: stripped.to_string(),
})
}
#[cfg(not(target_arch = "wasm32"))]
/// Convert a filesystem path to a [`Path`] relative to the filesystem root
///
/// This will return an error if the path contains illegal character sequences
/// as defined on the docstring for [`Path`] or does not exist
///
/// Note: this will canonicalize the provided path, resolving any symlinks
pub fn from_filesystem_path(path: impl AsRef<std::path::Path>) -> Result<Self, Error> {
let absolute = std::fs::canonicalize(&path).map_err(|source| {
let path = path.as_ref().into();
Error::Canonicalize { source, path }
})?;
Self::from_absolute_path(absolute)
}
#[cfg(not(target_arch = "wasm32"))]
/// Convert an absolute filesystem path to a [`Path`] relative to the filesystem root
///
/// This will return an error if the path contains illegal character sequences,
/// as defined on the docstring for [`Path`], or `base` is not an absolute path
pub fn from_absolute_path(path: impl AsRef<std::path::Path>) -> Result<Self, Error> {
Self::from_absolute_path_with_base(path, None)
}
#[cfg(not(target_arch = "wasm32"))]
/// Convert a filesystem path to a [`Path`] relative to the provided base
///
/// This will return an error if the path contains illegal character sequences,
/// as defined on the docstring for [`Path`], or `base` does not refer to a parent
/// path of `path`, or `base` is not an absolute path
pub(crate) fn from_absolute_path_with_base(
path: impl AsRef<std::path::Path>,
base: Option<&Url>,
) -> Result<Self, Error> {
let url = absolute_path_to_url(path)?;
let path = match base {
Some(prefix) => {
url.path()
.strip_prefix(prefix.path())
.ok_or_else(|| Error::PrefixMismatch {
path: url.path().to_string(),
prefix: prefix.to_string(),
})?
}
None => url.path(),
};
// Reverse any percent encoding performed by conversion to URL
Self::from_url_path(path)
}
/// Parse a url encoded string as a [`Path`], returning a [`Error`] if invalid
///
/// This will return an error if the path contains illegal character sequences
/// as defined on the docstring for [`Path`]
pub fn from_url_path(path: impl AsRef<str>) -> Result<Self, Error> {
let path = path.as_ref();
let decoded = percent_decode(path.as_bytes())
.decode_utf8()
.map_err(|source| {
let path = path.into();
Error::NonUnicode { source, path }
})?;
Self::parse(decoded)
}
/// Returns the [`PathPart`] of this [`Path`]
pub fn parts(&self) -> impl Iterator<Item = PathPart<'_>> {
self.raw
.split_terminator(DELIMITER)
.map(|s| PathPart { raw: s.into() })
}
/// Returns the last path segment containing the filename stored in this [`Path`]
pub fn filename(&self) -> Option<&str> {
match self.raw.is_empty() {
true => None,
false => self.raw.rsplit(DELIMITER).next(),
}
}
/// Returns the extension of the file stored in this [`Path`], if any
pub fn extension(&self) -> Option<&str> {
self.filename()
.and_then(|f| f.rsplit_once('.'))
.and_then(|(_, extension)| {
if extension.is_empty() {
None
} else {
Some(extension)
}
})
}
/// Returns an iterator of the [`PathPart`] of this [`Path`] after `prefix`
///
/// Returns `None` if the prefix does not match
pub fn prefix_match(&self, prefix: &Self) -> Option<impl Iterator<Item = PathPart<'_>> + '_> {
let mut stripped = self.raw.strip_prefix(&prefix.raw)?;
if !stripped.is_empty() && !prefix.raw.is_empty() {
stripped = stripped.strip_prefix(DELIMITER)?;
}
let iter = stripped
.split_terminator(DELIMITER)
.map(|x| PathPart { raw: x.into() });
Some(iter)
}
/// Returns true if this [`Path`] starts with `prefix`
pub fn prefix_matches(&self, prefix: &Self) -> bool {
self.prefix_match(prefix).is_some()
}
/// Creates a new child of this [`Path`]
pub fn child<'a>(&self, child: impl Into<PathPart<'a>>) -> Self {
let raw = match self.raw.is_empty() {
true => format!("{}", child.into().raw),
false => format!("{}{}{}", self.raw, DELIMITER, child.into().raw),
};
Self { raw }
}
}
impl AsRef<str> for Path {
fn as_ref(&self) -> &str {
&self.raw
}
}
impl From<&str> for Path {
fn from(path: &str) -> Self {
Self::from_iter(path.split(DELIMITER))
}
}
impl From<String> for Path {
fn from(path: String) -> Self {
Self::from_iter(path.split(DELIMITER))
}
}
impl From<Path> for String {
fn from(path: Path) -> Self {
path.raw
}
}
impl std::fmt::Display for Path {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
self.raw.fmt(f)
}
}
impl<'a, I> FromIterator<I> for Path
where
I: Into<PathPart<'a>>,
{
fn from_iter<T: IntoIterator<Item = I>>(iter: T) -> Self {
let raw = T::into_iter(iter)
.map(|s| s.into())
.filter(|s| !s.raw.is_empty())
.map(|s| s.raw)
.join(DELIMITER);
Self { raw }
}
}
#[cfg(not(target_arch = "wasm32"))]
/// Given an absolute filesystem path convert it to a URL representation without canonicalization
pub(crate) fn absolute_path_to_url(path: impl AsRef<std::path::Path>) -> Result<Url, Error> {
Url::from_file_path(&path).map_err(|_| Error::InvalidPath {
path: path.as_ref().into(),
})
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cloud_prefix_with_trailing_delimiter() {
// Use case: files exist in object storage named `foo/bar.json` and
// `foo_test.json`. A search for the prefix `foo/` should return
// `foo/bar.json` but not `foo_test.json'.
let prefix = Path::from_iter(["test"]);
assert_eq!(prefix.as_ref(), "test");
}
#[test]
fn push_encodes() {
let location = Path::from_iter(["foo/bar", "baz%2Ftest"]);
assert_eq!(location.as_ref(), "foo%2Fbar/baz%252Ftest");
}
#[test]
fn test_parse() {
assert_eq!(Path::parse("/").unwrap().as_ref(), "");
assert_eq!(Path::parse("").unwrap().as_ref(), "");
let err = Path::parse("//").unwrap_err();
assert!(matches!(err, Error::EmptySegment { .. }));
assert_eq!(Path::parse("/foo/bar/").unwrap().as_ref(), "foo/bar");
assert_eq!(Path::parse("foo/bar/").unwrap().as_ref(), "foo/bar");
assert_eq!(Path::parse("foo/bar").unwrap().as_ref(), "foo/bar");
let err = Path::parse("foo///bar").unwrap_err();
assert!(matches!(err, Error::EmptySegment { .. }));
}
#[test]
fn convert_raw_before_partial_eq() {
// dir and file_name
let cloud = Path::from("test_dir/test_file.json");
let built = Path::from_iter(["test_dir", "test_file.json"]);
assert_eq!(built, cloud);
// dir and file_name w/o dot
let cloud = Path::from("test_dir/test_file");
let built = Path::from_iter(["test_dir", "test_file"]);
assert_eq!(built, cloud);
// dir, no file
let cloud = Path::from("test_dir/");
let built = Path::from_iter(["test_dir"]);
assert_eq!(built, cloud);
// file_name, no dir
let cloud = Path::from("test_file.json");
let built = Path::from_iter(["test_file.json"]);
assert_eq!(built, cloud);
// empty
let cloud = Path::from("");
let built = Path::from_iter(["", ""]);
assert_eq!(built, cloud);
}
#[test]
fn parts_after_prefix_behavior() {
let existing_path = Path::from("apple/bear/cow/dog/egg.json");
// Prefix with one directory
let prefix = Path::from("apple");
let expected_parts: Vec<PathPart<'_>> = vec!["bear", "cow", "dog", "egg.json"]
.into_iter()
.map(Into::into)
.collect();
let parts: Vec<_> = existing_path.prefix_match(&prefix).unwrap().collect();
assert_eq!(parts, expected_parts);
// Prefix with two directories
let prefix = Path::from("apple/bear");
let expected_parts: Vec<PathPart<'_>> = vec!["cow", "dog", "egg.json"]
.into_iter()
.map(Into::into)
.collect();
let parts: Vec<_> = existing_path.prefix_match(&prefix).unwrap().collect();
assert_eq!(parts, expected_parts);
// Not a prefix
let prefix = Path::from("cow");
assert!(existing_path.prefix_match(&prefix).is_none());
// Prefix with a partial directory
let prefix = Path::from("ap");
assert!(existing_path.prefix_match(&prefix).is_none());
// Prefix matches but there aren't any parts after it
let existing = Path::from("apple/bear/cow/dog");
assert_eq!(existing.prefix_match(&existing).unwrap().count(), 0);
assert_eq!(Path::default().parts().count(), 0);
}
#[test]
fn prefix_matches() {
let haystack = Path::from_iter(["foo/bar", "baz%2Ftest", "something"]);
// self starts with self
assert!(
haystack.prefix_matches(&haystack),
"{haystack:?} should have started with {haystack:?}"
);
// a longer prefix doesn't match
let needle = haystack.child("longer now");
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} shouldn't have started with {needle:?}"
);
// one dir prefix matches
let needle = Path::from_iter(["foo/bar"]);
assert!(
haystack.prefix_matches(&needle),
"{haystack:?} should have started with {needle:?}"
);
// two dir prefix matches
let needle = needle.child("baz%2Ftest");
assert!(
haystack.prefix_matches(&needle),
"{haystack:?} should have started with {needle:?}"
);
// partial dir prefix doesn't match
let needle = Path::from_iter(["f"]);
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} should not have started with {needle:?}"
);
// one dir and one partial dir doesn't match
let needle = Path::from_iter(["foo/bar", "baz"]);
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} should not have started with {needle:?}"
);
// empty prefix matches
let needle = Path::from("");
assert!(
haystack.prefix_matches(&needle),
"{haystack:?} should have started with {needle:?}"
);
}
#[test]
fn prefix_matches_with_file_name() {
let haystack = Path::from_iter(["foo/bar", "baz%2Ftest", "something", "foo.segment"]);
// All directories match and file name is a prefix
let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "something", "foo"]);
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} should not have started with {needle:?}"
);
// All directories match but file name is not a prefix
let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "something", "e"]);
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} should not have started with {needle:?}"
);
// Not all directories match; file name is a prefix of the next directory; this
// does not match
let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "s"]);
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} should not have started with {needle:?}"
);
// Not all directories match; file name is NOT a prefix of the next directory;
// no match
let needle = Path::from_iter(["foo/bar", "baz%2Ftest", "p"]);
assert!(
!haystack.prefix_matches(&needle),
"{haystack:?} should not have started with {needle:?}"
);
}
#[test]
fn path_containing_spaces() {
let a = Path::from_iter(["foo bar", "baz"]);
let b = Path::from("foo bar/baz");
let c = Path::parse("foo bar/baz").unwrap();
assert_eq!(a.raw, "foo bar/baz");
assert_eq!(a.raw, b.raw);
assert_eq!(b.raw, c.raw);
}
#[test]
fn from_url_path() {
let a = Path::from_url_path("foo%20bar").unwrap();
let b = Path::from_url_path("foo/%2E%2E/bar").unwrap_err();
let c = Path::from_url_path("foo%2F%252E%252E%2Fbar").unwrap();
let d = Path::from_url_path("foo/%252E%252E/bar").unwrap();
let e = Path::from_url_path("%48%45%4C%4C%4F").unwrap();
let f = Path::from_url_path("foo/%FF/as").unwrap_err();
assert_eq!(a.raw, "foo bar");
assert!(matches!(b, Error::BadSegment { .. }));
assert_eq!(c.raw, "foo/%2E%2E/bar");
assert_eq!(d.raw, "foo/%2E%2E/bar");
assert_eq!(e.raw, "HELLO");
assert!(matches!(f, Error::NonUnicode { .. }));
}
#[test]
fn filename_from_path() {
let a = Path::from("foo/bar");
let b = Path::from("foo/bar.baz");
let c = Path::from("foo.bar/baz");
assert_eq!(a.filename(), Some("bar"));
assert_eq!(b.filename(), Some("bar.baz"));
assert_eq!(c.filename(), Some("baz"));
}
#[test]
fn file_extension() {
let a = Path::from("foo/bar");
let b = Path::from("foo/bar.baz");
let c = Path::from("foo.bar/baz");
let d = Path::from("foo.bar/baz.qux");
assert_eq!(a.extension(), None);
assert_eq!(b.extension(), Some("baz"));
assert_eq!(c.extension(), None);
assert_eq!(d.extension(), Some("qux"));
}
}
backend/vendor/object_store/src/path/parts.rs
Vendored
+175
View File
@@ -0,0 +1,175 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use percent_encoding::{percent_encode, AsciiSet, CONTROLS};
use std::borrow::Cow;
use crate::path::DELIMITER_BYTE;
/// Error returned by [`PathPart::parse`]
#[derive(Debug, thiserror::Error)]
#[error(
"Encountered illegal character sequence \"{}\" whilst parsing path segment \"{}\"",
illegal,
segment
)]
#[allow(missing_copy_implementations)]
pub struct InvalidPart {
segment: String,
illegal: String,
}
/// The PathPart type exists to validate the directory/file names that form part
/// of a path.
///
/// A [`PathPart`] is guaranteed to:
///
/// * Contain no ASCII control characters or `/`
/// * Not be a relative path segment, i.e. `.` or `..`
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Default, Hash)]
pub struct PathPart<'a> {
pub(super) raw: Cow<'a, str>,
}
impl<'a> PathPart<'a> {
/// Parse the provided path segment as a [`PathPart`] returning an error if invalid
pub fn parse(segment: &'a str) -> Result<Self, InvalidPart> {
if segment == "." || segment == ".." {
return Err(InvalidPart {
segment: segment.to_string(),
illegal: segment.to_string(),
});
}
for c in segment.chars() {
if c.is_ascii_control() || c == '/' {
return Err(InvalidPart {
segment: segment.to_string(),
// This is correct as only single byte characters up to this point
illegal: c.to_string(),
});
}
}
Ok(Self {
raw: segment.into(),
})
}
}
/// Characters we want to encode.
const INVALID: &AsciiSet = &CONTROLS
// The delimiter we are reserving for internal hierarchy
.add(DELIMITER_BYTE)
// Characters AWS recommends avoiding for object keys
// https://docs.aws.amazon.com/AmazonS3/latest/dev/UsingMetadata.html
.add(b'\\')
.add(b'{')
.add(b'^')
.add(b'}')
.add(b'%')
.add(b'`')
.add(b']')
.add(b'"') // " <-- my editor is confused about double quotes within single quotes
.add(b'>')
.add(b'[')
.add(b'~')
.add(b'<')
.add(b'#')
.add(b'|')
// Characters Google Cloud Storage recommends avoiding for object names
// https://cloud.google.com/storage/docs/naming-objects
.add(b'\r')
.add(b'\n')
.add(b'*')
.add(b'?');
impl<'a> From<&'a [u8]> for PathPart<'a> {
fn from(v: &'a [u8]) -> Self {
let inner = match v {
// We don't want to encode `.` generally, but we do want to disallow parts of paths
// to be equal to `.` or `..` to prevent file system traversal shenanigans.
b"." => "%2E".into(),
b".." => "%2E%2E".into(),
other => percent_encode(other, INVALID).into(),
};
Self { raw: inner }
}
}
impl<'a> From<&'a str> for PathPart<'a> {
fn from(v: &'a str) -> Self {
Self::from(v.as_bytes())
}
}
impl From<String> for PathPart<'static> {
fn from(s: String) -> Self {
Self {
raw: Cow::Owned(PathPart::from(s.as_str()).raw.into_owned()),
}
}
}
impl AsRef<str> for PathPart<'_> {
fn as_ref(&self) -> &str {
self.raw.as_ref()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn path_part_delimiter_gets_encoded() {
let part: PathPart<'_> = "foo/bar".into();
assert_eq!(part.raw, "foo%2Fbar");
}
#[test]
fn path_part_given_already_encoded_string() {
let part: PathPart<'_> = "foo%2Fbar".into();
assert_eq!(part.raw, "foo%252Fbar");
}
#[test]
fn path_part_cant_be_one_dot() {
let part: PathPart<'_> = ".".into();
assert_eq!(part.raw, "%2E");
}
#[test]
fn path_part_cant_be_two_dots() {
let part: PathPart<'_> = "..".into();
assert_eq!(part.raw, "%2E%2E");
}
#[test]
fn path_part_parse() {
PathPart::parse("foo").unwrap();
PathPart::parse("foo/bar").unwrap_err();
// Test percent-encoded path
PathPart::parse("foo%2Fbar").unwrap();
PathPart::parse("L%3ABC.parquet").unwrap();
// Test path containing bad escape sequence
PathPart::parse("%Z").unwrap();
PathPart::parse("%%").unwrap();
}
}
backend/vendor/object_store/src/payload.rs
Vendored
+321
View File
@@ -0,0 +1,321 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use bytes::Bytes;
use std::sync::Arc;
/// A cheaply cloneable, ordered collection of [`Bytes`]
#[derive(Debug, Clone)]
pub struct PutPayload(Arc<[Bytes]>);
impl Default for PutPayload {
fn default() -> Self {
Self(Arc::new([]))
}
}
impl PutPayload {
/// Create a new empty [`PutPayload`]
pub fn new() -> Self {
Self::default()
}
/// Creates a [`PutPayload`] from a static slice
pub fn from_static(s: &'static [u8]) -> Self {
s.into()
}
/// Creates a [`PutPayload`] from a [`Bytes`]
pub fn from_bytes(s: Bytes) -> Self {
s.into()
}
/// Returns the total length of the [`Bytes`] in this payload
pub fn content_length(&self) -> usize {
self.0.iter().map(|b| b.len()).sum()
}
/// Returns an iterator over the [`Bytes`] in this payload
pub fn iter(&self) -> PutPayloadIter<'_> {
PutPayloadIter(self.0.iter())
}
}
impl AsRef<[Bytes]> for PutPayload {
fn as_ref(&self) -> &[Bytes] {
self.0.as_ref()
}
}
impl<'a> IntoIterator for &'a PutPayload {
type Item = &'a Bytes;
type IntoIter = PutPayloadIter<'a>;
fn into_iter(self) -> Self::IntoIter {
self.iter()
}
}
impl IntoIterator for PutPayload {
type Item = Bytes;
type IntoIter = PutPayloadIntoIter;
fn into_iter(self) -> Self::IntoIter {
PutPayloadIntoIter {
payload: self,
idx: 0,
}
}
}
/// An iterator over [`PutPayload`]
#[derive(Debug)]
pub struct PutPayloadIter<'a>(std::slice::Iter<'a, Bytes>);
impl<'a> Iterator for PutPayloadIter<'a> {
type Item = &'a Bytes;
fn next(&mut self) -> Option<Self::Item> {
self.0.next()
}
fn size_hint(&self) -> (usize, Option<usize>) {
self.0.size_hint()
}
}
/// An owning iterator of [`PutPayload`]
#[derive(Debug)]
pub struct PutPayloadIntoIter {
payload: PutPayload,
idx: usize,
}
impl Iterator for PutPayloadIntoIter {
type Item = Bytes;
fn next(&mut self) -> Option<Self::Item> {
let p = self.payload.0.get(self.idx)?.clone();
self.idx += 1;
Some(p)
}
fn size_hint(&self) -> (usize, Option<usize>) {
let l = self.payload.0.len() - self.idx;
(l, Some(l))
}
}
impl From<Bytes> for PutPayload {
fn from(value: Bytes) -> Self {
Self(Arc::new([value]))
}
}
impl From<Vec<u8>> for PutPayload {
fn from(value: Vec<u8>) -> Self {
Self(Arc::new([value.into()]))
}
}
impl From<&'static str> for PutPayload {
fn from(value: &'static str) -> Self {
Bytes::from(value).into()
}
}
impl From<&'static [u8]> for PutPayload {
fn from(value: &'static [u8]) -> Self {
Bytes::from(value).into()
}
}
impl From<String> for PutPayload {
fn from(value: String) -> Self {
Bytes::from(value).into()
}
}
impl FromIterator<u8> for PutPayload {
fn from_iter<T: IntoIterator<Item = u8>>(iter: T) -> Self {
Bytes::from_iter(iter).into()
}
}
impl FromIterator<Bytes> for PutPayload {
fn from_iter<T: IntoIterator<Item = Bytes>>(iter: T) -> Self {
Self(iter.into_iter().collect())
}
}
impl From<PutPayload> for Bytes {
fn from(value: PutPayload) -> Self {
match value.0.len() {
0 => Self::new(),
1 => value.0[0].clone(),
_ => {
let mut buf = Vec::with_capacity(value.content_length());
value.iter().for_each(|x| buf.extend_from_slice(x));
buf.into()
}
}
}
}
/// A builder for [`PutPayload`] that avoids reallocating memory
///
/// Data is allocated in fixed blocks, which are flushed to [`Bytes`] once full.
/// Unlike [`Vec`] this avoids needing to repeatedly reallocate blocks of memory,
/// which typically involves copying all the previously written data to a new
/// contiguous memory region.
#[derive(Debug)]
pub struct PutPayloadMut {
len: usize,
completed: Vec<Bytes>,
in_progress: Vec<u8>,
block_size: usize,
}
impl Default for PutPayloadMut {
fn default() -> Self {
Self {
len: 0,
completed: vec![],
in_progress: vec![],
block_size: 8 * 1024,
}
}
}
impl PutPayloadMut {
/// Create a new [`PutPayloadMut`]
pub fn new() -> Self {
Self::default()
}
/// Configures the minimum allocation size
///
/// Defaults to 8KB
pub fn with_block_size(self, block_size: usize) -> Self {
Self { block_size, ..self }
}
/// Write bytes into this [`PutPayloadMut`]
///
/// If there is an in-progress block, data will be first written to it, flushing
/// it to [`Bytes`] once full. If data remains to be written, a new block of memory
/// of at least the configured block size will be allocated, to hold the remaining data.
pub fn extend_from_slice(&mut self, slice: &[u8]) {
let remaining = self.in_progress.capacity() - self.in_progress.len();
let to_copy = remaining.min(slice.len());
self.in_progress.extend_from_slice(&slice[..to_copy]);
if self.in_progress.capacity() == self.in_progress.len() {
let new_cap = self.block_size.max(slice.len() - to_copy);
let completed = std::mem::replace(&mut self.in_progress, Vec::with_capacity(new_cap));
if !completed.is_empty() {
self.completed.push(completed.into())
}
self.in_progress.extend_from_slice(&slice[to_copy..])
}
self.len += slice.len();
}
/// Append a [`Bytes`] to this [`PutPayloadMut`] without copying
///
/// This will close any currently buffered block populated by [`Self::extend_from_slice`],
/// and append `bytes` to this payload without copying.
pub fn push(&mut self, bytes: Bytes) {
if !self.in_progress.is_empty() {
let completed = std::mem::take(&mut self.in_progress);
self.completed.push(completed.into())
}
self.len += bytes.len();
self.completed.push(bytes);
}
/// Returns `true` if this [`PutPayloadMut`] contains no bytes
#[inline]
pub fn is_empty(&self) -> bool {
self.len == 0
}
/// Returns the total length of the [`Bytes`] in this payload
#[inline]
pub fn content_length(&self) -> usize {
self.len
}
/// Convert into [`PutPayload`]
pub fn freeze(mut self) -> PutPayload {
if !self.in_progress.is_empty() {
let completed = std::mem::take(&mut self.in_progress).into();
self.completed.push(completed);
}
PutPayload(self.completed.into())
}
}
impl From<PutPayloadMut> for PutPayload {
fn from(value: PutPayloadMut) -> Self {
value.freeze()
}
}
#[cfg(test)]
mod test {
use crate::PutPayloadMut;
#[test]
fn test_put_payload() {
let mut chunk = PutPayloadMut::new().with_block_size(23);
chunk.extend_from_slice(&[1; 16]);
chunk.extend_from_slice(&[2; 32]);
chunk.extend_from_slice(&[2; 5]);
chunk.extend_from_slice(&[2; 21]);
chunk.extend_from_slice(&[2; 40]);
chunk.extend_from_slice(&[0; 0]);
chunk.push("foobar".into());
let payload = chunk.freeze();
assert_eq!(payload.content_length(), 120);
let chunks = payload.as_ref();
assert_eq!(chunks.len(), 6);
assert_eq!(chunks[0].len(), 23);
assert_eq!(chunks[1].len(), 25); // 32 - (23 - 16)
assert_eq!(chunks[2].len(), 23);
assert_eq!(chunks[3].len(), 23);
assert_eq!(chunks[4].len(), 20);
assert_eq!(chunks[5].len(), 6);
}
#[test]
fn test_content_length() {
let mut chunk = PutPayloadMut::new();
chunk.push(vec![0; 23].into());
assert_eq!(chunk.content_length(), 23);
chunk.extend_from_slice(&[0; 4]);
assert_eq!(chunk.content_length(), 27);
chunk.push(vec![0; 121].into());
assert_eq!(chunk.content_length(), 148);
let payload = chunk.freeze();
assert_eq!(payload.content_length(), 148);
}
}
backend/vendor/object_store/src/prefix.rs
Vendored
+294
View File
@@ -0,0 +1,294 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! An object store wrapper handling a constant path prefix
use bytes::Bytes;
use futures::{stream::BoxStream, StreamExt, TryStreamExt};
use std::ops::Range;
use crate::path::Path;
use crate::{
GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore, PutMultipartOpts,
PutOptions, PutPayload, PutResult, Result,
};
/// Store wrapper that applies a constant prefix to all paths handled by the store.
#[derive(Debug, Clone)]
pub struct PrefixStore<T: ObjectStore> {
prefix: Path,
inner: T,
}
impl<T: ObjectStore> std::fmt::Display for PrefixStore<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "PrefixObjectStore({})", self.prefix.as_ref())
}
}
impl<T: ObjectStore> PrefixStore<T> {
/// Create a new instance of [`PrefixStore`]
pub fn new(store: T, prefix: impl Into<Path>) -> Self {
Self {
prefix: prefix.into(),
inner: store,
}
}
/// Create the full path from a path relative to prefix
fn full_path(&self, location: &Path) -> Path {
self.prefix.parts().chain(location.parts()).collect()
}
/// Strip the constant prefix from a given path
fn strip_prefix(&self, path: Path) -> Path {
// Note cannot use match because of borrow checker
if let Some(suffix) = path.prefix_match(&self.prefix) {
return suffix.collect();
}
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.
//
/// Strip the constant prefix from a given path
fn strip_prefix(prefix: &Path, path: Path) -> Path {
// Note cannot use match because of borrow checker
if let Some(suffix) = path.prefix_match(prefix) {
return suffix.collect();
}
path
}
/// Strip the constant prefix from a given ObjectMeta
fn strip_meta(prefix: &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<T: ObjectStore> ObjectStore for PrefixStore<T> {
async fn put(&self, location: &Path, payload: PutPayload) -> Result<PutResult> {
let full_path = self.full_path(location);
self.inner.put(&full_path, payload).await
}
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
let full_path = self.full_path(location);
self.inner.put_opts(&full_path, payload, opts).await
}
async fn put_multipart(&self, location: &Path) -> Result<Box<dyn MultipartUpload>> {
let full_path = self.full_path(location);
self.inner.put_multipart(&full_path).await
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
let full_path = self.full_path(location);
self.inner.put_multipart_opts(&full_path, opts).await
}
async fn get(&self, location: &Path) -> Result<GetResult> {
let full_path = self.full_path(location);
self.inner.get(&full_path).await
}
async fn get_range(&self, location: &Path, range: Range<u64>) -> Result<Bytes> {
let full_path = self.full_path(location);
self.inner.get_range(&full_path, range).await
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
let full_path = self.full_path(location);
self.inner.get_opts(&full_path, options).await
}
async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
let full_path = self.full_path(location);
self.inner.get_ranges(&full_path, ranges).await
}
async fn head(&self, location: &Path) -> Result<ObjectMeta> {
let full_path = self.full_path(location);
let meta = self.inner.head(&full_path).await?;
Ok(self.strip_meta(meta))
}
async fn delete(&self, location: &Path) -> Result<()> {
let full_path = self.full_path(location);
self.inner.delete(&full_path).await
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
let prefix = self.full_path(prefix.unwrap_or(&Path::default()));
let s = self.inner.list(Some(&prefix));
let slf_prefix = self.prefix.clone();
s.map_ok(move |meta| strip_meta(&slf_prefix, meta)).boxed()
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
let offset = self.full_path(offset);
let prefix = self.full_path(prefix.unwrap_or(&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, meta)).boxed()
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
let prefix = self.full_path(prefix.unwrap_or(&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(&self, from: &Path, to: &Path) -> Result<()> {
let full_from = self.full_path(from);
let full_to = self.full_path(to);
self.inner.copy(&full_from, &full_to).await
}
async fn rename(&self, from: &Path, to: &Path) -> Result<()> {
let full_from = self.full_path(from);
let full_to = self.full_path(to);
self.inner.rename(&full_from, &full_to).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
let full_from = self.full_path(from);
let full_to = self.full_path(to);
self.inner.copy_if_not_exists(&full_from, &full_to).await
}
async fn rename_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
let full_from = self.full_path(from);
let full_to = self.full_path(to);
self.inner.rename_if_not_exists(&full_from, &full_to).await
}
}
#[cfg(not(target_arch = "wasm32"))]
#[cfg(test)]
mod tests {
use super::*;
use crate::integration::*;
use crate::local::LocalFileSystem;
use tempfile::TempDir;
#[tokio::test]
async fn prefix_test() {
let root = TempDir::new().unwrap();
let inner = LocalFileSystem::new_with_prefix(root.path()).unwrap();
let integration = PrefixStore::new(inner, "prefix");
put_get_delete_list(&integration).await;
get_opts(&integration).await;
list_uses_directories_correctly(&integration).await;
list_with_delimiter(&integration).await;
rename_and_copy(&integration).await;
copy_if_not_exists(&integration).await;
stream_get(&integration).await;
}
#[tokio::test]
async fn prefix_test_applies_prefix() {
let tmpdir = TempDir::new().unwrap();
let local = LocalFileSystem::new_with_prefix(tmpdir.path()).unwrap();
let location = Path::from("prefix/test_file.json");
let data = Bytes::from("arbitrary data");
local.put(&location, data.clone().into()).await.unwrap();
let prefix = PrefixStore::new(local, "prefix");
let location_prefix = Path::from("test_file.json");
let content_list = flatten_list_stream(&prefix, None).await.unwrap();
assert_eq!(content_list, &[location_prefix.clone()]);
let root = Path::from("/");
let content_list = flatten_list_stream(&prefix, Some(&root)).await.unwrap();
assert_eq!(content_list, &[location_prefix.clone()]);
let read_data = prefix
.get(&location_prefix)
.await
.unwrap()
.bytes()
.await
.unwrap();
assert_eq!(&*read_data, data);
let target_prefix = Path::from("/test_written.json");
prefix
.put(&target_prefix, data.clone().into())
.await
.unwrap();
prefix.delete(&location_prefix).await.unwrap();
let local = LocalFileSystem::new_with_prefix(tmpdir.path()).unwrap();
let err = local.get(&location).await.unwrap_err();
assert!(matches!(err, crate::Error::NotFound { .. }), "{}", err);
let location = Path::from("prefix/test_written.json");
let read_data = local.get(&location).await.unwrap().bytes().await.unwrap();
assert_eq!(&*read_data, data)
}
}
backend/vendor/object_store/src/signer.rs
Vendored
+50
View File
@@ -0,0 +1,50 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Abstraction of signed URL generation for those object store implementations that support it
use crate::{path::Path, Result};
use async_trait::async_trait;
use reqwest::Method;
use std::{fmt, time::Duration};
use url::Url;
/// Universal API to generate presigned URLs from multiple object store services.
#[async_trait]
pub trait Signer: Send + Sync + fmt::Debug + 'static {
/// Given the intended [`Method`] and [`Path`] to use and the desired length of time for which
/// the URL should be valid, return a signed [`Url`] created with the object store
/// implementation's credentials such that the URL can be handed to something that doesn't have
/// access to the object store's credentials, to allow limited access to the object store.
async fn signed_url(&self, method: Method, path: &Path, expires_in: Duration) -> Result<Url>;
/// Generate signed urls for multiple paths.
///
/// See [`Signer::signed_url`] for more details.
async fn signed_urls(
&self,
method: Method,
paths: &[Path],
expires_in: Duration,
) -> Result<Vec<Url>> {
let mut urls = Vec::with_capacity(paths.len());
for path in paths {
urls.push(self.signed_url(method.clone(), path, expires_in).await?);
}
Ok(urls)
}
}
backend/vendor/object_store/src/tags.rs
Vendored
+60
View File
@@ -0,0 +1,60 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use url::form_urlencoded::Serializer;
/// A collection of key value pairs used to annotate objects
///
/// <https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-tagging.html>
/// <https://learn.microsoft.com/en-us/rest/api/storageservices/set-blob-tags>
#[derive(Debug, Clone, Default, Eq, PartialEq)]
pub struct TagSet(String);
impl TagSet {
/// Append a key value pair to this [`TagSet`]
///
/// Stores have different restrictions on what characters are permitted,
/// for portability it is recommended applications use no more than 10 tags,
/// and stick to alphanumeric characters, and `+ - = . _ : /`
///
/// <https://docs.aws.amazon.com/AmazonS3/latest/API/API_PutObjectTagging.html>
/// <https://learn.microsoft.com/en-us/rest/api/storageservices/set-blob-tags?tabs=azure-ad#request-body>
pub fn push(&mut self, key: &str, value: &str) {
Serializer::new(&mut self.0).append_pair(key, value);
}
/// Return this [`TagSet`] as a URL-encoded string
pub fn encoded(&self) -> &str {
&self.0
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_tag_set() {
let mut set = TagSet::default();
set.push("test/foo", "value sdlks");
set.push("foo", " sdf _ /+./sd");
assert_eq!(
set.encoded(),
"test%2Ffoo=value+sdlks&foo=+sdf+_+%2F%2B.%2Fsd"
);
}
}
backend/vendor/object_store/src/throttle.rs
Vendored
+659
View File
@@ -0,0 +1,659 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! A throttling object store wrapper
use parking_lot::Mutex;
use std::ops::Range;
use std::{convert::TryInto, sync::Arc};
use crate::multipart::{MultipartStore, PartId};
use crate::{
path::Path, GetResult, GetResultPayload, ListResult, MultipartId, MultipartUpload, ObjectMeta,
ObjectStore, PutMultipartOpts, PutOptions, PutPayload, PutResult, Result,
};
use crate::{GetOptions, UploadPart};
use async_trait::async_trait;
use bytes::Bytes;
use futures::{stream::BoxStream, FutureExt, StreamExt};
use std::time::Duration;
/// Configuration settings for throttled store
#[derive(Debug, Default, Clone, Copy)]
pub struct ThrottleConfig {
/// Sleep duration for every call to [`delete`](ThrottledStore::delete).
///
/// Sleeping is done before the underlying store is called and independently of the success of
/// the operation.
pub wait_delete_per_call: Duration,
/// Sleep duration for every byte received during [`get`](ThrottledStore::get).
///
/// Sleeping is performed after the underlying store returned and only for successful gets. The
/// sleep duration is additive to [`wait_get_per_call`](Self::wait_get_per_call).
///
/// Note that the per-byte sleep only happens as the user consumes the output bytes. Should
/// there be an intermediate failure (i.e. after partly consuming the output bytes), the
/// resulting sleep time will be partial as well.
pub wait_get_per_byte: Duration,
/// Sleep duration for every call to [`get`](ThrottledStore::get).
///
/// Sleeping is done before the underlying store is called and independently of the success of
/// the operation. The sleep duration is additive to
/// [`wait_get_per_byte`](Self::wait_get_per_byte).
pub wait_get_per_call: Duration,
/// Sleep duration for every call to [`list`](ThrottledStore::list).
///
/// Sleeping is done before the underlying store is called and independently of the success of
/// the operation. The sleep duration is additive to
/// [`wait_list_per_entry`](Self::wait_list_per_entry).
pub wait_list_per_call: Duration,
/// Sleep duration for every entry received during [`list`](ThrottledStore::list).
///
/// Sleeping is performed after the underlying store returned and only for successful lists.
/// The sleep duration is additive to [`wait_list_per_call`](Self::wait_list_per_call).
///
/// Note that the per-entry sleep only happens as the user consumes the output entries. Should
/// there be an intermediate failure (i.e. after partly consuming the output entries), the
/// resulting sleep time will be partial as well.
pub wait_list_per_entry: Duration,
/// Sleep duration for every call to
/// [`list_with_delimiter`](ThrottledStore::list_with_delimiter).
///
/// Sleeping is done before the underlying store is called and independently of the success of
/// the operation. The sleep duration is additive to
/// [`wait_list_with_delimiter_per_entry`](Self::wait_list_with_delimiter_per_entry).
pub wait_list_with_delimiter_per_call: Duration,
/// Sleep duration for every entry received during
/// [`list_with_delimiter`](ThrottledStore::list_with_delimiter).
///
/// Sleeping is performed after the underlying store returned and only for successful gets. The
/// sleep duration is additive to
/// [`wait_list_with_delimiter_per_call`](Self::wait_list_with_delimiter_per_call).
pub wait_list_with_delimiter_per_entry: Duration,
/// Sleep duration for every call to [`put`](ThrottledStore::put).
///
/// Sleeping is done before the underlying store is called and independently of the success of
/// the operation.
pub wait_put_per_call: Duration,
}
/// Sleep only if non-zero duration
async fn sleep(duration: Duration) {
if !duration.is_zero() {
tokio::time::sleep(duration).await
}
}
/// Store wrapper that wraps an inner store with some `sleep` calls.
///
/// This can be used for performance testing.
///
/// **Note that the behavior of the wrapper is deterministic and might not reflect real-world
/// conditions!**
#[derive(Debug)]
pub struct ThrottledStore<T> {
inner: T,
config: Arc<Mutex<ThrottleConfig>>,
}
impl<T> ThrottledStore<T> {
/// Create new wrapper with zero waiting times.
pub fn new(inner: T, config: ThrottleConfig) -> Self {
Self {
inner,
config: Arc::new(Mutex::new(config)),
}
}
/// Mutate config.
pub fn config_mut<F>(&self, f: F)
where
F: Fn(&mut ThrottleConfig),
{
let mut guard = self.config.lock();
f(&mut guard)
}
/// Return copy of current config.
pub fn config(&self) -> ThrottleConfig {
*self.config.lock()
}
}
impl<T: ObjectStore> std::fmt::Display for ThrottledStore<T> {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "ThrottledStore({})", self.inner)
}
}
#[async_trait]
impl<T: ObjectStore> ObjectStore for ThrottledStore<T> {
async fn put(&self, location: &Path, payload: PutPayload) -> Result<PutResult> {
sleep(self.config().wait_put_per_call).await;
self.inner.put(location, payload).await
}
async fn put_opts(
&self,
location: &Path,
payload: PutPayload,
opts: PutOptions,
) -> Result<PutResult> {
sleep(self.config().wait_put_per_call).await;
self.inner.put_opts(location, payload, opts).await
}
async fn put_multipart(&self, location: &Path) -> Result<Box<dyn MultipartUpload>> {
let upload = self.inner.put_multipart(location).await?;
Ok(Box::new(ThrottledUpload {
upload,
sleep: self.config().wait_put_per_call,
}))
}
async fn put_multipart_opts(
&self,
location: &Path,
opts: PutMultipartOpts,
) -> Result<Box<dyn MultipartUpload>> {
let upload = self.inner.put_multipart_opts(location, opts).await?;
Ok(Box::new(ThrottledUpload {
upload,
sleep: self.config().wait_put_per_call,
}))
}
async fn get(&self, location: &Path) -> Result<GetResult> {
sleep(self.config().wait_get_per_call).await;
// need to copy to avoid moving / referencing `self`
let wait_get_per_byte = self.config().wait_get_per_byte;
let result = self.inner.get(location).await?;
Ok(throttle_get(result, wait_get_per_byte))
}
async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
sleep(self.config().wait_get_per_call).await;
// need to copy to avoid moving / referencing `self`
let wait_get_per_byte = self.config().wait_get_per_byte;
let result = self.inner.get_opts(location, options).await?;
Ok(throttle_get(result, wait_get_per_byte))
}
async fn get_range(&self, location: &Path, range: Range<u64>) -> Result<Bytes> {
let config = self.config();
let sleep_duration =
config.wait_get_per_call + config.wait_get_per_byte * (range.end - range.start) as u32;
sleep(sleep_duration).await;
self.inner.get_range(location, range).await
}
async fn get_ranges(&self, location: &Path, ranges: &[Range<u64>]) -> Result<Vec<Bytes>> {
let config = self.config();
let total_bytes: u64 = ranges.iter().map(|range| range.end - range.start).sum();
let sleep_duration =
config.wait_get_per_call + config.wait_get_per_byte * total_bytes as u32;
sleep(sleep_duration).await;
self.inner.get_ranges(location, ranges).await
}
async fn head(&self, location: &Path) -> Result<ObjectMeta> {
sleep(self.config().wait_put_per_call).await;
self.inner.head(location).await
}
async fn delete(&self, location: &Path) -> Result<()> {
sleep(self.config().wait_delete_per_call).await;
self.inner.delete(location).await
}
fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
let stream = self.inner.list(prefix);
let config = Arc::clone(&self.config);
futures::stream::once(async move {
let config = *config.lock();
let wait_list_per_entry = config.wait_list_per_entry;
sleep(config.wait_list_per_call).await;
throttle_stream(stream, move |_| wait_list_per_entry)
})
.flatten()
.boxed()
}
fn list_with_offset(
&self,
prefix: Option<&Path>,
offset: &Path,
) -> BoxStream<'static, Result<ObjectMeta>> {
let stream = self.inner.list_with_offset(prefix, offset);
let config = Arc::clone(&self.config);
futures::stream::once(async move {
let config = *config.lock();
let wait_list_per_entry = config.wait_list_per_entry;
sleep(config.wait_list_per_call).await;
throttle_stream(stream, move |_| wait_list_per_entry)
})
.flatten()
.boxed()
}
async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
sleep(self.config().wait_list_with_delimiter_per_call).await;
match self.inner.list_with_delimiter(prefix).await {
Ok(list_result) => {
let entries_len = usize_to_u32_saturate(list_result.objects.len());
sleep(self.config().wait_list_with_delimiter_per_entry * entries_len).await;
Ok(list_result)
}
Err(err) => Err(err),
}
}
async fn copy(&self, from: &Path, to: &Path) -> Result<()> {
sleep(self.config().wait_put_per_call).await;
self.inner.copy(from, to).await
}
async fn rename(&self, from: &Path, to: &Path) -> Result<()> {
sleep(self.config().wait_put_per_call).await;
self.inner.rename(from, to).await
}
async fn copy_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
sleep(self.config().wait_put_per_call).await;
self.inner.copy_if_not_exists(from, to).await
}
async fn rename_if_not_exists(&self, from: &Path, to: &Path) -> Result<()> {
sleep(self.config().wait_put_per_call).await;
self.inner.rename_if_not_exists(from, to).await
}
}
/// Saturated `usize` to `u32` cast.
fn usize_to_u32_saturate(x: usize) -> u32 {
x.try_into().unwrap_or(u32::MAX)
}
fn throttle_get(result: GetResult, wait_get_per_byte: Duration) -> GetResult {
#[allow(clippy::infallible_destructuring_match)]
let s = match result.payload {
GetResultPayload::Stream(s) => s,
#[cfg(all(feature = "fs", not(target_arch = "wasm32")))]
GetResultPayload::File(_, _) => unimplemented!(),
};
let stream = throttle_stream(s, move |bytes| {
let bytes_len: u32 = usize_to_u32_saturate(bytes.len());
wait_get_per_byte * bytes_len
});
GetResult {
payload: GetResultPayload::Stream(stream),
..result
}
}
fn throttle_stream<T: Send + 'static, E: Send + 'static, F>(
stream: BoxStream<'_, Result<T, E>>,
delay: F,
) -> BoxStream<'_, Result<T, E>>
where
F: Fn(&T) -> Duration + Send + Sync + 'static,
{
stream
.then(move |result| {
let delay = result.as_ref().ok().map(&delay).unwrap_or_default();
sleep(delay).then(|_| futures::future::ready(result))
})
.boxed()
}
#[async_trait]
impl<T: MultipartStore> MultipartStore for ThrottledStore<T> {
async fn create_multipart(&self, path: &Path) -> Result<MultipartId> {
self.inner.create_multipart(path).await
}
async fn put_part(
&self,
path: &Path,
id: &MultipartId,
part_idx: usize,
data: PutPayload,
) -> Result<PartId> {
sleep(self.config().wait_put_per_call).await;
self.inner.put_part(path, id, part_idx, data).await
}
async fn complete_multipart(
&self,
path: &Path,
id: &MultipartId,
parts: Vec<PartId>,
) -> Result<PutResult> {
self.inner.complete_multipart(path, id, parts).await
}
async fn abort_multipart(&self, path: &Path, id: &MultipartId) -> Result<()> {
self.inner.abort_multipart(path, id).await
}
}
#[derive(Debug)]
struct ThrottledUpload {
upload: Box<dyn MultipartUpload>,
sleep: Duration,
}
#[async_trait]
impl MultipartUpload for ThrottledUpload {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
let duration = self.sleep;
let put = self.upload.put_part(data);
Box::pin(async move {
sleep(duration).await;
put.await
})
}
async fn complete(&mut self) -> Result<PutResult> {
self.upload.complete().await
}
async fn abort(&mut self) -> Result<()> {
self.upload.abort().await
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{integration::*, memory::InMemory, GetResultPayload};
use futures::TryStreamExt;
use tokio::time::Duration;
use tokio::time::Instant;
const WAIT_TIME: Duration = Duration::from_millis(100);
const ZERO: Duration = Duration::from_millis(0); // Duration::default isn't constant
macro_rules! assert_bounds {
($d:expr, $lower:expr) => {
assert_bounds!($d, $lower, $lower + 2);
};
($d:expr, $lower:expr, $upper:expr) => {
let d = $d;
let lower = $lower * WAIT_TIME;
let upper = $upper * WAIT_TIME;
assert!(d >= lower, "{:?} must be >= than {:?}", d, lower);
assert!(d < upper, "{:?} must be < than {:?}", d, upper);
};
}
#[tokio::test]
async fn throttle_test() {
let inner = InMemory::new();
let store = ThrottledStore::new(inner, ThrottleConfig::default());
put_get_delete_list(&store).await;
list_uses_directories_correctly(&store).await;
list_with_delimiter(&store).await;
rename_and_copy(&store).await;
copy_if_not_exists(&store).await;
stream_get(&store).await;
multipart(&store, &store).await;
}
#[tokio::test]
async fn delete_test() {
let inner = InMemory::new();
let store = ThrottledStore::new(inner, ThrottleConfig::default());
assert_bounds!(measure_delete(&store, None).await, 0);
assert_bounds!(measure_delete(&store, Some(0)).await, 0);
assert_bounds!(measure_delete(&store, Some(10)).await, 0);
store.config_mut(|cfg| cfg.wait_delete_per_call = WAIT_TIME);
assert_bounds!(measure_delete(&store, None).await, 1);
assert_bounds!(measure_delete(&store, Some(0)).await, 1);
assert_bounds!(measure_delete(&store, Some(10)).await, 1);
}
#[tokio::test]
// macos github runner is so slow it can't complete within WAIT_TIME*2
#[cfg(target_os = "linux")]
async fn get_test() {
let inner = InMemory::new();
let store = ThrottledStore::new(inner, ThrottleConfig::default());
assert_bounds!(measure_get(&store, None).await, 0);
assert_bounds!(measure_get(&store, Some(0)).await, 0);
assert_bounds!(measure_get(&store, Some(10)).await, 0);
store.config_mut(|cfg| cfg.wait_get_per_call = WAIT_TIME);
assert_bounds!(measure_get(&store, None).await, 1);
assert_bounds!(measure_get(&store, Some(0)).await, 1);
assert_bounds!(measure_get(&store, Some(10)).await, 1);
store.config_mut(|cfg| {
cfg.wait_get_per_call = ZERO;
cfg.wait_get_per_byte = WAIT_TIME;
});
assert_bounds!(measure_get(&store, Some(2)).await, 2);
store.config_mut(|cfg| {
cfg.wait_get_per_call = WAIT_TIME;
cfg.wait_get_per_byte = WAIT_TIME;
});
assert_bounds!(measure_get(&store, Some(2)).await, 3);
}
#[tokio::test]
// macos github runner is so slow it can't complete within WAIT_TIME*2
#[cfg(target_os = "linux")]
async fn list_test() {
let inner = InMemory::new();
let store = ThrottledStore::new(inner, ThrottleConfig::default());
assert_bounds!(measure_list(&store, 0).await, 0);
assert_bounds!(measure_list(&store, 10).await, 0);
store.config_mut(|cfg| cfg.wait_list_per_call = WAIT_TIME);
assert_bounds!(measure_list(&store, 0).await, 1);
assert_bounds!(measure_list(&store, 10).await, 1);
store.config_mut(|cfg| {
cfg.wait_list_per_call = ZERO;
cfg.wait_list_per_entry = WAIT_TIME;
});
assert_bounds!(measure_list(&store, 2).await, 2);
store.config_mut(|cfg| {
cfg.wait_list_per_call = WAIT_TIME;
cfg.wait_list_per_entry = WAIT_TIME;
});
assert_bounds!(measure_list(&store, 2).await, 3);
}
#[tokio::test]
// macos github runner is so slow it can't complete within WAIT_TIME*2
#[cfg(target_os = "linux")]
async fn list_with_delimiter_test() {
let inner = InMemory::new();
let store = ThrottledStore::new(inner, ThrottleConfig::default());
assert_bounds!(measure_list_with_delimiter(&store, 0).await, 0);
assert_bounds!(measure_list_with_delimiter(&store, 10).await, 0);
store.config_mut(|cfg| cfg.wait_list_with_delimiter_per_call = WAIT_TIME);
assert_bounds!(measure_list_with_delimiter(&store, 0).await, 1);
assert_bounds!(measure_list_with_delimiter(&store, 10).await, 1);
store.config_mut(|cfg| {
cfg.wait_list_with_delimiter_per_call = ZERO;
cfg.wait_list_with_delimiter_per_entry = WAIT_TIME;
});
assert_bounds!(measure_list_with_delimiter(&store, 2).await, 2);
store.config_mut(|cfg| {
cfg.wait_list_with_delimiter_per_call = WAIT_TIME;
cfg.wait_list_with_delimiter_per_entry = WAIT_TIME;
});
assert_bounds!(measure_list_with_delimiter(&store, 2).await, 3);
}
#[tokio::test]
async fn put_test() {
let inner = InMemory::new();
let store = ThrottledStore::new(inner, ThrottleConfig::default());
assert_bounds!(measure_put(&store, 0).await, 0);
assert_bounds!(measure_put(&store, 10).await, 0);
store.config_mut(|cfg| cfg.wait_put_per_call = WAIT_TIME);
assert_bounds!(measure_put(&store, 0).await, 1);
assert_bounds!(measure_put(&store, 10).await, 1);
store.config_mut(|cfg| cfg.wait_put_per_call = ZERO);
assert_bounds!(measure_put(&store, 0).await, 0);
}
async fn place_test_object(store: &ThrottledStore<InMemory>, n_bytes: Option<usize>) -> Path {
let path = Path::from("foo");
if let Some(n_bytes) = n_bytes {
let data: Vec<_> = std::iter::repeat(1u8).take(n_bytes).collect();
store.put(&path, data.into()).await.unwrap();
} else {
// ensure object is absent
store.delete(&path).await.unwrap();
}
path
}
#[allow(dead_code)]
async fn place_test_objects(store: &ThrottledStore<InMemory>, n_entries: usize) -> Path {
let prefix = Path::from("foo");
// clean up store
let entries: Vec<_> = store.list(Some(&prefix)).try_collect().await.unwrap();
for entry in entries {
store.delete(&entry.location).await.unwrap();
}
// create new entries
for i in 0..n_entries {
let path = prefix.child(i.to_string().as_str());
store.put(&path, "bar".into()).await.unwrap();
}
prefix
}
async fn measure_delete(store: &ThrottledStore<InMemory>, n_bytes: Option<usize>) -> Duration {
let path = place_test_object(store, n_bytes).await;
let t0 = Instant::now();
store.delete(&path).await.unwrap();
t0.elapsed()
}
#[allow(dead_code)]
#[cfg(target_os = "linux")]
async fn measure_get(store: &ThrottledStore<InMemory>, n_bytes: Option<usize>) -> Duration {
let path = place_test_object(store, n_bytes).await;
let t0 = Instant::now();
let res = store.get(&path).await;
if n_bytes.is_some() {
// need to consume bytes to provoke sleep times
let s = match res.unwrap().payload {
GetResultPayload::Stream(s) => s,
GetResultPayload::File(_, _) => unimplemented!(),
};
s.map_ok(|b| bytes::BytesMut::from(&b[..]))
.try_concat()
.await
.unwrap();
} else {
assert!(res.is_err());
}
t0.elapsed()
}
#[allow(dead_code)]
async fn measure_list(store: &ThrottledStore<InMemory>, n_entries: usize) -> Duration {
let prefix = place_test_objects(store, n_entries).await;
let t0 = Instant::now();
store
.list(Some(&prefix))
.try_collect::<Vec<_>>()
.await
.unwrap();
t0.elapsed()
}
#[allow(dead_code)]
async fn measure_list_with_delimiter(
store: &ThrottledStore<InMemory>,
n_entries: usize,
) -> Duration {
let prefix = place_test_objects(store, n_entries).await;
let t0 = Instant::now();
store.list_with_delimiter(Some(&prefix)).await.unwrap();
t0.elapsed()
}
async fn measure_put(store: &ThrottledStore<InMemory>, n_bytes: usize) -> Duration {
let data: Vec<_> = std::iter::repeat(1u8).take(n_bytes).collect();
let t0 = Instant::now();
store.put(&Path::from("foo"), data.into()).await.unwrap();
t0.elapsed()
}
}
backend/vendor/object_store/src/upload.rs
Vendored
+341
View File
@@ -0,0 +1,341 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
use std::task::{Context, Poll};
use crate::{PutPayload, PutPayloadMut, PutResult, Result};
use async_trait::async_trait;
use bytes::Bytes;
use futures::future::BoxFuture;
use futures::ready;
use tokio::task::JoinSet;
/// An upload part request
pub type UploadPart = BoxFuture<'static, Result<()>>;
/// A trait allowing writing an object in fixed size chunks
///
/// Consecutive chunks of data can be written by calling [`MultipartUpload::put_part`] and polling
/// the returned futures to completion. Multiple futures returned by [`MultipartUpload::put_part`]
/// may be polled in parallel, allowing for concurrent uploads.
///
/// Once all part uploads have been polled to completion, the upload can be completed by
/// calling [`MultipartUpload::complete`]. This will make the entire uploaded object visible
/// as an atomic operation.It is implementation behind behaviour if [`MultipartUpload::complete`]
/// is called before all [`UploadPart`] have been polled to completion.
#[async_trait]
pub trait MultipartUpload: Send + std::fmt::Debug {
/// Upload the next part
///
/// Most stores require that all parts excluding the last are at least 5 MiB, and some
/// further require that all parts excluding the last be the same size, e.g. [R2].
/// Clients wanting to maximise compatibility should therefore perform writes in
/// fixed size blocks larger than 5 MiB.
///
/// Implementations may invoke this method multiple times and then await on the
/// returned futures in parallel
///
/// ```no_run
/// # use futures::StreamExt;
/// # use object_store::MultipartUpload;
/// #
/// # async fn test() {
/// #
/// let mut upload: Box<&dyn MultipartUpload> = todo!();
/// let p1 = upload.put_part(vec![0; 10 * 1024 * 1024].into());
/// let p2 = upload.put_part(vec![1; 10 * 1024 * 1024].into());
/// futures::future::try_join(p1, p2).await.unwrap();
/// upload.complete().await.unwrap();
/// # }
/// ```
///
/// [R2]: https://developers.cloudflare.com/r2/objects/multipart-objects/#limitations
fn put_part(&mut self, data: PutPayload) -> UploadPart;
/// Complete the multipart upload
///
/// It is implementation defined behaviour if this method is called before polling
/// all [`UploadPart`] returned by [`MultipartUpload::put_part`] to completion. Additionally,
/// it is implementation defined behaviour to call [`MultipartUpload::complete`]
/// on an already completed or aborted [`MultipartUpload`].
async fn complete(&mut self) -> Result<PutResult>;
/// Abort the multipart upload
///
/// If a [`MultipartUpload`] is dropped without calling [`MultipartUpload::complete`],
/// some object stores will automatically clean up any previously uploaded parts.
/// However, some stores, such as S3 and GCS, cannot perform cleanup on drop.
/// As such [`MultipartUpload::abort`] can be invoked to perform this cleanup.
///
/// It will not be possible to call `abort` in all failure scenarios, for example
/// non-graceful shutdown of the calling application. It is therefore recommended
/// object stores are configured with lifecycle rules to automatically cleanup
/// unused parts older than some threshold. See [crate::aws] and [crate::gcp]
/// for more information.
///
/// It is implementation defined behaviour to call [`MultipartUpload::abort`]
/// on an already completed or aborted [`MultipartUpload`]
async fn abort(&mut self) -> Result<()>;
}
#[async_trait]
impl<W: MultipartUpload + ?Sized> MultipartUpload for Box<W> {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
(**self).put_part(data)
}
async fn complete(&mut self) -> Result<PutResult> {
(**self).complete().await
}
async fn abort(&mut self) -> Result<()> {
(**self).abort().await
}
}
/// A synchronous write API for uploading data in parallel in fixed size chunks
///
/// Uses multiple tokio tasks in a [`JoinSet`] to multiplex upload tasks in parallel
///
/// The design also takes inspiration from [`Sink`] with [`WriteMultipart::wait_for_capacity`]
/// allowing back pressure on producers, prior to buffering the next part. However, unlike
/// [`Sink`] this back pressure is optional, allowing integration with synchronous producers
///
/// [`Sink`]: futures::sink::Sink
#[derive(Debug)]
pub struct WriteMultipart {
upload: Box<dyn MultipartUpload>,
buffer: PutPayloadMut,
chunk_size: usize,
tasks: JoinSet<Result<()>>,
}
impl WriteMultipart {
/// Create a new [`WriteMultipart`] that will upload using 5MB chunks
pub fn new(upload: Box<dyn MultipartUpload>) -> Self {
Self::new_with_chunk_size(upload, 5 * 1024 * 1024)
}
/// Create a new [`WriteMultipart`] that will upload in fixed `chunk_size` sized chunks
pub fn new_with_chunk_size(upload: Box<dyn MultipartUpload>, chunk_size: usize) -> Self {
Self {
upload,
chunk_size,
buffer: PutPayloadMut::new(),
tasks: Default::default(),
}
}
/// Polls for there to be less than `max_concurrency` [`UploadPart`] in progress
///
/// See [`Self::wait_for_capacity`] for an async version of this function
pub fn poll_for_capacity(
&mut self,
cx: &mut Context<'_>,
max_concurrency: usize,
) -> Poll<Result<()>> {
while !self.tasks.is_empty() && self.tasks.len() >= max_concurrency {
ready!(self.tasks.poll_join_next(cx)).unwrap()??
}
Poll::Ready(Ok(()))
}
/// Wait until there are less than `max_concurrency` [`UploadPart`] in progress
///
/// See [`Self::poll_for_capacity`] for a [`Poll`] version of this function
pub async fn wait_for_capacity(&mut self, max_concurrency: usize) -> Result<()> {
futures::future::poll_fn(|cx| self.poll_for_capacity(cx, max_concurrency)).await
}
/// Write data to this [`WriteMultipart`]
///
/// Data is buffered using [`PutPayloadMut::extend_from_slice`]. Implementations looking to
/// write data from owned buffers may prefer [`Self::put`] as this avoids copying.
///
/// Note this method is synchronous (not `async`) and will immediately
/// start new uploads as soon as the internal `chunk_size` is hit,
/// regardless of how many outstanding uploads are already in progress.
///
/// Back pressure can optionally be applied to producers by calling
/// [`Self::wait_for_capacity`] prior to calling this method
pub fn write(&mut self, mut buf: &[u8]) {
while !buf.is_empty() {
let remaining = self.chunk_size - self.buffer.content_length();
let to_read = buf.len().min(remaining);
self.buffer.extend_from_slice(&buf[..to_read]);
if to_read == remaining {
let buffer = std::mem::take(&mut self.buffer);
self.put_part(buffer.into())
}
buf = &buf[to_read..]
}
}
/// Put a chunk of data into this [`WriteMultipart`] without copying
///
/// Data is buffered using [`PutPayloadMut::push`]. Implementations looking to
/// perform writes from non-owned buffers should prefer [`Self::write`] as this
/// will allow multiple calls to share the same underlying allocation.
///
/// See [`Self::write`] for information on backpressure
pub fn put(&mut self, mut bytes: Bytes) {
while !bytes.is_empty() {
let remaining = self.chunk_size - self.buffer.content_length();
if bytes.len() < remaining {
self.buffer.push(bytes);
return;
}
self.buffer.push(bytes.split_to(remaining));
let buffer = std::mem::take(&mut self.buffer);
self.put_part(buffer.into())
}
}
pub(crate) fn put_part(&mut self, part: PutPayload) {
self.tasks.spawn(self.upload.put_part(part));
}
/// Abort this upload, attempting to clean up any successfully uploaded parts
pub async fn abort(mut self) -> Result<()> {
self.tasks.shutdown().await;
self.upload.abort().await
}
/// Flush final chunk, and await completion of all in-flight requests
pub async fn finish(mut self) -> Result<PutResult> {
if !self.buffer.is_empty() {
let part = std::mem::take(&mut self.buffer);
self.put_part(part.into())
}
self.wait_for_capacity(0).await?;
match self.upload.complete().await {
Err(e) => {
self.tasks.shutdown().await;
self.upload.abort().await?;
Err(e)
}
Ok(result) => Ok(result),
}
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::Duration;
use futures::FutureExt;
use parking_lot::Mutex;
use rand::prelude::StdRng;
use rand::{Rng, SeedableRng};
use crate::memory::InMemory;
use crate::path::Path;
use crate::throttle::{ThrottleConfig, ThrottledStore};
use crate::ObjectStore;
use super::*;
#[tokio::test]
async fn test_concurrency() {
let config = ThrottleConfig {
wait_put_per_call: Duration::from_millis(1),
..Default::default()
};
let path = Path::from("foo");
let store = ThrottledStore::new(InMemory::new(), config);
let upload = store.put_multipart(&path).await.unwrap();
let mut write = WriteMultipart::new_with_chunk_size(upload, 10);
for _ in 0..20 {
write.write(&[0; 5]);
}
assert!(write.wait_for_capacity(10).now_or_never().is_none());
write.wait_for_capacity(10).await.unwrap()
}
#[derive(Debug, Default)]
struct InstrumentedUpload {
chunks: Arc<Mutex<Vec<PutPayload>>>,
}
#[async_trait]
impl MultipartUpload for InstrumentedUpload {
fn put_part(&mut self, data: PutPayload) -> UploadPart {
self.chunks.lock().push(data);
futures::future::ready(Ok(())).boxed()
}
async fn complete(&mut self) -> Result<PutResult> {
Ok(PutResult {
e_tag: None,
version: None,
})
}
async fn abort(&mut self) -> Result<()> {
unimplemented!()
}
}
#[tokio::test]
async fn test_write_multipart() {
let mut rng = StdRng::seed_from_u64(42);
for method in [0.0, 0.5, 1.0] {
for _ in 0..10 {
for chunk_size in [1, 17, 23] {
let upload = Box::<InstrumentedUpload>::default();
let chunks = Arc::clone(&upload.chunks);
let mut write = WriteMultipart::new_with_chunk_size(upload, chunk_size);
let mut expected = Vec::with_capacity(1024);
for _ in 0..50 {
let chunk_size = rng.random_range(0..30);
let data: Vec<_> = (0..chunk_size).map(|_| rng.random()).collect();
expected.extend_from_slice(&data);
match rng.random_bool(method) {
true => write.put(data.into()),
false => write.write(&data),
}
}
write.finish().await.unwrap();
let chunks = chunks.lock();
let actual: Vec<_> = chunks.iter().flatten().flatten().copied().collect();
assert_eq!(expected, actual);
for chunk in chunks.iter().take(chunks.len() - 1) {
assert_eq!(chunk.content_length(), chunk_size)
}
let last_chunk = chunks.last().unwrap().content_length();
assert!(last_chunk <= chunk_size, "{chunk_size}");
}
}
}
}
}
backend/vendor/object_store/src/util.rs
Vendored
+491
View File
@@ -0,0 +1,491 @@
// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License. You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied. See the License for the
// specific language governing permissions and limitations
// under the License.
//! Common logic for interacting with remote object stores
use std::{
fmt::Display,
ops::{Range, RangeBounds},
};
use super::Result;
use bytes::Bytes;
use futures::{stream::StreamExt, Stream, TryStreamExt};
#[cfg(any(feature = "azure", feature = "http"))]
pub(crate) static RFC1123_FMT: &str = "%a, %d %h %Y %T GMT";
// deserialize dates according to rfc1123
#[cfg(any(feature = "azure", feature = "http"))]
pub(crate) fn deserialize_rfc1123<'de, D>(
deserializer: D,
) -> Result<chrono::DateTime<chrono::Utc>, D::Error>
where
D: serde::Deserializer<'de>,
{
let s: String = serde::Deserialize::deserialize(deserializer)?;
let naive =
chrono::NaiveDateTime::parse_from_str(&s, RFC1123_FMT).map_err(serde::de::Error::custom)?;
Ok(chrono::TimeZone::from_utc_datetime(&chrono::Utc, &naive))
}
#[cfg(any(feature = "aws", feature = "azure"))]
pub(crate) fn hmac_sha256(secret: impl AsRef<[u8]>, bytes: impl AsRef<[u8]>) -> ring::hmac::Tag {
let key = ring::hmac::Key::new(ring::hmac::HMAC_SHA256, secret.as_ref());
ring::hmac::sign(&key, bytes.as_ref())
}
/// Collect a stream into [`Bytes`] avoiding copying in the event of a single chunk
pub async fn collect_bytes<S, E>(mut stream: S, size_hint: Option<u64>) -> Result<Bytes, E>
where
E: Send,
S: Stream<Item = Result<Bytes, E>> + Send + Unpin,
{
let first = stream.next().await.transpose()?.unwrap_or_default();
// Avoid copying if single response
match stream.next().await.transpose()? {
None => Ok(first),
Some(second) => {
let size_hint = size_hint.unwrap_or_else(|| first.len() as u64 + second.len() as u64);
let mut buf = Vec::with_capacity(size_hint as usize);
buf.extend_from_slice(&first);
buf.extend_from_slice(&second);
while let Some(maybe_bytes) = stream.next().await {
buf.extend_from_slice(&maybe_bytes?);
}
Ok(buf.into())
}
}
}
#[cfg(all(feature = "fs", not(target_arch = "wasm32")))]
/// Takes a function and spawns it to a tokio blocking pool if available
pub(crate) async fn maybe_spawn_blocking<F, T>(f: F) -> Result<T>
where
F: FnOnce() -> Result<T> + Send + 'static,
T: Send + 'static,
{
match tokio::runtime::Handle::try_current() {
Ok(runtime) => runtime.spawn_blocking(f).await?,
Err(_) => f(),
}
}
/// Range requests with a gap less than or equal to this,
/// will be coalesced into a single request by [`coalesce_ranges`]
pub const OBJECT_STORE_COALESCE_DEFAULT: u64 = 1024 * 1024;
/// Up to this number of range requests will be performed in parallel by [`coalesce_ranges`]
pub(crate) const OBJECT_STORE_COALESCE_PARALLEL: usize = 10;
/// Takes a function `fetch` that can fetch a range of bytes and uses this to
/// fetch the provided byte `ranges`
///
/// To improve performance it will:
///
/// * Combine ranges less than `coalesce` bytes apart into a single call to `fetch`
/// * Make multiple `fetch` requests in parallel (up to maximum of 10)
///
pub async fn coalesce_ranges<F, E, Fut>(
ranges: &[Range<u64>],
fetch: F,
coalesce: u64,
) -> Result<Vec<Bytes>, E>
where
F: Send + FnMut(Range<u64>) -> Fut,
E: Send,
Fut: std::future::Future<Output = Result<Bytes, E>> + Send,
{
let fetch_ranges = merge_ranges(ranges, coalesce);
let fetched: Vec<_> = futures::stream::iter(fetch_ranges.iter().cloned())
.map(fetch)
.buffered(OBJECT_STORE_COALESCE_PARALLEL)
.try_collect()
.await?;
Ok(ranges
.iter()
.map(|range| {
let idx = fetch_ranges.partition_point(|v| v.start <= range.start) - 1;
let fetch_range = &fetch_ranges[idx];
let fetch_bytes = &fetched[idx];
let start = range.start - fetch_range.start;
let end = range.end - fetch_range.start;
let range = (start as usize)..(end as usize).min(fetch_bytes.len());
fetch_bytes.slice(range)
})
.collect())
}
/// Returns a sorted list of ranges that cover `ranges`
fn merge_ranges(ranges: &[Range<u64>], coalesce: u64) -> Vec<Range<u64>> {
if ranges.is_empty() {
return vec![];
}
let mut ranges = ranges.to_vec();
ranges.sort_unstable_by_key(|range| range.start);
let mut ret = Vec::with_capacity(ranges.len());
let mut start_idx = 0;
let mut end_idx = 1;
while start_idx != ranges.len() {
let mut range_end = ranges[start_idx].end;
while end_idx != ranges.len()
&& ranges[end_idx]
.start
.checked_sub(range_end)
.map(|delta| delta <= coalesce)
.unwrap_or(true)
{
range_end = range_end.max(ranges[end_idx].end);
end_idx += 1;
}
let start = ranges[start_idx].start;
let end = range_end;
ret.push(start..end);
start_idx = end_idx;
end_idx += 1;
}
ret
}
/// Request only a portion of an object's bytes
///
/// These can be created from [usize] ranges, like
///
/// ```rust
/// # use object_store::GetRange;
/// let range1: GetRange = (50..150).into();
/// let range2: GetRange = (50..=150).into();
/// let range3: GetRange = (50..).into();
/// let range4: GetRange = (..150).into();
/// ```
///
/// Implementations may wish to inspect [`GetResult`] for the exact byte
/// range returned.
///
/// [`GetResult`]: crate::GetResult
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum GetRange {
/// Request a specific range of bytes
///
/// If the given range is zero-length or starts after the end of the object,
/// an error will be returned. Additionally, if the range ends after the end
/// of the object, the entire remainder of the object will be returned.
/// Otherwise, the exact requested range will be returned.
///
/// Note that range is u64 (i.e., not usize),
/// as `object_store` supports 32-bit architectures such as WASM
Bounded(Range<u64>),
/// Request all bytes starting from a given byte offset
Offset(u64),
/// Request up to the last n bytes
Suffix(u64),
}
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum InvalidGetRange {
#[error("Wanted range starting at {requested}, but object was only {length} bytes long")]
StartTooLarge { requested: u64, length: u64 },
#[error("Range started at {start} and ended at {end}")]
Inconsistent { start: u64, end: u64 },
#[error("Range {requested} is larger than system memory limit {max}")]
TooLarge { requested: u64, max: u64 },
}
impl GetRange {
/// Check if the range is valid.
pub fn is_valid(&self) -> Result<(), InvalidGetRange> {
if let Self::Bounded(r) = self {
if r.end <= r.start {
return Err(InvalidGetRange::Inconsistent {
start: r.start,
end: r.end,
});
}
if (r.end - r.start) > usize::MAX as u64 {
return Err(InvalidGetRange::TooLarge {
requested: r.start,
max: usize::MAX as u64,
});
}
}
Ok(())
}
/// Convert to a [`Range`] if [valid](Self::is_valid).
pub fn as_range(&self, len: u64) -> Result<Range<u64>, InvalidGetRange> {
self.is_valid()?;
match self {
Self::Bounded(r) => {
if r.start >= len {
Err(InvalidGetRange::StartTooLarge {
requested: r.start,
length: len,
})
} else if r.end > len {
Ok(r.start..len)
} else {
Ok(r.clone())
}
}
Self::Offset(o) => {
if *o >= len {
Err(InvalidGetRange::StartTooLarge {
requested: *o,
length: len,
})
} else {
Ok(*o..len)
}
}
Self::Suffix(n) => Ok(len.saturating_sub(*n)..len),
}
}
}
impl Display for GetRange {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Bounded(r) => write!(f, "bytes={}-{}", r.start, r.end - 1),
Self::Offset(o) => write!(f, "bytes={o}-"),
Self::Suffix(n) => write!(f, "bytes=-{n}"),
}
}
}
impl<T: RangeBounds<u64>> From<T> for GetRange {
fn from(value: T) -> Self {
use std::ops::Bound::*;
let first = match value.start_bound() {
Included(i) => *i,
Excluded(i) => i + 1,
Unbounded => 0,
};
match value.end_bound() {
Included(i) => Self::Bounded(first..(i + 1)),
Excluded(i) => Self::Bounded(first..*i),
Unbounded => Self::Offset(first),
}
}
}
// http://docs.aws.amazon.com/general/latest/gr/sigv4-create-canonical-request.html
//
// Do not URI-encode any of the unreserved characters that RFC 3986 defines:
// A-Z, a-z, 0-9, hyphen ( - ), underscore ( _ ), period ( . ), and tilde ( ~ ).
#[cfg(any(feature = "aws", feature = "gcp"))]
pub(crate) const STRICT_ENCODE_SET: percent_encoding::AsciiSet = percent_encoding::NON_ALPHANUMERIC
.remove(b'-')
.remove(b'.')
.remove(b'_')
.remove(b'~');
/// Computes the SHA256 digest of `body` returned as a hex encoded string
#[cfg(any(feature = "aws", feature = "gcp"))]
pub(crate) fn hex_digest(bytes: &[u8]) -> String {
let digest = ring::digest::digest(&ring::digest::SHA256, bytes);
hex_encode(digest.as_ref())
}
/// Returns `bytes` as a lower-case hex encoded string
#[cfg(any(feature = "aws", feature = "gcp"))]
pub(crate) fn hex_encode(bytes: &[u8]) -> String {
use std::fmt::Write;
let mut out = String::with_capacity(bytes.len() * 2);
for byte in bytes {
// String writing is infallible
let _ = write!(out, "{byte:02x}");
}
out
}
#[cfg(test)]
mod tests {
use crate::Error;
use super::*;
use rand::{rng, Rng};
use std::ops::Range;
/// Calls coalesce_ranges and validates the returned data is correct
///
/// Returns the fetched ranges
async fn do_fetch(ranges: Vec<Range<u64>>, coalesce: u64) -> Vec<Range<u64>> {
let max = ranges.iter().map(|x| x.end).max().unwrap_or(0);
let src: Vec<_> = (0..max).map(|x| x as u8).collect();
let mut fetches = vec![];
let coalesced = coalesce_ranges::<_, Error, _>(
&ranges,
|range| {
fetches.push(range.clone());
let start = usize::try_from(range.start).unwrap();
let end = usize::try_from(range.end).unwrap();
futures::future::ready(Ok(Bytes::from(src[start..end].to_vec())))
},
coalesce,
)
.await
.unwrap();
assert_eq!(ranges.len(), coalesced.len());
for (range, bytes) in ranges.iter().zip(coalesced) {
assert_eq!(
bytes.as_ref(),
&src[usize::try_from(range.start).unwrap()..usize::try_from(range.end).unwrap()]
);
}
fetches
}
#[tokio::test]
async fn test_coalesce_ranges() {
let fetches = do_fetch(vec![], 0).await;
assert!(fetches.is_empty());
let fetches = do_fetch(vec![0..3; 1], 0).await;
assert_eq!(fetches, vec![0..3]);
let fetches = do_fetch(vec![0..2, 3..5], 0).await;
assert_eq!(fetches, vec![0..2, 3..5]);
let fetches = do_fetch(vec![0..1, 1..2], 0).await;
assert_eq!(fetches, vec![0..2]);
let fetches = do_fetch(vec![0..1, 2..72], 1).await;
assert_eq!(fetches, vec![0..72]);
let fetches = do_fetch(vec![0..1, 56..72, 73..75], 1).await;
assert_eq!(fetches, vec![0..1, 56..75]);
let fetches = do_fetch(vec![0..1, 5..6, 7..9, 2..3, 4..6], 1).await;
assert_eq!(fetches, vec![0..9]);
let fetches = do_fetch(vec![0..1, 5..6, 7..9, 2..3, 4..6], 1).await;
assert_eq!(fetches, vec![0..9]);
let fetches = do_fetch(vec![0..1, 6..7, 8..9, 10..14, 9..10], 4).await;
assert_eq!(fetches, vec![0..1, 6..14]);
}
#[tokio::test]
async fn test_coalesce_fuzz() {
let mut rand = rng();
for _ in 0..100 {
let object_len = rand.random_range(10..250);
let range_count = rand.random_range(0..10);
let ranges: Vec<_> = (0..range_count)
.map(|_| {
let start = rand.random_range(0..object_len);
let max_len = 20.min(object_len - start);
let len = rand.random_range(0..max_len);
start..start + len
})
.collect();
let coalesce = rand.random_range(1..5);
let fetches = do_fetch(ranges.clone(), coalesce).await;
for fetch in fetches.windows(2) {
assert!(
fetch[0].start <= fetch[1].start,
"fetches should be sorted, {:?} vs {:?}",
fetch[0],
fetch[1]
);
let delta = fetch[1].end - fetch[0].end;
assert!(
delta > coalesce,
"fetches should not overlap by {}, {:?} vs {:?} for {:?}",
coalesce,
fetch[0],
fetch[1],
ranges
);
}
}
}
#[test]
fn getrange_str() {
assert_eq!(GetRange::Offset(0).to_string(), "bytes=0-");
assert_eq!(GetRange::Bounded(10..19).to_string(), "bytes=10-18");
assert_eq!(GetRange::Suffix(10).to_string(), "bytes=-10");
}
#[test]
fn getrange_from() {
assert_eq!(Into::<GetRange>::into(10..15), GetRange::Bounded(10..15),);
assert_eq!(Into::<GetRange>::into(10..=15), GetRange::Bounded(10..16),);
assert_eq!(Into::<GetRange>::into(10..), GetRange::Offset(10),);
assert_eq!(Into::<GetRange>::into(..=15), GetRange::Bounded(0..16));
}
#[test]
fn test_as_range() {
let range = GetRange::Bounded(2..5);
assert_eq!(range.as_range(5).unwrap(), 2..5);
let range = range.as_range(4).unwrap();
assert_eq!(range, 2..4);
let range = GetRange::Bounded(3..3);
let err = range.as_range(2).unwrap_err().to_string();
assert_eq!(err, "Range started at 3 and ended at 3");
let range = GetRange::Bounded(2..2);
let err = range.as_range(3).unwrap_err().to_string();
assert_eq!(err, "Range started at 2 and ended at 2");
let range = GetRange::Suffix(3);
assert_eq!(range.as_range(3).unwrap(), 0..3);
assert_eq!(range.as_range(2).unwrap(), 0..2);
let range = GetRange::Suffix(0);
assert_eq!(range.as_range(0).unwrap(), 0..0);
let range = GetRange::Offset(2);
let err = range.as_range(2).unwrap_err().to_string();
assert_eq!(
err,
"Wanted range starting at 2, but object was only 2 bytes long"
);
let err = range.as_range(1).unwrap_err().to_string();
assert_eq!(
err,
"Wanted range starting at 2, but object was only 1 bytes long"
);
let range = GetRange::Offset(1);
assert_eq!(range.as_range(2).unwrap(), 1..2);
}
}