dynoxide/partiql/

executor.rs

//! PartiQL statement executor.
//!
//! Maps parsed PartiQL statements to internal DynamoDB operations.

use crate::errors::{DynoxideError, Result};
use crate::partiql::parser::{
    CompOp, PartiqlValue, ReturningVariant, SetValue, Statement, WhereClause, WhereCondition,
};
use crate::storage_backend::StorageBackend;
use crate::types::{AttributeValue, Item};
use std::collections::HashMap;

/// Execute a parsed PartiQL statement.
///
/// Returns `Some(items)` for SELECT (may be empty) and for a DELETE or UPDATE
/// carrying a `RETURNING` clause (the deleted item or the requested projection);
/// `None` for a write with no `RETURNING` clause. An optional `limit` restricts
/// how many items a SELECT returns.
pub async fn execute<S: StorageBackend>(
    storage: &S,
    stmt: &Statement,
    parameters: &[AttributeValue],
    limit: Option<usize>,
) -> Result<Option<Vec<Item>>> {
    Ok(execute_measured(storage, stmt, parameters, limit).await?.0)
}

/// Like [`execute`], but also returns the total item byte size the statement
/// touched, for `ConsumedCapacity` accounting. SELECT reports the summed size of
/// the rows returned; INSERT/UPDATE/DELETE report the affected item's size (0
/// when the statement was a no-op, e.g. a missing DELETE target).
pub async fn execute_measured<S: StorageBackend>(
    storage: &S,
    stmt: &Statement,
    parameters: &[AttributeValue],
    limit: Option<usize>,
) -> Result<(Option<Vec<Item>>, usize)> {
    match stmt {
        Statement::Select {
            table_name,
            projections,
            where_clause,
        } => {
            let items = execute_select(
                storage,
                table_name,
                projections,
                where_clause.as_ref(),
                parameters,
                limit,
            )
            .await?;
            let size = items
                .as_ref()
                .map(|rows| rows.iter().map(crate::types::item_size).sum())
                .unwrap_or(0);
            Ok((items, size))
        }
        Statement::Insert {
            table_name,
            item,
            if_not_exists,
        } => {
            let size =
                execute_insert(storage, table_name, item, parameters, *if_not_exists).await?;
            Ok((None, size))
        }
        Statement::Update {
            table_name,
            set_clauses,
            remove_paths,
            where_clause,
            returning,
        } => {
            let (projection, size) = execute_update(
                storage,
                table_name,
                set_clauses,
                remove_paths,
                where_clause.as_ref(),
                parameters,
                *returning,
            )
            .await?;
            // RETURNING surfaces the requested projection of the updated item;
            // without a clause an UPDATE returns no items. An empty MODIFIED
            // projection surfaces as a present but empty Items array (no row),
            // matching DynamoDB, rather than a row holding an empty object.
            let items = projection.map(|item| {
                if item.is_empty() {
                    Vec::new()
                } else {
                    vec![item]
                }
            });
            Ok((items, size))
        }
        Statement::Delete {
            table_name,
            where_clause,
            returning,
        } => {
            // DynamoDB permits only RETURNING ALL OLD * on DELETE; the other
            // well-formed variants are rejected with a ValidationException whose
            // message echoes the offending variant.
            if let Some(variant) = returning {
                if *variant != ReturningVariant::AllOld {
                    return Err(DynoxideError::ValidationException(format!(
                        "Invalid returning clause: RETURNING {} *. Only RETURNING ALL OLD * is allowed in DELETE statements.",
                        variant.as_sql()
                    )));
                }
            }
            let (old_item, size) =
                execute_delete(storage, table_name, where_clause.as_ref(), parameters).await?;
            // RETURNING ALL OLD * always surfaces an Items array: the deleted
            // item on a hit, an empty array on a miss (a no-op success). This
            // differs from the classic DeleteItem ReturnValues path, which omits
            // Attributes on a miss.
            let items = if returning.is_some() {
                Some(old_item.map(|item| vec![item]).unwrap_or_default())
            } else {
                None
            };
            Ok((items, size))
        }
    }
}

/// Insert a projected value into a result item.
///
/// For dotted paths (e.g. `a.b.c`), DynamoDB PartiQL returns the resolved value
/// keyed by the leaf segment name (`c`), not the full path or reconstructed
/// nested structure. For simple paths and array index paths, the key is used as-is.
fn insert_nested_projection(result: &mut Item, path: &str, val: AttributeValue) {
    let parts: Vec<&str> = path.split('.').collect();
    // Use the leaf segment as the key
    let key = parts.last().unwrap();
    result.insert(key.to_string(), val);
}

async fn execute_select<S: StorageBackend>(
    storage: &S,
    table_name: &str,
    projections: &[String],
    where_clause: Option<&WhereClause>,
    parameters: &[AttributeValue],
    limit: Option<usize>,
) -> Result<Option<Vec<Item>>> {
    let meta = require_table(storage, table_name).await?;
    let key_schema = crate::actions::helpers::parse_key_schema(&meta)?;

    // Check for COUNT(*) projection
    if projections.len() == 1 && projections[0] == "COUNT(*)" {
        let items = collect_matching_items(
            storage,
            table_name,
            where_clause,
            parameters,
            &key_schema,
            None,
        )
        .await?;
        let count = items.len();
        let mut result = HashMap::new();
        result.insert("Count".to_string(), AttributeValue::N(count.to_string()));
        return Ok(Some(vec![result]));
    }

    let items = collect_matching_items(
        storage,
        table_name,
        where_clause,
        parameters,
        &key_schema,
        limit,
    )
    .await?;

    // Apply projections
    let items = if projections.is_empty() {
        items
    } else {
        items
            .into_iter()
            .map(|item| {
                let mut projected = HashMap::new();
                for proj in projections {
                    if let Some(val) = resolve_nested_path(&item, proj) {
                        insert_nested_projection(&mut projected, proj, val.clone());
                    }
                }
                projected
            })
            .collect()
    };

    Ok(Some(items))
}

