linkrs_query/context/
query_context_builder.rs

1//! QueryContext Builder
2//!
3//! Provide a fluent API for building QueryContext, simplifying the process of creating complex configurations.
4//!
5//! # Example
6//!
7//! ```rust,ignore
8//! use linkrs::query::context::{QueryContextBuilder, QueryRequestContext};
9//! use std::sync::Arc;
10//!
11//! let rctx = Arc::new(QueryRequestContext::new("MATCH (n) RETURN n".to_string()));
12//!
13//! let query_context = QueryContextBuilder::new(rctx)
14//!     .with_space_info(space_info)
15//!     .with_charset_info(charset_info)
16//!     .with_arena()
17//!     .build();
18//! ```
19
20use crate::executor::streaming::query_registry::CancelToken;
21use linkrs_core::types::{CharsetInfo, SpaceInfo, SpaceSummary, Timestamp};
22use linkrs_core::IdGenerator;
23use std::sync::Arc;
24
25use super::{QueryContext, QueryRequestContext};
26
27/// QueryContext Builder
28///
29/// Provides a fluent API for constructing QueryContext with complex configurations.
30///
31/// # Example
32///
33/// ```rust,ignore
34/// let rctx = Arc::new(QueryRequestContext::new("MATCH (n) RETURN n".to_string()));
35///
36/// let query_context = QueryContextBuilder::new(rctx)
37///     .with_space_info(space_info)
38///     .with_charset_info(charset_info)
39///     .with_start_id(100)
40///     .with_arena()
41///     .build();
42/// ```
43#[derive(Default)]
44pub struct QueryContextBuilder {
45    rctx: Option<Arc<QueryRequestContext>>,
46    id_gen: Option<IdGenerator>,
47    space_info: Option<SpaceInfo>,
48    charset_info: Option<Box<CharsetInfo>>,
49    snapshot_ts: Option<Timestamp>,
50    isolation_level: Option<linkrs_core::types::TransactionIsolationLevel>,
51    arena_enabled: bool,
52    password_history_depth: usize,
53}
54
55impl QueryContextBuilder {
56    /// Create a new builder with the required QueryRequestContext.
57    pub fn new(rctx: Arc<QueryRequestContext>) -> Self {
58        Self {
59            rctx: Some(rctx),
60            id_gen: None,
61            space_info: None,
62            charset_info: None,
63            snapshot_ts: None,
64            isolation_level: None,
65            arena_enabled: false,
66            password_history_depth: 0,
67        }
68    }
69
70    /// Create a builder from session context.
71    ///
72    /// This simplifies the common pattern of creating a QueryContext
73    /// from a ClientSession, automatically extracting space information.
74    pub fn from_session(rctx: Arc<QueryRequestContext>, space: Option<SpaceSummary>) -> Self {
75        Self {
76            rctx: Some(rctx),
77            id_gen: None,
78            space_info: space.map(SpaceInfo::from),
79            charset_info: None,
80            snapshot_ts: None,
81            isolation_level: None,
82            arena_enabled: false,
83            password_history_depth: 0,
84        }
85    }
86
87    /// Set the space information.
88    pub fn with_space_info(mut self, space_info: SpaceInfo) -> Self {
89        self.space_info = Some(space_info);
90        self
91    }
92
93    /// Set the charset information.
94    pub fn with_charset_info(mut self, charset_info: CharsetInfo) -> Self {
95        self.charset_info = Some(Box::new(charset_info));
96        self
97    }
98
99    /// Set the initial value for the ID generator.
100    pub fn with_start_id(mut self, start_id: i64) -> Self {
101        self.id_gen = Some(IdGenerator::new(start_id));
102        self
103    }
104
105    /// Enable arena allocation for query execution.
106    ///
107    /// Arena allocation is beneficial for queries that create many
108    /// temporary data structures during execution. The concrete arena is
109    /// owned by the execution context; this flag only records the intent.
110    pub fn with_arena(mut self) -> Self {
111        self.arena_enabled = true;
112        self
113    }
114
115    /// Set the MVCC snapshot timestamp for consistent reads within an
116    /// explicit transaction. When set, execution reads data as of this
117    /// timestamp; when absent (`None`), auto-commit current-version reads.
118    pub fn with_snapshot_ts(mut self, snapshot_ts: Timestamp) -> Self {
119        self.snapshot_ts = Some(snapshot_ts);
120        self
121    }
122
123    /// The configured MVCC snapshot timestamp (if any).
124    pub fn snapshot_ts(&self) -> Option<Timestamp> {
125        self.snapshot_ts
126    }
127
128    /// Set the transaction isolation level for an execution inside an
129    /// explicit transaction. Absent for auto-commit statements.
130    pub fn with_isolation_level(
131        mut self,
132        isolation_level: linkrs_core::types::TransactionIsolationLevel,
133    ) -> Self {
134        self.isolation_level = Some(isolation_level);
135        self
136    }
137
138    /// The configured transaction isolation level (if any).
139    pub fn isolation_level(&self) -> Option<linkrs_core::types::TransactionIsolationLevel> {
140        self.isolation_level
141    }
142
143    /// Set the retained password-history depth applied by user-management
144    /// operations planned under this context.
145    pub fn with_password_history_depth(mut self, depth: usize) -> Self {
146        self.password_history_depth = depth;
147        self
148    }
149
150    /// Build the QueryContext.
151    ///
152    /// # Panics
153    ///
154    /// Panics if QueryRequestContext was not provided.
155    pub fn build(self) -> QueryContext {
156        let rctx = self.rctx.expect("QueryRequestContext is required");
157        let id_gen = self.id_gen.unwrap_or_else(|| IdGenerator::new(0));
158
159        QueryContext::from_builder(
160            rctx,
161            super::query_context::ContextParams {
162                cancel_token: CancelToken::new(),
163                id_gen,
164                space_info: self.space_info,
165                charset_info: self.charset_info,
166                snapshot_ts: self.snapshot_ts,
167                isolation_level: self.isolation_level,
168                arena_enabled: self.arena_enabled,
169                password_history_depth: self.password_history_depth,
170            },
171        )
172    }
173}
174
175#[cfg(test)]
176mod tests {
177    use super::*;
178    use linkrs_core::types::{DataType, EngineType, IsolationLevel, MetadataVersion, SpaceStatus};
179    use std::collections::HashMap;
180
181    #[test]
182    fn test_builder_basic() {
183        let rctx = Arc::new(QueryRequestContext {
184            session_id: None,
185            user_name: None,
186            space_name: None,
187            query: "MATCH (n) RETURN n".to_string(),
188            parameters: HashMap::new(),
189            ..Default::default()
190        });
191
192        let query_context = QueryContextBuilder::new(rctx).build();
193
194        assert_eq!(query_context.query(), "MATCH (n) RETURN n");
195    }
196
197    #[test]
198    fn test_builder_with_space_info() {
199        let rctx = Arc::new(QueryRequestContext {
200            session_id: None,
201            user_name: None,
202            space_name: None,
203            query: "MATCH (n) RETURN n".to_string(),
204            parameters: HashMap::new(),
205            ..Default::default()
206        });
207
208        let space_info = SpaceInfo {
209            space_id: 1,
210            space_name: "test_space".to_string(),
211            vid_type: DataType::BigInt,
212            tags: Vec::new(),
213            edge_types: Vec::new(),
214            version: MetadataVersion::default(),
215            comment: None,
216            storage_path: None,
217            isolation_level: IsolationLevel::default(),
218            partition_num: 100,
219            replica_factor: 1,
220            engine_type: EngineType::default(),
221            status: SpaceStatus::Online,
222        };
223
224        let query_context = QueryContextBuilder::new(rctx)
225            .with_space_info(space_info)
226            .build();
227
228        assert_eq!(query_context.space_id(), Some(1));
229        assert_eq!(query_context.space_name(), Some("test_space".to_string()));
230    }
231
232    #[test]
233    fn test_builder_with_start_id() {
234        let rctx = Arc::new(QueryRequestContext {
235            session_id: None,
236            user_name: None,
237            space_name: None,
238            query: "MATCH (n) RETURN n".to_string(),
239            parameters: HashMap::new(),
240            ..Default::default()
241        });
242
243        let query_context = QueryContextBuilder::new(rctx).with_start_id(100).build();
244
245        assert_eq!(query_context.current_id(), 100);
246    }
247
248    #[test]
249    fn test_builder_with_arena() {
250        let rctx = Arc::new(QueryRequestContext {
251            session_id: None,
252            user_name: None,
253            space_name: None,
254            query: "MATCH (n) RETURN n".to_string(),
255            parameters: HashMap::new(),
256            ..Default::default()
257        });
258
259        let query_context = QueryContextBuilder::new(rctx).with_arena().build();
260
261        assert!(query_context.has_arena());
262    }
263
264    #[test]
265    fn test_builder_with_snapshot_ts() {
266        let rctx = Arc::new(QueryRequestContext {
267            session_id: None,
268            user_name: None,
269            space_name: None,
270            query: "MATCH (n) RETURN n".to_string(),
271            parameters: HashMap::new(),
272            ..Default::default()
273        });
274
275        let builder = QueryContextBuilder::new(rctx.clone());
276        assert_eq!(builder.snapshot_ts(), None);
277
278        let query_context = builder.with_snapshot_ts(42).build();
279
280        assert_eq!(query_context.snapshot_ts(), Some(42));
281
282        let default_context = QueryContextBuilder::new(rctx).build();
283        assert_eq!(default_context.snapshot_ts(), None);
284    }
285
286    #[test]
287    fn test_builder_chaining() {
288        let rctx = Arc::new(QueryRequestContext {
289            session_id: None,
290            user_name: None,
291            space_name: None,
292            query: "MATCH (n) RETURN n".to_string(),
293            parameters: HashMap::new(),
294            ..Default::default()
295        });
296
297        let space_info = SpaceInfo {
298            space_id: 1,
299            space_name: "test_space".to_string(),
300            vid_type: DataType::BigInt,
301            tags: Vec::new(),
302            edge_types: Vec::new(),
303            version: MetadataVersion::default(),
304            comment: None,
305            storage_path: None,
306            isolation_level: IsolationLevel::default(),
307            partition_num: 100,
308            replica_factor: 1,
309            engine_type: EngineType::default(),
310            status: SpaceStatus::Online,
311        };
312
313        let query_context = QueryContextBuilder::new(rctx)
314            .with_space_info(space_info)
315            .with_start_id(100)
316            .with_arena()
317            .build();
318
319        assert_eq!(query_context.space_id(), Some(1));
320        assert_eq!(query_context.current_id(), 100);
321        assert!(query_context.has_arena());
322    }
323}
linkrs_query/context/query_context_builder.rs

linkrs_query/context/
query_context_builder.rs
