Skip to main content

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}