toasty_driver_integration_suite/tests/

crud_partitioned.rs

use crate::prelude::*;

use toasty_core::{
    driver::{Operation, Rows},
    stmt::{Assignment, Expr, Source, Statement, UpdateTarget},
};

/// Test update on a model with a partitioned composite primary key using the
/// partition-key-only filter.
///
/// `Todo::filter_by_user_id(user_id).update()` uses only the partition key in the
/// filter expression. For NoSQL (DynamoDB), this requires a `QueryPk` to find all
/// matching records and then an `UpdateItem` for each — not just a bare `QueryPk`
/// that silently discards the mutation.
#[driver_test]
pub async fn update_by_partition_key(test: &mut Test) {
    #[derive(Debug, toasty::Model)]
    #[key(partition = user_id, local = id)]
    struct Todo {
        #[auto]
        id: uuid::Uuid,

        user_id: String,

        title: String,
    }

    let mut db = test.setup_db(models!(Todo)).await;

    let todo_table_id = table_id(&db, "todos");
    let is_sql = test.capability().sql;

    let todo1 = Todo::create()
        .user_id("alice")
        .title("original1")
        .exec(&mut db)
        .await
        .unwrap();

    let todo2 = Todo::create()
        .user_id("alice")
        .title("original2")
        .exec(&mut db)
        .await
        .unwrap();

    test.log().clear();

    // Update all todos for "alice" using only the partition key filter.
    Todo::filter_by_user_id("alice")
        .update()
        .title("updated")
        .exec(&mut db)
        .await
        .unwrap();

    if is_sql {
        let (op, resp) = test.log().pop();

        // Column index 2 = title (id=0, user_id=1, title=2).
        assert_struct!(op, Operation::QuerySql({
            stmt: Statement::Update({
                target: UpdateTarget::Table(== todo_table_id),
                assignments: #{ [2]: Assignment::Set(Expr::Arg({ position: 0 }))},
            }),
            params: [{ value: == "updated" }, ..],
            ret: None,
        }));

        assert_struct!(resp, {
            values: Rows::Count(_),
        });
    } else {
        // NoSQL: first a QueryPk to collect all matching PKs, then one
        // single-key UpdateByKey per matched record (the engine shreds
        // multi-key updates so each key is adjudicated independently).
        let (op, _) = test.log().pop();

        assert_struct!(op, Operation::QueryPk({
            table: == todo_table_id,
            select.len(): 2,
            filter: None,
        }));

        // Column index 2 = title (id=0, user_id=1, title=2).
        for _ in 0..2 {
            let (op, resp) = test.log().pop();

            assert_struct!(op, Operation::UpdateByKey({
                table: == todo_table_id,
                keys.len(): 1,
                assignments: #{ [2]: Assignment::Set(== "updated")},
                filter: None,
                returning: None,
            }));

            assert_struct!(resp, {
                values: Rows::Count(1),
            });
        }
    }

    assert!(test.log().is_empty(), "log should be empty after update");

    test.log().clear();
    let reloaded1 = Todo::get_by_user_id_and_id(&mut db, &todo1.user_id, todo1.id)
        .await
        .unwrap();
    assert_eq!(reloaded1.title, "updated");

    test.log().clear();
    let reloaded2 = Todo::get_by_user_id_and_id(&mut db, &todo2.user_id, todo2.id)
        .await
        .unwrap();
    assert_eq!(reloaded2.title, "updated");
}

/// Test delete on a model with a partitioned composite primary key using the
/// partition-key-only filter.
///
/// `Todo::filter_by_user_id(user_id).delete()` must delete all matching records,
/// not silently skip the deletion by issuing only a read-only `QueryPk`.
#[driver_test]
pub async fn delete_by_partition_key(test: &mut Test) {
    #[derive(Debug, toasty::Model)]
    #[key(partition = user_id, local = id)]
    struct Todo {
        #[auto]
        id: uuid::Uuid,

        user_id: String,

        title: String,
    }

    let mut db = test.setup_db(models!(Todo)).await;

    let todo_table_id = table_id(&db, "todos");
    let is_sql = test.capability().sql;

    let todo1 = Todo::create()
        .user_id("alice")
        .title("todo1")
        .exec(&mut db)
        .await
        .unwrap();

    let todo2 = Todo::create()
        .user_id("alice")
        .title("todo2")
        .exec(&mut db)
        .await
        .unwrap();

    let user_id = todo1.user_id.clone();
    let id1 = todo1.id;
    let id2 = todo2.id;

    test.log().clear();

    // Delete all todos for "alice" using only the partition key filter.
    Todo::filter_by_user_id("alice")
        .delete()
        .exec(&mut db)
        .await
        .unwrap();

    if is_sql {
        let (op, resp) = test.log().pop();

        assert_struct!(op, Operation::QuerySql({
            stmt: Statement::Delete({
                from: Source::Table({
                    tables: [== todo_table_id, ..],
                }),
            }),
        }));

        assert_struct!(resp, {
            values: Rows::Count(_),
        });
    } else {
        // NoSQL: first a QueryPk to collect all matching PKs, then one
        // DeleteByKey per matched record (the engine fans out individually).
        let (op, _) = test.log().pop();

        assert_struct!(op, Operation::QueryPk({
            table: == todo_table_id,
            select.len(): 2,
            filter: None,
        }));

        for _ in 0..2 {
            let (op, resp) = test.log().pop();

            assert_struct!(op, Operation::DeleteByKey({
                table: == todo_table_id,
                keys.len(): 1,
                filter: None,
            }));

            assert_struct!(resp, {
                values: Rows::Count(1),
            });
        }
    }

    assert!(test.log().is_empty(), "log should be empty after delete");

    assert_err!(Todo::get_by_user_id_and_id(&mut db, &user_id, id1).await);
    assert_err!(Todo::get_by_user_id_and_id(&mut db, &user_id, id2).await);
}