feat(cluster): connector-driven object store registration on executors (spiceai#10414)

Author: phillipleblanc

Cargo.lock

@@ -38,6 +38,7 @@ snafu.workspace = true
 token_provider = { path = "../../token_provider" }
 tokio.workspace = true
 tracing.workspace = true
+url.workspace = true
 
 [features]
 default = []

crates/data-connectors/connector-databricks/Cargo.toml

@@ -35,6 +35,8 @@ use data_components::unity_catalog::{
 };
 use data_components::{Read, RefreshableCatalogProvider};
 use datafusion::datasource::TableProvider;
+use datafusion::datasource::listing::ListingTableUrl;
+use datafusion::execution::runtime_env::RuntimeEnv;
 use datafusion::sql::TableReference;
 use opentelemetry::KeyValue;
 use runtime::Runtime;
@@ -53,7 +55,6 @@ use runtime::token_providers::databricks::{
     AuthCredentials, DatabricksM2MTokenProvider, DatabricksU2MTokenProvider,
 };
 use runtime_secrets::get_params_with_secrets;
-#[cfg(feature = "spark")]
 use secrecy::ExposeSecret;
 use secrecy::SecretString;
 use snafu::prelude::*;
@@ -222,6 +223,16 @@ pub struct Databricks {
     /// Unity Catalog client for table type detection and permission checking.
     /// Present when the connector was created with enough information to call UC APIs.
     uc_client: Option<Arc<UnityCatalogClient>>,
+    /// Typed handle to the Delta read provider, present only in `delta_lake`
+    /// mode. Used by `register_object_stores` to resolve table storage
+    /// locations (which are only known after a UC round-trip) so the
+    /// underlying object store can be registered on the cluster executor's
+    /// runtime env.
+    delta_provider: Option<Arc<DatabricksDelta>>,
+    /// Original connector params, retained so `register_object_stores` can
+    /// build the storage URL fragment understood by `SpiceObjectStoreRegistry`.
+    /// Present only in `delta_lake` mode.
+    storage_params: Option<Parameters>,
 }
 
 impl std::fmt::Debug for Databricks {
@@ -313,9 +324,12 @@ impl Databricks {
                     initialization,
                     metrics,
                     uc_client,
+                    delta_provider: None,
+                    storage_params: None,
                 })
             }
             "delta_lake" => {
+                let storage_params = params.clone();
                 let storage_options = params.to_secret_map();
                 let token_provider: Arc<dyn TokenProvider> = match auth_credentials {
                     AuthCredentials::Token(token) => {
@@ -358,12 +372,15 @@ impl Databricks {
                     token_provider,
                     io_runtime,
                 );
+                let delta_provider = Arc::new(read_provider);
 
                 Ok(Self {
-                    read_provider: Arc::new(read_provider),
+                    read_provider: Arc::clone(&delta_provider) as Arc<dyn Read>,
                     initialization,
                     metrics: None,
                     uc_client,
+                    delta_provider: Some(delta_provider),
+                    storage_params: Some(storage_params),
                 })
             }
             #[cfg(feature = "spark")]
@@ -537,6 +554,8 @@ impl Databricks {
             initialization: ComponentInitialization::default(),
             metrics: None,
             uc_client: None,
+            delta_provider: None,
+            storage_params: None,
         })
     }
 
@@ -855,6 +874,78 @@ impl DataConnector for Databricks {
             }) as Arc<dyn MetricsProvider>
         })
     }
+
+    async fn register_object_stores(
+        &self,
+        dataset: &Dataset,
+        runtime_env: &Arc<RuntimeEnv>,
+    ) -> DataConnectorResult<()> {
+        // Only `delta_lake` mode produces object-store-backed scans on the
+        // executor. `sql_warehouse` and `spark_connect` execute on Databricks
+        // and surface as Flight/Arrow streams; nothing to register.
+        let (Some(delta), Some(params)) = (&self.delta_provider, &self.storage_params) else {
+            return Ok(());
+        };
+
+        // Resolve the underlying storage location via Unity Catalog. This is
+        // the bare URL (e.g. `s3://databricks-workspace-stack-bfa88-bucket/...`)
+        // that DataFusion will look up in `runtime_env().object_store(url)`
+        // when executing the decoded `ParquetSource` on the executor.
+        let table_reference = TableReference::from(dataset.path());
+        let storage_location =
+            delta
+                .resolve_table_uri(table_reference)
+                .await
+                .map_err(|source| DataConnectorError::UnableToConnectInternal {
+                    dataconnector: "databricks".to_string(),
+                    connector_component: ConnectorComponent::from(dataset),
+                    source,
+                })?;
+
+        let mut parsed = url::Url::parse(&storage_location).map_err(|source| {
+            DataConnectorError::UnableToConnectInternal {
+                dataconnector: "databricks".to_string(),
+                connector_component: ConnectorComponent::from(dataset),
+                source: Box::new(source),
+            }
+        })?;
+
+        // Encode the connector's storage params as the URL fragment so
+        // `SpiceObjectStoreRegistry::get_store` can build the right object
+        // store. `storage_registry_params` returns just the AWS/Azure/GCS
+        // entries with their prefixed names rewritten to the registry's
+        // canonical names; Databricks-internal params (`endpoint`, `token`)
+        // are excluded.
+        let mut fragment_builder = url::form_urlencoded::Serializer::new(String::new());
+        for (key, value) in params.storage_registry_params() {
+            fragment_builder.append_pair(&key, value.expose_secret());
+        }
+        parsed.set_fragment(Some(fragment_builder.finish().as_str()));
+
+        let listing_url = ListingTableUrl::parse(parsed).map_err(|source| {
+            DataConnectorError::UnableToConnectInternal {
+                dataconnector: "databricks".to_string(),
+                connector_component: ConnectorComponent::from(dataset),
+                source: Box::new(source),
+            }
+        })?;
+
+        runtime_env.object_store(&listing_url).map_err(|source| {
+            DataConnectorError::UnableToConnectInternal {
+                dataconnector: "databricks".to_string(),
+                connector_component: ConnectorComponent::from(dataset),
+                source: Box::new(source),
+            }
+        })?;
+
+        let mut redacted = <ListingTableUrl as AsRef<url::Url>>::as_ref(&listing_url).clone();
+        redacted.set_fragment(None);
+        tracing::debug!(
+            "Configured object storage for Databricks Dataset {} ({redacted})",
+            dataset.name,
+        );
+        Ok(())
+    }
 }
 
 /// Classifies a table-provider error, promoting Databricks-specific
