linkrs_query/planning/statements/dml/
assignment_planner.rs

1//! Assignment Statement Planner
2//!
3//! Query planning for handling variable assignment statements.
4//! Supports syntax like: $var = GO FROM 1 OVER friend
5
6use crate::parser::ast::stmt::{AssignmentStmt, Stmt};
7use crate::planning::plan::SubPlan;
8use crate::planning::planner::{Planner, PlannerEnum, PlannerError, ValidatedStatement};
9use crate::QueryContext;
10use std::sync::Arc;
11
12/// Assignment Statement Planner
13/// Responsible for converting variable assignment statements into execution plans.
14/// The assignment statement executes the right-hand side query and binds the result
15/// to the specified variable name in the query context.
16#[derive(Debug, Clone)]
17pub struct AssignmentPlanner;
18
19impl AssignmentPlanner {
20    pub fn new() -> Self {
21        Self
22    }
23
24    fn extract_assignment_stmt<'a>(
25        &self,
26        stmt: &'a Stmt,
27    ) -> Result<&'a AssignmentStmt, PlannerError> {
28        match stmt {
29            Stmt::Assignment(assignment_stmt) => Ok(assignment_stmt),
30            _ => Err(PlannerError::PlanGenerationFailed(
31                "statement does not contain the Assignment".to_string(),
32            )),
33        }
34    }
35}
36
37impl Planner for AssignmentPlanner {
38    fn plan_bound(
39        &mut self,
40        ctx: &crate::planning::context::PlanContext<'_>,
41    ) -> Result<SubPlan, PlannerError> {
42        let bound = ctx.bound;
43        let qctx = ctx.qctx.clone();
44        let metadata = ctx.metadata;
45        let validated = ctx.validated;
46        let assignment = match bound {
47            crate::binder::BoundStatement::Assignment(a) => a,
48            // Legacy AST-only statements still use the AST path.
49            _ => return self.transform(validated, qctx),
50        };
51
52        let mut inner_planner = PlannerEnum::from_bound_statement(&assignment.statement)
53            .ok_or_else(|| {
54                PlannerError::NoSuitablePlanner(format!(
55                    "assignment inner statement: {}",
56                    assignment.statement.kind()
57                ))
58            })?;
59
60        let inner_validated = ctx.derive_validated(&assignment.statement);
61        let inner_ctx = crate::planning::context::PlanContext {
62            bound: &assignment.statement,
63            qctx: qctx.clone(),
64            metadata,
65            validated: inner_validated.as_ref().unwrap_or(validated),
66        };
67        let inner_plan = inner_planner.plan_bound(&inner_ctx)?;
68
69        log::debug!(
70            "AssignmentPlanner: variable '{}' bound to inner plan",
71            assignment.variable
72        );
73
74        Ok(inner_plan)
75    }
76
77    fn transform(
78        &mut self,
79        validated: &ValidatedStatement,
80        qctx: Arc<QueryContext>,
81    ) -> Result<SubPlan, PlannerError> {
82        let assignment_stmt = self.extract_assignment_stmt(validated.stmt())?;
83
84        let inner_validated = ValidatedStatement::new(
85            Arc::new(crate::parser::ast::stmt::Ast::new(
86                (*assignment_stmt.statement).clone(),
87                validated.ast.expr_context().clone(),
88            )),
89            validated.validation_info.clone(),
90        );
91
92        let mut inner_planner =
93            PlannerEnum::from_stmt_ref(inner_validated.stmt()).ok_or_else(|| {
94                PlannerError::NoSuitablePlanner(format!(
95                    "assignment inner statement: {}",
96                    assignment_stmt.statement.kind()
97                ))
98            })?;
99
100        let inner_plan = inner_planner.transform(&inner_validated, qctx)?;
101
102        log::debug!(
103            "AssignmentPlanner: variable '{}' bound to inner plan",
104            assignment_stmt.variable
105        );
106
107        Ok(inner_plan)
108    }
109
110    fn match_planner(&self, stmt: &Stmt) -> bool {
111        matches!(stmt, Stmt::Assignment(_))
112    }
113}
114
115impl Default for AssignmentPlanner {
116    fn default() -> Self {
117        Self::new()
118    }
119}
linkrs_query/planning/statements/dml/assignment_planner.rs

linkrs_query/planning/statements/dml/
assignment_planner.rs
