fix: make push to drainer generic and add application metrics for KV (#2563)

crates/router/src/consts.rs

@@ -46,6 +46,3 @@ pub(crate) const QR_IMAGE_DATA_SOURCE_STRING: &str = "data:image/png;base64";
 pub(crate) const MERCHANT_ID_FIELD_EXTENSION_ID: &str = "1.2.840.113635.100.6.32";
 
 pub(crate) const METRICS_HOST_TAG_NAME: &str = "host";
-
-// TTL for KV setup
-pub(crate) const KV_TTL: u32 = 300;

crates/router/src/db.rs

@@ -24,17 +24,14 @@ pub mod payouts;
 pub mod refund;
 pub mod reverse_lookup;
 
-use std::fmt::Debug;
-
 use data_models::payments::{
     payment_attempt::PaymentAttemptInterface, payment_intent::PaymentIntentInterface,
 };
 use masking::PeekInterface;
 use redis_interface::errors::RedisError;
-use serde::de;
 use storage_impl::{redis::kv_store::RedisConnInterface, MockDb};
 
-use crate::{consts, errors::CustomResult, services::Store};
+use crate::{errors::CustomResult, services::Store};
 
 #[derive(PartialEq, Eq)]
 pub enum StorageImpl {
@@ -134,68 +131,4 @@ where
         .change_context(redis_interface::errors::RedisError::JsonDeserializationFailed)
 }
 
-pub enum KvOperation<'a, S: serde::Serialize + Debug> {
-    Hset((&'a str, String)),
-    SetNx(S),
-    HSetNx(&'a str, S),
-    Get(&'a str),
-    Scan(&'a str),
-}
-
-#[derive(router_derive::TryGetEnumVariant)]
-#[error(RedisError(UnknownResult))]
-pub enum KvResult<T: de::DeserializeOwned> {
-    Get(T),
-    Hset(()),
-    SetNx(redis_interface::SetnxReply),
-    HSetNx(redis_interface::HsetnxReply),
-    Scan(Vec<T>),
-}
-
-pub async fn kv_wrapper<'a, T, S>(
-    store: &Store,
-    op: KvOperation<'a, S>,
-    key: impl AsRef<str>,
-) -> CustomResult<KvResult<T>, RedisError>
-where
-    T: de::DeserializeOwned,
-    S: serde::Serialize + Debug,
-{
-    let redis_conn = store.get_redis_conn()?;
-
-    let key = key.as_ref();
-    let type_name = std::any::type_name::<T>();
-
-    match op {
-        KvOperation::Hset(value) => {
-            redis_conn
-                .set_hash_fields(key, value, Some(consts::KV_TTL))
-                .await?;
-            Ok(KvResult::Hset(()))
-        }
-        KvOperation::Get(field) => {
-            let result = redis_conn
-                .get_hash_field_and_deserialize(key, field, type_name)
-                .await?;
-            Ok(KvResult::Get(result))
-        }
-        KvOperation::Scan(pattern) => {
-            let result: Vec<T> = redis_conn.hscan_and_deserialize(key, pattern, None).await?;
-            Ok(KvResult::Scan(result))
-        }
-        KvOperation::HSetNx(field, value) => {
-            let result = redis_conn
-                .serialize_and_set_hash_field_if_not_exist(key, field, value, Some(consts::KV_TTL))
-                .await?;
-            Ok(KvResult::HSetNx(result))
-        }
-        KvOperation::SetNx(value) => {
-            let result = redis_conn
-                .serialize_and_set_key_if_not_exist(key, value, Some(consts::KV_TTL.into()))
-                .await?;
-            Ok(KvResult::SetNx(result))
-        }
-    }
-}
-
 dyn_clone::clone_trait_object!(StorageInterface);

crates/router/src/db/address.rs

@@ -280,7 +280,7 @@ mod storage {
     use error_stack::{IntoReport, ResultExt};
     use redis_interface::HsetnxReply;
     use router_env::{instrument, tracing};
-    use storage_impl::redis::kv_store::{kv_wrapper, KvOperation, PartitionKey};
+    use storage_impl::redis::kv_store::{kv_wrapper, KvOperation};
 
     use super::AddressInterface;
     use crate::{
@@ -419,15 +419,6 @@ mod storage {
                     let redis_value = serde_json::to_string(&updated_address)
                         .into_report()
                         .change_context(errors::StorageError::KVError)?;
-                    kv_wrapper::<(), _, _>(
-                        self,
-                        KvOperation::Hset::<storage_types::Address>((&field, redis_value)),
-                        &key,
-                    )
-                    .await
-                    .change_context(errors::StorageError::KVError)?
-                    .try_into_hset()
-                    .change_context(errors::StorageError::KVError)?;
 
                     let redis_entry = kv::TypedSql {
                         op: kv::DBOperation::Update {
@@ -440,15 +431,19 @@ mod storage {
                         },
                     };
 
-                    self.push_to_drainer_stream::<storage_types::Address>(
+                    kv_wrapper::<(), _, _>(
-                        redis_entry,
+                        self,
-                        PartitionKey::MerchantIdPaymentId {
+                        KvOperation::Hset::<storage_types::Address>(
-                            merchant_id: &updated_address.merchant_id,
+                            (&field, redis_value),
-                            payment_id: &payment_id,
+                            redis_entry,
-                        },
+                        ),
+                        &key,
                     )
                     .await
+                    .change_context(errors::StorageError::KVError)?
+                    .try_into_hset()
                     .change_context(errors::StorageError::KVError)?;
+
                     updated_address
                         .convert(key_store.key.get_inner())
                         .await
@@ -510,9 +505,19 @@ mod storage {
                         payment_id: address_new.payment_id.clone(),
                     };
 
+                    let redis_entry = kv::TypedSql {
+                        op: kv::DBOperation::Insert {
+                            insertable: kv::Insertable::Address(Box::new(address_new)),
+                        },
+                    };
+
                     match kv_wrapper::<diesel_models::Address, _, _>(
                         self,
-                        KvOperation::HSetNx(&field, &created_address),
+                        KvOperation::HSetNx::<diesel_models::Address>(
+                            &field,
+                            &created_address,
+                            redis_entry,
+                        ),
                         &key,
                     )
                     .await
@@ -521,30 +526,13 @@ mod storage {
                     {
                         Ok(HsetnxReply::KeyNotSet) => Err(errors::StorageError::DuplicateValue {
                             entity: "address",
-                            key: Some(address_new.address_id),
+                            key: Some(created_address.address_id),
                         })
                         .into_report(),
-                        Ok(HsetnxReply::KeySet) => {
+                        Ok(HsetnxReply::KeySet) => Ok(created_address
-                            let redis_entry = kv::TypedSql {
+                            .convert(key_store.key.get_inner())
-                                op: kv::DBOperation::Insert {
-                                    insertable: kv::Insertable::Address(Box::new(address_new)),
-                                },
-                            };
-                            self.push_to_drainer_stream::<diesel_models::Address>(
-                                redis_entry,
-                                PartitionKey::MerchantIdPaymentId {
-                                    merchant_id: &merchant_id,
-                                    payment_id,
-                                },
-                            )
                             .await
-                            .change_context(errors::StorageError::KVError)?;
+                            .change_context(errors::StorageError::DecryptionError)?),
-
-                            Ok(created_address
-                                .convert(key_store.key.get_inner())
-                                .await
-                                .change_context(errors::StorageError::DecryptionError)?)
-                        }
                         Err(er) => Err(er).change_context(errors::StorageError::KVError),
                     }
                 }

crates/router/src/db/connector_response.rs

@@ -100,7 +100,7 @@ mod storage {
     use error_stack::{IntoReport, ResultExt};
     use redis_interface::HsetnxReply;
     use router_env::{instrument, tracing};
-    use storage_impl::redis::kv_store::{kv_wrapper, KvOperation, PartitionKey};
+    use storage_impl::redis::kv_store::{kv_wrapper, KvOperation};
 
     use super::Store;
     use crate::{
@@ -149,9 +149,17 @@ mod storage {
                         encoded_data: connector_response.encoded_data.clone(),
                     };
 
+                    let redis_entry = kv::TypedSql {
+                        op: kv::DBOperation::Insert {
+                            insertable: kv::Insertable::ConnectorResponse(
+                                connector_response.clone(),
+                            ),
+                        },
+                    };
+
                     match kv_wrapper::<storage_type::ConnectorResponse, _, _>(
                         self,
-                        KvOperation::HSetNx(&field, &created_connector_resp),
+                        KvOperation::HSetNx(&field, &created_connector_resp, redis_entry),
                         &key,
                     )
                     .await
@@ -163,25 +171,7 @@ mod storage {
                             key: Some(key),
                         })
                         .into_report(),
-                        Ok(HsetnxReply::KeySet) => {
+                        Ok(HsetnxReply::KeySet) => Ok(created_connector_resp),
-                            let redis_entry = kv::TypedSql {
-                                op: kv::DBOperation::Insert {
-                                    insertable: kv::Insertable::ConnectorResponse(
-                                        connector_response.clone(),
-                                    ),
-                                },
-                            };
-                            self.push_to_drainer_stream::<diesel_models::ConnectorResponse>(
-                                redis_entry,
-                                PartitionKey::MerchantIdPaymentId {
-                                    merchant_id,
-                                    payment_id,
-                                },
-                            )
-                            .await
-                            .change_context(errors::StorageError::KVError)?;
-                            Ok(created_connector_resp)
-                        }
                         Err(er) => Err(er).change_context(errors::StorageError::KVError),
                     }
                 }
@@ -259,16 +249,6 @@ mod storage {
                         &updated_connector_response.attempt_id
                     );
 
-                    kv_wrapper::<(), _, _>(
-                        self,
-                        KvOperation::Hset::<storage_type::ConnectorResponse>((&field, redis_value)),
-                        &key,
-                    )
-                    .await
-                    .change_context(errors::StorageError::KVError)?
-                    .try_into_hset()
-                    .change_context(errors::StorageError::KVError)?;
-
                     let redis_entry = kv::TypedSql {
                         op: kv::DBOperation::Update {
                             updatable: kv::Updateable::ConnectorResponseUpdate(
@@ -280,15 +260,19 @@ mod storage {
                         },
                     };
 
-                    self.push_to_drainer_stream::<storage_type::ConnectorResponse>(
+                    kv_wrapper::<(), _, _>(
-                        redis_entry,
+                        self,
-                        PartitionKey::MerchantIdPaymentId {
+                        KvOperation::Hset::<storage_type::ConnectorResponse>(
-                            merchant_id: &updated_connector_response.merchant_id,
+                            (&field, redis_value),
-                            payment_id: &updated_connector_response.payment_id,
+                            redis_entry,
-                        },
+                        ),
+                        &key,
                     )
                     .await
+                    .change_context(errors::StorageError::KVError)?
+                    .try_into_hset()
                     .change_context(errors::StorageError::KVError)?;
+
                     Ok(updated_connector_response)
                 }
             }

crates/router/src/db/refund.rs

@@ -280,7 +280,7 @@ mod storage {
         logger,
         services::Store,
         types::storage::{self as storage_types, enums, kv},
-        utils::{self, db_utils, storage_partitioning::PartitionKey},
+        utils::{self, db_utils},
     };
     #[async_trait::async_trait]
     impl RefundInterface for Store {
@@ -373,9 +373,20 @@ mod storage {
                         "pa_{}_ref_{}",
                         &created_refund.attempt_id, &created_refund.refund_id
                     );
+
+                    let redis_entry = kv::TypedSql {
+                        op: kv::DBOperation::Insert {
+                            insertable: kv::Insertable::Refund(new),
+                        },
+                    };
+
                     match kv_wrapper::<storage_types::Refund, _, _>(
                         self,
-                        KvOperation::HSetNx(&field, &created_refund),
+                        KvOperation::<storage_types::Refund>::HSetNx(
+                            &field,
+                            &created_refund,
+                            redis_entry,
+                        ),
                         &key,
                     )
                     .await
@@ -431,21 +442,6 @@ mod storage {
 
                             futures::future::try_join_all(rev_look).await?;
 
-                            let redis_entry = kv::TypedSql {
-                                op: kv::DBOperation::Insert {
-                                    insertable: kv::Insertable::Refund(new),
-                                },
-                            };
-                            self.push_to_drainer_stream::<storage_types::Refund>(
-                                redis_entry,
-                                PartitionKey::MerchantIdPaymentId {
-                                    merchant_id: &created_refund.merchant_id,
-                                    payment_id: &created_refund.payment_id,
-                                },
-                            )
-                            .await
-                            .change_context(errors::StorageError::KVError)?;
-
                             Ok(created_refund)
                         }
                         Err(er) => Err(er).change_context(errors::StorageError::KVError),
@@ -531,16 +527,6 @@ mod storage {
                         )
                         .change_context(errors::StorageError::SerializationFailed)?;
 
-                    kv_wrapper::<(), _, _>(
-                        self,
-                        KvOperation::Hset::<storage_types::Refund>((&field, redis_value)),
-                        &key,
-                    )
-                    .await
-                    .change_context(errors::StorageError::KVError)?
-                    .try_into_hset()
-                    .change_context(errors::StorageError::KVError)?;
-
                     let redis_entry = kv::TypedSql {
                         op: kv::DBOperation::Update {
                             updatable: kv::Updateable::RefundUpdate(kv::RefundUpdateMems {
@@ -549,15 +535,20 @@ mod storage {
                             }),
                         },
                     };
-                    self.push_to_drainer_stream::<storage_types::Refund>(
+
-                        redis_entry,
+                    kv_wrapper::<(), _, _>(
-                        PartitionKey::MerchantIdPaymentId {
+                        self,
-                            merchant_id: &updated_refund.merchant_id,
+                        KvOperation::Hset::<storage_types::Refund>(
-                            payment_id: &updated_refund.payment_id,
+                            (&field, redis_value),
-                        },
+                            redis_entry,
+                        ),
+                        &key,
                     )
                     .await
+                    .change_context(errors::StorageError::KVError)?
+                    .try_into_hset()
                     .change_context(errors::StorageError::KVError)?;
+
                     Ok(updated_refund)
                 }
             }

crates/router/src/db/reverse_lookup.rs

@@ -74,7 +74,7 @@ mod storage {
             enums, kv,
             reverse_lookup::{ReverseLookup, ReverseLookupNew},
         },
-        utils::{db_utils, storage_partitioning::PartitionKey},
+        utils::db_utils,
     };
 
     #[async_trait::async_trait]
@@ -96,31 +96,22 @@ mod storage {
                         pk_id: new.pk_id.clone(),
                         source: new.source.clone(),
                     };
-                    let combination = &created_rev_lookup.pk_id;
+                    let redis_entry = kv::TypedSql {
+                        op: kv::DBOperation::Insert {
+                            insertable: kv::Insertable::ReverseLookUp(new),
+                        },
+                    };
+
                     match kv_wrapper::<ReverseLookup, _, _>(
                         self,
-                        KvOperation::SetNx(&created_rev_lookup),
+                        KvOperation::SetNx(&created_rev_lookup, redis_entry),
                         format!("reverse_lookup_{}", &created_rev_lookup.lookup_id),
                     )
                     .await
                     .change_context(errors::StorageError::KVError)?
                     .try_into_setnx()
                     {
-                        Ok(SetnxReply::KeySet) => {
+                        Ok(SetnxReply::KeySet) => Ok(created_rev_lookup),
-                            let redis_entry = kv::TypedSql {
-                                op: kv::DBOperation::Insert {
-                                    insertable: kv::Insertable::ReverseLookUp(new),
-                                },
-                            };
-                            self.push_to_drainer_stream::<ReverseLookup>(
-                                redis_entry,
-                                PartitionKey::MerchantIdPaymentIdCombination { combination },
-                            )
-                            .await
-                            .change_context(errors::StorageError::KVError)?;
-
-                            Ok(created_rev_lookup)
-                        }
                         Ok(SetnxReply::KeyNotSet) => Err(errors::StorageError::DuplicateValue {
                             entity: "reverse_lookup",
                             key: Some(created_rev_lookup.lookup_id.clone()),

crates/storage_impl/src/lookup.rs

@@ -11,7 +11,7 @@ use redis_interface::SetnxReply;
 
 use crate::{
     diesel_error_to_data_error,
-    redis::kv_store::{kv_wrapper, KvOperation, PartitionKey},
+    redis::kv_store::{kv_wrapper, KvOperation},
     utils::{self, try_redis_get_else_try_database_get},
     DatabaseStore, KVRouterStore, RouterStore,
 };
@@ -84,31 +84,22 @@ impl<T: DatabaseStore> ReverseLookupInterface for KVRouterStore<T> {
                     pk_id: new.pk_id.clone(),
                     source: new.source.clone(),
                 };
-                let combination = &created_rev_lookup.pk_id;
+                let redis_entry = kv::TypedSql {
+                    op: kv::DBOperation::Insert {
+                        insertable: kv::Insertable::ReverseLookUp(new),
+                    },
+                };
+
                 match kv_wrapper::<DieselReverseLookup, _, _>(
                     self,
-                    KvOperation::SetNx(&created_rev_lookup),
+                    KvOperation::SetNx(&created_rev_lookup, redis_entry),
                     format!("reverse_lookup_{}", &created_rev_lookup.lookup_id),
                 )
                 .await
                 .change_context(errors::StorageError::KVError)?
                 .try_into_setnx()
                 {
-                    Ok(SetnxReply::KeySet) => {
+                    Ok(SetnxReply::KeySet) => Ok(created_rev_lookup),
-                        let redis_entry = kv::TypedSql {
-                            op: kv::DBOperation::Insert {
-                                insertable: kv::Insertable::ReverseLookUp(new),
-                            },
-                        };
-                        self.push_to_drainer_stream::<DieselReverseLookup>(
-                            redis_entry,
-                            PartitionKey::MerchantIdPaymentIdCombination { combination },
-                        )
-                        .await
-                        .change_context(errors::StorageError::KVError)?;
-
-                        Ok(created_rev_lookup)
-                    }
                     Ok(SetnxReply::KeyNotSet) => Err(errors::StorageError::DuplicateValue {
                         entity: "reverse_lookup",
                         key: Some(created_rev_lookup.lookup_id.clone()),

crates/storage_impl/src/metrics.rs

@@ -4,3 +4,7 @@ metrics_context!(CONTEXT);
 global_meter!(GLOBAL_METER, "ROUTER_API");
 
 counter_metric!(KV_MISS, GLOBAL_METER); // No. of KV misses
+
+// Metrics for KV
+counter_metric!(KV_OPERATION_SUCCESSFUL, GLOBAL_METER);
+counter_metric!(KV_OPERATION_FAILED, GLOBAL_METER);

crates/storage_impl/src/payments/payment_attempt.rs

@@ -28,7 +28,7 @@ use router_env::{instrument, tracing};
 use crate::{
     diesel_error_to_data_error,
     lookup::ReverseLookupInterface,
-    redis::kv_store::{kv_wrapper, KvOperation, PartitionKey},
+    redis::kv_store::{kv_wrapper, KvOperation},
     utils::{pg_connection_read, pg_connection_write, try_redis_get_else_try_database_get},
     DataModelExt, DatabaseStore, KVRouterStore, RouterStore,
 };
@@ -363,9 +363,21 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
 
                 let field = format!("pa_{}", created_attempt.attempt_id);
 
+                let redis_entry = kv::TypedSql {
+                    op: kv::DBOperation::Insert {
+                        insertable: kv::Insertable::PaymentAttempt(
+                            payment_attempt.to_storage_model(),
+                        ),
+                    },
+                };
+
                 match kv_wrapper::<PaymentAttempt, _, _>(
                     self,
-                    KvOperation::HSetNx(&field, &created_attempt),
+                    KvOperation::HSetNx(
+                        &field,
+                        &created_attempt.clone().to_storage_model(),
+                        redis_entry,
+                    ),
                     &key,
                 )
                 .await
@@ -391,22 +403,6 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                         self.insert_reverse_lookup(reverse_lookup, storage_scheme)
                             .await?;
 
-                        let redis_entry = kv::TypedSql {
-                            op: kv::DBOperation::Insert {
-                                insertable: kv::Insertable::PaymentAttempt(
-                                    payment_attempt.to_storage_model(),
-                                ),
-                            },
-                        };
-                        self.push_to_drainer_stream::<DieselPaymentAttempt>(
-                            redis_entry,
-                            PartitionKey::MerchantIdPaymentId {
-                                merchant_id: &created_attempt.merchant_id,
-                                payment_id: &created_attempt.payment_id,
-                            },
-                        )
-                        .await
-                        .change_context(errors::StorageError::KVError)?;
                         Ok(created_attempt)
                     }
                     Err(error) => Err(error.change_context(errors::StorageError::KVError)),
@@ -444,9 +440,20 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                     .change_context(errors::StorageError::KVError)?;
                 let field = format!("pa_{}", updated_attempt.attempt_id);
 
+                let redis_entry = kv::TypedSql {
+                    op: kv::DBOperation::Update {
+                        updatable: kv::Updateable::PaymentAttemptUpdate(
+                            kv::PaymentAttemptUpdateMems {
+                                orig: this.clone().to_storage_model(),
+                                update_data: payment_attempt.to_storage_model(),
+                            },
+                        ),
+                    },
+                };
+
                 kv_wrapper::<(), _, _>(
                     self,
-                    KvOperation::Hset::<PaymentAttempt>((&field, redis_value)),
+                    KvOperation::Hset::<DieselPaymentAttempt>((&field, redis_value), redis_entry),
                     &key,
                 )
                 .await
@@ -513,25 +520,6 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                     (_, _) => {}
                 }
 
-                let redis_entry = kv::TypedSql {
-                    op: kv::DBOperation::Update {
-                        updatable: kv::Updateable::PaymentAttemptUpdate(
-                            kv::PaymentAttemptUpdateMems {
-                                orig: this.to_storage_model(),
-                                update_data: payment_attempt.to_storage_model(),
-                            },
-                        ),
-                    },
-                };
-                self.push_to_drainer_stream::<DieselPaymentAttempt>(
-                    redis_entry,
-                    PartitionKey::MerchantIdPaymentId {
-                        merchant_id: &updated_attempt.merchant_id,
-                        payment_id: &updated_attempt.payment_id,
-                    },
-                )
-                .await
-                .change_context(errors::StorageError::KVError)?;
                 Ok(updated_attempt)
             }
         }
@@ -565,7 +553,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
 
                 try_redis_get_else_try_database_get(
                     async {
-                        kv_wrapper(self, KvOperation::<PaymentAttempt>::HGet(&lookup.sk_id), key).await?.try_into_hget()
+                        kv_wrapper(self, KvOperation::<DieselPaymentAttempt>::HGet(&lookup.sk_id), key).await?.try_into_hget()
                     },
                         || async {self.router_store.find_payment_attempt_by_connector_transaction_id_payment_id_merchant_id(connector_transaction_id, payment_id, merchant_id, storage_scheme).await},
                     )
@@ -597,7 +585,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                 let redis_fut = async {
                     let kv_result = kv_wrapper::<PaymentAttempt, _, _>(
                         self,
-                        KvOperation::<PaymentAttempt>::Scan(pattern),
+                        KvOperation::<DieselPaymentAttempt>::Scan(pattern),
                         key,
                     )
                     .await?
@@ -645,7 +633,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                     async {
                         kv_wrapper(
                             self,
-                            KvOperation::<PaymentAttempt>::HGet(&lookup.sk_id),
+                            KvOperation::<DieselPaymentAttempt>::HGet(&lookup.sk_id),
                             key,
                         )
                         .await?
@@ -690,7 +678,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                 let field = format!("pa_{attempt_id}");
                 try_redis_get_else_try_database_get(
                     async {
-                        kv_wrapper(self, KvOperation::<PaymentAttempt>::HGet(&field), key)
+                        kv_wrapper(self, KvOperation::<DieselPaymentAttempt>::HGet(&field), key)
                             .await?
                             .try_into_hget()
                     },
@@ -736,7 +724,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                     async {
                         kv_wrapper(
                             self,
-                            KvOperation::<PaymentAttempt>::HGet(&lookup.sk_id),
+                            KvOperation::<DieselPaymentAttempt>::HGet(&lookup.sk_id),
                             key,
                         )
                         .await?
@@ -784,7 +772,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
                     async {
                         kv_wrapper(
                             self,
-                            KvOperation::<PaymentAttempt>::HGet(&lookup.sk_id),
+                            KvOperation::<DieselPaymentAttempt>::HGet(&lookup.sk_id),
                             key,
                         )
                         .await?
@@ -824,7 +812,7 @@ impl<T: DatabaseStore> PaymentAttemptInterface for KVRouterStore<T> {
             MerchantStorageScheme::RedisKv => {
                 let key = format!("mid_{merchant_id}_pid_{payment_id}");
 
-                kv_wrapper(self, KvOperation::<PaymentAttempt>::Scan("pa_*"), key)
+                kv_wrapper(self, KvOperation::<DieselPaymentAttempt>::Scan("pa_*"), key)
                     .await
                     .change_context(errors::StorageError::KVError)?
                     .try_into_scan()

crates/storage_impl/src/payments/payment_intent.rs

@@ -35,7 +35,7 @@ use router_env::{instrument, tracing};
 
 use crate::{
     diesel_error_to_data_error,
-    redis::kv_store::{kv_wrapper, KvOperation, PartitionKey},
+    redis::kv_store::{kv_wrapper, KvOperation},
     utils::{pg_connection_read, pg_connection_write},
     DataModelExt, DatabaseStore, KVRouterStore,
 };
@@ -93,11 +93,19 @@ impl<T: DatabaseStore> PaymentIntentInterface for KVRouterStore<T> {
                     payment_link_id: new.payment_link_id.clone(),
                     payment_confirm_source: new.payment_confirm_source,
                 };
-                let diesel_intent = created_intent.clone().to_storage_model();
+                let redis_entry = kv::TypedSql {
+                    op: kv::DBOperation::Insert {
+                        insertable: kv::Insertable::PaymentIntent(new.to_storage_model()),
+                    },
+                };
 
                 match kv_wrapper::<DieselPaymentIntent, _, _>(
                     self,
-                    KvOperation::HSetNx(&field, &diesel_intent),
+                    KvOperation::<DieselPaymentIntent>::HSetNx(
+                        &field,
+                        &created_intent.clone().to_storage_model(),
+                        redis_entry,
+                    ),
                     &key,
                 )
                 .await
@@ -109,23 +117,7 @@ impl<T: DatabaseStore> PaymentIntentInterface for KVRouterStore<T> {
                         key: Some(key),
                     })
                     .into_report(),
-                    Ok(HsetnxReply::KeySet) => {
+                    Ok(HsetnxReply::KeySet) => Ok(created_intent),
-                        let redis_entry = kv::TypedSql {
-                            op: kv::DBOperation::Insert {
-                                insertable: kv::Insertable::PaymentIntent(new.to_storage_model()),
-                            },
-                        };
-                        self.push_to_drainer_stream::<DieselPaymentIntent>(
-                            redis_entry,
-                            PartitionKey::MerchantIdPaymentId {
-                                merchant_id: &created_intent.merchant_id,
-                                payment_id: &created_intent.payment_id,
-                            },
-                        )
-                        .await
-                        .change_context(StorageError::KVError)?;
-                        Ok(created_intent)
-                    }
                     Err(error) => Err(error.change_context(StorageError::KVError)),
                 }
             }
@@ -157,16 +149,6 @@ impl<T: DatabaseStore> PaymentIntentInterface for KVRouterStore<T> {
                     Encode::<DieselPaymentIntent>::encode_to_string_of_json(&diesel_intent)
                         .change_context(StorageError::SerializationFailed)?;
 
-                kv_wrapper::<(), _, _>(
-                    self,
-                    KvOperation::<DieselPaymentIntent>::Hset((&field, redis_value)),
-                    &key,
-                )
-                .await
-                .change_context(StorageError::KVError)?
-                .try_into_hset()
-                .change_context(StorageError::KVError)?;
-
                 let redis_entry = kv::TypedSql {
                     op: kv::DBOperation::Update {
                         updatable: kv::Updateable::PaymentIntentUpdate(
@@ -178,15 +160,16 @@ impl<T: DatabaseStore> PaymentIntentInterface for KVRouterStore<T> {
                     },
                 };
 
-                self.push_to_drainer_stream::<DieselPaymentIntent>(
+                kv_wrapper::<(), _, _>(
-                    redis_entry,
+                    self,
-                    PartitionKey::MerchantIdPaymentId {
+                    KvOperation::<DieselPaymentIntent>::Hset((&field, redis_value), redis_entry),
-                        merchant_id: &updated_intent.merchant_id,
+                    &key,
-                        payment_id: &updated_intent.payment_id,
-                    },
                 )
                 .await
+                .change_context(StorageError::KVError)?
+                .try_into_hset()
                 .change_context(StorageError::KVError)?;
+
                 Ok(updated_intent)
             }
         }

crates/storage_impl/src/redis/kv_store.rs

@@ -3,9 +3,10 @@ use std::{fmt::Debug, sync::Arc};
 use common_utils::errors::CustomResult;
 use redis_interface::errors::RedisError;
 use router_derive::TryGetEnumVariant;
+use router_env::logger;
 use serde::de;
 
-use crate::{consts, KVRouterStore};
+use crate::{consts, metrics, store::kv::TypedSql, KVRouterStore};
 
 pub trait KvStorePartition {
     fn partition_number(key: PartitionKey<'_>, num_partitions: u8) -> u32 {
@@ -48,10 +49,11 @@ pub trait RedisConnInterface {
     ) -> error_stack::Result<Arc<redis_interface::RedisConnectionPool>, RedisError>;
 }
 
+/// An enum to represent what operation to do on
 pub enum KvOperation<'a, S: serde::Serialize + Debug> {
-    Hset((&'a str, String)),
+    Hset((&'a str, String), TypedSql),
-    SetNx(S),
+    SetNx(&'a S, TypedSql),
-    HSetNx(&'a str, S),
+    HSetNx(&'a str, &'a S, TypedSql),
     HGet(&'a str),
     Get,
     Scan(&'a str),
@@ -68,6 +70,22 @@ pub enum KvResult<T: de::DeserializeOwned> {
     Scan(Vec<T>),
 }
 
+impl<T> std::fmt::Display for KvOperation<'_, T>
+where
+    T: serde::Serialize + Debug,
+{
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            KvOperation::Hset(_, _) => f.write_str("Hset"),
+            KvOperation::SetNx(_, _) => f.write_str("Setnx"),
+            KvOperation::HSetNx(_, _, _) => f.write_str("HSetNx"),
+            KvOperation::HGet(_) => f.write_str("Hget"),
+            KvOperation::Get => f.write_str("Get"),
+            KvOperation::Scan(_) => f.write_str("Scan"),
+        }
+    }
+}
+
 pub async fn kv_wrapper<'a, T, D, S>(
     store: &KVRouterStore<D>,
     op: KvOperation<'a, S>,
@@ -76,45 +94,101 @@ pub async fn kv_wrapper<'a, T, D, S>(
 where
     T: de::DeserializeOwned,
     D: crate::database::store::DatabaseStore,
-    S: serde::Serialize + Debug,
+    S: serde::Serialize + Debug + KvStorePartition,
 {
     let redis_conn = store.get_redis_conn()?;
 
     let key = key.as_ref();
     let type_name = std::any::type_name::<T>();
+    let operation = op.to_string();
 
-    match op {
+    let partition_key = PartitionKey::MerchantIdPaymentIdCombination { combination: key };
-        KvOperation::Hset(value) => {
+
-            redis_conn
+    let result = async {
-                .set_hash_fields(key, value, Some(consts::KV_TTL))
+        match op {
-                .await?;
+            KvOperation::Hset(value, sql) => {
-            Ok(KvResult::Hset(()))
+                logger::debug!("Operation: {operation} value: {value:?}");
+
+                redis_conn
+                    .set_hash_fields(key, value, Some(consts::KV_TTL))
+                    .await?;
+
+                store
+                    .push_to_drainer_stream::<S>(sql, partition_key)
+                    .await?;
+
+                Ok(KvResult::Hset(()))
+            }
+
+            KvOperation::HGet(field) => {
+                let result = redis_conn
+                    .get_hash_field_and_deserialize(key, field, type_name)
+                    .await?;
+                Ok(KvResult::HGet(result))
+            }
+
+            KvOperation::Scan(pattern) => {
+                let result: Vec<T> = redis_conn.hscan_and_deserialize(key, pattern, None).await?;
+                Ok(KvResult::Scan(result))
+            }
+
+            KvOperation::HSetNx(field, value, sql) => {
+                logger::debug!("Operation: {operation} value: {value:?}");
+
+                let result = redis_conn
+                    .serialize_and_set_hash_field_if_not_exist(
+                        key,
+                        field,
+                        value,
+                        Some(consts::KV_TTL),
+                    )
+                    .await?;
+
+                if matches!(result, redis_interface::HsetnxReply::KeySet) {
+                    store
+                        .push_to_drainer_stream::<S>(sql, partition_key)
+                        .await?;
+                }
+                Ok(KvResult::HSetNx(result))
+            }
+
+            KvOperation::SetNx(value, sql) => {
+                logger::debug!("Operation: {operation} value: {value:?}");
+
+                let result = redis_conn
+                    .serialize_and_set_key_if_not_exist(key, value, Some(consts::KV_TTL.into()))
+                    .await?;
+
+                if matches!(result, redis_interface::SetnxReply::KeySet) {
+                    store
+                        .push_to_drainer_stream::<S>(sql, partition_key)
+                        .await?;
+                }
+
+                Ok(KvResult::SetNx(result))
+            }
+
+            KvOperation::Get => {
+                let result = redis_conn.get_and_deserialize_key(key, type_name).await?;
+                Ok(KvResult::Get(result))
+            }
         }
-        KvOperation::HGet(field) => {
+    };
-            let result = redis_conn
+
-                .get_hash_field_and_deserialize(key, field, type_name)
+    result
-                .await?;
+        .await
-            Ok(KvResult::HGet(result))
+        .map(|result| {
-        }
+            logger::debug!("KvOperation {operation} succeeded");
-        KvOperation::Scan(pattern) => {
+            let keyvalue = router_env::opentelemetry::KeyValue::new("operation", operation.clone());
-            let result: Vec<T> = redis_conn.hscan_and_deserialize(key, pattern, None).await?;
+
-            Ok(KvResult::Scan(result))
+            metrics::KV_OPERATION_SUCCESSFUL.add(&metrics::CONTEXT, 1, &[keyvalue]);
-        }
+            result
-        KvOperation::HSetNx(field, value) => {
+        })
-            let result = redis_conn
+        .map_err(|err| {
-                .serialize_and_set_hash_field_if_not_exist(key, field, value, Some(consts::KV_TTL))
+            logger::error!("KvOperation for {operation} failed with {err:?}");
-                .await?;
+            let keyvalue = router_env::opentelemetry::KeyValue::new("operation", operation);
-            Ok(KvResult::HSetNx(result))
+
-        }
+            metrics::KV_OPERATION_FAILED.add(&metrics::CONTEXT, 1, &[keyvalue]);
-        KvOperation::SetNx(value) => {
+            err
-            let result = redis_conn
+        })
-                .serialize_and_set_key_if_not_exist(key, value, Some(consts::KV_TTL.into()))
-                .await?;
-            Ok(KvResult::SetNx(result))
-        }
-        KvOperation::Get => {
-            let result = redis_conn.get_and_deserialize_key(key, type_name).await?;
-            Ok(KvResult::Get(result))
-        }
-    }
 }