feat: add Txn for pg kv backend (#5266)

* feat: txn for pg kv backend

* chore: clippy

* fix: txn uses one client

* test: clean up and txn test

* test: clean up

* test: change lock_id to avoid conflict in test

* test: use different prefix in pg election test

* fix(test): just a fix

* test: aggregate multiple test to avoid concurrency problem

* test: use uuid instead of rng

* perf: batch cmp in txn

* perf: batch same op in txn

src/cli/Cargo.toml

@@ -4,6 +4,9 @@ version.workspace = true
 edition.workspace = true
 license.workspace = true
 
+[features]
+pg_kvbackend = ["common-meta/pg_kvbackend"]
+
 [lints]
 workspace = true
 

src/cli/src/bench.rs

@@ -22,6 +22,9 @@ use clap::Parser;
 use common_error::ext::BoxedError;
 use common_meta::key::{TableMetadataManager, TableMetadataManagerRef};
 use common_meta::kv_backend::etcd::EtcdStore;
+use common_meta::kv_backend::memory::MemoryKvBackend;
+#[cfg(feature = "pg_kvbackend")]
+use common_meta::kv_backend::postgres::PgStore;
 use common_meta::peer::Peer;
 use common_meta::rpc::router::{Region, RegionRoute};
 use common_telemetry::info;
@@ -55,18 +58,32 @@ where
 #[derive(Debug, Default, Parser)]
 pub struct BenchTableMetadataCommand {
     #[clap(long)]
-    etcd_addr: String,
+    etcd_addr: Option<String>,
+    #[cfg(feature = "pg_kvbackend")]
+    #[clap(long)]
+    postgres_addr: Option<String>,
     #[clap(long)]
     count: u32,
 }
 
 impl BenchTableMetadataCommand {
     pub async fn build(&self) -> std::result::Result<Box<dyn Tool>, BoxedError> {
-        let etcd_store = EtcdStore::with_endpoints([&self.etcd_addr], 128)
-            .await
-            .unwrap();
+        let kv_backend = if let Some(etcd_addr) = &self.etcd_addr {
+            info!("Using etcd as kv backend");
+            EtcdStore::with_endpoints([etcd_addr], 128).await.unwrap()
+        } else {
+            Arc::new(MemoryKvBackend::new())
+        };
+
+        #[cfg(feature = "pg_kvbackend")]
+        let kv_backend = if let Some(postgres_addr) = &self.postgres_addr {
+            info!("Using postgres as kv backend");
+            PgStore::with_url(postgres_addr, 128).await.unwrap()
+        } else {
+            kv_backend
+        };
 
-        let table_metadata_manager = Arc::new(TableMetadataManager::new(etcd_store));
+        let table_metadata_manager = Arc::new(TableMetadataManager::new(kv_backend));
 
         let tool = BenchTableMetadata {
             table_metadata_manager,

src/common/meta/src/error.rs

@@ -683,6 +683,16 @@ pub enum Error {
         location: Location,
     },
 
+    #[cfg(feature = "pg_kvbackend")]
+    #[snafu(display("Failed to {} Postgres transaction", operation))]
+    PostgresTransaction {
+        #[snafu(source)]
+        error: tokio_postgres::Error,
+        #[snafu(implicit)]
+        location: Location,
+        operation: String,
+    },
+
     #[snafu(display(
         "Datanode table info not found, table id: {}, datanode id: {}",
         table_id,
@@ -794,9 +804,10 @@ impl ErrorExt for Error {
             | EmptyDdlTasks { .. } => StatusCode::InvalidArguments,
 
             #[cfg(feature = "pg_kvbackend")]
-            PostgresExecution { .. } | CreatePostgresPool { .. } | GetPostgresConnection { .. } => {
-                StatusCode::Internal
-            }
+            PostgresExecution { .. }
+            | CreatePostgresPool { .. }
+            | GetPostgresConnection { .. }
+            | PostgresTransaction { .. } => StatusCode::Internal,
             Error::DatanodeTableInfoNotFound { .. } => StatusCode::Internal,
         }
     }

src/common/meta/src/kv_backend/etcd.rs

@@ -542,6 +542,8 @@ mod tests {
         prepare_kv_with_prefix, test_kv_batch_delete_with_prefix, test_kv_batch_get_with_prefix,
         test_kv_compare_and_put_with_prefix, test_kv_delete_range_with_prefix,
         test_kv_put_with_prefix, test_kv_range_2_with_prefix, test_kv_range_with_prefix,
+        test_txn_compare_equal, test_txn_compare_greater, test_txn_compare_less,
+        test_txn_compare_not_equal, test_txn_one_compare_op, text_txn_multi_compare_op,
         unprepare_kv,
     };
 
@@ -628,4 +630,17 @@ mod tests {
             test_kv_batch_delete_with_prefix(kv_backend, prefix.to_vec()).await;
         }
     }
+
+    #[tokio::test]
+    async fn test_etcd_txn() {
+        if let Some(kv_backend) = build_kv_backend().await {
+            let kv_backend_ref = Arc::new(kv_backend);
+            test_txn_one_compare_op(kv_backend_ref.clone()).await;
+            text_txn_multi_compare_op(kv_backend_ref.clone()).await;
+            test_txn_compare_equal(kv_backend_ref.clone()).await;
+            test_txn_compare_greater(kv_backend_ref.clone()).await;
+            test_txn_compare_less(kv_backend_ref.clone()).await;
+            test_txn_compare_not_equal(kv_backend_ref).await;
+        }
+    }
 }

src/common/meta/src/kv_backend/memory.rs

@@ -325,7 +325,9 @@ mod tests {
     use crate::error::Error;
     use crate::kv_backend::test::{
         prepare_kv, test_kv_batch_delete, test_kv_batch_get, test_kv_compare_and_put,
-        test_kv_delete_range, test_kv_put, test_kv_range, test_kv_range_2,
+        test_kv_delete_range, test_kv_put, test_kv_range, test_kv_range_2, test_txn_compare_equal,
+        test_txn_compare_greater, test_txn_compare_less, test_txn_compare_not_equal,
+        test_txn_one_compare_op, text_txn_multi_compare_op,
     };
 
     async fn mock_mem_store_with_data() -> MemoryKvBackend<Error> {
@@ -383,4 +385,15 @@ mod tests {
 
         test_kv_batch_delete(kv_backend).await;
     }
+
+    #[tokio::test]
+    async fn test_memory_txn() {
+        let kv_backend = Arc::new(MemoryKvBackend::<Error>::new());
+        test_txn_one_compare_op(kv_backend.clone()).await;
+        text_txn_multi_compare_op(kv_backend.clone()).await;
+        test_txn_compare_equal(kv_backend.clone()).await;
+        test_txn_compare_greater(kv_backend.clone()).await;
+        test_txn_compare_less(kv_backend.clone()).await;
+        test_txn_compare_not_equal(kv_backend).await;
+    }
 }