refactor: SessionBuilder to return Result<_> (#1138)

* refactor: SessionBuilder to return Result<_> * Update ballista/core/src/utils.rs Co-authored-by: Andy Grove <[email protected]> --------- Co-authored-by: Andy Grove <[email protected]>
apache · Nov 27, 2024 · 020d29d · 020d29d
1 parent 683dede
commit 020d29d
Show file tree

Hide file tree

Showing 8 changed files with 33 additions and 27 deletions.
diff --git a/ballista/core/src/utils.rs b/ballista/core/src/utils.rs
@@ -62,12 +62,14 @@ use tonic::codegen::StdError;
 use tonic::transport::{Channel, Error, Server};
 
 /// Default session builder using the provided configuration
-pub fn default_session_builder(config: SessionConfig) -> SessionState {
-    SessionStateBuilder::new()
+pub fn default_session_builder(
+    config: SessionConfig,
+) -> datafusion::common::Result<SessionState> {
+    Ok(SessionStateBuilder::new()
         .with_default_features()
         .with_config(config)
-        .with_runtime_env(Arc::new(RuntimeEnv::new(RuntimeConfig::default()).unwrap()))
-        .build()
+        .with_runtime_env(Arc::new(RuntimeEnv::new(RuntimeConfig::default())?))
+        .build())
 }
 
 pub fn default_config_producer() -> SessionConfig {

diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs
@@ -408,7 +408,7 @@ impl JobState for InMemoryJobState {
         &self,
         config: &SessionConfig,
     ) -> Result<Arc<SessionContext>> {
-        let session = create_datafusion_context(config, self.session_builder.clone());
+        let session = create_datafusion_context(config, self.session_builder.clone())?;
         self.sessions.insert(session.session_id(), session.clone());
 
         Ok(session)
@@ -419,7 +419,7 @@ impl JobState for InMemoryJobState {
         session_id: &str,
         config: &SessionConfig,
     ) -> Result<Arc<SessionContext>> {
-        let session = create_datafusion_context(config, self.session_builder.clone());
+        let session = create_datafusion_context(config, self.session_builder.clone())?;
         self.sessions
             .insert(session_id.to_string(), session.clone());
 

diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs
@@ -56,7 +56,8 @@ mod external_scaler;
 mod grpc;
 pub(crate) mod query_stage_scheduler;
 
-pub type SessionBuilder = Arc<dyn Fn(SessionConfig) -> SessionState + Send + Sync>;
+pub type SessionBuilder =
+    Arc<dyn Fn(SessionConfig) -> datafusion::common::Result<SessionState> + Send + Sync>;
 
 #[derive(Clone)]
 pub struct SchedulerServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> {

diff --git a/ballista/scheduler/src/standalone.rs b/ballista/scheduler/src/standalone.rs
@@ -57,9 +57,11 @@ pub async fn new_standalone_scheduler_from_state(
     let session_config = session_state.config().clone();
     let session_state = session_state.clone();
     let session_builder = Arc::new(move |c: SessionConfig| {
-        SessionStateBuilder::new_from_existing(session_state.clone())
-            .with_config(c)
-            .build()
+        Ok(
+            SessionStateBuilder::new_from_existing(session_state.clone())
+                .with_config(c)
+                .build(),
+        )
     });
 
     let config_producer = Arc::new(move || session_config.clone());

diff --git a/ballista/scheduler/src/state/session_manager.rs b/ballista/scheduler/src/state/session_manager.rs
@@ -67,18 +67,18 @@ impl SessionManager {
 pub fn create_datafusion_context(
     session_config: &SessionConfig,
     session_builder: SessionBuilder,
-) -> Arc<SessionContext> {
+) -> datafusion::common::Result<Arc<SessionContext>> {
     let session_state = if session_config.round_robin_repartition() {
         let session_config = session_config
             .clone()
             // should we disable catalog on the scheduler side
             .with_round_robin_repartition(false);
 
         log::warn!("session manager will override `datafusion.optimizer.enable_round_robin_repartition` to `false` ");
-        session_builder(session_config)
+        session_builder(session_config)?
     } else {
-        session_builder(session_config.clone())
+        session_builder(session_config.clone())?
     };
 
-    Arc::new(SessionContext::new_with_state(session_state))
+    Ok(Arc::new(SessionContext::new_with_state(session_state)))
 }
diff --git a/examples/examples/custom-client.rs b/examples/examples/custom-client.rs
@@ -62,7 +62,7 @@ async fn main() -> Result<()> {
 
     // new sessions state with required custom session configuration and runtime environment
     let state =
-        custom_session_state_with_s3_support(custom_session_config_with_s3_options());
+        custom_session_state_with_s3_support(custom_session_config_with_s3_options())?;
 
     let ctx: SessionContext =
         SessionContext::remote_with_state("df://localhost:50050", state).await?;

diff --git a/examples/src/object_store.rs b/examples/src/object_store.rs
@@ -88,13 +88,13 @@ pub fn custom_runtime_env_with_s3_support(
 /// and [RuntimeEnv].
 pub fn custom_session_state_with_s3_support(
     session_config: SessionConfig,
-) -> SessionState {
-    let runtime_env = custom_runtime_env_with_s3_support(&session_config).unwrap();
+) -> datafusion::common::Result<SessionState> {
+    let runtime_env = custom_runtime_env_with_s3_support(&session_config)?;
 
-    SessionStateBuilder::new()
+    Ok(SessionStateBuilder::new()
         .with_runtime_env(runtime_env)
         .with_config(session_config)
-        .build()
+        .build())
 }
 
 /// Custom [ObjectStoreRegistry] which will create

diff --git a/examples/tests/object_store.rs b/examples/tests/object_store.rs
@@ -298,7 +298,7 @@ mod custom_s3_config {
         // object store registry.
 
         let session_builder = Arc::new(produce_state);
-        let state = session_builder(config_producer());
+        let state = session_builder(config_producer())?;
 
         // setting up ballista cluster with new runtime, configuration, and session state producers
         let (host, port) = crate::common::setup_test_cluster_with_builders(
@@ -416,7 +416,7 @@ mod custom_s3_config {
         // object store registry.
 
         let session_builder = Arc::new(produce_state);
-        let state = session_builder(config_producer());
+        let state = session_builder(config_producer())?;
 
         // // establishing cluster connection,
         let ctx: SessionContext = SessionContext::standalone_with_state(state).await?;
@@ -480,24 +480,25 @@ mod custom_s3_config {
         Ok(())
     }
 
-    fn produce_state(session_config: SessionConfig) -> SessionState {
+    fn produce_state(
+        session_config: SessionConfig,
+    ) -> datafusion::common::Result<SessionState> {
         let s3options = session_config
             .options()
             .extensions
             .get::<S3Options>()
             .ok_or(DataFusionError::Configuration(
                 "S3 Options not set".to_string(),
-            ))
-            .unwrap();
+            ))?;
 
         let config = RuntimeConfig::new().with_object_store_registry(Arc::new(
             CustomObjectStoreRegistry::new(s3options.clone()),
         ));
-        let runtime_env = RuntimeEnv::new(config).unwrap();
+        let runtime_env = RuntimeEnv::new(config)?;
 
-        SessionStateBuilder::new()
+        Ok(SessionStateBuilder::new()
             .with_runtime_env(runtime_env.into())
             .with_config(session_config)
-            .build()
+            .build())
     }
 }