fixes for handling session context

Signed-off-by: bharath-techie <bharath78910@gmail.com>

Author: bharath-techie

sandbox/plugins/analytics-backend-datafusion/rust/src/api.rs

@@ -69,6 +69,10 @@ pub struct QueryStreamHandle {
     /// Held for its `Drop` impl — marks the query completed when the
     /// stream is closed.
     _query_tracking_context: QueryTrackingContext,
+    /// Keeps the SessionContext alive while the stream is being consumed.
+    /// The physical plan may reference state (e.g. RuntimeEnv, caches) owned
+    /// by the session; dropping it prematurely causes use-after-free.
+    _session_ctx: Option<datafusion::prelude::SessionContext>,
 }
 
 impl QueryStreamHandle {
@@ -79,6 +83,19 @@ impl QueryStreamHandle {
         Self {
             stream,
             _query_tracking_context: query_context,
+            _session_ctx: None,
+        }
+    }
+
+    pub fn with_session_context(
+        stream: RecordBatchStreamAdapter<CrossRtStream>,
+        query_context: QueryTrackingContext,
+        ctx: datafusion::prelude::SessionContext,
+    ) -> Self {
+        Self {
+            stream,
+            _query_tracking_context: query_context,
+            _session_ctx: Some(ctx),
         }
     }
 }

sandbox/plugins/analytics-backend-datafusion/rust/src/ffm.rs

@@ -624,28 +624,17 @@ pub unsafe extern "C" fn df_execute_with_context(
     plan_ptr: *const u8,
     plan_len: i64,
 ) -> i64 {
-    // Consume the session context handle on entry. Ownership transfers here
-    // regardless of whether the remainder of this function succeeds, returns an
-    // error via `?`, or panics — RAII (or `catch_unwind` drop-during-unwind)
-    // drops `session_handle` and frees the underlying SessionContext resources.
-    //
-    // This matches the Java-side contract: SessionContextHandle.markConsumed() is
-    // invoked in a `finally` after the FFM downcall, so every observable path from
-    // Java's perspective ("call.invoke ran") maps to "Rust consumed the handle".
-    // If we were to run fallible or panic-prone code (e.g. `get_rt_manager()?`)
-    // before Box::from_raw, the handle would leak on those paths.
     let session_handle = *Box::from_raw(session_ctx_ptr as *mut crate::session_context::SessionContextHandle);
 
     let mgr = get_rt_manager()?;
     let plan_bytes = slice::from_raw_parts(plan_ptr, plan_len as usize);
     let cpu_executor = mgr.cpu_executor();
 
-    // Route based on whether the session was configured for indexed execution
-    let handle_ref = &*(session_ctx_ptr as *const crate::session_context::SessionContextHandle);
-    if handle_ref.indexed_config.is_some() {
+    if session_handle.indexed_config.is_some() {
+        let ptr = Box::into_raw(Box::new(session_handle)) as i64;
         mgr.io_runtime
             .block_on(crate::indexed_executor::execute_indexed_with_context(
-                session_ctx_ptr,
+                ptr,
                 plan_bytes.to_vec(),
                 cpu_executor,
             ))

sandbox/plugins/analytics-backend-datafusion/rust/src/indexed_executor.rs

@@ -692,6 +692,6 @@ pub async unsafe fn execute_indexed_with_context(
     let cross_rt_stream = CrossRtStream::new_with_df_error_stream(df_stream, cpu_executor);
     let schema = cross_rt_stream.schema();
     let wrapped = RecordBatchStreamAdapter::new(schema, cross_rt_stream);
-    let stream_handle = crate::api::QueryStreamHandle::new(wrapped, query_context);
+    let stream_handle = crate::api::QueryStreamHandle::with_session_context(wrapped, query_context, ctx);
     Ok(Box::into_raw(Box::new(stream_handle)) as i64)
 }

sandbox/plugins/analytics-backend-datafusion/rust/src/query_executor.rs

@@ -185,6 +185,6 @@ pub async fn execute_with_context(
         cross_rt_stream,
     );
 
-    let stream_handle = crate::api::QueryStreamHandle::new(wrapped, handle.query_context);
+    let stream_handle = crate::api::QueryStreamHandle::with_session_context(wrapped, handle.query_context, handle.ctx);
     Ok(Box::into_raw(Box::new(stream_handle)) as i64)
 }