Move TableEncryption into kms_encryption.rs

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

Author: corwinjoy

crates/core/src/operations/encryption.rs

@@ -59,8 +59,6 @@ pub mod constraints;
 #[cfg(feature = "datafusion")]
 pub mod delete;
 #[cfg(feature = "datafusion")]
-pub mod encryption;
-#[cfg(feature = "datafusion")]
 mod load;
 #[cfg(feature = "datafusion")]
 pub mod load_cdf;

crates/core/src/operations/mod.rs

@@ -1,11 +1,12 @@
-use crate::operations::encryption::TableEncryption;
 use crate::table::file_format_options::{
     FileFormatOptions, TableOptions, WriterPropertiesFactory, WriterPropertiesFactoryRef,
 };
 use crate::{crate_version, DeltaResult};
 use arrow_schema::Schema as ArrowSchema;
 use async_trait::async_trait;
 use datafusion::catalog::Session;
+use datafusion::config::{ConfigField, EncryptionFactoryOptions, ExtensionOptions};
+use datafusion::execution::parquet_encryption::EncryptionFactory;
 use object_store::path::Path;
 use parquet::basic::Compression;
 use parquet::file::properties::{WriterProperties, WriterPropertiesBuilder};
@@ -14,11 +15,71 @@ use std::fmt::{Debug, Formatter};
 use std::sync::Arc;
 use uuid::Uuid;
 
+pub type SchemaRef = Arc<ArrowSchema>;
+
+#[derive(Clone, Debug)]
+pub struct TableEncryption {
+    encryption_factory: Arc<dyn EncryptionFactory>,
+    configuration: EncryptionFactoryOptions,
+}
+
+impl TableEncryption {
+    pub fn new(
+        encryption_factory: Arc<dyn EncryptionFactory>,
+        configuration: EncryptionFactoryOptions,
+    ) -> Self {
+        Self {
+            encryption_factory,
+            configuration,
+        }
+    }
+
+    pub fn new_with_extension_options<T: ExtensionOptions>(
+        encryption_factory: Arc<dyn EncryptionFactory>,
+        options: &T,
+    ) -> DeltaResult<Self> {
+        let mut configuration = EncryptionFactoryOptions::default();
+        for entry in options.entries() {
+            if let Some(value) = &entry.value {
+                configuration.set(&entry.key, value)?;
+            }
+        }
+        Ok(Self {
+            encryption_factory,
+            configuration,
+        })
+    }
+
+    pub fn encryption_factory(&self) -> &Arc<dyn EncryptionFactory> {
+        &self.encryption_factory
+    }
+
+    pub fn configuration(&self) -> &EncryptionFactoryOptions {
+        &self.configuration
+    }
+
+    pub async fn update_writer_properties(
+        &self,
+        mut builder: WriterPropertiesBuilder,
+        file_path: &Path,
+        file_schema: &SchemaRef,
+    ) -> DeltaResult<WriterPropertiesBuilder> {
+        let encryption_properties = self
+            .encryption_factory
+            .get_file_encryption_properties(&self.configuration, file_schema, file_path)
+            .await?;
+        if let Some(encryption_properties) = encryption_properties {
+            builder = builder.with_file_encryption_properties(encryption_properties);
+        }
+        Ok(builder)
+    }
+}
+
 // More advanced factory with KMS support
 #[derive(Clone, Debug)]
 pub struct KMSWriterPropertiesFactory {
     writer_properties: WriterProperties,
-    encryption: Option<crate::operations::encryption::TableEncryption>,
+    encryption: Option<TableEncryption>,
 }
 
 impl KMSWriterPropertiesFactory {

crates/core/src/test_utils/kms_encryption.rs

@@ -12,10 +12,9 @@ use datafusion::{
 };
 use deltalake_core::kernel::{DataType, PrimitiveType, StructField};
 use deltalake_core::operations::collect_sendable_stream;
-use deltalake_core::operations::encryption::TableEncryption;
 use deltalake_core::parquet::encryption::decrypt::FileDecryptionProperties;
 use deltalake_core::table::file_format_options::{FileFormatRef, SimpleFileFormatOptions};
-use deltalake_core::test_utils::kms_encryption::KmsFileFormatOptions;
+use deltalake_core::test_utils::kms_encryption::{KmsFileFormatOptions, TableEncryption};
 use deltalake_core::{arrow, parquet, DeltaOps};
 use deltalake_core::{operations::optimize::OptimizeType, DeltaTable, DeltaTableError};
 use parquet_key_management::datafusion::{KmsEncryptionFactory, KmsEncryptionFactoryOptions};

crates/core/tests/commands_with_encryption.rs

@@ -14,9 +14,8 @@ use deltalake::kernel::{DataType, PrimitiveType, StructField};
 use deltalake::operations::collect_sendable_stream;
 use deltalake::parquet::encryption::decrypt::FileDecryptionProperties;
 use deltalake::{arrow, parquet, DeltaOps};
-use deltalake_core::operations::encryption::TableEncryption;
 use deltalake_core::table::file_format_options::{FileFormatRef, SimpleFileFormatOptions};
-use deltalake_core::test_utils::kms_encryption::KmsFileFormatOptions;
+use deltalake_core::test_utils::kms_encryption::{KmsFileFormatOptions, TableEncryption};
 use deltalake_core::{
     datafusion::common::test_util::format_batches, operations::optimize::OptimizeType, DeltaTable,
     DeltaTableError,