diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 0066523419de..ba0000bbdf0d 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -432,6 +432,12 @@ impl FileDecryptionProperties { } (column_names, column_keys) } + + /// Whether these decryption properties use a key retriever. + /// When false, keys can be retrieved without providing key metadata. + pub fn uses_key_retriever(&self) -> bool { + matches!(self.keys, DecryptionKeys::ViaRetriever(_)) + } } impl std::fmt::Debug for FileDecryptionProperties { diff --git a/parquet/tests/encryption/encryption.rs b/parquet/tests/encryption/encryption.rs index c26a0edee682..ace9a913f5ad 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -1262,3 +1262,29 @@ fn test_decrypt_page_index( Ok(()) } + +#[test] +fn test_decryption_properties_uses_key_retriever() { + let key_retriever = TestKeyRetriever::new() + .with_key("kf".to_owned(), "0123456789012345".as_bytes().to_vec()) + .with_key("kc1".to_owned(), "1234567890123450".as_bytes().to_vec()); + + let properties_with_retriever = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + assert!(properties_with_retriever.uses_key_retriever()); + + let footer_key = "0123456789012345".as_bytes(); + let column_1_key = "1234567890123450".as_bytes(); + let column_2_key = "1234567890123451".as_bytes(); + + let properties_with_keys = FileDecryptionProperties::builder(footer_key.to_vec()) + .with_column_key("double_field", column_1_key.to_vec()) + .with_column_key("float_field", column_2_key.to_vec()) + .build() + .unwrap(); + + assert!(!properties_with_keys.uses_key_retriever()); +}