expose operation_id and make session_id required

Author: milenkovicm

Cargo.lock

@@ -60,6 +60,7 @@ tokio = { workspace = true }
 tokio-stream = { workspace = true, features = ["net"] }
 tonic = { workspace = true }
 url = { workspace = true }
+uuid = { workspace = true }
 
 [dev-dependencies]
 tempfile = { workspace = true }

ballista/core/Cargo.toml

@@ -524,8 +524,12 @@ message ExecuteQueryParams {
     string sql = 2 [deprecated=true]; // I'd suggest to remove this, if SQL needed use `flight-sql`
   }
 
-  optional string session_id = 3;
+  string session_id = 3;
   repeated KeyValuePair settings = 4;
+  // operation_id is unique number for each request 
+  // client makes. it helps mapping requests between 
+  // client and scheduler
+  string operation_id = 5;
 }
 
 message CreateUpdateSessionParams {
@@ -563,6 +567,7 @@ message ExecuteQueryResult {
     ExecuteQuerySuccessResult success = 1;
     ExecuteQueryFailureResult failure = 2;
   }
+  string operation_id = 3;
 }
 
 message ExecuteQuerySuccessResult {

ballista/core/proto/ballista.proto

@@ -39,7 +39,7 @@ use datafusion_proto::logical_plan::{
     AsLogicalPlan, DefaultLogicalExtensionCodec, LogicalExtensionCodec,
 };
 use futures::{Stream, StreamExt, TryFutureExt, TryStreamExt};
-use log::{error, info};
+use log::{debug, error, info};
 use std::any::Any;
 use std::fmt::Debug;
 use std::marker::PhantomData;
@@ -202,11 +202,16 @@ impl<T: 'static + AsLogicalPlan> ExecutionPlan for DistributedQueryExec<T> {
                 },
             )
             .collect();
-
+        let operation_id = uuid::Uuid::now_v7().to_string();
+        debug!(
+            "Distributed query with session_id: {}, execution operation_id: {}",
+            self.session_id, operation_id
+        );
         let query = ExecuteQueryParams {
             query: Some(Query::LogicalPlan(buf)),
             settings,
-            session_id: Some(self.session_id.clone()),
+            session_id: self.session_id.clone(),
+            operation_id,
         };
 
         let stream = futures::stream::once(

ballista/core/src/execution_plans/distributed_query.rs

@@ -786,10 +786,12 @@ pub struct UpdateTaskStatusResult {
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ExecuteQueryParams {
-    #[prost(string, optional, tag = "3")]
-    pub session_id: ::core::option::Option<::prost::alloc::string::String>,
+    #[prost(string, tag = "3")]
+    pub session_id: ::prost::alloc::string::String,
     #[prost(message, repeated, tag = "4")]
     pub settings: ::prost::alloc::vec::Vec<KeyValuePair>,
+    #[prost(string, tag = "5")]
+    pub operation_id: ::prost::alloc::string::String,
     #[prost(oneof = "execute_query_params::Query", tags = "1, 2")]
     pub query: ::core::option::Option<execute_query_params::Query>,
 }
@@ -845,6 +847,8 @@ pub struct ExecuteSqlParams {
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct ExecuteQueryResult {
+    #[prost(string, tag = "3")]
+    pub operation_id: ::prost::alloc::string::String,
     #[prost(oneof = "execute_query_result::Result", tags = "1, 2")]
     pub result: ::core::option::Option<execute_query_result::Result>,
 }

ballista/core/src/serde/generated/ballista.rs

@@ -313,6 +313,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
         if let ExecuteQueryParams {
             query: Some(query),
             session_id,
+            operation_id,
             settings,
         } = query_params
         {
@@ -322,35 +323,26 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                 .and_then(|s| s.value.clone())
                 .unwrap_or_default();
 
-            let (session_id, session_ctx) = match session_id {
-                Some(session_id) => {
-                    let session_config = self.state.session_manager.produce_config();
-                    let session_config =
-                        session_config.update_from_key_value_pair(&settings);
+            let job_id = self.state.task_manager.generate_job_id();
 
-                    let ctx = self
-                        .state
-                        .session_manager
-                        .create_or_update_session(&session_id, &session_config)
-                        .await
-                        .map_err(|e| {
-                            Status::internal(format!(
-                                "Failed to create SessionContext: {e:?}"
-                            ))
-                        })?;
+            info!("execution query - session_id: {}, operation_id: {}, job_name: {}, job_id: {}", session_id, operation_id, job_name, job_id);
 
-                    (session_id, ctx)
-                }
-                _ => {
-                    error!("Client should set session_id");
-                    return Ok(Response::new(ExecuteQueryResult {
-                                result: Some(execute_query_result::Result::Failure(
-                                    ExecuteQueryFailureResult {
-                                        failure: Some(execute_query_failure_result::Failure::SessionNotFound("Client should set session_id".to_string())),
-                                    },
-                                )),
-                            }));
-                }
+            let (session_id, session_ctx) = {
+                let session_config = self.state.session_manager.produce_config();
+                let session_config = session_config.update_from_key_value_pair(&settings);
+
+                let ctx = self
+                    .state
+                    .session_manager
+                    .create_or_update_session(&session_id, &session_config)
+                    .await
+                    .map_err(|e| {
+                        Status::internal(format!(
+                            "Failed to create SessionContext: {e:?}"
+                        ))
+                    })?;
+
+                (session_id, ctx)
             };
 
             let plan = match query {
@@ -367,6 +359,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                                 format!("Could not parse logical plan protobuf: {e}");
                             error!("{}", msg);
                             return Ok(Response::new(ExecuteQueryResult {
+                                operation_id,
                                 result: Some(execute_query_result::Result::Failure(
                                     ExecuteQueryFailureResult {
                                         failure: Some(execute_query_failure_result::Failure::PlanParsingFailure(msg)),
@@ -387,6 +380,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                             let msg = format!("Error parsing SQL: {e}");
                             error!("{}", msg);
                             return Ok(Response::new(ExecuteQueryResult {
+                                operation_id,
                                 result: Some(execute_query_result::Result::Failure(
                                     ExecuteQueryFailureResult {
                                         failure: Some(execute_query_failure_result::Failure::PlanParsingFailure(msg)),
@@ -403,9 +397,6 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                 plan.display_indent()
             );
 
-            let job_id = self.state.task_manager.generate_job_id();
-
-            log::trace!("setting job name: {}", job_name);
             self.submit_job(&job_id, &job_name, session_ctx, &plan)
                 .await
                 .map_err(|e| {
@@ -417,6 +408,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerGrpc
                 })?;
 
             Ok(Response::new(ExecuteQueryResult {
+                operation_id,
                 result: Some(execute_query_result::Result::Success(
                     ExecuteQuerySuccessResult { job_id, session_id },
                 )),