d_engine_proto/exts/

client_ext.rs

//! Extensions for protobuf-generated client types
//!
//! This module provides additional methods for protobuf-generated types
//! that are not automatically generated by the protobuf compiler.

use bytes::Bytes;

use crate::client::ClientReadRequest;
use crate::client::ClientResponse;
use crate::client::ClientResult;
use crate::client::ReadConsistencyPolicy;
use crate::client::ReadResults;
use crate::client::WriteCommand;
use crate::client::client_response::SuccessResult;
use crate::client::write_command;
use crate::error::ErrorCode;
use crate::error::ErrorMetadata;

impl ClientReadRequest {
    /// Checks if the consistency_policy field is present in the request
    ///
    /// Returns true if the client explicitly specified a consistency policy,
    /// false if the field is absent (should use server default).
    pub fn has_consistency_policy(&self) -> bool {
        self.consistency_policy.is_some()
    }

    /// Gets the consistency policy value safely
    ///
    /// Returns Some(policy) if present, None if field is absent.
    /// Safer alternative that doesn't panic.
    pub fn get_consistency_policy(&self) -> Option<ReadConsistencyPolicy> {
        self.consistency_policy.and_then(|policy_i32| {
            match policy_i32 {
                x if x == ReadConsistencyPolicy::LeaseRead as i32 => {
                    Some(ReadConsistencyPolicy::LeaseRead)
                }
                x if x == ReadConsistencyPolicy::LinearizableRead as i32 => {
                    Some(ReadConsistencyPolicy::LinearizableRead)
                }
                _ => None, // Invalid value
            }
        })
    }
}

impl WriteCommand {
    /// Create write command for key-value pair
    ///
    /// # Parameters
    /// - `key`: Byte array for storage key
    /// - `value`: Byte array to be stored
    pub fn insert(
        key: impl Into<Bytes>,
        value: impl Into<Bytes>,
    ) -> Self {
        let cmd = write_command::Insert {
            key: key.into(),
            value: value.into(),
            ttl_secs: 0,
        };
        Self {
            operation: Some(write_command::Operation::Insert(cmd)),
        }
    }

    /// Create write command for key-value pair with TTL
    ///
    /// # Parameters
    /// - `key`: Byte array for storage key
    /// - `value`: Byte array to be stored
    /// - `ttl_secs`: Time-to-live in seconds
    pub fn insert_with_ttl(
        key: impl Into<Bytes>,
        value: impl Into<Bytes>,
        ttl_secs: u64,
    ) -> Self {
        let cmd = write_command::Insert {
            key: key.into(),
            value: value.into(),
            ttl_secs,
        };
        Self {
            operation: Some(write_command::Operation::Insert(cmd)),
        }
    }

    /// Create deletion command for specified key
    ///
    /// # Parameters
    /// - `key`: Byte array of key to delete
    pub fn delete(key: impl Into<Bytes>) -> Self {
        let cmd = write_command::Delete { key: key.into() };
        Self {
            operation: Some(write_command::Operation::Delete(cmd)),
        }
    }
}

impl ClientResponse {
    /// Build success response for write operations
    ///
    /// # Returns
    /// Response with NoError code and write confirmation
    pub fn write_success() -> Self {
        Self {
            error: ErrorCode::Success as i32,
            success_result: Some(SuccessResult::WriteAck(true)),
            metadata: None,
        }
    }

    /// Check if the write operation was successful
    ///
    /// # Returns
    /// - `true` if the response indicates a successful write operation
    /// - `false` if the response indicates a failed write operation or is not a write response
    pub fn is_write_success(&self) -> bool {
        self.error == ErrorCode::Success as i32
            && matches!(self.success_result, Some(SuccessResult::WriteAck(true)))
    }

    /// Build success response for read operations
    ///
    /// # Parameters
    /// - `results`: Vector of retrieved key-value pairs
    pub fn read_results(results: Vec<ClientResult>) -> Self {
        Self {
            error: ErrorCode::Success as i32,
            success_result: Some(SuccessResult::ReadData(ReadResults { results })),
            metadata: None,
        }
    }

    /// Build generic error response for any operation type
    ///
    /// # Parameters
    /// - `error_code`: Predefined client request error code
    pub fn client_error(error_code: ErrorCode) -> Self {
        Self {
            error: error_code as i32,
            success_result: None,
            metadata: None,
        }
    }

    /// Build NOT_LEADER error response with leader metadata
    ///
    /// # Parameters
    /// - `leader_id`: Optional leader node ID
    /// - `leader_address`: Optional leader address
    pub fn not_leader(
        leader_id: Option<String>,
        leader_address: Option<String>,
    ) -> Self {
        let metadata = if leader_id.is_some() || leader_address.is_some() {
            Some(ErrorMetadata {
                retry_after_ms: None,
                leader_id,
                leader_address,
                debug_message: None,
            })
        } else {
            None
        };

        Self {
            error: ErrorCode::NotLeader as i32,
            success_result: None,
            metadata,
        }
    }

    /// Check if this response indicates the leader's term is outdated
    pub fn is_term_outdated(&self) -> bool {
        ErrorCode::try_from(self.error).map(|e| e.is_term_outdated()).unwrap_or(false)
    }

    /// Check if this response indicates a quorum timeout or failure to receive majority responses
    pub fn is_quorum_timeout_or_failure(&self) -> bool {
        ErrorCode::try_from(self.error)
            .map(|e| e.is_quorum_timeout_or_failure())
            .unwrap_or(false)
    }

    /// Check if this response indicates a failure to receive majority responses
    pub fn is_propose_failure(&self) -> bool {
        ErrorCode::try_from(self.error).map(|e| e.is_propose_failure()).unwrap_or(false)
    }

    /// Check if this response indicates a a retry required
    pub fn is_retry_required(&self) -> bool {
        ErrorCode::try_from(self.error).map(|e| e.is_retry_required()).unwrap_or(false)
    }
}

impl ErrorCode {
    /// Check if this error indicates the leader's term is outdated
    pub fn is_term_outdated(&self) -> bool {
        matches!(self, ErrorCode::TermOutdated)
    }

    /// Check if this error indicates a quorum timeout or failure to receive majority responses
    pub fn is_quorum_timeout_or_failure(&self) -> bool {
        matches!(
            self,
            ErrorCode::ConnectionTimeout | ErrorCode::ProposeFailed | ErrorCode::ClusterUnavailable
        )
    }

    /// Check if this error indicates a failure to receive majority responses
    pub fn is_propose_failure(&self) -> bool {
        matches!(self, ErrorCode::ProposeFailed)
    }

    /// Check if this error indicates a retry required
    pub fn is_retry_required(&self) -> bool {
        matches!(self, ErrorCode::RetryRequired)
    }
}