diff --git a/object_store/src/aws/client.rs b/object_store/src/aws/client.rs index 2cf808a9d49..b054baf36f4 100644 --- a/object_store/src/aws/client.rs +++ b/object_store/src/aws/client.rs @@ -389,6 +389,11 @@ impl Request<'_> { 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() diff --git a/object_store/src/aws/mod.rs b/object_store/src/aws/mod.rs index 0625ae140b7..7c46b54c3b7 100644 --- a/object_store/src/aws/mod.rs +++ b/object_store/src/aws/mod.rs @@ -159,15 +159,23 @@ impl ObjectStore for AmazonS3 { payload: PutPayload, opts: PutOptions, ) -> Result { + let PutOptions { + mode, + tags, + attributes, + extensions, + } = opts; + let request = self .client .request(Method::PUT, location) .with_payload(payload) - .with_attributes(opts.attributes) - .with_tags(opts.tags) + .with_attributes(attributes) + .with_tags(tags) + .with_extensions(extensions) .with_encryption_headers(); - match (opts.mode, &self.client.config.conditional_put) { + match (mode, &self.client.config.conditional_put) { (PutMode::Overwrite, _) => request.idempotent(true).do_put().await, (PutMode::Create | PutMode::Update(_), None) => Err(Error::NotImplemented), (PutMode::Create, Some(S3ConditionalPut::ETagMatch)) => { diff --git a/object_store/src/azure/client.rs b/object_store/src/azure/client.rs index 13e40bbf968..c4d026bcb8a 100644 --- a/object_store/src/azure/client.rs +++ b/object_store/src/azure/client.rs @@ -257,6 +257,11 @@ impl PutRequest<'_> { Self { builder, ..self } } + fn with_extensions(self, extensions: ::http::Extensions) -> Self { + let builder = self.builder.extensions(extensions); + Self { builder, ..self } + } + async fn send(self) -> Result { let credential = self.config.get_credential().await?; let sensitive = credential @@ -540,12 +545,20 @@ impl AzureClient { payload: PutPayload, opts: PutOptions, ) -> Result { + let PutOptions { + mode, + tags, + attributes, + extensions, + } = opts; + let builder = self .put_request(path, payload) - .with_attributes(opts.attributes) - .with_tags(opts.tags); + .with_attributes(attributes) + .with_extensions(extensions) + .with_tags(tags); - let builder = match &opts.mode { + let builder = match &mode { PutMode::Overwrite => builder.idempotent(true), PutMode::Create => builder.header(&IF_NONE_MATCH, "*"), PutMode::Update(v) => { diff --git a/object_store/src/gcp/client.rs b/object_store/src/gcp/client.rs index a52ad3663fd..e514624f8f7 100644 --- a/object_store/src/gcp/client.rs +++ b/object_store/src/gcp/client.rs @@ -231,6 +231,11 @@ impl Request<'_> { } } + fn with_extensions(self, extensions: ::http::Extensions) -> Self { + let builder = self.builder.extensions(extensions); + Self { builder, ..self } + } + async fn send(self) -> Result { let credential = self.config.credentials.get_credential().await?; let resp = self @@ -384,12 +389,21 @@ impl GoogleCloudStorageClient { payload: PutPayload, opts: PutOptions, ) -> Result { + let PutOptions { + mode, + // not supported by GCP + tags: _, + attributes, + extensions, + } = opts; + let builder = self .request(Method::PUT, path) .with_payload(payload) - .with_attributes(opts.attributes); + .with_attributes(attributes) + .with_extensions(extensions); - let builder = match &opts.mode { + let builder = match &mode { PutMode::Overwrite => builder.idempotent(true), PutMode::Create => builder.header(&VERSION_MATCH, "0"), PutMode::Update(v) => { @@ -398,7 +412,7 @@ impl GoogleCloudStorageClient { } }; - match (opts.mode, builder.do_put().await) { + match (mode, builder.do_put().await) { (PutMode::Create, Err(crate::Error::Precondition { path, source })) => { Err(crate::Error::AlreadyExists { path, source }) } diff --git a/object_store/src/lib.rs b/object_store/src/lib.rs index 21352f5761e..8f05fb39119 100644 --- a/object_store/src/lib.rs +++ b/object_store/src/lib.rs @@ -1154,7 +1154,7 @@ impl From for UpdateVersion { } /// Options for a put request -#[derive(Debug, Clone, PartialEq, Eq, Default)] +#[derive(Debug, Clone, Default)] pub struct PutOptions { /// Configure the [`PutMode`] for this operation pub mode: PutMode, @@ -1166,8 +1166,35 @@ pub struct PutOptions { /// /// Implementations that don't support an attribute should return an error pub attributes: Attributes, + /// 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. + /// + /// They are also eclused from [`PartialEq`] and [`Eq`]. + pub extensions: ::http::Extensions, } +impl PartialEq for PutOptions { + fn eq(&self, other: &Self) -> bool { + let Self { + mode, + tags, + attributes, + extensions: _, + } = self; + let Self { + mode: other_mode, + tags: other_tags, + attributes: other_attributes, + extensions: _, + } = other; + (mode == other_mode) && (tags == other_tags) && (attributes == other_attributes) + } +} + +impl Eq for PutOptions {} + impl From for PutOptions { fn from(mode: PutMode) -> Self { Self {