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

google_cloud_spanner/
statement.rs
