Skip to main content

linkrs_query/planning/statements/dql/
assign_variable_planner.rs

1//! LET (`AssignVariable`) Statement Planner
2//!
3//! Plans a `LET [$]name = expr` session-variable assignment as a single-row,
4//! single-column expression evaluation plan (Start -> Project), reusing the
5//! RETURN expression evaluation chain. The API layer evaluates the plan and
6//! stores the produced value in the session; the query engine itself never
7//! touches session state.
8
9use crate::binder::BoundStatement;
10use crate::parser::ast::stmt::Stmt;
11use crate::planning::plan::core::nodes::{ProjectNode, StartNode};
12use crate::planning::plan::{PlanNodeEnum, SubPlan};
13use crate::planning::planner::{Planner, PlannerError, ValidatedStatement};
14use crate::planning::statements::plan_combiner::{logical_start_root, wrap_logical_project};
15use crate::QueryContext;
16use linkrs_core::YieldColumn;
17use std::sync::Arc;
18
19/// LET statement planner
20/// Converts the LET statement into an execution plan.
21#[derive(Debug, Clone)]
22pub struct AssignVariablePlanner;
23
24impl AssignVariablePlanner {
25    /// Create a new LET planner.
26    pub fn new() -> Self {
27        Self
28    }
29}
30
31impl Planner for AssignVariablePlanner {
32    fn transform(
33        &mut self,
34        validated: &ValidatedStatement,
35        _qctx: Arc<QueryContext>,
36    ) -> Result<SubPlan, PlannerError> {
37        let assign = match validated.stmt() {
38            Stmt::AssignVariable(assign) => assign,
39            _ => {
40                return Err(PlannerError::PlanGenerationFailed(
41                    "statement does not contain a LET assignment".to_string(),
42                ));
43            }
44        };
45
46        let start_node = StartNode::new();
47        let current_node = PlanNodeEnum::Start(start_node.clone());
48
49        let yield_column = YieldColumn {
50            expression: assign.expression.clone(),
51            alias: assign.name.clone(),
52        };
53        let project_node = ProjectNode::new(current_node.clone(), vec![yield_column.clone()])
54            .map_err(|e| {
55                PlannerError::PlanGenerationFailed(format!("Failed to create ProjectNode: {}", e))
56            })?;
57        let current_node = PlanNodeEnum::Project(project_node);
58        let logical_root = wrap_logical_project(
59            logical_start_root(),
60            vec![yield_column],
61            current_node.col_names().to_vec(),
62        );
63
64        let sub_plan = SubPlan {
65            root: Some(current_node),
66            tail: Some(PlanNodeEnum::Start(start_node)),
67            logical_root: Some(logical_root),
68        };
69        Ok(sub_plan)
70    }
71
72    fn plan_bound(
73        &mut self,
74        ctx: &crate::planning::context::PlanContext<'_>,
75    ) -> Result<SubPlan, PlannerError> {
76        let bound = ctx.bound;
77        let qctx = ctx.qctx.clone();
78        let metadata = ctx.metadata;
79        let validated = ctx.validated;
80        let _ = (&bound, &qctx, &metadata, &validated);
81        let assign = match bound {
82            BoundStatement::AssignVariable(a) => a,
83            _ => {
84                return Err(PlannerError::PlanGenerationFailed(
85                    "statement does not contain a LET assignment".to_string(),
86                ));
87            }
88        };
89
90        let expr_ctx = Arc::new(
91            linkrs_core::types::expr::expression_context::ExpressionAnalysisContext::new(),
92        );
93        let ctx_expr =
94            crate::binder::expr_converter::bound_expr_to_contextual(&assign.expression, &expr_ctx)
95                .map_err(PlannerError::PlanGenerationFailed)?;
96
97        let start_node = StartNode::new();
98        let current_node = PlanNodeEnum::Start(start_node.clone());
99
100        let yield_column = YieldColumn {
101            expression: ctx_expr,
102            alias: assign.name.clone(),
103        };
104        let project_node = ProjectNode::new(current_node.clone(), vec![yield_column.clone()])
105            .map_err(|e| {
106                PlannerError::PlanGenerationFailed(format!("Failed to create ProjectNode: {}", e))
107            })?;
108        let current_node = PlanNodeEnum::Project(project_node);
109        let logical_root = wrap_logical_project(
110            logical_start_root(),
111            vec![yield_column],
112            current_node.col_names().to_vec(),
113        );
114
115        let sub_plan = SubPlan {
116            root: Some(current_node),
117            tail: Some(PlanNodeEnum::Start(start_node)),
118            logical_root: Some(logical_root),
119        };
120        Ok(sub_plan)
121    }
122
123    fn match_planner(&self, stmt: &Stmt) -> bool {
124        matches!(stmt, Stmt::AssignVariable(_))
125    }
126}
127
128impl Default for AssignVariablePlanner {
129    fn default() -> Self {
130        Self::new()
131    }
132}