/// Collect items that match the WHERE clause, optionally limited.
async fn collect_matching_items<S: StorageBackend>(
    storage: &S,
    table_name: &str,
    where_clause: Option<&WhereClause>,
    parameters: &[AttributeValue],
    key_schema: &crate::actions::helpers::KeySchema,
    limit: Option<usize>,
) -> Result<Vec<Item>> {
    // Try to use Query if the WHERE clause constrains the partition key
    let pk_condition = where_clause.and_then(|wc| find_pk_condition(wc, &key_schema.partition_key));

    let items: Vec<Item> = if let Some(pk_cond) = pk_condition {
        let pk_val = resolve_value(&pk_cond.value, parameters)?;
        let pk_str = pk_val
            .to_key_string()
            .ok_or_else(|| DynoxideError::ValidationException("Invalid key value".to_string()))?;

        let rows = storage
            .query_items(table_name, &pk_str, &Default::default())
            .await?;

        let iter = rows
            .into_iter()
            .filter_map(|(_, _, json)| serde_json::from_str::<Item>(&json).ok())
            .filter(|item| matches_where(item, where_clause, parameters));

        if let Some(lim) = limit {
            iter.take(lim).collect()
        } else {
            iter.collect()
        }
    } else {
        let rows = storage.scan_items(table_name, &Default::default()).await?;

        let iter = rows
            .into_iter()
            .filter_map(|(_, _, json)| serde_json::from_str::<Item>(&json).ok())
            .filter(|item| matches_where(item, where_clause, parameters));

        if let Some(lim) = limit {
            iter.take(lim).collect()
        } else {
            iter.collect()
        }
    };

    Ok(items)
}

/// Find a partition key equality condition, searching across all OR groups.
fn find_pk_condition<'a>(
    wc: &'a WhereClause,
    pk_name: &str,
) -> Option<&'a crate::partiql::parser::Condition> {
    // Only optimise to a Query when there is a single OR group
    // (multi-group OR with pk in only one group would need a union approach).
    if wc.groups.len() == 1 {
        wc.groups[0].iter().find_map(|c| match c {
            WhereCondition::Comparison(cond) if cond.path == pk_name && cond.op == CompOp::Eq => {
                Some(cond)
            }
            _ => None,
        })
    } else {
        None
    }
}

/// Returns the inserted item's size in bytes (0 when an `if_not_exists`
/// duplicate makes the insert a no-op), for `ConsumedCapacity` accounting.
async fn execute_insert<S: StorageBackend>(
    storage: &S,
    table_name: &str,
    item_template: &HashMap<String, PartiqlValue>,
    parameters: &[AttributeValue],
    if_not_exists: bool,
) -> Result<usize> {
    // Resolve any parameter placeholders in the item
    let mut item = HashMap::new();
    for (k, v) in item_template {
        let resolved = match v {
            PartiqlValue::Literal(av) => av.clone(),
            PartiqlValue::Parameter(idx) => parameters.get(*idx).cloned().ok_or_else(|| {
                DynoxideError::ValidationException(format!(
                    "Parameter index {idx} out of range (have {} parameters)",
                    parameters.len()
                ))
            })?,
        };
        item.insert(k.clone(), resolved);
    }

    let meta = require_table(storage, table_name).await?;
    let key_schema = crate::actions::helpers::parse_key_schema(&meta)?;

    // Validate keys present
    crate::actions::helpers::validate_item_keys(&item, &key_schema, &meta)?;
    crate::validation::validate_item_attribute_values(&item)?;

    // Deduplicate sets
    crate::validation::normalize_item_sets(&mut item);

    // TODO: validation must precede this call -- if reaching this line, caller has already validated keys.
    let (pk, sk) = crate::actions::helpers::extract_key_strings(&item, &key_schema)?;

    // PartiQL INSERT must reject duplicates (unlike PutItem which overwrites)
    let existing = storage.get_item(table_name, &pk, &sk).await?;
    if existing.is_some() {
        if if_not_exists {
            // Silently succeed — no-op
            return Ok(0);
        }
        return Err(DynoxideError::DuplicateItemException(
            "Duplicate primary key exists in table".to_string(),
        ));
    }

    let item_json = serde_json::to_string(&item)
        .map_err(|e| DynoxideError::InternalServerError(e.to_string()))?;
    let item_size = crate::types::item_size(&item);

    let hash_prefix = item
        .get(&key_schema.partition_key)
        .map(crate::storage::compute_hash_prefix)
        .unwrap_or_default();
    let old_json = storage
        .put_item_with_hash(table_name, &pk, &sk, &item_json, item_size, &hash_prefix)
        .await?;

    // GSI maintenance
    let table_sk_attr = key_schema.sort_key.as_deref();
    let _ = crate::actions::gsi::maintain_gsis_after_write(
        storage,
        table_name,
        &meta,
        &pk,
        &sk,
        &item,
        &key_schema.partition_key,
        table_sk_attr,
    )
    .await?;

    // LSI maintenance
    crate::actions::lsi::maintain_lsis_after_write(
        storage,
        table_name,
        &meta,
        &pk,
        &sk,
        &item,
        &key_schema.partition_key,
        table_sk_attr,
    )
    .await?;

    // Stream record
    let old_item: Option<Item> = old_json.as_ref().and_then(|j| serde_json::from_str(j).ok());
    crate::streams::record_stream_event(storage, &meta, old_item.as_ref(), Some(&item)).await?;

    Ok(item_size)
}

