google_cloud_bigquery/write/

client_builder.rs

// Copyright 2026 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use super::pool::StreamPoolOptions;
use super::retry_policy::RetryOptions;
use crate::ClientBuilderResult as BuilderResult;
use crate::client::Write;
use gaxi::options::ClientConfig;
use google_cloud_auth::credentials::Credentials;
use google_cloud_gax::backoff_policy::BackoffPolicyArg;
use google_cloud_gax::retry_policy::RetryPolicyArg;
use std::time::Duration;

/// A builder for [Write].
///
/// # Example
/// ```
/// # use google_cloud_bigquery::client::Write;
/// # async fn sample() -> anyhow::Result<()> {
/// let builder = Write::builder();
/// let client = builder
///     .with_endpoint("https://bigquerystorage.googleapis.com")
///     .build()
///     .await?;
/// # Ok(()) }
/// ```
#[derive(Debug)]
pub struct ClientBuilder {
    pub(super) config: ClientConfig,
    pub(super) pool_options: StreamPoolOptions,
    pub(super) retry_options: RetryOptions,
}

impl ClientBuilder {
    pub(super) fn new() -> Self {
        Self {
            config: ClientConfig::default(),
            pool_options: StreamPoolOptions::default(),
            retry_options: RetryOptions::default(),
        }
    }

    /// Creates a new client.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let client = Write::builder().build().await?;
    /// # Ok(()) }
    /// ```
    pub async fn build(self) -> BuilderResult<Write> {
        Write::new(self).await
    }

    /// Sets the endpoint.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let client = Write::builder()
    ///     .with_endpoint("https://private.googleapis.com")
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    pub fn with_endpoint<V: Into<String>>(mut self, v: V) -> Self {
        self.config.endpoint = Some(v.into());
        self
    }

    /// Configure the universe domain.
    ///
    /// The universe domain is the default service domain for a given cloud universe.
    /// The default value is "googleapis.com".
    ///
    /// # Example
    /// ```
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let client = Write::builder()
    ///     .with_universe_domain("googleapis.com")
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    pub fn with_universe_domain<V: Into<String>>(mut self, v: V) -> Self {
        self.config.universe_domain = Some(v.into());
        self
    }

    /// Configures the authentication credentials.
    ///
    /// More information about valid credentials types can be found in the
    /// [google-cloud-auth] crate documentation.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use google_cloud_auth::credentials::mds;
    /// let client = Write::builder()
    ///     .with_credentials(
    ///         mds::Builder::default()
    ///             .with_scopes(["https://www.googleapis.com/auth/bigquery"])
    ///             .build()?)
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    ///
    /// [google-cloud-auth]: https://docs.rs/google-cloud-auth
    pub fn with_credentials<V: Into<Credentials>>(mut self, v: V) -> Self {
        self.config.cred = Some(v.into());
        self
    }

    /// Configure the number of subchannels used by the client.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let count = std::thread::available_parallelism()?.get();
    /// let client = Write::builder()
    ///     .with_grpc_subchannel_count(count)
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    ///
    /// gRPC-based clients may exhibit high latency if many requests need to be
    /// demuxed over a single HTTP/2 connection (often called a *subchannel* in
    /// gRPC).
    ///
    /// Consider using more subchannels if your application creates many
    /// writers. Consider using fewer subchannels if your application needs the
    /// file descriptors for other purposes.
    pub fn with_grpc_subchannel_count(mut self, v: usize) -> Self {
        self.config.grpc_subchannel_count = Some(v);
        self
    }

