reddb_server/storage/query/step/
source.rs1use super::{Step, StepResult, Traverser, TraverserRequirement, TraverserValue};
13use std::any::Any;
14
15pub trait SourceStep: Step {
17 fn generate_traversers(&self) -> Vec<Traverser>;
19}
20
21#[derive(Debug, Clone)]
23pub struct VertexSourceStep {
24 id: String,
25 labels: Vec<String>,
26 vertex_ids: Vec<String>,
28 vertex_type: Option<String>,
30}
31
32impl VertexSourceStep {
33 pub fn new() -> Self {
35 Self {
36 id: "V_0".to_string(),
37 labels: Vec::new(),
38 vertex_ids: Vec::new(),
39 vertex_type: None,
40 }
41 }
42
43 pub fn with_ids(ids: Vec<String>) -> Self {
45 Self {
46 id: format!("V_{}", ids.first().unwrap_or(&"0".to_string())),
47 labels: Vec::new(),
48 vertex_ids: ids,
49 vertex_type: None,
50 }
51 }
52
53 pub fn with_type(vertex_type: String) -> Self {
55 Self {
56 id: format!("V_{}", vertex_type),
57 labels: Vec::new(),
58 vertex_ids: Vec::new(),
59 vertex_type: Some(vertex_type),
60 }
61 }
62
63 pub fn with_id(mut self, id: String) -> Self {
65 self.id = id;
66 self
67 }
68
69 pub fn vertex_ids(&self) -> &[String] {
71 &self.vertex_ids
72 }
73
74 pub fn vertex_type(&self) -> Option<&str> {
76 self.vertex_type.as_deref()
77 }
78}
79
80impl Default for VertexSourceStep {
81 fn default() -> Self {
82 Self::new()
83 }
84}
85
86impl Step for VertexSourceStep {
87 fn id(&self) -> &str {
88 &self.id
89 }
90
91 fn name(&self) -> &str {
92 "VertexSourceStep"
93 }
94
95 fn labels(&self) -> &[String] {
96 &self.labels
97 }
98
99 fn add_label(&mut self, label: String) {
100 if !self.labels.contains(&label) {
101 self.labels.push(label);
102 }
103 }
104
105 fn requirements(&self) -> &[TraverserRequirement] {
106 &[] }
108
109 fn process_traverser(&self, _traverser: Traverser) -> StepResult {
110 StepResult::Filter
113 }
114
115 fn reset(&mut self) {
116 }
118
119 fn clone_step(&self) -> Box<dyn Step> {
120 Box::new(self.clone())
121 }
122
123 fn as_any(&self) -> &dyn Any {
124 self
125 }
126
127 fn as_any_mut(&mut self) -> &mut dyn Any {
128 self
129 }
130}
131
132impl SourceStep for VertexSourceStep {
133 fn generate_traversers(&self) -> Vec<Traverser> {
134 self.vertex_ids
137 .iter()
138 .map(|id| Traverser::new(id))
139 .collect()
140 }
141}
142
143#[derive(Debug, Clone)]
145pub struct EdgeSourceStep {
146 id: String,
147 labels: Vec<String>,
148 edge_ids: Vec<String>,
150 edge_type: Option<String>,
152}
153
154impl EdgeSourceStep {
155 pub fn new() -> Self {
157 Self {
158 id: "E_0".to_string(),
159 labels: Vec::new(),
160 edge_ids: Vec::new(),
161 edge_type: None,
162 }
163 }
164
165 pub fn with_ids(ids: Vec<String>) -> Self {
167 Self {
168 id: format!("E_{}", ids.first().unwrap_or(&"0".to_string())),
169 labels: Vec::new(),
170 edge_ids: ids,
171 edge_type: None,
172 }
173 }
174
175 pub fn with_type(edge_type: String) -> Self {
177 Self {
178 id: format!("E_{}", edge_type),
179 labels: Vec::new(),
180 edge_ids: Vec::new(),
181 edge_type: Some(edge_type),
182 }
183 }
184
185 pub fn with_id(mut self, id: String) -> Self {
187 self.id = id;
188 self
189 }
190
191 pub fn edge_ids(&self) -> &[String] {
193 &self.edge_ids
194 }
195
196 pub fn edge_type(&self) -> Option<&str> {
198 self.edge_type.as_deref()
199 }
200}
201
202impl Default for EdgeSourceStep {
203 fn default() -> Self {
204 Self::new()
205 }
206}
207
208impl Step for EdgeSourceStep {
209 fn id(&self) -> &str {
210 &self.id
211 }
212
213 fn name(&self) -> &str {
214 "EdgeSourceStep"
215 }
216
217 fn labels(&self) -> &[String] {
218 &self.labels
219 }
220
221 fn add_label(&mut self, label: String) {
222 if !self.labels.contains(&label) {
223 self.labels.push(label);
224 }
225 }
226
227 fn requirements(&self) -> &[TraverserRequirement] {
228 &[]
229 }
230
231 fn process_traverser(&self, _traverser: Traverser) -> StepResult {
232 StepResult::Filter
233 }
234
235 fn reset(&mut self) {
236 }
238
239 fn clone_step(&self) -> Box<dyn Step> {
240 Box::new(self.clone())
241 }
242
243 fn as_any(&self) -> &dyn Any {
244 self
245 }
246
247 fn as_any_mut(&mut self) -> &mut dyn Any {
248 self
249 }
250}
251
252impl SourceStep for EdgeSourceStep {
253 fn generate_traversers(&self) -> Vec<Traverser> {
254 self.edge_ids
256 .iter()
257 .map(|id| {
258 Traverser::with_value(TraverserValue::Edge {
259 id: id.clone(),
260 source: String::new(), target: String::new(),
262 label: self.edge_type.clone().unwrap_or_default(),
263 })
264 })
265 .collect()
266 }
267}
268
269#[derive(Debug, Clone)]
271pub struct InjectStep {
272 id: String,
273 labels: Vec<String>,
274 values: Vec<TraverserValue>,
275}
276
277impl InjectStep {
278 pub fn new(values: Vec<TraverserValue>) -> Self {
280 Self {
281 id: "inject_0".to_string(),
282 labels: Vec::new(),
283 values,
284 }
285 }
286
287 pub fn with_id(mut self, id: String) -> Self {
289 self.id = id;
290 self
291 }
292}
293
294impl Step for InjectStep {
295 fn id(&self) -> &str {
296 &self.id
297 }
298
299 fn name(&self) -> &str {
300 "InjectStep"
301 }
302
303 fn labels(&self) -> &[String] {
304 &self.labels
305 }
306
307 fn add_label(&mut self, label: String) {
308 if !self.labels.contains(&label) {
309 self.labels.push(label);
310 }
311 }
312
313 fn requirements(&self) -> &[TraverserRequirement] {
314 &[]
315 }
316
317 fn process_traverser(&self, traverser: Traverser) -> StepResult {
318 let mut result = vec![traverser];
320 for value in &self.values {
321 result.push(Traverser::with_value(value.clone()));
322 }
323 StepResult::Emit(result)
324 }
325
326 fn reset(&mut self) {
327 }
329
330 fn clone_step(&self) -> Box<dyn Step> {
331 Box::new(self.clone())
332 }
333
334 fn as_any(&self) -> &dyn Any {
335 self
336 }
337
338 fn as_any_mut(&mut self) -> &mut dyn Any {
339 self
340 }
341}
342
343impl SourceStep for InjectStep {
344 fn generate_traversers(&self) -> Vec<Traverser> {
345 self.values
346 .iter()
347 .map(|v| Traverser::with_value(v.clone()))
348 .collect()
349 }
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355
356 #[test]
357 fn test_vertex_source_all() {
358 let step = VertexSourceStep::new();
359 assert_eq!(step.name(), "VertexSourceStep");
360 assert!(step.vertex_ids().is_empty());
361 assert!(step.vertex_type().is_none());
362 }
363
364 #[test]
365 fn test_vertex_source_with_ids() {
366 let step = VertexSourceStep::with_ids(vec!["v1".to_string(), "v2".to_string()]);
367 assert_eq!(step.vertex_ids().len(), 2);
368
369 let traversers = step.generate_traversers();
370 assert_eq!(traversers.len(), 2);
371 }
372
373 #[test]
374 fn test_vertex_source_with_type() {
375 let step = VertexSourceStep::with_type("Host".to_string());
376 assert_eq!(step.vertex_type(), Some("Host"));
377 }
378
379 #[test]
380 fn test_edge_source() {
381 let step = EdgeSourceStep::new();
382 assert_eq!(step.name(), "EdgeSourceStep");
383 }
384
385 #[test]
386 fn test_inject_step() {
387 let step = InjectStep::new(vec![
388 TraverserValue::String("hello".to_string()),
389 TraverserValue::Integer(42),
390 ]);
391
392 let traversers = step.generate_traversers();
393 assert_eq!(traversers.len(), 2);
394 }
395
396 #[test]
397 fn test_step_labels() {
398 let mut step = VertexSourceStep::new();
399 step.add_label("a".to_string());
400 step.add_label("b".to_string());
401 assert_eq!(step.labels().len(), 2);
402
403 step.add_label("a".to_string());
405 assert_eq!(step.labels().len(), 2);
406 }
407}