nodedb-lite 0.1.0

// SPDX-License-Identifier: Apache-2.0

//! `UPDATE ... FROM` lowering.

use std::collections::HashMap;

use nodedb_physical::physical_plan::document::types::UpdateValue;
use nodedb_sql::UpdateFromVisitArgs;
use nodedb_sql::types::SqlPlan;
use nodedb_sql::types::filter::Filter;
use nodedb_types::result::QueryResult;
use nodedb_types::value::Value;

use crate::error::LiteError;
use crate::query::assignments::{
    convert_assignments_with_source, joined_row, literal_updates, resolve_updates,
};
use crate::query::document_ops::sets::{collect_ids_pub, fetch_document_value_pub};
use crate::query::document_ops::writes::point_update_admitted;
use crate::query::engine::LiteQueryEngine;
use crate::query::value_utils::value_to_string;
use crate::storage::engine::StorageEngine;

use crate::query::visitor::adapter::LiteFut;
use crate::query::visitor::scan_post::filter_mask;

use super::rows::result_to_maps;

// ── UpdateFrom ───────────────────────────────────────────────────────────────

/// `UPDATE target SET ... FROM source WHERE target.col = source.col`.
///
/// Each assignment evaluates against the matched target row, with source
/// fields readable by their qualified name. Every post-image is resolved
/// before the first write, so an assignment that fails leaves every row
/// unchanged.
pub(in crate::query::visitor) fn lower_update_from<'a, S: StorageEngine + 'a>(
    engine: &'a LiteQueryEngine<S>,
    permit: &'a crate::engine::fts::coordinator::TextMutationPermit,
    args: UpdateFromVisitArgs<'_>,
) -> Result<LiteFut<'a>, LiteError> {
    let UpdateFromVisitArgs {
        collection,
        engine: _,
        source,
        source_alias,
        target_join_col,
        source_join_col,
        assignments,
        target_filters,
        returning: _,
    } = args;
    let target = collection.to_string();
    let source_refs = source_refs(source_alias, source);
    let source_alias = source_alias.to_string();
    let source = source.clone();
    let t_join = target_join_col.to_string();
    let s_join = source_join_col.to_string();
    let updates = convert_assignments_with_source(assignments, &source_refs)?;
    let target_filters = target_filters.to_vec();

    Ok(Box::pin(async move {
        let source_result = engine.execute_plan_admitted(&source, Some(permit)).await?;
        let source_maps = result_to_maps(source_result);

        let mut source_index: HashMap<String, HashMap<String, Value>> = HashMap::new();
        for row in source_maps {
            if let Some(key_val) = row.get(&s_join) {
                source_index.insert(value_to_string(key_val), row);
            }
        }

        let target_ids = collect_ids_pub(engine, &target).await?;
        let mut planned: Vec<(String, Vec<(String, UpdateValue)>)> = Vec::new();
        for doc_id in &target_ids {
            let target_val = fetch_document_value_pub(engine, &target, doc_id).await?;
            if !target_matches(&target_val, &target_filters)? {
                continue;
            }
            let Some(join_key) = target_val.get(&t_join).map(value_to_string) else {
                continue;
            };
            let Some(source_val) = source_index.get(&join_key) else {
                continue;
            };
            let row = update_from_row(&target_val, source_val, &source_alias);
            let resolved = resolve_updates(&updates, &Value::Object(row))?;
            planned.push((doc_id.clone(), literal_updates(resolved)?));
        }

        for (doc_id, resolved) in &planned {
            point_update_admitted(engine, permit, &target, doc_id, resolved).await?;
        }

        Ok(QueryResult {
            columns: Vec::new(),
            rows: Vec::new(),
            rows_affected: planned.len() as u64,
            command: Some("UPDATE".into()),
        })
    }))
}

/// The names a column qualifier can use for the source relation. The
/// statement's name for the source comes first: it is the key prefix
/// [`joined_row`] gives each source field. A scan source also answers to its
/// alias and its collection.
fn source_refs(source_alias: &str, plan: &SqlPlan) -> Vec<String> {
    let scan_names: Vec<String> = match plan {
        SqlPlan::Scan {
            collection, alias, ..
        } => alias
            .iter()
            .cloned()
            .chain(std::iter::once(collection.clone()))
            .collect(),
        SqlPlan::TimeseriesScan(scan) => vec![scan.collection.clone()],
        _ => Vec::new(),
    };
    let mut refs = vec![source_alias.to_string()];
    refs.extend(scan_names.into_iter().filter(|name| name != source_alias));
    refs
}

/// Whether a target row satisfies the WHERE conjuncts that read only the
/// target. Every row passes when `filters` is empty.
fn target_matches(row: &HashMap<String, Value>, filters: &[Filter]) -> Result<bool, LiteError> {
    if filters.is_empty() {
        return Ok(true);
    }
    let (columns, values): (Vec<String>, Vec<Value>) =
        row.iter().map(|(k, v)| (k.clone(), v.clone())).unzip();
    let single = QueryResult {
        columns,
        rows: vec![values],
        rows_affected: 0,
        command: None,
    };
    Ok(filter_mask(&single, filters)?
        .first()
        .copied()
        .unwrap_or(false))
}

/// The row an `UPDATE ... FROM` assignment evaluates against: the
/// [`joined_row`], plus each source field the target lacks under its bare
/// name. An unqualified column the target does not hold names the source.
fn update_from_row(
    target: &HashMap<String, Value>,
    source: &HashMap<String, Value>,
    source_alias: &str,
) -> HashMap<String, Value> {
    let mut row = joined_row(target, source, source_alias);
    for (field, value) in source {
        row.entry(field.clone()).or_insert_with(|| value.clone());
    }
    row
}