linkrs_query/planning/statements/dql/
assign_variable_planner.rs1use 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#[derive(Debug, Clone)]
22pub struct AssignVariablePlanner;
23
24impl AssignVariablePlanner {
25 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}