ScyllaDB Rust Driver: How do I create a Vec of batch values?

[Guy]

Originally from the User Slack

@Terry_Davis: How can you create a Vec of batch values in v0.15 like you used to be able to do in scylla rust driver v0.14? The examples only show how to do it using a tuple with hardcoded values: https://rust-driver.docs.scylladb.com/stable/queries/batch.html#batch-statement

#[derive(Debug, Deserialize, DeserializeRow, ValueList)]
pub struct PlayerRow {
    pub id: i32,
    pub account_id: i32,
}

fn generate_populate_queries(json_values: &JsonValues) -> (Batch, Vec<LegacySerializedValues>) {
    let mut batch = Batch::default();
    let mut batch_values: Vec<LegacySerializedValues> = Vec::new();

    if let Some(players) = json_values.players.as_ref() {
        for player in players.iter() {
            batch.append_statement("INSERT INTO players (id, account_id) VALUES (?, ?)");
            batch_values.push(player.serialized().unwrap().into_owned());
        }
    }

    // more values...

    (batch, batch_values)
}

I can’t create something like Vec<BatchValues> or Vec<dyn BatchValues> or Vec<Box<dyn BatchValues>> or anything like that.

Do I have to type out every single field manually like this? batch_values.push((player.id, player.account_id)); actually even that won’t work because tuples differ in shape.

@Karol_Baryła: 1. You should not use LegacySerializedValues - it is a not type safe remnant of previous serialization framework and will be removed in the next version. In your case you should probably use Player (or whatever is the type of player var).
2. fn generate_populate_queries(json_values: &JsonValues) -> (Batch, impl BatchValues) should work.

@Terry_Davis: 1. This is actually my objective to move away from LegacySerializedValues (0.14 -> 0.15 update)
2. The issue arises when I try to push another value of a different shape to batch_values

let mut batch_values: Vec<??> = Vec::new();

if let Some(players) = json_values.players.as_ref() {
    for player in players.iter() {
        batch.append_statement("INSERT INTO players (id, account_id) VALUES (?, ?)");
        batch_values.push(player);
    }
}

if let Some(accounts) = json_values.accounts.as_ref() {
    for account in accounts.iter() {
        batch.append_statement("INSERT INTO accounts (id, hash, created_at) VALUES (?, ?, ?)");
        batch_values.push(account);
    }
}

error:

46   |             batch_values.push(*player);
     |                          ---- ^^^^^^^ expected `Account`, found `Player

Previously, the "polymorphism" was achieved thanks toValueList`, but now I don’t know how to type that vector.

@Karol_Baryła: This would be a great use case for “BoundStatement” which we currently lack, but will be added in the future.
For now you can:
• Use 2 batches instead of one, so that the type of values in the batch is constant.
• Box your values and use Vec<Box<dyn SerializeRow>>
I think you could also make your own enum, with variants for your types, and implement SerializeRow for it (and then use Vec<YourEnum>) but I’m not sure because I didn’t try to do this before. I can try in a week (I am going for a leave now), or maybe @Mikołaj_Uzarski or @Wojciech_Przytuła will be able to help.

@Terry_Davis: The enum route will work for now. I had to implement the BatchValues trait which is required when calling batch. If someone has a better solution please lmk. This is horrid but works:

struct BatchValuesList(Vec<BatchItem>);

struct BatchValuesListIterator<'a> {
    values: std::slice::Iter<'a, BatchItem>,
}

impl BatchValues for BatchValuesList {
    type BatchValuesIter<'r>
        = BatchValuesListIterator<'r>
    where
        Self: 'r;

    fn batch_values_iter(&self) -> Self::BatchValuesIter<'_> {
        BatchValuesListIterator {
            values: self.0.iter(),
        }
    }
}

#[derive(Clone)]
enum BatchItem {
    PlayerRow(PlayerRow),
    AccountRow(AccountRow),
}

