hsiang-c commented on code in PR #2447:
URL: https://github.com/apache/datafusion-comet/pull/2447#discussion_r2373409278
##########
native/core/src/parquet/parquet_exec.rs:
##########
@@ -122,9 +145,131 @@ pub(crate) fn init_datasource_exec(
Ok(Arc::new(DataSourceExec::new(Arc::new(file_scan_config))))
}
+pub const ENCRYPTION_FACTORY_ID: &str = "comet.jni_kms_encryption";
+
+// Options used to configure our example encryption factory
+extensions_options! {
+ struct CometEncryptionConfig {
+ url_base: String, default = "file:///".into()
+ }
+}
+#[derive(Debug)]
+pub struct CometEncryptionFactory {
+ pub(crate) key_unwrapper: GlobalRef,
+}
+
+/// `EncryptionFactory` is a DataFusion trait for types that generate
+/// file encryption and decryption properties.
+#[async_trait]
+impl EncryptionFactory for CometEncryptionFactory {
+ async fn get_file_encryption_properties(
+ &self,
+ _options: &EncryptionFactoryOptions,
+ _schema: &SchemaRef,
+ _file_path: &Path,
+ ) -> Result<Option<FileEncryptionProperties>, DataFusionError> {
+ Err(DataFusionError::NotImplemented(
+ "Comet does not support Parquet encryption yet."
+ .parse()
+ .unwrap(),
+ ))
+ }
+
+ /// Generate file decryption properties to use when reading a Parquet file.
+ /// Rather than provide the AES keys directly for decryption, we set a
`KeyRetriever`
+ /// that can determine the keys using the encryption metadata.
+ async fn get_file_decryption_properties(
+ &self,
+ options: &EncryptionFactoryOptions,
+ file_path: &Path,
+ ) -> Result<Option<FileDecryptionProperties>, DataFusionError> {
+ let config: CometEncryptionConfig = options.to_extension_options()?;
+
+ let full_path: String = config.url_base + file_path.as_ref();
+ let key_retriever = CometKeyRetriever::new(&full_path,
self.key_unwrapper.clone())
+ .map_err(|e| DataFusionError::External(Box::new(e)))?;
+ let decryption_properties =
+
FileDecryptionProperties::with_key_retriever(Arc::new(key_retriever)).build()?;
+ Ok(Some(decryption_properties))
+ }
+}
+
+struct CometKeyRetriever {
+ file_path: String,
+ key_unwrapper: GlobalRef,
+ get_key_method_id: JMethodID,
+}
+
+impl CometKeyRetriever {
+ fn new(file_path: &str, key_unwrapper: GlobalRef) -> Result<Self,
ExecutionError> {
+ // Get JNI environment
+ let mut env = JVMClasses::get_env()?;
+
+ Ok(CometKeyRetriever {
+ file_path: file_path.to_string(),
+ key_unwrapper,
+ get_key_method_id: env
+ .get_method_id(
+ "org/apache/comet/parquet/CometFileKeyUnwrapper",
+ "getKey",
+ "(Ljava/lang/String;[B)[B",
Review Comment:
The return type of `getKey()` is `byte[]`?
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