@@ -1618,6 +1709,8 @@ mod tests {
                 UnityCatalogClient::new(Endpoint(endpoint), None, None)
                     .expect("mock Unity Catalog client should be created"),
             )),
+            delta_provider: None,
+            storage_params: None,
         };
         let dataset = make_dataset(dataset_from, "tpch_sf400_part").await;
 

crates/data-connectors/connector-databricks/src/lib.rs

@@ -358,6 +358,30 @@ impl Parameters {
 
         self.params = params.into_iter().collect();
     }
+
+    /// Returns the subset of params that map to `SpiceObjectStoreRegistry`
+    /// configuration keys, with the prefixed names (`aws_*`, `azure_storage_*`,
+    /// `google_*`) rewritten to the registry-facing names (`key`, `secret`,
+    /// `account`, etc.).
+    ///
+    /// Connector-internal params that aren't part of an object store mapping
+    /// (e.g. a Databricks workspace `endpoint`/`token`) are excluded, so this
+    /// is safe to encode directly into an object-store URL fragment.
+    #[must_use]
+    pub fn storage_registry_params(&self) -> Vec<(String, SecretString)> {
+        let map: HashMap<_, _> = self.params.iter().cloned().collect();
+        let mut out = Vec::new();
+        for (prefixed_key, registry_key) in AWS_PREFIXED_FRAGMENT_PARAMS
+            .iter()
+            .chain(AZURE_PREFIXED_FRAGMENT_PARAMS.iter())
+            .chain(GCS_PREFIXED_FRAGMENT_PARAMS.iter())
+        {
+            if let Some(value) = map.get(*prefixed_key) {
+                out.push(((*registry_key).to_string(), value.clone()));
+            }
+        }
+        out
+    }
 }
 
 #[derive(Clone)]

crates/runtime-parameters/src/lib.rs

@@ -21,8 +21,6 @@ use crate::cluster::partition::{
     scheduler_task::{PartitionManagementConfig, PartitionManagementTask},
 };
 use crate::config::{ClusterConfig, ClusterRole};
-use crate::dataconnector::listing;
-use crate::dataconnector::parameters::ConnectorParamsBuilder;
 use crate::jobs::JobExecutor;
 use crate::status::ComponentStatus;
 use crate::{
@@ -55,7 +53,6 @@ use ballista_scheduler::scheduler_server::SchedulerServer;
 use ballista_scheduler::state::execution_graph::RunningTaskInfo;
 use datafusion::codec::spice_logical_codec::SpiceLogicalCodec;
 use datafusion::codec::spice_physical_codec::SpicePhysicalCodec;
-use datafusion_datasource::ListingTableUrl;
 use datafusion_expr::Expr;
 use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
 use futures::future::try_join_all;
@@ -1803,8 +1800,15 @@ async fn executor_bind_app(
     Ok(())
 }
 
-/// Traverses dataset definitions and reifies `ListingTableUrl`s, triggering object store
-/// registration for each.
+/// For each registered dataset on the cluster executor, asks its data
+/// connector to register any object stores it needs against the executor's
+/// runtime env.
+///
+/// On the executor, decoded `ParquetSource` (and other file-source) plans
+/// arrive without their `parquet_file_reader_factory`, so `DataFusion` falls
+/// back to `runtime_env().object_store(url)`. This function gives each
+/// connector a chance to populate that registry using the dataset's
+/// already-secret-expanded params.
 async fn executor_bind_object_stores(rt: Arc<Runtime>) -> crate::Result<()> {
     let app = rt.app();
     let app = app.read().await;
@@ -1813,66 +1817,31 @@ async fn executor_bind_object_stores(rt: Arc<Runtime>) -> crate::Result<()> {
             source: "Runtime did not bind an App.".into(),
         });
     };
