xanderbailey commented on code in PR #2701:
URL: https://github.com/apache/iceberg-rust/pull/2701#discussion_r3668076694
##########
crates/integrations/datafusion/src/physical_plan/write.rs:
##########
@@ -601,4 +606,100 @@ mod tests {
Ok(())
}
+
+ #[tokio::test]
+ async fn test_iceberg_write_exec_encrypted() -> Result<()> {
+ let iceberg_catalog = get_iceberg_catalog().await;
+ let namespace = NamespaceIdent::new("test_namespace_enc".to_string());
+ iceberg_catalog
+ .create_namespace(&namespace, HashMap::new())
+ .await?;
+
+ // Create a table with encryption enabled via `encryption.key-id`.
+ let creation = TableCreation::builder()
+ .location(temp_path())
+ .name("test_table_enc".to_string())
+ .properties(HashMap::from([(
+ TableProperties::PROPERTY_ENCRYPTION_KEY_ID.to_string(),
+ "test-key".to_string(),
+ )]))
+ .schema(get_test_schema()?)
+ .build();
+ let table = iceberg_catalog.create_table(&namespace, creation).await?;
+
+ // Input data fed to the planner.
+ let arrow_schema = Arc::new(ArrowSchema::new(vec![
+ Field::new("id", DataType::Int32,
false).with_metadata(HashMap::from([(
+ PARQUET_FIELD_ID_META_KEY.to_string(),
+ "1".to_string(),
+ )])),
+ Field::new("name", DataType::Utf8,
false).with_metadata(HashMap::from([(
+ PARQUET_FIELD_ID_META_KEY.to_string(),
+ "2".to_string(),
+ )])),
+ ]));
+ let id_array = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef;
+ let name_array = Arc::new(StringArray::from(vec!["Alice", "Bob",
"Charlie"])) as ArrayRef;
+ let batch = RecordBatch::try_new(arrow_schema.clone(), vec![id_array,
name_array])
+ .map_err(|e| Error::new(ErrorKind::Unexpected, format!("record
batch: {e}")))?;
+
+ let input_plan = Arc::new(MockExecutionPlan::new(arrow_schema.clone(),
vec![batch]));
+ let write_exec = IcebergWriteExec::new(table.clone(), input_plan,
arrow_schema);
+
+ // Execute the planner and collect the returned (serialized) data
files.
+ let task_ctx = Arc::new(TaskContext::default());
+ let mut stream = write_exec
+ .execute(0, task_ctx)
+ .map_err(|e| Error::new(ErrorKind::Unexpected, format!("execute:
{e}")))?;
+ let mut results = vec![];
+ while let Some(batch) = stream.next().await {
+ results
+ .push(batch.map_err(|e| Error::new(ErrorKind::Unexpected,
format!("batch: {e}")))?);
+ }
+ assert_eq!(results.len(), 1, "expected one result batch");
+
+ let data_file_json = results[0]
+ .column(0)
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("expected StringArray")
+ .value(0);
+ let data_file = deserialize_data_file_from_json(
+ data_file_json,
+ table.metadata().default_partition_spec_id(),
+ table.metadata().default_partition_type(),
+ table.metadata().current_schema(),
+ )?;
+
+ // The planner encrypted the file.
Review Comment:
Not sure what this means
--
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]
