google_cloud_spanner/
batch_dml.rs

1// Copyright 2026 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
15use crate::error::{BatchUpdateError, internal_error};
16use crate::model::result_set_stats::RowCount;
17use crate::model::{ExecuteBatchDmlResponse, RequestOptions};
18use crate::statement::Statement;
19use google_cloud_gax::backoff_policy::BackoffPolicyArg;
20use google_cloud_gax::error::rpc::Code;
21use google_cloud_gax::error::rpc::Status as RpcStatus;
22use google_cloud_gax::options::RequestOptions as GaxRequestOptions;
23use google_cloud_gax::retry_policy::RetryPolicyArg;
24use std::time::Duration;
25
26/// A builder for [BatchDml].
27#[derive(Clone, Default, Debug)]
28pub struct BatchDmlBuilder {
29    statements: Vec<Statement>,
30    request_options: Option<RequestOptions>,
31    last_statements: bool,
32    gax_options: GaxRequestOptions,
33}
34
35impl BatchDmlBuilder {
36    /// Creates a new empty BatchDmlBuilder.
37    pub fn new() -> Self {
38        BatchDmlBuilder::default()
39    }
40
41    /// Adds a statement to the batch.
42    pub fn add_statement(mut self, statement: impl Into<Statement>) -> Self {
43        self.statements.push(statement.into());
44        self
45    }
46
47    /// Sets the request tag for this batch.
48    ///
49    /// # Example
50    /// ```
51    /// # use google_cloud_spanner::statement::Statement;
52    /// # use google_cloud_spanner::batch::BatchDml;
53    /// let statement1 = Statement::builder("UPDATE users SET active = true WHERE id = 1").build();
54    /// let batch = BatchDml::builder()
55    ///     .add_statement(statement1)
56    ///     .set_request_tag("my-tag")
57    ///     .build();
58    /// ```
59    ///
60    /// See also: [Troubleshooting with tags](https://docs.cloud.google.com/spanner/docs/introspection/troubleshooting-with-tags)
61    pub fn set_request_tag(mut self, tag: impl Into<String>) -> Self {
62        self.request_options
63            .get_or_insert_with(RequestOptions::default)
64            .request_tag = tag.into();
65        self
66    }
67
68    /// Sets the per-attempt timeout for this batch DML request.
69    pub fn with_attempt_timeout(mut self, timeout: Duration) -> Self {
70        self.gax_options.set_attempt_timeout(timeout);
71        self
72    }
73
74    /// Sets the retry policy for this batch DML request.
75    pub fn with_retry_policy(mut self, policy: impl Into<RetryPolicyArg>) -> Self {
76        self.gax_options.set_retry_policy(policy);
77        self
78    }
79
80    /// Sets the backoff policy for this batch DML request.
81    pub fn with_backoff_policy(mut self, policy: impl Into<BackoffPolicyArg>) -> Self {
82        self.gax_options.set_backoff_policy(policy);
83        self
84    }
85
86    /// Sets whether this batch is the last statement in a read/write transaction.
87    ///
88    /// # Example
89    /// ```
90    /// # use google_cloud_spanner::statement::Statement;
91    /// # use google_cloud_spanner::batch::BatchDml;
92    /// let statement1 = Statement::builder("UPDATE users SET active = true WHERE id = 1").build();
93    /// let batch = BatchDml::builder()
94    ///     .add_statement(statement1)
95    ///     .set_last_statements(true)
96    ///     .build();
97    /// ```
98    ///
99    /// If true, indicates that this request marks the end of the transaction, which
100    /// allows Spanner to optimize execution.
101    pub fn set_last_statements(mut self, last_statements: bool) -> Self {
102        self.last_statements = last_statements;
103        self
104    }
105
106    /// Builds and returns the finalized BatchDml object.
107    pub fn build(self) -> BatchDml {
108        BatchDml {
109            statements: self.statements,
110            request_options: self.request_options,
111            last_statements: self.last_statements,
112            gax_options: self.gax_options,
113        }
114    }
115}
116
117/// A batch of DML statements to be executed in a single round-trip to Spanner.
118#[derive(Clone, Debug)]
119pub struct BatchDml {
120    pub(crate) statements: Vec<Statement>,
121    pub(crate) request_options: Option<RequestOptions>,
122    pub(crate) last_statements: bool,
123    pub(crate) gax_options: GaxRequestOptions,
124}
125
126impl BatchDml {
127    /// Creates a new builder for constructing a [`BatchDml`] request.
128    pub fn builder() -> BatchDmlBuilder {
129        BatchDmlBuilder::new()
130    }
131}
132
133impl From<BatchDmlBuilder> for BatchDml {
134    fn from(builder: BatchDmlBuilder) -> Self {
135        builder.build()
136    }
137}
138
139impl<T: Into<Statement>> From<Vec<T>> for BatchDml {
140    fn from(statements: Vec<T>) -> Self {
141        BatchDml {
142            statements: statements.into_iter().map(Into::into).collect(),
143            request_options: None,
144            last_statements: false,
145            gax_options: GaxRequestOptions::default(),
146        }
147    }
148}
149
150/// Processes an ExecuteBatchDmlResponse and returns the success counts, or an error.
151pub(crate) fn process_response(response: ExecuteBatchDmlResponse) -> crate::Result<Vec<i64>> {
152    let mut update_counts = Vec::with_capacity(response.result_sets.len());
153    for result_set in response.result_sets {
154        if let Some(stats) = result_set.stats {
155            let exact_count = match stats.row_count {
156                Some(RowCount::RowCountExact(c)) => c,
157                _ => {
158                    return Err(internal_error(
159                        "ExecuteBatchDml returned an invalid or missing row count type",
160                    ));
161                }
162            };
163            update_counts.push(exact_count);
164        }
165    }
166
167    // If a non-zero status is present, it halted the batch somewhere in the middle of the batch.
168    if let Some(status) = response.status.filter(|s| s.code != Code::Ok as i32) {
169        let grpc_status = RpcStatus::default()
170            .set_code(status.code)
171            .set_message(status.message);
172
173        // If the error code is Aborted, then we propagate a 'normal' service error.
174        // The TransactionRunner will then retry the transaction.
175        if status.code == Code::Aborted as i32 {
176            return Err(crate::Error::service(grpc_status));
177        }
178        return Err(BatchUpdateError::build_error(update_counts, grpc_status));
179    }
180
181    Ok(update_counts)
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use crate::model::{ResultSet, ResultSetStats};
188    use google_cloud_rpc::model::Status;
189    use static_assertions::assert_impl_all;
190
191    #[test]
192    fn auto_traits() {
193        assert_impl_all!(BatchDml: Send, Sync, Clone, std::fmt::Debug);
194        assert_impl_all!(BatchDmlBuilder: Send, Sync, Clone, std::fmt::Debug);
195    }
196
197    #[test]
198    fn builder() {
199        let stmt1 = Statement::builder("UPDATE t SET c = 1 WHERE id = 1").build();
200        let stmt2 = Statement::builder("UPDATE t SET c = 2 WHERE id = 2").build();
201
202        let batch = BatchDml::builder()
203            .add_statement(stmt1)
204            .add_statement(stmt2)
205            .set_last_statements(true)
206            .build();
207
208        assert_eq!(batch.statements.len(), 2);
209        assert_eq!(batch.statements[0].sql, "UPDATE t SET c = 1 WHERE id = 1");
210        assert_eq!(batch.statements[1].sql, "UPDATE t SET c = 2 WHERE id = 2");
211        assert!(batch.last_statements);
212        assert!(
213            batch.request_options.is_none(),
214            "Unexpected request_options set: {:#?}",
215            batch.request_options
216        );
217    }
218
219    #[test]
220    fn builder_with_gax_options() {
221        use google_cloud_gax::backoff_policy::BackoffPolicy;
222        use google_cloud_gax::retry_policy::Aip194Strict;
223        use google_cloud_gax::retry_state::RetryState;
224        use std::time::Duration;
225
226        #[derive(Debug)]
227        struct DummyBackoff;
228        impl BackoffPolicy for DummyBackoff {
229            fn on_failure(&self, _state: &RetryState) -> Duration {
230                Duration::ZERO
231            }
232        }
233
234        let stmt = Statement::builder("UPDATE t SET c = 1 WHERE id = 1").build();
235
236        let batch = BatchDml::builder()
237            .add_statement(stmt)
238            .with_attempt_timeout(Duration::from_secs(5))
239            .with_retry_policy(Aip194Strict)
240            .with_backoff_policy(DummyBackoff)
241            .build();
242
243        assert_eq!(
244            *batch.gax_options.attempt_timeout(),
245            Some(Duration::from_secs(5))
246        );
247        assert!(batch.gax_options.retry_policy().is_some());
248        assert!(batch.gax_options.backoff_policy().is_some());
249    }
250
251    #[test]
252    fn builder_with_request_tag() {
253        let stmt = Statement::builder("UPDATE t SET c = 1 WHERE id = 1").build();
254
255        let batch = BatchDml::builder()
256            .add_statement(stmt)
257            .set_request_tag("tag1")
258            .build();
259
260        assert_eq!(batch.statements.len(), 1);
261        assert_eq!(
262            batch
263                .request_options
264                .expect("request options missing")
265                .request_tag,
266            "tag1"
267        );
268    }
269
270    #[test]
271    fn process_response_success() -> anyhow::Result<()> {
272        let stats1 = ResultSetStats {
273            row_count: Some(RowCount::RowCountExact(5)),
274            ..Default::default()
275        };
276        let stats2 = ResultSetStats {
277            row_count: Some(RowCount::RowCountExact(10)),
278            ..Default::default()
279        };
280
281        let rs1 = ResultSet {
282            stats: Some(stats1),
283            ..Default::default()
284        };
285        let rs2 = ResultSet {
286            stats: Some(stats2),
287            ..Default::default()
288        };
289
290        let response = ExecuteBatchDmlResponse {
291            result_sets: vec![rs1, rs2],
292            status: None,
293            ..Default::default()
294        };
295
296        let counts = process_response(response)?;
297        assert_eq!(counts, vec![5, 10]);
298        Ok(())
299    }
300
301    #[test]
302    fn process_response_grpc_error() {
303        let stats = ResultSetStats {
304            row_count: Some(RowCount::RowCountExact(3)),
305            ..Default::default()
306        };
307        let rs = ResultSet {
308            stats: Some(stats),
309            ..Default::default()
310        };
311
312        // Note: crate::model::Status is the common type for status embedded in generated grpc responses.
313        let err_status = Status::default()
314            .set_code(Code::InvalidArgument as i32)
315            .set_message("Bad query");
316
317        let response = ExecuteBatchDmlResponse {
318            result_sets: vec![rs],
319            status: Some(err_status),
320            ..Default::default()
321        };
322
323        let result = process_response(response);
324        let err = result.expect_err("should return error");
325        let batch_err = BatchUpdateError::extract(&err).expect("should extract BatchUpdateError");
326
327        assert_eq!(batch_err.update_counts, vec![3]);
328        assert_eq!(
329            batch_err.status.status().expect("status").code,
330            Code::InvalidArgument
331        );
332        assert_eq!(
333            batch_err.status.status().expect("status").message,
334            "Bad query"
335        );
336    }
337
338    #[test]
339    fn process_response_aborted() {
340        let stats = ResultSetStats {
341            row_count: Some(RowCount::RowCountExact(3)),
342            ..Default::default()
343        };
344        let rs = ResultSet {
345            stats: Some(stats),
346            ..Default::default()
347        };
348
349        let err_status = Status::default()
350            .set_code(Code::Aborted as i32)
351            .set_message("transaction aborted");
352
353        let response = ExecuteBatchDmlResponse {
354            result_sets: vec![rs],
355            status: Some(err_status),
356            ..Default::default()
357        };
358
359        let result = process_response(response);
360        let err = result.expect_err("should return error");
361        let batch_err = BatchUpdateError::extract(&err);
362        assert!(
363            batch_err.is_none(),
364            "Unexpected BatchUpdateError: {batch_err:?}"
365        );
366        assert_eq!(err.status().expect("status").code, Code::Aborted);
367        assert_eq!(err.status().expect("status").message, "transaction aborted");
368    }
369
370    #[test]
371    fn process_response_missing_stats() {
372        let rs = ResultSet {
373            stats: None,
374            ..Default::default()
375        };
376
377        let response = ExecuteBatchDmlResponse {
378            result_sets: vec![rs],
379            ..Default::default()
380        };
381
382        let result = process_response(response).expect("should return empty update counts");
383        assert!(result.is_empty());
384    }
385
386    #[test]
387    fn process_response_missing_row_count_type() {
388        let stats = ResultSetStats {
389            row_count: None,
390            ..Default::default()
391        };
392
393        let rs = ResultSet {
394            stats: Some(stats),
395            ..Default::default()
396        };
397
398        let response = ExecuteBatchDmlResponse {
399            result_sets: vec![rs],
400            ..Default::default()
401        };
402
403        let result = process_response(response);
404        let err = result.expect_err("should fail");
405        assert!(
406            err.to_string()
407                .contains("invalid or missing row count type")
408        );
409    }
410
411    #[test]
412    fn from_vector_of_strings() {
413        let statements = vec!["UPDATE table SET col = 1", "UPDATE table SET col = 2"];
414        let batch: BatchDml = statements.into();
415        assert_eq!(batch.statements.len(), 2);
416        assert_eq!(batch.statements[0].sql, "UPDATE table SET col = 1");
417        assert_eq!(batch.statements[1].sql, "UPDATE table SET col = 2");
418    }
419
420    #[test]
421    fn from_vector_of_statements() {
422        let statement1 = Statement::builder("UPDATE table SET col = 1").build();
423        let statement2 = Statement::builder("UPDATE table SET col = 2").build();
424        let statements = vec![statement1, statement2];
425        let batch: BatchDml = statements.into();
426        assert_eq!(batch.statements.len(), 2);
427        assert_eq!(batch.statements[0].sql, "UPDATE table SET col = 1");
428        assert_eq!(batch.statements[1].sql, "UPDATE table SET col = 2");
429    }
430
431    #[test]
432    fn from_builder() {
433        let builder = BatchDml::builder().add_statement("UPDATE table SET col = 1");
434        let batch: BatchDml = builder.into();
435        assert_eq!(batch.statements.len(), 1);
436        assert_eq!(batch.statements[0].sql, "UPDATE table SET col = 1");
437    }
438
439    #[test]
440    fn process_response_metadata_no_stats_grpc_error() {
441        let rs = ResultSet {
442            metadata: Some(crate::model::ResultSetMetadata {
443                transaction: Some(crate::model::Transaction {
444                    id: vec![7, 7, 7].into(),
445                    ..Default::default()
446                }),
447                ..Default::default()
448            }),
449            stats: None,
450            ..Default::default()
451        };
452
453        let err_status = Status::default()
454            .set_code(Code::InvalidArgument as i32)
455            .set_message("Table not found or syntax invalid");
456
457        let response = ExecuteBatchDmlResponse {
458            result_sets: vec![rs],
459            status: Some(err_status),
460            ..Default::default()
461        };
462
463        let result = process_response(response);
464        let err = result.expect_err("should return error");
465        let batch_err = BatchUpdateError::extract(&err)
466            .expect("should extract BatchUpdateError cleanly and not return internal error");
467
468        assert_eq!(
469            batch_err.update_counts,
470            Vec::<i64>::new(),
471            "Update counts should be completely empty"
472        );
473        assert_eq!(
474            batch_err.status.status().expect("status").code,
475            Code::InvalidArgument
476        );
477        assert_eq!(
478            batch_err.status.status().expect("status").message,
479            "Table not found or syntax invalid"
480        );
481    }
482}
google_cloud_spanner/batch_dml.rs

google_cloud_spanner/
batch_dml.rs
