google_cloud_storage/
builder_ext.rs

1// Copyright 2025 Google LLC
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     https://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15//! Extends [builder][crate::builder] with types that improve type safety and/or
16//! ergonomics.
17
18/// An extension trait for `RewriteObject` to provide a convenient way
19/// to poll a rewrite operation until it is complete.
20#[async_trait::async_trait]
21pub trait RewriteObjectExt {
22    /// Sends the request and polls the operation until it is complete.
23    ///
24    /// This helper function simplifies the process of handling a
25    /// [StorageControl::rewrite_object][crate::client::StorageControl::rewrite_object]
26    /// operation, which may require multiple requests to complete. It automatically
27    /// handles the logic of sending the
28    /// [rewrite_token][crate::generated::gapic::model::RewriteObjectRequest::rewrite_token]
29    /// from one response in the next request.
30    ///
31    /// For more details on this loop, see the "Rewriting objects" section of the
32    /// user guide:
33    /// <https://googleapis.github.io/google-cloud-rust/storage/rewrite_object.html>
34    ///
35    /// # Example
36    ///
37    /// ```
38    /// # use google_cloud_storage::client::StorageControl;
39    /// # use google_cloud_storage::builder_ext::RewriteObjectExt;
40    /// # async fn sample(client: &StorageControl) -> anyhow::Result<()> {
41    /// const SOURCE_NAME: &str = "object-to-copy";
42    /// const DEST_NAME: &str = "copied-object";
43    /// let source_bucket_id = "source-bucket";
44    /// let dest_bucket_id = "dest-bucket";
45    /// let copied = client
46    ///     .rewrite_object()
47    ///     .set_source_bucket(format!("projects/_/buckets/{source_bucket_id}"))
48    ///     .set_source_object(SOURCE_NAME)
49    ///     .set_destination_bucket(format!("projects/_/buckets/{dest_bucket_id}"))
50    ///     .set_destination_name(DEST_NAME)
51    ///     .rewrite_until_done()
52    ///     .await?;
53    /// # Ok(())
54    /// # }
55    /// ```
56    async fn rewrite_until_done(self) -> crate::Result<crate::model::Object>;
57}
58
59#[async_trait::async_trait]
60impl RewriteObjectExt for crate::builder::storage_control::RewriteObject {
61    async fn rewrite_until_done(mut self) -> crate::Result<crate::model::Object> {
62        loop {
63            let resp = self.clone().send().await?;
64            if resp.done {
65                return Ok(resp
66                    .resource
67                    .expect("an object is always returned when the rewrite operation is done"));
68            }
69            self = self.set_rewrite_token(resp.rewrite_token);
70        }
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use crate::client::StorageControl;
78    use crate::model::{Object, RewriteObjectRequest, RewriteResponse};
79    use gax::error::rpc::{Code, Status};
80    use gax::options::RequestOptions;
81    use gax::response::Response;
82
83    mockall::mock! {
84        #[derive(Debug)]
85        StorageControl {}
86        impl crate::stub::StorageControl for StorageControl {
87            async fn rewrite_object( &self, _req: RewriteObjectRequest, _options: RequestOptions) -> gax::Result<Response<RewriteResponse>>;
88        }
89    }
90
91    #[tokio::test]
92    async fn test_rewrite_until_done() -> anyhow::Result<()> {
93        let mut mock = MockStorageControl::new();
94        let final_object = Object::new().set_name("final-object");
95
96        let mut seq = mockall::Sequence::new();
97        mock.expect_rewrite_object()
98            .withf(|req: &RewriteObjectRequest, _| req.rewrite_token.is_empty())
99            .times(1)
100            .in_sequence(&mut seq)
101            .returning(|_, _| {
102                Ok(Response::from(
103                    RewriteResponse::new()
104                        .set_done(false)
105                        .set_rewrite_token("token1"),
106                ))
107            });
108
109        mock.expect_rewrite_object()
110            .withf(|req: &RewriteObjectRequest, _| req.rewrite_token == "token1")
111            .times(1)
112            .in_sequence(&mut seq)
113            .returning({
114                let obj = final_object.clone();
115                move |_, _| {
116                    Ok(Response::from(
117                        RewriteResponse::new()
118                            .set_done(true)
119                            .set_resource(obj.clone()),
120                    ))
121                }
122            });
123
124        let client = StorageControl::from_stub(mock);
125        let result = client.rewrite_object().rewrite_until_done().await?;
126
127        assert_eq!(result, final_object);
128        Ok(())
129    }
130
131    #[tokio::test]
132    async fn test_rewrite_until_done_error() -> anyhow::Result<()> {
133        let mut mock = MockStorageControl::new();
134        mock.expect_rewrite_object()
135            .withf(|req: &RewriteObjectRequest, _| req.rewrite_token.is_empty())
136            .times(1)
137            .returning(|_, _| {
138                Err(gax::error::Error::service(
139                    Status::default().set_code(Code::Unavailable),
140                ))
141            });
142
143        let client = StorageControl::from_stub(mock);
144        let _ = client
145            .rewrite_object()
146            .rewrite_until_done()
147            .await
148            .unwrap_err();
149
150        Ok(())
151    }
152}
google_cloud_storage/builder_ext.rs

google_cloud_storage/
builder_ext.rs
