feat(spanner): inline BeginTransaction with the first statement (work in progress, do not review)

deny.toml

@@ -115,6 +115,7 @@ wrappers = [
   # Use in tests is fine.
   "grpc-server",
   "integration-tests-o11y",
+  "integration-tests-spanner",
   "pubsub-grpc-mock",
   "spanner-grpc-mock",
   "storage-grpc-mock",

src/spanner/grpc-mock/src/lib.rs

@@ -65,6 +65,8 @@ pub mod google {
     }
     pub mod spanner {
         pub mod v1 {
+            #![allow(rustdoc::broken_intra_doc_links)]
+            #![allow(rustdoc::bare_urls)]
             include!("generated/protos/google.spanner.v1.rs");
         }
     }

src/spanner/src/batch_read_only_transaction.rs

@@ -16,7 +16,8 @@ use crate::database_client::DatabaseClient;
 use crate::model::PartitionOptions;
 use crate::precommit::PrecommitTokenTracker;
 use crate::read_only_transaction::{
-    MultiUseReadOnlyTransaction, MultiUseReadOnlyTransactionBuilder,
+    BeginTransactionOption, MultiUseReadOnlyTransaction, MultiUseReadOnlyTransactionBuilder,
+    ReadContextTransactionSelector,
 };
 use crate::result_set::{ResultSet, StreamOperation};
 use crate::statement::Statement;
@@ -44,7 +45,8 @@ pub struct BatchReadOnlyTransactionBuilder {
 impl BatchReadOnlyTransactionBuilder {
     pub(crate) fn new(client: DatabaseClient) -> Self {
         Self {
-            inner: MultiUseReadOnlyTransactionBuilder::new(client),
+            inner: MultiUseReadOnlyTransactionBuilder::new(client)
+                .with_begin_transaction_option(BeginTransactionOption::ExplicitBegin),
         }
     }
 
@@ -143,12 +145,13 @@ impl BatchReadOnlyTransaction {
         statement: T,
         options: PartitionOptions,
     ) -> crate::Result<Vec<Partition>> {
+        let selector = self.inner.context.transaction_selector.selector().await?;
         let statement = statement.into();
         let request = statement
             .clone()
             .into_partition_query_request()
             .set_session(self.inner.context.client.session.name.clone())
-            .set_transaction(self.inner.context.transaction_selector.clone())
+            .set_transaction(selector.clone())
             .set_partition_options(options);
 
         let response = self
@@ -165,7 +168,7 @@ impl BatchReadOnlyTransaction {
             .map(|p| Partition {
                 inner: PartitionedOperation::Query {
                     partition_token: p.partition_token,
-                    transaction_selector: self.inner.context.transaction_selector.clone(),
+                    transaction_selector: selector.clone(),
                     session_name: self.inner.context.client.session.name.clone(),
                     statement: statement.clone(),
                 },
@@ -198,12 +201,13 @@ impl BatchReadOnlyTransaction {
         read: T,
         options: PartitionOptions,
     ) -> crate::Result<Vec<Partition>> {
+        let selector = self.inner.context.transaction_selector.selector().await?;
         let read = read.into();
         let request = read
             .clone()
             .into_partition_read_request()
             .set_session(self.inner.context.client.session.name.clone())
-            .set_transaction(self.inner.context.transaction_selector.clone())
+            .set_transaction(selector.clone())
             .set_partition_options(options);
 
         let response = self
@@ -220,7 +224,7 @@ impl BatchReadOnlyTransaction {
             .map(|p| Partition {
                 inner: PartitionedOperation::Read {
                     partition_token: p.partition_token,
-                    transaction_selector: self.inner.context.transaction_selector.clone(),
+                    transaction_selector: selector.clone(),
                     session_name: self.inner.context.client.session.name.clone(),
                     read_request: read.clone(),
                 },
@@ -344,6 +348,10 @@ impl Partition {
 
         Ok(ResultSet::new(
             stream,
+            Some(ReadContextTransactionSelector::Fixed(
+                transaction_selector.clone(),
+                None,
+            )),
             PrecommitTokenTracker::new_noop(),
             client.clone(),
             StreamOperation::Query(request),
@@ -373,6 +381,10 @@ impl Partition {
 
         Ok(ResultSet::new(
             stream,
+            Some(ReadContextTransactionSelector::Fixed(
+                transaction_selector.clone(),
+                None,
+            )),
             PrecommitTokenTracker::new_noop(),
             client.clone(),
             StreamOperation::Read(request),

src/spanner/src/client.rs

@@ -26,6 +26,7 @@ pub use crate::mutation::{Mutation, ValueBinder, WriteBuilder};
 pub use crate::read::ConfiguredReadRequestBuilder;
 pub use crate::read::ReadRequest;
 pub use crate::read::ReadRequestBuilder;
+pub use crate::read_only_transaction::BeginTransactionOption;
 pub use crate::read_only_transaction::MultiUseReadOnlyTransaction;
 pub use crate::read_only_transaction::MultiUseReadOnlyTransactionBuilder;
 pub use crate::read_only_transaction::SingleUseReadOnlyTransaction;