From e2f20935627b7391b66472e24ef940f58eb41581 Mon Sep 17 00:00:00 2001
From: Chethan <chethan.rao@juspay.in>
Date: Fri, 20 Dec 2024 13:47:50 +0530
Subject: [PATCH] add a trait to include headers into tonic request

---
 crates/external_services/src/grpc_client.rs   | 22 ++++++++++++-------
 .../elimination_rate_client.rs                |  8 +++----
 .../dynamic_routing/success_rate_client.rs    |  8 +++----
 3 files changed, 22 insertions(+), 16 deletions(-)

diff --git a/crates/external_services/src/grpc_client.rs b/crates/external_services/src/grpc_client.rs
index cb822d22f9b5..404685025ed3 100644
--- 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
index 908faf362623..bc5ce4997270 100644
--- 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
index 7d5a57262324..3cf06ab63beb 100644
--- 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()