Skip to main content

linkrs_query/parser/parsing/
util_stmt_parser.rs

1//! Tool Statement Parsing Module
2//!
3//! Responsible for parsing statements related to tool functions, including USE, SHOW, EXPLAIN, FETCH, LOOKUP, UNWIND, RETURN, WITH, YIELD, SET, REMOVE, etc.
4
5use crate::parser::ast::stmt::*;
6use crate::parser::ast::types::OrderDirection;
7use crate::parser::core::error::{ParseError, ParseErrorKind};
8use crate::parser::parsing::clause_parser::ClauseParser;
9use crate::parser::parsing::expr_parser::parse_expression_with_context;
10use crate::parser::parsing::parse_context::ParseContext;
11use crate::parser::TokenKind;
12use linkrs_core::types::expr::contextual::ContextualExpression;
13
14/// Tool Syntax Parser
15pub struct UtilStmtParser;
16
17impl UtilStmtParser {
18    pub fn new() -> Self {
19        Self
20    }
21
22    /// Analysis of the USE statement
23    pub fn parse_use_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
24        let start_span = ctx.current_span();
25        ctx.expect_token(TokenKind::Use)?;
26
27        // `USE GRAPH <name>` is a namespace alias for `USE <name>`.
28        // `SPACE` is intentionally not consumed here: it stays a legal
29        // space name in USE position, unlike CREATE/DROP where SPACE is
30        // the object keyword (see ddl_parser::consume_space_or_graph).
31        // A lone `USE graph` names the space `graph`: GRAPH only acts as
32        // the alias keyword when another identifier follows it.
33        if ctx.check_keyword("GRAPH") {
34            let graph_word = ctx.current_token().lexeme.clone();
35            ctx.consume_keyword("GRAPH")?;
36            match ctx.expect_identifier() {
37                Ok(space) => {
38                    return Ok(Stmt::Use(UseStmt {
39                        span: start_span,
40                        space,
41                    }));
42                }
43                Err(_) => {
44                    return Ok(Stmt::Use(UseStmt {
45                        span: start_span,
46                        space: graph_word,
47                    }));
48                }
49            }
50        }
51
52        let space = ctx.expect_identifier()?;
53
54        Ok(Stmt::Use(UseStmt {
55            span: start_span,
56            space,
57        }))
58    }
59
60    /// Analyzing the internal methods of the SHOW CREATE statement
61    pub fn parse_show_create_internal(
62        &mut self,
63        ctx: &mut ParseContext,
64        start_span: crate::parser::ast::types::Span,
65    ) -> Result<Stmt, ParseError> {
66        ctx.expect_token(TokenKind::Create)?;
67
68        let target = if ctx.match_token(TokenKind::Space) {
69            ShowCreateTarget::Space(ctx.expect_identifier()?)
70        } else if ctx.match_token(TokenKind::Tag) {
71            ShowCreateTarget::Tag(ctx.expect_identifier()?)
72        } else if ctx.match_token(TokenKind::Edge) {
73            ShowCreateTarget::Edge(ctx.expect_identifier()?)
74        } else if ctx.match_token(TokenKind::Index) {
75            ShowCreateTarget::Index(ctx.expect_identifier()?)
76        } else {
77            return Err(ParseError::new(
78                ParseErrorKind::UnexpectedToken,
79                "Expected SPACE, TAG, EDGE, or INDEX after SHOW CREATE".to_string(),
80                ctx.current_position(),
81            ));
82        };
83
84        let span = ctx.span_since(start_span.start);
85
86        Ok(Stmt::ShowCreate(ShowCreateStmt { span, target }))
87    }
88
89    /// Analysis of the internal methods of the SHOW ROLES command
90    pub fn parse_show_roles_internal(
91        &mut self,
92        ctx: &mut ParseContext,
93        start_span: crate::parser::ast::types::Span,
94    ) -> Result<Stmt, ParseError> {
95        ctx.expect_token(TokenKind::Roles)?;
96
97        // Optional IN <space_name> clause; space names accept the same
98        // hyphenated/quoted forms as other DCL object names.
99        let space_name = if ctx.match_token(TokenKind::In) {
100            Some(ctx.expect_dcl_name()?)
101        } else {
102            None
103        };
104
105        let span = ctx.span_since(start_span.start);
106
107        Ok(Stmt::ShowRoles(ShowRolesStmt { span, space_name }))
108    }
109
110    /// Analyzing the EXPLAIN statement
111    pub fn parse_explain_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
112        let _start_span = ctx.current_span();
113        ctx.expect_token(TokenKind::Explain)?;
114
115        // After “EXPLAIN”, a substatement needs to be parsed.
116        // Here, we need to invoke the main parser, but due to circular dependency issues, we are returning a placeholder.
117        // The actual parsing will be handled within the StmtParser.
118        Err(ParseError::new(
119            ParseErrorKind::SyntaxError,
120            "EXPLAIN should be handled by main parser".to_string(),
121            ctx.current_position(),
122        ))
123    }
124
125    /// Analysis of the FETCH statement
126    pub fn parse_fetch_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
127        let start_span = ctx.current_span();
128        ctx.expect_token(TokenKind::Fetch)?;
129
130        // The FETCH PROP ON <tag> <ids> syntax is supported.
131        let _with_props = ctx.match_token(TokenKind::Prop);
132
133        // `->` is an edge arrow inside FETCH, not the JSON access operator
134        // (`expr -> 'key'` postfix JSON get), so expression parsing must not
135        // consume `src -> dst` as a JSON get.
136        let target = ctx.with_edge_syntax_mode(|ctx| {
137            if ctx.match_token(TokenKind::On) {
138                // Check for * (all tags)
139                let tag_name = if ctx.match_token(TokenKind::Star) {
140                    None // None means all tags
141                } else {
142                    Some(ctx.expect_identifier()?)
143                };
144
145                let first_expr = self.parse_expression(ctx)?;
146                if ctx.check_token(TokenKind::Arrow) {
147                    ctx.expect_token(TokenKind::Arrow)?;
148                    let dst = self.parse_expression(ctx)?;
149                    let rank = if ctx.match_token(TokenKind::At) {
150                        Some(self.parse_expression(ctx)?)
151                    } else {
152                        None
153                    };
154                    let edge_type = tag_name.ok_or_else(|| {
155                        ParseError::new(
156                            ParseErrorKind::SyntaxError,
157                            "FETCH edge form requires an edge type; '*' is only valid for vertices"
158                                .to_string(),
159                            ctx.current_position(),
160                        )
161                    })?;
162                    Ok(FetchTarget::Edges {
163                        src: first_expr,
164                        dst,
165                        edge_type,
166                        rank,
167                        properties: None,
168                    })
169                } else {
170                    let mut ids = vec![first_expr];
171                    while ctx.match_token(TokenKind::Comma) {
172                        ids.push(self.parse_expression(ctx)?);
173                    }
174                    Ok(FetchTarget::Vertices {
175                        tag_name,
176                        ids,
177                        properties: None,
178                    })
179                }
180            } else if ctx.match_token(TokenKind::Tag) {
181                let tag_name = Some(ctx.expect_identifier()?);
182                let ids = self.parse_expression_list(ctx)?;
183                Ok(FetchTarget::Vertices {
184                    tag_name,
185                    ids,
186                    properties: None,
187                })
188            } else if ctx.match_token(TokenKind::Edge) {
189                let edge_type = ctx.expect_identifier()?;
190                let src = self.parse_expression(ctx)?;
191                ctx.expect_token(TokenKind::Arrow)?;
192                let dst = self.parse_expression(ctx)?;
193                let rank = if ctx.match_token(TokenKind::At) {
194                    Some(self.parse_expression(ctx)?)
195                } else {
196                    None
197                };
198                Ok(FetchTarget::Edges {
199                    src,
200                    dst,
201                    edge_type,
202                    rank,
203                    properties: None,
204                })
205            } else {
206                let ids = self.parse_expression_list(ctx)?;
207                Ok(FetchTarget::Vertices {
208                    tag_name: None,
209                    ids,
210                    properties: None,
211                })
212            }
213        })?;
214
215        let yield_clause = if ctx.match_token(TokenKind::Yield) {
216            Some(ClauseParser::new().parse_yield_clause(ctx)?)
217        } else {
218            None
219        };
220
221        Ok(Stmt::Fetch(FetchStmt {
222            span: start_span,
223            target,
224            yield_clause,
225        }))
226    }
227
228    /// Analysis of the LOOKUP statement
229    pub fn parse_lookup_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
230        let start_span = ctx.current_span();
231        ctx.expect_token(TokenKind::Lookup)?;
232
233        let target = if ctx.match_token(TokenKind::On) {
234            // Check if it's an EDGE or TAG keyword followed by identifier
235            if ctx.current_token().kind == TokenKind::Edge {
236                // LOOKUP ON EDGE <name>
237                ctx.next_token(); // consume EDGE
238                let name = ctx.expect_identifier()?;
239                LookupTarget::Edge(name)
240            } else if ctx.current_token().kind == TokenKind::Tag {
241                // LOOKUP ON TAG <name>
242                ctx.next_token(); // consume TAG
243                let name = ctx.expect_identifier()?;
244                LookupTarget::Tag(name)
245            } else {
246                // LOOKUP ON <name> - type will be resolved during validation
247                let name = ctx.expect_identifier()?;
248                LookupTarget::Unspecified(name)
249            }
250        } else {
251            LookupTarget::Unspecified(String::new())
252        };
253
254        let where_clause = if ctx.match_token(TokenKind::Where) {
255            ctx.recover_clause(|_| Ok(None), |c| self.parse_expression(c).map(Some))?
256        } else {
257            None
258        };
259
260        let yield_clause = if ctx.match_token(TokenKind::Yield) {
261            Some(ClauseParser::new().parse_yield_clause(ctx)?)
262        } else {
263            None
264        };
265
266        Ok(Stmt::Lookup(LookupStmt {
267            span: start_span,
268            target,
269            where_clause,
270            yield_clause,
271        }))
272    }
273
274    /// Analysis of the UNWIND statement
275    pub fn parse_unwind_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
276        let start_span = ctx.current_span();
277        ctx.expect_token(TokenKind::Unwind)?;
278
279        let expression = self.parse_expression(ctx)?;
280
281        ctx.match_token(TokenKind::As);
282
283        let variable = ctx.expect_identifier()?;
284
285        let return_clause = if ctx.match_token(TokenKind::Return) {
286            Some(ClauseParser::new().parse_return_clause(ctx)?)
287        } else {
288            None
289        };
290
291        let order_by = None;
292        let limit = None;
293        let skip = None;
294
295        Ok(Stmt::Unwind(UnwindStmt {
296            span: start_span,
297            expression,
298            variable,
299            return_clause,
300            order_by,
301            limit,
302            skip,
303        }))
304    }
305
306    /// Analysis of the RETURN statement
307    pub fn parse_return_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
308        let start_span = ctx.current_span();
309        ctx.expect_token(TokenKind::Return)?;
310
311        let return_clause = ClauseParser::new().parse_return_clause(ctx)?;
312
313        Ok(Stmt::Return(ReturnStmt {
314            span: start_span,
315            items: return_clause.items,
316            distinct: return_clause.distinct,
317            order_by: return_clause.order_by,
318            skip: return_clause.skip,
319            limit: return_clause.limit,
320        }))
321    }
322
323    /// Analysis of the WITH statement
324    pub fn parse_with_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
325        let start_span = ctx.current_span();
326        ctx.expect_token(TokenKind::With)?;
327
328        let recursive = ctx.match_token(TokenKind::Recursive);
329
330        // Common table expressions: `name AS (<query>) [, ...]`.
331        // A `IDENT AS (` sequence can never start a plain WITH item (an
332        // alias after AS must be an identifier), so detection is unambiguous.
333        let mut ctes = Vec::new();
334        while let Some(cte) = Self::try_parse_cte_def(ctx)? {
335            ctes.push(cte);
336            if !ctx.match_token(TokenKind::Comma) {
337                break;
338            }
339        }
340
341        let mut items = Vec::new();
342        let distinct = ctx.match_token(TokenKind::Distinct);
343
344        loop {
345            let expr = self.parse_expression(ctx)?;
346            let alias = if ctx.match_token(TokenKind::As) {
347                Some(ctx.expect_identifier()?)
348            } else {
349                None
350            };
351            items.push(ReturnItem::Expression {
352                expression: expr,
353                alias,
354            });
355            if !ctx.match_token(TokenKind::Comma) {
356                break;
357            }
358        }
359
360        let where_clause = if ctx.match_token(TokenKind::Where) {
361            ctx.recover_clause(|_| Ok(None), |c| self.parse_expression(c).map(Some))?
362        } else {
363            None
364        };
365
366        // Explanation of `ORDER BY`
367        let order_by = if ctx.match_token(TokenKind::Order) {
368            ctx.recover_clause(
369                |_| Ok(None),
370                |c| {
371                    c.expect_token(TokenKind::By)?;
372                    self.parse_order_by_clause(c).map(Some)
373                },
374            )?
375        } else {
376            None
377        };
378
379        // Analysis of SKIP
380        let skip = if ctx.match_token(TokenKind::Skip) {
381            ctx.recover_clause(
382                |_| Ok(None),
383                |c| {
384                    let count = c.expect_integer_literal()? as usize;
385                    Ok(Some(crate::parser::ast::types::SkipClause {
386                        span: c.current_span(),
387                        count,
388                    }))
389                },
390            )?
391        } else {
392            None
393        };
394
395        // Analysis of the LIMIT clause
396        let limit = if ctx.match_token(TokenKind::Limit) {
397            ctx.recover_clause(
398                |_| Ok(None),
399                |c| {
400                    let count = c.expect_integer_literal()? as usize;
401                    Ok(Some(crate::parser::ast::types::LimitClause {
402                        span: c.current_span(),
403                        count,
404                    }))
405                },
406            )?
407        } else {
408            None
409        };
410
411        Ok(Stmt::With(WithStmt {
412            span: start_span,
413            items,
414            where_clause,
415            distinct,
416            order_by,
417            skip,
418            limit,
419            recursive,
420            ctes,
421        }))
422    }
423
424    /// Try to parse one `name AS (<query>)` CTE definition.
425    ///
426    /// Returns `Ok(None)` (without consuming input) when the upcoming tokens
427    /// do not form a CTE header.
428    fn try_parse_cte_def(ctx: &mut ParseContext) -> Result<Option<CteDef>, ParseError> {
429        let ckpt = ctx.checkpoint();
430        let name = match ctx.expect_identifier() {
431            Ok(name) => name,
432            Err(_) => {
433                ctx.restore(ckpt);
434                return Ok(None);
435            }
436        };
437        // A CTE definition is always `name AS (`; `name (` starts a function
438        // call item instead (e.g. `WITH count(*) AS total`), so back off and
439        // let the expression parser handle it.
440        if ctx.check_token(TokenKind::LParen) {
441            ctx.restore(ckpt);
442            return Ok(None);
443        }
444        if !ctx.check_keyword("AS") {
445            ctx.restore(ckpt);
446            return Ok(None);
447        }
448        ctx.consume_keyword("AS")?;
449        if !ctx.check_token(TokenKind::LParen) {
450            ctx.restore(ckpt);
451            return Ok(None);
452        }
453        let cte_span = ctx.current_span();
454        ctx.expect_token(TokenKind::LParen)?;
455        let body = crate::parser::parsing::StmtParser::parse_statement(ctx).map_err(|e| {
456            ParseError::new(
457                ParseErrorKind::SyntaxError,
458                format!("Invalid CTE body for '{}': {}", name, e),
459                ctx.current_position(),
460            )
461        })?;
462        ctx.expect_token(TokenKind::RParen)?;
463        Ok(Some(CteDef {
464            span: cte_span,
465            name,
466            body: Box::new(body),
467        }))
468    }
469
470    /// Analysis of the YIELD statement
471    pub fn parse_yield_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
472        let start_span = ctx.current_span();
473        ctx.expect_token(TokenKind::Yield)?;
474
475        let yield_clause = ClauseParser::new().parse_yield_clause(ctx)?;
476
477        Ok(Stmt::Yield(YieldStmt {
478            span: start_span,
479            items: yield_clause.items,
480            where_clause: yield_clause.where_clause,
481            distinct: yield_clause.distinct,
482            order_by: yield_clause.order_by,
483            skip: yield_clause.skip,
484            limit: yield_clause.limit,
485            sample: yield_clause.sample,
486        }))
487    }
488
489    /// Analyzing the SET statement
490    pub fn parse_set_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
491        let start_span = ctx.current_span();
492        ctx.expect_token(TokenKind::Set)?;
493
494        let set_clause = ClauseParser::new().parse_set_clause(ctx)?;
495
496        Ok(Stmt::Set(SetStmt {
497            span: start_span,
498            assignments: set_clause.assignments,
499        }))
500    }
501
502    /// Analysis of the REMOVE statement
503    pub fn parse_remove_statement(&mut self, ctx: &mut ParseContext) -> Result<Stmt, ParseError> {
504        let start_span = ctx.current_span();
505        ctx.expect_token(TokenKind::Remove)?;
506
507        let mut items = Vec::new();
508        loop {
509            items.push(self.parse_expression(ctx)?);
510            if !ctx.match_token(TokenKind::Comma) {
511                break;
512            }
513        }
514
515        Ok(Stmt::Remove(RemoveStmt {
516            span: start_span,
517            items,
518        }))
519    }
520
521    /// Parse the list of expressions
522    fn parse_expression_list(
523        &mut self,
524        ctx: &mut ParseContext,
525    ) -> Result<Vec<ContextualExpression>, ParseError> {
526        let mut expressions = Vec::new();
527
528        loop {
529            expressions.push(self.parse_expression(ctx)?);
530            if !ctx.match_token(TokenKind::Comma) {
531                break;
532            }
533        }
534
535        Ok(expressions)
536    }
537
538    /// Analyzing the expression
539    fn parse_expression(
540        &mut self,
541        ctx: &mut ParseContext,
542    ) -> Result<ContextualExpression, ParseError> {
543        parse_expression_with_context(ctx, ctx.expression_context_clone())
544    }
545
546    /// Analysis of the ORDER BY clause
547    fn parse_order_by_clause(
548        &mut self,
549        ctx: &mut ParseContext,
550    ) -> Result<OrderByClause, ParseError> {
551        let span = ctx.current_span();
552        let mut items = Vec::new();
553
554        loop {
555            let expr = self.parse_expression(ctx)?;
556            let direction = if ctx.match_token(TokenKind::Asc) {
557                OrderDirection::Asc
558            } else if ctx.match_token(TokenKind::Desc) {
559                OrderDirection::Desc
560            } else {
561                OrderDirection::Asc
562            };
563            items.push(OrderByItem {
564                expression: expr,
565                direction,
566            });
567            if !ctx.match_token(TokenKind::Comma) {
568                break;
569            }
570        }
571
572        Ok(OrderByClause { span, items })
573    }
574}
575
576impl Default for UtilStmtParser {
577    fn default() -> Self {
578        Self::new()
579    }
580}