linkrs-query 0.1.0

//! QueryContext Builder
//!
//! Provide a fluent API for building QueryContext, simplifying the process of creating complex configurations.
//!
//! # Example
//!
//! ```rust,ignore
//! use linkrs::query::context::{QueryContextBuilder, QueryRequestContext};
//! use std::sync::Arc;
//!
//! let rctx = Arc::new(QueryRequestContext::new("MATCH (n) RETURN n".to_string()));
//!
//! let query_context = QueryContextBuilder::new(rctx)
//!     .with_space_info(space_info)
//!     .with_charset_info(charset_info)
//!     .with_arena()
//!     .build();
//! ```

use crate::executor::streaming::query_registry::CancelToken;
use linkrs_core::types::{CharsetInfo, SpaceInfo, SpaceSummary, Timestamp};
use linkrs_core::IdGenerator;
use std::sync::Arc;

use super::{QueryContext, QueryRequestContext};

/// QueryContext Builder
///
/// Provides a fluent API for constructing QueryContext with complex configurations.
///
/// # Example
///
/// ```rust,ignore
/// let rctx = Arc::new(QueryRequestContext::new("MATCH (n) RETURN n".to_string()));
///
/// let query_context = QueryContextBuilder::new(rctx)
///     .with_space_info(space_info)
///     .with_charset_info(charset_info)
///     .with_start_id(100)
///     .with_arena()
///     .build();
/// ```
#[derive(Default)]
pub struct QueryContextBuilder {
    rctx: Option<Arc<QueryRequestContext>>,
    id_gen: Option<IdGenerator>,
    space_info: Option<SpaceInfo>,
    charset_info: Option<Box<CharsetInfo>>,
    snapshot_ts: Option<Timestamp>,
    isolation_level: Option<linkrs_core::types::TransactionIsolationLevel>,
    arena_enabled: bool,
    password_history_depth: usize,
}

impl QueryContextBuilder {
    /// Create a new builder with the required QueryRequestContext.
    pub fn new(rctx: Arc<QueryRequestContext>) -> Self {
        Self {
            rctx: Some(rctx),
            id_gen: None,
            space_info: None,
            charset_info: None,
            snapshot_ts: None,
            isolation_level: None,
            arena_enabled: false,
            password_history_depth: 0,
        }
    }

    /// Create a builder from session context.
    ///
    /// This simplifies the common pattern of creating a QueryContext
    /// from a ClientSession, automatically extracting space information.
    pub fn from_session(rctx: Arc<QueryRequestContext>, space: Option<SpaceSummary>) -> Self {
        Self {
            rctx: Some(rctx),
            id_gen: None,
            space_info: space.map(SpaceInfo::from),
            charset_info: None,
            snapshot_ts: None,
            isolation_level: None,
            arena_enabled: false,
            password_history_depth: 0,
        }
    }

    /// Set the space information.
    pub fn with_space_info(mut self, space_info: SpaceInfo) -> Self {
        self.space_info = Some(space_info);
        self
    }

    /// Set the charset information.
    pub fn with_charset_info(mut self, charset_info: CharsetInfo) -> Self {
        self.charset_info = Some(Box::new(charset_info));
        self
    }

    /// Set the initial value for the ID generator.
    pub fn with_start_id(mut self, start_id: i64) -> Self {
        self.id_gen = Some(IdGenerator::new(start_id));
        self
    }

    /// Enable arena allocation for query execution.
    ///
    /// Arena allocation is beneficial for queries that create many
    /// temporary data structures during execution. The concrete arena is
    /// owned by the execution context; this flag only records the intent.
    pub fn with_arena(mut self) -> Self {
        self.arena_enabled = true;
        self
    }

    /// Set the MVCC snapshot timestamp for consistent reads within an
    /// explicit transaction. When set, execution reads data as of this
    /// timestamp; when absent (`None`), auto-commit current-version reads.
    pub fn with_snapshot_ts(mut self, snapshot_ts: Timestamp) -> Self {
        self.snapshot_ts = Some(snapshot_ts);
        self
    }

    /// The configured MVCC snapshot timestamp (if any).
    pub fn snapshot_ts(&self) -> Option<Timestamp> {
        self.snapshot_ts
    }

    /// Set the transaction isolation level for an execution inside an
    /// explicit transaction. Absent for auto-commit statements.
    pub fn with_isolation_level(
        mut self,
        isolation_level: linkrs_core::types::TransactionIsolationLevel,
    ) -> Self {
        self.isolation_level = Some(isolation_level);
        self
    }

    /// The configured transaction isolation level (if any).
    pub fn isolation_level(&self) -> Option<linkrs_core::types::TransactionIsolationLevel> {
        self.isolation_level
    }

    /// Set the retained password-history depth applied by user-management
    /// operations planned under this context.
    pub fn with_password_history_depth(mut self, depth: usize) -> Self {
        self.password_history_depth = depth;
        self
    }

