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