Feature: transaction api

common/management/tests/it/user.rs

@@ -29,6 +29,8 @@ use common_meta_types::Operation;
 use common_meta_types::PasswordHashMethod;
 use common_meta_types::PrefixListReply;
 use common_meta_types::SeqV;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_meta_types::UpsertKVAction;
 use common_meta_types::UpsertKVActionReply;
 use common_meta_types::UserIdentity;
@@ -53,6 +55,9 @@ mock! {
         ) -> Result<MGetKVActionReply,MetaError>;
 
         async fn prefix_list_kv(&self, prefix: &str) -> Result<PrefixListReply, MetaError>;
+
+        async fn transaction(&self, txn: TransactionReq) -> Result<TransactionResponse, MetaError>;
+
         }
 }
 

common/meta/api/src/kv_api.rs

@@ -20,6 +20,8 @@ use common_meta_types::GetKVActionReply;
 use common_meta_types::MGetKVActionReply;
 use common_meta_types::MetaError;
 use common_meta_types::PrefixListReply;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_meta_types::UpsertKVAction;
 use common_meta_types::UpsertKVActionReply;
 
@@ -44,6 +46,8 @@ pub trait KVApi: Send + Sync {
     async fn mget_kv(&self, key: &[String]) -> Result<MGetKVActionReply, MetaError>;
 
     async fn prefix_list_kv(&self, prefix: &str) -> Result<PrefixListReply, MetaError>;
+
+    async fn transaction(&self, txn: TransactionReq) -> Result<TransactionResponse, MetaError>;
 }
 
 #[async_trait]
@@ -63,4 +67,8 @@ impl<U: KVApi, T: Deref<Target = U> + Send + Sync> KVApi for T {
     async fn prefix_list_kv(&self, prefix: &str) -> Result<PrefixListReply, MetaError> {
         self.deref().prefix_list_kv(prefix).await
     }
+
+    async fn transaction(&self, txn: TransactionReq) -> Result<TransactionResponse, MetaError> {
+        self.deref().transaction(txn).await
+    }
 }

common/meta/embedded/src/kv_api_impl.rs

@@ -19,6 +19,8 @@ use common_meta_types::GetKVActionReply;
 use common_meta_types::MGetKVActionReply;
 use common_meta_types::MetaError;
 use common_meta_types::PrefixListReply;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_meta_types::UpsertKVAction;
 use common_meta_types::UpsertKVActionReply;
 
@@ -45,4 +47,9 @@ impl KVApi for MetaEmbedded {
         let sm = self.inner.lock().await;
         sm.prefix_list_kv(prefix).await
     }
+
+    async fn transaction(&self, txn: TransactionReq) -> Result<TransactionResponse, MetaError> {
+        let sm = self.inner.lock().await;
+        sm.transaction(txn).await
+    }
 }

common/meta/grpc/src/grpc_client.rs

@@ -30,9 +30,12 @@ use common_meta_types::protobuf::meta_service_client::MetaServiceClient;
 use common_meta_types::protobuf::HandshakeRequest;
 use common_meta_types::protobuf::RaftReply;
 use common_meta_types::protobuf::RaftRequest;
+use common_meta_types::protobuf::TxnRequest;
 use common_meta_types::ConnectionError;
 use common_meta_types::MetaError;
 use common_meta_types::MetaNetworkError;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_tracing::tracing;
 use futures::stream::StreamExt;
 use prost::Message;
@@ -265,6 +268,42 @@ impl MetaGrpcClient {
         let res: std::result::Result<R, MetaError> = raft_reply.into();
         res
     }
+
+    //#[tracing::instrument(level = "debug", skip(self))]
+    pub(crate) async fn transaction(
+        &self,
+        req: TransactionReq,
+    ) -> std::result::Result<TransactionResponse, MetaError> {
+        let txn: TxnRequest = req.to_pb();
+
+        let req: Request<TxnRequest> = Request::new(txn.clone());
+        let req = common_tracing::inject_span_to_tonic_request(req);
+
+        let mut client = self.make_client().await?;
+        let result = client.transaction(req).await;
+        let result = match result {
+            Ok(r) => Ok(r),
+            Err(s) => {
+                if status_is_retryable(&s) {
+                    {
+                        let mut token = self.token.write().await;
+                        *token = None;
+                    }
+                    let mut client = self.make_client().await?;
+                    let req: Request<TxnRequest> = Request::new(txn);
+                    let req = common_tracing::inject_span_to_tonic_request(req);
+                    let ret = client.transaction(req).await?.into_inner();
+                    return Ok(TransactionResponse::new(ret));
+                } else {
+                    Err(s)
+                }
+            }
+        };
+
+        let reply = result?.into_inner();
+
+        Ok(TransactionResponse::new(reply))
+    }
 }
 
 fn status_is_retryable(status: &Status) -> bool {

common/meta/grpc/src/kv_api_impl.rs

@@ -17,6 +17,8 @@ use common_meta_types::GetKVActionReply;
 use common_meta_types::MGetKVActionReply;
 use common_meta_types::MetaError;
 use common_meta_types::PrefixListReply;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_meta_types::UpsertKVAction;
 use common_meta_types::UpsertKVActionReply;
 
@@ -51,4 +53,9 @@ impl KVApi for MetaGrpcClient {
         let reply = self.do_read(PrefixListReq(prefix.to_string())).await?;
         Ok(reply)
     }
+
+    async fn transaction(&self, txn: TransactionReq) -> Result<TransactionResponse, MetaError> {
+        let reply = self.transaction(txn).await?;
+        Ok(reply)
+    }
 }

