add a trait to include headers into tonic request

juspay · Dec 20, 2024 · e2f2093 · e2f2093
1 parent 4473608
commit e2f2093
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 16 deletions.
diff --git a/crates/external_services/src/grpc_client.rs b/crates/external_services/src/grpc_client.rs
@@ -90,15 +90,21 @@ pub struct GrpcHeaders {
     pub request_id: Option<String>,
 }
 
-impl GrpcHeaders {
-    #[cfg(feature = "dynamic_routing")]
-    /// Method to add necessary headers to the tonic Request
+#[cfg(feature = "dynamic_routing")]
+/// Trait to add necessary headers to the tonic Request
+pub(crate) trait AddHeaders {
+    /// Add necessary header fields to the tonic Request
+    fn add_headers_to_grpc_request(&mut self, headers: GrpcHeaders);
+}
+
+#[cfg(feature = "dynamic_routing")]
+impl<T> AddHeaders for tonic::Request<T> {
     #[track_caller]
-    pub fn add_headers_to_grpc_request<T>(&self, request: &mut tonic::Request<T>) {
-        self.tenant_id
+    fn add_headers_to_grpc_request(&mut self, headers: GrpcHeaders) {
+        headers.tenant_id
             .parse()
             .map(|tenant_id| {
-                request
+                self
                     .metadata_mut()
                     .append(consts::TENANT_HEADER, tenant_id)
             })
@@ -107,11 +113,11 @@ impl GrpcHeaders {
             )
             .ok();
 
-        self.request_id.as_ref().map(|request_id| {
+        headers.request_id.map(|request_id| {
             request_id
                 .parse()
                 .map(|request_id| {
-                    request
+                    self
                         .metadata_mut()
                         .append(consts::X_REQUEST_ID, request_id)
                 })

diff --git a/crates/external_services/src/grpc_client/dynamic_routing/elimination_rate_client.rs b/crates/external_services/src/grpc_client/dynamic_routing/elimination_rate_client.rs
@@ -21,7 +21,7 @@ pub mod elimination_rate {
 }
 
 use super::{Client, DynamicRoutingError, DynamicRoutingResult};
-use crate::grpc_client::GrpcHeaders;
+use crate::grpc_client::{AddHeaders, GrpcHeaders};
 
 /// The trait Elimination Based Routing would have the functions required to support performance, calculation and invalidation bucket
 #[async_trait::async_trait]
@@ -76,7 +76,7 @@ impl EliminationBasedRouting for EliminationAnalyserClient<Client> {
             config,
         });
 
-        headers.add_headers_to_grpc_request(&mut request);
+        request.add_headers_to_grpc_request(headers);
 
         let response = self
             .clone()
@@ -117,7 +117,7 @@ impl EliminationBasedRouting for EliminationAnalyserClient<Client> {
             config,
         });
 
-        headers.add_headers_to_grpc_request(&mut request);
+        request.add_headers_to_grpc_request(headers);
 
         let response = self
             .clone()
@@ -136,7 +136,7 @@ impl EliminationBasedRouting for EliminationAnalyserClient<Client> {
     ) -> DynamicRoutingResult<InvalidateBucketResponse> {
         let mut request = tonic::Request::new(InvalidateBucketRequest { id });
 
-        headers.add_headers_to_grpc_request(&mut request);
+        request.add_headers_to_grpc_request(headers);
 
         let response = self
             .clone()

diff --git a/crates/external_services/src/grpc_client/dynamic_routing/success_rate_client.rs b/crates/external_services/src/grpc_client/dynamic_routing/success_rate_client.rs
@@ -21,7 +21,7 @@ pub mod success_rate {
     tonic::include_proto!("success_rate");
 }
 use super::{Client, DynamicRoutingError, DynamicRoutingResult};
-use crate::grpc_client::GrpcHeaders;
+use crate::grpc_client::{AddHeaders, GrpcHeaders};
 /// The trait Success Based Dynamic Routing would have the functions required to support the calculation and updation window
 #[async_trait::async_trait]
 pub trait SuccessBasedDynamicRouting: dyn_clone::DynClone + Send + Sync {
@@ -78,7 +78,7 @@ impl SuccessBasedDynamicRouting for SuccessRateCalculatorClient<Client> {
             config,
         });
 
-        headers.add_headers_to_grpc_request(&mut request);
+        request.add_headers_to_grpc_request(headers);
 
         let response = self
             .clone()
@@ -120,7 +120,7 @@ impl SuccessBasedDynamicRouting for SuccessRateCalculatorClient<Client> {
             config,
         });
 
-        headers.add_headers_to_grpc_request(&mut request);
+        request.add_headers_to_grpc_request(headers);
 
         let response = self
             .clone()
@@ -140,7 +140,7 @@ impl SuccessBasedDynamicRouting for SuccessRateCalculatorClient<Client> {
     ) -> DynamicRoutingResult<InvalidateWindowsResponse> {
         let mut request = tonic::Request::new(InvalidateWindowsRequest { id });
 
-        headers.add_headers_to_grpc_request(&mut request);
+        request.add_headers_to_grpc_request(headers);
 
         let response = self
             .clone()