/// Applies an UPDATE and returns the `RETURNING` projection (or `None` when the
/// statement carried no `RETURNING` clause) and the updated item's new size in
/// bytes (0 when the update resolves to an empty item and is skipped), for
/// `ConsumedCapacity` accounting.
async fn execute_update<S: StorageBackend>(
    storage: &S,
    table_name: &str,
    set_clauses: &[crate::partiql::parser::SetClause],
    remove_paths: &[String],
    where_clause: Option<&WhereClause>,
    parameters: &[AttributeValue],
    returning: Option<ReturningVariant>,
) -> Result<(Option<Item>, usize)> {
    let meta = require_table(storage, table_name).await?;
    let key_schema = crate::actions::helpers::parse_key_schema(&meta)?;

    // WHERE clause is required for UPDATE to identify the item
    let wc = where_clause.ok_or_else(|| {
        DynoxideError::ValidationException("UPDATE requires a WHERE clause".to_string())
    })?;

    // DynamoDB does not support OR in UPDATE WHERE clauses
    if wc.groups.len() > 1 {
        return Err(DynoxideError::ValidationException(
            "UPDATE does not support OR conditions in WHERE clause".to_string(),
        ));
    }

    // Extract partition key from WHERE (must be in first/only group for key lookup)
    let pk_cond =
        find_comparison_in_groups(&wc.groups, &key_schema.partition_key).ok_or_else(|| {
            DynoxideError::ValidationException(
                "Where clause does not contain a mandatory equality on all key attributes"
                    .to_string(),
            )
        })?;

    let pk_val = resolve_value(&pk_cond.value, parameters)?;
    let pk_str = pk_val
        .to_key_string()
        .ok_or_else(|| DynoxideError::ValidationException("Invalid key value".to_string()))?;

    let sk_str = if let Some(ref sk_name) = key_schema.sort_key {
        let sk_cond = find_comparison_in_groups(&wc.groups, sk_name);
        if sk_cond.is_none() {
            return Err(DynoxideError::ValidationException(
                "Where clause does not contain a mandatory equality on all key attributes"
                    .to_string(),
            ));
        }
        sk_cond
            .map(|c| resolve_value(&c.value, parameters))
            .transpose()?
            .and_then(|v| v.to_key_string())
            .unwrap_or_default()
    } else {
        String::new()
    };

    // Get existing item
    let existing_json = storage.get_item(table_name, &pk_str, &sk_str).await?;
    let mut item: Item = existing_json
        .as_ref()
        .and_then(|j| serde_json::from_str(j).ok())
        .unwrap_or_default();

    let old_item = item.clone();

    // PartiQL UPDATE is not an upsert: the target item must already exist, so a
    // missing item fails ConditionalCheckFailedException and creates nothing.
    // Non-key WHERE predicates act as a further condition on the existing item;
    // if that predicate is false the update fails the same way. Neither writes.
    if existing_json.is_none() || !matches_where(&old_item, where_clause, parameters) {
        return Err(DynoxideError::ConditionalCheckFailedException(
            "The conditional request failed".to_string(),
            None,
        ));
    }

    let before_item = item.clone();

    // Apply SET clauses with nested path support
    for clause in set_clauses {
        let val = resolve_set_value(&clause.value, &item, parameters)?;
        set_nested_value(&mut item, &clause.path, val)?;
    }

    // Apply REMOVE clauses
    for path in remove_paths {
        remove_nested_value(&mut item, path);
    }

    // Ensure keys are present
    if item.is_empty() {
        return Ok((None, 0));
    }

    // Validate attribute values after SET clauses applied
    crate::validation::validate_item_attribute_values(&item)?;
    crate::validation::normalize_item_sets(&mut item);

    // Reject an index key this update set to an invalid value (see helpers).
    crate::actions::helpers::validate_updated_index_keys(&before_item, &item, &meta)?;

    let item_json = serde_json::to_string(&item)
        .map_err(|e| DynoxideError::InternalServerError(e.to_string()))?;
    let item_size = crate::types::item_size(&item);

    let hash_prefix = item
        .get(&key_schema.partition_key)
        .map(crate::storage::compute_hash_prefix)
        .unwrap_or_default();
    storage
        .put_item_with_hash(
            table_name,
            &pk_str,
            &sk_str,
            &item_json,
            item_size,
            &hash_prefix,
        )
        .await?;

    // GSI maintenance
    let table_sk_attr = key_schema.sort_key.as_deref();
    let _ = crate::actions::gsi::maintain_gsis_after_write(
        storage,
        table_name,
        &meta,
        &pk_str,
        &sk_str,
        &item,
        &key_schema.partition_key,
        table_sk_attr,
    )
    .await?;

    // LSI maintenance
    crate::actions::lsi::maintain_lsis_after_write(
        storage,
        table_name,
        &meta,
        &pk_str,
        &sk_str,
        &item,
        &key_schema.partition_key,
        table_sk_attr,
    )
    .await?;

    // Stream record
    let old_ref = if existing_json.is_some() {
        Some(&old_item)
    } else {
        None
    };
    crate::streams::record_stream_event(storage, &meta, old_ref, Some(&item)).await?;

    // Build the RETURNING projection from the item's before/after states. The
    // MODIFIED variants project each touched path (a nested `a.b` yields just
    // the changed leaf, never the whole `a` attribute and never the key),
    // resolved against the relevant item.
    let projection = returning.map(|variant| {
        let modified: std::collections::BTreeSet<String> = set_clauses
            .iter()
            .map(|c| c.path.clone())
            .chain(remove_paths.iter().cloned())
            .collect();
        project_returning(variant, &old_item, &item, &modified)
    });

    Ok((projection, item_size))
}

