diff --git a/.github/workflows/integration-mongodb.yml b/.github/workflows/integration-mongodb.yml index 03223f967..b35233402 100644 --- a/.github/workflows/integration-mongodb.yml +++ b/.github/workflows/integration-mongodb.yml @@ -117,7 +117,7 @@ jobs: - name: Build release (mongodb backend) run: cargo build --release --no-default-features --features mongodb-test-hooks - name: Run MongoDB rust integration suite - run: devtools/run-mongodb-tests -- --rust --rust-integration + run: devtools/run-mongodb-tests --rust-live -- --rust --rust-integration mongodb-production-build: runs-on: extenddb_ubuntu-2404_4-core diff --git a/crates/storage-mongodb/src/data/key_fields.rs b/crates/storage-mongodb/src/data/key_fields.rs new file mode 100644 index 000000000..972503c76 --- /dev/null +++ b/crates/storage-mongodb/src/data/key_fields.rs @@ -0,0 +1,461 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Key attributes of a base-table document, for the two places that build a +//! document from something other than a full item. +//! +//! [`item_to_document`] writes a base-table document as `{ _id, pk, sk_*, +//! item_data }`, with `_id` the key text and the key attributes inside +//! `item_data`. It writes no `_v`: the version is added by the UpdateItem +//! write itself, as `$inc` on the native path. An UpdateItem that creates an +//! item through a native MongoDB upsert starts from the key filter (`pk` and +//! `sk_*`) and the update operators, so it needs the `_id` and the +//! `item_data` key attributes added with `$setOnInsert`: +//! [`insert_only_key_fields`] builds that document. +//! +//! Before that was done, such an upsert stored a document with a generated +//! ObjectId `_id` and no key attributes in `item_data` (issue #436), and for +//! a REMOVE-only update no `item_data` at all. +//! [`repaired_document`] rebuilds a document of either shape into the shape +//! `item_to_document` writes, reading the key values back from the top-level +//! `pk` and `sk_*` fields. + +use base64::Engine as _; +use bson::{Bson, Document}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType, +}; +use extenddb_core::validation::number::validate_and_normalize_number; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::sk_info; + +use super::{document_to_item, item_to_document, sk_suffix}; + +/// The `$setOnInsert` document for an upsert whose filter is [`super::pk_filter`] +/// on `key`: `_id` as `item_to_document` writes it, and every key attribute +/// under `item_data.`. +/// +/// The filter's own equality fields (`pk`, `sk_*`) are not repeated; MongoDB +/// copies them into the inserted document itself. +pub fn insert_only_key_fields( + key: &Item, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], +) -> Result { + let key_doc = item_to_document(key, key_schema, attribute_definitions)?; + let mut fields = Document::new(); + let id = key_doc + .get("_id") + .cloned() + .ok_or_else(|| StorageError::Internal("key document has no _id".to_owned()))?; + fields.insert("_id", id); + for ks in key_schema { + let value = key.get(&ks.attribute_name).ok_or_else(|| { + StorageError::Internal(format!("missing key attribute {}", ks.attribute_name)) + })?; + let value_json = + serde_json::to_value(value).map_err(|e| StorageError::Internal(e.to_string()))?; + let value_bson = + bson::to_bson(&value_json).map_err(|e| StorageError::Internal(e.to_string()))?; + fields.insert(format!("item_data.{}", ks.attribute_name), value_bson); + } + Ok(fields) +} + +/// Whether `doc`'s `_id` is an ObjectId. +/// +/// That is the mark of the upsert described in the module comment, and the +/// whole test: `item_to_document` writes `_id` as the key text for every +/// base-table document, so a base-table document whose `_id` is an ObjectId +/// was generated by MongoDB on an insert that supplied no `_id`, which is +/// the pre-fix upsert. The `item_data` contents are not examined. +#[must_use] +pub fn is_stored_without_key(doc: &Document) -> bool { + matches!(doc.get("_id"), Some(Bson::ObjectId(_))) +} + +/// Rebuild a document stored without its key into the shape +/// `item_to_document` writes: the same `item_data` with the key attributes +/// added, `_id` as the key text, and `_v` carried over. +/// +/// A document with no `item_data` at all rebuilds into the key-only item. +/// That is the shape the pre-fix upsert stored for a REMOVE-only creating +/// UpdateItem, where the filter supplied `pk` and `sk_*`, the `$unset` +/// contributed nothing and the `$inc` gave `_v`; it is also what the fixed +/// fast path now writes for the same request. +/// +/// The key values come from the top-level `pk` text and `sk_s`/`sk_n`/`sk_b` +/// field, typed by `attribute_definitions`. A document whose top-level key +/// fields are missing or of the wrong type is an error; the caller leaves it +/// in place. +pub fn repaired_document( + doc: &Document, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], +) -> Result { + let mut item = if doc.contains_key("item_data") { + document_to_item(doc)? + } else { + Item::new() + }; + + let pk_name = key_schema + .iter() + .find(|ks| ks.key_type == KeyType::Hash) + .map(|ks| ks.attribute_name.as_str()) + .ok_or_else(|| StorageError::Internal("key schema has no partition key".to_owned()))?; + let pk_type = attribute_type(attribute_definitions, pk_name)?; + let pk_text = doc + .get_str("pk") + .map_err(|_| StorageError::Internal("document has no pk text".to_owned()))?; + item.insert(pk_name.to_owned(), pk_value(pk_type, pk_text)?); + + if let Some((sk_name, sk_type)) = sk_info(key_schema, attribute_definitions) { + let field = format!("sk_{}", sk_suffix(sk_type)); + let value = doc + .get(&field) + .ok_or_else(|| StorageError::Internal(format!("document has no {field}")))?; + item.insert(sk_name.to_owned(), sk_value(sk_type, value, &field)?); + } + + let mut repaired = item_to_document(&item, key_schema, attribute_definitions)?; + if let Some(version) = doc.get("_v") { + repaired.insert("_v", version.clone()); + } + Ok(repaired) +} + +fn attribute_type( + attribute_definitions: &[AttributeDefinition], + name: &str, +) -> Result { + attribute_definitions + .iter() + .find(|d| d.attribute_name == name) + .map(|d| d.attribute_type) + .ok_or_else(|| StorageError::Internal(format!("no attribute definition for {name}"))) +} + +/// The partition key attribute from its `pk` text, inverting +/// `extenddb_storage::util::pk_to_text`. +fn pk_value(pk_type: ScalarAttributeType, text: &str) -> Result { + Ok(match pk_type { + ScalarAttributeType::S => AttributeValue::S(text.to_owned()), + ScalarAttributeType::N => AttributeValue::N(text.to_owned()), + ScalarAttributeType::B => AttributeValue::B( + base64::engine::general_purpose::STANDARD + .decode(text) + .map_err(|e| StorageError::Internal(format!("pk text is not base64: {e}")))?, + ), + }) +} + +/// The sort key attribute from its typed `sk_*` field, inverting +/// `insert_typed_sk`. +fn sk_value( + sk_type: ScalarAttributeType, + value: &Bson, + field: &str, +) -> Result { + match (sk_type, value) { + (ScalarAttributeType::S, Bson::String(s)) => Ok(AttributeValue::S(s.clone())), + (ScalarAttributeType::N, Bson::Decimal128(d)) => { + Ok(AttributeValue::N(plain_number_text(&d.to_string())?)) + } + (ScalarAttributeType::B, Bson::String(hex)) => Ok(AttributeValue::B(hex_to_bytes(hex)?)), + _ => Err(StorageError::Internal(format!( + "{field} holds a {} value, not the sort key type", + value.element_type() as u8 + ))), + } +} + +/// The text a request carries for a numeric sort key, from the display form +/// of its stored `Decimal128`. +/// +/// `Decimal128` displays a value below 1e-6 in exponent form (0.0000001 is +/// `1E-7`), while a request's number is the plain decimal that +/// `validate_and_normalize_number` produces, and PutItem builds the `_id` +/// from that text. A document rebuilt with the exponent form in its `_id` +/// would make every later PutItem on the key fail on the immutable `_id`. +fn plain_number_text(display: &str) -> Result { + validate_and_normalize_number(display) + .map_err(|e| StorageError::Internal(format!("sort key is not a number: {e}"))) +} + +fn hex_to_bytes(hex: &str) -> Result, StorageError> { + let bytes = hex.as_bytes(); + if !bytes.len().is_multiple_of(2) { + return Err(StorageError::Internal("sk_b hex has odd length".to_owned())); + } + bytes + .chunks(2) + .map(|pair| { + let s = std::str::from_utf8(pair).map_err(|e| StorageError::Internal(e.to_string()))?; + u8::from_str_radix(s, 16) + .map_err(|e| StorageError::Internal(format!("sk_b is not hex: {e}"))) + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + use bson::doc; + + fn key(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } + } + + fn def(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + } + } + + fn hash_schema() -> (Vec, Vec) { + ( + vec![key("pk", KeyType::Hash)], + vec![def("pk", ScalarAttributeType::S)], + ) + } + + fn range_schema( + sk_type: ScalarAttributeType, + ) -> (Vec, Vec) { + ( + vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)], + vec![def("pk", ScalarAttributeType::S), def("sk", sk_type)], + ) + } + + fn item(pairs: &[(&str, AttributeValue)]) -> Item { + pairs + .iter() + .map(|(k, v)| ((*k).to_owned(), v.clone())) + .collect() + } + + #[test] + fn insert_only_fields_carry_the_id_and_the_key_attributes_for_a_hash_key() { + let (schema, defs) = hash_schema(); + let key = item(&[("pk", AttributeValue::S("u1".to_owned()))]); + let fields = insert_only_key_fields(&key, &schema, &defs).expect("fields"); + assert_eq!(fields, doc! { "_id": "u1", "item_data.pk": { "S": "u1" } }); + } + + #[test] + fn insert_only_fields_carry_both_key_attributes_for_a_range_key() { + let (schema, defs) = range_schema(ScalarAttributeType::N); + let key = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("sk", AttributeValue::N("42".to_owned())), + ]); + let fields = insert_only_key_fields(&key, &schema, &defs).expect("fields"); + let full = item_to_document(&key, &schema, &defs).expect("document"); + assert_eq!(fields.get("_id"), full.get("_id")); + assert_eq!( + fields.get("item_data.pk"), + Some(&bson::bson!({ "S": "u1" })) + ); + assert_eq!( + fields.get("item_data.sk"), + Some(&bson::bson!({ "N": "42" })) + ); + assert_eq!( + fields.len(), + 3, + "no pk or sk_n field: the filter supplies those" + ); + } + + /// The document MongoDB inserts for the pre-fix upsert on a hash table, + /// as measured on MongoDB 7.0: the filter's `pk`, the `$set` attribute, + /// the `$inc` version, and a generated `_id`. + fn stored_without_key_hash() -> Document { + doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "item_data": { "w": { "S": "upd" } }, + "_v": 3_i64, + } + } + + #[test] + fn a_generated_object_id_marks_a_document_stored_without_its_key() { + let (schema, defs) = hash_schema(); + assert!(is_stored_without_key(&stored_without_key_hash())); + let good = item_to_document( + &item(&[("pk", AttributeValue::S("u1".to_owned()))]), + &schema, + &defs, + ) + .expect("document"); + assert!(!is_stored_without_key(&good)); + } + + #[test] + fn repair_restores_a_hash_key_and_keeps_attributes_and_version() { + let (schema, defs) = hash_schema(); + let repaired = + repaired_document(&stored_without_key_hash(), &schema, &defs).expect("repair"); + let want_item = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("w", AttributeValue::S("upd".to_owned())), + ]); + let mut want = item_to_document(&want_item, &schema, &defs).expect("document"); + want.insert("_v", 3_i64); + assert_eq!(repaired, want); + } + + #[test] + fn repair_restores_numeric_and_binary_partition_keys() { + // The partition key text is stored as written, so a number below 1e-6 + // keeps its plain form there. + let cases = [ + (ScalarAttributeType::N, AttributeValue::N("42".to_owned())), + ( + ScalarAttributeType::N, + AttributeValue::N("0.0000001".to_owned()), + ), + ( + ScalarAttributeType::B, + AttributeValue::B(vec![0x01, 0xff, 0x10]), + ), + ]; + for (pk_type, pk) in cases { + let schema = vec![key("pk", KeyType::Hash)]; + let defs = vec![def("pk", pk_type)]; + let want_item = item(&[ + ("pk", pk.clone()), + ("w", AttributeValue::S("upd".to_owned())), + ]); + let mut want = item_to_document(&want_item, &schema, &defs).expect("document"); + want.insert("_v", 2_i64); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": want.get("pk").cloned().expect("pk text"), + "item_data": { "w": { "S": "upd" } }, + "_v": 2_i64, + }; + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + assert_eq!(repaired, want, "{pk:?} partition key"); + assert_eq!( + document_to_item(&repaired).expect("item"), + want_item, + "{pk:?} partition key item" + ); + } + } + + /// A number as a request carries it: `{"N": raw}` deserialised, which + /// normalises the text to the plain decimal that PutItem stores. + fn request_number(raw: &str) -> AttributeValue { + serde_json::from_value(serde_json::json!({ "N": raw })).expect("attribute value") + } + + #[test] + fn repair_restores_string_numeric_and_binary_sort_keys() { + let cases = [ + (ScalarAttributeType::S, AttributeValue::S("s1".to_owned())), + (ScalarAttributeType::N, AttributeValue::N("12.5".to_owned())), + // `Decimal128` displays the next three in exponent form + // (`1E-7`, `-1.2345E-7`, `1E-30`), and the `_id` PutItem builds + // from them is the plain decimal. + ( + ScalarAttributeType::N, + AttributeValue::N("0.0000001".to_owned()), + ), + ( + ScalarAttributeType::N, + AttributeValue::N("-0.00000012345".to_owned()), + ), + (ScalarAttributeType::N, request_number("1E-30")), + // The same value as a request spells it. + (ScalarAttributeType::N, request_number("1E-7")), + // Plain below the switch to exponent form. + ( + ScalarAttributeType::N, + AttributeValue::N("0.000001".to_owned()), + ), + ( + ScalarAttributeType::B, + AttributeValue::B(vec![0x01, 0xff, 0x10]), + ), + ]; + for (sk_type, sk) in cases { + let (schema, defs) = range_schema(sk_type); + let want_item = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("sk", sk.clone()), + ("w", AttributeValue::S("upd".to_owned())), + ]); + let want = item_to_document(&want_item, &schema, &defs).expect("document"); + // The pre-fix upsert copied the filter's pk and sk_* fields and + // set only the non-key attribute. + let sk_field = format!("sk_{}", sk_suffix(sk_type)); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + &sk_field: want.get(&sk_field).cloned().expect("sort key field"), + "item_data": { "w": { "S": "upd" } }, + }; + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + assert_eq!(repaired, want, "{sk_type:?} sort key {sk:?}"); + assert_eq!( + document_to_item(&repaired).expect("item"), + want_item, + "{sk_type:?} sort key {sk:?}: item" + ); + } + } + + #[test] + fn repair_rebuilds_a_document_with_no_item_data_as_the_key_only_item() { + let (schema, defs) = hash_schema(); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "_v": 1_i64, + }; + assert!(is_stored_without_key(&bad)); + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + let want_item = item(&[("pk", AttributeValue::S("u1".to_owned()))]); + let mut want = item_to_document(&want_item, &schema, &defs).expect("document"); + want.insert("_v", 1_i64); + assert_eq!(repaired, want); + assert_eq!(document_to_item(&repaired).expect("item"), want_item); + } + + #[test] + fn repair_rebuilds_a_range_key_document_with_no_item_data() { + let (schema, defs) = range_schema(ScalarAttributeType::S); + let want_item = item(&[ + ("pk", AttributeValue::S("u1".to_owned())), + ("sk", AttributeValue::S("s1".to_owned())), + ]); + let want = item_to_document(&want_item, &schema, &defs).expect("document"); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "sk_s": "s1", + }; + let repaired = repaired_document(&bad, &schema, &defs).expect("repair"); + assert_eq!(repaired, want); + } + + #[test] + fn repair_refuses_a_document_whose_sort_key_field_has_the_wrong_type() { + let (schema, defs) = range_schema(ScalarAttributeType::N); + let bad = doc! { + "_id": bson::oid::ObjectId::new(), + "pk": "u1", + "sk_n": "not a decimal", + "item_data": { "w": { "S": "upd" } }, + }; + assert!(repaired_document(&bad, &schema, &defs).is_err()); + } +} diff --git a/crates/storage-mongodb/src/data/mod.rs b/crates/storage-mongodb/src/data/mod.rs index 45465ba61..d1678862a 100644 --- a/crates/storage-mongodb/src/data/mod.rs +++ b/crates/storage-mongodb/src/data/mod.rs @@ -5,6 +5,8 @@ //! //! Contains document conversion, collection naming, and key extraction utilities. +pub mod key_fields; + use bson::{Document, doc}; #[cfg(test)] diff --git a/crates/storage-mongodb/src/data_engine.rs b/crates/storage-mongodb/src/data_engine.rs index dd8f4cc9f..871a04254 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -47,7 +47,7 @@ fn native_attribute_name(path: &[PathElement], maps: &ExpressionMaps) -> Option< return None; }; let attr_name = resolve_name_ref(raw_name, maps).ok()?; - (!attr_name.contains('.') && !attr_name.starts_with('$')).then(|| attr_name.into_owned()) + update_create::is_plain_path_name(&attr_name).then(|| attr_name.into_owned()) } #[derive(Clone, Copy)] @@ -800,24 +800,25 @@ impl MongoEngine { // indexed attribute is $set to a new value or $unset (RFC-0003 // §2.2). The GSI cache lets us avoid a catalog query in the // common (no-GSI) case; when the cache is stale-or-unknown - // we fall through to the slow path which is authoritative. + // we fall through to the slow path which is authoritative. Key + // attribute names that are not plain path names (a dot, a leading + // `$`) also take the slow path, because the upsert writes the key + // attributes as `item_data.` (see `update_create`). if condition.is_none() && !return_old && stream.is_none() && self.gsi_cache_get_fresh(&key_info.table_id) == Some(false) + && update_create::key_names_are_plain_paths(&key_info.key_schema) && let Some(mongo_update) = Self::try_build_native_update(actions, maps) { let (fast_filter, took_fast) = match mongo_update { NativeUpdate::Doc(d) => { - let opts = mongodb::options::FindOneAndUpdateOptions::builder() - .upsert(true) - .return_document(ReturnDocument::After) - .build(); - let result = coll - .find_one_and_update(key_filter.clone(), d) - .with_options(opts) - .await - .map_err(|e| StorageError::Internal(e.to_string()))?; + // The upsert's insert branch must write the document + // `item_to_document` would write (issue #436); see + // `update_create::upsert_native_update`. + let result = self + .upsert_native_update(&coll, &key_filter, d, key, key_info) + .await?; (Some(result), true) } NativeUpdate::Pipeline { @@ -4059,6 +4060,16 @@ fn project_item( } } +#[path = "data_engine_key_repair.rs"] +pub(crate) mod key_repair; + +#[path = "data_engine_update_create.rs"] +mod update_create; + +#[cfg(test)] +#[path = "data_engine_live_tests.rs"] +mod live_tests; + #[cfg(test)] mod tests { use super::*; diff --git a/crates/storage-mongodb/src/data_engine_key_repair.rs b/crates/storage-mongodb/src/data_engine_key_repair.rs new file mode 100644 index 000000000..a3cb710cc --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_key_repair.rs @@ -0,0 +1,513 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Startup repair of base-table documents stored without their key. +//! +//! Before issue #436 was fixed, an UpdateItem that created an item on the +//! native fast path upserted a document with a generated ObjectId `_id` and +//! an `item_data` without the key attributes. Such a document answers +//! GetItem without the key, and every PutItem on the key fails on the +//! immutable `_id`. `_id` cannot be changed in place, so the repair replaces +//! each such document with the one `item_to_document` writes: delete the old +//! document by its ObjectId, insert the rebuilt one. The scan is a bracket +//! query on the `_id` index (`{_id: {$type: "objectId"}}`), so a collection +//! with no such document costs one index probe. +//! +//! Up to [`REPAIR_BATCH`] documents are rebuilt in one transaction. A +//! transaction's commit is one majority-acknowledged round trip and dominates +//! the cost of a rebuild, so batching the rebuilds behind one commit is what +//! sets the rate; each batch is one in-transaction read, one delete and one +//! insert. The batch is atomic, so a process killed mid-repair leaves every +//! document either as it was or rebuilt, and the next start finishes the +//! remainder. A batch whose transaction cannot run is rebuilt one document +//! per transaction, so one document another writer holds, or one that cannot +//! be rebuilt at all, leaves the rest of the batch repaired. +//! +//! The documents are re-read inside the transaction rather than taken from +//! the scan, so a client that writes to a document between the scan and the +//! rebuild does not lose that write: the rebuild either sees it or hits a +//! write conflict and runs again. +//! +//! [`spawn_key_repair`] runs the whole pass as a background task, so `serve` +//! opens its listener without waiting for it (see `server_components_factory` +//! in `lib.rs`). Until the repair reaches a given document, requests on that +//! document's key answer as they did before this repair existed. Progress is +//! logged every [`PROGRESS_EVERY`] documents and once at the end. A failure +//! on one document is logged and skipped; the summary reports it. +//! +//! A scan that stops on an error, as it does when a table is dropped while the +//! pass runs, is logged, and the pass goes on to the next table. The summary +//! counts it, and the documents that scan had not reached wait for the next +//! start. The scan projects `_id` only, so the `_id` index answers it without +//! reading a document. +//! +//! A client write to a key whose document a repair transaction holds waits for +//! that transaction, and a foreground write sets no `maxTimeMS`. Normally the +//! wait is one batch commit. A process killed mid-batch leaves its transaction +//! open until the server expires it (`transactionLifetimeLimitSeconds`, 60 s +//! by default), and writes to the up to [`REPAIR_BATCH`] keys of that batch +//! wait that long. + +use std::sync::Arc; + +use bson::{Bson, Document, doc}; +use extenddb_core::types::{AttributeDefinition, KeySchemaElement}; +use extenddb_storage::error::StorageError; +use futures::TryStreamExt; + +use super::{CommitOutcome, TRANSIENT_RETRY_ATTEMPTS, backoff_sleep, is_transient_write_conflict}; +use crate::MongoEngine; +use crate::data::data_collection_name; +use crate::data::key_fields::{is_stored_without_key, repaired_document}; + +/// Documents rebuilt in one transaction. +const REPAIR_BATCH: usize = 500; + +/// Repaired documents between progress lines. +const PROGRESS_EVERY: u64 = 10_000; + +/// How many times a single-document rebuild waits for a writer that holds its +/// document, and how long each wait is. +/// +/// A process killed mid-repair leaves its last transaction holding that +/// batch's documents until the server expires it, which takes up to +/// `transactionLifetimeLimitSeconds` (60 s by default). The repair runs in the +/// background and nothing waits on it, so it waits the locks out instead of +/// counting those documents failed and leaving them for the next start. +const CONFLICT_WAITS: u32 = 40; +const CONFLICT_WAIT: std::time::Duration = std::time::Duration::from_secs(2); + +/// What a repair run did. +#[derive(Debug, Default, Clone, PartialEq, Eq)] +pub struct KeyRepairSummary { + /// Tables whose data collection was scanned. + pub tables_scanned: u64, + /// Tables among those whose scan stopped on an error. The documents that + /// scan had not reached wait for the next start; each stop is logged. + pub tables_failed: u64, + /// Documents rewritten with their key. + pub documents_repaired: u64, + /// Documents left as they were because they could not be rebuilt or + /// written; each is logged. + pub documents_failed: u64, +} + +impl KeyRepairSummary { + #[must_use] + pub fn is_clean(&self) -> bool { + self.documents_repaired == 0 && self.documents_failed == 0 && self.tables_failed == 0 + } +} + +impl std::fmt::Display for KeyRepairSummary { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "{} tables scanned, {} stopped by an error, {} documents stored without their key \ + repaired, {} failed", + self.tables_scanned, self.tables_failed, self.documents_repaired, self.documents_failed + ) + } +} + +impl MongoEngine { + /// Repair every base-table document stored without its key, across all + /// tables in the catalog. See the module comment. + pub async fn repair_items_stored_without_key(&self) -> Result { + let tables = self.catalog_db.collection::("tables"); + let mut cursor = tables + .find(doc! {}) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let mut summary = KeyRepairSummary::default(); + while let Some(table_doc) = cursor + .try_next() + .await + .map_err(|e| StorageError::Internal(e.to_string()))? + { + let (table_id, key_schema, attribute_definitions) = match table_key(&table_doc) { + Ok(parsed) => parsed, + Err(e) => { + tracing::warn!("key repair: skipping a catalog table row: {e}"); + continue; + } + }; + summary.tables_scanned += 1; + if let Err(e) = self + .repair_table(&table_id, &key_schema, &attribute_definitions, &mut summary) + .await + { + tracing::error!( + "key repair: table {table_id}: the scan stopped ({e}); the documents it \ + had not reached wait for the next start" + ); + summary.tables_failed += 1; + } + } + Ok(summary) + } + + /// Repair the documents of one table's data collection. + pub(crate) async fn repair_table( + &self, + table_id: &str, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + summary: &mut KeyRepairSummary, + ) -> Result<(), StorageError> { + let coll = self + .data_db + .collection::(&data_collection_name(table_id)); + let mut cursor = coll + .find(doc! { "_id": { "$type": "objectId" } }) + .projection(doc! { "_id": 1 }) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + // The rebuilt documents have the key text as their `_id`, outside the + // ObjectId bracket this cursor reads, so inserting them cannot feed + // the scan its own output. + let mut ids: Vec = Vec::with_capacity(REPAIR_BATCH); + while let Some(bad) = cursor + .try_next() + .await + .map_err(|e| StorageError::Internal(e.to_string()))? + { + if !is_stored_without_key(&bad) { + continue; + } + if let Some(id) = bad.get("_id") { + ids.push(id.clone()); + } + if ids.len() >= REPAIR_BATCH { + self.repair_batch( + &coll, + table_id, + key_schema, + attribute_definitions, + &ids, + summary, + ) + .await; + ids.clear(); + } + } + if !ids.is_empty() { + self.repair_batch( + &coll, + table_id, + key_schema, + attribute_definitions, + &ids, + summary, + ) + .await; + } + Ok(()) + } + + /// Rebuild one batch, and when the batch transaction cannot run, rebuild + /// its documents one transaction at a time so that one document that + /// cannot be written leaves the rest of the batch repaired. + async fn repair_batch( + &self, + coll: &mongodb::Collection, + table_id: &str, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + ids: &[Bson], + summary: &mut KeyRepairSummary, + ) { + let before = summary.documents_repaired; + match self + .replace_batch(coll, ids, key_schema, attribute_definitions) + .await + { + Ok(outcome) => { + summary.documents_repaired += outcome.repaired; + summary.documents_failed += outcome.failed.len() as u64; + for (id, reason) in outcome.failed { + tracing::error!( + "key repair: table {table_id}: cannot rebuild document {id:?}: {reason}" + ); + } + } + Err(e) => { + tracing::warn!( + "key repair: table {table_id}: a batch of {} documents could not be \ + rewritten together ({e}); rewriting them one at a time", + ids.len() + ); + for id in ids { + self.repair_one( + coll, + table_id, + id, + key_schema, + attribute_definitions, + summary, + ) + .await; + } + } + } + log_progress(before, summary); + } + + /// Rebuild the document at `id` in its own transaction, waiting out a + /// writer that holds it rather than leaving it for the next start. + async fn repair_one( + &self, + coll: &mongodb::Collection, + table_id: &str, + id: &Bson, + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + summary: &mut KeyRepairSummary, + ) { + let one = std::slice::from_ref(id); + for waited in 0..=CONFLICT_WAITS { + match self + .replace_batch(coll, one, key_schema, attribute_definitions) + .await + { + Ok(outcome) => { + if waited > 0 && outcome.repaired > 0 { + tracing::info!( + "key repair: table {table_id}: document {id:?} rewritten after \ + waiting {} s for the writer holding it", + waited * CONFLICT_WAIT.as_secs() as u32 + ); + } + summary.documents_repaired += outcome.repaired; + summary.documents_failed += outcome.failed.len() as u64; + for (failed_id, reason) in outcome.failed { + tracing::error!( + "key repair: table {table_id}: cannot rebuild document \ + {failed_id:?}: {reason}" + ); + } + return; + } + Err(StorageError::TransactionConflict(_)) if waited < CONFLICT_WAITS => { + tokio::time::sleep(CONFLICT_WAIT).await; + } + Err(e) => { + tracing::error!( + "key repair: table {table_id}: cannot rewrite document {id:?}: {e}" + ); + summary.documents_failed += 1; + return; + } + } + } + } + + /// Rebuild the documents at `ids` in one transaction: read them in the + /// transaction, delete them, insert the rebuilt documents. Ids whose + /// document is gone or already carries its key are skipped. Ids whose + /// document cannot be rebuilt are reported in `failed` and left in place. + async fn replace_batch( + &self, + coll: &mongodb::Collection, + ids: &[Bson], + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], + ) -> Result { + let mut session = self + .client + .start_session() + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + let tx_options = self.write_transaction_options(); + for attempt in 0..TRANSIENT_RETRY_ATTEMPTS { + session + .start_transaction() + .with_options(tx_options.clone()) + .await + .map_err(|e| StorageError::Internal(e.to_string()))?; + match rebuild_in_session(coll, &mut session, ids, key_schema, attribute_definitions) + .await + { + Ok(outcome) if outcome.repaired == 0 => { + let _ = session.abort_transaction().await; + return Ok(outcome); + } + Ok(outcome) => match self.commit_write_transaction(&mut session, None).await? { + CommitOutcome::Committed => return Ok(outcome), + CommitOutcome::RerunBody => backoff_sleep(attempt).await, + }, + Err(BatchError::Mongo(e)) if is_transient_write_conflict(&e) => { + let _ = session.abort_transaction().await; + backoff_sleep(attempt).await; + } + Err(e) => { + let _ = session.abort_transaction().await; + return Err(StorageError::Internal(e.to_string())); + } + } + } + Err(StorageError::TransactionConflict( + "key repair: too many concurrent write conflicts, giving up".to_owned(), + )) + } +} + +/// `table_id`, key schema and attribute definitions of a catalog `tables` row. +fn table_key( + table_doc: &Document, +) -> Result<(String, Vec, Vec), StorageError> { + let table_id = table_doc + .get_str("table_id") + .map_err(|_| StorageError::Internal("missing table_id".to_owned()))? + .to_owned(); + let key_schema = table_doc + .get("key_schema") + .ok_or_else(|| StorageError::Internal(format!("table {table_id}: missing key_schema")))?; + let key_schema: Vec = bson::from_bson(key_schema.clone()) + .map_err(|e| StorageError::Internal(format!("table {table_id}: key_schema: {e}")))?; + let attribute_definitions = table_doc.get("attribute_definitions").ok_or_else(|| { + StorageError::Internal(format!("table {table_id}: missing attribute_definitions")) + })?; + let attribute_definitions: Vec = + bson::from_bson(attribute_definitions.clone()).map_err(|e| { + StorageError::Internal(format!("table {table_id}: attribute_definitions: {e}")) + })?; + Ok((table_id, key_schema, attribute_definitions)) +} + +/// What one batch transaction did. +#[derive(Debug, Default)] +struct BatchOutcome { + /// Documents the transaction rewrote. + repaired: u64, + /// Documents it left in place, with the reason. + failed: Vec<(Bson, String)>, +} + +/// Why a batch transaction could not run. +#[derive(Debug)] +enum BatchError { + /// A MongoDB error, which may be a transient write conflict. + Mongo(mongodb::error::Error), + /// The delete did not match the documents the transaction had just read, + /// so the batch cannot be inserted without risking a duplicate item. A + /// transaction holds its snapshot, so this does not happen; the caller + /// falls back to one transaction per document rather than trusting it. + Counts { deleted: u64, read: usize }, +} + +impl From for BatchError { + fn from(e: mongodb::error::Error) -> Self { + Self::Mongo(e) + } +} + +impl std::fmt::Display for BatchError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Mongo(e) => write!(f, "{e}"), + Self::Counts { deleted, read } => { + write!( + f, + "deleted {deleted} of {read} documents read in the transaction" + ) + } + } + } +} + +/// Read the documents at `ids` in `session`'s transaction, delete them and +/// insert the rebuilt documents. The read is inside the transaction, so a +/// write that landed after the scan is either rebuilt with the document or +/// turns the delete into a write conflict. +async fn rebuild_in_session( + coll: &mongodb::Collection, + session: &mut mongodb::ClientSession, + ids: &[Bson], + key_schema: &[KeySchemaElement], + attribute_definitions: &[AttributeDefinition], +) -> Result { + let mut outcome = BatchOutcome::default(); + let mut stale: Vec = Vec::with_capacity(ids.len()); + let mut rebuilt: Vec = Vec::with_capacity(ids.len()); + let mut cursor = coll + .find(doc! { "_id": { "$in": ids.to_vec() } }) + .session(&mut *session) + .await?; + while let Some(bad) = cursor.next(&mut *session).await { + let bad = bad?; + let Some(id) = bad.get("_id").cloned() else { + continue; + }; + if !is_stored_without_key(&bad) { + continue; + } + match repaired_document(&bad, key_schema, attribute_definitions) { + Ok(d) => { + stale.push(id); + rebuilt.push(d); + } + Err(e) => outcome.failed.push((id, e.to_string())), + } + } + if rebuilt.is_empty() { + return Ok(outcome); + } + let deleted = coll + .delete_many(doc! { "_id": { "$in": stale } }) + .session(&mut *session) + .await?; + if deleted.deleted_count != rebuilt.len() as u64 { + return Err(BatchError::Counts { + deleted: deleted.deleted_count, + read: rebuilt.len(), + }); + } + coll.insert_many(&rebuilt).session(session).await?; + outcome.repaired = rebuilt.len() as u64; + Ok(outcome) +} + +/// Log a progress line when the repaired count has crossed a multiple of +/// `PROGRESS_EVERY`. A batch is smaller than that interval, so each crossing +/// is logged once. +fn log_progress(before: u64, summary: &KeyRepairSummary) { + if before / PROGRESS_EVERY != summary.documents_repaired / PROGRESS_EVERY { + tracing::info!("key repair in progress: {summary}"); + } +} + +/// Run a whole repair pass in the background and log the outcome. +/// +/// `serve` starts this instead of awaiting the pass, so the listener opens at +/// once however many documents need rebuilding. Requests on a key whose +/// document the pass has not reached yet answer as they did before the repair +/// existed, so a long pass delays the repair of a key rather than the start +/// of the server. +pub(crate) fn spawn_key_repair(engine: Arc) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let started = std::time::Instant::now(); + match engine.repair_items_stored_without_key().await { + Ok(repair) if repair.is_clean() => { + tracing::debug!("startup key repair: {repair}"); + } + Ok(repair) => { + let seconds = started.elapsed().as_secs_f64(); + let rate = if seconds > 0.0 { + (repair.documents_repaired as f64) / seconds + } else { + 0.0 + }; + let line = format!( + "startup key repair finished: {repair}, {seconds:.2} s, \ + {rate:.0} documents a second" + ); + if repair.documents_failed == 0 && repair.tables_failed == 0 { + tracing::info!("{line}"); + } else { + tracing::warn!("{line}"); + } + } + Err(e) => tracing::error!("Failed to repair items stored without their key: {e}"), + } + }) +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests.rs b/crates/storage-mongodb/src/data_engine_live_tests.rs new file mode 100644 index 000000000..4ea858a58 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests.rs @@ -0,0 +1,503 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the startup repair of documents that an earlier build +//! stored without their key (issue #436), and the helpers the live test +//! modules below share. The tests of the document an UpdateItem writes when it +//! creates an item are in `data_engine_live_tests_create.rs`. +//! +//! They read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d so `cargo test --lib` passes without one; run them with +//! `cargo test -p extenddb-storage-mongodb --lib -- --ignored live_`. Each test +//! uses its own catalog and data databases and drops them at the end. + +use std::collections::HashMap; + +use bson::{Document, doc}; +use extenddb_core::expression::{ExpressionMaps, PathElement}; +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType, + TableKeyInfo, +}; + +use super::key_repair::KeyRepairSummary; +use crate::MongoEngine; +use crate::data::{data_collection_name, document_to_item, item_to_document, pk_filter}; + +#[path = "data_engine_live_tests_create.rs"] +mod create; + +#[path = "data_engine_live_tests_key_names.rs"] +mod key_names; + +#[path = "data_engine_live_tests_repair_keys.rs"] +mod repair_keys; + +#[path = "data_engine_live_tests_txn_create.rs"] +mod txn_create; + +#[path = "data_engine_live_tests_repair_pass.rs"] +mod repair_pass; + +fn uri() -> String { + std::env::var("EXTENDDB_TEST_MONGODB_URI").expect( + "set EXTENDDB_TEST_MONGODB_URI to a MongoDB replica set, e.g. \ + mongodb://127.0.0.1:27741/?directConnection=true", + ) +} + +/// An engine on fresh catalog and data databases. +async fn scratch_engine() -> MongoEngine { + let base = MongoEngine::new( + &uri(), + "us-east-1", + None, + mongodb::options::ReadConcern::snapshot(), + ) + .await + .expect("engine"); + let tag = uuid::Uuid::new_v4().simple().to_string(); + MongoEngine { + catalog_db: base.client.database(&format!("k{}_cat", &tag[..12])), + data_db: base.client.database(&format!("k{}_data", &tag[..12])), + ..base + } +} + +async fn finish(engine: MongoEngine) { + engine.catalog_db.drop().await.expect("drop catalog"); + engine.data_db.drop().await.expect("drop data"); +} + +fn key(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } +} + +fn def(name: &str, attribute_type: ScalarAttributeType) -> AttributeDefinition { + AttributeDefinition { + attribute_name: name.to_owned(), + attribute_type, + } +} + +/// Key info for a table with a string hash key `pk` and, with `range`, a +/// string sort key `sk`. The GSI cache is set to "no indexes" so an +/// unconditional UpdateItem takes the native fast path, as it does on a +/// served table without indexes once the first write has filled the cache. +fn table(engine: &MongoEngine, table_id: &str, range: bool) -> TableKeyInfo { + engine.gsi_cache_set(table_id, false); + let mut key_schema = vec![key("pk", KeyType::Hash)]; + let mut attribute_definitions = vec![def("pk", ScalarAttributeType::S)]; + if range { + key_schema.push(key("sk", KeyType::Range)); + attribute_definitions.push(def("sk", ScalarAttributeType::S)); + } + TableKeyInfo { + table_id: table_id.to_owned(), + table_name: table_id.to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions, + ..Default::default() + } +} + +fn item_key(range: bool) -> Item { + let mut key = Item::new(); + key.insert("pk".to_owned(), AttributeValue::S("u1".to_owned())); + if range { + key.insert("sk".to_owned(), AttributeValue::S("s1".to_owned())); + } + key +} + +fn with(mut item: Item, name: &str, value: AttributeValue) -> Item { + item.insert(name.to_owned(), value); + item +} + +fn s(v: &str) -> AttributeValue { + AttributeValue::S(v.to_owned()) +} + +fn n(v: &str) -> AttributeValue { + AttributeValue::N(v.to_owned()) +} + +fn path(name: &str) -> Vec { + vec![PathElement::Attribute(name.to_owned())] +} + +fn maps(value: AttributeValue) -> ExpressionMaps { + let mut values = HashMap::new(); + values.insert(":v".to_owned(), value); + ExpressionMaps::new(HashMap::new(), values) +} + +async fn stored(engine: &MongoEngine, key_info: &TableKeyInfo, key: &Item) -> Document { + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + let filter = + pk_filter(key, &key_info.key_schema, &key_info.attribute_definitions).expect("filter"); + coll.find_one(filter) + .await + .expect("read back") + .expect("the item's document exists") +} + +/// The document `item_to_document` writes for `item`, with version `v`. +fn expected_document(item: &Item, key_info: &TableKeyInfo, v: i64) -> Document { + let mut doc = item_to_document(item, &key_info.key_schema, &key_info.attribute_definitions) + .expect("document"); + doc.insert("_v", v); + doc +} + +/// Store a document the way the fast path did before the fix: an upsert whose +/// filter is the key and whose update sets one attribute, or, when `w` is +/// `None`, removes one. MongoDB copies the filter's `pk` and `sk_*` fields, +/// generates an ObjectId `_id`, and writes no key attribute into `item_data`. +/// The `None` form is main's REMOVE-only creating update: `$unset` +/// contributes nothing to the inserted document, so it has no `item_data` at +/// all. +async fn store_without_key( + engine: &MongoEngine, + key_info: &TableKeyInfo, + key: &Item, + w: Option<&str>, +) { + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + let filter = + pk_filter(key, &key_info.key_schema, &key_info.attribute_definitions).expect("filter"); + let update = match w { + Some(w) => doc! { "$set": { "item_data.w": { "S": w } }, "$inc": { "_v": 1_i64 } }, + None => doc! { "$unset": { "item_data.w": "" }, "$inc": { "_v": 1_i64 } }, + }; + coll.find_one_and_update(filter, update) + .upsert(true) + .await + .expect("upsert"); +} + +/// The repair rewrites a document stored without its key into the document +/// `item_to_document` builds, keeping the attributes and carrying `_v` over. +/// The item then reads with its key and a PutItem on it succeeds. Documents +/// with their key are left alone, and a second run finds nothing. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_rewrites_documents_stored_without_their_key() { + let engine = scratch_engine().await; + for range in [false, true] { + let table_id = format!("repair-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + store_without_key(&engine, &key_info, &key, Some("upd")).await; + store_without_key(&engine, &key_info, &key, Some("upd2")).await; + let bad = stored(&engine, &key_info, &key).await; + assert!( + bad.get_object_id("_id").is_ok(), + "{table_id}: the pre-fix shape has an ObjectId _id" + ); + assert!( + !document_to_item(&bad).expect("item").contains_key("pk"), + "{table_id}: and no key" + ); + + // A healthy item in the same collection is not touched. + let other = with(with(item_key(range), "w", s("fine")), "pk", s("u2")); + engine + .put_item_impl( + &key_info, + other.clone(), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + .expect("put the healthy item"); + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + tables_failed: 0, + documents_repaired: 1, + documents_failed: 0, + }, + "{table_id}" + ); + + let want = with(key.clone(), "w", s("upd2")); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 2), + "{table_id}: repaired document" + ); + assert_eq!( + document_to_item(&stored(&engine, &key_info, &other).await).expect("item"), + other, + "{table_id}: the healthy item is unchanged" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after repair: {e:?}")); + assert_eq!(old, Some(want), "{table_id}: put returns the repaired item"); + + let mut again = KeyRepairSummary::default(); + engine + .repair_table( + &table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut again, + ) + .await + .expect("second repair"); + assert!( + again.is_clean(), + "{table_id}: nothing left to repair: {again}" + ); + } + finish(engine).await; +} + +/// A pre-fix document with no `item_data` at all, which is what main's fast +/// path stored for a REMOVE-only creating UpdateItem, reads as an internal +/// error before the repair and as the key-only item after it. The repaired +/// document is the one the fixed fast path writes for the same request, and a +/// PutItem on the key then succeeds. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_rewrites_a_document_stored_with_no_item_data() { + let engine = scratch_engine().await; + for range in [false, true] { + let table_id = format!("repair-noitem-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + store_without_key(&engine, &key_info, &key, None).await; + + let bad = stored(&engine, &key_info, &key).await; + assert!( + bad.get_object_id("_id").is_ok(), + "{table_id}: the pre-fix shape has an ObjectId _id" + ); + assert!( + !bad.contains_key("item_data"), + "{table_id}: a REMOVE-only creating update stored no item_data: {bad:?}" + ); + assert!( + engine.get_item_impl(&key_info, &key).await.is_err(), + "{table_id}: the key reads as an error until the repair runs" + ); + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + tables_failed: 0, + documents_repaired: 1, + documents_failed: 0, + }, + "{table_id}" + ); + + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&key, &key_info, 1), + "{table_id}: repaired into the key-only item" + ); + assert_eq!( + engine.get_item_impl(&key_info, &key).await.expect("get"), + Some(key.clone()), + "{table_id}: the key reads back" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after repair: {e:?}")); + assert_eq!( + old, + Some(key.clone()), + "{table_id}: put returns the repaired key-only item" + ); + } + finish(engine).await; +} + +/// Insert `count` documents in the pre-fix shape directly, one per key +/// `u0..u`, with no key attribute in `item_data`. The upsert path in +/// `store_without_key` is one round trip per document, which is too slow for +/// the counts the batched repair needs. +async fn store_many_without_key(engine: &MongoEngine, key_info: &TableKeyInfo, count: usize) { + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + let docs: Vec = (0..count) + .map(|i| { + doc! { + "_id": bson::oid::ObjectId::new(), + "pk": format!("u{i}"), + "item_data": { "w": { "S": format!("value-{i}") } }, + "_v": 1_i64, + } + }) + .collect(); + coll.insert_many(docs).await.expect("plant"); +} + +/// A repair pass that spans more than one batch rebuilds every document. The +/// batch is 500 documents, so 1,100 crosses two boundaries and ends with a +/// partial batch. Every key reads back with its value, none is duplicated, +/// and a second pass finds nothing. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_spans_more_than_one_batch() { + const COUNT: usize = 1_100; + let engine = scratch_engine().await; + let key_info = table(&engine, "repair-batches", false); + store_many_without_key(&engine, &key_info, COUNT).await; + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + tables_failed: 0, + documents_repaired: COUNT as u64, + documents_failed: 0, + } + ); + + let coll = engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)); + assert_eq!( + coll.count_documents(doc! {}).await.expect("count"), + COUNT as u64, + "no document was duplicated or lost" + ); + assert_eq!( + coll.count_documents(doc! { "_id": { "$type": "objectId" } }) + .await + .expect("count"), + 0, + "every document carries its key text as _id" + ); + for i in [0_usize, 499, 500, 999, 1_000, COUNT - 1] { + let mut key = Item::new(); + key.insert("pk".to_owned(), s(&format!("u{i}"))); + assert_eq!( + engine.get_item_impl(&key_info, &key).await.expect("get"), + Some(with(key.clone(), "w", s(&format!("value-{i}")))), + "key u{i} reads back" + ); + } + + let mut again = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut again, + ) + .await + .expect("second repair"); + assert!(again.is_clean(), "nothing left to repair: {again}"); + finish(engine).await; +} + +/// The whole-catalog entry point finds the table through its catalog row. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_reads_the_key_schema_from_the_catalog() { + let engine = scratch_engine().await; + let key_info = table(&engine, "catalog-hash", false); + engine + .catalog_db + .collection::("tables") + .insert_one(doc! { + "_id": { "account_id": "123456789012", "table_name": "t" }, + "table_id": "catalog-hash", + "table_status": "ACTIVE", + "key_schema": bson::to_bson(&key_info.key_schema).expect("bson"), + "attribute_definitions": bson::to_bson(&key_info.attribute_definitions).expect("bson"), + }) + .await + .expect("catalog row"); + let key = item_key(false); + store_without_key(&engine, &key_info, &key, Some("upd")).await; + + let summary = engine + .repair_items_stored_without_key() + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 1, + tables_failed: 0, + documents_repaired: 1, + documents_failed: 0, + } + ); + assert_eq!( + document_to_item(&stored(&engine, &key_info, &key).await).expect("item"), + with(key, "w", s("upd")) + ); + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests_create.rs b/crates/storage-mongodb/src/data_engine_live_tests_create.rs new file mode 100644 index 000000000..0f76faa39 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_create.rs @@ -0,0 +1,227 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the document an UpdateItem writes when it creates an +//! item (issue #436), on the native fast path and on the transaction path. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use bson::{Document, doc}; +use extenddb_core::expression::{Expr, ExpressionMaps, UpdateAction}; +use extenddb_core::types::{AttributeValue, Item}; +use extenddb_storage::error::StorageError; + +use super::{ + expected_document, finish, item_key, maps, n, path, s, scratch_engine, stored, table, with, +}; +use crate::data::{data_collection_name, document_to_item}; + +/// An UpdateItem that creates the item, on the native fast path (SET with a +/// placeholder, REMOVE) and on the transaction path (numeric ADD), stores the +/// document `item_to_document` builds, plus `_v` 1: `_id` as the key text and +/// the key attributes in `item_data`. For SET and ADD, ALL_NEW returns the +/// item with its key, and a PutItem on the key replaces the item and returns +/// it as the old image. For REMOVE the test pins the stored document and that +/// a PutItem on the key succeeds. The service creates no item for a +/// REMOVE-only update and its ALL_NEW answer for one is not measured, so the +/// REMOVE case asserts neither the ALL_NEW answer nor the old image. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_that_creates_an_item_stores_its_key_like_put_item() { + let engine = scratch_engine().await; + let cases: [(&str, Vec, AttributeValue, Item); 3] = [ + ( + "set", + vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }], + s("upd"), + with(Item::new(), "w", s("upd")), + ), + ( + "remove", + vec![UpdateAction::Remove { path: path("w") }], + s("unused"), + Item::new(), + ), + ( + "add", + vec![UpdateAction::Add { + path: path("c"), + value: Expr::Placeholder(":v".to_owned()), + }], + n("5"), + with(Item::new(), "c", n("5")), + ), + ]; + for range in [false, true] { + for (name, actions, value, extra) in &cases { + let table_id = format!("{name}-{}", if range { "range" } else { "hash" }); + let key_info = table(&engine, &table_id, range); + let key = item_key(range); + let mut want = key.clone(); + want.extend(extra.clone()); + + let (old, new) = engine + .update_item_impl( + &key_info, + &key, + actions, + false, + true, + None, + &maps(value.clone()), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: update: {e:?}")); + assert_eq!(old, None, "{table_id}: no old image on creation"); + if *name != "remove" { + assert_eq!( + new, + Some(want.clone()), + "{table_id}: ALL_NEW carries the key" + ); + } + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 1), + "{table_id}: stored document" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put after update: {e:?}")); + if *name != "remove" { + assert_eq!(old, Some(want), "{table_id}: put returns the created item"); + } + let after = stored(&engine, &key_info, &key).await; + assert_eq!( + document_to_item(&after).expect("item"), + put, + "{table_id}: put replaced the item" + ); + } + } + finish(engine).await; +} + +/// Sixty-four UpdateItem calls creating the same item at once, on the fast +/// path, from eight runtime threads: every call succeeds, the stored +/// document has the shape PutItem writes, and `_v` counts every call, so +/// each one applied to the one document after losing the create race. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_concurrent_creating_updates_on_one_key_all_apply() { + const WRITERS: i64 = 64; + let engine = std::sync::Arc::new(scratch_engine().await); + for range in [false, true] { + let table_id = format!("race-{}", if range { "range" } else { "hash" }); + let key_info = std::sync::Arc::new(table(&engine, &table_id, range)); + let key = item_key(range); + let mut tasks = Vec::new(); + for i in 0..WRITERS { + let engine = engine.clone(); + let key_info = key_info.clone(); + let key = key.clone(); + tasks.push(tokio::spawn(async move { + let actions = vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }]; + engine + .update_item_impl( + &key_info, + &key, + &actions, + false, + true, + None, + &maps(s(&format!("w{i}"))), + None, + ) + .await + })); + } + for task in tasks { + let (old, new) = task + .await + .expect("task") + .unwrap_or_else(|e| panic!("{table_id}: a creating update failed: {e:?}")); + assert_eq!(old, None, "{table_id}"); + let new = new.expect("ALL_NEW"); + for (name, value) in &key { + assert_eq!( + new.get(name), + Some(value), + "{table_id}: ALL_NEW has the key" + ); + } + } + let doc = stored(&engine, &key_info, &key).await; + let item = document_to_item(&doc).expect("item"); + let w = item.get("w").cloned().expect("w was set"); + let mut want = key.clone(); + want.insert("w".to_owned(), w); + assert_eq!( + doc, + expected_document(&want, &key_info, WRITERS), + "{table_id}" + ); + } + let engine = std::sync::Arc::try_unwrap(engine) + .ok() + .expect("all tasks done"); + finish(engine).await; +} + +/// A creating upsert whose insert always collides on `_id` (another document +/// holds the key text as its `_id` under a different `pk`) exhausts the +/// retry budget and answers `TransactionConflict`, which the server maps to +/// `TransactionConflictException`, not an internal error: the duplicate-key +/// command error `findAndModify` returns is recognised. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_creating_update_that_always_loses_on_id_answers_transaction_conflict() { + let engine = scratch_engine().await; + let key_info = table(&engine, "collide-hash", false); + let key = item_key(false); + engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)) + .insert_one(doc! { "_id": "u1", "pk": "other", "item_data": { "pk": { "S": "other" } } }) + .await + .expect("colliding document"); + let actions = vec![UpdateAction::Set { + path: path("w"), + value: Expr::Placeholder(":v".to_owned()), + }]; + let err = engine + .update_item_impl( + &key_info, + &key, + &actions, + false, + true, + None, + &maps(s("upd")), + None, + ) + .await + .expect_err("the insert cannot succeed"); + assert!( + matches!(err, StorageError::TransactionConflict(_)), + "expected TransactionConflict, got {err:?}" + ); + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests_key_names.rs b/crates/storage-mongodb/src/data_engine_live_tests_key_names.rs new file mode 100644 index 000000000..ebf058fd6 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_key_names.rs @@ -0,0 +1,266 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for tables whose key attribute names are not plain +//! field names. +//! +//! The native UpdateItem upsert writes each key attribute as the path +//! `item_data.`, and MongoDB reads a dot in a path as a step into a +//! sub-document. A table with a key attribute named `a.b` or `s.k` must take +//! the session path instead, which stores the field `"a.b"` the way PutItem +//! does (see `data_engine_update_create.rs`). +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use std::collections::HashMap; + +use extenddb_core::expression::{Expr, ExpressionMaps, KeyCondition, UpdateAction}; +use extenddb_core::types::{AttributeValue, Item, KeyType, ScalarAttributeType, TableKeyInfo}; + +use super::{def, expected_document, finish, key, maps, path, s, scratch_engine, stored, with}; +use crate::MongoEngine; + +/// Key info for a table with a string hash key named `hash` and, when given, +/// a string range key named `range`. The GSI cache is set to "no indexes", so +/// an unconditional UpdateItem would take the native fast path if the key +/// names allowed it. +fn named_table( + engine: &MongoEngine, + table_id: &str, + hash: &str, + range: Option<&str>, +) -> TableKeyInfo { + engine.gsi_cache_set(table_id, false); + let mut key_schema = vec![key(hash, KeyType::Hash)]; + let mut attribute_definitions = vec![def(hash, ScalarAttributeType::S)]; + if let Some(range) = range { + key_schema.push(key(range, KeyType::Range)); + attribute_definitions.push(def(range, ScalarAttributeType::S)); + } + TableKeyInfo { + table_id: table_id.to_owned(), + table_name: table_id.to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions, + ..Default::default() + } +} + +fn named_key(hash: &str, range: Option<&str>) -> Item { + let mut key = Item::new(); + key.insert(hash.to_owned(), s("u1")); + if let Some(range) = range { + key.insert(range.to_owned(), s("s1")); + } + key +} + +/// `SET = :v`. +fn set(name: &str) -> Vec { + vec![UpdateAction::Set { + path: path(name), + value: Expr::Placeholder(":v".to_owned()), + }] +} + +/// The key condition ` = :p` and the map that supplies `:p`. +fn partition_condition(hash: &str) -> (KeyCondition, ExpressionMaps) { + let condition = KeyCondition { + pk_path: path(hash), + pk_value: Expr::Placeholder(":p".to_owned()), + extra_pk_conditions: Vec::new(), + sk_condition: None, + extra_sk_conditions: Vec::new(), + }; + let values = HashMap::from([(":p".to_owned(), s("u1"))]); + (condition, ExpressionMaps::new(HashMap::new(), values)) +} + +/// An UpdateItem that creates an item on a table with these key names +/// answers Ok, returns the item with its key under ALL_NEW and stores the +/// document PutItem stores. GetItem, Query and Scan return the item. An +/// UpdateItem that sets an attribute named `prefix`, which is a path prefix +/// of a key name, succeeds on the created item, and PutItem and DeleteItem +/// then read it as any other item. +async fn creation_matches_put_item( + engine: &MongoEngine, + table_id: &str, + hash: &str, + range: Option<&str>, + prefix: &str, +) { + let key_info = named_table(engine, table_id, hash, range); + let key = named_key(hash, range); + let created = with(key.clone(), "w", s("upd")); + + let (old, new) = engine + .update_item_impl( + &key_info, + &key, + &set("w"), + false, + true, + None, + &maps(s("upd")), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: creating update: {e:?}")); + assert_eq!(old, None, "{table_id}: no old image on creation"); + assert_eq!( + new, + Some(created.clone()), + "{table_id}: ALL_NEW has the key" + ); + assert_eq!( + stored(engine, &key_info, &key).await, + expected_document(&created, &key_info, 1), + "{table_id}: stored document" + ); + assert_eq!( + engine + .get_item_impl(&key_info, &key) + .await + .unwrap_or_else(|e| panic!("{table_id}: get: {e:?}")), + Some(created.clone()), + "{table_id}: GetItem returns the created item" + ); + + let (condition, query_maps) = partition_condition(hash); + let (query, _) = engine + .query_impl(&key_info, &condition, &query_maps, true, None, None, None) + .await + .unwrap_or_else(|e| panic!("{table_id}: query: {e:?}")); + assert_eq!(query, vec![created.clone()], "{table_id}: Query"); + let (scan, _) = engine + .scan_impl(&key_info, None, None, None, None, None) + .await + .unwrap_or_else(|e| panic!("{table_id}: scan: {e:?}")); + assert_eq!(scan, vec![created.clone()], "{table_id}: Scan"); + + let updated = with(created, prefix, s("p")); + let (_, new) = engine + .update_item_impl( + &key_info, + &key, + &set(prefix), + false, + true, + None, + &maps(s("p")), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: SET {prefix} on the created item: {e:?}")); + assert_eq!(new, Some(updated.clone()), "{table_id}: SET {prefix}"); + assert_eq!( + stored(engine, &key_info, &key).await, + expected_document(&updated, &key_info, 2), + "{table_id}: stored document after SET {prefix}" + ); + + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put: {e:?}")); + assert_eq!( + old, + Some(updated), + "{table_id}: PutItem returns the old item" + ); + let deleted = engine + .delete_item_impl( + &key_info, + &key, + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: delete: {e:?}")); + assert_eq!( + deleted, + Some(put), + "{table_id}: DeleteItem returns the item" + ); +} + +/// A hash key named `a.b`: the creating UpdateItem would have stored +/// `item_data: { a: { b: ... } }`, which no read parses. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_creates_an_item_when_the_hash_key_name_has_a_dot() { + let engine = scratch_engine().await; + creation_matches_put_item(&engine, "dot-hash", "a.b", None, "a").await; + finish(engine).await; +} + +/// A range key named `s.k`, with a plain hash key. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_creates_an_item_when_the_range_key_name_has_a_dot() { + let engine = scratch_engine().await; + creation_matches_put_item(&engine, "dot-range", "pk", Some("s.k"), "s").await; + finish(engine).await; +} + +/// An UpdateItem on an item that already exists may set an attribute named +/// like a path prefix of a key name (`a` for the key `a.b`). Through the +/// native upsert the `$set` of `item_data.a` and the `$setOnInsert` of +/// `item_data.a.b` conflict, and MongoDB refuses the update before it looks +/// for the document, so this worked on main and must keep working. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_update_item_may_set_an_attribute_named_like_a_prefix_of_a_key_name() { + let engine = scratch_engine().await; + let cases: [(&str, Option<&str>, &str); 2] = [("a.b", None, "a"), ("pk", Some("s.k"), "s")]; + for (hash, range, prefix) in cases { + let table_id = format!("prefix-{prefix}"); + let key_info = named_table(&engine, &table_id, hash, range); + let key = named_key(hash, range); + let existing = with(key.clone(), "w", s("put")); + engine + .put_item_impl( + &key_info, + existing.clone(), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: put: {e:?}")); + + let (_, new) = engine + .update_item_impl( + &key_info, + &key, + &set(prefix), + false, + true, + None, + &maps(AttributeValue::S("p".to_owned())), + None, + ) + .await + .unwrap_or_else(|e| panic!("{table_id}: SET {prefix} on a stored item: {e:?}")); + let updated = with(existing, prefix, s("p")); + assert_eq!(new, Some(updated.clone()), "{table_id}: ALL_NEW"); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&updated, &key_info, 1), + "{table_id}: stored document" + ); + } + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs b/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs new file mode 100644 index 000000000..039368c37 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_repair_keys.rs @@ -0,0 +1,117 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the startup repair of documents stored without their +//! key, for key shapes the tests in the parent module do not cover. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use extenddb_core::expression::ExpressionMaps; +use extenddb_core::types::{AttributeValue, Item, KeyType, ScalarAttributeType, TableKeyInfo}; + +use super::{ + KeyRepairSummary, def, finish, key, s, scratch_engine, store_without_key, stored, with, +}; +use crate::MongoEngine; +use crate::data::{document_to_item, item_to_document}; + +/// Key info for a table with a string hash key `pk` and a numeric sort key +/// `sk`, with the GSI cache set to "no indexes". +fn numeric_range_table(engine: &MongoEngine, table_id: &str) -> TableKeyInfo { + engine.gsi_cache_set(table_id, false); + let key_schema = vec![key("pk", KeyType::Hash), key("sk", KeyType::Range)]; + TableKeyInfo { + table_id: table_id.to_owned(), + table_name: table_id.to_owned(), + key_schema: key_schema.clone(), + base_key_schema: key_schema, + attribute_definitions: vec![ + def("pk", ScalarAttributeType::S), + def("sk", ScalarAttributeType::N), + ], + ..Default::default() + } +} + +/// The key `pk = "u1"`, `sk = `, with the number in the plain decimal +/// form a request carries. +fn number_key(number: &str) -> Item { + let mut key = Item::new(); + key.insert("pk".to_owned(), s("u1")); + key.insert("sk".to_owned(), AttributeValue::N(number.to_owned())); + key +} + +/// A numeric sort key below 1e-6, which `Decimal128` displays in exponent +/// form (`1E-7`), is repaired into the `_id` that PutItem builds from the +/// plain decimal. A PutItem on each repaired key then succeeds and returns +/// the repaired item as the old image; built on the exponent form, the repair +/// would leave an `_id` that PutItem cannot replace. The values at and above +/// 1e-6 are repaired the same way and show the other side of the switch. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_keeps_a_small_numeric_sort_key_writable() { + const NUMBERS: [&str; 4] = ["0.0000001", "-0.00000012345", "0.000001", "12.5"]; + let engine = scratch_engine().await; + let key_info = numeric_range_table(&engine, "repair-small-numbers"); + for number in NUMBERS { + store_without_key(&engine, &key_info, &number_key(number), Some("upd")).await; + } + + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 0, + tables_failed: 0, + documents_repaired: NUMBERS.len() as u64, + documents_failed: 0, + } + ); + + for number in NUMBERS { + let key = number_key(number); + let repaired = with(key.clone(), "w", s("upd")); + let put = with(key.clone(), "w", s("put")); + let old = engine + .put_item_impl( + &key_info, + put.clone(), + true, + None, + &ExpressionMaps::default(), + None, + ) + .await + .unwrap_or_else(|e| panic!("{number}: PutItem on the repaired key: {e:?}")); + assert_eq!( + old, + Some(repaired), + "{number}: PutItem returns the repaired item" + ); + + let after = stored(&engine, &key_info, &key).await; + let want = item_to_document(&put, &key_info.key_schema, &key_info.attribute_definitions) + .expect("document"); + assert_eq!( + after.get("_id"), + want.get("_id"), + "{number}: the _id is the one PutItem builds" + ); + assert_eq!( + document_to_item(&after).expect("item"), + put, + "{number}: PutItem replaced the item" + ); + } + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs b/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs new file mode 100644 index 000000000..3d9cd5d8d --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_repair_pass.rs @@ -0,0 +1,200 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for the pass that finds the documents to repair: how it +//! reads each table, how it goes on after a table it cannot read, and the task +//! that `serve` starts to run it. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use std::sync::Arc; + +use bson::{Bson, Document, doc}; +use extenddb_core::types::TableKeyInfo; +use futures::TryStreamExt; + +use super::{ + KeyRepairSummary, expected_document, finish, item_key, s, scratch_engine, + store_many_without_key, store_without_key, stored, table, with, +}; +use crate::data::{data_collection_name, document_to_item}; +use crate::data_engine::key_repair::spawn_key_repair; + +/// A catalog `tables` row for `table_id` with the key schema of `key_info`. +fn catalog_row(table_name: &str, table_id: &str, key_info: &TableKeyInfo) -> Document { + doc! { + "_id": { "account_id": "123456789012", "table_name": table_name }, + "table_id": table_id, + "table_status": "ACTIVE", + "key_schema": bson::to_bson(&key_info.key_schema).expect("bson"), + "attribute_definitions": bson::to_bson(&key_info.attribute_definitions).expect("bson"), + } +} + +/// An integer field of a profiler entry, which MongoDB writes as int32 or +/// int64. +fn count(entry: &Document, field: &str) -> i64 { + match entry.get(field) { + Some(Bson::Int32(n)) => i64::from(*n), + Some(Bson::Int64(n)) => *n, + other => panic!("{field} in {entry:?} is {other:?}, not an integer"), + } +} + +/// A scan that stops on an error does not stop the pass. The first catalog +/// row carries a table id with a `$`, which MongoDB refuses in a collection +/// name, so the scan of its data collection fails. The second row is an +/// ordinary table with one document stored without its key. The pass counts +/// the first table as stopped by an error, repairs the second, and returns Ok. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_goes_on_to_the_next_table_when_a_scan_fails() { + let engine = scratch_engine().await; + let key_info = table(&engine, "pass-good", false); + let key = item_key(false); + let tables = engine.catalog_db.collection::("tables"); + tables + .insert_one(catalog_row("unreadable", "pass-bad$id", &key_info)) + .await + .expect("catalog row of the table that cannot be scanned"); + tables + .insert_one(catalog_row("readable", &key_info.table_id, &key_info)) + .await + .expect("catalog row of the readable table"); + store_without_key(&engine, &key_info, &key, Some("upd")).await; + + let summary = engine + .repair_items_stored_without_key() + .await + .expect("a table that cannot be scanned does not fail the pass"); + assert_eq!( + summary, + KeyRepairSummary { + tables_scanned: 2, + tables_failed: 1, + documents_repaired: 1, + documents_failed: 0, + } + ); + assert!(!summary.is_clean(), "a stopped scan is not a clean pass"); + assert_eq!( + document_to_item(&stored(&engine, &key_info, &key).await).expect("item"), + with(key, "w", s("upd")), + "the table after the failing one was repaired" + ); + finish(engine).await; +} + +/// The scan reads `_id` only. With the projection the `_id` index answers +/// it, and the server's profiler records that the find returned every +/// document it found and examined none; without the projection it examines +/// each one to return its fields. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_repair_scan_is_answered_from_the_id_index() { + const DOCUMENTS: usize = 20; + let engine = scratch_engine().await; + let key_info = table(&engine, "pass-scan", false); + store_many_without_key(&engine, &key_info, DOCUMENTS).await; + + engine + .data_db + .run_command(doc! { "profile": 2 }) + .await + .expect("profiler on"); + let mut summary = KeyRepairSummary::default(); + engine + .repair_table( + &key_info.table_id, + &key_info.key_schema, + &key_info.attribute_definitions, + &mut summary, + ) + .await + .expect("repair"); + engine + .data_db + .run_command(doc! { "profile": 0 }) + .await + .expect("profiler off"); + assert_eq!(summary.documents_repaired, DOCUMENTS as u64); + + let namespace = format!( + "{}.{}", + engine.data_db.name(), + data_collection_name(&key_info.table_id) + ); + let queries: Vec = engine + .data_db + .collection::("system.profile") + .find(doc! { "ns": &namespace, "op": "query" }) + .await + .expect("profile entries") + .try_collect() + .await + .expect("profile entries"); + // The scan filters on the `_id` type; the batch transaction's own read + // filters on a list of ids. + let scans: Vec<&Document> = queries + .iter() + .filter(|entry| { + entry + .get_document("command") + .and_then(|command| command.get_document("filter")) + .and_then(|filter| filter.get_document("_id")) + .is_ok_and(|id| id.contains_key("$type")) + }) + .collect(); + assert_eq!(scans.len(), 1, "one scan of {namespace}: {queries:?}"); + let scan = scans[0]; + assert_eq!(count(scan, "nreturned"), DOCUMENTS as i64, "{scan:?}"); + assert_eq!( + count(scan, "docsExamined"), + 0, + "the scan is answered from the _id index: {scan:?}" + ); + finish(engine).await; +} + +/// `spawn_key_repair`, the task `serve` starts (`lib.rs`), runs the whole +/// pass: with a catalog table whose collection holds a document stored +/// without its key, the document is rebuilt by the time the task finishes, +/// and the task does finish. The other tests call `repair_table` and +/// `repair_items_stored_without_key` directly. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_the_spawned_key_repair_rewrites_documents_stored_without_their_key() { + let engine = scratch_engine().await; + let key_info = table(&engine, "pass-spawned", false); + let key = item_key(false); + engine + .catalog_db + .collection::("tables") + .insert_one(catalog_row("spawned", &key_info.table_id, &key_info)) + .await + .expect("catalog row"); + store_without_key(&engine, &key_info, &key, Some("upd")).await; + assert!( + stored(&engine, &key_info, &key) + .await + .get_object_id("_id") + .is_ok(), + "planted in the shape the pre-fix upsert stored" + ); + + let engine = Arc::new(engine); + spawn_key_repair(engine.clone()) + .await + .expect("the repair task finishes"); + + let want = with(key.clone(), "w", s("upd")); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, 1), + "the task repaired the document" + ); + let engine = Arc::try_unwrap(engine) + .ok() + .expect("the finished task dropped its engine"); + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs b/crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs new file mode 100644 index 000000000..c057786aa --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_live_tests_txn_create.rs @@ -0,0 +1,113 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Live MongoDB tests for UpdateItem calls that create an item on the +//! transaction path. A numeric ADD takes it: the call finds no item inside a +//! snapshot transaction and inserts with a `replace_one` upsert. +//! +//! The tests read the replica set URI from `EXTENDDB_TEST_MONGODB_URI` and are +//! `#[ignore]`d, like the tests in the parent module. + +use std::sync::Arc; + +use bson::{Document, doc}; +use extenddb_core::expression::{Expr, UpdateAction}; +use extenddb_storage::error::StorageError; + +use super::{ + expected_document, finish, item_key, maps, n, path, scratch_engine, stored, table, with, +}; +use crate::data::data_collection_name; + +fn add_one() -> Vec { + vec![UpdateAction::Add { + path: path("c"), + value: Expr::Placeholder(":v".to_owned()), + }] +} + +/// Eight numeric ADD calls creating the same item at once all apply. Each +/// finds no item, and the creators that lose to the first insert get a write +/// conflict from MongoDB, which the engine retries: the retry reads the +/// winner's item and adds to it. The stored `c` is the sum and `_v` counts +/// every call. +#[tokio::test(flavor = "multi_thread", worker_threads = 8)] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_concurrent_numeric_adds_that_create_one_item_all_apply() { + const WRITERS: i64 = 8; + let engine = Arc::new(scratch_engine().await); + let key_info = Arc::new(table(&engine, "race-add", false)); + let key = item_key(false); + let mut tasks = Vec::new(); + for _ in 0..WRITERS { + let engine = engine.clone(); + let key_info = key_info.clone(); + let key = key.clone(); + tasks.push(tokio::spawn(async move { + engine + .update_item_impl( + &key_info, + &key, + &add_one(), + false, + true, + None, + &maps(n("1")), + None, + ) + .await + })); + } + for task in tasks { + task.await + .expect("task") + .unwrap_or_else(|e| panic!("a creating ADD failed: {e:?}")); + } + let want = with(key.clone(), "c", n(&WRITERS.to_string())); + assert_eq!( + stored(&engine, &key_info, &key).await, + expected_document(&want, &key_info, WRITERS) + ); + let engine = Arc::try_unwrap(engine).ok().expect("all tasks done"); + finish(engine).await; +} + +/// A creating numeric ADD whose `_id` is already held by a document with a +/// different `pk` fails with the duplicate-key error, and the call does not +/// run again. The other document is visible in the transaction's snapshot, +/// so every rerun would meet it again; a creator that loses a race for the +/// key is a write conflict instead, which the test above covers. +#[tokio::test] +#[ignore = "needs a MongoDB replica set in EXTENDDB_TEST_MONGODB_URI"] +async fn live_creating_add_that_collides_on_id_answers_the_duplicate_key_error() { + let engine = scratch_engine().await; + let key_info = table(&engine, "collide-add", false); + let key = item_key(false); + engine + .data_db + .collection::(&data_collection_name(&key_info.table_id)) + .insert_one(doc! { "_id": "u1", "pk": "other", "item_data": { "pk": { "S": "other" } } }) + .await + .expect("colliding document"); + + let err = engine + .update_item_impl( + &key_info, + &key, + &add_one(), + false, + true, + None, + &maps(n("1")), + None, + ) + .await + .expect_err("the insert cannot succeed"); + match err { + StorageError::Internal(message) => assert!( + message.contains("E11000 duplicate key error"), + "expected the server's duplicate-key error, got {message}" + ), + other => panic!("expected Internal with the duplicate-key error, got {other:?}"), + } + finish(engine).await; +} diff --git a/crates/storage-mongodb/src/data_engine_update_create.rs b/crates/storage-mongodb/src/data_engine_update_create.rs new file mode 100644 index 000000000..a7f4e7e01 --- /dev/null +++ b/crates/storage-mongodb/src/data_engine_update_create.rs @@ -0,0 +1,176 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! The native UpdateItem upsert, which creates the item when it does not +//! exist yet. +//! +//! `update_item_impl` runs an unconditional single-attribute update on a +//! table without indexes as one `findAndModify` with `upsert: true`, filtered +//! on the key fields `pk` and `sk_*`. When no document matches, MongoDB builds +//! the inserted document from those filter fields and the update operators. +//! Without more, that document gets a generated ObjectId `_id` and an +//! `item_data` without the key attributes, while PutItem writes `_id` as the +//! key text and the whole item (`item_to_document`). Every later PutItem on +//! the key then fails on the immutable `_id` (issue #436). +//! [`MongoEngine::upsert_native_update`] adds a `$setOnInsert` with the `_id` +//! and the key attributes, so the created document is the one +//! `item_to_document` builds for PutItem, plus `_v` 1 from the update's +//! `$inc`. +//! +//! Two writers creating the same item at once each find no document and each +//! insert; the second insert fails with a duplicate-key error on `_id`. The +//! data collection also has a unique index on the key fields, and MongoDB +//! retries an upsert itself only when the error names an index whose fields +//! the filter matches by equality. The `_id` collision is reported first and +//! the filter does not name `_id`, so the server does not retry it and this +//! module does: the item now exists, so the next attempt updates it. +//! +//! The `$setOnInsert` names each key attribute as the path +//! `item_data.`, and MongoDB reads a dot in that path as a +//! step into a sub-document. A key attribute named `a.b` would be inserted as +//! `item_data: { a: { b: ... } }`, where PutItem stores the field `"a.b"`, +//! and the same path conflicts with an update of an attribute named `a`. +//! A table with a key attribute name that `key_names_are_plain_paths` +//! rejects therefore does not take the native upsert; it takes the session +//! path, which writes the document `item_to_document` builds. + +use bson::Document; +use extenddb_core::types::{Item, KeySchemaElement, TableKeyInfo}; +use extenddb_storage::error::StorageError; +use mongodb::options::{FindOneAndUpdateOptions, ReturnDocument}; + +use super::{TRANSIENT_RETRY_ATTEMPTS, backoff_sleep}; +use crate::MongoEngine; +use crate::data::key_fields::insert_only_key_fields; + +/// Whether `name` is usable as one field name in a MongoDB update path. A dot +/// would split it into nested fields and a leading `$` would read as an +/// operator or positional path. +pub(super) fn is_plain_path_name(name: &str) -> bool { + !name.contains('.') && !name.starts_with('$') +} + +/// Whether every key attribute of `key_schema` is a plain path name, so that +/// [`MongoEngine::upsert_native_update`] can write it as `item_data.`. +pub(super) fn key_names_are_plain_paths(key_schema: &[KeySchemaElement]) -> bool { + key_schema + .iter() + .all(|ks| is_plain_path_name(&ks.attribute_name)) +} + +/// Whether `e` is a duplicate-key error (E11000) as `findAndModify` reports +/// it: a command error, where `insert_one` reports a write error. +fn is_duplicate_key_command_error(e: &mongodb::error::Error) -> bool { + matches!(*e.kind, mongodb::error::ErrorKind::Command(ref c) if c.code == 11000) +} + +impl MongoEngine { + /// Apply the native `update` (a `$set`/`$unset`/`$inc` operator document) + /// to the item at `key_filter`, creating the item when it does not exist, + /// and return the document after the update. + /// + /// `key` and `key_info` supply the `$setOnInsert` fields that make a + /// created document match what PutItem writes. Attempts that lose the + /// create race share the transient retry budget and back off between + /// them; when every attempt loses, the error is the same + /// `TransactionConflict` the transaction path returns. + pub(crate) async fn upsert_native_update( + &self, + coll: &mongodb::Collection, + key_filter: &Document, + mut update: Document, + key: &Item, + key_info: &TableKeyInfo, + ) -> Result, StorageError> { + update.insert( + "$setOnInsert", + insert_only_key_fields(key, &key_info.key_schema, &key_info.attribute_definitions)?, + ); + let opts = FindOneAndUpdateOptions::builder() + .upsert(true) + .return_document(ReturnDocument::After) + .build(); + for attempt in 0..TRANSIENT_RETRY_ATTEMPTS { + match coll + .find_one_and_update(key_filter.clone(), update.clone()) + .with_options(opts.clone()) + .await + { + Ok(doc) => return Ok(doc), + Err(e) if is_duplicate_key_command_error(&e) => backoff_sleep(attempt).await, + Err(e) => return Err(StorageError::Internal(e.to_string())), + } + } + Err(StorageError::TransactionConflict( + "UpdateItem: too many concurrent write conflicts, giving up".to_owned(), + )) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use bson::doc; + use extenddb_core::types::KeyType; + + fn command_error(code: i32) -> mongodb::error::Error { + let command_error: mongodb::error::CommandError = bson::from_document(doc! { + "code": code, + "codeName": "SyntheticTestError", + "errmsg": "synthetic command error", + }) + .expect("static CommandError document deserializes"); + mongodb::error::ErrorKind::Command(command_error).into() + } + + fn write_error(code: i32) -> mongodb::error::Error { + let write_error: mongodb::error::WriteError = bson::from_document(doc! { + "code": code, + "codeName": "SyntheticTestError", + "errmsg": "synthetic write error", + }) + .expect("static WriteError document deserializes"); + mongodb::error::ErrorKind::Write(mongodb::error::WriteFailure::WriteError(write_error)) + .into() + } + + #[test] + fn only_the_command_form_of_e11000_is_a_lost_create_race() { + assert!(is_duplicate_key_command_error(&command_error(11000))); + assert!(!is_duplicate_key_command_error(&command_error(112))); + assert!(!is_duplicate_key_command_error(&write_error(11000))); + } + + fn schema(names: &[(&str, KeyType)]) -> Vec { + names + .iter() + .map(|(name, key_type)| KeySchemaElement { + attribute_name: (*name).to_owned(), + key_type: *key_type, + }) + .collect() + } + + #[test] + fn a_dot_or_a_leading_dollar_in_a_key_name_leaves_the_native_upsert() { + let plain = [("pk", KeyType::Hash), ("sk", KeyType::Range)]; + assert!(key_names_are_plain_paths(&schema(&plain))); + assert!(key_names_are_plain_paths(&schema(&[( + "a$b", + KeyType::Hash + )]))); + + for hash in ["a.b", "$k", ".", "pk."] { + assert!( + !key_names_are_plain_paths(&schema(&[(hash, KeyType::Hash)])), + "hash key {hash:?}" + ); + } + for range in ["s.k", "$s"] { + let keys = [("pk", KeyType::Hash), (range, KeyType::Range)]; + assert!( + !key_names_are_plain_paths(&schema(&keys)), + "range key {range:?}" + ); + } + } +} diff --git a/crates/storage-mongodb/src/lib.rs b/crates/storage-mongodb/src/lib.rs index b27cc0baf..37d9ed246 100644 --- a/crates/storage-mongodb/src/lib.rs +++ b/crates/storage-mongodb/src/lib.rs @@ -139,6 +139,17 @@ fn server_components_factory( let engine = Arc::new(engine); + // An UpdateItem on a build before the fix for issue #436 could create + // an item whose document has no key. Rewrite any such document, so + // that reads return the key and PutItem on it succeeds (see + // `data_engine::key_repair`). This runs in the background: the pass + // costs time in proportion to the number of affected documents, and + // the listener must not wait for it. A key the pass has not reached + // yet answers as it did before the repair existed, and each batch is + // its own transaction, so running this while another process serves + // the same data is safe. + data_engine::key_repair::spawn_key_repair(engine.clone()); + // Create one shared catalog/management client. MongoDB Client clones // share the underlying pool, so max_catalog_connections limits the // combined catalog and auth traffic rather than creating two diff --git a/devtools/run-mongodb-tests b/devtools/run-mongodb-tests index 2a287798a..4a3b03885 100755 --- a/devtools/run-mongodb-tests +++ b/devtools/run-mongodb-tests @@ -23,6 +23,10 @@ # For tests that talk to MongoDB directly (the storage # crate's snapshot tests). Tear down with # docker rm -f extenddb-runtests-mongo-. +# --rust-live Also run the Rust live tests of the storage-mongodb +# crate against this runner's replica set. They are +# `#[ignore]`d, because they need a replica set, and +# they are the only tests of the startup key repair. # -h, --help Show this help # # Arguments after `--` are passed to `devtools/run-tests` verbatim. @@ -32,6 +36,7 @@ # devtools/run-mongodb-tests # devtools/run-mongodb-tests -- --pytest --filter test_put_item # devtools/run-mongodb-tests --keep -- --pytest --filter test_rfc0003 +# devtools/run-mongodb-tests --rust-live -- --rust --rust-integration # # Prerequisites: # - Docker running @@ -54,6 +59,7 @@ MONGO_PORT=27021 OUTPUT_DIR="" KEEP=false MONGO_ONLY=false +RUST_LIVE=false RUN_TESTS_ARGS=() while [[ $# -gt 0 ]]; do @@ -63,8 +69,9 @@ while [[ $# -gt 0 ]]; do --output) OUTPUT_DIR="$2"; shift 2 ;; --keep) KEEP=true; shift ;; --mongo-only) MONGO_ONLY=true; KEEP=true; shift ;; + --rust-live) RUST_LIVE=true; shift ;; --) shift; RUN_TESTS_ARGS=("$@"); break ;; - -h|--help) sed -n '5,40p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; + -h|--help) sed -n '5,45p' "$0" | sed 's/^# \{0,1\}//'; exit 0 ;; *) echo "error: unknown option: $1" >&2; exit 1 ;; esac done @@ -317,7 +324,29 @@ export EXTENDDB_CONFIG="$CONFIG" # tempfiles land in an allowed location on both Linux and macOS. export TMPDIR="$TMP_CANON" +# The Rust live tests in crates/storage-mongodb connect to the replica set +# directly, rather than through the server, so they read its URI from here. +export EXTENDDB_TEST_MONGODB_URI="mongodb://localhost:$MONGO_PORT/?replicaSet=rs0&directConnection=true" + RC=0 devtools/run-tests --extenddb --backend mongodb --config "$CONFIG" "${RUN_TESTS_ARGS[@]}" || RC=$? +# The live tests of the storage-mongodb crate are `#[ignore]`d, so +# `cargo test` does not run them: they need the replica set this runner +# started. They are the only tests of the startup key repair, which is why +# they run here rather than in the plain `cargo test --workspace` job. Each +# one works in its own pair of databases and drops them at the end, so they +# do not touch the served instance's catalog or data. +if $RUST_LIVE; then + echo "" + echo "=== MongoDB live Rust tests (storage-mongodb, --ignored) ===" + if cargo test -p extenddb-storage-mongodb --lib -- \ + --ignored --test-threads=1 data_engine::live_tests::; then + echo " ✓ MongoDB live Rust tests passed" + else + echo " ✗ MongoDB live Rust tests failed" + RC=1 + fi +fi + exit $RC diff --git a/tests/test_update_item_creates_key.py b/tests/test_update_item_creates_key.py new file mode 100644 index 000000000..cee23776e --- /dev/null +++ b/tests/test_update_item_creates_key.py @@ -0,0 +1,415 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""An UpdateItem that creates an item stores the item with its key. + +Amazon DynamoDB creates the item from the key and the update expression: +GetItem returns the key attributes with the updated ones, ALL_NEW on creation +returns the whole new item including the key, UPDATED_NEW returns only the +updated attributes, ALL_OLD and UPDATED_OLD on creation return no Attributes, +and a later PutItem, UpdateItem, Query, Scan and DeleteItem on the key treat +the item as any other. On the MongoDB backend an unconditional UpdateItem on a +table without an index used to store the created item without its key +attributes, after which GetItem returned the item without the key and every +PutItem on the key answered HTTP 500 (issue #436). The race tests run the +issue's shape: writers creating the same keys with UpdateItem and PutItem at +once. + +Expected responses were measured against the service on 2026-10-23 (probe +W4-1 in the fix report). +""" + +from __future__ import annotations + +import os +import threading +import uuid +from dataclasses import dataclass, field + +import boto3 +import pytest +from botocore.config import Config +from botocore.exceptions import ClientError + +from conftest import scoped_table + +WRITERS = 8 +KEYS = 30 +RETURN_VALUES = ["ALL_NEW", "UPDATED_NEW", "ALL_OLD", "UPDATED_OLD", "NONE"] + + +# TEMPORARY: on MongoDB, a put without a condition on a table with no index and +# no stream returns HTTP 500 when it loses the create race to another writer's +# insert. The MongoDB write-race fix repairs it and lands separately. +# TODO: remove this marker once the MongoDB write-race fix is on main. +# Only an assertion counts as the expected failure. The marker is not strict, +# because the race does not fire in every run. The MongoDB test runner sets +# EXTENDDB_TEST_MONGODB_CONTAINER. +XFAIL_UNTIL_MONGODB_FIX = pytest.mark.xfail( + bool(os.environ.get("EXTENDDB_TEST_MONGODB_CONTAINER", "").strip()), + reason="MongoDB returns HTTP 500 for a lost create race, fixed by the MongoDB write-race fix", + raises=AssertionError, + strict=False, +) + + +def _s(name: str) -> dict: + return {"AttributeName": name, "AttributeType": "S"} + + +def _k(name: str, kind: str) -> dict: + return {"AttributeName": name, "KeyType": kind} + + +@pytest.fixture(scope="module") +def hash_table(dynamodb_client): + with scoped_table(dynamodb_client) as name: + yield name + + +@pytest.fixture(scope="module") +def range_table(dynamodb_client): + with scoped_table( + dynamodb_client, + attribute_definitions=[_s("pk"), _s("sk")], + key_schema=[_k("pk", "HASH"), _k("sk", "RANGE")], + ) as name: + yield name + + +@pytest.fixture(scope="module") +def dotted_hash_table(dynamodb_client): + """A table whose hash key name contains a dot.""" + with scoped_table( + dynamodb_client, + attribute_definitions=[_s("a.b")], + key_schema=[_k("a.b", "HASH")], + ) as name: + yield name + + +@pytest.fixture(scope="module") +def dotted_range_table(dynamodb_client): + """A table whose range key name contains a dot.""" + with scoped_table( + dynamodb_client, + attribute_definitions=[_s("pk"), _s("s.k")], + key_schema=[_k("pk", "HASH"), _k("s.k", "RANGE")], + ) as name: + yield name + + +@pytest.fixture(scope="module") +def raw_client(endpoint_url): + """A client that never retries, so every failure is seen.""" + kwargs: dict = { + "service_name": "dynamodb", + "region_name": os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + "config": Config( + retries={"total_max_attempts": 1, "mode": "standard"}, + max_pool_connections=WRITERS * 4, + ), + } + if endpoint_url: + kwargs["endpoint_url"] = endpoint_url + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +KEY_NAMES = { + "hash": ("pk", None), + "range": ("pk", "sk"), + "dotted_hash": ("a.b", None), + "dotted_range": ("pk", "s.k"), +} + + +def _key(table_kind: str, pk: str) -> dict: + hash_name, range_name = KEY_NAMES[table_kind] + key = {hash_name: {"S": pk}} + if range_name: + key[range_name] = {"S": "s"} + return key + + +def _get(client, table: str, key: dict) -> dict | None: + return client.get_item(TableName=table, Key=key, ConsistentRead=True).get("Item") + + +def _warm(client, table: str, table_kind: str) -> None: + """Write another key first. + + The MongoDB backend takes its native UpdateItem path only once a write has + recorded that the table has no index; the first write on a fresh table + takes the transaction path, which did not have the defect. The issue + reproduced it the same way. + """ + client.put_item(TableName=table, Item={**_key(table_kind, f"warm-{uuid.uuid4().hex}")}) + + +def _create(client, table: str, key: dict, **extra) -> dict: + """UpdateItem SET w = upd on a key that does not exist yet.""" + resp = client.update_item( + TableName=table, + Key=key, + UpdateExpression="SET w = :w", + ExpressionAttributeValues={":w": {"S": "upd"}}, + **extra, + ) + resp.pop("ResponseMetadata", None) + return resp + + +def _put_then_get(client, table: str, key: dict) -> None: + """PutItem on a key created by UpdateItem succeeds and replaces the item.""" + put = {**key, "w": {"S": "put"}} + client.put_item(TableName=table, Item=put) + assert _get(client, table, key) == put + + +@pytest.mark.parametrize("table_kind", ["hash", "range", "dotted_hash", "dotted_range"]) +def test_set_creates_the_item_with_its_key(request, dynamodb_client, table_kind): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"set-{uuid.uuid4().hex}") + assert _create(dynamodb_client, table, key) == {} + assert _get(dynamodb_client, table, key) == {**key, "w": {"S": "upd"}} + _put_then_get(dynamodb_client, table, key) + + +@pytest.mark.parametrize("table_kind", ["dotted_hash", "dotted_range"]) +def test_query_and_scan_return_an_item_created_on_a_dotted_key_table( + request, dynamodb_client, table_kind +): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"qs-{uuid.uuid4().hex}") + _create(dynamodb_client, table, key) + created = {**key, "w": {"S": "upd"}} + hash_name = next(iter(key)) + resp = dynamodb_client.query( + TableName=table, + KeyConditionExpression="#h = :p", + ExpressionAttributeNames={"#h": hash_name}, + ExpressionAttributeValues={":p": key[hash_name]}, + ConsistentRead=True, + ) + assert resp["Items"] == [created] + assert created in dynamodb_client.scan(TableName=table, ConsistentRead=True)["Items"] + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_add_number_creates_the_item_with_its_key(request, dynamodb_client, table_kind): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"add-{uuid.uuid4().hex}") + dynamodb_client.update_item( + TableName=table, + Key=key, + UpdateExpression="ADD c :one", + ExpressionAttributeValues={":one": {"N": "1"}}, + ) + assert _get(dynamodb_client, table, key) == {**key, "c": {"N": "1"}} + _put_then_get(dynamodb_client, table, key) + + +def test_add_set_creates_the_item_with_its_key(dynamodb_client, hash_table): + _warm(dynamodb_client, hash_table, "hash") + key = _key("hash", f"add-set-{uuid.uuid4().hex}") + dynamodb_client.update_item( + TableName=hash_table, + Key=key, + UpdateExpression="ADD tags :t", + ExpressionAttributeValues={":t": {"SS": ["a", "b"]}}, + ) + item = _get(dynamodb_client, hash_table, key) + assert item is not None + assert item["pk"] == key["pk"] + assert sorted(item["tags"]["SS"]) == ["a", "b"] + _put_then_get(dynamodb_client, hash_table, key) + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_remove_only_on_a_missing_item_succeeds_and_the_key_stays_writable( + request, dynamodb_client, table_kind +): + """The service answers 200 and creates no item from a REMOVE-only update. + + ExtendDB creates an item holding the key (measured on SQLite and on + MongoDB); that difference is older than this test and is not asserted + either way here. + On MongoDB the call itself answered 500 before the fix, because the + created document had no item_data at all. + """ + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"remove-{uuid.uuid4().hex}") + resp = dynamodb_client.update_item(TableName=table, Key=key, UpdateExpression="REMOVE w") + assert "Attributes" not in resp + assert _get(dynamodb_client, table, key) in (None, key) + _put_then_get(dynamodb_client, table, key) + + +def test_attribute_not_exists_condition_creates_the_item_once(dynamodb_client, hash_table): + _warm(dynamodb_client, hash_table, "hash") + key = _key("hash", f"cond-{uuid.uuid4().hex}") + resp = _create(dynamodb_client, hash_table, key, ConditionExpression="attribute_not_exists(pk)") + assert resp == {} + assert _get(dynamodb_client, hash_table, key) == {**key, "w": {"S": "upd"}} + with pytest.raises(ClientError) as exc: + _create(dynamodb_client, hash_table, key, ConditionExpression="attribute_not_exists(pk)") + assert exc.value.response["Error"]["Code"] == "ConditionalCheckFailedException" + assert exc.value.response["Error"]["Message"] == "The conditional request failed" + _put_then_get(dynamodb_client, hash_table, key) + + +@pytest.mark.parametrize("return_values", RETURN_VALUES) +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_return_values_on_creation(request, dynamodb_client, table_kind, return_values): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"rv-{return_values.lower()}-{uuid.uuid4().hex}") + resp = _create(dynamodb_client, table, key, ReturnValues=return_values) + created = {**key, "w": {"S": "upd"}} + if return_values == "ALL_NEW": + assert resp == {"Attributes": created} + elif return_values == "UPDATED_NEW": + assert resp == {"Attributes": {"w": {"S": "upd"}}} + else: + assert resp == {} + assert _get(dynamodb_client, table, key) == created + _put_then_get(dynamodb_client, table, key) + + +@pytest.mark.parametrize("table_kind", ["hash", "range"]) +def test_later_operations_on_a_key_created_by_update_item(request, dynamodb_client, table_kind): + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + key = _key(table_kind, f"later-{uuid.uuid4().hex}") + created = {**key, "w": {"S": "upd"}} + _create(dynamodb_client, table, key) + assert _get(dynamodb_client, table, key) == created + + put = {**key, "w": {"S": "put"}} + resp = dynamodb_client.put_item(TableName=table, Item=put, ReturnValues="ALL_OLD") + assert resp["Attributes"] == created + assert _get(dynamodb_client, table, key) == put + + resp = dynamodb_client.update_item( + TableName=table, + Key=key, + UpdateExpression="SET x = :x", + ExpressionAttributeValues={":x": {"N": "1"}}, + ReturnValues="ALL_NEW", + ) + final = {**put, "x": {"N": "1"}} + assert resp["Attributes"] == final + + resp = dynamodb_client.query( + TableName=table, + KeyConditionExpression="pk = :p", + ExpressionAttributeValues={":p": key["pk"]}, + ConsistentRead=True, + ) + assert resp["Count"] == 1 + assert resp["Items"] == [final] + + scanned = dynamodb_client.scan(TableName=table, ConsistentRead=True)["Items"] + assert final in scanned + assert all("pk" in item for item in scanned) + + resp = dynamodb_client.delete_item(TableName=table, Key=key, ReturnValues="ALL_OLD") + assert resp["Attributes"] == final + assert _get(dynamodb_client, table, key) is None + + resp = dynamodb_client.put_item(TableName=table, Item=put, ReturnValues="ALL_OLD") + assert "Attributes" not in resp + assert _get(dynamodb_client, table, key) == put + + +@dataclass +class RaceResult: + table: str + table_kind: str + keys: list[dict] + failures: list[tuple[str, str, str]] = field(default_factory=list) + + def failed(self, op: str) -> list[tuple[str, str, str]]: + return [f for f in self.failures if f[0] == op] + + +@pytest.fixture(scope="module", params=["hash", "range"]) +def race(request, dynamodb_client, raw_client) -> RaceResult: + """WRITERS UpdateItem and WRITERS PutItem per key, KEYS keys, all at once. + + The issue measured 128 of 240 PutItem failing with HTTP 500 and 16 of 30 + final items without `pk` on the MongoDB backend. + """ + table_kind = request.param + table = request.getfixturevalue(f"{table_kind}_table") + _warm(dynamodb_client, table, table_kind) + result = RaceResult( + table, table_kind, [_key(table_kind, f"race-{uuid.uuid4().hex}") for _ in range(KEYS)] + ) + start = threading.Barrier(2 * WRITERS, timeout=60) + errors: list[BaseException] = [] + lock = threading.Lock() + + def run(op: str, i: int): + try: + start.wait() + for key in result.keys: + try: + if op == "put": + raw_client.put_item(TableName=table, Item={**key, "w": {"S": f"put{i}"}}) + else: + raw_client.update_item( + TableName=table, + Key=key, + UpdateExpression="SET w = :w", + ExpressionAttributeValues={":w": {"S": f"upd{i}"}}, + ) + except ClientError as e: + with lock: + result.failures.append((op, key["pk"]["S"], e.response["Error"]["Code"])) + except BaseException as e: # noqa: BLE001 - surfaced below + errors.append(e) + + threads = [ + threading.Thread(target=run, args=(op, i)) + for op in ("put", "update") + for i in range(WRITERS) + ] + for t in threads: + t.start() + for t in threads: + t.join() + if errors: + raise errors[0] + return result + + +def test_racing_update_item_calls_all_succeed(race): + failed = race.failed("update") + assert not failed, f"{len(failed)} of {WRITERS * KEYS} UpdateItem failed: {failed[:5]}" + + +def test_racing_writers_leave_every_item_with_its_key(dynamodb_client, race): + without_key = [] + for key in race.keys: + item = _get(dynamodb_client, race.table, key) + if item is None or any(item.get(name) != value for name, value in key.items()): + without_key.append(key["pk"]["S"]) + assert not without_key, f"{len(without_key)} of {KEYS} items lack a key attribute" + + +def test_put_item_after_the_race_replaces_every_item(dynamodb_client, race): + for key in race.keys: + _put_then_get(dynamodb_client, race.table, key) + + +@XFAIL_UNTIL_MONGODB_FIX +def test_racing_put_item_calls_all_succeed(race): + failed = race.failed("put") + assert not failed, f"{len(failed)} of {WRITERS * KEYS} PutItem failed: {failed[:5]}"