Use async kms

Signed-off-by: Corwin Joy <[email protected]>

Author: jgiannuzzi

crates/core/Cargo.toml

@@ -94,7 +94,7 @@ datatest-stable = "0.3"
 deltalake-test = { path = "../test" }
 dotenvy = "0"
 fs_extra = "1.2.0"
-parquet-key-management = { version = "0.4.1", features = ["_test_utils", "datafusion"] }
+parquet-key-management = { version = "0.5.0", features = ["_test_utils", "datafusion", "tokio"] }
 paste = "1"
 pretty_assertions = "1.2.1"
 pretty_env_logger = "0.5.0"

crates/core/tests/commands_with_encryption.rs

@@ -19,11 +19,11 @@ use deltalake_core::table::file_format_options::{
 };
 use deltalake_core::{arrow, parquet, DeltaOps};
 use deltalake_core::{operations::optimize::OptimizeType, DeltaTable, DeltaTableError};
-use parquet_key_management::datafusion::{KmsEncryptionFactory, KmsEncryptionFactoryOptions};
 use parquet_key_management::{
     crypto_factory::{CryptoFactory, DecryptionConfiguration, EncryptionConfiguration},
+    datafusion::{KmsEncryptionFactory, KmsEncryptionFactoryOptions},
     kms::KmsConnectionConfig,
-    test_kms::TestKmsClientFactory,
+    test_kms::TestAsyncKmsClientFactory,
 };
 use paste::paste;
 use std::{fs, sync::Arc};
@@ -294,7 +294,8 @@ fn plain_crypto_format_bad_decryptor() -> Result<FileFormatRef, DeltaTableError>
     create_plain_crypto_format(encryption_key.clone(), decryption_key.clone())
 }
 fn kms_crypto_format() -> Result<FileFormatRef, DeltaTableError> {
-    let crypto_factory = CryptoFactory::new(TestKmsClientFactory::with_default_keys());
+    let crypto_factory =
+        CryptoFactory::new_async_with_tokio(TestAsyncKmsClientFactory::with_default_keys());
 
     let kms_connection_config = Arc::new(KmsConnectionConfig::default());
     let encryption_factory = Arc::new(KmsEncryptionFactory::new(
@@ -404,13 +405,13 @@ macro_rules! encryption_tests {
                 $runner(file_format_options, false).await;
             }
 
-            #[tokio::test]
+            #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
             async fn [<$runner _kms>]() {
                 let file_format_options = kms_crypto_format().unwrap();
                 $runner(file_format_options, true).await;
             }
 
-            #[tokio::test]
+            #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
             #[should_panic(expected = "Failed to read encrypted table")]
             async fn [<$runner _kms_no_decryptor>]() {
                 let file_format_options = kms_crypto_format().unwrap();

crates/deltalake/Cargo.toml

@@ -59,7 +59,7 @@ tokio = { version = "1", features = ["macros", "rt-multi-thread"] }
 chrono = { workspace = true, default-features = false, features = ["clock"] }
 tracing = { workspace = true }
 tempfile = "3.21.0"
-parquet-key-management = { version = "0.4.1", features = ["_test_utils", "datafusion"] }
+parquet-key-management = { version = "0.5.0", features = ["_test_utils", "datafusion", "tokio"] }
 url = { workspace = true }
 
 [[example]]

crates/deltalake/examples/basic_operations_encryption.rs

@@ -22,11 +22,11 @@ use deltalake_core::{
     datafusion::common::test_util::format_batches, operations::optimize::OptimizeType, DeltaTable,
     DeltaTableError,
 };
-use parquet_key_management::datafusion::{KmsEncryptionFactory, KmsEncryptionFactoryOptions};
 use parquet_key_management::{
     crypto_factory::{CryptoFactory, DecryptionConfiguration, EncryptionConfiguration},
+    datafusion::{KmsEncryptionFactory, KmsEncryptionFactoryOptions},
     kms::KmsConnectionConfig,
-    test_kms::TestKmsClientFactory,
+    test_kms::TestAsyncKmsClientFactory,
 };
 use std::{fs, sync::Arc};
 use tempfile::TempDir;
@@ -291,8 +291,10 @@ fn plain_crypto_format() -> Result<FileFormatRef, DeltaTableError> {
     let file_format_options = Arc::new(SimpleFileFormatOptions::new(tbl_options)) as FileFormatRef;
     Ok(file_format_options)
 }
+
 fn kms_crypto_format() -> Result<FileFormatRef, DeltaTableError> {
-    let crypto_factory = CryptoFactory::new(TestKmsClientFactory::with_default_keys());
+    let crypto_factory =
+        CryptoFactory::new_async_with_tokio(TestAsyncKmsClientFactory::with_default_keys());
 
     let kms_connection_config = Arc::new(KmsConnectionConfig::default());
     let encryption_factory = Arc::new(KmsEncryptionFactory::new(
@@ -332,7 +334,7 @@ async fn round_trip_test(
     Ok(())
 }
 
-#[tokio::main(flavor = "current_thread")]
+#[tokio::main(flavor = "multi_thread", worker_threads = 1)]
 async fn main() -> Result<(), DeltaTableError> {
     println!("====================");
     println!("Begin Plain encryption test");