diff --git a/parquet/src/encryption/decrypt.rs b/parquet/src/encryption/decrypt.rs index 0066523419de..bfa587277908 100644 --- a/parquet/src/encryption/decrypt.rs +++ b/parquet/src/encryption/decrypt.rs @@ -432,6 +432,14 @@ impl FileDecryptionProperties { } (column_names, column_keys) } + + /// Whether these decryption properties use a key retriever. + /// When false, explicit keys were provided up front and can + /// be retrieved without providing key metadata, rather than + /// resolved on demand. + 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 edd26f29619c..2ac94b8c3786 100644 --- a/parquet/tests/encryption/encryption.rs +++ b/parquet/tests/encryption/encryption.rs @@ -1474,3 +1474,36 @@ fn test_decrypt_page_index( Ok(()) } + +#[test] +fn test_decryption_properties_uses_key_retriever() { + let key_retriever = TestKeyRetriever::new() + .with_key( + AES_128_FOOTER_KEY_NAME.to_owned(), + AES_128_FOOTER_KEY.to_vec(), + ) + .with_key( + AES_128_KEY_NAMES[0].to_owned(), + AES_128_COLUMN_KEYS[0].to_vec(), + ); + + let properties_with_retriever = + FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)) + .build() + .unwrap(); + + assert!(properties_with_retriever.uses_key_retriever()); + + let properties_with_keys = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.to_vec()) + .with_column_key(AES_128_COLUMN_NAMES[0], AES_128_COLUMN_KEYS[0].to_vec()) + .build() + .unwrap(); + + assert!(!properties_with_keys.uses_key_retriever()); + + let uniform_properties = FileDecryptionProperties::builder(AES_128_FOOTER_KEY.to_vec()) + .build() + .unwrap(); + + assert!(!uniform_properties.uses_key_retriever()); +}