/// Build the `RETURNING` projection for an UPDATE from the item's before/after
/// states. `ALL` variants return the whole item (key included); `MODIFIED`
/// variants project only the touched paths, resolved against the old item
/// (`MODIFIED OLD *`) or the new item (`MODIFIED NEW *`), which never includes
/// the key. A `MODIFIED` projection can be empty (a path that no longer
/// resolves contributes nothing), which the caller surfaces as an empty `Items`
/// array.
fn project_returning(
    variant: ReturningVariant,
    old_item: &Item,
    new_item: &Item,
    modified: &std::collections::BTreeSet<String>,
) -> Item {
    match variant {
        ReturningVariant::AllOld => old_item.clone(),
        ReturningVariant::AllNew => new_item.clone(),
        ReturningVariant::ModifiedOld => project_modified(modified, old_item),
        ReturningVariant::ModifiedNew => project_modified(modified, new_item),
    }
}

/// An intermediate `MODIFIED` projection node. List positions are collected by
/// their real index in a `BTreeMap` so that, on conversion, the contributed
/// elements emerge as a dense list in ascending index order: DynamoDB does not
/// keep gaps, so `SET a[0], a[2]` projects `{a: [v0, v2]}`, not a sparse list.
enum ProjNode {
    Leaf(AttributeValue),
    Map(HashMap<String, ProjNode>),
    List(std::collections::BTreeMap<usize, ProjNode>),
}

/// Project the touched paths from `source` into a `MODIFIED` projection. Each
/// path is resolved against `source` (navigating map keys and real list
/// indices); a path that resolves contributes its value, one that does not
/// contributes nothing. This is why a map REMOVE yields nothing under
/// `MODIFIED NEW` (the key is gone) while a list REMOVE contributes the
/// shifted-in element (`tags[1]` still resolves after `REMOVE tags[1]`).
/// Contributed list elements pack densely in ascending index order. This is by
/// index, not statement order: for `SET a[2]=.., a[0]=..` real DynamoDB returns
/// `{a: [v0, v2]}`, which the `BTreeMap`-keyed pack matches.
fn project_modified(paths: &std::collections::BTreeSet<String>, source: &Item) -> Item {
    let mut root: HashMap<String, ProjNode> = HashMap::new();
    for path in paths {
        if let (Some(val), Some(segments)) =
            (resolve_nested_path(source, path), split_path_segments(path))
        {
            // The top-level of an item is always addressed by a map key.
            if let Some((PathSegment::Key(key), rest)) = segments.split_first() {
                let node = root
                    .entry((*key).to_string())
                    .or_insert_with(|| fresh_proj_node(rest));
                insert_proj_node(node, rest, val.clone());
            }
        }
    }
    root.into_iter()
        .map(|(k, node)| (k, proj_node_to_value(node)))
        .collect()
}

/// The container a projection node needs for its next segment: a list for an
/// index, a map otherwise. An empty `segments` is a leaf position, overwritten
/// immediately by the value, so the placeholder type there is immaterial.
fn fresh_proj_node(segments: &[PathSegment]) -> ProjNode {
    match segments.first() {
        Some(PathSegment::Index(_)) => ProjNode::List(std::collections::BTreeMap::new()),
        _ => ProjNode::Map(HashMap::new()),
    }
}

/// Insert `val` at `segments` within `node`, creating intermediate map/list
/// nodes as needed.
fn insert_proj_node(node: &mut ProjNode, segments: &[PathSegment], val: AttributeValue) {
    let Some((seg, rest)) = segments.split_first() else {
        *node = ProjNode::Leaf(val);
        return;
    };
    match seg {
        PathSegment::Key(k) => {
            if let ProjNode::Map(map) = node {
                let child = map
                    .entry((*k).to_string())
                    .or_insert_with(|| fresh_proj_node(rest));
                insert_proj_node(child, rest, val);
            }
        }
        PathSegment::Index(i) => {
            if let ProjNode::List(list) = node {
                let child = list.entry(*i).or_insert_with(|| fresh_proj_node(rest));
                insert_proj_node(child, rest, val);
            }
        }
    }
}

/// Convert a projection node into an `AttributeValue`. A `List` node's
/// `BTreeMap` yields its elements in ascending index order, densely.
fn proj_node_to_value(node: ProjNode) -> AttributeValue {
    match node {
        ProjNode::Leaf(v) => v,
        ProjNode::Map(map) => AttributeValue::M(
            map.into_iter()
                .map(|(k, n)| (k, proj_node_to_value(n)))
                .collect(),
        ),
        ProjNode::List(list) => {
            AttributeValue::L(list.into_values().map(proj_node_to_value).collect())
        }
    }
}

