refactor: primary_key_or

Signed-off-by: ekexium <ekexium@gmail.com>

src/transaction/buffer.rs

@@ -27,8 +27,9 @@ impl InnerBuffer {
         self.entry_map.insert(key, entry);
     }
 
-    pub fn get_primary_key_or_insert(&mut self, key: &Key) -> &Key {
-        self.primary_key.get_or_insert(key.clone())
+    /// Set the primary key if it is not set
+    pub fn primary_key_or(&mut self, key: &Key) {
+        self.primary_key.get_or_insert(key.clone());
     }
 }
 
@@ -49,13 +50,9 @@ impl Buffer {
         self.inner.lock().await.primary_key.clone()
     }
 
-    /// Get the primary key of the buffer, if not exists, use `key` as the primary key.
-    pub async fn get_primary_key_or_insert(&self, key: &Key) -> Key {
-        self.inner
-            .lock()
-            .await
-            .get_primary_key_or_insert(key)
-            .clone()
+    /// Set the primary key if it is not set
+    pub async fn primary_key_or(&self, key: &Key) {
+        self.inner.lock().await.primary_key_or(key);
     }
 
     /// Get a value from the buffer.

src/transaction/transaction.rs

@@ -674,7 +674,7 @@ impl<PdC: PdClient> Transaction<PdC> {
         let pairs = plan.execute().await;
 
         // primary key will be set here if needed
-        self.buffer.get_primary_key_or_insert(&first_key).await;
+        self.buffer.primary_key_or(&first_key).await;
 
         self.start_auto_heartbeat().await;