Bug: Make KV storage flow more generic



Model the KV storage flow as a series of common operations (For example, `set in redis -> get sql query -> push query to redis stream`) and define generic utilities that perform these operations for all KV-enabled storage interfaces (as opposed to rewriting code).
diff --git a/crates/router/src/db/payment_attempt.rs b/crates/router/src/db/payment_attempt.rs
index 3fb453d35be..8881682448f 100644
--- a/crates/router/src/db/payment_attempt.rs
+++ b/crates/router/src/db/payment_attempt.rs
@@ -315,7 +315,7 @@ impl PaymentAttemptInterface for MockDb {
 mod storage {
     use common_utils::date_time;
     use error_stack::{IntoReport, ResultExt};
-    use redis_interface::{HsetnxReply, RedisEntryId};
+    use redis_interface::HsetnxReply;
 
     use super::PaymentAttemptInterface;
     use crate::{
@@ -324,7 +324,6 @@ mod storage {
         db::reverse_lookup::ReverseLookupInterface,
         services::Store,
         types::storage::{enums, kv, payment_attempt::*, ReverseLookupNew},
-        utils::storage_partitioning::KvStorePartition,
     };
 
     #[async_trait::async_trait]
@@ -421,23 +420,14 @@ mod storage {
                                     insertable: kv::Insertable::PaymentAttempt(payment_attempt),
                                 },
                             };
-                            let stream_name = self.get_drainer_stream_name(&PaymentAttempt::shard_key(
+                            self.push_to_drainer_stream::<PaymentAttempt>(
+                                redis_entry,
                                 crate::utils::storage_partitioning::PartitionKey::MerchantIdPaymentId {
                                     merchant_id: &created_attempt.merchant_id,
                                     payment_id: &created_attempt.payment_id,
-                                },
-                                self.config.drainer_num_partitions,
-                            ));
-                            self.redis_conn
-                                .stream_append_entry(
-                                    &stream_name,
-                                    &RedisEntryId::AutoGeneratedID,
-                                    redis_entry
-                                        .to_field_value_pairs()
-                                        .change_context(errors::StorageError::KVError)?,
-                                )
-                                .await
-                                .change_context(errors::StorageError::KVError)?;
+                                }
+                            )
+                            .await?;
                             Ok(created_attempt)
                         }
                         Err(error) => Err(error.change_context(errors::StorageError::KVError)),
@@ -508,24 +498,14 @@ mod storage {
                             ),
                         },
                     };
-
-                    let stream_name = self.get_drainer_stream_name(&PaymentAttempt::shard_key(
+                    self.push_to_drainer_stream::<PaymentAttempt>(
+                        redis_entry,
                         crate::utils::storage_partitioning::PartitionKey::MerchantIdPaymentId {
                             merchant_id: &updated_attempt.merchant_id,
                             payment_id: &updated_attempt.payment_id,
                         },
-                        self.config.drainer_num_partitions,
-                    ));
-                    self.redis_conn
-                        .stream_append_entry(
-                            &stream_name,
-                            &RedisEntryId::AutoGeneratedID,
-                            redis_entry
-                                .to_field_value_pairs()
-                                .change_context(errors::StorageError::KVError)?,
-                        )
-                        .await
-                        .change_context(errors::StorageError::KVError)?;
+                    )
+                    .await?;
                     Ok(updated_attempt)
                 }
             }