/// Returns the deleted item (None when the target was missing and the delete
/// was a no-op) and its size in bytes for `ConsumedCapacity` accounting.
async fn execute_delete<S: StorageBackend>(
    storage: &S,
    table_name: &str,
    where_clause: Option<&WhereClause>,
    parameters: &[AttributeValue],
) -> Result<(Option<Item>, usize)> {
    let meta = require_table(storage, table_name).await?;
    let key_schema = crate::actions::helpers::parse_key_schema(&meta)?;

    let wc = where_clause.ok_or_else(|| {
        DynoxideError::ValidationException("DELETE requires a WHERE clause".to_string())
    })?;

    // DynamoDB does not support OR in DELETE WHERE clauses
    if wc.groups.len() > 1 {
        return Err(DynoxideError::ValidationException(
            "DELETE does not support OR conditions in WHERE clause".to_string(),
        ));
    }

    let pk_cond =
        find_comparison_in_groups(&wc.groups, &key_schema.partition_key).ok_or_else(|| {
            DynoxideError::ValidationException(
                "Where clause does not contain a mandatory equality on all key attributes"
                    .to_string(),
            )
        })?;

    let pk_val = resolve_value(&pk_cond.value, parameters)?;
    let pk_str = pk_val
        .to_key_string()
        .ok_or_else(|| DynoxideError::ValidationException("Invalid key value".to_string()))?;

    // I15: Validate that the sort key is present in the WHERE clause if the table has one
    if let Some(ref sk_name) = key_schema.sort_key {
        let has_sk_condition = wc.groups.iter().any(|group| {
            group.iter().any(|c| match c {
                WhereCondition::Comparison(comp) => comp.path == *sk_name && comp.op == CompOp::Eq,
                _ => false,
            })
        });
        if !has_sk_condition {
            return Err(DynoxideError::ValidationException(
                "Where clause does not contain a mandatory equality on all key attributes"
                    .to_string(),
            ));
        }
    }

    let sk_str = if let Some(ref sk_name) = key_schema.sort_key {
        find_comparison_in_groups(&wc.groups, sk_name)
            .map(|c| resolve_value(&c.value, parameters))
            .transpose()?
            .and_then(|v| v.to_key_string())
            .unwrap_or_default()
    } else {
        String::new()
    };

    // Non-key WHERE predicates act as a condition on the existing item, like a
    // conditional write. AWS raises ConditionalCheckFailedException when the item
    // is present but the condition is false, and a missing item is a silent
    // no-op (the condition is never evaluated). Re-running the full WHERE via
    // matches_where covers both the key equality (always true for the fetched
    // item) and any extra predicates.
    if let Some(json) = storage.get_item(table_name, &pk_str, &sk_str).await? {
        let existing: Item = serde_json::from_str(&json)
            .map_err(|e| DynoxideError::InternalServerError(format!("Bad item JSON: {e}")))?;
        if !matches_where(&existing, where_clause, parameters) {
            return Err(DynoxideError::ConditionalCheckFailedException(
                "The conditional request failed".to_string(),
                None,
            ));
        }
    }

    let old_json = storage.delete_item(table_name, &pk_str, &sk_str).await?;

    // GSI maintenance
    let _ = crate::actions::gsi::maintain_gsis_after_delete(
        storage, table_name, &meta, &pk_str, &sk_str,
    )
    .await?;

    // LSI maintenance
    crate::actions::lsi::maintain_lsis_after_delete(storage, table_name, &meta, &pk_str, &sk_str)
        .await?;

    // Stream record
    let old_item: Option<Item> = old_json.as_ref().and_then(|j| serde_json::from_str(j).ok());
    if old_item.is_some() {
        crate::streams::record_stream_event(storage, &meta, old_item.as_ref(), None).await?;
    }

    // A delete is charged for the size of the item it removed; a no-op delete
    // (missing target) reports 0.
    let deleted_size = old_item.as_ref().map(crate::types::item_size).unwrap_or(0);
    Ok((old_item, deleted_size))
}

// ---------------------------------------------------------------------------
// Helpers
// ---------------------------------------------------------------------------

async fn require_table<S: StorageBackend>(
    storage: &S,
    table_name: &str,
) -> Result<crate::storage::TableMetadata> {
    crate::actions::helpers::require_table(storage, table_name).await
}

/// Find a comparison condition matching a given path with Eq operator,
/// searching across all OR groups.
fn find_comparison_in_groups<'a>(
    groups: &'a [Vec<WhereCondition>],
    path: &str,
) -> Option<&'a crate::partiql::parser::Condition> {
    for group in groups {
        if let Some(cond) = find_comparison(group, path) {
            return Some(cond);
        }
    }
    None
}

/// Find a comparison condition matching a given path with Eq operator.
fn find_comparison<'a>(
    conditions: &'a [WhereCondition],
    path: &str,
) -> Option<&'a crate::partiql::parser::Condition> {
    conditions.iter().find_map(|c| match c {
        WhereCondition::Comparison(cond) if cond.path == path && cond.op == CompOp::Eq => {
            Some(cond)
        }
        _ => None,
    })
}

/// Resolve a PartiqlValue to a concrete AttributeValue.
fn resolve_value(val: &PartiqlValue, parameters: &[AttributeValue]) -> Result<AttributeValue> {
    match val {
        PartiqlValue::Literal(av) => Ok(av.clone()),
        PartiqlValue::Parameter(idx) => parameters.get(*idx).cloned().ok_or_else(|| {
            DynoxideError::ValidationException(format!(
                "Parameter index {idx} out of range (have {} parameters)",
                parameters.len()
            ))
        }),
    }
}

