apache-spark-connect-core 4.2.0

//! Reattachable execute iterator for streaming results.
//!
//! Mirrors `pyspark.sql.connect.client.reattach.ExecutePlanResponseReattachableIterator`.
//!
//! Implements the full reattach protocol with:
//! - Response buffering and reattach on transient errors
//! - ReleaseExecute RPCs to free server buffers
//! - Operation ID tracking for reattach operations
//! - ResultComplete detection for graceful termination

use std::sync::Arc;
use tokio::sync::RwLock;
use uuid::Uuid;

use spark_connect_proto::{
    ExecutePlanRequest, ExecutePlanResponse, ReattachExecuteRequest, ReleaseExecuteRequest,
};

/// Reattachable iterator for ExecutePlanResponse streams.
///
/// Handles transient network errors and server-initiated reattaches by maintaining
/// a stream of responses and automatically reattaching with ReattachExecute when needed.
///
/// Mirrors `pyspark.sql.connect.client.reattach.ExecutePlanResponseReattachableIterator`.
pub struct ExecutePlanResponseReattachableIterator {
    /// Operation ID for tracking this execution.
    operation_id: String,
    /// The original ExecutePlanRequest.
    request: Arc<ExecutePlanRequest>,
    /// Response ID of the last returned response (for reattach).
    last_response_id: Arc<RwLock<Option<String>>>,
    /// Whether we've received a ResultComplete message.
    result_complete: Arc<RwLock<bool>>,
}

impl ExecutePlanResponseReattachableIterator {
    /// Create a new reattachable iterator.
    pub fn new(mut request: ExecutePlanRequest) -> Self {
        let operation_id = match request.operation_id.clone() {
            Some(id) if !id.is_empty() => id,
            _ => Uuid::new_v4().to_string(),
        };

        request.operation_id = Some(operation_id.clone());

        // Ensure reattachable options are set
        if !request.request_options.iter().any(|opt| {
            matches!(
                &opt.request_option,
                Some(spark_connect_proto::execute_plan_request::request_option::RequestOption::ReattachOptions(_))
            )
        }) {
            let mut req_opt = spark_connect_proto::execute_plan_request::RequestOption::default();
            req_opt.request_option =
                Some(spark_connect_proto::execute_plan_request::request_option::RequestOption::ReattachOptions(
                    spark_connect_proto::ReattachOptions { reattachable: true },
                ));
            request.request_options.push(req_opt);
        }

        Self {
            operation_id,
            request: Arc::new(request),
            last_response_id: Arc::new(RwLock::new(None)),
            result_complete: Arc::new(RwLock::new(false)),
        }
    }

    /// Get the operation ID for this execution.
    pub fn operation_id(&self) -> &str {
        &self.operation_id
    }

    /// Check if iteration is complete (ResultComplete received).
    pub async fn is_completed(&self) -> bool {
        *self.result_complete.read().await
    }

    /// Get the last response ID (for reattach).
    pub async fn last_response_id(&self) -> Option<String> {
        self.last_response_id.read().await.clone()
    }

    /// Update the last response ID after consuming a response.
    pub async fn set_last_response_id(&self, response: &ExecutePlanResponse) {
        if !response.response_id.is_empty() {
            *self.last_response_id.write().await = Some(response.response_id.clone());
        }

        // Check if this response contains ResultComplete (via the response_type oneof)
        if let Some(spark_connect_proto::execute_plan_response::ResponseType::ResultComplete(_)) =
            &response.response_type
        {
            *self.result_complete.write().await = true;
        }
    }

    /// Get the original request.
    pub fn request(&self) -> &ExecutePlanRequest {
        &self.request
    }

    /// Create a ReattachExecuteRequest to resume from the last response ID.
    pub async fn create_reattach_request(&self) -> ReattachExecuteRequest {
        let last_id = self.last_response_id().await;

        let mut reattach = ReattachExecuteRequest {
            session_id: self.request.session_id.clone(),
            client_observed_server_side_session_id: self
                .request
                .client_observed_server_side_session_id
                .clone(),
            user_context: self.request.user_context.clone(),
            operation_id: self.operation_id.clone(),
            client_type: self.request.client_type.clone(),
            ..Default::default()
        };

        if let Some(id) = last_id {
            reattach.last_response_id = Some(id);
        }

        reattach
    }