diff --git a/crates/router/src/db/payment_intent.rs b/crates/router/src/db/payment_intent.rs
index e3cc2b1c23e..30c5a9285ef 100644
--- a/crates/router/src/db/payment_intent.rs
+++ b/crates/router/src/db/payment_intent.rs
@@ -41,7 +41,7 @@ pub trait PaymentIntentInterface {
 mod storage {
     use common_utils::date_time;
     use error_stack::{IntoReport, ResultExt};
-    use redis_interface::{HsetnxReply, RedisEntryId};
+    use redis_interface::HsetnxReply;
 
     use super::PaymentIntentInterface;
     #[cfg(feature = "olap")]
@@ -51,10 +51,7 @@ mod storage {
         core::errors::{self, CustomResult},
         services::Store,
         types::storage::{enums, kv, payment_intent::*},
-        utils::{
-            self,
-            storage_partitioning::{self, KvStorePartition},
-        },
+        utils::{self, storage_partitioning},
     };
 
     #[async_trait::async_trait]
@@ -113,24 +110,14 @@ mod storage {
                                     insertable: kv::Insertable::PaymentIntent(new),
                                 },
                             };
-                            let stream_name =
-                                self.get_drainer_stream_name(&PaymentIntent::shard_key(
-                                    storage_partitioning::PartitionKey::MerchantIdPaymentId {
-                                        merchant_id: &created_intent.merchant_id,
-                                        payment_id: &created_intent.payment_id,
-                                    },
-                                    self.config.drainer_num_partitions,
-                                ));
-                            self.redis_conn
-                                .stream_append_entry(
-                                    &stream_name,
-                                    &RedisEntryId::AutoGeneratedID,
-                                    redis_entry
-                                        .to_field_value_pairs()
-                                        .change_context(errors::StorageError::KVError)?,
-                                )
-                                .await
-                                .change_context(errors::StorageError::KVError)?;
+                            self.push_to_drainer_stream::<PaymentIntent>(
+                                redis_entry,
+                                storage_partitioning::PartitionKey::MerchantIdPaymentId {
+                                    merchant_id: &created_intent.merchant_id,
+                                    payment_id: &created_intent.payment_id,
+                                },
+                            )
+                            .await?;
                             Ok(created_intent)
                         }
                         Err(error) => Err(error.change_context(errors::StorageError::KVError)),
@@ -182,24 +169,14 @@ mod storage {
                         },
                     };
 
-                    let stream_name = self.get_drainer_stream_name(&PaymentIntent::shard_key(
+                    self.push_to_drainer_stream::<PaymentIntent>(
+                        redis_entry,
                         storage_partitioning::PartitionKey::MerchantIdPaymentId {
                             merchant_id: &updated_intent.merchant_id,
                             payment_id: &updated_intent.payment_id,
                         },
-                        self.config.drainer_num_partitions,
-                    ));
-
-                    self.redis_conn
-                        .stream_append_entry(
-                            &stream_name,
-                            &RedisEntryId::AutoGeneratedID,
-                            redis_entry
-                                .to_field_value_pairs()
-                                .change_context(errors::StorageError::KVError)?,
-                        )
-                        .await
-                        .change_context(errors::StorageError::KVError)?;
+                    )
+                    .await?;
                     Ok(updated_intent)
                 }
             }
