google_cloud_spanner/database_client.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::batch_read_only_transaction::BatchReadOnlyTransactionBuilder;
16use crate::batch_write_transaction::BatchWriteTransactionBuilder;
17use crate::client::Spanner;
18use crate::partitioned_dml_transaction::PartitionedDmlTransactionBuilder;
19use crate::read_only_transaction::{
20 MultiUseReadOnlyTransactionBuilder, SingleUseReadOnlyTransactionBuilder,
21};
22use crate::session_maintainer::ManagedSessionMaintainer;
23use std::sync::Arc;
24
25/// A client for interacting with a specific Spanner database.
26///
27/// `DatabaseClient` provides methods to execute transactions and queries.
28/// # Example
29/// ```
30/// # use google_cloud_spanner::client::Spanner;
31/// # async fn sample() -> anyhow::Result<()> {
32/// let spanner = Spanner::builder().build().await?;
33/// let database_client = spanner
34/// .database_client("projects/my-project/instances/my-instance/databases/my-db")
35/// .build()
36/// .await?;
37/// # Ok(())
38/// # }
39/// ```
40///
41/// `DatabaseClient` provides methods to execute transactions and queries.
42/// It holds a single multiplexed session for the database.
43///
44/// A `DatabaseClient` is intended to be a long-lived object, and normally an
45/// application will have a single `DatabaseClient` per database. The client is
46/// thread-safe and should be reused for all operations on the database.
47///
48/// Cloning a `DatabaseClient` is cheap, as it shares the underlying session and channel.
49#[derive(Clone, Debug)]
50pub struct DatabaseClient {
51 pub(crate) spanner: Spanner,
52 pub(crate) session_maintainer: Arc<ManagedSessionMaintainer>,
53 pub(crate) leader_aware_routing_enabled: bool,
54}
55
56impl DatabaseClient {
57 /// Returns a builder for a single-use read-only transaction.
58 ///
59 /// # Example
60 /// ```
61 /// # use google_cloud_spanner::client::Spanner;
62 /// # use google_cloud_spanner::statement::Statement;
63 /// # async fn run(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
64 /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
65 /// let tx = db_client.single_use().build();
66 /// let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
67 /// .add_param("id", &42)
68 /// .build();
69 /// let mut rs = tx.execute_query(stmt).await?;
70 /// # Ok(())
71 /// # }
72 /// ```
73 ///
74 /// A single-use read-only transaction is optimized for the case where only a single
75 /// read or query is needed. This is more efficient than using a read-only transaction
76 /// for a single read or query.
77 pub fn single_use(&self) -> SingleUseReadOnlyTransactionBuilder {
78 SingleUseReadOnlyTransactionBuilder::new(self.clone())
79 }
80
81 /// Returns a builder for a multi-use read-only transaction.
82 ///
83 /// # Example
84 /// ```
85 /// # use google_cloud_spanner::client::Spanner;
86 /// # use google_cloud_spanner::statement::Statement;
87 /// # async fn run(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
88 /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
89 /// let tx = db_client.read_only_transaction().build().await?;
90 /// let stmt = Statement::builder("SELECT * FROM users WHERE id = @id")
91 /// .add_param("id", &42)
92 /// .build();
93 /// let mut rs = tx.execute_query(stmt).await?;
94 /// # Ok(())
95 /// # }
96 /// ```
97 ///
98 /// A read-only transaction can be used to execute multiple reads or queries.
99 /// These transactions guarantee data consistency across multiple read operations,
100 /// but don't permit data modifications. Read-only transactions do not take locks.
101 pub fn read_only_transaction(&self) -> MultiUseReadOnlyTransactionBuilder {
102 MultiUseReadOnlyTransactionBuilder::new(self.clone())
103 }
104
105 /// Returns a builder for a batch read-only transaction.
106 ///
107 /// # Example
108 /// ```
109 /// # use google_cloud_spanner::client::Spanner;
110 /// # use google_cloud_spanner::statement::Statement;
111 /// # async fn build(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
112 /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
113 /// let transaction = db_client.batch_read_only_transaction().build().await?;
114 /// # Ok(())
115 /// # }
116 /// ```
117 ///
118 /// A batch read-only transaction is similar to a read-only transaction, but it allows for partitioning
119 /// a read or query request. Run tasks in parallel over the partitions to execute a large read or query.
120 pub fn batch_read_only_transaction(&self) -> BatchReadOnlyTransactionBuilder {
121 BatchReadOnlyTransactionBuilder::new(self.clone())
122 }
123
124 /// Returns a builder for a partitioned DML transaction.
125 ///
126 /// # Example
127 /// ```
128 /// # use google_cloud_spanner::client::Spanner;
129 /// # use google_cloud_spanner::statement::Statement;
130 /// # async fn run(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
131 /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
132 /// let transaction = db_client.partitioned_dml_transaction().build().await?;
133 /// let statement = Statement::builder("UPDATE users SET active = true WHERE TRUE").build();
134 /// let modified_rows = transaction.execute_update(statement).await?;
135 /// # Ok(())
136 /// # }
137 /// ```
138 ///
139 /// Partitioned DML is used to execute a single DML statement that may modify a large number
140 /// of rows. The execution of the statement will automatically be partitioned into smaller
141 /// transactions by Spanner, which may execute in parallel.
142 ///
143 /// See also: <https://docs.cloud.google.com/spanner/docs/dml-partitioned>
144 pub fn partitioned_dml_transaction(&self) -> PartitionedDmlTransactionBuilder {
145 PartitionedDmlTransactionBuilder::new(self.clone())
146 }
147
148 /// Returns a builder for a read-write transaction runner.
149 ///
150 /// # Example
151 /// ```
152 /// # use google_cloud_spanner::client::Spanner;
153 /// # use google_cloud_spanner::statement::Statement;
154 /// # async fn build(spanner: Spanner) -> Result<(), google_cloud_spanner::Error> {
155 /// let db_client = spanner.database_client("projects/p/instances/i/databases/d").build().await?;
156 /// let runner = db_client.read_write_transaction().build().await?;
157 /// let result = runner.run(async |transaction| {
158 /// let statement = Statement::builder("UPDATE users SET active = true WHERE id = 1").build();
159 /// transaction.execute_update(statement).await?;
160 /// Ok(())
161 /// }).await?;
162 /// # Ok(())
163 /// # }
164 /// ```
165 ///
166 /// Read-write transactions can be used to execute multiple queries and updates
167 /// atomically. If the transaction is aborted by Spanner, the `run` method will
168 /// automatically retry the transaction.
169 pub fn read_write_transaction(&self) -> crate::transaction_runner::TransactionRunnerBuilder {
170 crate::transaction_runner::TransactionRunnerBuilder::new(self.clone())
171 }
172
173 /// Returns a builder for a write-only transaction.
174 ///
175 /// # Example
176 /// ```rust
177 /// # use google_cloud_spanner::client::Spanner;
178 /// # use google_cloud_spanner::mutation::Mutation;
179 /// # async fn test_doc() -> Result<(), Box<dyn std::error::Error>> {
180 /// let client = Spanner::builder().build().await?;
181 /// let db = client.database_client("projects/p/instances/i/databases/d").build().await?;
182 ///
183 /// let mutation = Mutation::new_insert_builder("Users")
184 /// .set("UserId").to(&1)
185 /// .set("UserName").to(&"Alice")
186 /// .build();
187 ///
188 /// let response = db.write_only_transaction()
189 /// .set_transaction_tag("my-tag")
190 /// .build()
191 /// .write(vec![mutation])
192 /// .await?;
193 /// # Ok(())
194 /// # }
195 /// ```
196 ///
197 /// A write-only transaction is used to execute blind writes using mutations.
198 pub fn write_only_transaction(
199 &self,
200 ) -> crate::write_only_transaction::WriteOnlyTransactionBuilder {
201 crate::write_only_transaction::WriteOnlyTransactionBuilder::new(self.clone())
202 }
203
204 /// Returns a builder for a batch write transaction.
205 ///
206 /// # Example
207 /// ```
208 /// # use google_cloud_spanner::client::Spanner;
209 /// # use google_cloud_spanner::mutation::Mutation;
210 /// # use google_cloud_spanner::mutation::MutationGroup;
211 /// # use google_cloud_gax::error::rpc::Code;
212 /// # async fn sample() -> Result<(), Box<dyn std::error::Error>> {
213 /// let client = Spanner::builder().build().await?;
214 /// let db = client.database_client("projects/p/instances/i/databases/d").build().await?;
215 ///
216 /// let mutation1a = Mutation::new_insert_builder("Users")
217 /// .set("UserId").to(&1)
218 /// .build();
219 /// let mutation1b = Mutation::new_insert_builder("UserRoles")
220 /// .set("UserId").to(&1)
221 /// .set("Role").to(&"Admin")
222 /// .build();
223 /// let group1 = MutationGroup::new(vec![mutation1a, mutation1b]);
224 ///
225 /// let mutation2 = Mutation::new_insert_builder("Users")
226 /// .set("UserId").to(&2)
227 /// .build();
228 /// let group2 = MutationGroup::new(vec![mutation2]);
229 ///
230 /// let transaction = db.batch_write_transaction().build();
231 /// let mut stream = transaction.execute_streaming(vec![group1, group2]).await?;
232 ///
233 /// while let Some(response) = stream.next().await {
234 /// let response = response?;
235 /// if let Some(status) = response.status.as_ref().filter(|s| s.code != Code::Ok as i32) {
236 /// eprintln!("Error applying groups {:?}: {}", response.indexes, status.message);
237 /// } else {
238 /// println!("Applied groups: {:?}", response.indexes);
239 /// }
240 /// }
241 /// # Ok(())
242 /// # }
243 /// ```
244 ///
245 /// A batch write transaction is used to execute non-atomic writes using mutations.
246 /// Related mutations should be placed in a group. For example, two mutations inserting
247 /// rows with the same primary key prefix in both parent and child tables are related.
248 /// All mutations within a group are applied atomically, but the entire batch is not
249 /// guaranteed to be atomic.
250 pub fn batch_write_transaction(&self) -> BatchWriteTransactionBuilder {
251 BatchWriteTransactionBuilder::new(self.clone())
252 }
253
254 pub(crate) fn session_name(&self) -> String {
255 self.session_maintainer.session_name()
256 }
257}
258
259/// A builder for [DatabaseClient].
260pub struct DatabaseClientBuilder {
261 spanner: Spanner,
262 database_name: String,
263 database_role: Option<String>,
264 options: Option<crate::RequestOptions>,
265 leader_aware_routing_enabled: bool,
266}
267
268impl DatabaseClientBuilder {
269 pub(crate) fn new(spanner: Spanner, database_name: String) -> Self {
270 Self {
271 spanner,
272 database_name,
273 database_role: None,
274 options: None,
275 leader_aware_routing_enabled: true,
276 }
277 }
278
279 /// Sets the database role for the client.
280 ///
281 /// # Example
282 /// ```
283 /// # use google_cloud_spanner::client::Spanner;
284 /// # async fn sample() -> anyhow::Result<()> {
285 /// let spanner = Spanner::builder().build().await?;
286 /// let database_client = spanner
287 /// .database_client("projects/my-project/instances/my-instance/databases/my-db")
288 /// .with_database_role("my-role")
289 /// .build()
290 /// .await?;
291 /// # Ok(())
292 /// # }
293 /// ```
294 ///
295 /// Database roles are used for Fine-Grained Access Control (FGAC).
296 /// You can assign a database role to a session, and that role determines the permissions for that session.
297 /// For more information, see [Access with FGAC](https://docs.cloud.google.com/spanner/docs/access-with-fgac).
298 pub fn with_database_role(mut self, role: impl Into<String>) -> Self {
299 self.database_role = Some(role.into());
300 self
301 }
302
303 /// Sets the request options that will be used when creating the multiplexed
304 /// session for the client.
305 ///
306 /// # Example
307 /// ```
308 /// # use google_cloud_spanner::client::Spanner;
309 /// # use google_cloud_gax::options::RequestOptions;
310 /// # use std::time::Duration;
311 /// # async fn sample() -> anyhow::Result<()> {
312 /// let spanner = Spanner::builder().build().await?;
313 /// let mut options = RequestOptions::default();
314 /// options.set_attempt_timeout(Duration::from_secs(60));
315 /// let database_client = spanner
316 /// .database_client("projects/my-project/instances/my-instance/databases/my-db")
317 /// .with_request_options(options)
318 /// .build()
319 /// .await?;
320 /// # Ok(())
321 /// # }
322 /// ```
323 pub fn with_request_options(mut self, options: crate::RequestOptions) -> Self {
324 self.options = Some(options);
325 self
326 }
327
328 /// Sets whether Leader-Aware Routing (LAR) is enabled for read/write transactions.
329 ///
330 /// # Example
331 /// ```
332 /// # use google_cloud_spanner::client::Spanner;
333 /// # async fn sample() -> anyhow::Result<()> {
334 /// let spanner = Spanner::builder().build().await?;
335 /// let database_client = spanner
336 /// .database_client("projects/my-project/instances/my-instance/databases/my-db")
337 /// .with_leader_aware_routing(true)
338 /// .build()
339 /// .await?;
340 /// # Ok(())
341 /// # }
342 /// ```
343 ///
344 /// When LAR is enabled, modifying operations (Read-Write, Write-Only, and Partitioned DML
345 /// transactions) automatically route requests directly to the Spanner leader replica. This
346 /// eliminates internal forwarding hops between replicas and reduces overall transaction latency.
347 ///
348 /// Enabled by default.
349 ///
350 /// See also: <https://docs.cloud.google.com/spanner/docs/leader-aware-routing>
351 pub fn with_leader_aware_routing(mut self, enabled: bool) -> Self {
352 self.leader_aware_routing_enabled = enabled;
353 self
354 }
355
356 /// Builds the [DatabaseClient] and creates a single multiplexed session that
357 /// will be used for all operations on the database.
358 pub async fn build(self) -> crate::Result<DatabaseClient> {
359 let spanner_clone = self.spanner.clone();
360 let session_maintainer = ManagedSessionMaintainer::create_and_start_maintenance(
361 self.spanner,
362 self.database_name,
363 self.database_role.unwrap_or_default(),
364 self.options.unwrap_or_default(),
365 )
366 .await?;
367
368 Ok(DatabaseClient {
369 spanner: spanner_clone,
370 session_maintainer,
371 leader_aware_routing_enabled: self.leader_aware_routing_enabled,
372 })
373 }
374}
375
376#[cfg(test)]
377mod tests {
378 use super::*;
379 use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
380 use google_cloud_test_macros::tokio_test_no_panics;
381 use spanner_grpc_mock::{MockSpanner, start};
382
383 #[test]
384 fn test_auto_traits() {
385 use static_assertions::assert_impl_all;
386 assert_impl_all!(DatabaseClient: Send, Sync, Clone, std::fmt::Debug);
387 }
388
389 #[tokio_test_no_panics]
390 async fn test_database_client_builder() {
391 let mut mock = MockSpanner::new();
392 mock.expect_create_session().once().returning(|req| {
393 let req = req.into_inner();
394 let session = req.session.unwrap();
395 assert!(session.multiplexed);
396 assert_eq!(session.creator_role, "test-role");
397
398 Ok(gaxi::grpc::tonic::Response::new(
399 spanner_grpc_mock::google::spanner::v1::Session {
400 name: "projects/test-project/instances/test-instance/databases/test-db/sessions/123".to_string(),
401 multiplexed: true,
402 creator_role: "test-role".to_string(),
403 ..Default::default()
404 },
405 ))
406 });
407
408 let (address, _server) = start("0.0.0.0:0", mock)
409 .await
410 .expect("Failed to start mock server");
411 let spanner = Spanner::builder()
412 .with_endpoint(address)
413 .with_credentials(Anonymous::new().build())
414 .build()
415 .await
416 .expect("Failed to build client");
417
418 let db_client = spanner
419 .database_client("projects/test-project/instances/test-instance/databases/test-db")
420 .with_database_role("test-role")
421 .build()
422 .await
423 .expect("Failed to create DatabaseClient");
424
425 let session = db_client
426 .session_maintainer
427 .session
428 .read()
429 .expect("failed to read session")
430 .session
431 .clone();
432 assert_eq!(
433 session.name,
434 "projects/test-project/instances/test-instance/databases/test-db/sessions/123"
435 );
436 assert!(session.multiplexed);
437 assert_eq!(session.creator_role, "test-role");
438 }
439
440 #[tokio_test_no_panics]
441 async fn test_database_client_builder_with_options() {
442 let mut mock = MockSpanner::new();
443 let mut seq = mockall::Sequence::new();
444 mock.expect_create_session()
445 .once()
446 .in_sequence(&mut seq)
447 .returning(|_| Err(gaxi::grpc::tonic::Status::unavailable("unavailable")));
448 mock.expect_create_session()
449 .once()
450 .in_sequence(&mut seq)
451 .returning(|req| {
452 let req = req.into_inner();
453 let session = req.session.unwrap();
454 assert!(session.multiplexed);
455 Ok(gaxi::grpc::tonic::Response::new(
456 spanner_grpc_mock::google::spanner::v1::Session {
457 name: "projects/test-project/instances/test-instance/databases/test-db/sessions/123".to_string(),
458 multiplexed: true,
459 ..Default::default()
460 },
461 ))
462 });
463
464 let (address, _server) = start("0.0.0.0:0", mock)
465 .await
466 .expect("Failed to start mock server");
467 let spanner = Spanner::builder()
468 .with_endpoint(address)
469 .with_credentials(Anonymous::new().build())
470 .build()
471 .await
472 .expect("Failed to build client");
473
474 let mut options = crate::RequestOptions::default();
475 options.set_retry_policy(google_cloud_gax::retry_policy::Aip194Strict);
476 options.set_idempotency(true);
477
478 let db_client = spanner
479 .database_client("projects/test-project/instances/test-instance/databases/test-db")
480 .with_request_options(options)
481 .build()
482 .await
483 .expect("Failed to create DatabaseClient");
484
485 let session = db_client
486 .session_maintainer
487 .session
488 .read()
489 .expect("failed to read session")
490 .session
491 .clone();
492 assert_eq!(
493 session.name,
494 "projects/test-project/instances/test-instance/databases/test-db/sessions/123"
495 );
496 }
497
498 #[tokio_test_no_panics]
499 async fn test_database_client_builder_error() {
500 let mut mock = MockSpanner::new();
501 mock.expect_create_session().once().returning(|_| {
502 Err(gaxi::grpc::tonic::Status::permission_denied(
503 "permission denied",
504 ))
505 });
506
507 let (address, _server) = start("0.0.0.0:0", mock)
508 .await
509 .expect("Failed to start mock server");
510 let spanner = Spanner::builder()
511 .with_endpoint(address)
512 .with_credentials(Anonymous::new().build())
513 .build()
514 .await
515 .expect("Failed to build client");
516
517 let result = spanner
518 .database_client("projects/test-project/instances/test-instance/databases/test-db")
519 .build()
520 .await;
521
522 match result {
523 Ok(_) => panic!("Client creation should have failed"),
524 Err(e) => assert_eq!(
525 e.status().map(|s| s.code),
526 Some(google_cloud_gax::error::rpc::Code::PermissionDenied)
527 ),
528 }
529 }
530}