spark-connect 0.2.2

//! Internal gRPC client abstraction for the Spark Connect protocol.
//!
//! This module defines [`SparkClient`], the low-level asynchronous client that
//! manages communication with a Spark Connect server over gRPC.
//!
//! <div class="warning">
//! 
//! End users are advised <b>not</b> to construct or use `SparkClient` directly — use
//! [`SparkSession`](crate::SparkSession) instead, which provides a high-level API.
//! 
//! </div>
//!
//! # Overview
//!
//! `SparkClient` wraps the generated [`SparkConnectServiceClient`] and provides:
//!
//! - Connection setup (via [`ChannelBuilder`]);
//! - Metadata injection (via [`HeaderInterceptor`]);
//! - Structured request/response handling for:
//!   - [`Analyze` plan](crate::spark::analyze_plan_request::Analyze);
//!   - [Execution `Plan`](crate::spark::Plan);
//!   - [`Interrupt`](crate::spark::interrupt_request::Interrupt);
//!   - Reattach/release semantics.
//!
//! Each call validates the active Spark session and maps server responses into
//! safe Rust types or [`SparkError`] values.
mod builder;
mod error;
mod handlers;
mod middleware;

pub use self::builder::ChannelBuilder;
pub(crate) use self::error::{ClientError, ClientErrorKind};
pub use self::middleware::HeaderInterceptor;
use self::handlers::{AnalyzeHandler, ExecuteHandler, InterruptHandler};
use crate::spark;
use crate::spark::spark_connect_service_client::SparkConnectServiceClient;
use crate::spark::execute_plan_response::ResponseType;

use arrow::array::RecordBatch;
use std::sync::Arc;
use tokio::sync::RwLock;
use tonic::codec::Streaming;
use tonic::transport::Channel;
use uuid;

/// Utility type alias for a gRPC channel with an attached interceptor.
type InterceptedChannel = tonic::service::interceptor::InterceptedService<Channel, HeaderInterceptor>;

/// Asynchronous gRPC client for Spark Connect.
///
/// `SparkClient` manages RPC calls, session validation, and response
/// interpretation. It is used internally by [`SparkSession`](crate::SparkSession)
/// to execute plans and perform analysis or interrupt operations.
///
/// <div class="warning">
/// This struct is <b>not</b> intended for direct use; it exposes low-level details
/// of the Spark Connect protocol.
/// </div>
///
/// # Lifecycle
///
/// - Constructed indirectly through [`SparkSessionBuilder`](crate::SparkSessionBuilder);
/// - Maintains session context (e.g. `session_id`, `user_context`);
/// - Automatically attaches metadata headers.
#[derive(Clone, Debug)]
pub struct SparkClient {
    pub(crate) builder: ChannelBuilder,
    stub: Arc<RwLock<SparkConnectServiceClient<InterceptedChannel>>>,
    user_context: Option<spark::UserContext>,
    use_reattachable_execute: bool,
    session_id: String,
    operation_id: Option<String>,
    response_id: Option<String>,
    handler_analyze: AnalyzeHandler,
    handler_execute: ExecuteHandler,
    handler_interrupt: InterruptHandler,
}

impl SparkClient {
    /// Creates a new client from a gRPC stub and a configured [`ChannelBuilder`].
    ///
    /// Typically called internally by [`SparkSessionBuilder`](crate::SparkSessionBuilder).
    pub(crate) fn new(
        stub: Arc<RwLock<SparkConnectServiceClient<InterceptedChannel>>>,
        builder: ChannelBuilder,
    ) -> Self {
        let user_ref = builder.user_id.clone().unwrap_or("".to_string());
        let session_id = builder.session_id.to_string();

        Self {
            stub,
            builder,
            user_context: Some(spark::UserContext {
                user_id: user_ref.clone(),
                user_name: user_ref,
                extensions: vec![],
            }),
            session_id,
            operation_id: None,
            response_id: None,
            handler_analyze: AnalyzeHandler::default(),
            handler_execute: ExecuteHandler::default(),
            handler_interrupt: InterruptHandler::default(),
            use_reattachable_execute: true,
        }
    }

    /// Returns the session ID associated with this client.
    pub(crate) fn session_id(&self) -> String {
        self.session_id.to_string()
    }

    /// Returns the Spark version obtained from the last analyze request.
    pub(crate) fn spark_version(&self) -> Result<String, ClientError> {
        self.handler_analyze
            .spark_version
            .to_owned()
            .ok_or_else(|| ClientError::new(ClientErrorKind::AnalyzeResponseNotFound(
                "Spark version response is empty".to_string()
            )))
    }

    /// Returns the list of operation IDs that were interrupted.
    pub(crate) fn interrupted_ids(&self) -> Vec<String> {
        self.handler_interrupt.interrupted_ids.to_owned()
    }

