Propagate credentials for distributed execution

Author: ccciudatu

crates/core/src/delta_datafusion/mod.rs

@@ -78,7 +78,7 @@ use crate::delta_datafusion::expr::parse_predicate_expression;
 use crate::delta_datafusion::schema_adapter::DeltaSchemaAdapterFactory;
 use crate::errors::{DeltaResult, DeltaTableError};
 use crate::kernel::{Add, DataCheck, EagerSnapshot, Invariant, Snapshot, StructTypeExt};
-use crate::logstore::LogStoreRef;
+use crate::logstore::{logstore_for, object_store_url, LogStoreRef};
 use crate::table::builder::ensure_table_uri;
 use crate::table::state::DeltaTableState;
 use crate::table::Constraint;
@@ -342,6 +342,7 @@ pub struct DeltaScanConfigBuilder {
     enable_parquet_pushdown: bool,
     /// Schema to scan table with
     schema: Option<SchemaRef>,
+    options: HashMap<String, String>,
 }
 
 impl Default for DeltaScanConfigBuilder {
@@ -352,6 +353,7 @@ impl Default for DeltaScanConfigBuilder {
             wrap_partition_values: None,
             enable_parquet_pushdown: true,
             schema: None,
+            options: HashMap::new(),
         }
     }
 }
@@ -396,6 +398,11 @@ impl DeltaScanConfigBuilder {
         self
     }
 
+    pub fn with_options(mut self, options: HashMap<String, String>) -> Self {
+        self.options = options;
+        self
+    }
+
     /// Build a DeltaScanConfig and ensure no column name conflicts occur during downstream processing
     pub fn build(&self, snapshot: &DeltaTableState) -> DeltaResult<DeltaScanConfig> {
         let file_column_name = if self.include_file_column {
@@ -438,6 +445,7 @@ impl DeltaScanConfigBuilder {
             wrap_partition_values: self.wrap_partition_values.unwrap_or(true),
             enable_parquet_pushdown: self.enable_parquet_pushdown,
             schema: self.schema.clone(),
+            options: self.options.clone(),
         })
     }
 }
@@ -453,6 +461,8 @@ pub struct DeltaScanConfig {
     pub enable_parquet_pushdown: bool,
     /// Schema to read as
     pub schema: Option<SchemaRef>,
+
+    pub options: HashMap<String, String>,
 }
 
 pub(crate) struct DeltaScanBuilder<'a> {
@@ -730,11 +740,14 @@ impl TableProvider for DeltaTable {
     ) -> DataFusionResult<Arc<dyn ExecutionPlan>> {
         register_store(self.log_store(), session.runtime_env().clone());
         let filter_expr = conjunction(filters.iter().cloned());
+        let mut config = DeltaScanConfig::default();
+        config.options = self.config.options.clone();
 
         let scan = DeltaScanBuilder::new(self.snapshot()?, self.log_store(), session)
             .with_projection(projection)
             .with_limit(limit)
             .with_filter(filter_expr)
+            .with_scan_config(config)
             .build()
             .await?;
 
@@ -753,12 +766,15 @@ impl TableProvider for DeltaTable {
 
         register_store(self.log_store(), session_state.runtime_env().clone());
         let filter_expr = conjunction(filters.iter().cloned());
+        let mut config = DeltaScanConfig::default();
+        config.options = self.config.options.clone();
 
         let scan = DeltaScanBuilder::new(self.snapshot()?, self.log_store(), session_state)
             .with_projection(projection)
             .with_projection_deep(projection_deep)
             .with_limit(limit)
             .with_filter(filter_expr)
+            .with_scan_config(config)
             .build()
             .await?;
 
@@ -970,6 +986,22 @@ impl ExecutionPlan for DeltaScan {
         partition: usize,
         context: Arc<TaskContext>,
     ) -> DataFusionResult<SendableRecordBatchStream> {
+        let source_uri = Url::parse(self.table_uri.as_str())
+            .map_err(|e| DataFusionError::External(Box::new(e)))?;
+        let url_key = object_store_url(&source_uri);
+        let runtime = context.runtime_env();
+        if runtime.object_store(url_key).is_err() {
+            let source_store = logstore_for(source_uri, self.config.options.clone(), None)
+                .map_err(|e| DataFusionError::External(Box::new(e)))?;
+            let object_store_url = source_store.object_store_url();
+
+            // check if delta store is already registered
+            if runtime.object_store(object_store_url.clone()).is_err() {
+                runtime
+                    .register_object_store(object_store_url.as_ref(), source_store.object_store());
+            }
+        }
+        // Now that everything is set up, execute the inner Delta
         self.parquet_scan.execute(partition, context)
     }
 

crates/core/src/table/builder.rs

@@ -82,6 +82,8 @@ pub struct DeltaTableConfig {
     #[serde(skip_serializing, skip_deserializing)]
     /// When a runtime handler is provided, all IO tasks are spawn in that handle
     pub io_runtime: Option<IORuntime>,
+
+    pub options: HashMap<String, String>
 }
 
 impl Default for DeltaTableConfig {
@@ -92,6 +94,7 @@ impl Default for DeltaTableConfig {
             log_buffer_size: num_cpus::get() * 4,
             log_batch_size: 1024,
             io_runtime: None,
+            options: HashMap::new(),
         }
     }
 }
@@ -314,7 +317,10 @@ impl DeltaTableBuilder {
     /// This will not load the log, i.e. the table is not initialized. To get an initialized
     /// table use the `load` function
     pub fn build(self) -> DeltaResult<DeltaTable> {
-        Ok(DeltaTable::new(self.build_storage()?, self.table_config))
+        let log_store = self.build_storage()?;
+        let mut config = self.table_config;
+        config.options = self.storage_options.unwrap_or_default();
+        Ok(DeltaTable::new(log_store, config))
     }
 
     /// Build the [`DeltaTable`] and load its state

crates/deltalake/src/lib.rs

@@ -1,3 +1,4 @@
+use std::sync::Once;
 /**
  * The deltalake crate is currently just a meta-package shim for deltalake-core
  */
@@ -11,3 +12,17 @@ pub use deltalake_azure as azure;
 pub use deltalake_gcp as gcp;
 #[cfg(feature = "hdfs")]
 pub use deltalake_hdfs as hdfs;
+
+pub fn ensure_initialized() {
+    static INIT: Once = Once::new();
+    INIT.call_once(|| {
+        #[cfg(any(feature = "s3", feature = "s3-native-tls"))]
+        aws::register_handlers(None);
+        #[cfg(feature = "azure")]
+        azure::register_handlers(None);
+        #[cfg(feature = "gcs")]
+        gcp::register_handlers(None);
+        #[cfg(feature = "hdfs")]
+        hdfs::register_handlers(None);
+    })
+}