google-cloud-spanner 0.34.2-preview

// 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 google_cloud_gax::error::Error;
use google_cloud_gax::retry_policy::{Aip194Strict, RetryPolicy};
use google_cloud_gax::retry_result::RetryResult;
use google_cloud_gax::retry_state::RetryState;
use google_cloud_gax::throttle_result::ThrottleResult;
use std::time::Duration;

/// A custom retry policy for Cloud Spanner that decorates/extends `Aip194Strict`.
///
/// Spanner allows transport/connection errors to be retried on idempotent operations.
/// This policy explicitly allows retries on these transport/network errors if the request is idempotent.
#[derive(Clone, Debug)]
pub(crate) struct SpannerRetryPolicy {
    inner: Aip194Strict,
}

impl SpannerRetryPolicy {
    pub(crate) fn new() -> Self {
        Self {
            inner: Aip194Strict,
        }
    }
}

impl Default for SpannerRetryPolicy {
    fn default() -> Self {
        Self::new()
    }
}

impl RetryPolicy for SpannerRetryPolicy {
    fn on_error(&self, state: &RetryState, error: Error) -> RetryResult {
        // 1. Strict AIP-194 checks (Unavailable, is_transient_and_before_rpc)
        let result = self.inner.on_error(state, error);
        match result {
            // If the strict AIP-194 checks classified the error as permanent (such as a transport
            // error that occurred post-headers), we override it to Continue if the request is idempotent.
            RetryResult::Permanent(error)
                if state.idempotent && (error.is_transport() || error.is_io()) =>
            {
                RetryResult::Continue(error)
            }
            res => res,
        }
    }

    fn on_throttle(&self, state: &RetryState, error: Error) -> ThrottleResult {
        self.inner.on_throttle(state, error)
    }

    fn remaining_time(&self, state: &RetryState) -> Option<Duration> {
        self.inner.remaining_time(state)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use google_cloud_gax::error::Error as GaxError;
    use google_cloud_gax::error::rpc::{Code, Status};
    use http::HeaderMap;

    #[test]
    fn test_spanner_retry_policy_idempotent() {
        let policy = SpannerRetryPolicy::new();
        let state = RetryState::new(true); // idempotent = true

        // 1. Service UNAVAILABLE error should be retried (via inner AIP-194)
        let status = Status::default()
            .set_code(Code::Unavailable)
            .set_message("Service Unavailable");
        let err = GaxError::service(status);
        assert!(
            policy.on_error(&state, err).is_continue(),
            "Expected UNAVAILABLE to be retried when idempotent"
        );

        // 2. Service PERMISSION_DENIED error should not be retried
        let status = Status::default()
            .set_code(Code::PermissionDenied)
            .set_message("Denied");
        let err = GaxError::service(status);
        assert!(
            policy.on_error(&state, err).is_permanent(),
            "Expected PERMISSION_DENIED to not be retried"
        );

        // 3. IO/Transport error should be retried when idempotent
        let err = GaxError::transport(
            HeaderMap::new(),
            std::io::Error::new(std::io::ErrorKind::ConnectionReset, "connection closed"),
        );
        assert!(
            policy.on_error(&state, err).is_continue(),
            "Expected transport connection reset to be retried when idempotent"
        );
    }

    #[test]
    fn test_spanner_retry_policy_non_idempotent() {
        let policy = SpannerRetryPolicy::new();
        let state = RetryState::new(false); // idempotent = false

        // 1. Service UNAVAILABLE error should NOT be retried (AIP-194 requires idempotency)
        let status = Status::default()
            .set_code(Code::Unavailable)
            .set_message("Service Unavailable");
        let err = GaxError::service(status);
        assert!(
            policy.on_error(&state, err).is_permanent(),
            "Expected UNAVAILABLE to be permanent when non-idempotent"
        );

        // 2. IO/Transport error should NOT be retried when non-idempotent
        let err = GaxError::transport(
            HeaderMap::new(),
            std::io::Error::new(std::io::ErrorKind::ConnectionReset, "connection closed"),
        );
        assert!(
            policy.on_error(&state, err).is_permanent(),
            "Expected transport connection reset to not be retried when non-idempotent"
        );
    }
}