+    let runtime_env = rt.df.ctx.runtime_env();
     for dataset in Arc::clone(&rt).get_valid_datasets(app, LogErrors(true)) {
-        let mut params = ConnectorParamsBuilder::new(dataset.source().into(), (&dataset).into())
-            .build(Arc::clone(&rt.secrets), rt.tokio_io_runtime())
+        let connector = match Arc::clone(&rt)
+            .get_dataconnector_from_dataset(Arc::clone(&dataset))
             .await
-            .context(FailedToStartClusterExecutorSnafu)?;
-
-        // Either this is a URL with a scheme, or a URL with a connector name prefixing it
-        let url = match dataset.from.as_str().split_once(':') {
-            Some((_, rest)) if !rest.starts_with("//") => rest,
-            _ => dataset.from.as_str(),
-        };
-
-        let Ok(mut parsed) = Url::parse(url) else {
-            tracing::warn!("Unable to configure Dataset URL {}", url);
-            continue;
+        {
+            Ok(connector) => connector,
+            Err(error) => {
+                tracing::warn!(
+                    "Skipping object store registration for dataset {}: {error}",
+                    dataset.name
+                );
+                continue;
+            }
         };
 
-        if parsed.scheme() == "file" {
+        if let Err(error) = connector
+            .register_object_stores(&dataset, &runtime_env)
+            .await
+        {
             tracing::warn!(
-                "Dataset {} has a file:// scheme and may not be resolvable without a shared mount.",
+                "Failed to register object stores for dataset {}: {error}",
                 dataset.name
             );
-            continue;
         }
-
-        // Not all connectors have the same parameter structures for S3 -- this makes all fragment
-        // keys match the spec expected by the S3 connector and `SpiceObjectRegistry`.
-        params.parameters.canonicalize_s3_fragments();
-
-        // Canonicalize Azure parameters (e.g., `azure_storage_account_name` -> `account`)
-        // for Delta Lake and other connectors that use Azure-prefixed parameter names.
-        params.parameters.canonicalize_azure_fragments();
-
-        // Canonicalize GCS parameters (e.g., `google_service_account` -> `service_account`)
-        // for Delta Lake and other connectors that use GCS-prefixed parameter names.
-        params.parameters.canonicalize_gcs_fragments();
-
-        let unprefixed = params
-            .parameters
-            .into_iter()
-            .map(|(k, _)| k.as_str())
-            .collect::<Vec<_>>();
-
-        parsed.set_fragment(Some(
-            listing::build_fragments(&params.parameters, unprefixed).as_str(),
-        ));
-
-        let listing_table_url = ListingTableUrl::parse(parsed)
-            .boxed()
-            .context(FailedToStartClusterExecutorSnafu)?;
-
-        let _ = rt
-            .df
-            .ctx
-            .runtime_env()
-            .object_store(listing_table_url)
-            .boxed()
-            .context(FailedToStartClusterExecutorSnafu)?;
-
-        tracing::info!("Configured object storage for Dataset {}", dataset.name);
     }
 
     Ok(())

crates/runtime/src/cluster/mod.rs

@@ -1217,6 +1217,47 @@ impl<T: ListingTableConnector + Display> DataConnector for T {
         }
     }
 
+    async fn register_object_stores(
+        &self,
+        dataset: &Dataset,
+        runtime_env: &Arc<datafusion::execution::runtime_env::RuntimeEnv>,
+    ) -> DataConnectorResult<()> {
+        let url = self.get_object_store_url(dataset, None)?;
+        if url.scheme() == "file" {
+            tracing::warn!(
+                "Dataset {} has a file:// scheme and may not be resolvable on cluster executors without a shared mount.",
+                dataset.name
+            );
+            return Ok(());
+        }
+
+        let listing_url = ListingTableUrl::parse(url).boxed().context(
+            crate::dataconnector::UnableToConnectInternalSnafu {
+                dataconnector: format!("{self}"),
+                connector_component: ConnectorComponent::from(dataset),
+            },
+        )?;
+
+        // Triggers SpiceObjectStoreRegistry::get_store, which builds an object
+        // store from the URL fragment params (already secret-expanded by
+        // ConnectorParamsBuilder when the connector was created) and registers
+        // it on the runtime env keyed by the bare URL.
+        runtime_env.object_store(&listing_url).boxed().context(
+            crate::dataconnector::UnableToConnectInternalSnafu {
+                dataconnector: format!("{self}"),
+                connector_component: ConnectorComponent::from(dataset),
+            },
+        )?;
+
+        let mut redacted = <ListingTableUrl as AsRef<url::Url>>::as_ref(&listing_url).clone();
+        redacted.set_fragment(None);
+        tracing::debug!(
+            "Configured object storage for Dataset {} ({redacted})",
+            dataset.name,
+        );
+        Ok(())
+    }
+
     /// A hook that is called when an accelerated table is registered to the
     /// `DataFusion` context for this data connector.
     ///