Add public methods to SchedulerServer

Author: thinkharderdev

ballista/scheduler/src/scheduler_server/mod.rs

@@ -20,7 +20,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use ballista_core::error::Result;
 use ballista_core::event_loop::{EventLoop, EventSender};
-use ballista_core::serde::protobuf::{StopExecutorParams, TaskStatus};
+use ballista_core::serde::protobuf::{JobStatus, StopExecutorParams, TaskStatus};
 use ballista_core::serde::{AsExecutionPlan, BallistaCodec};
 use ballista_core::utils::default_session_builder;
 
@@ -171,15 +171,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         Ok(())
     }
 
-    pub(crate) fn pending_tasks(&self) -> usize {
+    pub fn pending_tasks(&self) -> usize {
         self.query_stage_scheduler.pending_tasks()
     }
 
-    pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector {
-        self.query_stage_scheduler.metrics_collector()
-    }
-
-    pub(crate) async fn submit_job(
+    pub async fn submit_job(
         &self,
         job_id: &str,
         job_name: &str,
@@ -198,6 +194,14 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             .await
     }
 
+    pub async fn get_active_job_status(&self, job_id: &str) -> Result<Option<JobStatus>> {
+        self.state.task_manager.get_job_status(job_id).await
+    }
+
+    pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector {
+        self.query_stage_scheduler.metrics_collector()
+    }
+
     /// It just send task status update event to the channel,
     /// and will not guarantee the event processing completed after return
     pub(crate) async fn update_task_status(