Skip to main content

dynoxide/actions/
batch_write_item.rs

1use crate::actions::helpers;
2use crate::errors::{DynoxideError, Result};
3use crate::storage_backend::StorageBackend;
4use crate::types::{self, AttributeValue, Item};
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8#[derive(Debug, Default, Deserialize)]
9pub struct BatchWriteItemRequest {
10    #[serde(rename = "RequestItems")]
11    pub request_items: HashMap<String, Vec<WriteRequest>>,
12    #[serde(rename = "ReturnConsumedCapacity", default)]
13    pub return_consumed_capacity: Option<String>,
14    #[serde(rename = "ReturnItemCollectionMetrics", default)]
15    pub return_item_collection_metrics: Option<String>,
16}
17
18#[derive(Debug, Default, Deserialize)]
19pub struct WriteRequest {
20    #[serde(rename = "PutRequest", default)]
21    pub put_request: Option<PutRequest>,
22    #[serde(rename = "DeleteRequest", default)]
23    pub delete_request: Option<DeleteRequest>,
24}
25
26#[derive(Debug, Default, Deserialize)]
27pub struct PutRequest {
28    #[serde(rename = "Item")]
29    pub item: Item,
30}
31
32#[derive(Debug, Default, Deserialize)]
33pub struct DeleteRequest {
34    #[serde(rename = "Key")]
35    pub key: HashMap<String, AttributeValue>,
36}
37
38#[derive(Debug, Default, Serialize)]
39pub struct BatchWriteItemResponse {
40    #[serde(rename = "UnprocessedItems")]
41    pub unprocessed_items: HashMap<String, serde_json::Value>,
42    #[serde(rename = "ConsumedCapacity", skip_serializing_if = "Option::is_none")]
43    pub consumed_capacity: Option<Vec<crate::types::ConsumedCapacity>>,
44    #[serde(
45        rename = "ItemCollectionMetrics",
46        skip_serializing_if = "Option::is_none"
47    )]
48    pub item_collection_metrics: Option<HashMap<String, Vec<crate::types::ItemCollectionMetrics>>>,
49}
50
51/// Fold one item's per-index units into the running per-table totals. A batch
52/// reports the sum across its items, and an index no item in the batch touched
53/// stays absent rather than appearing as a zero.
54fn accumulate_index_units(
55    totals: &mut HashMap<String, HashMap<String, f64>>,
56    table_name: &str,
57    units: &HashMap<String, f64>,
58) {
59    if units.is_empty() {
60        return;
61    }
62    let table_entry = totals.entry(table_name.to_string()).or_default();
63    for (index_name, units) in units {
64        *table_entry.entry(index_name.clone()).or_insert(0.0) += units;
65    }
66}
67
68pub async fn execute<S: StorageBackend>(
69    storage: &S,
70    mut request: BatchWriteItemRequest,
71) -> Result<BatchWriteItemResponse> {
72    const MAX_REQUEST_SIZE: usize = 16 * 1024 * 1024; // 16MB
73
74    // Validate RequestItems is not empty.
75    // AWS routes the empty-map case through a separate parameter-required path
76    // rather than the standard "N validation errors detected" envelope.
77    if request.request_items.is_empty() {
78        return Err(DynoxideError::ValidationException(
79            "The requestItems parameter is required for BatchWriteItem".to_string(),
80        ));
81    }
82
83    // Validate each table entry has at least one write request
84    for (table_name, wrs) in &request.request_items {
85        if wrs.is_empty() {
86            return Err(DynoxideError::ValidationException(format!(
87                "1 validation error detected: Value at 'requestItems.{table_name}.member' failed to satisfy constraint: Member must have length greater than or equal to 1"
88            )));
89        }
90    }
91
92    // Validate table name format for all tables before checking existence
93    for table_name in request.request_items.keys() {
94        crate::validation::validate_table_name(table_name)?;
95    }
96
97    // Validate total request count.
98    // AWS surfaces this as the standard "1 validation error detected" envelope
99    // and echoes the WriteRequest list inside `Value '{<table>=[<dump>]}'`. The
100    // conformance suite anchors a regex around the envelope and the constraint
101    // phrase but leaves the dump body unconstrained (because the AWS SDK's
102    // Java-toString shape adds new AttributeValue fields over time). We emit
103    // the table name verbatim and a Rust Debug dump of the WriteRequests so
104    // the envelope matches without coupling to a specific SDK version. If a
105    // future suite tightens the regex to pin the dump exactly, this site
106    // will need a follow-up change.
107    let total_requests: usize = request.request_items.values().map(|v| v.len()).sum();
108    if total_requests > 25 {
109        let empty: Vec<WriteRequest> = Vec::new();
110        let (table_name, requests) = request
111            .request_items
112            .iter()
113            .max_by_key(|(_, v)| v.len())
114            .map(|(name, v)| (name.as_str(), v))
115            .unwrap_or(("", &empty));
116        let dump = format!("{requests:?}");
117        return Err(DynoxideError::ValidationException(format!(
118            "1 validation error detected: Value '{{{table_name}=[{dump}]}}' at 'requestItems' failed to satisfy constraint: Map value must satisfy constraint: [Member must have length less than or equal to 25, Member must have length greater than or equal to 1]"
119        )));
120    }
121
122    // --- Pre-table validations ---
123    // DynamoDB validates attribute values, item size, and empty write requests
124    // BEFORE checking table existence.
125    for write_requests in request.request_items.values() {
126        for wr in write_requests {
127            if wr.put_request.is_none() && wr.delete_request.is_none() {
128                return Err(DynoxideError::ValidationException(
129                    "Supplied AttributeValue has more than one datatypes set, must contain exactly one of the supported datatypes".to_string(),
130                ));
131            }
132            if let Some(ref put_req) = wr.put_request {
133                // Validate attribute values (empty strings, empty sets, invalid numbers)
134                crate::validation::validate_item_attribute_values(&put_req.item)?;
135
136                // Validate item size before table lookup
137                let size = types::item_size(&put_req.item);
138                if size > types::MAX_ITEM_SIZE {
139                    return Err(DynoxideError::ValidationException(
140                        "Item size has exceeded the maximum allowed size".to_string(),
141                    ));
142                }
143            }
144            if let Some(ref del_req) = wr.delete_request {
145                crate::validation::validate_item_attribute_values(&del_req.key)?;
146            }
147        }
148    }
149
150    // Validate aggregate request size
151    let total_size: usize = request
152        .request_items
153        .values()
154        .flat_map(|wrs| wrs.iter())
155        .map(|wr| {
156            if let Some(ref put_req) = wr.put_request {
157                types::item_size(&put_req.item)
158            } else if let Some(ref del_req) = wr.delete_request {
159                types::item_size(&del_req.key)
160            } else {
161                0
162            }
163        })
164        .sum();
165    if total_size > MAX_REQUEST_SIZE {
166        return Err(DynoxideError::ValidationException(
167            "Item collection too large: aggregate size of items in BatchWriteItem exceeds 16MB limit".to_string(),
168        ));
169    }
170
171    // Validate: no duplicate keys across all operations
172    {
173        let mut seen_keys: std::collections::HashSet<(String, String, String)> =
174            std::collections::HashSet::new();
175        for (table_name, write_requests) in &request.request_items {
176            let meta = helpers::require_table_for_item_op(storage, table_name).await?;
177            let key_schema = helpers::parse_key_schema(&meta)?;
178            for wr in write_requests {
179                // Validate keys BEFORE extract_key_strings: that helper returns
180                // InternalServerError (HTTP 500) for a missing partition or sort
181                // key, but a key-less request is client input and must surface as
182                // a 400 ValidationException. Mirrors put_item.rs, which validates
183                // before extracting.
184                let key_item = if let Some(ref put) = wr.put_request {
185                    helpers::validate_item_keys_for_batch(&put.item, &key_schema, &meta)?;
186                    &put.item
187                } else if let Some(ref del) = wr.delete_request {
188                    helpers::validate_key_only(&del.key, &key_schema)?;
189                    &del.key
190                } else {
191                    continue;
192                };
193                let (pk, sk) = helpers::extract_key_strings(key_item, &key_schema)?;
194                let key = (table_name.clone(), pk, sk);
195                if !seen_keys.insert(key) {
196                    return Err(DynoxideError::ValidationException(
197                        "Provided list of item keys contains duplicates".to_string(),
198                    ));
199                }
200            }
201        }
202    }
203
204    // Track per-table secondary index capacity and affected partition keys for
205    // deferred metrics
206    let mut table_gsi_units: HashMap<String, HashMap<String, f64>> = HashMap::new();
207    let mut table_lsi_units: HashMap<String, HashMap<String, f64>> = HashMap::new();
208    // Vector replication is billed in bytes, so it accumulates on its own map
209    // and never joins the unit totals.
210    let mut table_vector_bytes: HashMap<String, HashMap<String, f64>> = HashMap::new();
211    // Track per-table WCU (table-level, excludes GSI)
212    let mut table_wcu: HashMap<String, f64> = HashMap::new();
213    // Collect unique (table, pk_str, pk_attr, pk_value) for deferred metrics computation
214    let mut affected_partitions: Vec<(String, String, String, AttributeValue)> = Vec::new();
215
216    for (table_name, write_requests) in &mut request.request_items {
217        let meta = helpers::require_table_for_item_op(storage, table_name).await?;
218        let key_schema = helpers::parse_key_schema(&meta)?;
219        // Parse the index definitions once per table: the maintain helpers'
220        // meta-accepting forms deserialise the JSON on every call, which for
221        // a 25-item batch against one table would repeat the work per item.
222        let gsi_defs = super::gsi::parse_gsi_defs(&meta)?;
223        let lsi_defs = super::lsi::parse_lsi_defs(&meta)?;
224        let vector_defs = super::vector_index::parse_vector_defs(&meta)?;
225        let attr_defs = super::vector_index::parse_attr_defs(&meta)?;
226
227        for wr in write_requests {
228            if let Some(ref mut put_req) = wr.put_request {
229                // Validate keys
230                helpers::validate_item_keys_for_batch(&put_req.item, &key_schema, &meta)?;
231
232                // Validate attribute values (empty strings, empty sets)
233                crate::validation::validate_item_attribute_values(&put_req.item)?;
234
235                // Normalize sets (deduplication)
236                crate::validation::normalize_item_sets(&mut put_req.item);
237
238                // Validate item size
239                let size = types::item_size(&put_req.item);
240                if size > types::MAX_ITEM_SIZE {
241                    return Err(DynoxideError::ValidationException(
242                        "Item size has exceeded the maximum allowed size".to_string(),
243                    ));
244                }
245
246                // TODO: validation must precede this call -- if reaching this line, caller has already validated keys.
247                let (pk, sk) = helpers::extract_key_strings(&put_req.item, &key_schema)?;
248                let item_json = serde_json::to_string(&put_req.item)
249                    .map_err(|e| DynoxideError::InternalServerError(e.to_string()))?;
250                let hash_prefix = put_req
251                    .item
252                    .get(&key_schema.partition_key)
253                    .map(crate::storage::compute_hash_prefix)
254                    .unwrap_or_default();
255                // Base write + index fan-out + stream are one atomic unit per
256                // item: a mid-fan-out failure rolls this item's write back.
257                // BatchWriteItem items are independent, so this is one
258                // transaction per write request.
259                let (old_size, gsi_units, lsi_units, vector_bytes) =
260                    helpers::with_write_transaction(storage, async {
261                        let old_json = storage
262                            .put_item_with_hash(
263                                table_name,
264                                &pk,
265                                &sk,
266                                &item_json,
267                                size,
268                                &hash_prefix,
269                            )
270                            .await?;
271                        let old_item: Option<Item> =
272                            old_json.and_then(|j| serde_json::from_str(&j).ok());
273                        let target = super::gsi::IndexWrite {
274                            table_name,
275                            pk: &pk,
276                            sk: &sk,
277                            pk_attr: &key_schema.partition_key,
278                            sk_attr: key_schema.sort_key.as_deref(),
279                            old_item: old_item.as_ref(),
280                            capacity_mode: request.return_consumed_capacity.as_deref(),
281                        };
282                        let gsi_units = super::gsi::maintain_gsis_after_write_with_defs(
283                            storage,
284                            &gsi_defs,
285                            &target,
286                            &put_req.item,
287                        )
288                        .await?;
289                        let lsi_units = super::lsi::maintain_lsis_after_write_with_defs(
290                            storage,
291                            &lsi_defs,
292                            &target,
293                            &put_req.item,
294                        )
295                        .await?;
296                        let vector_bytes =
297                            super::vector_index::maintain_vector_indexes_after_write_with_defs(
298                                storage,
299                                &vector_defs,
300                                &attr_defs,
301                                &target,
302                                &put_req.item,
303                                false,
304                            )
305                            .await?;
306                        crate::streams::record_stream_event(
307                            storage,
308                            &meta,
309                            old_item.as_ref(),
310                            Some(&put_req.item),
311                        )
312                        .await?;
313                        Ok((
314                            old_item.as_ref().map(types::item_size),
315                            gsi_units,
316                            lsi_units,
317                            vector_bytes,
318                        ))
319                    })
320                    .await?;
321
322                // Accumulate WCU on the larger of the displaced and new images
323                *table_wcu.entry(table_name.clone()).or_insert(0.0) +=
324                    types::table_write_capacity_units(old_size, Some(size));
325
326                // Accumulate secondary index units per table
327                accumulate_index_units(&mut table_gsi_units, table_name, &gsi_units);
328                accumulate_index_units(&mut table_lsi_units, table_name, &lsi_units);
329                accumulate_index_units(&mut table_vector_bytes, table_name, &vector_bytes);
330
331                // Track affected partition for deferred metrics
332                if let Some(pk_val) = put_req.item.get(&key_schema.partition_key) {
333                    affected_partitions.push((
334                        table_name.clone(),
335                        pk.clone(),
336                        key_schema.partition_key.clone(),
337                        pk_val.clone(),
338                    ));
339                }
340            } else if let Some(ref del_req) = wr.delete_request {
341                helpers::validate_key_only(&del_req.key, &key_schema)?;
342                // TODO: validation must precede this call -- if reaching this line, caller has already validated keys.
343                let (pk, sk) = helpers::extract_key_strings(&del_req.key, &key_schema)?;
344
345                // Base delete + index fan-out + stream are one atomic unit per item.
346                let (old_item, gsi_units, lsi_units, vector_bytes) =
347                    helpers::with_write_transaction(storage, async {
348                        let old_json = storage.delete_item(table_name, &pk, &sk).await?;
349                        let old_item: Option<Item> =
350                            old_json.as_ref().and_then(|j| serde_json::from_str(j).ok());
351                        let target = super::gsi::IndexWrite {
352                            table_name,
353                            pk: &pk,
354                            sk: &sk,
355                            pk_attr: &key_schema.partition_key,
356                            sk_attr: key_schema.sort_key.as_deref(),
357                            old_item: old_item.as_ref(),
358                            capacity_mode: request.return_consumed_capacity.as_deref(),
359                        };
360                        let gsi_units = super::gsi::maintain_gsis_after_delete_with_defs(
361                            storage, &gsi_defs, &target,
362                        )
363                        .await?;
364                        let lsi_units = super::lsi::maintain_lsis_after_delete_with_defs(
365                            storage, &lsi_defs, &target,
366                        )
367                        .await?;
368                        let vector_bytes =
369                            super::vector_index::maintain_vector_indexes_after_delete_with_defs(
370                                storage,
371                                &vector_defs,
372                                &attr_defs,
373                                &target,
374                            )
375                            .await?;
376                        if old_item.is_some() {
377                            crate::streams::record_stream_event(
378                                storage,
379                                &meta,
380                                old_item.as_ref(),
381                                None,
382                            )
383                            .await?;
384                        }
385                        Ok((old_item, gsi_units, lsi_units, vector_bytes))
386                    })
387                    .await?;
388
389                // Accumulate WCU: based on old item size if it existed, else 1 WCU
390                *table_wcu.entry(table_name.clone()).or_insert(0.0) +=
391                    types::table_write_capacity_units(
392                        old_item.as_ref().map(types::item_size),
393                        None,
394                    );
395
396                // Accumulate secondary index units per table
397                accumulate_index_units(&mut table_gsi_units, table_name, &gsi_units);
398                accumulate_index_units(&mut table_lsi_units, table_name, &lsi_units);
399                accumulate_index_units(&mut table_vector_bytes, table_name, &vector_bytes);
400
401                // Track affected partition for deferred metrics
402                if let Some(pk_val) = del_req.key.get(&key_schema.partition_key) {
403                    affected_partitions.push((
404                        table_name.clone(),
405                        pk.clone(),
406                        key_schema.partition_key.clone(),
407                        pk_val.clone(),
408                    ));
409                }
410            } else {
411                return Err(DynoxideError::ValidationException(
412                    "WriteRequest must contain either PutRequest or DeleteRequest".to_string(),
413                ));
414            }
415        }
416    }
417
418    // Build consumed capacity per table using pre-tracked WCU
419    let consumed_capacity = if matches!(
420        request.return_consumed_capacity.as_deref(),
421        Some("TOTAL") | Some("INDEXES")
422    ) {
423        let mut caps = Vec::new();
424        for table_name in request.request_items.keys() {
425            let total_wcu = table_wcu.get(table_name).copied().unwrap_or(0.0);
426            let gsi_units = table_gsi_units.get(table_name).cloned().unwrap_or_default();
427            let lsi_units = table_lsi_units.get(table_name).cloned().unwrap_or_default();
428            let vector_bytes = table_vector_bytes
429                .get(table_name)
430                .cloned()
431                .unwrap_or_default();
432            if let Some(cc) = crate::types::consumed_capacity_with_vector_indexes(
433                table_name,
434                total_wcu,
435                &gsi_units,
436                &lsi_units,
437                &vector_bytes,
438                &request.return_consumed_capacity,
439            ) {
440                caps.push(cc);
441            }
442        }
443        Some(caps)
444    } else {
445        None
446    };
447
448    // Item collection metrics are computed once per unique (table, pk), deferred from the write loop
449    let mut all_item_collection_metrics: HashMap<String, Vec<crate::types::ItemCollectionMetrics>> =
450        HashMap::new();
451    if matches!(
452        request.return_item_collection_metrics.as_deref(),
453        Some("SIZE")
454    ) {
455        // Deduplicate by (table, pk) to avoid redundant queries
456        let mut seen = std::collections::HashSet::new();
457        for (tbl, pk_str, pk_attr, pk_val) in &affected_partitions {
458            let key = (tbl.as_str(), pk_str.as_str());
459            if !seen.insert(key) {
460                continue;
461            }
462            let meta = helpers::require_table(storage, tbl).await?;
463            if let Some(icm) = helpers::build_item_collection_metrics(
464                storage,
465                &meta,
466                tbl,
467                pk_str,
468                pk_attr,
469                pk_val,
470                &request.return_item_collection_metrics,
471            )
472            .await?
473            {
474                all_item_collection_metrics
475                    .entry(tbl.clone())
476                    .or_default()
477                    .push(icm);
478            }
479        }
480    }
481    let item_collection_metrics = if all_item_collection_metrics.is_empty() {
482        None
483    } else {
484        Some(all_item_collection_metrics)
485    };
486
487    Ok(BatchWriteItemResponse {
488        unprocessed_items: HashMap::new(),
489        consumed_capacity,
490        item_collection_metrics,
491    })
492}
493
494#[cfg(test)]
495mod tests {
496    use crate::actions::{batch_write_item, create_table};
497    use crate::storage::Storage;
498    use crate::storage_backend::StorageBackend;
499
500    /// Each batch put is atomic with its own GSI fan-out: a mid-fan-out failure
501    /// rolls that item's base write back rather than leaving a torn index.
502    #[test]
503    fn batch_put_rolls_back_base_write_when_gsi_fan_out_fails() {
504        let storage = Storage::memory().unwrap();
505
506        let create = serde_json::from_value(serde_json::json!({
507            "TableName": "Orders",
508            "KeySchema": [{"AttributeName": "UserId", "KeyType": "HASH"}],
509            "AttributeDefinitions": [
510                {"AttributeName": "UserId", "AttributeType": "S"},
511                {"AttributeName": "Status", "AttributeType": "S"},
512                {"AttributeName": "Priority", "AttributeType": "S"}
513            ],
514            "GlobalSecondaryIndexes": [
515                {"IndexName": "StatusIndex", "KeySchema": [{"AttributeName": "Status", "KeyType": "HASH"}], "Projection": {"ProjectionType": "ALL"}},
516                {"IndexName": "PriorityIndex", "KeySchema": [{"AttributeName": "Priority", "KeyType": "HASH"}], "Projection": {"ProjectionType": "ALL"}}
517            ]
518        }))
519        .unwrap();
520        pollster::block_on(create_table::execute(&storage, create)).unwrap();
521
522        // Break the second GSI's fan-out by dropping its physical table.
523        storage.drop_gsi_table("Orders", "PriorityIndex").unwrap();
524
525        let batch = serde_json::from_value(serde_json::json!({
526            "RequestItems": {
527                "Orders": [
528                    {"PutRequest": {"Item": {"UserId": {"S": "u1"}, "Status": {"S": "SHIPPED"}, "Priority": {"S": "HIGH"}}}}
529                ]
530            }
531        }))
532        .unwrap();
533        let res = pollster::block_on(batch_write_item::execute(&storage, batch));
534        assert!(
535            res.is_err(),
536            "a mid-fan-out failure must surface as an error"
537        );
538
539        // The base write must roll back: no item landed.
540        let count =
541            pollster::block_on(<Storage as StorageBackend>::count_items(&storage, "Orders"))
542                .unwrap();
543        assert_eq!(
544            count, 0,
545            "batch put base write must roll back when fan-out fails"
546        );
547    }
548}