common/meta/grpc/tests/it/grpc_server.rs

@@ -23,6 +23,8 @@ use common_meta_types::protobuf::ExportedChunk;
 use common_meta_types::protobuf::HandshakeResponse;
 use common_meta_types::protobuf::RaftReply;
 use common_meta_types::protobuf::RaftRequest;
+use common_meta_types::protobuf::TxnReply;
+use common_meta_types::protobuf::TxnRequest;
 use common_meta_types::protobuf::WatchRequest;
 use common_meta_types::protobuf::WatchResponse;
 use futures::Stream;
@@ -84,6 +86,13 @@ impl MetaService for GrpcServiceForTestImpl {
     ) -> Result<Response<Self::WatchStream>, Status> {
         todo!()
     }
+
+    async fn transaction(
+        &self,
+        _request: Request<TxnRequest>,
+    ) -> Result<Response<TxnReply>, Status> {
+        todo!()
+    }
 }
 
 pub fn start_grpc_server() -> String {

common/meta/raft-store/src/state_machine/sm.rs

@@ -55,6 +55,17 @@ use common_meta_types::SeqV;
 use common_meta_types::ShareInfo;
 use common_meta_types::TableAlreadyExists;
 use common_meta_types::TableMeta;
+use common_meta_types::TransactionCondition;
+use common_meta_types::TransactionDeleteRequest;
+use common_meta_types::TransactionDeleteResponse;
+use common_meta_types::TransactionGetRequest;
+use common_meta_types::TransactionGetResponse;
+use common_meta_types::TransactionOperation;
+use common_meta_types::TransactionOperationResponse;
+use common_meta_types::TransactionPutRequest;
+use common_meta_types::TransactionPutResponse;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_meta_types::UnknownDatabase;
 use common_meta_types::UnknownDatabaseId;
 use common_meta_types::UnknownShare;
@@ -559,6 +570,173 @@ impl StateMachine {
         Ok(Change::new(prev, result).into())
     }
 
+    fn txn_execute_cond(&self, cond: &TransactionCondition) -> bool {
+        let is_key_condition = cond.is_key_condition();
+
+        let key = cond.key.generate_key();
+        let sv = self.kvs().get(&key);
+
+        if is_key_condition {
+            match sv {
+                Ok(sv) => match sv {
+                    Some(_) => cond.return_key_condition_result(true),
+                    None => cond.return_key_condition_result(false),
+                },
+                Err(_) => cond.return_key_condition_result(false),
+            }
+        } else {
+            match sv {
+                Ok(sv) => match sv {
+                    Some(_val) => cond.return_key_condition_result(true),
+                    None => false,
+                },
+                Err(_) => false,
+            }
+        }
+    }
+
+    fn txn_execute_get_operation(
+        &self,
+        get: &TransactionGetRequest,
+        resp: &mut TransactionResponse,
+    ) {
+        let key = get.key.generate_key();
+        let sv = self.kvs().get(&key);
+        let get = match sv {
+            Ok(v) => match v {
+                Some(v) => TransactionGetResponse {
+                    key: get.key.clone(),
+                    value: Some(v),
+                },
+                None => TransactionGetResponse {
+                    key: get.key.clone(),
+                    value: None,
+                },
+            },
+            Err(_v) => TransactionGetResponse {
+                key: get.key.clone(),
+                value: None,
+            },
+        };
+
+        resp.responses
+            .push(TransactionOperationResponse::TransactionGetResponse(get));
+    }
+
+    fn txn_execute_put_operation(
+        &self,
+        txn_tree: &TransactionSledTree,
+        put: &TransactionPutRequest,
+        resp: &mut TransactionResponse,
+    ) {
+        let key = put.key.generate_key();
+        let sub_tree = txn_tree.key_space::<GenericKV>();
+        let key_str = key.to_string();
+
+        if let Ok(result) = self.txn_sub_tree_upsert(
+            &sub_tree,
+            &key_str,
+            &MatchSeq::Any,
+            Operation::Update(put.value.clone()),
+            None,
+        ) {
+            let (prev, result) = result;
+            tracing::debug!("applied txn_execute_put_operation: {} {:?}", key, result);
+
+            let put_resp = TransactionPutResponse {
+                key: put.key.clone(),
+                prev_value: if put.prev_kv { prev } else { None },
+            };
+
+            resp.responses
+                .push(TransactionOperationResponse::TransactionPutResponse(
+                    put_resp,
+                ));
+        }
+    }
+
+    fn txn_execute_delete_operation(
+        &self,
+        txn_tree: &TransactionSledTree,
+        delete: &TransactionDeleteRequest,
+        resp: &mut TransactionResponse,
+    ) {
+        let key = delete.key.generate_key();
+        let sub_tree = txn_tree.key_space::<GenericKV>();
+        let key_str = key.to_string();
+
+        let del_resp = if let Ok(result) =
+            self.txn_sub_tree_upsert(&sub_tree, &key_str, &MatchSeq::Any, Operation::Delete, None)
+        {
+            let (prev, result) = result;
+            tracing::debug!("applied txn_execute_delete_operation: {} {:?}", key, result);
+
+            TransactionDeleteResponse {
+                key: delete.key.clone(),
+                success: true,
+                prev_value: if delete.prev_kv { prev } else { None },
+            }
+        } else {
+            TransactionDeleteResponse {
+                key: delete.key.clone(),
+                success: false,
+                prev_value: None,
+            }
+        };
+
+        resp.responses
+            .push(TransactionOperationResponse::TransactionDeleteResponse(
+                del_resp,
+            ));
+    }
+
+    fn txn_execute_operation(
+        &self,
+        txn_tree: &TransactionSledTree,
+        op: &TransactionOperation,
+        resp: &mut TransactionResponse,
+    ) {
+        match op {
+            TransactionOperation::TransactionGetRequest(get) => {
+                self.txn_execute_get_operation(get, resp);
+            }
+            TransactionOperation::TransactionPutRequest(put) => {
+                self.txn_execute_put_operation(txn_tree, put, resp);
+            }
+            TransactionOperation::TransactionDeleteRequest(delete) => {
+                self.txn_execute_delete_operation(txn_tree, delete, resp);
+            }
+        }
+    }
+
+    fn apply_txn_cmd(
+        &self,
+        req: &TransactionReq,
+        txn_tree: &TransactionSledTree,
+    ) -> MetaStorageResult<AppliedState> {
+        let cond = &req.cond;
+
+        let ops: &Vec<TransactionOperation>;
+        let success = if self.txn_execute_cond(cond) {
+            ops = &req.if_then;
+            true
+        } else {
+            ops = &req.else_then;
+            false
+        };
+
+        let mut resp: TransactionResponse = TransactionResponse {
+            success,
+            responses: vec![],
+        };
+
+        for op in ops {
+            self.txn_execute_operation(txn_tree, op, &mut resp);
+        }
+
+        Ok(AppliedState::TransactionResponse(resp))
+    }
+
     #[tracing::instrument(level = "debug", skip(self, txn_tree))]
     fn apply_upsert_table_options_cmd(
         &self,
@@ -655,6 +833,8 @@ impl StateMachine {
             } => self.apply_update_kv_cmd(key, seq, value_op, value_meta, txn_tree),
 
             Cmd::UpsertTableOptions(ref req) => self.apply_upsert_table_options_cmd(req, txn_tree),
+
+            Cmd::Transaction(txn) => self.apply_txn_cmd(txn, txn_tree),
         }
     }
 

