rok commented on code in PR #7286:
URL: https://github.com/apache/arrow-rs/pull/7286#discussion_r1997663765
##########
parquet/src/arrow/arrow_reader/mod.rs:
##########
@@ -706,30 +706,26 @@ impl<T: ChunkReader + 'static> Iterator for
ReaderPageIterator<T> {
.schema_descr()
.column(self.column_idx);
- if file_decryptor.is_column_encrypted(column_name.name()) {
- let data_decryptor =
file_decryptor.get_column_data_decryptor(column_name.name());
- let data_decryptor = match data_decryptor {
- Ok(data_decryptor) => data_decryptor,
- Err(err) => return Some(Err(err)),
- };
-
- let metadata_decryptor =
-
file_decryptor.get_column_metadata_decryptor(column_name.name());
- let metadata_decryptor = match metadata_decryptor {
- Ok(metadata_decryptor) => metadata_decryptor,
- Err(err) => return Some(Err(err)),
- };
-
- let crypto_context = CryptoContext::new(
- rg_idx,
- self.column_idx,
- data_decryptor,
- metadata_decryptor,
- file_decryptor.file_aad().clone(),
- );
- Some(Arc::new(crypto_context))
- } else {
- None
+ let crypto_metadata = self
+ .metadata
+ .row_group(rg_idx)
+ .column(self.column_idx)
+ .crypto_metadata();
+
+ match crypto_metadata {
+ Some(crypto_metadata) => {
+ match CryptoContext::for_column(
+ file_decryptor,
+ crypto_metadata,
+ column_name.name(),
+ rg_idx,
+ self.column_idx,
+ ) {
+ Ok(context) => Some(Arc::new(context)),
+ Err(err) => return Some(Err(err)),
+ }
+ }
+ None => None,
Review Comment:
This is much nicer now.
##########
parquet/src/encryption/decrypt.rs:
##########
@@ -136,125 +160,211 @@ impl CryptoContext {
}
}
-/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
-#[derive(Debug, Clone, PartialEq)]
-pub struct FileDecryptionProperties {
+#[derive(Clone, PartialEq)]
+struct ExplicitDecryptionKeys {
footer_key: Vec<u8>,
column_keys: HashMap<String, Vec<u8>>,
+}
+
+#[derive(Clone)]
+enum DecryptionKeys {
+ Explicit(ExplicitDecryptionKeys),
+ ViaRetriever(Arc<dyn KeyRetriever>),
+}
+
+/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
+#[derive(Clone)]
Review Comment:
```suggestion
impl PartialEq for DecryptionKeys {
fn eq(&self, other: &Self) -> bool {
match (self, other) {
(DecryptionKeys::Explicit(keys),
DecryptionKeys::Explicit(other_keys)) => {
keys.footer_key == other_keys.footer_key && keys.column_keys
== other_keys.column_keys
}
(DecryptionKeys::ViaRetriever(_),
DecryptionKeys::ViaRetriever(_)) => true,
_ => false,
}
}
}
/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
#[derive(Clone, PartialEq)]
pub struct FileDecryptionProperties {
```
##########
parquet/src/encryption/decrypt.rs:
##########
@@ -136,125 +160,211 @@ impl CryptoContext {
}
}
-/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
-#[derive(Debug, Clone, PartialEq)]
-pub struct FileDecryptionProperties {
+#[derive(Clone, PartialEq)]
+struct ExplicitDecryptionKeys {
footer_key: Vec<u8>,
column_keys: HashMap<String, Vec<u8>>,
+}
+
+#[derive(Clone)]
+enum DecryptionKeys {
+ Explicit(ExplicitDecryptionKeys),
+ ViaRetriever(Arc<dyn KeyRetriever>),
+}
+
+/// FileDecryptionProperties hold keys and AAD data required to decrypt a
Parquet file.
+#[derive(Clone)]
+pub struct FileDecryptionProperties {
+ keys: DecryptionKeys,
pub(crate) aad_prefix: Option<Vec<u8>>,
}
impl FileDecryptionProperties {
- /// Returns a new FileDecryptionProperties builder
+ /// Returns a new [`FileDecryptionProperties`] builder that will use the
provided key to
+ /// decrypt footer metadata.
pub fn builder(footer_key: Vec<u8>) -> DecryptionPropertiesBuilder {
DecryptionPropertiesBuilder::new(footer_key)
}
+
+ /// Returns a new [`FileDecryptionProperties`] builder that uses a
[`KeyRetriever`]
+ /// to get decryption keys based on key metadata.
+ pub fn with_key_retriever(key_retriever: Arc<dyn KeyRetriever>) ->
DecryptionPropertiesBuilder {
+ DecryptionPropertiesBuilder::new_with_key_retriever(key_retriever)
+ }
}
+impl std::fmt::Debug for FileDecryptionProperties {
+ fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+ write!(f, "FileDecryptionProperties {{ }}")
+ }
+}
+
+impl PartialEq for FileDecryptionProperties {
+ // FileDecryptionProperties needs to implement PartialEq to allow
+ // ParquetMetaData to implement PartialEq.
+ // We cannot compare a key retriever, but this isn't derived from the
metadata.
Review Comment:
I think implementing `PartialEq` is fine. Perhaps it'd be cleaner to
implement it on `DecryptionKeys` and derive it on `FileDecryptionProperties`?
See comment above.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