    /// Build the QueryContext.
    ///
    /// # Panics
    ///
    /// Panics if QueryRequestContext was not provided.
    pub fn build(self) -> QueryContext {
        let rctx = self.rctx.expect("QueryRequestContext is required");
        let id_gen = self.id_gen.unwrap_or_else(|| IdGenerator::new(0));

        QueryContext::from_builder(
            rctx,
            super::query_context::ContextParams {
                cancel_token: CancelToken::new(),
                id_gen,
                space_info: self.space_info,
                charset_info: self.charset_info,
                snapshot_ts: self.snapshot_ts,
                isolation_level: self.isolation_level,
                arena_enabled: self.arena_enabled,
                password_history_depth: self.password_history_depth,
            },
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use linkrs_core::types::{DataType, EngineType, IsolationLevel, MetadataVersion, SpaceStatus};
    use std::collections::HashMap;

    #[test]
    fn test_builder_basic() {
        let rctx = Arc::new(QueryRequestContext {
            session_id: None,
            user_name: None,
            space_name: None,
            query: "MATCH (n) RETURN n".to_string(),
            parameters: HashMap::new(),
            ..Default::default()
        });

        let query_context = QueryContextBuilder::new(rctx).build();

        assert_eq!(query_context.query(), "MATCH (n) RETURN n");
    }

    #[test]
    fn test_builder_with_space_info() {
        let rctx = Arc::new(QueryRequestContext {
            session_id: None,
            user_name: None,
            space_name: None,
            query: "MATCH (n) RETURN n".to_string(),
            parameters: HashMap::new(),
            ..Default::default()
        });

        let space_info = SpaceInfo {
            space_id: 1,
            space_name: "test_space".to_string(),
            vid_type: DataType::BigInt,
            tags: Vec::new(),
            edge_types: Vec::new(),
            version: MetadataVersion::default(),
            comment: None,
            storage_path: None,
            isolation_level: IsolationLevel::default(),
            partition_num: 100,
            replica_factor: 1,
            engine_type: EngineType::default(),
            status: SpaceStatus::Online,
        };

        let query_context = QueryContextBuilder::new(rctx)
            .with_space_info(space_info)
            .build();

        assert_eq!(query_context.space_id(), Some(1));
        assert_eq!(query_context.space_name(), Some("test_space".to_string()));
    }

    #[test]
    fn test_builder_with_start_id() {
        let rctx = Arc::new(QueryRequestContext {
            session_id: None,
            user_name: None,
            space_name: None,
            query: "MATCH (n) RETURN n".to_string(),
            parameters: HashMap::new(),
            ..Default::default()
        });

        let query_context = QueryContextBuilder::new(rctx).with_start_id(100).build();

        assert_eq!(query_context.current_id(), 100);
    }

    #[test]
    fn test_builder_with_arena() {
        let rctx = Arc::new(QueryRequestContext {
            session_id: None,
            user_name: None,
            space_name: None,
            query: "MATCH (n) RETURN n".to_string(),
            parameters: HashMap::new(),
            ..Default::default()
        });

        let query_context = QueryContextBuilder::new(rctx).with_arena().build();

        assert!(query_context.has_arena());
    }

    #[test]
    fn test_builder_with_snapshot_ts() {
        let rctx = Arc::new(QueryRequestContext {
            session_id: None,
            user_name: None,
            space_name: None,
            query: "MATCH (n) RETURN n".to_string(),
            parameters: HashMap::new(),
            ..Default::default()
        });

        let builder = QueryContextBuilder::new(rctx.clone());
        assert_eq!(builder.snapshot_ts(), None);

        let query_context = builder.with_snapshot_ts(42).build();

        assert_eq!(query_context.snapshot_ts(), Some(42));

        let default_context = QueryContextBuilder::new(rctx).build();
        assert_eq!(default_context.snapshot_ts(), None);
    }

    #[test]
    fn test_builder_chaining() {
        let rctx = Arc::new(QueryRequestContext {
            session_id: None,
            user_name: None,
            space_name: None,
            query: "MATCH (n) RETURN n".to_string(),
            parameters: HashMap::new(),
            ..Default::default()
        });

        let space_info = SpaceInfo {
            space_id: 1,
            space_name: "test_space".to_string(),
            vid_type: DataType::BigInt,
            tags: Vec::new(),
            edge_types: Vec::new(),
            version: MetadataVersion::default(),
            comment: None,
            storage_path: None,
            isolation_level: IsolationLevel::default(),
            partition_num: 100,
            replica_factor: 1,
            engine_type: EngineType::default(),
            status: SpaceStatus::Online,
        };

        let query_context = QueryContextBuilder::new(rctx)
            .with_space_info(space_info)
            .with_start_id(100)
            .with_arena()
            .build();

        assert_eq!(query_context.space_id(), Some(1));
        assert_eq!(query_context.current_id(), 100);
        assert!(query_context.has_arena());
    }
}