google_cloud_spanner/
statement.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::model::DirectedReadOptions;
16use crate::model::execute_sql_request::QueryMode;
17use crate::model::execute_sql_request::QueryOptions;
18use crate::model::request_options::Priority;
19use crate::to_value::ToValue;
20use crate::types::Type;
21use crate::value::Value;
22use google_cloud_gax::backoff_policy::BackoffPolicyArg;
23use google_cloud_gax::options::RequestOptions as GaxRequestOptions;
24use google_cloud_gax::retry_policy::RetryPolicyArg;
25use std::collections::BTreeMap;
26use std::time::Duration;
27
28/// A builder for [Statement].
29///
30/// # Example
31/// ```
32/// # use google_cloud_spanner::statement::Statement;
33/// let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
34///     .add_param("id", &42)
35///     .build();
36/// ```
37#[derive(Clone, Debug)]
38pub struct StatementBuilder {
39    sql: String,
40    params: BTreeMap<String, Value>,
41    param_types: BTreeMap<String, Type>,
42    request_options: Option<crate::model::RequestOptions>,
43    directed_read_options: Option<DirectedReadOptions>,
44    query_options: Option<QueryOptions>,
45    query_mode: Option<QueryMode>,
46    last_statement: bool,
47    gax_options: GaxRequestOptions,
48}
49
50impl StatementBuilder {
51    pub(crate) fn new(sql: impl Into<String>) -> Self {
52        Self {
53            sql: sql.into(),
54            params: BTreeMap::new(),
55            param_types: BTreeMap::new(),
56            request_options: None,
57            directed_read_options: None,
58            query_options: None,
59            query_mode: None,
60            last_statement: false,
61            gax_options: GaxRequestOptions::default(),
62        }
63    }
64
65    /// Adds a parameter value to this Statement.
66    ///
67    /// The parameter value is sent without an explicit type code to Spanner. This allows Spanner
68    /// to automatically infer the correct data type from the SQL string of the statement.
69    /// It is recommended to use untyped parameter values, unless you explicitly want Spanner to
70    /// verify that the type of the parameter value is exactly the same as the type that would
71    /// otherwise be inferred from the SQL string.
72    pub fn add_param<T: ToValue + ?Sized>(mut self, name: impl Into<String>, value: &T) -> Self {
73        self.params.insert(name.into(), value.to_value());
74        self
75    }
76
77    /// Adds a typed parameter value to this Statement.
78    ///
79    /// The parameter value is sent with an explicit type code to Spanner. The type code must
80    /// correspond with the expression in the SQL string that the query parameter is bound to.
81    pub fn add_typed_param<T: ToValue + ?Sized>(
82        mut self,
83        name: impl Into<String>,
84        value: &T,
85        param_type: Type,
86    ) -> Self {
87        let name = name.into();
88        self.params.insert(name.clone(), value.to_value());
89        self.param_types.insert(name, param_type);
90        self
91    }
92
93    /// Sets the request tag to use for this statement.
94    ///
95    /// # Example
96    /// ```
97    /// # use google_cloud_spanner::statement::Statement;
98    /// let statement = Statement::builder("SELECT * FROM users")
99    ///     .set_request_tag("my-tag")
100    ///     .build();
101    /// ```
102    ///
103    /// See also: [Troubleshooting with tags](https://docs.cloud.google.com/spanner/docs/introspection/troubleshooting-with-tags)
104    pub fn set_request_tag(mut self, tag: impl Into<String>) -> Self {
105        self.request_options
106            .get_or_insert_with(crate::model::RequestOptions::default)
107            .request_tag = tag.into();
108        self
109    }
110
111    /// Sets the RPC priority to use for this statement.
112    ///
113    /// # Example
114    /// ```
115    /// # use google_cloud_spanner::statement::Statement;
116    /// # use google_cloud_spanner::model::request_options::Priority;
117    /// let statement = Statement::builder("SELECT * FROM users")
118    ///     .set_priority(Priority::Low)
119    ///     .build();
120    /// ```
121    pub fn set_priority(mut self, priority: Priority) -> Self {
122        self.request_options
123            .get_or_insert_with(crate::model::RequestOptions::default)
124            .priority = priority;
125        self
126    }
127
128    /// Sets the directed read options for this statement.
129    ///
130    /// ```
131    /// # use google_cloud_spanner::statement::Statement;
132    /// # use google_cloud_spanner::model::DirectedReadOptions;
133    /// let dro = DirectedReadOptions::default();
134    /// let stmt = Statement::builder("SELECT * FROM users")
135    ///     .set_directed_read_options(dro)
136    ///     .build();
137    /// ```
138    ///
139    /// DirectedReadOptions can only be specified for a read-only transaction,
140    /// otherwise Spanner returns an INVALID_ARGUMENT error.
141    pub fn set_directed_read_options(mut self, options: DirectedReadOptions) -> Self {
142        self.directed_read_options = Some(options);
143        self
144    }
145
146    /// Sets the query options to use for this statement.
147    ///
148    /// # Example
149    /// ```
150    /// # use google_cloud_spanner::statement::Statement;
151    /// # use google_cloud_spanner::model::execute_sql_request::QueryOptions;
152    /// let options = QueryOptions::default()
153    ///     .set_optimizer_version("latest");
154    /// let statement = Statement::builder("SELECT * FROM users")
155    ///     .set_query_options(options)
156    ///     .build();
157    /// ```
158    pub fn set_query_options(mut self, options: QueryOptions) -> Self {
159        self.query_options = Some(options);
160        self
161    }
162
163    /// Sets the query mode to use for this statement.
164    ///
165    /// # Example
166    /// ```
167    /// # use google_cloud_spanner::statement::Statement;
168    /// # use google_cloud_spanner::model::execute_sql_request::QueryMode;
169    /// let statement = Statement::builder("SELECT * FROM users")
170    ///     .set_query_mode(QueryMode::Plan)
171    ///     .build();
172    /// ```
173    pub fn set_query_mode(mut self, mode: QueryMode) -> Self {
174        self.query_mode = Some(mode);
175        self
176    }
177
178    /// Sets the per-attempt timeout for this statement.
179    pub fn with_attempt_timeout(mut self, timeout: Duration) -> Self {
180        self.gax_options.set_attempt_timeout(timeout);
181        self
182    }
183
184    /// Sets whether this statement is the last statement in a read/write transaction.
185    ///
186    /// # Example
187    /// ```
188    /// # use google_cloud_spanner::client::Spanner;
189    /// # use google_cloud_spanner::statement::Statement;
190    /// # async fn run_tx(client: Spanner) -> Result<(), google_cloud_spanner::Error> {
191    /// let db_client = client.database_client("projects/p/instances/i/databases/d").build().await?;
192    /// let runner = db_client.read_write_transaction().build().await?;
193    ///
194    /// let result = runner.run(async |transaction| {
195    ///     let statement = Statement::builder("UPDATE MyTable SET MyColumn = 'MyValue' WHERE Id = 1")
196    ///         .set_last_statement(true)
197    ///         .build();
198    ///     transaction.execute_update(statement).await?;
199    ///     Ok(42)
200    /// }).await?;
201    /// # Ok(())
202    /// # }
203    /// ```
204    ///
205    /// If true, indicates that this is the last statement in a transaction, which allows
206    /// Spanner to optimize execution.
207    pub fn set_last_statement(mut self, last_statement: bool) -> Self {
208        self.last_statement = last_statement;
209        self
210    }
211
212    /// Sets the retry policy for this statement.
213    pub fn with_retry_policy(mut self, policy: impl Into<RetryPolicyArg>) -> Self {
214        self.gax_options.set_retry_policy(policy);
215        self
216    }
217
218    /// Sets the backoff policy for this statement.
219    pub fn with_backoff_policy(mut self, policy: impl Into<BackoffPolicyArg>) -> Self {
220        self.gax_options.set_backoff_policy(policy);
221        self
222    }
223
224    /// Builds and returns the finalized Statement object.
225    pub fn build(self) -> Statement {
226        Statement {
227            sql: self.sql,
228            params: self.params,
229            param_types: self.param_types,
230            request_options: self.request_options,
231            directed_read_options: self.directed_read_options,
232            query_options: self.query_options,
233            query_mode: self.query_mode,
234            last_statement: self.last_statement,
235            gax_options: self.gax_options,
236        }
237    }
238}
239
240/// A SQL statement for execution on Spanner.
241///
242/// # Example
243/// ```
244/// # use google_cloud_spanner::client::Spanner;
245/// # use google_cloud_spanner::statement::Statement;
246/// # async fn test_doc() -> Result<(), google_cloud_spanner::Error> {
247/// let client = Spanner::builder().build().await.unwrap();
248/// let db = client.database_client("projects/p/instances/i/databases/d").build().await.unwrap();
249///
250/// let tx = db.single_use().build();
251/// let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
252///     .add_param("id", &42)
253///     .build();
254/// let mut rs = tx.execute_query(stmt).await?;
255///
256/// while let Some(row) = rs.next().await {
257///     let row = row?;
258///     // process row
259/// }
260/// # Ok(())
261/// # }
262/// ```
263#[derive(Clone, Debug)]
264pub struct Statement {
265    pub(crate) sql: String,
266    pub(crate) params: BTreeMap<String, Value>,
267    pub(crate) param_types: BTreeMap<String, Type>,
268    pub(crate) request_options: Option<crate::model::RequestOptions>,
269    pub(crate) directed_read_options: Option<DirectedReadOptions>,
270    pub(crate) query_options: Option<QueryOptions>,
271    pub(crate) query_mode: Option<QueryMode>,
272    pub(crate) last_statement: bool,
273    pub(crate) gax_options: GaxRequestOptions,
274}
275
276impl Statement {
277    /// Creates a new statement builder.
278    pub fn builder(sql: impl Into<String>) -> StatementBuilder {
279        StatementBuilder::new(sql)
280    }
281
282    /// Returns the SQL query string of this statement.
283    pub fn sql(&self) -> &str {
284        &self.sql
285    }
286
287    pub(crate) fn gax_options(&self) -> &GaxRequestOptions {
288        &self.gax_options
289    }
290
291    /// Returns a new `Statement` with the given `GaxRequestOptions`.
292    pub(crate) fn with_gax_options(mut self, options: GaxRequestOptions) -> Self {
293        self.gax_options = options;
294        self
295    }
296
297    /// Sets the query mode to use for this statement.
298    ///
299    /// # Example
300    /// ```
301    /// # use google_cloud_spanner::statement::Statement;
302    /// # use google_cloud_spanner::model::execute_sql_request::QueryMode;
303    /// # use google_cloud_spanner::transaction::SingleUseReadOnlyTransaction;
304    /// # async fn test_doc(tx: SingleUseReadOnlyTransaction) -> Result<(), google_cloud_spanner::Error> {
305    /// let statement = Statement::builder("SELECT * FROM users WHERE id = @id")
306    ///     .add_param("id", &42)
307    ///     .build();
308    /// let mut query_plan = tx.execute_query(statement.clone().set_query_mode(QueryMode::Plan)).await?;
309    /// # Ok(())
310    /// # }
311    /// ```
312    ///
313    /// This method consumes the statement and returns a new one with the specified mode.
314    pub fn set_query_mode(mut self, mode: QueryMode) -> Self {
315        self.query_mode = Some(mode);
316        self
317    }
318
319    /// Sets whether this statement is the last statement in a read/write transaction.
320    ///
321    /// # Example
322    /// ```
323    /// # use google_cloud_spanner::client::Spanner;
324    /// # use google_cloud_spanner::statement::Statement;
325    /// # async fn run_tx(client: Spanner) -> Result<(), google_cloud_spanner::Error> {
326    /// let db_client = client.database_client("projects/p/instances/i/databases/d").build().await?;
327    /// let runner = db_client.read_write_transaction().build().await?;
328    ///
329    /// let result = runner.run(async |transaction| {
330    ///     let statement = Statement::builder("UPDATE MyTable SET MyColumn = 'MyValue' WHERE Id = 1")
331    ///         .build()
332    ///         .set_last_statement(true);
333    ///     transaction.execute_update(statement).await?;
334    ///     Ok(42)
335    /// }).await?;
336    /// # Ok(())
337    /// # }
338    /// ```
339    ///
340    /// If true, indicates that this is the last statement in a transaction, which allows
341    /// Spanner to optimize execution.
342    pub fn set_last_statement(mut self, last_statement: bool) -> Self {
343        self.last_statement = last_statement;
344        self
345    }
346
347    fn into_parts(
348        self,
349    ) -> (
350        String,
351        Option<wkt::Struct>,
352        std::collections::HashMap<String, crate::model::Type>,
353    ) {
354        let params: Option<wkt::Struct> = if self.params.is_empty() {
355            None
356        } else {
357            Some(
358                self.params
359                    .into_iter()
360                    .map(|(k, v)| (k, v.into_serde_value()))
361                    .collect(),
362            )
363        };
364        let param_types: std::collections::HashMap<String, crate::model::Type> = self
365            .param_types
366            .into_iter()
367            .map(|(k, v)| (k, v.0))
368            .collect();
369        (self.sql, params, param_types)
370    }
371
372    pub(crate) fn into_request(self) -> crate::model::ExecuteSqlRequest {
373        let request_options = self.request_options.clone();
374        let directed_read_options = self.directed_read_options.clone();
375        let query_options = self.query_options.clone();
376        let query_mode = self.query_mode.clone();
377        let last_statement = self.last_statement;
378        let (sql, params, param_types) = self.into_parts();
379        crate::model::ExecuteSqlRequest::default()
380            .set_sql(sql)
381            .set_or_clear_params(params)
382            .set_param_types(param_types)
383            .set_or_clear_request_options(request_options)
384            .set_or_clear_directed_read_options(directed_read_options)
385            .set_or_clear_query_options(query_options)
386            .set_query_mode(query_mode.unwrap_or_default())
387            .set_last_statement(last_statement)
388    }
389
390    pub(crate) fn into_batch_statement(self) -> crate::model::execute_batch_dml_request::Statement {
391        let (sql, params, param_types) = self.into_parts();
392        crate::model::execute_batch_dml_request::Statement::default()
393            .set_sql(sql)
394            .set_or_clear_params(params)
395            .set_param_types(param_types)
396    }
397
398    pub(crate) fn into_partition_query_request(self) -> crate::model::PartitionQueryRequest {
399        let (sql, params, param_types) = self.into_parts();
400        crate::model::PartitionQueryRequest::default()
401            .set_sql(sql)
402            .set_or_clear_params(params)
403            .set_param_types(param_types)
404    }
405}
406
407impl From<StatementBuilder> for Statement {
408    fn from(builder: StatementBuilder) -> Self {
409        builder.build()
410    }
411}
412
413impl From<String> for Statement {
414    fn from(sql: String) -> Self {
415        Statement::builder(sql).build()
416    }
417}
418
419impl From<&str> for Statement {
420    fn from(sql: &str) -> Self {
421        Statement::builder(sql).build()
422    }
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428    use anyhow::Context;
429
430    #[test]
431    fn test_auto_traits() {
432        static_assertions::assert_impl_all!(Statement: Clone, std::fmt::Debug, Send, Sync);
433        static_assertions::assert_impl_all!(StatementBuilder: Clone, std::fmt::Debug, Send, Sync);
434    }
435
436    #[test]
437    fn test_untyped_param() {
438        let stmt = Statement::builder("SELECT * FROM users WHERE age > @age")
439            .add_param("age", &21)
440            .build();
441
442        assert_eq!(stmt.sql, "SELECT * FROM users WHERE age > @age");
443        assert_eq!(stmt.param_types.len(), 0);
444        assert_eq!(stmt.params.len(), 1);
445        assert_eq!(stmt.request_options, None);
446
447        let val = stmt.params.get("age").unwrap();
448        assert_eq!(val.as_string(), "21");
449    }
450
451    #[test]
452    fn test_typed_param() {
453        use crate::types;
454        let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
455            .add_typed_param("id", &"user-123", types::string())
456            .build();
457
458        assert_eq!(stmt.param_types.len(), 1);
459        assert_eq!(stmt.param_types.get("id").unwrap(), &types::string());
460
461        assert_eq!(stmt.params.len(), 1);
462        let val = stmt.params.get("id").unwrap();
463        assert_eq!(val.as_string(), "user-123");
464    }
465
466    #[test]
467    fn test_multiple_params() {
468        use crate::types;
469        let stmt = Statement::builder("SELECT * FROM users WHERE age > @age AND role = @role")
470            .add_param("age", &21)
471            .add_typed_param("role", &"admin", types::string())
472            .build();
473
474        assert_eq!(stmt.params.len(), 2);
475        assert_eq!(stmt.param_types.len(), 1);
476    }
477
478    #[test]
479    fn test_from_string_conversions() {
480        let stmt_str: Statement = "SELECT 1".into();
481        let stmt_string: Statement = "SELECT 1".to_string().into();
482        assert_eq!(stmt_str.sql, "SELECT 1");
483        assert_eq!(stmt_string.sql, "SELECT 1");
484        assert!(stmt_str.params.is_empty());
485        assert!(stmt_string.params.is_empty());
486        assert!(stmt_str.param_types.is_empty());
487        assert!(stmt_string.param_types.is_empty());
488        assert!(stmt_str.request_options.is_none());
489        assert!(stmt_string.request_options.is_none());
490    }
491
492    #[test]
493    fn test_from_builder_conversion() {
494        use crate::types;
495        let builder = Statement::builder("SELECT * FROM users WHERE age > @age AND role = @role")
496            .add_param("age", &21)
497            .add_typed_param("role", &"admin", types::string());
498
499        let stmt: Statement = builder.into();
500        assert_eq!(
501            stmt.sql,
502            "SELECT * FROM users WHERE age > @age AND role = @role"
503        );
504        assert_eq!(stmt.params.len(), 2);
505        assert_eq!(stmt.param_types.len(), 1);
506    }
507
508    #[test]
509    fn test_into_request() {
510        use crate::types;
511        let stmt = Statement::builder("SELECT * FROM users WHERE age > @age AND role = @role")
512            .add_param("age", &21)
513            .add_typed_param("role", &"admin", types::string())
514            .build();
515
516        let req = stmt.into_request();
517
518        let params = req
519            .params
520            .expect("ExecuteSqlRequest parameters should be set after into_request conversion");
521        assert_eq!(params.len(), 2);
522        assert!(params.contains_key("age"));
523        assert!(params.contains_key("role"));
524
525        let param_types = req.param_types;
526        assert_eq!(param_types.len(), 1);
527        assert!(param_types.contains_key("role"));
528    }
529
530    #[test]
531    fn with_request_tag() {
532        let stmt = Statement::builder("SELECT * FROM users")
533            .set_request_tag("tag1")
534            .build();
535        assert_eq!(
536            stmt.request_options
537                .expect("request options missing")
538                .request_tag,
539            "tag1"
540        );
541    }
542
543    #[test]
544    fn with_priority() {
545        let stmt = Statement::builder("SELECT * FROM users")
546            .set_priority(Priority::High)
547            .build();
548        assert_eq!(
549            stmt.request_options
550                .expect("request options missing")
551                .priority,
552            Priority::High
553        );
554    }
555
556    #[test]
557    fn with_directed_read_options() {
558        let dro = DirectedReadOptions::default();
559        let stmt = Statement::builder("SELECT * FROM users")
560            .set_directed_read_options(dro.clone())
561            .build();
562        assert_eq!(stmt.directed_read_options, Some(dro));
563    }
564
565    #[test]
566    fn with_query_options() -> anyhow::Result<()> {
567        let query_options = QueryOptions::default().set_optimizer_version("1");
568        let stmt = Statement::builder("SELECT * FROM users")
569            .set_query_options(query_options.clone())
570            .build();
571        assert_eq!(
572            stmt.query_options
573                .as_ref()
574                .context("query options missing")?
575                .optimizer_version,
576            "1"
577        );
578
579        let req = stmt.into_request();
580        assert_eq!(
581            req.query_options
582                .context("query options missing in request")?
583                .optimizer_version,
584            "1"
585        );
586        Ok(())
587    }
588
589    #[test]
590    fn with_query_mode() -> anyhow::Result<()> {
591        let stmt = Statement::builder("SELECT * FROM users")
592            .set_query_mode(QueryMode::Plan)
593            .build();
594        assert_eq!(stmt.query_mode, Some(QueryMode::Plan));
595
596        let req = stmt.into_request();
597        assert_eq!(req.query_mode, QueryMode::Plan);
598        Ok(())
599    }
600
601    #[test]
602    fn statement_with_query_mode() -> anyhow::Result<()> {
603        let stmt = Statement::builder("SELECT * FROM users").build();
604        assert_eq!(stmt.query_mode, None);
605
606        let stmt = stmt.set_query_mode(QueryMode::Profile);
607        assert_eq!(stmt.query_mode, Some(QueryMode::Profile));
608
609        let req = stmt.into_request();
610        assert_eq!(req.query_mode, QueryMode::Profile);
611        Ok(())
612    }
613
614    #[test]
615    fn with_gax_options() -> anyhow::Result<()> {
616        use google_cloud_gax::exponential_backoff::ExponentialBackoff;
617        use google_cloud_gax::retry_policy::NeverRetry;
618        use std::time::Duration;
619
620        let stmt = Statement::builder("SELECT * FROM users")
621            .with_attempt_timeout(Duration::from_secs(10))
622            .with_retry_policy(NeverRetry)
623            .with_backoff_policy(ExponentialBackoff::default())
624            .build();
625
626        assert_eq!(
627            stmt.gax_options.attempt_timeout(),
628            &Some(Duration::from_secs(10))
629        );
630        assert!(stmt.gax_options.retry_policy().is_some());
631        assert!(stmt.gax_options.backoff_policy().is_some());
632
633        Ok(())
634    }
635
636    #[test]
637    fn set_last_statement() -> anyhow::Result<()> {
638        let stmt = Statement::builder("SELECT * FROM users").build();
639        assert!(!stmt.last_statement);
640
641        let req = stmt.into_request();
642        assert!(!req.last_statement);
643
644        let stmt = Statement::builder("SELECT * FROM users")
645            .set_last_statement(true)
646            .build();
647        assert!(stmt.last_statement);
648
649        let req = stmt.into_request();
650        assert!(req.last_statement);
651
652        let stmt = Statement::builder("SELECT * FROM users")
653            .build()
654            .set_last_statement(true);
655        assert!(stmt.last_statement);
656
657        let req = stmt.into_request();
658        assert!(req.last_statement);
659
660        Ok(())
661    }
662}
google_cloud_spanner/statement.rs

google_cloud_spanner/
statement.rs