    /// Configure the maximum streams in the client's multiplexed stream pool.
    ///
    /// # Example
    /// ```
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let count = std::thread::available_parallelism()?.get();
    /// let client = Write::builder()
    ///     .with_pool_size_limit(count)
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    ///
    /// This stream pool is shared by default writers with multiplexing enabled.
    ///
    /// The client scales the stream pool up to this limit as the streams in the
    /// pool encounter load.
    ///
    /// Note that until background stream rebalancing ([#6866]) is implemented,
    /// stream assignment and pool scale-up only occur when a writer is built.
    /// If multiple writers are created all at once before writes are in flight,
    /// they will all share the same initial stream. To scale across multiple
    /// streams, stagger writer creation so writes are in flight when new
    /// writers are built.
    ///
    /// The default is 8 streams.
    ///
    /// [#6866]: https://github.com/googleapis/google-cloud-rust/issues/6866
    pub fn with_pool_size_limit(mut self, v: usize) -> Self {
        self.pool_options.max_streams = v.max(1);
        self
    }

    // TODO(#6866) - expose when we have rebalancing
    #[cfg_attr(not(test), expect(dead_code))]
    /// Configure the maximum outstanding requests per stream in the client's
    /// multiplexed stream pool.
    ///
    /// # Example
    /// ```no_rust
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let client = Write::builder()
    ///     .with_max_outstanding_requests(200)
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    ///
    /// As streams in the stream pool approach this limit, the client
    /// dynamically adds more streams to the stream pool, up to the limit
    /// configured by `with_pool_size_limit`.
    ///
    /// Note that this is not a hard maximum. When the pool reaches the stream
    /// limit, the streams will continue accepting requests. Consider using a
    /// semaphore locally for flow control.
    ///
    /// The default is 1 request.
    pub(crate) fn with_max_outstanding_requests(mut self, v: u64) -> Self {
        self.pool_options.max_outstanding_requests = Some(v.max(1));
        self
    }

    // TODO(#6866) - expose when we have rebalancing
    #[cfg_attr(not(test), expect(dead_code))]
    /// Configure the maximum outstanding bytes per stream in the client's
    /// multiplexed stream pool.
    ///
    /// # Example
    /// ```no_rust
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// let client = Write::builder()
    ///     .with_max_outstanding_bytes(200_000)
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    ///
    /// As streams in the stream pool approach this limit, the client
    /// dynamically adds more streams to the stream pool, up to the limit
    /// configured by `with_pool_size_limit`.
    ///
    /// Note that this is not a hard maximum. When the pool reaches the stream
    /// limit, the streams will continue accepting requests. Consider using a
    /// semaphore locally for flow control.
    pub(crate) fn with_max_outstanding_bytes(mut self, v: u64) -> Self {
        self.pool_options.max_outstanding_bytes = Some(v.max(1));
        self
    }

    // TODO(#6851) - release when all streams support retries.
    #[cfg_attr(not(test), expect(dead_code))]
    /// Configure the retry policy.
    ///
    /// The client libraries can automatically retry operations that fail. The
    /// retry policy controls what errors are considered retryable, sets limits
    /// on the number of attempts or the time trying to make attempts.
    ///
    /// # Example
    /// ```no_rust
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use google_cloud_bigquery::write::retry_policy::RetryableErrors;
    /// use google_cloud_gax::retry_policy::RetryPolicyExt;
    /// let client = Write::builder()
    ///     .with_retry_policy(RetryableErrors.with_attempt_limit(3))
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    pub(crate) fn with_retry_policy<V: Into<RetryPolicyArg>>(mut self, v: V) -> Self {
        self.retry_options.retry_policy = v.into().into();
        self
    }

    // TODO(#6851) - release when all streams support retries.
    #[cfg_attr(not(test), expect(dead_code))]
    /// Configure the retry backoff policy.
    ///
    /// The client libraries can automatically retry operations that fail. The
    /// backoff policy controls how long to wait in between retry attempts.
    ///
    /// # Example
    /// ```no_rust
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use google_cloud_gax::exponential_backoff::ExponentialBackoff;
    /// let policy = ExponentialBackoff::default();
    /// let client = Write::builder()
    ///     .with_backoff_policy(policy)
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    pub(crate) fn with_backoff_policy<V: Into<BackoffPolicyArg>>(mut self, v: V) -> Self {
        self.retry_options.backoff_policy = v.into().into();
        self
    }

