Public method for getting execution graph

Author: thinkharderdev

ballista/scheduler/src/scheduler_server/mod.rs

@@ -36,6 +36,7 @@ use log::{error, warn};
 use crate::scheduler_server::event::QueryStageSchedulerEvent;
 use crate::scheduler_server::query_stage_scheduler::QueryStageScheduler;
 use crate::state::backend::StateBackendClient;
+use crate::state::execution_graph::ExecutionGraph;
 use crate::state::executor_manager::{
     ExecutorManager, ExecutorReservation, DEFAULT_EXECUTOR_TIMEOUT_SECONDS,
 };
@@ -198,6 +199,16 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         self.state.task_manager.get_job_status(job_id).await
     }
 
+    pub async fn get_execution_graph(
+        &self,
+        job_id: &str,
+    ) -> Result<Option<Arc<ExecutionGraph>>> {
+        self.state
+            .task_manager
+            .get_job_execution_graph(job_id)
+            .await
+    }
+
     pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector {
         self.query_stage_scheduler.metrics_collector()
     }