    /// Returns the last relation received in an [`ExecutePlanResponse`](crate::spark::ExecutePlanResponse).
    pub(crate) fn relation(&self) -> Result<spark::Relation, ClientError> {
        self.handler_execute
            .relation
            .to_owned()
            .ok_or_else(||  ClientError::new(ClientErrorKind::AnalyzeResponseNotFound(
                "relation response is empty".to_string()
            )))
    }

    /// Returns all record batches accumulated during the last execution.
    pub(crate) fn batches(&self) -> Vec<RecordBatch> {
        self.handler_execute.batches.to_owned()
    }

    /// Compares a session_id to the current session's id.
    fn validate_session(&self, session_id: &str) -> Result<(), ClientError> {
        if self.session_id() != session_id {
            return Err(ClientError::new(ClientErrorKind::SessionIDMismatch {
                client_session_id: self.builder.session_id.to_string(),
                request_session_id: session_id.to_string()
            }));
        }
        Ok(())
    }

    /// Sends an [`AnalyzePlanRequest`](crate::spark::AnalyzePlanRequest)
    /// to the Spark Connect server and updates the internal analysis handler.
    pub(crate) async fn analyze(
        &mut self,
        analyze: spark::analyze_plan_request::Analyze,
    ) -> Result<&mut Self, ClientError> {
        let request = spark::AnalyzePlanRequest {
            session_id: self.session_id(),
            user_context: self.user_context.clone(),
            client_type: self.builder.user_agent.clone(),
            analyze: Some(analyze),
        };
        
        let mut client = self.stub.write().await;
        let resp = client.analyze_plan(request.clone())
            .await
            .map_err(|status| {
                ClientError::new(ClientErrorKind::AnalyzeRequest { status, request })
            })?
            .into_inner();
        drop(client);
        
        self.handle_analyze_response(resp)?;
        
        Ok(self)
    }

    fn handle_analyze_response(
        &mut self,
        resp: spark::AnalyzePlanResponse,
    ) -> Result<(), ClientError> {
        self.validate_session(&resp.session_id)?;

        // clear out any prior responses
        self.handler_analyze = AnalyzeHandler::default();
        
        if let Some(result) = resp.result {
            match result {
                spark::analyze_plan_response::Result::Schema(schema) => {
                    self.handler_analyze.schema = schema.schema
                }
                // spark::analyze_plan_response::Result::Explain(explain) => {
                //     self.handler_analyze.explain = Some(explain.explain_string)
                // }
                // spark::analyze_plan_response::Result::TreeString(tree_string) => {
                //     self.handler_analyze.tree_string = Some(tree_string.tree_string)
                // }
                // spark::analyze_plan_response::Result::IsLocal(is_local) => {
                //     self.handler_analyze.is_local = Some(is_local.is_local)
                // }
                // spark::analyze_plan_response::Result::IsStreaming(is_streaming) => {
                //     self.handler_analyze.is_streaming = Some(is_streaming.is_streaming)
                // }
                // spark::analyze_plan_response::Result::InputFiles(input_files) => {
                //     self.handler_analyze.input_files = Some(input_files.files)
                // }
                spark::analyze_plan_response::Result::SparkVersion(spark_version) => {
                    self.handler_analyze.spark_version = Some(spark_version.version)
                }
                // spark::analyze_plan_response::Result::DdlParse(ddl_parse) => {
                //     self.handler_analyze.ddl_parse = ddl_parse.parsed
                // }
                // spark::analyze_plan_response::Result::SameSemantics(same_semantics) => {
                //     self.handler_analyze.same_semantics = Some(same_semantics.result)
                // }
                // spark::analyze_plan_response::Result::SemanticHash(semantic_hash) => {
                //     self.handler_analyze.semantic_hash = Some(semantic_hash.result)
                // }
                // spark::analyze_plan_response::Result::Persist(_) => {}
                // spark::analyze_plan_response::Result::Unpersist(_) => {}
                // spark::analyze_plan_response::Result::GetStorageLevel(level) => {
                //     self.handler_analyze.get_storage_level = level.storage_level
                // }
                _ => return Err(ClientError::new(ClientErrorKind::Unimplemented(format!(
                    "Handling of analyze response {result:?} not implemented!"
                ))))
            }
        }

        Ok(())
    }