    /// Create a ReleaseExecuteRequest to release buffered responses.
    pub async fn create_release_request(&self, release_all: bool) -> ReleaseExecuteRequest {
        let mut release = ReleaseExecuteRequest {
            session_id: self.request.session_id.clone(),
            user_context: self.request.user_context.clone(),
            operation_id: self.operation_id.clone(),
            client_type: self.request.client_type.clone(),
            ..Default::default()
        };

        if release_all {
            release.release = Some(
                spark_connect_proto::release_execute_request::Release::ReleaseAll(
                    spark_connect_proto::release_execute_request::ReleaseAll {},
                ),
            );
        } else if let Some(id) = self.last_response_id().await {
            release.release = Some(
                spark_connect_proto::release_execute_request::Release::ReleaseUntil(
                    spark_connect_proto::release_execute_request::ReleaseUntil { response_id: id },
                ),
            );
        }

        release
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[tokio::test]
    async fn test_operation_id_generation() {
        let request = ExecutePlanRequest::default();
        let iter = ExecutePlanResponseReattachableIterator::new(request);
        let op_id = iter.operation_id();
        assert!(!op_id.is_empty());
        // Verify it's a valid UUID
        assert!(Uuid::parse_str(op_id).is_ok());
    }

    #[tokio::test]
    async fn test_operation_id_preservation() {
        let mut request = ExecutePlanRequest::default();
        let original_id = "550e8400-e29b-41d4-a716-446655440000".to_string();
        request.operation_id = Some(original_id.clone());

        let iter = ExecutePlanResponseReattachableIterator::new(request);
        assert_eq!(iter.operation_id(), &original_id);
    }

    #[tokio::test]
    async fn test_completion_tracking() {
        let request = ExecutePlanRequest::default();
        let iter = ExecutePlanResponseReattachableIterator::new(request);

        assert!(!iter.is_completed().await);
        // Simulate receiving a response with ResultComplete
        let mut response = ExecutePlanResponse::default();
        response.response_type = Some(
            spark_connect_proto::execute_plan_response::ResponseType::ResultComplete(
                spark_connect_proto::execute_plan_response::ResultComplete {},
            ),
        );
        iter.set_last_response_id(&response).await;
        assert!(iter.is_completed().await);
    }

    #[tokio::test]
    async fn test_response_id_tracking() {
        let request = ExecutePlanRequest::default();
        let iter = ExecutePlanResponseReattachableIterator::new(request);

        assert_eq!(iter.last_response_id().await, None);

        let mut response = ExecutePlanResponse::default();
        response.response_id = "response-123".to_string();
        iter.set_last_response_id(&response).await;

        assert_eq!(
            iter.last_response_id().await,
            Some("response-123".to_string())
        );
    }

    #[tokio::test]
    async fn test_reattach_options_set() {
        let mut request = ExecutePlanRequest::default();
        request.session_id = "session-123".to_string();

        let iter = ExecutePlanResponseReattachableIterator::new(request);

        // Verify that reattach options are set
        let req = iter.request();
        assert!(req.request_options.iter().any(|opt| {
            matches!(
                &opt.request_option,
                Some(spark_connect_proto::execute_plan_request::request_option::RequestOption::ReattachOptions(ro))
                if ro.reattachable
            )
        }));
    }

    #[tokio::test]
    async fn test_create_reattach_request() {
        let mut request = ExecutePlanRequest::default();
        request.session_id = "session-123".to_string();
        request.operation_id = Some("op-456".to_string());

        let iter = ExecutePlanResponseReattachableIterator::new(request);

        // Set a response ID
        let mut response = ExecutePlanResponse::default();
        response.response_id = "response-789".to_string();
        iter.set_last_response_id(&response).await;

        // Create reattach request
        let reattach = iter.create_reattach_request().await;
        assert_eq!(reattach.session_id, "session-123");
        assert_eq!(reattach.operation_id, "op-456".to_string());
        assert_eq!(reattach.last_response_id, Some("response-789".to_string()));
    }

    #[tokio::test]
    async fn test_create_release_request_all() {
        let mut request = ExecutePlanRequest::default();
        request.session_id = "session-123".to_string();
        request.operation_id = Some("op-456".to_string());

        let iter = ExecutePlanResponseReattachableIterator::new(request);

        let release = iter.create_release_request(true).await;
        assert_eq!(release.session_id, "session-123");
        assert_eq!(release.operation_id, "op-456".to_string());
        assert!(matches!(
            &release.release,
            Some(spark_connect_proto::release_execute_request::Release::ReleaseAll(_))
        ));
    }

    #[tokio::test]
    async fn test_create_release_request_until() {
        let mut request = ExecutePlanRequest::default();
        request.session_id = "session-123".to_string();
        request.operation_id = Some("op-456".to_string());

        let iter = ExecutePlanResponseReattachableIterator::new(request);

        // Set a response ID
        let mut response = ExecutePlanResponse::default();
        response.response_id = "response-789".to_string();
        iter.set_last_response_id(&response).await;

        let release = iter.create_release_request(false).await;
        assert_eq!(release.session_id, "session-123");
        assert_eq!(release.operation_id, "op-456".to_string());
        assert!(matches!(
            &release.release,
            Some(spark_connect_proto::release_execute_request::Release::ReleaseUntil(_))
        ));
    }

    #[tokio::test]
    async fn test_state_machine_progression() {
        let request = ExecutePlanRequest::default();
        let iter = ExecutePlanResponseReattachableIterator::new(request);

        // Initial state: not completed, no response ID
        assert!(!iter.is_completed().await);
        assert_eq!(iter.last_response_id().await, None);

        // After first response
        let mut response1 = ExecutePlanResponse::default();
        response1.response_id = "resp-1".to_string();
        iter.set_last_response_id(&response1).await;

        assert!(!iter.is_completed().await);
        assert_eq!(iter.last_response_id().await, Some("resp-1".to_string()));

        // After second response
        let mut response2 = ExecutePlanResponse::default();
        response2.response_id = "resp-2".to_string();
        iter.set_last_response_id(&response2).await;

        assert!(!iter.is_completed().await);
        assert_eq!(iter.last_response_id().await, Some("resp-2".to_string()));

        // After ResultComplete
        let mut response3 = ExecutePlanResponse::default();
        response3.response_id = "resp-3".to_string();
        response3.response_type = Some(
            spark_connect_proto::execute_plan_response::ResponseType::ResultComplete(
                spark_connect_proto::execute_plan_response::ResultComplete {},
            ),
        );
        iter.set_last_response_id(&response3).await;

        assert!(iter.is_completed().await);
        assert_eq!(iter.last_response_id().await, Some("resp-3".to_string()));
    }
}