From 3277faf72266d9a149291ba4e6fef3c68f4634c2 Mon Sep 17 00:00:00 2001
From: Joshua Tan <joshuatan@google.com>
Date: Wed, 8 Jul 2026 03:33:27 +0000
Subject: [PATCH] feat(storage): make GrpcRustClient available as a transport
 option

---
 src/storage/src/storage/bidi.rs       | 13 ++++++++++---
 src/storage/src/storage/bidi_write.rs | 18 ++++++++++++++----
 2 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/src/storage/src/storage/bidi.rs b/src/storage/src/storage/bidi.rs
index b342e58a31..d3b21964b3 100644
--- a/src/storage/src/storage/bidi.rs
+++ b/src/storage/src/storage/bidi.rs
@@ -27,11 +27,18 @@ mod worker;
 
 use crate::google::storage::v2::{BidiReadObjectRequest, BidiReadObjectResponse};
 use crate::request_options::RequestOptions;
-use gaxi::grpc::tonic::{Extensions, Response as TonicResponse, Result as TonicResult, Streaming};
+use gaxi::grpc::tonic::{Extensions, Response as TonicResponse, Result as TonicResult};
 use tokio::sync::mpsc::Receiver;
 
+#[cfg(google_cloud_unstable_grpc_rust)]
+pub(crate) type GrpcClient = gaxi::grpc::GrpcRustClient;
+#[cfg(google_cloud_unstable_grpc_rust)]
+pub(crate) type GrpcStream = gaxi::grpc::GrpcRustStreaming<BidiReadObjectResponse>;
+
+#[cfg(not(google_cloud_unstable_grpc_rust))]
 pub(crate) type GrpcClient = gaxi::grpc::Client;
-pub(crate) type GrpcStream = Streaming<BidiReadObjectResponse>;
+#[cfg(not(google_cloud_unstable_grpc_rust))]
+pub(crate) type GrpcStream = gaxi::grpc::tonic::Streaming<BidiReadObjectResponse>;
 
 pub use super::open_object::OpenObject;
 
@@ -45,7 +52,7 @@ pub trait TonicStreaming: std::fmt::Debug + Send + 'static {
 }
 
 /// Implement [TonicStreaming] for the one `Streaming<T>` we use.
-impl TonicStreaming for Streaming<BidiReadObjectResponse> {
+impl TonicStreaming for GrpcStream {
     async fn next_message(&mut self) -> TonicResult<Option<BidiReadObjectResponse>> {
         self.message().await
     }
diff --git a/src/storage/src/storage/bidi_write.rs b/src/storage/src/storage/bidi_write.rs
index b6576d33be..c3280d3313 100644
--- a/src/storage/src/storage/bidi_write.rs
+++ b/src/storage/src/storage/bidi_write.rs
@@ -20,10 +20,20 @@ pub(crate) mod stub;
 
 use crate::google::storage::v2::{BidiWriteObjectRequest, BidiWriteObjectResponse};
 use crate::request_options::RequestOptions;
-use gaxi::grpc::tonic::{Extensions, Response as TonicResponse, Result as TonicResult, Streaming};
+use gaxi::grpc::tonic::{Extensions, Response as TonicResponse, Result as TonicResult};
 use std::future::Future;
 use tokio::sync::mpsc::Receiver;
 
+#[cfg(google_cloud_unstable_grpc_rust)]
+pub(crate) type GrpcClient = gaxi::grpc::GrpcRustClient;
+#[cfg(google_cloud_unstable_grpc_rust)]
+pub(crate) type GrpcStream = gaxi::grpc::GrpcRustStreaming<BidiWriteObjectResponse>;
+
+#[cfg(not(google_cloud_unstable_grpc_rust))]
+pub(crate) type GrpcClient = gaxi::grpc::Client;
+#[cfg(not(google_cloud_unstable_grpc_rust))]
+pub(crate) type GrpcStream = gaxi::grpc::tonic::Streaming<BidiWriteObjectResponse>;
+
 /// A trait to mock `Streaming<T>` in the unit tests.
 ///
 /// This is not a public trait, we only need this for our own testing.
@@ -35,7 +45,7 @@ pub(crate) trait TonicStreaming: std::fmt::Debug + Send + 'static {
 }
 
 /// Implement [TonicStreaming] for the one `Streaming<T>`` we use.
-impl TonicStreaming for Streaming<BidiWriteObjectResponse> {
+impl TonicStreaming for GrpcStream {
     async fn next_message(&mut self) -> TonicResult<Option<BidiWriteObjectResponse>> {
         self.message().await
     }
@@ -58,8 +68,8 @@ pub(crate) trait Client: std::fmt::Debug + Send + 'static {
     ) -> impl Future<Output = crate::Result<TonicResult<TonicResponse<Self::Stream>>>> + Send;
 }
 
-impl Client for gaxi::grpc::Client {
-    type Stream = Streaming<BidiWriteObjectResponse>;
+impl Client for GrpcClient {
+    type Stream = GrpcStream;
     async fn start(
         &self,
         extensions: Extensions,