uqa-planner 0.5.3

//
// Unified Query Algebra
//
// Copyright (c) 2023-2026 Cognica, Inc.
//

//! Select prepared variants and publish usage only against the captured logical identity.

use std::sync::Arc;
use uqa_sql::{
    plan::UnifiedPlan,
    prepared::{
        definition::{PreparedDefinition, PreparedDefinitionContext},
        entry::PreparedStatementPlan,
        planning::{PreparedPlanAnalysis, PreparedPlanSelection, PreparedPlanUpdate},
    },
    SQLError, SQLParam,
};

#[cfg(test)]
mod tests;

pub trait PreparedPlanSession {
    fn prepared_entry(&self, name: &str) -> Option<PreparedStatementPlan>;
    fn plan_cache_mode(&self) -> Result<String, SQLError>;
    fn publish_usage(
        &self,
        name: &str,
        logical_plan: &Arc<UnifiedPlan>,
        update: PreparedPlanUpdate,
    );
}

pub trait PreparedPlanOptimization {
    fn optimize_plan(&self, plan: UnifiedPlan) -> Result<UnifiedPlan, SQLError>;
    fn estimate_plan(&self, plan: &UnifiedPlan) -> Result<crate::plan_cost::PlanCost, SQLError>;
}

pub struct PreparedPlanningContext<'a> {
    pub session: &'a dyn PreparedPlanSession,
    pub analysis: PreparedDefinitionContext<'a>,
    pub optimization: &'a dyn PreparedPlanOptimization,
}

fn needs_analysis(
    context: &PreparedDefinitionContext<'_>,
    entry: &PreparedStatementPlan,
) -> Result<bool, SQLError> {
    if entry.needs_analysis {
        return Ok(true);
    }
    Ok(!uqa_sql::prepared::definition::analysis_is_current(
        context,
        entry.effective_search_path.as_ref(),
        &entry.dependencies,
        entry.dependency_snapshot.as_ref(),
    )?)
}

pub fn select_plan(
    context: &PreparedPlanningContext<'_>,
    name: &str,
    parameters: &[SQLParam],
) -> Result<Option<UnifiedPlan>, SQLError> {
    let Some(entry) = context.session.prepared_entry(name) else {
        return Ok(None);
    };
    let selected = select_entry(context, &entry, parameters)?;
    context
        .session
        .publish_usage(name, &entry.logical_plan, selected.update);
    Ok(Some(selected.plan))
}

/// Select a variant for an already captured definition, including routine-owned
/// entries that are not registered as named prepared statements.
pub fn select_entry(
    context: &PreparedPlanningContext<'_>,
    entry: &PreparedStatementPlan,
    parameters: &[SQLParam],
) -> Result<PreparedPlanSelection, SQLError> {
    select(context, entry, parameters, true)
}

/// A routine's owner has already checked or rebuilt its analyzed inputs and
/// result contract. Reuse only the plan policy here; named PREPARE descriptor
/// checks do not apply to procedural CALL and utility statements.
pub fn select_analyzed_entry(
    context: &PreparedPlanningContext<'_>,
    entry: &PreparedStatementPlan,
    parameters: &[SQLParam],
) -> Result<PreparedPlanSelection, SQLError> {
    select(context, entry, parameters, false)
}