    // TODO(#6851) - release when all streams support retries.
    #[cfg_attr(not(test), expect(dead_code))]
    /// Configure the timeout for a single write attempt.
    ///
    /// Without this limit, a write can block forever if the service accepts
    /// the stream but never responds. On a timeout, the client abandons the
    /// stream and the retry policy decides whether to make another attempt.
    ///
    /// # Example
    /// ```no_rust
    /// # use google_cloud_bigquery::client::Write;
    /// # async fn sample() -> anyhow::Result<()> {
    /// use std::time::Duration;
    /// let client = Write::builder()
    ///     .with_attempt_timeout(Duration::from_secs(10))
    ///     .build()
    ///     .await?;
    /// # Ok(()) }
    /// ```
    pub(crate) fn with_attempt_timeout(mut self, v: Duration) -> Self {
        self.retry_options.attempt_timeout = Some(v);
        self
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::write::test::NoBackoff;
    use google_cloud_auth::credentials::anonymous::Builder as Anonymous;
    use google_cloud_gax::retry_policy::NeverRetry;

    #[test]
    fn defaults() {
        let builder = ClientBuilder::new();
        assert!(builder.config.endpoint.is_none(), "{:?}", builder.config);
        assert!(builder.config.cred.is_none(), "{:?}", builder.config);
        assert!(
            builder.config.universe_domain.is_none(),
            "{:?}",
            builder.config
        );
        assert!(
            builder.config.grpc_subchannel_count.is_none(),
            "{:?}",
            builder.config
        );
        assert_eq!(builder.pool_options.max_streams, 8);
        assert_eq!(builder.pool_options.max_outstanding_requests, Some(1));
        assert_eq!(builder.pool_options.max_outstanding_bytes, None);
    }

    #[test]
    fn setters() {
        let builder = ClientBuilder::new()
            .with_endpoint("test-endpoint.com")
            .with_universe_domain("test-ud.com")
            .with_credentials(Anonymous::new().build())
            .with_grpc_subchannel_count(16)
            .with_pool_size_limit(10)
            .with_max_outstanding_requests(900)
            .with_max_outstanding_bytes(1_000_000)
            .with_retry_policy(NeverRetry)
            .with_backoff_policy(NoBackoff)
            .with_attempt_timeout(Duration::from_secs(10));
        assert_eq!(
            builder.config.endpoint,
            Some("test-endpoint.com".to_string())
        );
        assert_eq!(
            builder.config.universe_domain,
            Some("test-ud.com".to_string())
        );
        assert!(builder.config.cred.is_some(), "{:?}", builder.config);
        assert_eq!(builder.config.grpc_subchannel_count, Some(16));
        assert_eq!(builder.pool_options.max_streams, 10);
        assert_eq!(builder.pool_options.max_outstanding_requests, Some(900));
        assert_eq!(builder.pool_options.max_outstanding_bytes, Some(1_000_000));
        assert_eq!(
            builder.retry_options.attempt_timeout,
            Some(Duration::from_secs(10))
        );

        let fmt = format!("{:?}", builder.retry_options);
        assert!(fmt.contains("NeverRetry"), "{fmt}");
        assert!(fmt.contains("NoBackoff"), "{fmt}");
    }

    #[test]
    fn validate_pool_options() {
        let builder = ClientBuilder::new()
            .with_credentials(Anonymous::new().build())
            .with_pool_size_limit(0)
            .with_max_outstanding_requests(0)
            .with_max_outstanding_bytes(0);
        assert_eq!(builder.pool_options.max_streams, 1);
        assert_eq!(builder.pool_options.max_outstanding_requests, Some(1));
        assert_eq!(builder.pool_options.max_outstanding_bytes, Some(1));
    }
}