impl SerializeRow for BatchItem {
    fn serialize(
        &self,
        ctx: &RowSerializationContext<'_>,
        writer: &mut RowWriter,
    ) -> Result<(), SerializationError> {
        match self {
            BatchItem::PlayerRow(val) => val.serialize(ctx, writer),
            BatchItem::AccountRow(val) => val.serialize(ctx, writer),
        }
    }

    fn is_empty(&self) -> bool {
        false
    }
}

impl<'bv> BatchValuesIterator<'bv> for BatchValuesListIterator<'bv> {
    fn serialize_next(
        &mut self,
        ctx: &RowSerializationContext<'_>,
        writer: &mut RowWriter,
    ) -> Option<Result<(), SerializationError>> {
        self.values.next().map(|value| value.serialize(ctx, writer))
    }

    fn is_empty_next(&mut self) -> Option<bool> {
        self.values.next().map(|value| value.is_empty())
    }

    fn skip_next(&mut self) -> Option<()> {
        self.values.next().map(|_| ())
    }
}

fn generate_populate_queries(db_data: &DBData) -> (Batch, BatchValuesList) {
    let mut batch = Batch::default();
    let mut batch_values = Vec::new();

    if let Some(players) = db_data.players.as_ref() {
        for player in players.iter() {
            batch.append_statement(INSERT_PLAYER_QUERY);
            batch_values.push(BatchItem::CostumeRow(player.clone()));
        }
    }

    if let Some(accounts) = db_data.accounts.as_ref() {
        for account in accounts.iter() {
            batch.append_statement(INSERT_ACCOUNT_QUERY);
            batch_values.push(BatchItem::AccountRow(account.clone()));
        }
    }
}

pub async fn persist_db_data(db_data: &DBData, db_manager: &Arc<DBManager>) {
    let (batch, batch_values) = generate_populate_queries(db_data);

    db_manager
        .session
        .batch(&batch, batch_values)
        .await
        .unwrap();
}

@Karol_Baryła: Hi @Terry_Davis, I’m back now so I can help again. There should be no need to implement BatchValues in your case - there is already an impl for Vec<T: SerializeRow>. I wrote a quick example of using enum for batch with multiple row types:

#[derive(SerializeRow)]
struct CustomStructA {
    field_1: String,
    field_2: String,
}

#[derive(SerializeRow)]
struct CustomStructB {
    field_x: String,
    field_y: String,
}

enum CustomBatchRow {
    A(CustomStructA),
    B(CustomStructB),
}

impl SerializeRow for CustomBatchRow {
    fn serialize(
        &self,
        ctx: &scylla::serialize::row::RowSerializationContext<'_>,
        writer: &mut scylla::serialize::writers::RowWriter,
    ) -> std::result::Result<(), scylla::serialize::SerializationError> {
        match self {
            CustomBatchRow::A(s) => SerializeRow::serialize(s, ctx, writer),
            CustomBatchRow::B(s) => SerializeRow::serialize(s, ctx, writer),
        }
    }

    fn is_empty(&self) -> bool {
        match self {
            CustomBatchRow::A(s) => SerializeRow::is_empty(s),
            CustomBatchRow::B(s) => SerializeRow::is_empty(s),
        }
    }
}

fn generate_populate_queries(
    a_rows: impl Iterator<Item = CustomStructA>,
    b_rows: impl Iterator<Item = CustomStructB>,
) -> (Batch, Vec<CustomBatchRow>) {
    let mut batch = Batch::default();
    let mut batch_values: Vec<CustomBatchRow> = Vec::new();

    for a_row_values in a_rows {
        batch.append_statement("INSERT INTO table_A (field_1, field_2) VALUES (?, ?)");
        batch_values.push(CustomBatchRow::A(a_row_values));
    }

    for b_row_values in b_rows {
        batch.append_statement("INSERT INTO table_B (field_x, field_y) VALUES (?, ?)");
        batch_values.push(CustomBatchRow::B(b_row_values));
    }

    // more values...

    (batch, batch_values)
}

async fn execute_batch(s: &Session) {
    let (batch, values) = generate_populate_queries(vec![].into_iter(), vec![].into_iter());
    let result = s.batch(&batch, &values).await.unwrap();
}