/// Resolve a SetValue to a concrete AttributeValue, potentially using the current item.
fn resolve_set_value(
    val: &SetValue,
    item: &Item,
    parameters: &[AttributeValue],
) -> Result<AttributeValue> {
    match val {
        SetValue::Simple(pv) => resolve_value(pv, parameters),
        SetValue::Add(attr, pv) => {
            let current = resolve_nested_path(item, attr);
            let operand = resolve_value(pv, parameters)?;
            match (current, &operand) {
                (Some(AttributeValue::N(cur)), AttributeValue::N(add)) => {
                    use bigdecimal::BigDecimal;
                    use std::str::FromStr;
                    let a = BigDecimal::from_str(cur).map_err(|e| {
                        DynoxideError::ValidationException(format!("Invalid number: {e}"))
                    })?;
                    let b = BigDecimal::from_str(add).map_err(|e| {
                        DynoxideError::ValidationException(format!("Invalid number: {e}"))
                    })?;
                    let result = a + b;
                    Ok(AttributeValue::N(format_bigdecimal(&result)))
                }
                (None, AttributeValue::N(_)) => {
                    // Attribute doesn't exist yet — use the operand value
                    Ok(operand)
                }
                _ => Err(DynoxideError::ValidationException(
                    "SET expression add requires numeric attribute and operand".to_string(),
                )),
            }
        }
        SetValue::Sub(attr, pv) => {
            let current = resolve_nested_path(item, attr);
            let operand = resolve_value(pv, parameters)?;
            match (current, &operand) {
                (Some(AttributeValue::N(cur)), AttributeValue::N(sub)) => {
                    use bigdecimal::BigDecimal;
                    use std::str::FromStr;
                    let a = BigDecimal::from_str(cur).map_err(|e| {
                        DynoxideError::ValidationException(format!("Invalid number: {e}"))
                    })?;
                    let b = BigDecimal::from_str(sub).map_err(|e| {
                        DynoxideError::ValidationException(format!("Invalid number: {e}"))
                    })?;
                    let result = a - b;
                    Ok(AttributeValue::N(format_bigdecimal(&result)))
                }
                (None, AttributeValue::N(sub)) => {
                    // Attribute doesn't exist yet — treat as 0 - operand
                    use bigdecimal::BigDecimal;
                    use std::str::FromStr;
                    let b = BigDecimal::from_str(sub).map_err(|e| {
                        DynoxideError::ValidationException(format!("Invalid number: {e}"))
                    })?;
                    let result = -b;
                    Ok(AttributeValue::N(format_bigdecimal(&result)))
                }
                _ => Err(DynoxideError::ValidationException(
                    "SET expression subtract requires numeric attribute and operand".to_string(),
                )),
            }
        }
        SetValue::ListAppend(first, second) => {
            let a = resolve_value(first, parameters)?;
            let b = resolve_value(second, parameters)?;
            // At least one should be a list. If an attribute name was given,
            // resolve it from the item.
            let list_a = match &a {
                AttributeValue::S(name) => resolve_nested_path(item, name)
                    .cloned()
                    .unwrap_or(AttributeValue::L(Vec::new())),
                other => other.clone(),
            };
            let list_b = match &b {
                AttributeValue::S(name) => resolve_nested_path(item, name)
                    .cloned()
                    .unwrap_or(AttributeValue::L(Vec::new())),
                other => other.clone(),
            };
            match (list_a, list_b) {
                (AttributeValue::L(mut la), AttributeValue::L(lb)) => {
                    la.extend(lb);
                    Ok(AttributeValue::L(la))
                }
                _ => Err(DynoxideError::ValidationException(
                    "list_append requires list operands".to_string(),
                )),
            }
        }
    }
}

/// The `ValidationException` DynamoDB raises for a document path that cannot be
/// applied by an update (e.g. indexing a scalar, or a missing intermediate).
fn invalid_update_path() -> DynoxideError {
    DynoxideError::ValidationException(
        "The document path provided in the update expression is invalid for update".to_string(),
    )
}

/// Set a value at a document path, navigating both map keys and list indices,
/// so `SET tags[0] = :v` writes the real list element rather than a literal
/// `tags[0]` key. A list index at or beyond the end appends, matching DynamoDB.
/// Intermediate map keys are created if absent, preserving the prior behaviour
/// for dotted map paths.
fn set_nested_value(item: &mut Item, path: &str, val: AttributeValue) -> Result<()> {
    let segments = split_path_segments(path).ok_or_else(invalid_update_path)?;
    let (first, rest) = segments.split_first().ok_or_else(invalid_update_path)?;
    let key = match first {
        PathSegment::Key(k) => (*k).to_string(),
        // The top-level item is a map; it cannot be indexed.
        PathSegment::Index(_) => return Err(invalid_update_path()),
    };
    if rest.is_empty() {
        item.insert(key, val);
        return Ok(());
    }
    let entry = item
        .entry(key)
        .or_insert_with(|| AttributeValue::M(HashMap::new()));
    set_into_value(entry, rest, val)
}

/// Recursive helper for [`set_nested_value`]: apply the remaining path segments
/// to `current`.
fn set_into_value(
    current: &mut AttributeValue,
    segments: &[PathSegment],
    val: AttributeValue,
) -> Result<()> {
    let (seg, rest) = segments.split_first().expect("segments is non-empty");
    if rest.is_empty() {
        return match seg {
            PathSegment::Key(k) => match current {
                AttributeValue::M(map) => {
                    map.insert((*k).to_string(), val);
                    Ok(())
                }
                _ => Err(invalid_update_path()),
            },
            PathSegment::Index(i) => match current {
                AttributeValue::L(list) => {
                    if *i < list.len() {
                        list[*i] = val;
                    } else {
                        list.push(val);
                    }
                    Ok(())
                }
                _ => Err(invalid_update_path()),
            },
        };
    }
    match seg {
        PathSegment::Key(k) => match current {
            AttributeValue::M(map) => {
                let next = map
                    .entry((*k).to_string())
                    .or_insert_with(|| AttributeValue::M(HashMap::new()));
                set_into_value(next, rest, val)
            }
            _ => Err(invalid_update_path()),
        },
        PathSegment::Index(i) => match current {
            AttributeValue::L(list) => match list.get_mut(*i) {
                Some(next) => set_into_value(next, rest, val),
                None => Err(invalid_update_path()),
            },
            _ => Err(invalid_update_path()),
        },
    }
}