diff --git a/crates/router/src/db/refund.rs b/crates/router/src/db/refund.rs
index 52c5710d574..8896a0ad3aa 100644
--- a/crates/router/src/db/refund.rs
+++ b/crates/router/src/db/refund.rs
@@ -199,7 +199,7 @@ mod storage {
 mod storage {
     use common_utils::date_time;
     use error_stack::{IntoReport, ResultExt};
-    use redis_interface::{HsetnxReply, RedisEntryId};
+    use redis_interface::HsetnxReply;
 
     use super::RefundInterface;
     use crate::{
@@ -209,10 +209,7 @@ mod storage {
         logger,
         services::Store,
         types::storage::{self as storage_types, enums, kv},
-        utils::{
-            self, db_utils,
-            storage_partitioning::{KvStorePartition, PartitionKey},
-        },
+        utils::{self, db_utils, storage_partitioning::PartitionKey},
     };
     #[async_trait::async_trait]
     impl RefundInterface for Store {
@@ -344,25 +341,15 @@ mod storage {
                                     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?;
 
-                            let stream_name =
-                                self.get_drainer_stream_name(&storage_types::Refund::shard_key(
-                                    PartitionKey::MerchantIdPaymentId {
-                                        merchant_id: &created_refund.merchant_id,
-                                        payment_id: &created_refund.payment_id,
-                                    },
-                                    self.config.drainer_num_partitions,
-                                ));
-                            self.redis_conn
-                                .stream_append_entry(
-                                    &stream_name,
-                                    &RedisEntryId::AutoGeneratedID,
-                                    redis_entry
-                                        .to_field_value_pairs()
-                                        .change_context(errors::StorageError::KVError)?,
-                                )
-                                .await
-                                .change_context(errors::StorageError::KVError)?;
                             Ok(created_refund)
                         }
                         Err(er) => Err(er).change_context(errors::StorageError::KVError),
@@ -449,14 +436,6 @@ mod storage {
                         .await
                         .change_context(errors::StorageError::KVError)?;
 
-                    let stream_name =
-                        self.get_drainer_stream_name(&storage_types::Refund::shard_key(
-                            PartitionKey::MerchantIdPaymentId {
-                                merchant_id: &updated_refund.merchant_id,
-                                payment_id: &updated_refund.payment_id,
-                            },
-                            self.config.drainer_num_partitions,
-                        ));
                     let redis_entry = kv::TypedSql {
                         op: kv::DBOperation::Update {
                             updatable: kv::Updateable::RefundUpdate(kv::RefundUpdateMems {
@@ -465,16 +444,14 @@ mod storage {
                             }),
                         },
                     };
-                    self.redis_conn
-                        .stream_append_entry(
-                            &stream_name,
-                            &RedisEntryId::AutoGeneratedID,
-                            redis_entry
-                                .to_field_value_pairs()
-                                .change_context(errors::StorageError::KVError)?,
-                        )
-                        .await
-                        .change_context(errors::StorageError::KVError)?;
+                    self.push_to_drainer_stream::<storage_types::Refund>(
+                        redis_entry,
+                        PartitionKey::MerchantIdPaymentId {
+                            merchant_id: &updated_refund.merchant_id,
+                            payment_id: &updated_refund.payment_id,
+                        },
+                    )
+                    .await?;
                     Ok(updated_refund)
                 }
             }
diff --git a/crates/router/src/services.rs b/crates/router/src/services.rs
index fba45c00ba6..9d40d39cda7 100644
--- a/crates/router/src/services.rs
+++ b/crates/router/src/services.rs
@@ -45,4 +45,29 @@ impl Store {
         // Example: {shard_5}_drainer_stream
         format!("{{{}}}_{}", shard_key, self.config.drainer_stream_name,)
     }
+
+    #[cfg(feature = "kv_store")]
+    pub(crate) async fn push_to_drainer_stream<T>(
+        &self,
+        redis_entry: storage_models::kv::TypedSql,
+        partition_key: crate::utils::storage_partitioning::PartitionKey<'_>,
+    ) -> crate::core::errors::CustomResult<(), crate::core::errors::StorageError>
+    where
+        T: crate::utils::storage_partitioning::KvStorePartition,
+    {
+        use error_stack::ResultExt;
+
+        let shard_key = T::shard_key(partition_key, self.config.drainer_num_partitions);
+        let stream_name = self.get_drainer_stream_name(&shard_key);
+        self.redis_conn
+            .stream_append_entry(
+                &stream_name,
+                &redis_interface::RedisEntryId::AutoGeneratedID,
+                redis_entry
+                    .to_field_value_pairs()
+                    .change_context(crate::core::errors::StorageError::KVError)?,
+            )
+            .await
+            .change_context(crate::core::errors::StorageError::KVError)
+    }
 }

## Type of Change
<!-- Put an `x` in the boxes that apply -->
- [x] Refactoring

## Description
<!-- Describe your changes in detail -->
Use a generic function to push sql queries into drainer stream.

## Motivation and Context
<!--
Why is this change required? What problem does it solve?
If it fixes an open issue, please link to the issue here.

If you don't have an issue, we'd recommend starting with one first so the PR
can focus on the implementation (unless its an obvious bug or documentation fix
that will have little conversation).
-->
Code reusability for pushing to drainer stream. This Closes #97.

## How did you test it?
<!--
Did you write an integration/unit/API test to verify the code changes?
Or did you test this change manually (provide relevant screenshots)?
-->
Manual

## Checklist
<!-- Put an `x` in the boxes that apply -->

- [x] I formatted the code `cargo +nightly fmt --all`
- [x] I addressed lints thrown by `cargo clippy`
- [x] I reviewed submitted code