linkrs-query 0.1.0

//! Assignment Statement Planner
//!
//! Query planning for handling variable assignment statements.
//! Supports syntax like: $var = GO FROM 1 OVER friend

use crate::parser::ast::stmt::{AssignmentStmt, Stmt};
use crate::planning::plan::SubPlan;
use crate::planning::planner::{Planner, PlannerEnum, PlannerError, ValidatedStatement};
use crate::QueryContext;
use std::sync::Arc;

/// Assignment Statement Planner
/// Responsible for converting variable assignment statements into execution plans.
/// The assignment statement executes the right-hand side query and binds the result
/// to the specified variable name in the query context.
#[derive(Debug, Clone)]
pub struct AssignmentPlanner;

impl AssignmentPlanner {
    pub fn new() -> Self {
        Self
    }

    fn extract_assignment_stmt<'a>(
        &self,
        stmt: &'a Stmt,
    ) -> Result<&'a AssignmentStmt, PlannerError> {
        match stmt {
            Stmt::Assignment(assignment_stmt) => Ok(assignment_stmt),
            _ => Err(PlannerError::PlanGenerationFailed(
                "statement does not contain the Assignment".to_string(),
            )),
        }
    }
}

impl Planner for AssignmentPlanner {
    fn plan_bound(
        &mut self,
        ctx: &crate::planning::context::PlanContext<'_>,
    ) -> Result<SubPlan, PlannerError> {
        let bound = ctx.bound;
        let qctx = ctx.qctx.clone();
        let metadata = ctx.metadata;
        let validated = ctx.validated;
        let assignment = match bound {
            crate::binder::BoundStatement::Assignment(a) => a,
            // Legacy AST-only statements still use the AST path.
            _ => return self.transform(validated, qctx),
        };

        let mut inner_planner = PlannerEnum::from_bound_statement(&assignment.statement)
            .ok_or_else(|| {
                PlannerError::NoSuitablePlanner(format!(
                    "assignment inner statement: {}",
                    assignment.statement.kind()
                ))
            })?;

        let inner_validated = ctx.derive_validated(&assignment.statement);
        let inner_ctx = crate::planning::context::PlanContext {
            bound: &assignment.statement,
            qctx: qctx.clone(),
            metadata,
            validated: inner_validated.as_ref().unwrap_or(validated),
        };
        let inner_plan = inner_planner.plan_bound(&inner_ctx)?;

        log::debug!(
            "AssignmentPlanner: variable '{}' bound to inner plan",
            assignment.variable
        );

        Ok(inner_plan)
    }

    fn transform(
        &mut self,
        validated: &ValidatedStatement,
        qctx: Arc<QueryContext>,
    ) -> Result<SubPlan, PlannerError> {
        let assignment_stmt = self.extract_assignment_stmt(validated.stmt())?;

        let inner_validated = ValidatedStatement::new(
            Arc::new(crate::parser::ast::stmt::Ast::new(
                (*assignment_stmt.statement).clone(),
                validated.ast.expr_context().clone(),
            )),
            validated.validation_info.clone(),
        );

        let mut inner_planner =
            PlannerEnum::from_stmt_ref(inner_validated.stmt()).ok_or_else(|| {
                PlannerError::NoSuitablePlanner(format!(
                    "assignment inner statement: {}",
                    assignment_stmt.statement.kind()
                ))
            })?;

        let inner_plan = inner_planner.transform(&inner_validated, qctx)?;

        log::debug!(
            "AssignmentPlanner: variable '{}' bound to inner plan",
            assignment_stmt.variable
        );

        Ok(inner_plan)
    }

    fn match_planner(&self, stmt: &Stmt) -> bool {
        matches!(stmt, Stmt::Assignment(_))
    }
}

impl Default for AssignmentPlanner {
    fn default() -> Self {
        Self::new()
    }
}