/// Remove the value at a document path, navigating both map keys and list
/// indices, so `REMOVE tags[0]` deletes the list element (shifting the rest)
/// rather than a literal `tags[0]` key.
fn remove_nested_value(item: &mut Item, path: &str) {
    let Some(segments) = split_path_segments(path) else {
        return;
    };
    let Some((first, rest)) = segments.split_first() else {
        return;
    };
    let PathSegment::Key(key) = first else {
        return; // the top-level item cannot be indexed
    };
    if rest.is_empty() {
        item.remove(*key);
        return;
    }
    if let Some(current) = item.get_mut(*key) {
        remove_from_value(current, rest);
    }
}

/// Recursive helper for [`remove_nested_value`]. A path that does not exist or
/// whose type does not match is a no-op, mirroring DynamoDB's tolerant REMOVE.
fn remove_from_value(current: &mut AttributeValue, segments: &[PathSegment]) {
    let (seg, rest) = segments.split_first().expect("segments is non-empty");
    if rest.is_empty() {
        match seg {
            PathSegment::Key(k) => {
                if let AttributeValue::M(map) = current {
                    map.remove(*k);
                }
            }
            PathSegment::Index(i) => {
                if let AttributeValue::L(list) = current {
                    if *i < list.len() {
                        list.remove(*i);
                    }
                }
            }
        }
        return;
    }
    match seg {
        PathSegment::Key(k) => {
            if let AttributeValue::M(map) = current {
                if let Some(next) = map.get_mut(*k) {
                    remove_from_value(next, rest);
                }
            }
        }
        PathSegment::Index(i) => {
            if let AttributeValue::L(list) = current {
                if let Some(next) = list.get_mut(*i) {
                    remove_from_value(next, rest);
                }
            }
        }
    }
}

/// Check if an item matches a WHERE clause (with OR-group support).
fn matches_where(
    item: &Item,
    where_clause: Option<&WhereClause>,
    parameters: &[AttributeValue],
) -> bool {
    let wc = match where_clause {
        Some(wc) => wc,
        None => return true,
    };

    // OR semantics: any group matching is sufficient
    wc.groups
        .iter()
        .any(|group| matches_conditions(item, group, parameters))
}

/// Check if an item matches all conditions in a group (AND semantics).
fn matches_conditions(
    item: &Item,
    conditions: &[WhereCondition],
    parameters: &[AttributeValue],
) -> bool {
    for cond in conditions {
        match cond {
            WhereCondition::Comparison(c) => {
                let item_val = match resolve_nested_path(item, &c.path) {
                    Some(v) => v,
                    None => return false,
                };
                let target = match resolve_value(&c.value, parameters) {
                    Ok(v) => v,
                    Err(_) => return false,
                };
                if !compare_values(item_val, &c.op, &target) {
                    return false;
                }
            }
            WhereCondition::Exists(path) | WhereCondition::IsNotMissing(path) => {
                if resolve_nested_path(item, path).is_none() {
                    return false;
                }
            }
            WhereCondition::NotExists(path) | WhereCondition::IsMissing(path) => {
                if resolve_nested_path(item, path).is_some() {
                    return false;
                }
            }
            WhereCondition::BeginsWith(path, prefix_val) => {
                let item_val = match resolve_nested_path(item, path) {
                    Some(v) => v,
                    None => return false,
                };
                let prefix = match resolve_value(prefix_val, parameters) {
                    Ok(v) => v,
                    Err(_) => return false,
                };
                match (item_val, &prefix) {
                    (AttributeValue::S(s), AttributeValue::S(p)) => {
                        if !s.starts_with(p.as_str()) {
                            return false;
                        }
                    }
                    _ => return false,
                }
            }
            WhereCondition::NotBeginsWith(path, prefix_val) => {
                // Logical negation of begins_with: the row matches unless the
                // value is a string that starts with the prefix. A missing or
                // non-string attribute does not begin with the prefix, so it is
                // kept.
                if let Some(item_val) = resolve_nested_path(item, path) {
                    let prefix = match resolve_value(prefix_val, parameters) {
                        Ok(v) => v,
                        Err(_) => return false,
                    };
                    if let (AttributeValue::S(s), AttributeValue::S(p)) = (item_val, &prefix) {
                        if s.starts_with(p.as_str()) {
                            return false;
                        }
                    }
                }
            }
            WhereCondition::Between(path, low, high) => {
                let item_val = match resolve_nested_path(item, path) {
                    Some(v) => v,
                    None => return false,
                };
                let low_val = match resolve_value(low, parameters) {
                    Ok(v) => v,
                    Err(_) => return false,
                };
                let high_val = match resolve_value(high, parameters) {
                    Ok(v) => v,
                    Err(_) => return false,
                };
                if !compare_values(item_val, &CompOp::Ge, &low_val)
                    || !compare_values(item_val, &CompOp::Le, &high_val)
                {
                    return false;
                }
            }
            WhereCondition::In(path, values) => {
                let item_val = match resolve_nested_path(item, path) {
                    Some(v) => v,
                    None => return false,
                };
                let matched = values.iter().any(|v| {
                    resolve_value(v, parameters)
                        .map(|target| compare_values(item_val, &CompOp::Eq, &target))
                        .unwrap_or(false)
                });
                if !matched {
                    return false;
                }
            }
            WhereCondition::Contains(path, substr_val) => {
                let item_val = match resolve_nested_path(item, path) {
                    Some(v) => v,
                    None => return false,
                };
                let substr = match resolve_value(substr_val, parameters) {
                    Ok(v) => v,
                    Err(_) => return false,
                };
                match (item_val, &substr) {
                    (AttributeValue::S(s), AttributeValue::S(sub)) => {
                        if !s.contains(sub.as_str()) {
                            return false;
                        }
                    }
                    (AttributeValue::SS(set), AttributeValue::S(val)) => {
                        if !set.contains(val) {
                            return false;
                        }
                    }
                    (AttributeValue::NS(set), AttributeValue::N(val)) => {
                        if !set.contains(val) {
                            return false;
                        }
                    }
                    (AttributeValue::L(list), target) => {
                        if !list.contains(target) {
                            return false;
                        }
                    }
                    _ => return false,
                }
            }
        }
    }

    true
}