common/meta/raft-store/src/state_machine/sm_kv_api_impl.rs

@@ -19,6 +19,8 @@ use common_meta_types::GetKVActionReply;
 use common_meta_types::MGetKVActionReply;
 use common_meta_types::MetaError;
 use common_meta_types::SeqV;
+use common_meta_types::TransactionReq;
+use common_meta_types::TransactionResponse;
 use common_meta_types::UpsertKVAction;
 use common_meta_types::UpsertKVActionReply;
 use common_tracing::tracing;
@@ -48,6 +50,22 @@ impl KVApi for StateMachine {
         }
     }
 
+    async fn transaction(&self, txn: TransactionReq) -> Result<TransactionResponse, MetaError> {
+        let cmd = Cmd::Transaction(txn);
+
+        let res = self.sm_tree.txn(true, |t| {
+            let r = self.apply_cmd(&cmd, &t).unwrap();
+            Ok(r)
+        })?;
+
+        match res {
+            AppliedState::TransactionResponse(x) => Ok(x),
+            _ => {
+                panic!("expect AppliedState::TransactionResponse");
+            }
+        }
+    }
+
     async fn get_kv(&self, key: &str) -> Result<GetKVActionReply, MetaError> {
         // TODO(xp) refine get(): a &str is enough for key
         let sv = self.kvs().get(&key.to_string())?;