    /// Sends an [`InterruptRequest`](crate::spark::InterruptRequest) to Spark.
    ///
    /// Used to stop long-running operations or cancel all running executions.
    pub(crate) async fn interrupt(
        &mut self,
        interrupt_type: spark::interrupt_request::InterruptType,
        id_or_tag: Option<String>,
    ) -> Result<&mut Self, ClientError> {
        let mut request = spark::InterruptRequest {
            session_id: self.session_id(),
            user_context: self.user_context.clone(),
            client_type: self.builder.user_agent.clone(),
            interrupt_type: 0,
            interrupt: None,
        };

        match interrupt_type {
            spark::interrupt_request::InterruptType::All => {
                request.interrupt_type = interrupt_type.into();
            }
            spark::interrupt_request::InterruptType::Tag => {
                return Err(ClientError::new(ClientErrorKind::Unimplemented(
                    "Tag interrupts are not implemented!".to_string()
                )))
            }
            spark::interrupt_request::InterruptType::OperationId => {
                let op_id = id_or_tag.expect("Operation ID can not be empty");
                let interrupt = spark::interrupt_request::Interrupt::OperationId(op_id);
                request.interrupt_type = interrupt_type.into();
                request.interrupt = Some(interrupt);
            }
            spark::interrupt_request::InterruptType::Unspecified => {
                return Err(ClientError::new(ClientErrorKind::UnspecifiedInterruptRequest))
            }
        };

        let mut client = self.stub.write().await;
        let resp = client
            .interrupt(request.clone())
            .await
            .map_err(|status| {
                ClientError::new(ClientErrorKind::InterruptRequest { status, request })
            })?
            .into_inner();
        drop(client);
        
        self.handler_interrupt = InterruptHandler::default();
        self.handler_interrupt.interrupted_ids = resp.interrupted_ids;
        
        Ok(self)
    }
    

    /// Executes a query plan and streams results from Spark.
    ///
    /// This method handles deserialization, streaming,
    /// and optional *reattachment* for fault-tolerant execution.
    ///
    /// The resulting record batches can be retrieved with [`batches()`](Self::batches).
    pub(crate) async fn execute_plan(
        &mut self,
        plan: spark::Plan
    ) -> Result<&mut Self, ClientError> {
        let mut request = self.new_execute_plan_request();
        request.plan = Some(plan);

        let mut client = self.stub.write().await;
        let mut stream = client
            .execute_plan(request.clone())
            .await
            .map_err(|status| {
                ClientError::new(ClientErrorKind::ExecutePlanRequest { status, request })
            })?
            .into_inner();
        drop(client);

        self.handler_execute = ExecuteHandler::default();
        self.process_stream(&mut stream).await?;
        
        if self.use_reattachable_execute && self.handler_execute.result_complete {
            self.release_all().await?;
        }
        
        Ok(self)
    }

    fn new_execute_plan_request(&mut self) -> spark::ExecutePlanRequest {
        let operation_id = uuid::Uuid::new_v4().to_string();

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

        spark::ExecutePlanRequest {
            session_id: self.session_id(),
            user_context: self.user_context.clone(),
            operation_id: Some(operation_id),
            plan: None,
            client_type: self.builder.user_agent.clone(),
            request_options: vec![spark::execute_plan_request::RequestOption {
                request_option: Some(
                    spark::execute_plan_request::request_option::RequestOption::ReattachOptions(
                        spark::ReattachOptions { reattachable: self.use_reattachable_execute },
                    ),
                ),
            }],
            tags: vec![],
        }
    }
    
    fn handle_execute_response(
        &mut self,
        resp: spark::ExecutePlanResponse
    ) -> Result<(), ClientError> {
        self.validate_session(&resp.session_id)?;

        self.operation_id = Some(resp.operation_id);
        self.response_id = Some(resp.response_id);

        if let Some(data) = resp.response_type {
            match data {
                ResponseType::ArrowBatch(res) => {
                    let (batches, total_count) = crate::io::deserialize(
                        res.data.as_slice(), res.row_count
                    )?;

                    self.handler_execute.batches.extend(batches);
                    self.handler_execute.total_count += total_count;
                }
                ResponseType::SqlCommandResult(sql_cmd) => {
                    self.handler_execute.relation = sql_cmd.clone().relation
                }
                // ResponseType::WriteStreamOperationStartResult(write_stream_op) => {
                //     self.handler.write_stream_operation_start_result = Some(write_stream_op)
                // }
                // ResponseType::StreamingQueryCommandResult(stream_qry_cmd) => {
                //     self.handler.streaming_query_command_result = Some(stream_qry_cmd)
                // }
                // ResponseType::GetResourcesCommandResult(resource_cmd) => {
                //     self.handler.get_resources_command_result = Some(resource_cmd)
                // }
                // ResponseType::StreamingQueryManagerCommandResult(stream_qry_mngr_cmd) => {
                //     self.handler.streaming_query_manager_command_result = Some(stream_qry_mngr_cmd)
                // }
                ResponseType::ResultComplete(_) => self.handler_execute.result_complete = true,
                _ => return Err(ClientError::new(ClientErrorKind::Unimplemented(
                    format!("Handling of plan response {data:?} not implemented!")
                )))
            }
        }
        Ok(())
    }