/// Resolve a dotted/indexed path to a nested attribute value.
///
/// Supports paths like `"a"`, `"a.b.c"`, and `"a[0].b"`.
fn resolve_nested_path<'a>(item: &'a Item, path: &str) -> Option<&'a AttributeValue> {
    // Fast path: no dots or brackets means a simple top-level lookup
    if !path.contains('.') && !path.contains('[') {
        return item.get(path);
    }

    let segments = split_path_segments(path)?;
    if segments.is_empty() {
        return None;
    }

    // First segment must be a map key on the top-level item
    let mut current = match &segments[0] {
        PathSegment::Key(k) => item.get(*k)?,
        PathSegment::Index(_) => return None,
    };

    for seg in &segments[1..] {
        current = match seg {
            PathSegment::Key(k) => match current {
                AttributeValue::M(map) => map.get(*k)?,
                _ => return None,
            },
            PathSegment::Index(idx) => match current {
                AttributeValue::L(list) => list.get(*idx)?,
                _ => return None,
            },
        };
    }

    Some(current)
}

enum PathSegment<'a> {
    Key(&'a str),
    Index(usize),
}

/// Split a path like `"a.b[0].c"` into segments.
/// Returns None if the path contains malformed bracket expressions (e.g. `a[xyz]`).
fn split_path_segments(path: &str) -> Option<Vec<PathSegment<'_>>> {
    let mut segments = Vec::new();
    let bytes = path.as_bytes();
    let mut start = 0;
    let mut i = 0;

    while i < bytes.len() {
        match bytes[i] {
            b'.' => {
                if start < i {
                    segments.push(PathSegment::Key(&path[start..i]));
                }
                i += 1;
                start = i;
            }
            b'[' => {
                if start < i {
                    segments.push(PathSegment::Key(&path[start..i]));
                }
                i += 1;
                let idx_start = i;
                while i < bytes.len() && bytes[i] != b']' {
                    i += 1;
                }
                let idx = path[idx_start..i].parse::<usize>().ok()?;
                segments.push(PathSegment::Index(idx));
                if i < bytes.len() {
                    i += 1; // skip ']'
                }
                start = i;
                // Skip a trailing dot after ']' (e.g. `a[0].b`)
                if i < bytes.len() && bytes[i] == b'.' {
                    i += 1;
                    start = i;
                }
            }
            _ => {
                i += 1;
            }
        }
    }

    if start < bytes.len() {
        segments.push(PathSegment::Key(&path[start..]));
    }

    Some(segments)
}

/// Compare two AttributeValues using a comparison operator.
fn compare_values(left: &AttributeValue, op: &CompOp, right: &AttributeValue) -> bool {
    match (left, right) {
        (AttributeValue::S(a), AttributeValue::S(b)) => compare_ord(a, op, b),
        (AttributeValue::N(a), AttributeValue::N(b)) => {
            use bigdecimal::BigDecimal;
            use std::str::FromStr;
            match (BigDecimal::from_str(a), BigDecimal::from_str(b)) {
                (Ok(da), Ok(db)) => compare_ord(&da, op, &db),
                _ => false,
            }
        }
        (AttributeValue::BOOL(a), AttributeValue::BOOL(b)) => match op {
            CompOp::Eq => a == b,
            CompOp::Ne => a != b,
            _ => false,
        },
        _ => match op {
            CompOp::Eq => false,
            CompOp::Ne => true,
            _ => false,
        },
    }
}

/// Format a BigDecimal number, stripping unnecessary trailing zeros.
fn format_bigdecimal(n: &bigdecimal::BigDecimal) -> String {
    let normalized = n.normalized();
    if normalized.as_bigint_and_exponent().1 < 0 {
        normalized.with_scale(0).to_string()
    } else {
        normalized.to_string()
    }
}

fn compare_ord<T: PartialOrd>(a: &T, op: &CompOp, b: &T) -> bool {
    match op {
        CompOp::Eq => a == b,
        CompOp::Ne => a != b,
        CompOp::Lt => a < b,
        CompOp::Le => a <= b,
        CompOp::Gt => a > b,
        CompOp::Ge => a >= b,
    }
}