fn select(
    context: &PreparedPlanningContext<'_>,
    entry: &PreparedStatementPlan,
    parameters: &[SQLParam],
    check_analysis: bool,
) -> Result<PreparedPlanSelection, SQLError> {
    let mode = context.session.plan_cache_mode()?;
    let needs_analysis = check_analysis && needs_analysis(&context.analysis, entry)?;
    let mut generic_plan = if needs_analysis {
        None
    } else {
        entry.plan.clone()
    };
    let mut generic_cost = entry.generic_cost;
    let reanalyzed = needs_analysis
        .then(|| reanalyze(context, entry))
        .transpose()?;
    let logical_plan = reanalyzed
        .as_ref()
        .map_or(entry.logical_plan.as_ref(), |definition| {
            &definition.logical_plan
        });
    let inputs = reanalyzed
        .as_ref()
        .map_or(&entry.composite_inputs, |definition| {
            &definition.composite_inputs
        });
    if generic_plan.is_some() && inputs.generic_requires_rebuild(context.analysis.types)? {
        generic_plan = None;
        generic_cost = None;
    }
    let usage = super::PreparedPlanUsage {
        has_parameters: !entry.parameter_types.is_empty(),
        custom_plans: entry.custom_plans,
        total_custom_cost: entry.total_custom_cost,
    };
    let mut custom = super::choose_custom_plan(usage, &mode, generic_cost);
    let generic_projection = !custom;
    let source_logical_plan = logical_plan;
    let projected = if custom {
        inputs.project_logical(logical_plan, context.analysis.types)?
    } else if generic_plan.is_none() {
        inputs.project_for_generic(logical_plan, context.analysis.types)?
    } else {
        None
    };
    let logical_plan = projected.as_ref().unwrap_or(logical_plan);
    let mut composite_inputs = inputs.clone();
    if check_analysis && (custom || generic_plan.is_none() || reanalyzed.is_some()) {
        validate_result_schema(context, entry, logical_plan, reanalyzed.as_ref())?;
    }
    if !custom && generic_plan.is_none() {
        let plan = context.optimization.optimize_plan(logical_plan.clone())?;
        composite_inputs = inputs.with_generic(&plan, context.analysis.types)?;
        generic_cost = Some(context.optimization.estimate_plan(&plan)?.execution);
        generic_plan = Some(plan);
        // Building the first generic plan supplies its previously unknown cost.
        // Recheck before execution so an expensive generic plan is never used just to measure it.
        custom = super::choose_custom_plan(usage, &mode, generic_cost);
    }
    let (plan, custom_cost) = if custom {
        let mut plan = if generic_projection {
            inputs
                .project_logical(source_logical_plan, context.analysis.types)?
                .unwrap_or_else(|| source_logical_plan.clone())
        } else {
            logical_plan.clone()
        };
        super::specialize_parameters(&mut plan, parameters);
        let plan = context.optimization.optimize_plan(plan)?;
        let cost = context
            .optimization
            .estimate_plan(&plan)?
            .including_planning(&crate::CostEstimator::default());
        (plan, Some(cost))
    } else {
        let cached = generic_plan.as_ref().expect("generic plan was built");
        (
            composite_inputs
                .project_generic(cached, context.analysis.types)?
                .unwrap_or_else(|| cached.clone()),
            None,
        )
    };
    Ok(PreparedPlanSelection {
        plan,
        update: PreparedPlanUpdate {
            composite_inputs,
            reanalyzed: reanalyzed.map(|definition| PreparedPlanAnalysis {
                logical_plan: Arc::new(definition.logical_plan),
                effective_search_path: definition.effective_search_path,
                dependencies: definition.dependencies,
                dependency_snapshot: definition.dependency_snapshot,
            }),
            generic_plan,
            generic_cost,
            custom_cost,
        },
    })
}

fn validate_result_schema(
    context: &PreparedPlanningContext<'_>,
    entry: &PreparedStatementPlan,
    logical_plan: &UnifiedPlan,
    reanalyzed: Option<&PreparedDefinition>,
) -> Result<(), SQLError> {
    let result_schema = match reanalyzed {
        Some(definition) => definition.result_schema.clone(),
        None => uqa_sql::prepared::definition::analyze_result_schema(
            &context.analysis,
            logical_plan,
            &entry.parameter_types,
        )?,
    };
    if !uqa_sql::prepared::prepared_result_schema_matches(
        entry.result_schema.as_ref(),
        result_schema.as_ref(),
    ) {
        return Err(SQLError::Routine {
            sqlstate: "0A000".into(),
            message: "cached plan must not change result type".into(),
        });
    }
    Ok(())
}

fn reanalyze(
    context: &PreparedPlanningContext<'_>,
    entry: &PreparedStatementPlan,
) -> Result<uqa_sql::prepared::definition::PreparedDefinition, SQLError> {
    let declared = entry
        .parameter_types
        .iter()
        .flatten()
        .cloned()
        .collect::<Vec<_>>();
    uqa_sql::prepared::definition::analyze_definition(
        &context.analysis,
        (*entry.source_plan).clone(),
        &declared,
    )
}