[arrow-ballista] branch main updated: added match arms to create ClusterStorageConfig (#766)

[ya...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

yangjiang pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-ballista.git


The following commit(s) were added to refs/heads/main by this push:
     new fce15a5a added match arms to create ClusterStorageConfig (#766)
fce15a5a is described below

commit fce15a5a9d9cc1a280ba4a7acc3ea54f478969d4
Author: Bokarev Nikita <57...@users.noreply.github.com>
AuthorDate: Sat May 6 10:12:11 2023 +0300

    added match arms to create ClusterStorageConfig (#766)
    
    * added match arms to create ClusterStorageConfig
    
    * removed extra itertools dependency
    
    * lint fix
---
 ballista/scheduler/src/bin/main.rs | 25 +++++++++++++++++++++----
 1 file changed, 21 insertions(+), 4 deletions(-)

diff --git a/ballista/scheduler/src/bin/main.rs b/ballista/scheduler/src/bin/main.rs
index 896313fb..44ff79d8 100644
--- a/ballista/scheduler/src/bin/main.rs
+++ b/ballista/scheduler/src/bin/main.rs
@@ -21,13 +21,13 @@ use std::{env, io};
 
 use anyhow::Result;
 
-use ballista_core::print_version;
-use ballista_scheduler::scheduler_process::start_server;
-
 use crate::config::{Config, ResultExt};
 use ballista_core::config::LogRotationPolicy;
+use ballista_core::print_version;
 use ballista_scheduler::cluster::BallistaCluster;
+use ballista_scheduler::cluster::ClusterStorage;
 use ballista_scheduler::config::{ClusterStorageConfig, SchedulerConfig};
+use ballista_scheduler::scheduler_process::start_server;
 use tracing_subscriber::EnvFilter;
 
 #[macro_use]
@@ -103,6 +103,23 @@ async fn main() -> Result<()> {
     let addr = format!("{}:{}", opt.bind_host, opt.bind_port);
     let addr = addr.parse()?;
 
+    let cluster_storage_config = match opt.cluster_backend {
+        ClusterStorage::Memory => ClusterStorageConfig::Memory,
+        ClusterStorage::Etcd => ClusterStorageConfig::Etcd(
+            opt.etcd_urls
+                .split_whitespace()
+                .map(|s| s.to_string())
+                .collect(),
+        ),
+        ClusterStorage::Sled => {
+            if opt.sled_dir.is_empty() {
+                ClusterStorageConfig::Sled(None)
+            } else {
+                ClusterStorageConfig::Sled(Some(opt.sled_dir))
+            }
+        }
+    };
+
     let config = SchedulerConfig {
         namespace: opt.namespace,
         external_host: opt.external_host,
@@ -115,7 +132,7 @@ async fn main() -> Result<()> {
         finished_job_state_clean_up_interval_seconds: opt
             .finished_job_state_clean_up_interval_seconds,
         advertise_flight_sql_endpoint: opt.advertise_flight_sql_endpoint,
-        cluster_storage: ClusterStorageConfig::Memory,
+        cluster_storage: cluster_storage_config,
         job_resubmit_interval_ms: (opt.job_resubmit_interval_ms > 0)
             .then_some(opt.job_resubmit_interval_ms),
         executor_termination_grace_period: opt.executor_termination_grace_period,