google_cloud_spanner/
batch_read_only_transaction.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::database_client::DatabaseClient;
16use crate::model::PartitionOptions;
17use crate::precommit::PrecommitTokenTracker;
18use crate::read_only_transaction::{
19    BeginTransactionOption, MultiUseReadOnlyTransaction, MultiUseReadOnlyTransactionBuilder,
20    ReadContextTransactionSelector,
21};
22use crate::result_set::{ResultSet, ResultSetParams, StreamOperation};
23use crate::statement::Statement;
24use crate::timestamp_bound::TimestampBound;
25use google_cloud_gax::backoff_policy::BackoffPolicyArg;
26use google_cloud_gax::options::RequestOptions as GaxRequestOptions;
27use google_cloud_gax::retry_policy::RetryPolicyArg;
28use serde::{Deserialize, Serialize};
29use std::time::Duration;
30
31/// A builder for [BatchReadOnlyTransaction].
32///
33/// # Example
34/// ```
35/// # use google_cloud_spanner::client::Spanner;
36/// # use google_cloud_spanner::transaction::TimestampBound;
37/// # async fn build_tx(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
38/// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
39/// let read_only_transaction = db_client.batch_read_only_transaction()
40///     .set_timestamp_bound(TimestampBound::strong())
41///     .build()
42///     .await?;
43/// # Ok(())
44/// # }
45/// ```
46pub struct BatchReadOnlyTransactionBuilder {
47    inner: MultiUseReadOnlyTransactionBuilder,
48}
49
50impl BatchReadOnlyTransactionBuilder {
51    pub(crate) fn new(client: DatabaseClient) -> Self {
52        Self {
53            inner: MultiUseReadOnlyTransactionBuilder::new(client)
54                .with_begin_transaction_option(BeginTransactionOption::ExplicitBegin),
55        }
56    }
57
58    /// Sets the timestamp bound for the read-only transaction.
59    ///
60    /// # Example
61    /// ```
62    /// # use google_cloud_spanner::client::Spanner;
63    /// # use google_cloud_spanner::transaction::TimestampBound;
64    /// # async fn set_bound(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
65    /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
66    /// let builder = db_client.batch_read_only_transaction().set_timestamp_bound(TimestampBound::strong());
67    /// # Ok(())
68    /// # }
69    /// ```
70    pub fn set_timestamp_bound(self, bound: TimestampBound) -> Self {
71        Self {
72            inner: self.inner.set_timestamp_bound(bound),
73        }
74    }
75
76    /// Builds the [BatchReadOnlyTransaction] and starts the transaction
77    /// by calling the `BeginTransaction` RPC.
78    ///
79    /// # Example
80    /// ```
81    /// # use google_cloud_spanner::client::Spanner;
82    /// # async fn build(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
83    /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
84    /// let transaction = db_client.batch_read_only_transaction().build().await?;
85    /// # Ok(())
86    /// # }
87    /// ```
88    pub async fn build(self) -> crate::Result<BatchReadOnlyTransaction> {
89        let inner = self.inner.build().await?;
90        Ok(BatchReadOnlyTransaction { inner })
91    }
92}
93
94/// A read-only transaction that can be used to partition reads and queries
95/// and execute these in parallel across multiple workers.
96///
97/// # Example
98/// ```
99/// # use google_cloud_spanner::client::Spanner;
100/// # use google_cloud_spanner::statement::Statement;
101/// # use google_cloud_spanner::model::PartitionOptions;
102///
103/// # async fn run(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
104/// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
105/// let transaction = db_client.batch_read_only_transaction().build().await?;
106/// let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
107///     .add_param("id", &42)
108///     .build();
109/// let options = PartitionOptions::default()
110///     .set_max_partitions(10);
111/// let partitions = transaction.partition_query(stmt, options).await?;
112///
113/// // partitions can be sent to other workers for parallel execution
114/// # Ok(())
115/// # }
116/// ```
117#[derive(Debug)]
118pub struct BatchReadOnlyTransaction {
119    inner: MultiUseReadOnlyTransaction,
120}
121
122impl BatchReadOnlyTransaction {
123    /// Returns the read timestamp chosen for the transaction.
124    pub fn read_timestamp(&self) -> Option<wkt::Timestamp> {
125        self.inner.read_timestamp()
126    }
127
128    /// Creates a set of partitions that can be used to execute a query in parallel.
129    ///
130    /// # Example
131    /// ```
132    /// # use google_cloud_spanner::client::Spanner;
133    /// # use google_cloud_spanner::statement::Statement;
134    /// # use google_cloud_spanner::model::PartitionOptions;
135    /// # async fn run(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
136    /// let db = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
137    /// let transaction = db.batch_read_only_transaction().build().await?;
138    ///
139    /// let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
140    ///     .add_param("id", &42)
141    ///     .build();
142    /// let options = PartitionOptions::default()
143    ///     .set_max_partitions(10);
144    /// let partitions = transaction.partition_query(stmt, options).await?;
145    /// # Ok(())
146    /// # }
147    /// ```
148    pub async fn partition_query<T: Into<Statement>>(
149        &self,
150        statement: T,
151        options: PartitionOptions,
152    ) -> crate::Result<Vec<Partition>> {
153        let selector = self.inner.context.transaction_selector.selector().await?;
154        let statement = statement.into();
155        let request = statement
156            .clone()
157            .into_partition_query_request()
158            .set_session(self.inner.context.session_name.clone())
159            .set_transaction(selector.clone())
160            .set_partition_options(options);
161
162        let response = self
163            .inner
164            .context
165            .client
166            .spanner
167            .partition_query(
168                request,
169                crate::RequestOptions::default(),
170                self.inner.context.channel_hint,
171                &self.inner.context.client.o11y,
172            )
173            .await?;
174
175        Ok(response
176            .partitions
177            .into_iter()
178            .map(|p| {
179                let mut req = statement.clone().into_request();
180                req.session = self.inner.context.session_name.clone();
181                req.transaction = Some(selector.clone());
182                req.partition_token = p.partition_token;
183
184                Partition {
185                    inner: PartitionedOperation::Query(req),
186                    gax_options: GaxRequestOptions::default(),
187                }
188            })
189            .collect())
190    }
191
192    /// Creates a set of partitions that can be used to execute a read in parallel.
193    ///
194    /// # Example
195    /// ```
196    /// # use google_cloud_spanner::client::Spanner;
197    /// # use google_cloud_spanner::key::KeySet;
198    /// # use google_cloud_spanner::read::ReadRequest;
199    /// # use google_cloud_spanner::model::PartitionOptions;
200    /// # async fn run(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
201    /// let db = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
202    /// let transaction = db.batch_read_only_transaction().build().await?;
203    ///
204    /// let read = ReadRequest::builder("users", vec!["id".to_string(), "name".to_string()])
205    ///     .with_keys(KeySet::all())
206    ///     .build();
207    /// let options = PartitionOptions::default()
208    ///     .set_max_partitions(10);
209    /// let partitions = transaction.partition_read(read, options).await?;
210    /// # Ok(())
211    /// # }
212    /// ```
213    pub async fn partition_read<T: Into<crate::read::ReadRequest>>(
214        &self,
215        read: T,
216        options: PartitionOptions,
217    ) -> crate::Result<Vec<Partition>> {
218        let selector = self.inner.context.transaction_selector.selector().await?;
219        let read = read.into();
220        let request = read
221            .clone()
222            .into_partition_read_request()
223            .set_session(self.inner.context.session_name.clone())
224            .set_transaction(selector.clone())
225            .set_partition_options(options);
226
227        let response = self
228            .inner
229            .context
230            .client
231            .spanner
232            .partition_read(
233                request,
234                crate::RequestOptions::default(),
235                self.inner.context.channel_hint,
236                &self.inner.context.client.o11y,
237            )
238            .await?;
239
240        Ok(response
241            .partitions
242            .into_iter()
243            .map(|p| {
244                let mut req = read.clone().into_request();
245                req.session = self.inner.context.session_name.clone();
246                req.transaction = Some(selector.clone());
247                req.partition_token = p.partition_token;
248
249                Partition {
250                    inner: PartitionedOperation::Read(req),
251                    gax_options: GaxRequestOptions::default(),
252                }
253            })
254            .collect())
255    }
256}
257
258/// Defines the segments of data to be read in a partitioned read or query.
259/// These partitions can be serialized and processed across several
260/// different machines or processes.
261#[derive(Clone, Debug, Serialize, Deserialize)]
262pub struct Partition {
263    pub(crate) inner: PartitionedOperation,
264    #[serde(skip)]
265    pub(crate) gax_options: GaxRequestOptions,
266}
267
268impl Partition {
269    /// Sets whether Data Boost is enabled for this partition.
270    ///
271    /// # Example
272    /// ```
273    /// # use google_cloud_spanner::client::Spanner;
274    /// # use google_cloud_spanner::statement::Statement;
275    /// # use google_cloud_spanner::model::PartitionOptions;
276    /// # async fn run_query(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
277    /// # let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
278    /// # let transaction = db_client.batch_read_only_transaction().build().await?;
279    /// # let partitions = transaction.partition_query(Statement::builder("SELECT * FROM Users").build(), PartitionOptions::default()).await?;
280    /// // On a worker receiving a partition, execute it with Data Boost:
281    /// let mut result_set = partitions[0].clone()
282    ///     .set_data_boost(true)
283    ///     .execute(&db_client)
284    ///     .await?;
285    /// # Ok(())
286    /// # }
287    /// ```
288    pub fn set_data_boost(mut self, enabled: bool) -> Self {
289        match &mut self.inner {
290            PartitionedOperation::Query(req) => req.data_boost_enabled = enabled,
291            PartitionedOperation::Read(req) => req.data_boost_enabled = enabled,
292        }
293        self
294    }
295
296    /// Sets the per-attempt timeout for this partition execution.
297    ///
298    /// **Note:** This field is **not serialized**. Each host that executes a partition must set its own attempt timeout.
299    pub fn with_attempt_timeout(mut self, timeout: Duration) -> Self {
300        self.gax_options.set_attempt_timeout(timeout);
301        self
302    }
303
304    /// Sets the retry policy for this partition execution.
305    ///
306    /// **Note:** This field is **not serialized**. Each host that executes a partition must set its own retry policy.
307    pub fn with_retry_policy(mut self, policy: impl Into<RetryPolicyArg>) -> Self {
308        self.gax_options.set_retry_policy(policy);
309        self
310    }
311
312    /// Sets the backoff policy for this partition execution.
313    ///
314    /// **Note:** This field is **not serialized**. Each host that executes a partition must set its own backoff policy.
315    pub fn with_backoff_policy(mut self, policy: impl Into<BackoffPolicyArg>) -> Self {
316        self.gax_options.set_backoff_policy(policy);
317        self
318    }
319
320    /// Executes this partition and returns a [ResultSet] that
321    /// contains the rows that belong to this partition.
322    ///
323    /// # Example: executing a query partition
324    /// ```
325    /// # use google_cloud_spanner::client::Spanner;
326    /// # use google_cloud_spanner::statement::Statement;
327    /// # use google_cloud_spanner::model::PartitionOptions;
328    /// # async fn run_query(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
329    /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
330    /// let transaction = db_client.batch_read_only_transaction().build().await?;
331    /// let partitions = transaction.partition_query(
332    ///     Statement::builder("SELECT * FROM Users").build(),
333    ///     PartitionOptions::default()
334    /// ).await?;
335    ///
336    /// // ... send partitions to other workers ...
337    ///
338    /// // On a worker receiving a partition, execute it:
339    /// let mut result_set = partitions[0].execute(&db_client).await?;
340    /// while let Some(row) = result_set.next().await.transpose()? {
341    ///     // process row
342    /// }
343    /// # Ok(())
344    /// # }
345    /// ```
346    /// # Example: executing a read partition
347    /// ```
348    /// # use google_cloud_spanner::client::Spanner;
349    /// # use google_cloud_spanner::key::KeySet;
350    /// # use google_cloud_spanner::read::ReadRequest;
351    /// # use google_cloud_spanner::model::PartitionOptions;
352    /// # async fn run_read(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
353    /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
354    /// let transaction = db_client.batch_read_only_transaction().build().await?;
355    /// let req = ReadRequest::builder("Users", vec!["Id", "Name"]).with_keys(KeySet::all()).build();
356    /// let partitions = transaction.partition_read(req, PartitionOptions::default()).await?;
357    ///
358    /// // ... send partitions to other workers ...
359    ///
360    /// // On a worker receiving a partition, execute it:
361    /// let mut result_set = partitions[0].execute(&db_client).await?;
362    /// while let Some(row) = result_set.next().await.transpose()? {
363    ///     // process row
364    /// }
365    /// # Ok(())
366    /// # }
367    /// ```
368    ///
369    /// A partition can be executed by any `DatabaseClient` that is connected to
370    /// the database that the partitions belong to.
371    pub async fn execute(&self, client: &DatabaseClient) -> crate::Result<ResultSet> {
372        match &self.inner {
373            PartitionedOperation::Query(req) => {
374                Self::execute_query(client, req, self.gax_options.clone()).await
375            }
376            PartitionedOperation::Read(req) => {
377                Self::execute_read(client, req, self.gax_options.clone()).await
378            }
379        }
380    }
381
382    async fn execute_query(
383        client: &DatabaseClient,
384        req: &crate::model::ExecuteSqlRequest,
385        gax_options: GaxRequestOptions,
386    ) -> crate::Result<ResultSet> {
387        let channel_hint = client.spanner.next_channel_hint();
388        let stream = client
389            .spanner
390            .execute_streaming_sql(req.clone(), gax_options.clone(), channel_hint)
391            .send()
392            .await?;
393
394        ResultSet::create(ResultSetParams {
395            stream,
396            transaction_selector: Some(ReadContextTransactionSelector::Fixed(
397                req.transaction
398                    .clone()
399                    .expect("transaction must be set in partition request"),
400                None,
401            )),
402            precommit_token_tracker: PrecommitTokenTracker::new_noop(),
403            client: client.clone(),
404            session_name: req.session.clone(),
405            transaction_tag: None,
406            operation: StreamOperation::Query(req.clone()),
407            channel_hint,
408            gax_options,
409        })
410        .await
411    }
412
413    async fn execute_read(
414        client: &DatabaseClient,
415        req: &crate::model::ReadRequest,
416        gax_options: GaxRequestOptions,
417    ) -> crate::Result<ResultSet> {
418        let channel_hint = client.spanner.next_channel_hint();
419        let stream = client
420            .spanner
421            .streaming_read(req.clone(), gax_options.clone(), channel_hint)
422            .send()
423            .await?;
424
425        ResultSet::create(ResultSetParams {
426            stream,
427            transaction_selector: Some(ReadContextTransactionSelector::Fixed(
428                req.transaction
429                    .clone()
430                    .expect("transaction must be set in partition request"),
431                None,
432            )),
433            precommit_token_tracker: PrecommitTokenTracker::new_noop(),
434            client: client.clone(),
435            session_name: req.session.clone(),
436            transaction_tag: None,
437            operation: StreamOperation::Read(req.clone()),
438            channel_hint,
439            gax_options,
440        })
441        .await
442    }
443}
444
445#[derive(Clone, Debug, Serialize, Deserialize)]
446pub(crate) enum PartitionedOperation {
447    Query(crate::model::ExecuteSqlRequest),
448    Read(crate::model::ReadRequest),
449}
450
451#[cfg(test)]
452pub(crate) mod tests {
453    use super::*;
454    use crate::key::KeySet;
455    use crate::model::transaction_selector::Selector;
456    use crate::model::{ExecuteSqlRequest, ReadRequest as GrpcReadRequest, TransactionSelector};
457    use crate::read::ReadRequest as SpannerReadRequest;
458    use crate::read_only_transaction::tests::{create_session_mock, setup_db_client};
459    use crate::statement::Statement;
460    use crate::transaction::TimestampBound;
461    use gaxi::grpc::tonic::Response;
462    use google_cloud_test_macros::tokio_test_no_panics;
463    use prost_types::Timestamp;
464    use spanner_grpc_mock::google::spanner::v1::{
465        PartialResultSet, Partition as MockPartition, PartitionResponse, ResultSetMetadata,
466        StructType, Transaction,
467    };
468    use static_assertions::assert_impl_all;
469    use std::fmt::Debug;
470
471    #[test]
472    fn auto_traits() {
473        assert_impl_all!(BatchReadOnlyTransactionBuilder: Send, Sync);
474        assert_impl_all!(BatchReadOnlyTransaction: Send, Sync, Debug);
475        assert_impl_all!(Partition: Send, Sync, Debug);
476    }
477
478    #[test]
479    fn serialize_partition_skips_gax_options() -> anyhow::Result<()> {
480        use std::time::Duration;
481
482        let req = crate::model::ExecuteSqlRequest::new()
483            .set_sql("SELECT 1")
484            .set_partition_token(b"token".to_vec());
485
486        let mut gax_options = GaxRequestOptions::default();
487        gax_options.set_attempt_timeout(Duration::from_secs(5));
488        let partition = Partition {
489            inner: PartitionedOperation::Query(req),
490            gax_options,
491        };
492
493        let serialized = serde_json::to_string(&partition)?;
494        let deserialized: Partition = serde_json::from_str(&serialized)?;
495
496        // Verify that gax_options was NOT preserved (it uses default, which is None timeout)
497        assert_eq!(*deserialized.gax_options.attempt_timeout(), None);
498
499        Ok(())
500    }
501
502    fn setup_select1() -> PartialResultSet {
503        PartialResultSet {
504            metadata: Some(ResultSetMetadata {
505                row_type: Some(StructType {
506                    fields: vec![Default::default()],
507                }),
508                ..Default::default()
509            }),
510            values: vec![prost_types::Value {
511                kind: Some(prost_types::value::Kind::StringValue("1".to_string())),
512            }],
513            last: true,
514            ..Default::default()
515        }
516    }
517
518    #[tokio_test_no_panics]
519    async fn partition_execute_respects_options() -> anyhow::Result<()> {
520        use gaxi::grpc::tonic::Response;
521        use std::time::Duration;
522
523        let mut mock = create_session_mock();
524
525        mock.expect_execute_streaming_sql().once().returning(|req| {
526            let timeout = req.metadata().get("grpc-timeout");
527            assert!(timeout.is_some(), "Missing grpc-timeout header");
528            assert_eq!(timeout.unwrap(), "5000000u"); // 5 seconds in micros
529
530            Ok(Response::from(crate::result_set::tests::adapt([Ok(
531                setup_select1(),
532            )])))
533        });
534
535        let (db_client, _server) = setup_db_client(mock).await;
536
537        let req = crate::model::ExecuteSqlRequest::new()
538            .set_session("projects/p/instances/i/databases/d/sessions/123")
539            .set_transaction(crate::model::TransactionSelector {
540                selector: Some(Selector::Id(b"tx_id_1".to_vec().into())),
541                ..Default::default()
542            })
543            .set_sql("SELECT 1")
544            .set_partition_token(b"token".to_vec());
545
546        let partition = Partition {
547            inner: PartitionedOperation::Query(req),
548            gax_options: GaxRequestOptions::default(),
549        };
550
551        let partition = partition.with_attempt_timeout(Duration::from_secs(5));
552
553        let _result_set = partition.execute(&db_client).await?;
554
555        Ok(())
556    }
557
558    #[test]
559    fn serialize_partition_query() -> anyhow::Result<()> {
560        let req = crate::model::ExecuteSqlRequest::new()
561            .set_session("projects/p/instances/i/databases/d/sessions/123")
562            .set_transaction(crate::model::TransactionSelector {
563                selector: Some(crate::model::transaction_selector::Selector::Id(
564                    b"tx_id_1".to_vec().into(),
565                )),
566                ..Default::default()
567            })
568            .set_sql("SELECT * FROM Users")
569            .set_partition_token(b"partition_token_123".to_vec());
570
571        let partition = Partition {
572            inner: PartitionedOperation::Query(req),
573            gax_options: GaxRequestOptions::default(),
574        };
575
576        let serialized = serde_json::to_string(&partition)?;
577        let deserialized: Partition = serde_json::from_str(&serialized)?;
578
579        match &deserialized.inner {
580            PartitionedOperation::Query(r) => {
581                assert_eq!(r.partition_token.as_ref(), b"partition_token_123");
582                assert_eq!(r.sql, "SELECT * FROM Users");
583                assert_eq!(r.session, "projects/p/instances/i/databases/d/sessions/123");
584            }
585            _ => panic!("Expected Query partition"),
586        }
587        Ok(())
588    }
589
590    #[test]
591    fn serialize_partition_read() -> anyhow::Result<()> {
592        let req = crate::model::ReadRequest::new()
593            .set_session("projects/p/instances/i/databases/d/sessions/456")
594            .set_transaction(crate::model::TransactionSelector {
595                selector: Some(crate::model::transaction_selector::Selector::Id(
596                    b"tx_id_2".to_vec().into(),
597                )),
598                ..Default::default()
599            })
600            .set_table("Users")
601            .set_columns(vec!["Id"])
602            .set_partition_token(b"partition_token_456".to_vec());
603
604        let partition = Partition {
605            inner: PartitionedOperation::Read(req),
606            gax_options: GaxRequestOptions::default(),
607        };
608
609        let serialized = serde_json::to_string(&partition)?;
610        let deserialized: Partition = serde_json::from_str(&serialized)?;
611
612        match &deserialized.inner {
613            PartitionedOperation::Read(r) => {
614                assert_eq!(r.partition_token.as_ref(), b"partition_token_456");
615                assert_eq!(r.table, "Users");
616                assert_eq!(r.session, "projects/p/instances/i/databases/d/sessions/456");
617            }
618            _ => panic!("Expected Read partition"),
619        }
620        Ok(())
621    }
622
623    #[tokio_test_no_panics]
624    async fn execute_query() -> anyhow::Result<()> {
625        let mut mock = create_session_mock();
626
627        mock.expect_execute_streaming_sql().once().returning(|req| {
628            let req = req.into_inner();
629            // Verify the partition details were properly stamped onto the request
630            assert_eq!(
631                req.session,
632                "projects/p/instances/i/databases/d/sessions/123"
633            );
634            assert_eq!(req.partition_token, b"partition_token_123".as_slice());
635            assert!(req.transaction.is_some());
636            assert_eq!(req.sql, "SELECT * FROM Users");
637
638            Ok(Response::from(crate::result_set::tests::adapt([Ok(
639                setup_select1(),
640            )])))
641        });
642
643        let (db_client, _server) = setup_db_client(mock).await;
644
645        let req = crate::model::ExecuteSqlRequest::new()
646            .set_session("projects/p/instances/i/databases/d/sessions/123")
647            .set_transaction(crate::model::TransactionSelector {
648                selector: Some(crate::model::transaction_selector::Selector::Id(
649                    b"tx_id_1".to_vec().into(),
650                )),
651                ..Default::default()
652            })
653            .set_sql("SELECT * FROM Users")
654            .set_partition_token(b"partition_token_123".to_vec());
655
656        let partition = Partition {
657            inner: PartitionedOperation::Query(req),
658            gax_options: GaxRequestOptions::default(),
659        };
660
661        let _result_set = partition.execute(&db_client).await?;
662
663        Ok(())
664    }
665
666    #[tokio_test_no_panics]
667    async fn execute_read() -> anyhow::Result<()> {
668        let mut mock = create_session_mock();
669
670        mock.expect_streaming_read().once().returning(|req| {
671            let req = req.into_inner();
672            // Verify the partition details were properly stamped onto the request
673            assert_eq!(
674                req.session,
675                "projects/p/instances/i/databases/d/sessions/456"
676            );
677            assert_eq!(req.partition_token, b"partition_token_456".as_slice());
678            assert!(req.transaction.is_some());
679            assert_eq!(req.table, "Users");
680
681            Ok(Response::from(crate::result_set::tests::adapt([Ok(
682                setup_select1(),
683            )])))
684        });
685
686        let (db_client, _server) = setup_db_client(mock).await;
687
688        let req = crate::model::ReadRequest::new()
689            .set_session("projects/p/instances/i/databases/d/sessions/456")
690            .set_transaction(crate::model::TransactionSelector {
691                selector: Some(crate::model::transaction_selector::Selector::Id(
692                    b"tx_id_2".to_vec().into(),
693                )),
694                ..Default::default()
695            })
696            .set_table("Users")
697            .set_columns(vec!["Id"])
698            .set_partition_token(b"partition_token_456".to_vec());
699
700        let partition = Partition {
701            inner: PartitionedOperation::Read(req),
702            gax_options: GaxRequestOptions::default(),
703        };
704
705        let _result_set = partition.execute(&db_client).await?;
706
707        Ok(())
708    }
709
710    #[tokio_test_no_panics]
711    async fn partition_query() -> anyhow::Result<()> {
712        let mut mock = create_session_mock();
713
714        mock.expect_begin_transaction().once().returning(|req| {
715            let req = req.into_inner();
716            assert_eq!(
717                req.session,
718                "projects/p/instances/i/databases/d/sessions/123"
719            );
720            Ok(Response::new(Transaction {
721                id: vec![1, 2, 3],
722                read_timestamp: Some(Timestamp {
723                    seconds: 123456789,
724                    nanos: 0,
725                }),
726                ..Default::default()
727            }))
728        });
729
730        mock.expect_partition_query().once().returning(|req| {
731            let req = req.into_inner();
732            assert_eq!(
733                req.session,
734                "projects/p/instances/i/databases/d/sessions/123"
735            );
736            assert_eq!(req.sql, "SELECT 1");
737            Ok(Response::new(PartitionResponse {
738                partitions: vec![
739                    MockPartition {
740                        partition_token: vec![10],
741                    },
742                    MockPartition {
743                        partition_token: vec![20],
744                    },
745                ],
746                transaction: None,
747            }))
748        });
749
750        let (db_client, _server) = setup_db_client(mock).await;
751
752        let tx = db_client
753            .batch_read_only_transaction()
754            .set_timestamp_bound(TimestampBound::strong())
755            .build()
756            .await?;
757
758        let ts = tx.read_timestamp().expect("Missing read timestamp");
759        assert_eq!(ts.seconds(), 123456789);
760        assert_eq!(ts.nanos(), 0);
761
762        let partitions = tx
763            .partition_query(
764                Statement::builder("SELECT 1").build(),
765                PartitionOptions::default(),
766            )
767            .await?;
768
769        assert_eq!(partitions.len(), 2);
770
771        match &partitions[0].inner {
772            PartitionedOperation::Query(req) => {
773                assert_eq!(req.partition_token.as_ref(), &[10]);
774                assert_eq!(req.sql, "SELECT 1");
775            }
776            _ => panic!("Expected Query partition"),
777        }
778        Ok(())
779    }
780
781    #[tokio_test_no_panics]
782    async fn partition_read() -> anyhow::Result<()> {
783        let mut mock = create_session_mock();
784
785        mock.expect_begin_transaction().once().returning(|req| {
786            let req = req.into_inner();
787            assert_eq!(
788                req.session,
789                "projects/p/instances/i/databases/d/sessions/123"
790            );
791            Ok(Response::new(Transaction {
792                id: vec![1, 2, 3],
793                read_timestamp: Some(Timestamp {
794                    seconds: 123456789,
795                    nanos: 0,
796                }),
797                ..Default::default()
798            }))
799        });
800
801        mock.expect_partition_read().once().returning(|req| {
802            let req = req.into_inner();
803            assert_eq!(
804                req.session,
805                "projects/p/instances/i/databases/d/sessions/123"
806            );
807            assert_eq!(req.table, "Users");
808            Ok(Response::new(PartitionResponse {
809                partitions: vec![MockPartition {
810                    partition_token: vec![30],
811                }],
812                transaction: None,
813            }))
814        });
815
816        let (db_client, _server) = setup_db_client(mock).await;
817
818        let transaction = db_client.batch_read_only_transaction().build().await?;
819
820        let read = SpannerReadRequest::builder("Users", vec!["Id", "Name"])
821            .with_keys(KeySet::all())
822            .build();
823        let partitions = transaction
824            .partition_read(read, PartitionOptions::default())
825            .await?;
826
827        assert_eq!(partitions.len(), 1);
828
829        match &partitions[0].inner {
830            PartitionedOperation::Read(req) => {
831                assert_eq!(req.partition_token.as_ref(), &[30]);
832                assert_eq!(req.table, "Users");
833            }
834            _ => panic!("Expected Read partition"),
835        }
836        Ok(())
837    }
838
839    #[tokio_test_no_panics]
840    async fn execute_query_with_data_boost() -> anyhow::Result<()> {
841        let mut mock = create_session_mock();
842
843        mock.expect_execute_streaming_sql().once().returning(|req| {
844            let req = req.into_inner();
845            assert!(req.data_boost_enabled, "data_boost_enabled should be true");
846            Ok(Response::from(crate::result_set::tests::adapt([Ok(
847                setup_select1(),
848            )])))
849        });
850
851        let (db_client, _server) = setup_db_client(mock).await;
852
853        let req = ExecuteSqlRequest::new()
854            .set_session("projects/p/instances/i/databases/d/sessions/123")
855            .set_transaction(TransactionSelector {
856                selector: Some(Selector::Id(b"tx_id_1".to_vec().into())),
857                ..Default::default()
858            })
859            .set_sql("SELECT * FROM Users")
860            .set_partition_token(b"partition_token_123".to_vec());
861
862        let partition = Partition {
863            inner: PartitionedOperation::Query(req),
864            gax_options: GaxRequestOptions::default(),
865        };
866
867        let _result_set = partition.set_data_boost(true).execute(&db_client).await?;
868
869        Ok(())
870    }
871
872    #[tokio_test_no_panics]
873    async fn execute_read_with_data_boost() -> anyhow::Result<()> {
874        let mut mock = create_session_mock();
875
876        mock.expect_streaming_read().once().returning(|req| {
877            let req = req.into_inner();
878            assert!(req.data_boost_enabled, "data_boost_enabled should be true");
879            Ok(Response::from(crate::result_set::tests::adapt([Ok(
880                setup_select1(),
881            )])))
882        });
883
884        let (db_client, _server) = setup_db_client(mock).await;
885
886        let req = GrpcReadRequest::new()
887            .set_session("projects/p/instances/i/databases/d/sessions/123")
888            .set_transaction(TransactionSelector {
889                selector: Some(Selector::Id(b"tx_id_2".to_vec().into())),
890                ..Default::default()
891            })
892            .set_table("Users")
893            .set_columns(vec!["Id".to_string(), "Name".to_string()])
894            .set_partition_token(b"partition_token_456".to_vec());
895
896        let partition = Partition {
897            inner: PartitionedOperation::Read(req),
898            gax_options: GaxRequestOptions::default(),
899        };
900
901        let _result_set = partition.set_data_boost(true).execute(&db_client).await?;
902
903        Ok(())
904    }
905}
google_cloud_spanner/batch_read_only_transaction.rs

google_cloud_spanner/
batch_read_only_transaction.rs
