Add TLS support to scheduler flight proxy service

- Update BallistaFlightProxyService to accept use_tls and customize_endpoint parameters
- Add use_tls field to SchedulerConfig with with_use_tls() builder method
- Unify EndpointOverrideFn type across crates to use ballista_core::extension definition
- Update flight proxy to use https/http scheme based on TLS configuration
- Apply custom endpoint configuration for TLS certificate setup

Author: phillipleblanc

ballista/executor/src/executor_process.rs

@@ -42,7 +42,7 @@ use datafusion::execution::runtime_env::RuntimeEnvBuilder;
 
 use ballista_core::config::{LogRotationPolicy, TaskSchedulingPolicy};
 use ballista_core::error::BallistaError;
-use ballista_core::extension::SessionConfigExt;
+use ballista_core::extension::{EndpointOverrideFn, SessionConfigExt};
 use ballista_core::serde::protobuf::executor_resource::Resource;
 use ballista_core::serde::protobuf::executor_status::Status;
 use ballista_core::serde::protobuf::{
@@ -57,11 +57,6 @@ use ballista_core::utils::{
     default_config_producer, get_time_before,
 };
 use ballista_core::{BALLISTA_VERSION, ConfigProducer, RuntimeProducer};
-use tonic::transport::{Endpoint, Error as TonicTransportError};
-
-/// Type alias for the endpoint override function used in gRPC client configuration
-pub type EndpointOverrideFn =
-    Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>;
 
 use crate::execution_engine::ExecutionEngine;
 use crate::executor::{Executor, TasksDrainedFuture};

ballista/executor/src/executor_server.rs

@@ -30,14 +30,11 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 use tokio::sync::mpsc;
 
 use log::{debug, error, info, warn};
-use tonic::transport::{Channel, Endpoint, Error as TonicTransportError};
+use tonic::transport::Channel;
 use tonic::{Request, Response, Status};
 
-/// Type alias for the endpoint override function used in gRPC client configuration
-pub type EndpointOverrideFn =
-    Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>;
-
 use ballista_core::error::BallistaError;
+use ballista_core::extension::EndpointOverrideFn;
 use ballista_core::serde::BallistaCodec;
 use ballista_core::serde::protobuf::{
     CancelTasksParams, CancelTasksResult, ExecutorMetric, ExecutorStatus,
@@ -276,7 +273,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> ExecutorServer<T,
             let mut endpoint = create_grpc_client_endpoint(scheduler_url, None)?;
 
             if let Some(ref override_fn) = self.override_create_grpc_client_endpoint {
-                endpoint = override_fn(endpoint)?;
+                endpoint = override_fn(endpoint).map_err(|e| {
+                    BallistaError::GrpcConnectionError(format!(
+                        "Failed to customize endpoint for scheduler {scheduler_id}: {e}"
+                    ))
+                })?;
             }
 
             let connection = endpoint.connect().await?;

ballista/scheduler/src/config.rs

@@ -27,16 +27,12 @@
 
 use crate::SessionBuilder;
 use crate::cluster::DistributionPolicy;
+use ballista_core::extension::EndpointOverrideFn;
 use ballista_core::{ConfigProducer, config::TaskSchedulingPolicy};
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use std::fmt::Display;
 use std::sync::Arc;
-use tonic::transport::{Endpoint, Error as TonicTransportError};
-
-/// Type alias for the endpoint override function used in gRPC client configuration
-pub type EndpointOverrideFn =
-    Arc<dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync>;
 
 /// Command-line configuration for the scheduler binary.
 #[cfg(feature = "build-binary")]
@@ -247,6 +243,8 @@ pub struct SchedulerConfig {
     pub override_physical_codec: Option<Arc<dyn PhysicalExtensionCodec>>,
     /// Override function for customizing gRPC client endpoints before they are used
     pub override_create_grpc_client_endpoint: Option<EndpointOverrideFn>,
+    /// Whether to use TLS when connecting to executors (for flight proxy)
+    pub use_tls: bool,
 }
 
 impl Default for SchedulerConfig {
@@ -274,6 +272,7 @@ impl Default for SchedulerConfig {
             override_logical_codec: None,
             override_physical_codec: None,
             override_create_grpc_client_endpoint: None,
+            use_tls: false,
         }
     }
 }
@@ -398,13 +397,17 @@ impl SchedulerConfig {
     /// This allows configuring TLS, timeouts, and other transport settings.
     pub fn with_override_create_grpc_client_endpoint(
         mut self,
-        override_fn: Arc<
-            dyn Fn(Endpoint) -> Result<Endpoint, TonicTransportError> + Send + Sync,
-        >,
+        override_fn: EndpointOverrideFn,
     ) -> Self {
         self.override_create_grpc_client_endpoint = Some(override_fn);
         self
     }
+
+    /// Sets whether TLS should be used when connecting to executors (for flight proxy).
+    pub fn with_use_tls(mut self, use_tls: bool) -> Self {
+        self.use_tls = use_tls;
+        self
+    }
 }
 
 /// Policy of distributing tasks to available executor slots
@@ -516,6 +519,7 @@ impl TryFrom<Config> for SchedulerConfig {
             override_physical_codec: None,
             override_session_builder: None,
             override_create_grpc_client_endpoint: None,
+            use_tls: false,
         };
 
         Ok(config)

ballista/scheduler/src/flight_proxy_service.rs

@@ -22,13 +22,15 @@ use arrow_flight::{
     HandshakeRequest, HandshakeResponse, PollInfo, PutResult, SchemaResult, Ticket,
 };
 use ballista_core::error::BallistaError;
+use ballista_core::extension::BallistaConfigGrpcEndpoint;
 use ballista_core::serde::decode_protobuf;
 use ballista_core::serde::scheduler::Action as BallistaAction;
-use ballista_core::utils::{GrpcClientConfig, create_grpc_client_connection};
+use ballista_core::utils::{GrpcClientConfig, create_grpc_client_endpoint};
 
 use futures::{Stream, TryFutureExt};
 use log::debug;
 use std::pin::Pin;
+use std::sync::Arc;
 use tonic::{Request, Response, Status, Streaming};
 
 /// Service implementing a proxy from scheduler to executor Apache Arrow Flight Protocol
@@ -40,16 +42,24 @@ use tonic::{Request, Response, Status, Streaming};
 pub struct BallistaFlightProxyService {
     max_decoding_message_size: usize,
     max_encoding_message_size: usize,
+    /// Whether to use TLS when connecting to executors
+    use_tls: bool,
+    /// Optional function to customize gRPC endpoint configuration (e.g., for TLS)
+    customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
 }
 
 impl BallistaFlightProxyService {
     pub fn new(
         max_decoding_message_size: usize,
         max_encoding_message_size: usize,
+        use_tls: bool,
+        customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
     ) -> Self {
         Self {
             max_decoding_message_size,
             max_encoding_message_size,
+            use_tls,
+            customize_endpoint,
         }
     }
 }
@@ -120,6 +130,8 @@ impl FlightService for BallistaFlightProxyService {
                     *port,
                     self.max_decoding_message_size,
                     self.max_encoding_message_size,
+                    self.use_tls,
+                    self.customize_endpoint.clone(),
                 )
                 .map_err(|e| from_ballista_err(&e))
                 .await?;
@@ -169,16 +181,34 @@ async fn get_flight_client(
     port: u16,
     max_decoding_message_size: usize,
     max_encoding_message_size: usize,
+    use_tls: bool,
+    customize_endpoint: Option<Arc<BallistaConfigGrpcEndpoint>>,
 ) -> Result<FlightServiceClient<tonic::transport::channel::Channel>, BallistaError> {
-    let addr = format!("http://{host}:{port}");
+    let scheme = if use_tls { "https" } else { "http" };
+    let addr = format!("{scheme}://{host}:{port}");
     let grpc_config = GrpcClientConfig::default();
-    let connection = create_grpc_client_connection(addr.clone(), &grpc_config)
-        .await
+
+    let mut endpoint = create_grpc_client_endpoint(addr.clone(), Some(&grpc_config))
         .map_err(|e| {
             BallistaError::GrpcConnectionError(format!(
-                "Error connecting to Ballista scheduler or executor at {addr}: {e:?}"
+                "Error creating endpoint for Ballista executor at {addr}: {e:?}"
+            ))
+        })?;
+
+    if let Some(ref customize) = customize_endpoint {
+        endpoint = customize.configure_endpoint(endpoint).map_err(|e| {
+            BallistaError::GrpcConnectionError(format!(
+                "Error customizing endpoint for Ballista executor at {addr}: {e}"
             ))
         })?;
+    }
+
+    let connection = endpoint.connect().await.map_err(|e| {
+        BallistaError::GrpcConnectionError(format!(
+            "Error connecting to Ballista executor at {addr}: {e:?}"
+        ))
+    })?;
+
     let flight_client = FlightServiceClient::new(connection)
         .max_decoding_message_size(max_decoding_message_size)
         .max_encoding_message_size(max_encoding_message_size);

ballista/scheduler/src/scheduler_process.rs

@@ -20,6 +20,7 @@ use crate::flight_proxy_service::BallistaFlightProxyService;
 use arrow_flight::flight_service_server::FlightServiceServer;
 use ballista_core::BALLISTA_VERSION;
 use ballista_core::error::BallistaError;
+use ballista_core::extension::BallistaConfigGrpcEndpoint;
 use ballista_core::serde::protobuf::scheduler_grpc_server::SchedulerGrpcServer;
 use ballista_core::serde::{
     BallistaCodec, BallistaLogicalExtensionCodec, BallistaPhysicalExtensionCodec,
@@ -102,9 +103,17 @@ pub async fn start_grpc_service<
     match &config.advertise_flight_sql_endpoint {
         Some(proxy) if proxy.is_empty() => {
             info!("Adding embedded flight proxy service on scheduler");
+            // Wrap the endpoint override function in BallistaConfigGrpcEndpoint
+            let customize_endpoint = config
+                .override_create_grpc_client_endpoint
+                .clone()
+                .map(|f| Arc::new(BallistaConfigGrpcEndpoint::new(f)));
+
             let flight_proxy = FlightServiceServer::new(BallistaFlightProxyService::new(
                 config.grpc_server_max_encoding_message_size as usize,
                 config.grpc_server_max_decoding_message_size as usize,
+                config.use_tls,
+                customize_endpoint,
             ))
             .max_decoding_message_size(
                 config.grpc_server_max_decoding_message_size as usize,

ballista/scheduler/src/state/executor_manager.rs

@@ -481,7 +481,11 @@ impl ExecutorManager {
             if let Some(ref override_fn) =
                 self.config.override_create_grpc_client_endpoint
             {
-                endpoint = override_fn(endpoint)?;
+                endpoint = override_fn(endpoint).map_err(|e| {
+                    BallistaError::GrpcConnectionError(format!(
+                        "Failed to customize endpoint for executor {executor_id}: {e}"
+                    ))
+                })?;
             }
 
             let connection = endpoint.connect().await?;