    /// Execute an [execution reattachment request](crate::spark::ReattachExecuteRequest).
    async fn reattach(&mut self) -> Result<(), ClientError> {
        let request = spark::ReattachExecuteRequest {
            session_id: self.session_id(),
            user_context: self.user_context.clone(),
            operation_id: self.operation_id.clone().unwrap(),
            client_type: self.builder.user_agent.clone(),
            last_response_id: self.response_id.clone(),
        };

        let mut client = self.stub.write().await;
        let mut stream = client
            .reattach_execute(request.clone())
            .await
            .map_err(|status| {
                ClientError::new(ClientErrorKind::ReattachExecuteRequest { status, request })
            })?
            .into_inner();
        drop(client);

        self.process_stream(&mut stream).await?;
        
        if self.use_reattachable_execute && self.handler_execute.result_complete {
            self.release_all().await?;
        }

        Ok(())
    }
    
    async fn process_stream(
        &mut self, stream: &mut Streaming<spark::ExecutePlanResponse>
    ) -> Result<(), ClientError> {
        while let Some(_resp) = match stream.message().await {
            Ok(Some(msg)) => {
                self.handle_execute_response(msg.clone())?;
                Some(msg)
            }
            Ok(None) => {
                if self.use_reattachable_execute && !self.handler_execute.result_complete {
                    Box::pin(self.reattach()).await?;
                }
                None
            }
            Err(status) => {
                if self.use_reattachable_execute && self.response_id.is_some() {
                    self.release_until().await?;
                }
                return Err(ClientError::new(ClientErrorKind::Stream(status)));
            }
        } {}

        Ok(())
    }

    async fn release_until(&mut self) -> Result<(), ClientError> {
        let release_until = spark::release_execute_request::ReleaseUntil {
            response_id: self.response_id.clone().unwrap(),
        };

        self.release_execute(spark::release_execute_request::Release::ReleaseUntil(
            release_until,
        )).await
    }

    async fn release_all(&mut self) -> Result<(), ClientError> {
        let release_all = spark::release_execute_request::ReleaseAll {};

        self.release_execute(spark::release_execute_request::Release::ReleaseAll(
            release_all,
        )).await
    }

    async fn release_execute(
        &mut self,
        release: spark::release_execute_request::Release,
    ) -> Result<(), ClientError> {
        let mut client = self.stub.write().await;

        let request = spark::ReleaseExecuteRequest {
            session_id: self.session_id(),
            user_context: self.user_context.clone(),
            operation_id: self.operation_id.clone().unwrap(),
            client_type: self.builder.user_agent.clone(),
            release: Some(release),
        };

        let _resp = client
            .release_execute(request.clone())
            .await
            .map_err(|status| {
                ClientError::new(ClientErrorKind::ReleaseExecuteRequest { status, request })
            })?
            .into_inner();

        Ok(())
    }
}


#[cfg(test)]
mod tests {
    use crate::test_utils::test_utils::setup_session;
    use crate::spark;
    
    /// Verifies that the client correctly handles and reports errors, such as
    /// a session validation failure.
    #[tokio::test]
    async fn test_validate_session_error() {
        // Arrange: Start server and create a session
        let session = setup_session().await.expect("Failed to create Spark session");

        // Create a clone of the client and manually corrupt its session ID
        let mut client_with_bad_session = session.client().clone();
        client_with_bad_session.session_id = "invalid-session-id".to_string();

        // Act: Attempt to use the corrupted client. This will cause the real server
        // to return an error that Spark Connect may not map directly to a session
        // ID mismatch, but it will be an error nonetheless.
        let result = client_with_bad_session
            .analyze(spark::analyze_plan_request::Analyze::SparkVersion(
                spark::analyze_plan_request::SparkVersion {},
            ))
            .await;

        // Assert: The operation should fail.
        assert!(
            result.is_err(),
            "Expected an error due to invalid session ID"
        );
    }
    
    /// Verifies that the client can send an interrupt request without errors.
    /// This tests the `SparkClient::interrupt_request` method.
    #[tokio::test]
    async fn test_interrupt_all_request() {
        // Arrange: Start server and create a session
        let session = setup_session().await.expect("Failed to create Spark session");
        
        // Act: Send an "interrupt all" request. The server should accept this
        // command gracefully even if nothing is running.
        let mut client = session.client();
        let result = client
            .interrupt(spark::interrupt_request::InterruptType::All, None)
            .await
            .unwrap();
            
        // Assert: The request should succeed. The response may be empty.
        assert_eq!(result.session_id(), session.session_id());
    }
}