code refactoring

Author: prasunna09

crates/router/src/core/proxy.rs

@@ -16,21 +16,17 @@ pub async fn proxy_core(
     req: proxy_api_models::ProxyRequest,
 ) -> RouterResponse<proxy_api_models::ProxyResponse> {
     let req_wrapper = utils::ProxyRequestWrapper(req.clone());
-    let vault_id = req_wrapper
-        .get_vault_id(
+    let proxy_record = req_wrapper
+        .get_payment_method_or_tokenization_record(
             &state,
             merchant_context.get_merchant_key_store(),
             merchant_context.get_merchant_account().storage_scheme,
         )
         .await?;
 
-    let customer_id = req_wrapper
-        .get_customer_id(
-            &state,
-            merchant_context.get_merchant_key_store(),
-            merchant_context.get_merchant_account().storage_scheme,
-        )
-        .await?;
+    let vault_id = proxy_record.get_vault_id()?;
+    let customer_id = proxy_record.get_customer_id()?;
+
     let vault_response =
         super::payment_methods::vault::retrieve_payment_method_from_vault_internal(
             &state,

crates/router/src/core/proxy/utils.rs

@@ -19,14 +19,18 @@ use crate::{
 };
 
 pub struct ProxyRequestWrapper(pub proxy_api_models::ProxyRequest);
+pub enum ProxyRecord {
+    PaymentMethodRecord(domain::PaymentMethod),
+    TokenizationRecord(domain::Tokenization),
+}
 
 impl ProxyRequestWrapper {
-    pub async fn get_vault_id(
+    pub async fn get_payment_method_or_tokenization_record(
         &self,
         state: &SessionState,
         key_store: &domain::MerchantKeyStore,
         storage_scheme: common_enums::enums::MerchantStorageScheme,
-    ) -> RouterResult<payment_methods::VaultId> {
+    ) -> RouterResult<ProxyRecord> {
         let token = &self.0.token;
 
         match self.0.token_type {
@@ -39,15 +43,13 @@ impl ProxyRequestWrapper {
                         .change_context(errors::ApiErrorResponse::InternalServerError)
                         .attach_printable("Unable to generate GlobalPaymentMethodId")?;
 
-                state
+                let payment_method_record = state
                     .store
                     .find_payment_method(&((state).into()), key_store, &pm_id, storage_scheme)
                     .await
-                    .change_context(errors::ApiErrorResponse::PaymentMethodNotFound)?
-                    .locker_id
-                    .get_required_value("vault_id")
-                    .change_context(errors::ApiErrorResponse::InternalServerError)
-                    .attach_printable("Locker id not present in Payment Method Entry")
+                    .change_context(errors::ApiErrorResponse::PaymentMethodNotFound)?;
+                Ok(ProxyRecord::PaymentMethodRecord(payment_method_record))
+
             }
             proxy_api_models::TokenType::TokenizationId => {
                 Err(report!(errors::ApiErrorResponse::NotImplemented {
@@ -112,6 +114,47 @@ impl ProxyRequestWrapper {
     }
 }
 
+impl ProxyRecord{
+    pub fn get_vault_id(
+        &self,
+    ) -> RouterResult<payment_methods::VaultId>{
+
+        match self {
+            Self::PaymentMethodRecord(payment_method) => {
+                payment_method.locker_id.clone()
+                .get_required_value("vault_id")
+                .change_context(errors::ApiErrorResponse::InternalServerError)
+                .attach_printable("Locker id not present in Payment Method Entry")
+
+            }
+            Self::TokenizationRecord(_) => {
+                Err(report!(errors::ApiErrorResponse::NotImplemented {
+                    message: NotImplementedMessage::Reason(
+                        "Proxy flow using tokenization id".to_string(),
+                    ),
+                }))
+            }
+        }
+    }
+
+    pub fn get_customer_id(
+        &self,
+    ) -> RouterResult<id_type::GlobalCustomerId> {
+        match self {
+            Self::PaymentMethodRecord(payment_method) => {
+                Ok(payment_method.customer_id.clone())
+            }
+            Self::TokenizationRecord(_) => {
+                Err(report!(errors::ApiErrorResponse::NotImplemented {
+                    message: NotImplementedMessage::Reason(
+                        "Proxy flow using tokenization id".to_string(),
+                    ),
+                }))
+            }
+        }
+    }
+}
+
 #[derive(Debug)]
 pub struct TokenReference {
     pub field: String,

crates/router/src/types/domain.rs

@@ -56,6 +56,12 @@ pub mod vault {
     pub use hyperswitch_domain_models::vault::*;
 }
 
+#[cfg(feature = "v2")]
+pub mod tokenization {
+    pub use hyperswitch_domain_models::tokenization::*;
+}
+
+
 mod routing {
     pub use hyperswitch_domain_models::routing::*;
 }
@@ -83,3 +89,6 @@ pub use user::*;
 pub use user_key_store::*;
 #[cfg(feature = "v2")]
 pub use vault::*;
+
+#[cfg(feature = "v2")]
+pub use tokenization::*;