add sessionconfigext customization for endpoint create

Author: mach-kernel

ballista/core/src/execution_plans/distributed_query.rs

@@ -17,7 +17,9 @@
 
 use crate::client::BallistaClient;
 use crate::config::BallistaConfig;
-use crate::extension::BallistaGrpcMetadataInterceptor;
+use crate::extension::{
+    BallistaConfigGrpcEndpoint, BallistaGrpcMetadataInterceptor, SessionConfigExt,
+};
 use crate::serde::protobuf::SuccessfulJob;
 use crate::serde::protobuf::{
     execute_query_params::Query, execute_query_result, job_status,
@@ -50,6 +52,7 @@ use std::fmt::Debug;
 use std::marker::PhantomData;
 use std::sync::Arc;
 use std::time::Duration;
+use tonic::transport::Endpoint;
 
 /// This operator sends a logical plan to a Ballista scheduler for execution and
 /// polls the scheduler until the query is complete and then fetches the resulting
@@ -234,10 +237,11 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
         let metric_total_bytes =
             MetricBuilder::new(&self.metrics).counter("transferred_bytes", partition);
 
-        let interceptor = context
+        let interceptor = context.session_config().ballista_grpc_interceptor();
+
+        let customize_endpoint = context
             .session_config()
-            .get_extension::<BallistaGrpcMetadataInterceptor>()
-            .unwrap_or_default();
+            .ballista_override_create_grpc_client_endpoint();
 
         let stream = futures::stream::once(
             execute_query(
@@ -246,6 +250,7 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
                 query,
                 self.config.clone(),
                 interceptor,
+                customize_endpoint,
             )
             .map_err(|e| ArrowError::ExternalError(Box::new(e))),
         )
@@ -283,11 +288,20 @@ async fn execute_query(
     query: ExecuteQueryParams,
     config: BallistaConfig,
     grpc_interceptor: Arc<BallistaGrpcMetadataInterceptor>,
+    customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
 ) -> Result<impl Stream<Item = Result<RecordBatch>> + Send> {
     info!("Connecting to Ballista scheduler at {scheduler_url}");
     // TODO reuse the scheduler to avoid connecting to the Ballista scheduler again and again
-    let connection = create_grpc_client_endpoint(scheduler_url)
-        .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?
+    let mut endpoint = create_grpc_client_endpoint(scheduler_url)
+        .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
+
+    if let Some(customize) = customize_endpoint {
+        endpoint = customize
+            .configure_endpoint(endpoint)
+            .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;
+    }
+
+    let connection = endpoint
         .connect()
         .await
         .map_err(|e| DataFusionError::Execution(format!("{e:?}")))?;

ballista/core/src/extension.rs

@@ -30,10 +30,12 @@ use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use datafusion_proto::protobuf::LogicalPlanNode;
 use std::collections::HashMap;
+use std::error::Error;
 use std::sync::Arc;
 use tonic::codegen::http::HeaderName;
 use tonic::metadata::MetadataMap;
 use tonic::service::Interceptor;
+use tonic::transport::Endpoint;
 use tonic::{Request, Status};
 
 /// Provides methods which adapt [SessionState]
@@ -146,7 +148,24 @@ pub trait SessionConfigExt {
         prefer_flight: bool,
     ) -> Self;
 
+    /// Set user defined metadata keys in Ballista gRPC requests
     fn with_ballista_grpc_metadata(self, metadata: HashMap<String, String>) -> Self;
+
+    /// Get a `tonic` interceptor configured to decorate the provided metadata keys
+    fn ballista_grpc_interceptor(&self) -> Arc<BallistaGrpcMetadataInterceptor>;
+
+    fn with_ballista_override_create_grpc_client_endpoint(
+        self,
+        override_f: Arc<
+            dyn Fn(Endpoint) -> Result<Endpoint, Box<dyn Error + Send + Sync>>
+                + Send
+                + Sync,
+        >,
+    ) -> Self;
+
+    fn ballista_override_create_grpc_client_endpoint(
+        &self,
+    ) -> Option<Arc<BallistaConfigGrpcEndpoint>>;
 }
 
 /// [SessionConfigHelperExt] is set of [SessionConfig] extension methods
@@ -398,6 +417,29 @@ impl SessionConfigExt for SessionConfig {
         let extension = BallistaGrpcMetadataInterceptor::new(metadata);
         self.with_extension(Arc::new(extension))
     }
+
+    fn ballista_grpc_interceptor(&self) -> Arc<BallistaGrpcMetadataInterceptor> {
+        self.get_extension::<BallistaGrpcMetadataInterceptor>()
+            .unwrap_or_default()
+    }
+
+    fn with_ballista_override_create_grpc_client_endpoint(
+        self,
+        override_f: Arc<
+            dyn Fn(Endpoint) -> Result<Endpoint, Box<dyn Error + Send + Sync>>
+                + Send
+                + Sync,
+        >,
+    ) -> Self {
+        let extension = BallistaConfigGrpcEndpoint::new(override_f);
+        self.with_extension(Arc::new(extension))
+    }
+
+    fn ballista_override_create_grpc_client_endpoint(
+        &self,
+    ) -> Option<Arc<BallistaConfigGrpcEndpoint>> {
+        self.get_extension::<BallistaConfigGrpcEndpoint>()
+    }
 }
 
 impl SessionConfigHelperExt for SessionConfig {
@@ -570,6 +612,32 @@ impl Interceptor for BallistaGrpcMetadataInterceptor {
     }
 }
 
+#[derive(Clone)]
+pub struct BallistaConfigGrpcEndpoint {
+    override_f: Arc<
+        dyn Fn(Endpoint) -> Result<Endpoint, Box<dyn Error + Send + Sync>> + Send + Sync,
+    >,
+}
+
+impl BallistaConfigGrpcEndpoint {
+    pub fn new(
+        override_f: Arc<
+            dyn Fn(Endpoint) -> Result<Endpoint, Box<dyn Error + Send + Sync>>
+                + Send
+                + Sync,
+        >,
+    ) -> Self {
+        Self { override_f }
+    }
+
+    pub fn configure_endpoint(
+        &self,
+        endpoint: Endpoint,
+    ) -> Result<Endpoint, Box<dyn Error + Send + Sync>> {
+        (self.override_f)(endpoint)
+    }
+}
+
 #[cfg(test)]
 mod test {
     use datafusion::{