diff --git a/crates/storage-mongodb/src/data/mod.rs b/crates/storage-mongodb/src/data/mod.rs index 45465ba61..2036586ce 100644 --- a/crates/storage-mongodb/src/data/mod.rs +++ b/crates/storage-mongodb/src/data/mod.rs @@ -13,7 +13,9 @@ use extenddb_core::types::{ AttributeDefinition, AttributeValue, Item, KeySchemaElement, ScalarAttributeType, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{composite_pk_to_text, encode_netstring_composite, sk_info}; +use extenddb_storage::util::{ + composite_pk_to_text, encode_netstring_composite, escape_json_keys, sk_info, unescape_json_keys, +}; /// Returns the `MongoDB` collection name for a `DynamoDB` table. pub fn data_collection_name(table_id: &str) -> String { @@ -32,6 +34,42 @@ pub fn composite_id(pk_text: &str, sk_text: &str) -> String { encode_netstring_composite(&[pk_text.to_owned(), sk_text.to_owned()]) } +/// Convert an item (base item or projected index item) to the BSON stored +/// under `item_data`. +/// +/// BSON field names are C strings and cannot contain U+0000, while `DynamoDB` +/// accepts it in attribute names and map keys. Every object key in the item's +/// JSON form is run through the shared order-preserving escape (U+0000 and +/// U+0001 become two-character sequences, see +/// `extenddb_storage::util::escape_control`) before the BSON conversion, so +/// the stored document never carries a NUL field name. String values are left +/// alone; BSON strings are length-prefixed and hold U+0000 verbatim. +/// +/// Every conversion of an item to `item_data` BSON must go through this +/// function or [`json_to_item_data`]; a direct `bson::to_bson` on item JSON +/// fails in the driver at write time when a name contains U+0000. +/// [`item_data_to_json`] reverses the escape on read. +pub(crate) fn item_to_item_data(item: &Item) -> Result { + let item_json = + serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + json_to_item_data(item_json) +} + +/// JSON-tree variant of [`item_to_item_data`], for payloads that carry items +/// nested inside another structure (stream records with key/image maps). +pub(crate) fn json_to_item_data(json: serde_json::Value) -> Result { + let escaped = escape_json_keys(json); + bson::to_bson(&escaped).map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Reverse of [`json_to_item_data`]: BSON read from storage back to JSON with +/// the object-key escape undone. +pub(crate) fn item_data_to_json(value: &bson::Bson) -> Result { + let json: serde_json::Value = bson::from_bson(value.clone()) + .map_err(|e| StorageError::Internal(format!("BSON to JSON conversion error: {e}")))?; + Ok(unescape_json_keys(json)) +} + /// Convert a `DynamoDB` Item to a `MongoDB` BSON document for storage. /// /// Document structure: `{ _id, pk, sk_s/sk_n/sk_b, item_data }` @@ -42,10 +80,9 @@ pub fn item_to_document( ) -> Result { let pk_text = composite_pk_to_text(item, key_schema)?; - // Serialize the full item as item_data - let item_json = - serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; - let item_bson = bson::to_bson(&item_json).map_err(|e| StorageError::Internal(e.to_string()))?; + // Serialize the full item as item_data (object keys escaped, see + // item_to_item_data). + let item_bson = item_to_item_data(item)?; let mut doc = Document::new(); @@ -192,9 +229,7 @@ pub fn index_document( insert_typed_sk(&mut doc, &field, sk_type, sk_value)?; } - let item_json = - serde_json::to_value(projected).map_err(|e| StorageError::Internal(e.to_string()))?; - let item_bson = bson::to_bson(&item_json).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_bson = item_to_item_data(projected)?; doc.insert("item_data", item_bson); Ok(doc) @@ -254,8 +289,7 @@ pub fn document_to_item(doc: &Document) -> Result { .get("item_data") .ok_or_else(|| StorageError::Internal("Document missing item_data field".to_string()))?; - let json_value: serde_json::Value = bson::from_bson(item_data.clone()) - .map_err(|e| StorageError::Internal(format!("BSON to JSON conversion error: {e}")))?; + let json_value = item_data_to_json(item_data)?; let item: Item = serde_json::from_value(json_value) .map_err(|e| StorageError::Internal(format!("JSON to Item conversion error: {e}")))?; @@ -434,6 +468,89 @@ mod tests { assert!(matches!(err, StorageError::Validation(_))); } + /// Walk a BSON tree and assert no document field name contains U+0000. + fn assert_no_nul_field_names(value: &bson::Bson) { + match value { + bson::Bson::Document(doc) => { + for (k, v) in doc { + assert!(!k.contains('\u{0}'), "field name contains NUL: {k:?}"); + assert_no_nul_field_names(v); + } + } + bson::Bson::Array(items) => { + for item in items { + assert_no_nul_field_names(item); + } + } + _ => {} + } + } + + #[test] + fn item_document_round_trips_control_characters_in_names() { + use std::collections::BTreeMap; + + let (schema, attrs) = schema_pk_str_sk_num(); + + let mut inner = BTreeMap::new(); + inner.insert( + "deep\u{0}key".to_owned(), + AttributeValue::S("v\u{0}".to_owned()), + ); + let mut mid = BTreeMap::new(); + mid.insert("mid\u{1}key".to_owned(), AttributeValue::M(inner)); + let mut in_list = BTreeMap::new(); + in_list.insert( + "list\u{0}\u{1}key".to_owned(), + AttributeValue::S("w".to_owned()), + ); + + let mut item = Item::new(); + item.insert("pk".to_owned(), AttributeValue::S("p".to_owned())); + item.insert("sk".to_owned(), AttributeValue::N("1".to_owned())); + item.insert( + "top\u{0}\u{1}name".to_owned(), + AttributeValue::S("x".to_owned()), + ); + item.insert("map".to_owned(), AttributeValue::M(mid)); + item.insert( + "list".to_owned(), + AttributeValue::L(vec![AttributeValue::M(in_list)]), + ); + + let doc = item_to_document(&item, &schema, &attrs).unwrap(); + assert_no_nul_field_names(&bson::Bson::Document(doc.clone())); + + let back = document_to_item(&doc).unwrap(); + assert_eq!(back, item); + } + + #[test] + fn plain_item_document_stores_raw_field_names() { + use std::collections::BTreeMap; + + let (schema, attrs) = schema_pk_str_sk_num(); + + let mut inner = BTreeMap::new(); + inner.insert("k".to_owned(), AttributeValue::S("v".to_owned())); + let mut item = Item::new(); + item.insert("pk".to_owned(), AttributeValue::S("p".to_owned())); + item.insert("sk".to_owned(), AttributeValue::N("1".to_owned())); + item.insert("attr".to_owned(), AttributeValue::M(inner)); + + let doc = item_to_document(&item, &schema, &attrs).unwrap(); + let item_data = doc.get_document("item_data").unwrap(); + let names: Vec<&str> = item_data.keys().map(String::as_str).collect(); + assert_eq!(names, vec!["attr", "pk", "sk"]); + let attr_m = item_data + .get_document("attr") + .unwrap() + .get_document("M") + .unwrap(); + let map_keys: Vec<&str> = attr_m.keys().map(String::as_str).collect(); + assert_eq!(map_keys, vec!["k"]); + } + #[test] fn composite_id_disambiguates_delimiter_in_pk_or_sk() { // Two items whose naive "{pk}#{sk}" strings would collide must diff --git a/crates/storage-mongodb/src/data_engine.rs b/crates/storage-mongodb/src/data_engine.rs index 3895c4562..af40046fd 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -18,7 +18,7 @@ use extenddb_core::types::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::{ - composite_pk_to_text, encode_netstring_composite, pk_to_text, sk_info, + composite_pk_to_text, encode_netstring_composite, needs_escape, pk_to_text, sk_info, }; use extenddb_storage::{ DataEngine, IdempotencyKey, ItemPairResult, QueryResult, StreamCapture, TransactGetOp, @@ -29,7 +29,7 @@ use crate::MongoEngine; use crate::condition::condition_to_filter; use crate::data::{ binary_sk_to_hex, composite_id, data_collection_name, document_to_item, index_document, - index_entry_filter, item_to_document, pk_filter, sk_field_name, sk_suffix, + index_entry_filter, item_to_document, json_to_item_data, pk_filter, sk_field_name, sk_suffix, }; use crate::pushdown::{Pushable, is_pushable}; @@ -37,8 +37,9 @@ use extenddb_core::types::{Projection, ProjectionType}; /// Resolve a single-component update path for the native MongoDB fast path. /// Literal attribute names containing dots or beginning with `$` cannot be -/// represented safely by MongoDB's ordinary update paths, so they must use -/// the Rust/session path. +/// represented safely by MongoDB's ordinary update paths, and names +/// containing U+0000 or U+0001 differ from the stored (escaped) field name, +/// so they must use the Rust/session path. fn native_attribute_name(path: &[PathElement], maps: &ExpressionMaps) -> Option { if path.len() != 1 { return None; @@ -47,7 +48,8 @@ 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()) + (!attr_name.contains('.') && !attr_name.starts_with('$') && !needs_escape(&attr_name)) + .then(|| attr_name.into_owned()) } #[derive(Clone, Copy)] @@ -1651,8 +1653,11 @@ impl MongoEngine { _ => return None, // complex expressions (if_not_exists, list_append, arithmetic) }; let field = format!("item_data.{attr_name}"); + // The value lands inside item_data, so its nested map + // keys must be in stored (escaped) form, matching what + // item_to_document writes. let val_json = serde_json::to_value(val).ok()?; - let val_bson = bson::to_bson(&val_json).ok()?; + let val_bson = json_to_item_data(val_json).ok()?; set_doc.insert(field, val_bson); } UpdateAction::Remove { path } => { @@ -2033,8 +2038,10 @@ impl MongoEngine { let record_json = serde_json::to_value(&record).map_err(|e| StorageError::Internal(e.to_string()))?; - let record_bson = - bson::to_bson(&record_json).map_err(|e| StorageError::Internal(e.to_string()))?; + // The record's key and image maps carry attribute names as object + // keys, which become BSON field names; convert through the same + // escaping helper as item_data. + let record_bson = json_to_item_data(record_json)?; // key_info already carries table_id — no need to re-read the catalog // just to resolve it, and re-reading inside the session against the @@ -4385,6 +4392,21 @@ transaction_read_concern = "majority""#, assert_eq!(native_attribute_name(&path, &maps), None); } + #[test] + fn native_update_rejects_control_character_literal_attribute_names() { + // Stored field names carry the escaped form of U+0000 and U+0001, + // so a native update path built from the raw name would target a + // different field; these names must use the Rust/session path. + for name in ["a\u{0}b", "a\u{1}b"] { + let mut names = std::collections::HashMap::new(); + names.insert("name".to_owned(), name.to_owned()); + let maps = ExpressionMaps::new(names, std::collections::HashMap::new()); + let path = vec![PathElement::Attribute("#name".to_owned())]; + + assert_eq!(native_attribute_name(&path, &maps), None, "{name:?}"); + } + } + #[test] fn native_update_keeps_simple_attribute_names_on_fast_path() { let maps = ExpressionMaps::new( diff --git a/crates/storage-mongodb/src/metadata_engine.rs b/crates/storage-mongodb/src/metadata_engine.rs index 7fe21f9c3..3dac4241b 100644 --- a/crates/storage-mongodb/src/metadata_engine.rs +++ b/crates/storage-mongodb/src/metadata_engine.rs @@ -18,6 +18,7 @@ use extenddb_storage::error::StorageError; use crate::MongoEngine; use crate::data::{data_collection_name, document_to_item}; +use extenddb_storage::util::escape_control; fn extract_id_fields(doc: &Document) -> (String, String) { let id = doc.get_document("_id").ok(); @@ -35,13 +36,15 @@ fn extract_id_fields(doc: &Document) -> (String, String) { /// Build the expression that reads a DynamoDB attribute from `item_data` as a /// literal field name. MongoDB normally interprets dots in a query field path /// as nested-document traversal, but DynamoDB permits dots in attribute names. +/// The stored field name is the escaped form of the attribute name (see +/// `extenddb_storage::util::escape_control`), so the lookup targets that form. fn literal_ttl_value_expression(ttl_attribute: &str) -> Document { doc! { "$getField": { "field": "N", "input": { "$getField": { - "field": ttl_attribute, + "field": escape_control(ttl_attribute).as_ref(), "input": "$item_data", } } @@ -69,7 +72,7 @@ fn literal_ttl_epoch_expression(ttl_attribute: &str) -> Document { /// also performs the expiry comparison in MongoDB, so non-expired low-`_id` /// documents cannot starve expired documents later in the collection. fn ttl_candidate_query(ttl_attribute: &str, now_epoch: i64) -> (Document, Document) { - let ttl_field = format!("item_data.{ttl_attribute}.N"); + let ttl_field = format!("item_data.{}.N", escape_control(ttl_attribute)); if ttl_attribute.contains('.') { let ttl_epoch = literal_ttl_epoch_expression(ttl_attribute); ( @@ -374,8 +377,12 @@ impl MetadataEngine for MongoEngine { let coll_name = data_collection_name(table_id); let data_coll = self.data_db.collection::(&coll_name); - let index_name = format!("idx_ttl_{ttl_attribute}"); - let index_key = format!("item_data.{ttl_attribute}.N"); + // The stored field name is the escaped attribute name, and the + // escape also keeps U+0000 out of the index name. Names without + // U+0000 or U+0001 escape to themselves, so existing index names + // do not change. + let index_name = format!("idx_ttl_{}", escape_control(&ttl_attribute)); + let index_key = format!("item_data.{}.N", escape_control(&ttl_attribute)); let index = IndexModel::builder() .keys(doc! { &index_key: 1 }) @@ -433,7 +440,9 @@ impl MetadataEngine for MongoEngine { .get_str("table_id") .map_err(|_| StorageError::Internal("missing table_id".to_string()))?; - let index_name = format!("idx_ttl_{ttl_attribute}"); + // Same escaped form as create_ttl_index so the drop targets the + // index that was created. + let index_name = format!("idx_ttl_{}", escape_control(&ttl_attribute)); let coll_name = data_collection_name(table_id); let data_coll = self.data_db.collection::(&coll_name); @@ -513,9 +522,11 @@ impl MetadataEngine for MongoEngine { if items.len() >= limit { break; } - // Parse the TTL value and check if expired + // Parse the TTL value and check if expired. The stored field + // name is the escaped attribute name. if let Ok(item_data) = doc.get_document("item_data") - && let Ok(ttl_obj) = item_data.get_document(&ttl_attribute) + && let Ok(ttl_obj) = + item_data.get_document(escape_control(&ttl_attribute).as_ref()) && let Ok(n_str) = ttl_obj.get_str("N") && let Ok(ttl_val) = n_str.parse::() && ttl_val >= 1 diff --git a/crates/storage-mongodb/src/pushdown.rs b/crates/storage-mongodb/src/pushdown.rs index 616ef7508..05f96dac9 100644 --- a/crates/storage-mongodb/src/pushdown.rs +++ b/crates/storage-mongodb/src/pushdown.rs @@ -44,6 +44,7 @@ use extenddb_core::expression::{CompareOp, Expr, ExpressionMaps, PathElement}; use extenddb_core::types::AttributeValue; +use extenddb_storage::util::needs_escape; /// Outcome of the pushdown analyzer. #[derive(Debug, Clone, PartialEq, Eq)] @@ -76,7 +77,10 @@ pub fn is_pushable(expr: &Expr, maps: &ExpressionMaps) -> Pushable { /// Return why an expression contains a literal DynamoDB attribute name that /// cannot safely be emitted as a MongoDB field path. A dot would mean nested /// document traversal in MongoDB, while a leading `$` would be interpreted as -/// an operator or otherwise have special meaning in a MongoDB path. +/// an operator or otherwise have special meaning in a MongoDB path. A name +/// containing U+0000 or U+0001 is stored under its escaped form (see +/// `extenddb_storage::util::escape_control`), so the raw name would not match +/// the stored field. fn unsafe_attribute_reason(expr: &Expr, maps: &ExpressionMaps) -> Option<&'static str> { match expr { Expr::Path(elements) => elements.iter().find_map(|element| match element { @@ -91,6 +95,9 @@ fn unsafe_attribute_reason(expr: &Expr, maps: &ExpressionMaps) -> Option<&'stati Some(name) if name.starts_with('$') => { Some("literal attribute name begins with '$'") } + Some(name) if needs_escape(name) => { + Some("literal attribute name contains U+0000 or U+0001") + } Some(_) => None, None => Some("attribute name alias is unresolved"), } @@ -393,6 +400,41 @@ mod tests { ); } + #[test] + fn nul_attribute_name_is_not_pushable() { + // The stored field name is the escaped form, so a raw path on a + // name containing U+0000 would never match; must fall back. + let expr = Expr::Compare { + left: Box::new(path("a\u{0}b")), + op: CompareOp::Eq, + right: Box::new(Expr::Placeholder(":value".into())), + }; + let maps = maps_with(&[(":value", AttributeValue::S("value".into()))]); + + assert_eq!( + is_pushable(&expr, &maps), + Pushable::No("literal attribute name contains U+0000 or U+0001") + ); + } + + #[test] + fn control_character_attribute_alias_is_not_pushable() { + for resolved in ["a\u{0}b", "a\u{1}b"] { + let expr = Expr::Compare { + left: Box::new(path("#name")), + op: CompareOp::Eq, + right: Box::new(Expr::Placeholder(":value".into())), + }; + let maps = aliased_maps(resolved, AttributeValue::S("value".into())); + + assert_eq!( + is_pushable(&expr, &maps), + Pushable::No("literal attribute name contains U+0000 or U+0001"), + "{resolved:?}" + ); + } + } + #[test] fn nested_document_path_without_literal_dot_remains_pushable() { let expr = Expr::Compare { diff --git a/crates/storage-mongodb/src/stream_engine.rs b/crates/storage-mongodb/src/stream_engine.rs index e976a667a..b57180c18 100644 --- a/crates/storage-mongodb/src/stream_engine.rs +++ b/crates/storage-mongodb/src/stream_engine.rs @@ -368,8 +368,9 @@ impl StreamEngine for MongoEngine { let record_bson = d .get("record_data") .ok_or_else(|| StorageError::Internal("Missing record_data".to_owned()))?; - let json_val: serde_json::Value = bson::from_bson(record_bson.clone()) - .map_err(|e| StorageError::Internal(e.to_string()))?; + // Reverses the object-key escape applied when the record + // was written; see json_to_item_data. + let json_val = crate::data::item_data_to_json(record_bson)?; serde_json::from_value(json_val) .map_err(|e| StorageError::Internal(e.to_string())) }) diff --git a/crates/storage-postgres/src/bootstrapper.rs b/crates/storage-postgres/src/bootstrapper.rs index a8827df35..e25ae3f02 100755 --- a/crates/storage-postgres/src/bootstrapper.rs +++ b/crates/storage-postgres/src/bootstrapper.rs @@ -110,6 +110,9 @@ impl PostgresBootstrapper { .username(&self.config.app_user) .password(&self.config.app_password) .database(database) + // Lets a migration tell its own connections apart from a server's + // in pg_stat_activity (see `migrations::MIGRATE_APPLICATION_NAME`). + .application_name(migrations::MIGRATE_APPLICATION_NAME) } /// Build the connection URL for the application user and a named database. diff --git a/crates/storage-postgres/src/data/delete_item.rs b/crates/storage-postgres/src/data/delete_item.rs index 1ea9dd1e8..e350194ac 100755 --- a/crates/storage-postgres/src/data/delete_item.rs +++ b/crates/storage-postgres/src/data/delete_item.rs @@ -7,7 +7,9 @@ use extenddb_core::expression::{Expr, ExpressionMaps}; use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{SortKeyValue, parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{SortKeyValue, sk_column, sk_info}; + +use super::key_text::{parse_sk, pk_to_text}; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; diff --git a/crates/storage-postgres/src/data/index.rs b/crates/storage-postgres/src/data/index.rs index 551861dbf..07152b2e3 100644 --- a/crates/storage-postgres/src/data/index.rs +++ b/crates/storage-postgres/src/data/index.rs @@ -13,9 +13,11 @@ use extenddb_core::types::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::SortKeyValue; -use extenddb_storage::util::{composite_pk_to_text, parse_sk, sk_column, sk_column_n}; +use extenddb_storage::util::{sk_column, sk_column_n}; -use super::{all_sort_key_info, index_table_name}; +use super::key_text::{composite_pk_to_text, parse_sk}; + +use super::{all_sort_key_info, index_table_name, item_to_json}; use crate::gsi_queue::{GsiApplyContext, GsiIndexDef, enqueue_gsi_pending}; /// Map a `sqlx` error to `StorageError`, preserving the SQLSTATE code in the @@ -351,8 +353,7 @@ pub(crate) async fn insert_index_row_multi( let idx_pk_text = composite_pk_to_text(item, index_ks)?; let base_pk_text = composite_pk_to_text(item, base_ks)?; - let item_json = - serde_json::to_value(projected).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(projected)?; // Build column list dynamically let mut cols = vec!["pk".to_owned()]; diff --git a/crates/storage-postgres/src/data/key_text.rs b/crates/storage-postgres/src/data/key_text.rs new file mode 100644 index 000000000..a433907af --- /dev/null +++ b/crates/storage-postgres/src/data/key_text.rs @@ -0,0 +1,139 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Key text for the `PostgreSQL` key columns. +//! +//! The shared helpers in `extenddb_storage::util` turn key attribute values into +//! the text the SQL backends bind against `pk`, `sk_s`, `base_pk`, and the other +//! string key columns. `PostgreSQL` `TEXT` cannot hold the byte 0x00, which +//! DynamoDB allows anywhere in a string key, so this module wraps those helpers +//! and passes every string component through the order-preserving escape from +//! `extenddb_storage::util::escape_control`. Numbers and base64 binaries contain +//! neither U+0000 nor U+0001 and are unchanged. The escape preserves byte order +//! and prefixes, so the `COLLATE "C"` comparisons, `BETWEEN`, `begins_with`, +//! and row-comparison cursors that run on these columns return the same rows +//! they would on the raw text. +//! +//! Every site in this crate that produces key column text imports these three +//! functions instead of the shared ones, so the write path, the read path, index +//! rows, cursors, and the propagation queue all agree on the stored form. + +use std::borrow::Cow; + +use extenddb_core::types::{AttributeValue, Item, KeySchemaElement, KeyType, ScalarAttributeType}; +use extenddb_storage::error::StorageError; +use extenddb_storage::util::{self, SortKeyValue, encode_netstring_composite, escape_control}; + +/// Partition key attribute value as escaped column text. +pub(crate) fn pk_to_text(value: &AttributeValue) -> Result, StorageError> { + Ok(match util::pk_to_text(value)? { + Cow::Borrowed(s) => escape_control(s), + Cow::Owned(s) => Cow::Owned(escape_control(&s).into_owned()), + }) +} + +/// Composite partition key as escaped column text. +/// +/// Mirrors `extenddb_storage::util::composite_pk_to_text`: a single HASH +/// attribute is its own text, several are netstring-joined. Each part is +/// escaped before joining, which is also what `query_scan` does when it builds +/// the same text from a key condition, so the two agree byte for byte. +pub(crate) fn composite_pk_to_text( + item: &Item, + key_schema: &[KeySchemaElement], +) -> Result { + let hash_elements: Vec<_> = key_schema + .iter() + .filter(|ks| ks.key_type == KeyType::Hash) + .collect(); + if hash_elements.len() == 1 { + let val = item + .get(&hash_elements[0].attribute_name) + .ok_or_else(|| StorageError::Internal("missing partition key".to_owned()))?; + return Ok(pk_to_text(val)?.into_owned()); + } + let mut parts = Vec::with_capacity(hash_elements.len()); + for ks in &hash_elements { + let val = item.get(&ks.attribute_name).ok_or_else(|| { + StorageError::Internal(format!( + "missing partition key attribute {}", + ks.attribute_name + )) + })?; + parts.push(pk_to_text(val)?.into_owned()); + } + Ok(encode_netstring_composite(&parts)) +} + +/// Sort key attribute value as a typed bind value, string keys escaped. +pub(crate) fn parse_sk( + value: &AttributeValue, + sk_type: ScalarAttributeType, +) -> Result { + Ok(match util::parse_sk(value, sk_type)? { + SortKeyValue::S(s) => SortKeyValue::S(escape_control(&s).into_owned()), + other => other, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn ks(name: &str) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type: KeyType::Hash, + } + } + + #[test] + fn string_keys_are_escaped_and_other_types_are_not() { + assert_eq!( + pk_to_text(&AttributeValue::S("a\u{0}b".into())).unwrap(), + "a\u{1}\u{1}b" + ); + assert_eq!( + pk_to_text(&AttributeValue::S("plain".into())).unwrap(), + "plain" + ); + assert_eq!( + pk_to_text(&AttributeValue::N("12.5".into())).unwrap(), + "12.5" + ); + assert_eq!( + pk_to_text(&AttributeValue::B(vec![0, 1, 2])).unwrap(), + util::pk_to_text(&AttributeValue::B(vec![0, 1, 2])).unwrap() + ); + match parse_sk(&AttributeValue::S("\u{1}".into()), ScalarAttributeType::S).unwrap() { + SortKeyValue::S(s) => assert_eq!(s, "\u{1}\u{2}"), + _ => panic!("string sort key"), + } + match parse_sk(&AttributeValue::B(vec![0]), ScalarAttributeType::B).unwrap() { + SortKeyValue::B(b) => assert_eq!(b, vec![0]), + _ => panic!("binary sort key"), + } + } + + #[test] + fn composite_text_escapes_each_part_before_joining() { + let mut item = Item::new(); + item.insert("h1".to_owned(), AttributeValue::S("a\u{0}".into())); + item.insert("h2".to_owned(), AttributeValue::S("b".into())); + let text = composite_pk_to_text(&item, &[ks("h1"), ks("h2")]).unwrap(); + // The escaped part is three bytes, and the netstring length says so. + assert_eq!(text, "3:a\u{1}\u{1},1:b,"); + assert!(!text.contains('\u{0}')); + // Single hash attribute: the escaped value itself. + let text = composite_pk_to_text(&item, &[ks("h1")]).unwrap(); + assert_eq!(text, "a\u{1}\u{1}"); + } + + #[test] + fn stored_text_never_contains_nul() { + for raw in ["\u{0}", "a\u{0}b", "\u{0}\u{1}\u{2}", "x\u{1}\u{1}y"] { + let av = AttributeValue::S(raw.into()); + let text = pk_to_text(&av).unwrap(); + assert!(!text.contains('\u{0}'), "{raw:?} -> {text:?}"); + } + } +} diff --git a/crates/storage-postgres/src/data/mod.rs b/crates/storage-postgres/src/data/mod.rs index d64ceb539..b5ffbc0cc 100755 --- a/crates/storage-postgres/src/data/mod.rs +++ b/crates/storage-postgres/src/data/mod.rs @@ -9,6 +9,7 @@ use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, ScalarAttributeType}; use extenddb_storage::error::StorageError; +use extenddb_storage::util::{escape_json_strings, unescape_json_strings}; /// SQL table name for a Virtual `DynamoDB` table. /// @@ -60,8 +61,44 @@ pub(crate) fn all_sort_key_info<'a>( } /// Deserialize an `item_data` JSONB value into an `Item`. +/// +/// Reverses [`item_to_json`]: every string in the stored tree is unescaped +/// before deserialization, so an item written with U+0000 anywhere in it comes +/// back byte-identical. pub(crate) fn json_to_item(v: serde_json::Value) -> Result { - serde_json::from_value(v).map_err(|e| StorageError::Internal(e.to_string())) + serde_json::from_value(unescape_json_strings(v)) + .map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Serialize an `Item` for an `item_data` JSONB column. +/// +/// PostgreSQL `jsonb` rejects the `\u0000` escape, so every string in the tree +/// (attribute names, map keys, string values) goes through the order-preserving +/// escape from `extenddb_storage::util` before it reaches the column. Items +/// without U+0000 or U+0001 serialize to exactly what they did before. +pub(crate) fn item_to_json(item: &Item) -> Result { + serde_json::to_value(item) + .map(escape_json_strings) + .map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Serialize any value for a JSONB column that may carry item strings (stream +/// records, queued index updates, index contexts). Same escape as +/// [`item_to_json`]; the inverse is [`stored_json_to`]. +pub(crate) fn to_stored_json( + value: &T, +) -> Result { + serde_json::to_value(value) + .map(escape_json_strings) + .map_err(|e| StorageError::Internal(e.to_string())) +} + +/// Deserialize a JSONB value written by [`to_stored_json`]. +pub(crate) fn stored_json_to( + v: serde_json::Value, +) -> Result { + serde_json::from_value(unescape_json_strings(v)) + .map_err(|e| StorageError::Internal(e.to_string())) } /// Bind a `SortKeyValue` to a positional parameter in a sqlx query and execute it. @@ -133,6 +170,7 @@ mod data_engine; mod ddl; mod delete_item; pub(crate) mod index; +pub(crate) mod key_text; mod put_item; mod query; mod query_scan; diff --git a/crates/storage-postgres/src/data/put_item.rs b/crates/storage-postgres/src/data/put_item.rs index e307df660..6ce8ec2f2 100755 --- a/crates/storage-postgres/src/data/put_item.rs +++ b/crates/storage-postgres/src/data/put_item.rs @@ -7,12 +7,14 @@ use extenddb_core::expression::{Expr, ExpressionMaps}; use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{composite_pk_to_text, parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{sk_column, sk_info}; + +use super::key_text::{composite_pk_to_text, parse_sk, pk_to_text}; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; -use super::{data_table_name, json_to_item}; +use super::{data_table_name, item_to_json, json_to_item}; use crate::PostgresEngine; impl PostgresEngine { @@ -30,8 +32,7 @@ impl PostgresEngine { let pk_text = composite_pk_to_text(&item, &key_info.key_schema)?; - let item_json = - serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(&item)?; // Both index families in one catalog visit (D-4: sync + async split for the // secondary indexes). diff --git a/crates/storage-postgres/src/data/query.rs b/crates/storage-postgres/src/data/query.rs index 6a84dc1b7..ca3c8ddcc 100755 --- a/crates/storage-postgres/src/data/query.rs +++ b/crates/storage-postgres/src/data/query.rs @@ -12,7 +12,9 @@ use extenddb_core::types::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::SortKeyValue; -use extenddb_storage::util::{parse_sk, pk_to_text, sk_info}; +use extenddb_storage::util::sk_info; + +use super::key_text::{parse_sk, pk_to_text}; /// Extra pagination bind values for index queries. /// diff --git a/crates/storage-postgres/src/data/query_scan.rs b/crates/storage-postgres/src/data/query_scan.rs index c0f0c8b05..ef4525f9c 100644 --- a/crates/storage-postgres/src/data/query_scan.rs +++ b/crates/storage-postgres/src/data/query_scan.rs @@ -6,9 +6,9 @@ use extenddb_core::expression::{ExpressionMaps, KeyCondition}; use extenddb_core::types::{Item, ScalarAttributeType, TableKeyInfo}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{ - encode_netstring_composite, parse_sk, pk_to_text, sk_column, sk_column_n, sk_info, -}; +use extenddb_storage::util::{encode_netstring_composite, sk_column, sk_column_n, sk_info}; + +use super::key_text::{parse_sk, pk_to_text}; use super::query::{ PaginationBinds, build_key, build_sk_sql, execute_query_sql, execute_scan_sql, diff --git a/crates/storage-postgres/src/data/tx_helpers.rs b/crates/storage-postgres/src/data/tx_helpers.rs index ed30a1bf3..aa1454555 100755 --- a/crates/storage-postgres/src/data/tx_helpers.rs +++ b/crates/storage-postgres/src/data/tx_helpers.rs @@ -12,9 +12,11 @@ use extenddb_core::types::{ }; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{SortKeyValue, parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{SortKeyValue, sk_column, sk_info}; -use super::{data_table_name, json_to_item}; +use super::key_text::{parse_sk, pk_to_text}; + +use super::{data_table_name, item_to_json, json_to_item, to_stored_json}; /// Fetch a single item within an existing transaction. pub(super) async fn fetch_item_in_tx( @@ -105,8 +107,7 @@ pub(super) async fn upsert_item_in_tx( .get(pk_name) .ok_or_else(|| StorageError::Internal("missing partition key".to_owned()))?; let pk_text = pk_to_text(pk_value)?; - let item_json = - serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(item)?; if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) { @@ -157,8 +158,7 @@ pub(super) async fn insert_item_if_absent_in_tx( .get(pk_name) .ok_or_else(|| StorageError::Internal("missing partition key".to_owned()))?; let pk_text = pk_to_text(pk_value)?; - let item_json = - serde_json::to_value(item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(item)?; let rows_affected = if let Some((sk_name, sk_type)) = sk_info(&key_info.key_schema, &key_info.attribute_definitions) @@ -366,8 +366,7 @@ pub(super) async fn write_stream_record_in_tx( user_identity: capture.user_identity.clone(), }; - let record_json = - serde_json::to_value(&record).map_err(|e| StorageError::Internal(e.to_string()))?; + let record_json = to_stored_json(&record)?; sqlx::query( "INSERT INTO stream_records (sequence_number, shard_id, table_id, event_name, record_data) \ diff --git a/crates/storage-postgres/src/data/update_item.rs b/crates/storage-postgres/src/data/update_item.rs index 72ce6edb9..7a731309d 100755 --- a/crates/storage-postgres/src/data/update_item.rs +++ b/crates/storage-postgres/src/data/update_item.rs @@ -8,9 +8,12 @@ use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_core::validation; use extenddb_storage::StreamCapture; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{parse_sk, pk_to_text, sk_column, sk_info}; +use extenddb_storage::util::{sk_column, sk_info}; + +use super::key_text::{parse_sk, pk_to_text}; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; +use super::item_to_json; use super::query::check_condition; use super::tx_helpers::write_stream_record_in_tx; use super::{data_table_name, json_to_item}; @@ -192,8 +195,7 @@ impl PostgresEngine { let new_item = if return_new { Some(item.clone()) } else { None }; - let item_json = - serde_json::to_value(&item).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(&item)?; if old_json.is_some() { // Row existed and is locked by the read above, so update in place. diff --git a/crates/storage-postgres/src/data/vector_index.rs b/crates/storage-postgres/src/data/vector_index.rs index 0bc0a18b5..500225ed8 100644 --- a/crates/storage-postgres/src/data/vector_index.rs +++ b/crates/storage-postgres/src/data/vector_index.rs @@ -32,15 +32,15 @@ use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, ScalarAttributeType}; use extenddb_core::validation::vector_item::{vector_components, vector_norm}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{ - SortKeyValue, composite_pk_to_text, parse_sk, sk_column, sk_column_n, -}; +use extenddb_storage::util::{SortKeyValue, sk_column, sk_column_n}; + +use super::key_text::{composite_pk_to_text, parse_sk}; use extenddb_storage::vector_lifecycle::{ VectorApplyContext, VectorIndexMeta, item_is_indexable, item_partition, projected_payload, }; use pgvector::Vector; -use super::{all_sort_key_info, vector_table_name}; +use super::{all_sort_key_info, item_to_json, vector_table_name}; /// Read the vector index metadata for a table from the catalog. /// @@ -403,8 +403,7 @@ impl<'a> VectorInsertPlan<'a> { // vector attribute are the shared payload rules, so a live-written row and a // backfilled one cannot differ in shape. let projected = projected_payload(item, base_key_schema, meta); - let item_json = - serde_json::to_value(&projected).map_err(|e| StorageError::Internal(e.to_string()))?; + let item_json = item_to_json(&projected)?; // A plain INSERT, deliberately, where the GSI sibling upserts. Two callers reach // it and only one deletes first: the live write path through diff --git a/crates/storage-postgres/src/gsi_queue.rs b/crates/storage-postgres/src/gsi_queue.rs index 39e60d443..ab567dbb2 100644 --- a/crates/storage-postgres/src/gsi_queue.rs +++ b/crates/storage-postgres/src/gsi_queue.rs @@ -23,9 +23,9 @@ use std::sync::Arc; +use crate::data::key_text::composite_pk_to_text; use extenddb_core::types::{AttributeDefinition, Item, KeySchemaElement, Projection}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::composite_pk_to_text; use serde::{Deserialize, Serialize}; use sqlx::PgPool; use tokio::sync::Notify; @@ -223,16 +223,9 @@ pub(crate) async fn enqueue_gsi_pending( delay_ms: u64, context: &PendingApplyContext, ) -> Result<(), StorageError> { - let old_json = old_item - .map(serde_json::to_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let new_json = new_item - .map(serde_json::to_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let context_json = - serde_json::to_value(context).map_err(|e| StorageError::Internal(e.to_string()))?; + let old_json = old_item.map(crate::data::item_to_json).transpose()?; + let new_json = new_item.map(crate::data::item_to_json).transpose()?; + let context_json = crate::data::to_stored_json(context)?; // Route all updates for a given base item to one worker (per-key FIFO). // The base key is immutable over an item's lifetime; `new_item` carries it @@ -426,16 +419,9 @@ async fn apply_claimed_row( new_json: Option, ctx_json: serde_json::Value, ) -> Result<(), StorageError> { - let old_item: Option = old_json - .map(serde_json::from_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let new_item: Option = new_json - .map(serde_json::from_value) - .transpose() - .map_err(|e| StorageError::Internal(e.to_string()))?; - let context: PendingApplyContext = - serde_json::from_value(ctx_json).map_err(|e| StorageError::Internal(e.to_string()))?; + let old_item: Option = old_json.map(crate::data::json_to_item).transpose()?; + let new_item: Option = new_json.map(crate::data::json_to_item).transpose()?; + let context: PendingApplyContext = crate::data::stored_json_to(ctx_json)?; // One index per row. Guard the apply with a savepoint so a dropped-index // race (42P01) can be recovered and the row still consumed; aborting the diff --git a/crates/storage-postgres/src/lib.rs b/crates/storage-postgres/src/lib.rs index 0d7a51c0d..3dfeb03b3 100755 --- a/crates/storage-postgres/src/lib.rs +++ b/crates/storage-postgres/src/lib.rs @@ -406,6 +406,38 @@ impl PostgresEngine { Ok(()) } + /// Refuse to serve a data database whose stored strings have not been + /// brought to this build's encoding. + /// + /// The escape in `extenddb_storage::util::escape_control` stores U+0001 as a + /// two-character sequence, and the migration that re-encodes rows written by + /// older builds (`004_escape_control_chars`) rewrites every row that contains + /// U+0001. It cannot tell a row this build wrote from a legacy one, so this + /// build must not write anything before the migration has run. `extenddb + /// init` records the migration on a fresh deployment; an upgrade runs + /// `extenddb migrate --yes` after stopping every older server and before + /// starting this one. + /// + /// # Errors + /// + /// Returns `StorageError::Internal` naming the migration and the command to + /// run when the data database's `schema_history` does not record it. + pub async fn check_data_migrations_applied(&self) -> Result<(), StorageError> { + let missing = migrations::unapplied_required_data_migrations(&self.data_pool) + .await + .map_err(|e| StorageError::Connection(format!("{e:?}")))?; + if missing.is_empty() { + return Ok(()); + } + Err(StorageError::Internal(format!( + "Data migration not applied: {}. This build stores strings in a different \ + encoding from earlier builds, so it must not serve until the existing rows \ + have been re-encoded. Stop every older ExtendDB server that uses this data \ + database, run 'extenddb migrate --yes', then start this server.", + missing.join(", ") + ))) + } + /// Query the data database name from the catalog for the startup banner (REQ-LOG-001). /// /// Returns `"(not configured)"` if no data database has been registered. @@ -656,6 +688,18 @@ fn server_components_factory( _ => BackendError::InitializationFailed(e.to_string()), })?; + // This build stores strings through the control-character escape, and the + // re-encoding migration cannot tell a row this build wrote from a row an + // older build wrote (both can hold U+0001). Serving before the migration + // has run would let the migration later rewrite this build's own rows, so + // refuse to start until it is recorded. `extenddb init` records it on a + // fresh deployment; an upgraded deployment runs `extenddb migrate --yes` + // with every older server stopped. + engine + .check_data_migrations_applied() + .await + .map_err(|e| BackendError::InitializationFailed(e.to_string()))?; + // Rebuild any vector index a crash left CREATING, before serving: an index // in that state is not searchable, and nothing else will repair it. match crate::data::vector_index::reconcile_incomplete_vector_indexes(&engine).await { diff --git a/crates/storage-postgres/src/metadata_engine.rs b/crates/storage-postgres/src/metadata_engine.rs index 322a9015f..bd553b3e9 100755 --- a/crates/storage-postgres/src/metadata_engine.rs +++ b/crates/storage-postgres/src/metadata_engine.rs @@ -273,7 +273,9 @@ impl MetadataEngine for PostgresEngine { ) -> BoxFuture<'_, Result<(), StorageError>> { let account_id = account_id.to_string(); let table_name = table_name.to_string(); - let ttl_attribute = ttl_attribute.to_string(); + // Stored attribute names carry the same escape as the rest of item_data, + // so the name SQL addresses must be the stored form. + let ttl_attribute = extenddb_storage::util::escape_control(ttl_attribute).into_owned(); Box::pin(async move { Self::validate_account_id(&account_id)?; let (table_id,): (String,) = sqlx::query_as( @@ -366,7 +368,9 @@ impl MetadataEngine for PostgresEngine { ) -> BoxFuture<'_, Result, StorageError>> { let account_id = account_id.to_string(); let table_name = table_name.to_string(); - let ttl_attribute = ttl_attribute.to_string(); + // Stored attribute names carry the same escape as the rest of item_data, + // so the name SQL addresses must be the stored form. + let ttl_attribute = extenddb_storage::util::escape_control(ttl_attribute).into_owned(); Box::pin(async move { Self::validate_account_id(&account_id)?; let (table_id,): (String,) = sqlx::query_as( diff --git a/crates/storage-postgres/src/migrations.rs b/crates/storage-postgres/src/migrations.rs index 45af16167..eaa5ffb36 100755 --- a/crates/storage-postgres/src/migrations.rs +++ b/crates/storage-postgres/src/migrations.rs @@ -129,7 +129,91 @@ pub(crate) async fn run_data_migrations(pool: &PgPool) -> OpResult<()> { /// outside a transaction (they use `CREATE INDEX CONCURRENTLY`), so they cannot /// be expressed as SQL in [`DATA_MIGRATIONS`]. Applied by `extenddb migrate` /// after the SQL migrations, so the operator controls when the change happens. -pub(crate) const DATA_CODE_MIGRATIONS: &[&str] = &["003_gsi_base_key_index"]; +pub(crate) const DATA_CODE_MIGRATIONS: &[&str] = + &["003_gsi_base_key_index", "004_escape_control_chars"]; + +/// Code migrations a server must see recorded before it serves. `003` only adds +/// an index and a server runs correctly without it; `004` changes the stored +/// form of strings, and a server that wrote through the new escape before the +/// migration ran would have its own rows rewritten by it (see +/// [`escape_legacy_control_chars`]). `PostgresEngine::check_data_migrations_applied` +/// refuses to start on a data database missing any of these. +pub(crate) const REQUIRED_DATA_CODE_MIGRATIONS: &[&str] = &["004_escape_control_chars"]; + +/// `application_name` the init and migrate commands set on their connections, +/// so [`refuse_if_other_clients_connected`] can leave them out of its count. +pub(crate) const MIGRATE_APPLICATION_NAME: &str = "extenddb-migrate"; + +/// Environment variable that skips [`refuse_if_other_clients_connected`], for a +/// data database whose only other connections are known to be idle (a +/// connection pooler holding server-side connections open, for example). +pub(crate) const IGNORE_CONNECTIONS_ENV: &str = "EXTENDDB_MIGRATE_IGNORE_CONNECTIONS"; + +/// Refuse to run a data-rewriting migration while anything else is connected +/// to the data database. +/// +/// `004_escape_control_chars` rewrites rows in place and cannot tell a row a +/// running server writes from a legacy row: an older server still writing +/// leaves unescaped rows behind, and a server of this release writing before +/// the migration finishes gets its rows double-escaped. The upgrade manual +/// requires every server to be stopped first; this check catches the case +/// where one was not. Connections carrying [`MIGRATE_APPLICATION_NAME`] are +/// this process's own and are ignored; every other client backend on the +/// database counts. Managed environments where a pooler keeps idle +/// connections open can set [`IGNORE_CONNECTIONS_ENV`] to skip the check. +pub(crate) async fn refuse_if_other_clients_connected(pool: &PgPool) -> OpResult<()> { + let others: Vec<(String, Option, Option)> = sqlx::query_as( + "SELECT usename::text, application_name, client_addr::text \ + FROM pg_stat_activity \ + WHERE datname = current_database() \ + AND backend_type = 'client backend' \ + AND pid <> pg_backend_pid() \ + AND coalesce(application_name, '') <> $1", + ) + .bind(MIGRATE_APPLICATION_NAME) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("Check for other connections: {e}")))?; + if others.is_empty() { + return Ok(()); + } + let mut described: Vec = others + .iter() + .map(|(user, app, addr)| { + format!( + "{user} from {} ({})", + addr.as_deref().unwrap_or("local socket"), + match app.as_deref() { + Some("") | None => "no application name", + Some(a) => a, + } + ) + }) + .collect(); + described.sort(); + described.dedup(); + Err(OpError::Internal(format!( + "{} other connection(s) hold the data database: {}. This migration rewrites \ + stored rows and must run with every ExtendDB server stopped and every other \ + session closed. Stop them and run 'extenddb migrate --yes' again. If the only \ + other connections are idle ones held by a connection pooler, set \ + {IGNORE_CONNECTIONS_ENV}=1 to skip this check.", + others.len(), + described.join("; ") + ))) +} + +/// Names in [`REQUIRED_DATA_CODE_MIGRATIONS`] not recorded in this data +/// database's `schema_history`. +pub(crate) async fn unapplied_required_data_migrations(pool: &PgPool) -> OpResult> { + let mut missing = Vec::new(); + for name in REQUIRED_DATA_CODE_MIGRATIONS { + if !is_migration_applied(pool, name).await? { + missing.push((*name).to_owned()); + } + } + Ok(missing) +} /// Run programmatic data migrations, skipping already-applied ones. /// @@ -151,6 +235,15 @@ pub(crate) async fn run_data_code_migrations( "003_gsi_base_key_index" => { ensure_gsi_base_key_indexes(catalog_pool, data_pool).await?; } + "004_escape_control_chars" => { + // Rows written by a server that is still running would be missed + // (an older release) or double-escaped (this release); see + // `refuse_if_other_clients_connected`. + if std::env::var_os(IGNORE_CONNECTIONS_ENV).is_none() { + refuse_if_other_clients_connected(data_pool).await?; + } + escape_legacy_control_chars(catalog_pool, data_pool).await?; + } other => { return Err(OpError::Internal(format!( "Unknown data code migration: {other}" @@ -219,6 +312,355 @@ async fn ensure_gsi_base_key_indexes(catalog_pool: &PgPool, data_pool: &PgPool) Ok(()) } +/// Progress ledger for multi-table data code migrations: one row per +/// (migration, table) pair, committed in the same transaction as that table's +/// rewrite, so a re-run after a crash knows exactly which tables are done. +const CODE_MIGRATION_PROGRESS_TABLE: &str = "data_code_migration_progress"; + +/// Create the progress ledger if this database does not have one yet. +async fn ensure_progress_table(pool: &PgPool) -> OpResult<()> { + sqlx::query(&format!( + "CREATE TABLE IF NOT EXISTS {CODE_MIGRATION_PROGRESS_TABLE} (\ + migration_name TEXT NOT NULL, \ + table_name TEXT NOT NULL, \ + completed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), \ + PRIMARY KEY (migration_name, table_name)\ + )" + )) + .execute(pool) + .await + .map_err(|e| OpError::Internal(format!("Create migration progress table: {e}")))?; + Ok(()) +} + +/// Whether the progress ledger says `table` is already done for `migration`. +async fn is_table_marked(pool: &PgPool, migration: &str, table: &str) -> OpResult { + let marked: bool = sqlx::query_scalar(&format!( + "SELECT EXISTS(SELECT 1 FROM {CODE_MIGRATION_PROGRESS_TABLE} \ + WHERE migration_name = $1 AND table_name = $2)" + )) + .bind(migration) + .bind(table) + .fetch_one(pool) + .await + .map_err(|e| OpError::Internal(format!("Check migration progress: {e}")))?; + Ok(marked) +} + +/// One physical table's rewrite work for `004_escape_control_chars`: the text +/// key columns that get the single-statement SQL replace, and the jsonb +/// columns whose rows are re-encoded one by one. +struct EscapeWork { + /// Bare (unquoted) table name, e.g. `_ddb_` or `gsi_pending`. + table: String, + text_columns: Vec, + jsonb_columns: Vec, +} + +/// Rewrite work for a dynamically named `_ddb_*` table, read from +/// `information_schema.columns` so that exactly the sort key columns this +/// table has (`sk_s`, `sk_s`, `base_sk_s`, ...) are covered. Returns +/// `None` when the physical table does not exist (a catalog row mid-create or +/// mid-delete): such a table has no legacy rows to re-encode. +async fn dynamic_table_work(pool: &PgPool, table: String) -> OpResult> { + let text_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'text' \ + AND (column_name IN ('pk', 'base_pk') OR column_name::text ~ '^(base_)?sk[0-9]*_s$') \ + ORDER BY ordinal_position", + ) + .bind(&table) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List text key columns of {table}: {e}")))?; + + let jsonb_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'jsonb' \ + AND column_name = 'item_data'", + ) + .bind(&table) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List jsonb columns of {table}: {e}")))?; + + if text_columns.is_empty() && jsonb_columns.is_empty() { + return Ok(None); + } + Ok(Some(EscapeWork { + table, + text_columns, + jsonb_columns, + })) +} + +/// Rewrite work for a fixed-name table, with the candidate columns filtered +/// through `information_schema.columns` so a column absent on this deployment +/// is skipped. Returns `None` when the table itself is absent. +async fn fixed_table_work( + pool: &PgPool, + table: &str, + text_candidates: &[&str], + jsonb_candidates: &[&str], +) -> OpResult> { + if !table_exists(pool, table).await? { + return Ok(None); + } + let text_candidates: Vec = text_candidates.iter().map(|c| (*c).to_owned()).collect(); + let jsonb_candidates: Vec = jsonb_candidates.iter().map(|c| (*c).to_owned()).collect(); + let text_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'text' \ + AND column_name::text = ANY($2) \ + ORDER BY ordinal_position", + ) + .bind(table) + .bind(&text_candidates) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List text columns of {table}: {e}")))?; + let jsonb_columns: Vec = sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = 'jsonb' \ + AND column_name::text = ANY($2) \ + ORDER BY ordinal_position", + ) + .bind(table) + .bind(&jsonb_candidates) + .fetch_all(pool) + .await + .map_err(|e| OpError::Internal(format!("List jsonb columns of {table}: {e}")))?; + Ok(Some(EscapeWork { + table: table.to_owned(), + text_columns, + jsonb_columns, + })) +} + +/// The re-encoded form of a stored jsonb value, or `None` when re-encoding +/// does not change it. `escape_json_strings` only changes strings and object +/// keys containing U+0000 or U+0001, so a plain document and a document whose +/// strings merely contain the literal six-character text backslash-u-0001 both +/// come back equal and are skipped. +fn reencoded_json(value: &serde_json::Value) -> Option { + let escaped = extenddb_storage::util::escape_json_strings(value.clone()); + (escaped != *value).then_some(escaped) +} + +/// Re-encode the rows of one jsonb column whose stored document contains +/// U+0001. Runs inside the caller's per-table transaction. +/// +/// The candidate ctids are snapshotted once, before any row of this column is +/// touched, so no row is ever visited twice; the updates relocate rows (a ctid +/// names a physical row version), but always after that row's value was read. +/// A parameterized `DECLARE CURSOR` cannot be issued through the extended +/// protocol, so the locked-snapshot form of a streamed fetch is used instead, +/// and the values are then fetched in batches of 500 by ctid so a large table +/// never loads into memory. The matching set is tiny in practice; the table +/// scan, not the writes, is the cost. +async fn escape_jsonb_rows( + tx: &mut sqlx::Transaction<'_, sqlx::Postgres>, + quoted_table: &str, + col: &str, +) -> OpResult { + // '\u0001' (six characters) is how jsonb::text renders a raw U+0001 inside + // a string or key. It also matches a stored string whose content is a + // literal backslash followed by u0001 (rendered '\\u0001'), which is + // harmless: that row re-encodes to itself and is skipped below. jsonb text + // rendering is why SQL-level replace() is not used here: it would corrupt + // exactly that literal-backslash case. The backslash is spelled chr(92) so + // the predicate reads the same under either value of + // standard_conforming_strings; a plain '\u0001' literal would become the + // single byte 0x01 with the setting off and match nothing. + let select_ctids = format!( + "SELECT ctid::text FROM {quoted_table} \ + WHERE position((chr(92) || 'u0001') IN \"{col}\"::text) > 0 FOR UPDATE" + ); + let ctids: Vec = sqlx::query_scalar(&select_ctids) + .fetch_all(&mut **tx) + .await + .map_err(|e| OpError::Internal(format!("Find rows to re-encode in {quoted_table}: {e}")))?; + + let fetch_batch = + format!("SELECT ctid::text, \"{col}\" FROM {quoted_table} WHERE ctid = ANY($1::tid[])"); + let update_row = format!("UPDATE {quoted_table} SET \"{col}\" = $1 WHERE ctid = $2::tid"); + let mut changed: u64 = 0; + for batch in ctids.chunks(500) { + let rows: Vec<(String, serde_json::Value)> = sqlx::query_as(&fetch_batch) + .bind(batch.to_vec()) + .fetch_all(&mut **tx) + .await + .map_err(|e| OpError::Internal(format!("Fetch rows from {quoted_table}: {e}")))?; + for (ctid, value) in rows { + if let Some(escaped) = reencoded_json(&value) { + sqlx::query(&update_row) + .bind(escaped) + .bind(&ctid) + .execute(&mut **tx) + .await + .map_err(|e| { + OpError::Internal(format!("Re-encode a row of {quoted_table}: {e}")) + })?; + changed += 1; + } + } + } + Ok(changed) +} + +/// Rewrite one table in one transaction and mark it done in the progress +/// ledger before committing, so the rewrite and the marker are atomic. +async fn escape_rows_in_table(pool: &PgPool, work: &EscapeWork) -> OpResult { + let mut tx = pool + .begin() + .await + .map_err(|e| OpError::Internal(format!("Begin transaction for {}: {e}", work.table)))?; + let quoted = format!("\"{}\"", work.table); + let mut rewritten: u64 = 0; + + // Key text columns: TEXT has no escaping layer, so a single SQL statement + // per column is exact. NULL columns never match (position() is NULL). + for col in &work.text_columns { + let sql = format!( + "UPDATE {quoted} SET \"{col}\" = replace(\"{col}\", chr(1), chr(1) || chr(2)) \ + WHERE position(chr(1) IN \"{col}\") > 0" + ); + let done = sqlx::query(&sql) + .execute(&mut *tx) + .await + .map_err(|e| OpError::Internal(format!("Re-encode {}.{col}: {e}", work.table)))?; + rewritten += done.rows_affected(); + } + + for col in &work.jsonb_columns { + rewritten += escape_jsonb_rows(&mut tx, "ed, col).await?; + } + + sqlx::query(&format!( + "INSERT INTO {CODE_MIGRATION_PROGRESS_TABLE} (migration_name, table_name) \ + VALUES ($1, $2) ON CONFLICT DO NOTHING" + )) + .bind("004_escape_control_chars") + .bind(&work.table) + .execute(&mut *tx) + .await + .map_err(|e| OpError::Internal(format!("Mark {} as re-encoded: {e}", work.table)))?; + + tx.commit() + .await + .map_err(|e| OpError::Internal(format!("Commit re-encode of {}: {e}", work.table)))?; + Ok(rewritten) +} + +/// Re-encode legacy rows that contain a raw U+0001 (`004_escape_control_chars`). +/// +/// Strings are stored through the order-preserving escape in +/// `extenddb_storage::util::escape_control` (U+0000 becomes U+0001 U+0001, +/// U+0001 becomes U+0001 U+0002). Rows written before the escape existed hold +/// their strings raw. The decoder keeps a stray U+0001 literal, so the only +/// legacy rows it can misread are those where a raw U+0001 is immediately +/// followed by U+0001 or U+0002; rewriting every row containing U+0001 removes +/// that case. Legacy rows never contain U+0000 (PostgreSQL refused it), so for +/// legacy data applying the escape to the raw stored value is exactly the +/// re-encode, and rows without U+0001 encode to themselves and are left alone. +/// +/// Crash safety. This migration is NOT idempotent at the row level: escaped +/// data also contains U+0001 (the first character of every escape pair), and a +/// legacy U+0001 U+0002 is indistinguishable from an escaped U+0001, so +/// applying the escape to an already-escaped row would double-escape it. A +/// partially rewritten table must therefore never be rescanned. Each table is +/// processed in one transaction: its matching rows are rewritten, a marker row +/// (migration_name, table_name) is inserted into the small progress table +/// created at the start of this migration, and the transaction commits. A +/// re-run after a crash skips tables that carry a marker and resumes with the +/// rest; when every table is marked, the caller records the migration in +/// `schema_history` exactly as `003` does, and the usual +/// `is_migration_applied` guard makes later runs no-ops. Marker rows are kept +/// forever: deleting them before the `schema_history` row is durable would +/// reopen the rescan window. +/// +/// Rows written by THIS build are indistinguishable from legacy rows for the +/// same reason (an escaped U+0001 is stored as U+0001 U+0002), so the server +/// refuses to start on a data database that does not record this migration +/// (`PostgresEngine::check_data_migrations_applied`). The migration therefore +/// only ever sees rows written by older builds, provided every older server +/// is stopped before it runs; `extenddb init` records it on a fresh deployment. `backup_items` lives in the catalog database +/// while every other table lives in the data database, and a transaction +/// cannot span two databases, so each database gets its own progress table and +/// each marker commits with the table it marks. +async fn escape_legacy_control_chars(catalog_pool: &PgPool, data_pool: &PgPool) -> OpResult<()> { + const NAME: &str = "004_escape_control_chars"; + ensure_progress_table(data_pool).await?; + ensure_progress_table(catalog_pool).await?; + + let mut work: Vec<(&PgPool, EscapeWork)> = Vec::new(); + + // Every data table, index table, and vector index table, enumerated from + // the catalog. Their text key columns are read from information_schema so + // single-sort-key, multi-sort-key, and base_* layouts are all covered. + let table_ids: Vec = + sqlx::query_scalar("SELECT table_id FROM tables ORDER BY table_id") + .fetch_all(catalog_pool) + .await + .map_err(|e| OpError::Internal(format!("Enumerate tables: {e}")))?; + for id in table_ids { + if let Some(w) = dynamic_table_work(data_pool, format!("_ddb_{id}")).await? { + work.push((data_pool, w)); + } + } + let index_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes ORDER BY index_id") + .fetch_all(catalog_pool) + .await + .map_err(|e| OpError::Internal(format!("Enumerate indexes: {e}")))?; + for id in index_ids { + if let Some(w) = dynamic_table_work(data_pool, format!("_ddb_{id}")).await? { + work.push((data_pool, w)); + } + } + let vector_ids: Vec = + sqlx::query_scalar("SELECT index_id FROM vector_indexes ORDER BY index_id") + .fetch_all(catalog_pool) + .await + .map_err(|e| OpError::Internal(format!("Enumerate vector indexes: {e}")))?; + for id in vector_ids { + if let Some(w) = dynamic_table_work(data_pool, format!("_ddb_vec_{id}")).await? { + work.push((data_pool, w)); + } + } + + // Fixed tables holding item documents. gsi_pending and stream_records live + // in the data database; backup_items lives in the catalog database. + if let Some(w) = fixed_table_work( + data_pool, + "gsi_pending", + &[], + &["old_item", "new_item", "index_context"], + ) + .await? + { + work.push((data_pool, w)); + } + if let Some(w) = fixed_table_work(data_pool, "stream_records", &[], &["record_data"]).await? { + work.push((data_pool, w)); + } + if let Some(w) = + fixed_table_work(catalog_pool, "backup_items", &["pk", "sk"], &["item_data"]).await? + { + work.push((catalog_pool, w)); + } + + for (pool, w) in &work { + if is_table_marked(pool, NAME, &w.table).await? { + println!(" {}: already re-encoded, skipping.", w.table); + continue; + } + let rewritten = escape_rows_in_table(pool, w).await?; + println!(" {}: {rewritten} row update(s)", w.table); + } + Ok(()) +} + /// Check if a table exists in the public schema. pub(crate) async fn table_exists(pool: &PgPool, name: &str) -> OpResult { let exists: bool = sqlx::query_scalar( @@ -299,6 +741,30 @@ mod tests { use super::CATALOG_MIGRATIONS; use crate::CATALOG_VERSION; + /// `reencoded_json` is what decides whether a stored jsonb row is + /// rewritten by `004_escape_control_chars`, so it must change exactly the + /// rows holding a raw control character and nothing else. + #[test] + fn reencoded_json_changes_exactly_the_rows_holding_raw_control_chars() { + use serde_json::json; + + // Raw U+0001 in a string value and in an object key: rewritten, with + // each U+0001 becoming U+0001 U+0002 wherever it appears. + let raw = json!({"k\u{1}": {"S": "v\u{1}"}, "plain": {"N": "1"}}); + let expected = json!({"k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}, "plain": {"N": "1"}}); + assert_eq!(super::reencoded_json(&raw), Some(expected)); + + // A string whose content is the literal six-character text + // backslash-u-0001: the SQL candidate filter over-matches it, and this + // check is what keeps it byte-identical. + let literal = json!({"note": {"S": "x\\u0001y"}}); + assert_eq!(super::reencoded_json(&literal), None); + + // A plain document: untouched. + let plain = json!({"pk": {"S": "p"}, "m": {"M": {"k": {"S": "v"}}}}); + assert_eq!(super::reencoded_json(&plain), None); + } + /// The catalog version and the migration list must move together. /// /// A migration that creates its schema without moving the version leaves a @@ -361,3 +827,572 @@ mod tests { } } } + +#[cfg(test)] +mod live_reencode { + //! Live checks for `004_escape_control_chars` against a scratch database. + //! + //! Follows the convention of `tests/key_collation.rs`: each test needs + //! `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no database name, + //! for example `postgresql://postgres:postgres@127.0.0.1:5432`), builds a + //! throwaway database, and drops it when it passes. Without the variable + //! every test here reports a skip and passes. One database serves as both + //! catalog and data, as the storage-level tests do. + + use serde_json::json; + use sqlx::PgPool; + use sqlx::postgres::PgPoolOptions; + + use super::{ + CATALOG_MIGRATIONS, DATA_MIGRATIONS, IGNORE_CONNECTIONS_ENV, MIGRATE_APPLICATION_NAME, + pending_data_migrations, refuse_if_other_clients_connected, run_data_code_migrations, + unapplied_required_data_migrations, + }; + + struct Scratch { + db: PgPool, + admin: PgPool, + db_name: String, + } + + impl Scratch { + async fn cleanup(self) { + let Scratch { db, admin, db_name } = self; + db.close().await; + sqlx::query(&format!( + "DROP DATABASE IF EXISTS \"{db_name}\" WITH (FORCE)" + )) + .execute(&admin) + .await + .expect("drop the scratch database"); + admin.close().await; + } + } + + fn base_conn() -> Option { + let conn = std::env::var("EXTENDDB_TEST_PG_CONNECTION_STRING").ok()?; + (!conn.trim().is_empty()).then(|| conn.trim_end_matches('/').to_owned()) + } + + fn skip(test: &str) { + eprintln!( + "SKIP {test}: EXTENDDB_TEST_PG_CONNECTION_STRING is not set, so there is no \ + PostgreSQL to build a scratch database in." + ); + } + + async fn scratch() -> Scratch { + let base = base_conn().expect("caller checks base_conn() first"); + let db_name = format!("eddb_esc_{}", uuid::Uuid::new_v4().simple())[..24].to_owned(); + let admin = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/postgres")) + .await + .expect("connect to the postgres maintenance database"); + sqlx::query(&format!("CREATE DATABASE \"{db_name}\"")) + .execute(&admin) + .await + .expect("create the scratch database"); + // Named like the migrate command, so the connection guard on 004 treats + // this pool as the migration's own. + let db = PgPoolOptions::new() + .max_connections(2) + .connect_with( + format!("{base}/{db_name}") + .parse::() + .expect("scratch connection string") + .application_name(MIGRATE_APPLICATION_NAME), + ) + .await + .expect("connect to the scratch database"); + for (_, sql) in CATALOG_MIGRATIONS.iter().chain(DATA_MIGRATIONS) { + sqlx::raw_sql(sql) + .execute(&db) + .await + .expect("apply a shipped migration"); + } + // 003 creates its indexes from parsed catalog key schemas and is not + // what these tests exercise, so record it as already applied. + sqlx::query("INSERT INTO schema_history (filename) VALUES ('003_gsi_base_key_index')") + .execute(&db) + .await + .expect("record 003 as applied"); + Scratch { db, admin, db_name } + } + + /// Seed a data table, an index table, `gsi_pending`, `stream_records`, and + /// `backup_items` with three row kinds: raw U+0001 in key columns and + /// inside jsonb strings and object keys, the literal six-character text + /// backslash-u-0001 inside a jsonb string, and plain rows. + async fn seed(db: &PgPool) { + let ks = json!([ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ]); + let ad = json!([ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"} + ]); + sqlx::query("INSERT INTO accounts (account_id, account_name) VALUES ($1, $2)") + .bind("111122223333") + .bind("esc-check") + .execute(db) + .await + .expect("seed the account row"); + sqlx::query( + "INSERT INTO tables (account_id, table_name, key_schema, attribute_definitions, \ + table_status, table_arn, table_id) \ + VALUES ($1, 't', $2, $3, 'ACTIVE', \ + 'arn:aws:dynamodb:us-east-1:111122223333:table/t', 't1')", + ) + .bind("111122223333") + .bind(&ks) + .bind(&ad) + .execute(db) + .await + .expect("seed the tables row"); + sqlx::query( + "INSERT INTO indexes (table_id, index_id, index_name, index_type, key_schema, \ + projection) VALUES ('t1', 'i1', 'gsi1', 'GSI', $1, $2)", + ) + .bind(json!([{"AttributeName": "gpk", "KeyType": "HASH"}])) + .bind(json!({"ProjectionType": "ALL"})) + .execute(db) + .await + .expect("seed the indexes row"); + + sqlx::raw_sql( + r#"CREATE TABLE "_ddb_t1" ( + pk TEXT COLLATE "C" NOT NULL, + sk_s TEXT COLLATE "C", + sk_n NUMERIC, + sk_b BYTEA, + item_data JSONB NOT NULL, + PRIMARY KEY (pk, sk_s) + ); + CREATE TABLE "_ddb_i1" ( + pk TEXT COLLATE "C" NOT NULL, + sk_s TEXT COLLATE "C", + sk_n NUMERIC, + sk_b BYTEA, + base_pk TEXT COLLATE "C" NOT NULL, + base_sk_s TEXT COLLATE "C", + base_sk_n NUMERIC, + base_sk_b BYTEA, + item_data JSONB NOT NULL, + PRIMARY KEY (pk, base_pk, base_sk_s) + );"#, + ) + .execute(db) + .await + .expect("create the physical data and index tables"); + + for (pk, sk, item) in [ + ( + "a\u{1}b", + "s\u{1}", + json!({"pk": {"S": "a\u{1}b"}, "k\u{1}": {"S": "v\u{1}"}}), + ), + ("lit", "s", json!({"note": {"S": "x\\u0001y"}})), + ("plain", "s", json!({"pk": {"S": "plain"}})), + ] { + sqlx::query("INSERT INTO \"_ddb_t1\" (pk, sk_s, item_data) VALUES ($1, $2, $3)") + .bind(pk) + .bind(sk) + .bind(item) + .execute(db) + .await + .expect("seed a data table row"); + } + for (pk, sk, base_pk, base_sk, item) in [ + ( + "g\u{1}", + "gs\u{1}", + "a\u{1}b", + "s\u{1}", + json!({"gpk": {"S": "g\u{1}"}}), + ), + ("zz", "zs", "plain", "s", json!({"gpk": {"S": "zz"}})), + ] { + sqlx::query( + "INSERT INTO \"_ddb_i1\" (pk, sk_s, base_pk, base_sk_s, item_data) \ + VALUES ($1, $2, $3, $4, $5)", + ) + .bind(pk) + .bind(sk) + .bind(base_pk) + .bind(base_sk) + .bind(item) + .execute(db) + .await + .expect("seed an index table row"); + } + + for (new_item, context) in [ + (json!({"pk": {"S": "a\u{1}b"}}), json!({"k\u{1}": 1})), + (json!({"pk": {"S": "p"}}), json!({"k": 1})), + ] { + sqlx::query( + "INSERT INTO gsi_pending (table_id, worker_partition, old_item, new_item, \ + index_context) VALUES ('t1', 0, NULL, $1, $2)", + ) + .bind(new_item) + .bind(context) + .execute(db) + .await + .expect("seed a gsi_pending row"); + } + + sqlx::query( + "INSERT INTO stream_shards (shard_id, table_id, starting_sequence_number) \ + VALUES ('shard-1', 't1', '1')", + ) + .execute(db) + .await + .expect("seed the stream shard"); + for (seq, record) in [ + ("1", json!({"Keys": {"pk": {"S": "a\u{1}b"}}})), + ("2", json!({"Keys": {"pk": {"S": "p"}}})), + ] { + sqlx::query( + "INSERT INTO stream_records (shard_id, sequence_number, table_id, event_name, \ + record_data) VALUES ('shard-1', $1, 't1', 'INSERT', $2)", + ) + .bind(seq) + .bind(record) + .execute(db) + .await + .expect("seed a stream record"); + } + + sqlx::query( + "INSERT INTO backups (backup_arn, backup_name, table_id, table_name, account_id, \ + key_schema, attribute_definitions) \ + VALUES ('arn:aws:dynamodb:us-east-1:111122223333:table/t/backup/b1', 'b1', 't1', \ + 't', '111122223333', $1, $2)", + ) + .bind(&ks) + .bind(&ad) + .execute(db) + .await + .expect("seed the backups row"); + for (pk, sk, item) in [ + ("p\u{1}", Some("s\u{1}"), json!({"pk": {"S": "p\u{1}"}})), + ("plain", None, json!({"pk": {"S": "plain"}})), + ] { + sqlx::query( + "INSERT INTO backup_items (backup_arn, pk, sk, item_data) \ + VALUES ('arn:aws:dynamodb:us-east-1:111122223333:table/t/backup/b1', $1, $2, $3)", + ) + .bind(pk) + .bind(sk) + .bind(item) + .execute(db) + .await + .expect("seed a backup item"); + } + } + + #[tokio::test] + async fn reencodes_legacy_rows_once_and_records_itself() { + if base_conn().is_none() { + skip("reencodes_legacy_rows_once_and_records_itself"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let pending = pending_data_migrations(&s.db).await.expect("pending list"); + assert!( + pending.iter().any(|p| p == "004_escape_control_chars"), + "004 must be pending before the run: {pending:?}" + ); + + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run the data code migrations"); + + // Data table: raw U+0001 became U+0001 U+0002 in the key columns and + // inside every jsonb string and object key; the literal-backslash row + // and the plain row are byte-identical. + let rows: Vec<(String, String, serde_json::Value)> = + sqlx::query_as("SELECT pk, sk_s, item_data FROM \"_ddb_t1\" ORDER BY pk") + .fetch_all(&s.db) + .await + .expect("read the data table"); + assert_eq!( + rows, + vec![ + ( + "a\u{1}\u{2}b".to_owned(), + "s\u{1}\u{2}".to_owned(), + json!({"pk": {"S": "a\u{1}\u{2}b"}, "k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}}), + ), + ( + "lit".to_owned(), + "s".to_owned(), + json!({"note": {"S": "x\\u0001y"}}), + ), + ( + "plain".to_owned(), + "s".to_owned(), + json!({"pk": {"S": "plain"}}), + ), + ] + ); + + // Index table: pk, sk_s, base_pk, base_sk_s and item_data all re-encoded. + let rows: Vec<(String, String, String, String, serde_json::Value)> = sqlx::query_as( + "SELECT pk, sk_s, base_pk, base_sk_s, item_data FROM \"_ddb_i1\" ORDER BY pk", + ) + .fetch_all(&s.db) + .await + .expect("read the index table"); + assert_eq!( + rows, + vec![ + ( + "g\u{1}\u{2}".to_owned(), + "gs\u{1}\u{2}".to_owned(), + "a\u{1}\u{2}b".to_owned(), + "s\u{1}\u{2}".to_owned(), + json!({"gpk": {"S": "g\u{1}\u{2}"}}), + ), + ( + "zz".to_owned(), + "zs".to_owned(), + "plain".to_owned(), + "s".to_owned(), + json!({"gpk": {"S": "zz"}}), + ), + ] + ); + + // gsi_pending: NULL old_item stays NULL, the raw row is re-encoded in + // both jsonb columns, the plain row is untouched. + let rows: Vec<( + Option, + serde_json::Value, + serde_json::Value, + )> = + sqlx::query_as("SELECT old_item, new_item, index_context FROM gsi_pending ORDER BY id") + .fetch_all(&s.db) + .await + .expect("read gsi_pending"); + assert_eq!( + rows, + vec![ + ( + None, + json!({"pk": {"S": "a\u{1}\u{2}b"}}), + json!({"k\u{1}\u{2}": 1}), + ), + (None, json!({"pk": {"S": "p"}}), json!({"k": 1})), + ] + ); + + // stream_records. + let rows: Vec<(serde_json::Value,)> = + sqlx::query_as("SELECT record_data FROM stream_records ORDER BY sequence_number") + .fetch_all(&s.db) + .await + .expect("read stream_records"); + assert_eq!( + rows, + vec![ + (json!({"Keys": {"pk": {"S": "a\u{1}\u{2}b"}}}),), + (json!({"Keys": {"pk": {"S": "p"}}}),), + ] + ); + + // backup_items, including the NULL sk row. + let rows: Vec<(String, Option, serde_json::Value)> = + sqlx::query_as("SELECT pk, sk, item_data FROM backup_items ORDER BY pk") + .fetch_all(&s.db) + .await + .expect("read backup_items"); + assert_eq!( + rows, + vec![ + ( + "p\u{1}\u{2}".to_owned(), + Some("s\u{1}\u{2}".to_owned()), + json!({"pk": {"S": "p\u{1}\u{2}"}}), + ), + ("plain".to_owned(), None, json!({"pk": {"S": "plain"}}),), + ] + ); + + // Recorded in schema_history, one progress marker per table, and no + // longer pending. + let recorded: bool = sqlx::query_scalar( + "SELECT EXISTS(SELECT 1 FROM schema_history \ + WHERE filename = '004_escape_control_chars')", + ) + .fetch_one(&s.db) + .await + .expect("read schema_history"); + assert!(recorded, "schema_history must record the migration"); + let markers: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM data_code_migration_progress \ + WHERE migration_name = '004_escape_control_chars'", + ) + .fetch_one(&s.db) + .await + .expect("count progress markers"); + assert_eq!(markers, 5, "one marker per processed table"); + let pending = pending_data_migrations(&s.db).await.expect("pending list"); + assert!( + !pending.iter().any(|p| p == "004_escape_control_chars"), + "004 must not be pending after the run: {pending:?}" + ); + + // A second run is skipped by the is_migration_applied guard and does + // not double-escape anything. + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run the data code migrations again"); + let (pk, item): (String, serde_json::Value) = + sqlx::query_as("SELECT pk, item_data FROM \"_ddb_t1\" WHERE sk_s = $1") + .bind("s\u{1}\u{2}") + .fetch_one(&s.db) + .await + .expect("re-read the re-encoded row"); + assert_eq!(pk, "a\u{1}\u{2}b"); + assert_eq!( + item, + json!({"pk": {"S": "a\u{1}\u{2}b"}, "k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}}) + ); + + s.cleanup().await; + } + + /// The jsonb candidate predicate must not depend on the session's + /// standard_conforming_strings: with it off, a plain '\u0001' literal is a + /// single 0x01 byte that never appears in jsonb text and matches nothing, + /// so every jsonb row would be skipped while the migration recorded itself. + #[tokio::test] + async fn reencodes_jsonb_rows_with_standard_conforming_strings_off() { + if base_conn().is_none() { + skip("reencodes_jsonb_rows_with_standard_conforming_strings_off"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let base = base_conn().unwrap(); + let scs_off = PgPoolOptions::new() + .max_connections(2) + .after_connect(|conn, _| { + Box::pin(async move { + sqlx::Executor::execute(&mut *conn, "SET standard_conforming_strings = off") + .await?; + Ok(()) + }) + }) + .connect_with( + format!("{base}/{}", s.db_name) + .parse::() + .unwrap() + .application_name(MIGRATE_APPLICATION_NAME), + ) + .await + .expect("connect with standard_conforming_strings off"); + + run_data_code_migrations(&scs_off, &scs_off) + .await + .expect("run 004"); + + let (pk, sk, item): (String, String, serde_json::Value) = + sqlx::query_as("SELECT pk, sk_s, item_data FROM \"_ddb_t1\" WHERE pk LIKE 'a%'") + .fetch_one(&s.db) + .await + .unwrap(); + assert_eq!(pk, "a\u{1}\u{2}b"); + assert_eq!(sk, "s\u{1}\u{2}"); + assert_eq!( + item, + json!({"pk": {"S": "a\u{1}\u{2}b"}, "k\u{1}\u{2}": {"S": "v\u{1}\u{2}"}}), + "jsonb rows must be re-encoded whatever the session setting" + ); + scs_off.close().await; + s.cleanup().await; + } + + /// The server must not serve a data database that has not been through + /// 004: rows it wrote through the escape would be rewritten by the + /// migration later. The required list is empty only once 004 is recorded. + #[tokio::test] + async fn required_migrations_are_reported_until_004_is_recorded() { + if base_conn().is_none() { + skip("required_migrations_are_reported_until_004_is_recorded"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let missing = unapplied_required_data_migrations(&s.db).await.unwrap(); + assert_eq!(missing, vec!["004_escape_control_chars".to_owned()]); + + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run 004"); + + let missing = unapplied_required_data_migrations(&s.db).await.unwrap(); + assert!(missing.is_empty(), "{missing:?}"); + s.cleanup().await; + } + + /// A server (or any other session) still connected to the data database + /// blocks 004; the migrate command's own connections do not. + #[tokio::test] + async fn refuses_004_while_another_client_is_connected() { + if base_conn().is_none() { + skip("refuses_004_while_another_client_is_connected"); + return; + } + let s = scratch().await; + seed(&s.db).await; + let base = base_conn().unwrap(); + + // An unnamed pool stands in for a running server. + let server_pool = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/{}", s.db_name)) + .await + .expect("connect as a server would"); + sqlx::query("SELECT 1").execute(&server_pool).await.unwrap(); + + let err = refuse_if_other_clients_connected(&s.db) + .await + .expect_err("an unnamed connection must block the migration"); + let text = format!("{err:?}"); + assert!( + text.contains("other connection(s) hold the data database"), + "{text}" + ); + assert!(text.contains(IGNORE_CONNECTIONS_ENV), "{text}"); + let err = run_data_code_migrations(&s.db, &s.db) + .await + .expect_err("004 must not run while the server is connected"); + assert!(format!("{err:?}").contains("other connection(s)")); + assert!( + unapplied_required_data_migrations(&s.db) + .await + .unwrap() + .len() + == 1, + "nothing may be recorded when the guard refuses" + ); + + // Server stopped: only the migrate command's own connections remain. + server_pool.close().await; + refuse_if_other_clients_connected(&s.db) + .await + .expect("no foreign connection left"); + run_data_code_migrations(&s.db, &s.db) + .await + .expect("run 004"); + s.cleanup().await; + } +} diff --git a/crates/storage-postgres/src/stream_engine.rs b/crates/storage-postgres/src/stream_engine.rs index 151a90180..059b26ce2 100755 --- a/crates/storage-postgres/src/stream_engine.rs +++ b/crates/storage-postgres/src/stream_engine.rs @@ -94,8 +94,7 @@ impl StreamEngine for PostgresEngine { let shard_id = shard_id.to_string(); let table_name = table_name.to_string(); Box::pin(async move { - let record_json = - serde_json::to_value(&record).map_err(|e| StorageError::Internal(e.to_string()))?; + let record_json = crate::data::to_stored_json(&record)?; let table_id: String = sqlx::query_scalar( "SELECT table_id FROM tables WHERE account_id = $1 AND table_name = $2", @@ -197,9 +196,7 @@ impl StreamEngine for PostgresEngine { let records: Vec = rows .into_iter() - .map(|(data,)| { - serde_json::from_value(data).map_err(|e| StorageError::Internal(e.to_string())) - }) + .map(|(data,)| crate::data::stored_json_to(data)) .collect::, _>>()?; let last_seq = records.last().map(|r| r.dynamodb.sequence_number.clone()); diff --git a/crates/storage-postgres/src/vector_search.rs b/crates/storage-postgres/src/vector_search.rs index 5b278d467..5f4106acf 100644 --- a/crates/storage-postgres/src/vector_search.rs +++ b/crates/storage-postgres/src/vector_search.rs @@ -277,10 +277,17 @@ impl VectorSearchEngine for PostgresEngine { .bind(partition.into_bytes()) .bind(top_k) .bind(query_norm); + // The stored item carries the same escape as every other item_data + // column, so the filter's attribute name and value must be compared + // in stored form. for (name, value) in &filters { - let value_json = serde_json::to_string(value) - .map_err(|e| StorageError::Internal(format!("filter value: {e}")))?; - query = query.bind(name).bind(value_json); + let value_json = serde_json::to_value(value) + .map(extenddb_storage::util::escape_json_strings) + .map_err(|e| StorageError::Internal(format!("filter value: {e}")))? + .to_string(); + query = query + .bind(extenddb_storage::util::escape_control(name).into_owned()) + .bind(value_json); } let rows = query.fetch_all(&self.data_pool).await.map_err(|e| { @@ -298,7 +305,7 @@ impl VectorSearchEngine for PostgresEngine { let mut hits = Vec::with_capacity(rows.len()); for (ordered, embedding, item_json) in rows { - let mut item: extenddb_core::types::Item = serde_json::from_value(item_json) + let mut item: extenddb_core::types::Item = crate::data::json_to_item(item_json) .map_err(|e| StorageError::Internal(format!("stored item: {e}")))?; // Reinstated from the stored f32s rather than from a second copy in // the payload, so what comes back is the narrowed value that was diff --git a/crates/storage/src/util/control_chars.rs b/crates/storage/src/util/control_chars.rs new file mode 100644 index 000000000..07db91310 --- /dev/null +++ b/crates/storage/src/util/control_chars.rs @@ -0,0 +1,332 @@ +// SPDX-License-Identifier: Apache-2.0 + +//! Order-preserving escape for the two characters a storage engine may not be +//! able to hold verbatim. +//! +//! DynamoDB accepts the character U+0000 anywhere a string appears: partition +//! and sort keys, index keys, attribute values, map keys, attribute names. Two +//! of the storage engines cannot store it as is. PostgreSQL `TEXT` rejects the +//! byte (`invalid byte sequence for encoding "UTF8": 0x00`) and `jsonb` rejects +//! the `\u0000` escape. BSON field names are C strings and end at the first NUL. +//! +//! The escape maps each affected character to a two-character sequence that +//! contains no NUL: +//! +//! | input | output | +//! |--------|---------------| +//! | U+0000 | U+0001 U+0001 | +//! | U+0001 | U+0001 U+0002 | +//! +//! Every other character passes through unchanged, so a string without either +//! character encodes to itself and the stored form of ordinary data does not +//! change. The mapping is applied per character and the two output sequences +//! share the prefix U+0001, so three properties hold, each pinned by a test +//! below: +//! +//! - **Round trip.** `unescape(escape(s)) == s` for every `s`. +//! - **Byte order.** `escape(a) < escape(b)` exactly when `a < b` under UTF-8 +//! byte comparison. PostgreSQL compares key columns under `COLLATE "C"`, so +//! Query ordering, `BETWEEN`, and cursor comparisons on escaped keys give the +//! same answers as on the originals. +//! - **Prefix.** `escape(p)` is a prefix of `escape(s)` exactly when `p` is a +//! prefix of `s`, so `begins_with` on escaped values is exact. +//! +//! Decoding is lenient about a U+0001 that is not followed by U+0001 or U+0002: +//! it is returned as a literal U+0001. Rows written before this escape existed +//! can hold such a character, and a lenient decoder reads them unchanged. A +//! legacy string that contained U+0001 immediately followed by U+0001 or U+0002 +//! would decode differently; the operator migration `004_escape_control_chars` +//! re-encodes such rows so that case cannot arise after it has run. + +use std::borrow::Cow; + +const ESC: char = '\u{1}'; +const NUL_TAIL: char = '\u{1}'; +const ESC_TAIL: char = '\u{2}'; + +/// Whether `s` contains a character the escape changes. +#[must_use] +pub fn needs_escape(s: &str) -> bool { + s.bytes().any(|b| b == 0 || b == 1) +} + +/// Escape U+0000 and U+0001 in `s`. Borrows when there is nothing to change. +#[must_use] +pub fn escape_control(s: &str) -> Cow<'_, str> { + if !needs_escape(s) { + return Cow::Borrowed(s); + } + let mut out = String::with_capacity(s.len() + 8); + for ch in s.chars() { + match ch { + '\u{0}' => { + out.push(ESC); + out.push(NUL_TAIL); + } + '\u{1}' => { + out.push(ESC); + out.push(ESC_TAIL); + } + other => out.push(other), + } + } + Cow::Owned(out) +} + +/// Reverse [`escape_control`]. Borrows when there is nothing to change. A +/// U+0001 not followed by U+0001 or U+0002 is kept as a literal U+0001 (see +/// the module documentation for why). +#[must_use] +pub fn unescape_control(s: &str) -> Cow<'_, str> { + if !s.bytes().any(|b| b == 1) { + return Cow::Borrowed(s); + } + let mut out = String::with_capacity(s.len()); + let mut chars = s.chars().peekable(); + while let Some(ch) = chars.next() { + if ch == ESC { + match chars.peek() { + Some(&NUL_TAIL) => { + chars.next(); + out.push('\u{0}'); + } + Some(&ESC_TAIL) => { + chars.next(); + out.push('\u{1}'); + } + _ => out.push('\u{1}'), + } + } else { + out.push(ch); + } + } + Cow::Owned(out) +} + +/// Apply [`escape_control`] to every string in a JSON tree: object keys and +/// string values, at every depth. Used where the whole document lands in a +/// store that rejects U+0000 in any string (PostgreSQL `jsonb`). +#[must_use] +pub fn escape_json_strings(v: serde_json::Value) -> serde_json::Value { + map_json(v, true, &|s| escape_control(s).into_owned()) +} + +/// Reverse [`escape_json_strings`]. +#[must_use] +pub fn unescape_json_strings(v: serde_json::Value) -> serde_json::Value { + map_json(v, true, &|s| unescape_control(s).into_owned()) +} + +/// Apply [`escape_control`] to every object key in a JSON tree, leaving string +/// values alone. Used where only field names are restricted (BSON). +#[must_use] +pub fn escape_json_keys(v: serde_json::Value) -> serde_json::Value { + map_json(v, false, &|s| escape_control(s).into_owned()) +} + +/// Reverse [`escape_json_keys`]. +#[must_use] +pub fn unescape_json_keys(v: serde_json::Value) -> serde_json::Value { + map_json(v, false, &|s| unescape_control(s).into_owned()) +} + +fn map_json( + v: serde_json::Value, + values_too: bool, + f: &dyn Fn(&str) -> String, +) -> serde_json::Value { + use serde_json::Value; + match v { + Value::String(s) if values_too => { + if needs_escape(&s) { + Value::String(f(&s)) + } else { + Value::String(s) + } + } + Value::Array(items) => Value::Array( + items + .into_iter() + .map(|item| map_json(item, values_too, f)) + .collect(), + ), + Value::Object(map) => { + let mut out = serde_json::Map::with_capacity(map.len()); + for (k, val) in map { + let key = if k.bytes().any(|b| b == 0 || b == 1) { + f(&k) + } else { + k + }; + out.insert(key, map_json(val, values_too, f)); + } + Value::Object(out) + } + other => other, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::json; + + /// Every string over the alphabet {NUL, U+0001, U+0002, 'a'} up to the + /// given length. Small alphabet, exhaustive: the escape only looks at + /// these three code points and everything else is identity. + fn corpus(max_len: usize) -> Vec { + let alphabet = ['\u{0}', '\u{1}', '\u{2}', 'a']; + let mut out = vec![String::new()]; + let mut frontier = vec![String::new()]; + for _ in 0..max_len { + let mut next = Vec::new(); + for s in &frontier { + for ch in alphabet { + let mut t = s.clone(); + t.push(ch); + next.push(t); + } + } + out.extend(next.iter().cloned()); + frontier = next; + } + out + } + + #[test] + fn plain_strings_are_borrowed_unchanged() { + for s in ["", "a", "hello", "\u{2}\u{3}", "\u{e9}\u{4e2d}", "a\tb\nc"] { + assert!(matches!(escape_control(s), Cow::Borrowed(_)), "{s:?}"); + assert!(matches!(unescape_control(s), Cow::Borrowed(_)), "{s:?}"); + assert_eq!(escape_control(s), s); + assert_eq!(unescape_control(s), s); + } + } + + #[test] + fn table_of_the_two_mappings() { + assert_eq!(escape_control("\u{0}"), "\u{1}\u{1}"); + assert_eq!(escape_control("\u{1}"), "\u{1}\u{2}"); + assert_eq!(escape_control("a\u{0}b"), "a\u{1}\u{1}b"); + assert_eq!(escape_control("\u{0}\u{1}"), "\u{1}\u{1}\u{1}\u{2}"); + assert_eq!(unescape_control("\u{1}\u{1}"), "\u{0}"); + assert_eq!(unescape_control("\u{1}\u{2}"), "\u{1}"); + } + + #[test] + fn escaped_output_never_contains_nul() { + for s in corpus(5) { + assert!(!escape_control(&s).contains('\u{0}'), "{s:?}"); + } + } + + #[test] + fn round_trips_every_string_in_the_corpus() { + for s in corpus(5) { + assert_eq!(unescape_control(&escape_control(&s)), s, "{s:?}"); + } + for s in [ + "caf\u{e9}\u{0}\u{4e2d}", + "\u{0}\u{0}\u{0}", + "\u{1}\u{1}\u{1}", + "x\u{1}\u{0}\u{2}y", + ] { + assert_eq!(unescape_control(&escape_control(s)), s, "{s:?}"); + } + } + + #[test] + fn escape_preserves_utf8_byte_order() { + // The whole point: PostgreSQL compares the escaped column under + // COLLATE "C", and every ordering answer must equal the answer on + // the originals. + let corpus = corpus(4); + for a in &corpus { + for b in &corpus { + let want = a.as_bytes().cmp(b.as_bytes()); + let got = escape_control(a) + .as_bytes() + .cmp(escape_control(b).as_bytes()); + assert_eq!(want, got, "a={a:?} b={b:?}"); + } + } + } + + #[test] + fn escape_preserves_prefixes() { + // begins_with on escaped values must be exact in both directions. + let corpus = corpus(4); + for p in &corpus { + for s in &corpus { + let want = s.as_bytes().starts_with(p.as_bytes()); + let got = escape_control(s) + .as_bytes() + .starts_with(escape_control(p).as_bytes()); + assert_eq!(want, got, "p={p:?} s={s:?}"); + } + } + } + + #[test] + fn escaped_forms_are_distinct_for_distinct_inputs() { + let corpus = corpus(4); + let mut seen = std::collections::HashMap::new(); + for s in &corpus { + let e = escape_control(s).into_owned(); + if let Some(prev) = seen.insert(e.clone(), s.clone()) { + panic!("collision: {prev:?} and {s:?} both escape to {e:?}"); + } + } + } + + #[test] + fn lenient_decode_keeps_a_stray_escape_character_literal() { + // A row written before the escape existed can carry a raw U+0001 + // followed by anything other than U+0001 or U+0002. + assert_eq!(unescape_control("a\u{1}b"), "a\u{1}b"); + assert_eq!(unescape_control("\u{1}"), "\u{1}"); + assert_eq!(unescape_control("\u{1}\u{3}"), "\u{1}\u{3}"); + assert_eq!(unescape_control("x\u{1}"), "x\u{1}"); + } + + #[test] + fn json_strings_variant_touches_keys_and_values_at_every_depth() { + let item = json!({ + "a\u{0}b": {"S": "v\u{0}"}, + "m": {"M": {"k\u{1}": {"S": "\u{0}"}, "plain": {"N": "1"}}}, + "l": {"L": [{"S": "\u{0}x"}, {"SS": ["\u{0}", "y"]}]}, + "b": {"B": "AAEC"} + }); + let escaped = escape_json_strings(item.clone()); + assert!(!escaped.to_string().contains("\\u0000"), "{escaped}"); + assert_eq!(escaped["a\u{1}\u{1}b"]["S"], json!("v\u{1}\u{1}")); + assert_eq!(escaped["m"]["M"]["k\u{1}\u{2}"]["S"], json!("\u{1}\u{1}")); + assert_eq!(escaped["l"]["L"][1]["SS"][0], json!("\u{1}\u{1}")); + assert_eq!(escaped["b"]["B"], json!("AAEC")); + assert_eq!(unescape_json_strings(escaped), item); + } + + #[test] + fn json_keys_variant_leaves_values_alone() { + let item = json!({ + "a\u{0}b": {"S": "v\u{0}"}, + "m": {"M": {"k\u{0}": {"L": [{"S": "\u{0}"}]}}} + }); + let escaped = escape_json_keys(item.clone()); + assert_eq!(escaped["a\u{1}\u{1}b"]["S"], json!("v\u{0}")); + assert_eq!( + escaped["m"]["M"]["k\u{1}\u{1}"]["L"][0]["S"], + json!("\u{0}") + ); + assert_eq!(unescape_json_keys(escaped), item); + } + + #[test] + fn json_without_control_characters_is_unchanged() { + let item = json!({"pk": {"S": "p"}, "n": {"N": "1"}, "m": {"M": {"k": {"S": "v"}}}}); + assert_eq!(escape_json_strings(item.clone()), item); + assert_eq!(escape_json_keys(item.clone()), item); + assert_eq!(unescape_json_strings(item.clone()), item); + assert_eq!(unescape_json_keys(item.clone()), item); + } +} diff --git a/crates/storage/src/util/key.rs b/crates/storage/src/util/key.rs index 1519afd9f..4453b794b 100755 --- a/crates/storage/src/util/key.rs +++ b/crates/storage/src/util/key.rs @@ -23,9 +23,10 @@ pub enum SortKeyValue { /// /// For single-attribute keys, returns the value directly (no encoding). /// For multi-attribute keys, uses netstring encoding: each part is encoded as -/// `:,` and concatenated. This is provably collision-free -/// regardless of value content, and compatible with `PostgreSQL` TEXT columns -/// (no null bytes). +/// `:,` and concatenated. This is collision-free +/// regardless of value content. The parts are the raw attribute strings; a +/// backend whose columns cannot hold every character (PostgreSQL, see +/// `escape_control`) escapes its key text in its own wrapper. pub fn composite_pk_to_text( item: &Item, key_schema: &[KeySchemaElement], @@ -286,7 +287,7 @@ pub fn recover_sort_key_definitions( /// Encode multiple string parts into a single netstring-encoded composite key. /// /// Format: `:,:,...` — e.g., `"abc"` + `"de"` → `"3:abc,2:de,"`. -/// This encoding is unambiguous for arbitrary byte content and contains no null bytes. +/// This encoding is unambiguous for arbitrary part content; the parts are copied as given. #[must_use] pub fn encode_netstring_composite(parts: &[String]) -> String { let mut out = String::new(); diff --git a/crates/storage/src/util/mod.rs b/crates/storage/src/util/mod.rs index 81e2ae850..3f1d08a06 100755 --- a/crates/storage/src/util/mod.rs +++ b/crates/storage/src/util/mod.rs @@ -7,9 +7,14 @@ //! id validation. mod arn; +mod control_chars; mod key; pub use arn::{index_arn, parse_stream_arn, stream_arn, table_arn}; +pub use control_chars::{ + escape_control, escape_json_keys, escape_json_strings, needs_escape, unescape_control, + unescape_json_keys, unescape_json_strings, +}; pub use key::SortKeyValue; pub use key::{ composite_pk_to_text, effective_attribute_definitions, encode_netstring_composite, diff --git a/docs/design/04-component-storage.md b/docs/design/04-component-storage.md index b85c60232..5b6be7fb5 100755 --- a/docs/design/04-component-storage.md +++ b/docs/design/04-component-storage.md @@ -402,6 +402,25 @@ The PostgreSQL backend uses two categories of tables: attributes, matching the DynamoDB model where key attributes are part of the item. +- **U+0000 in strings**: DynamoDB accepts the character U+0000 anywhere a + string appears (measured 2026-09-18: partition and sort keys, index keys, + values, strings in lists, map keys, attribute names; it sorts as the byte + 0x00 and is a different key from the six-character text `\u0000`). + PostgreSQL `TEXT` rejects the byte and `jsonb` rejects the `\u0000` escape, + so every string that reaches a key column or an `item_data`, `gsi_pending`, + `stream_records`, or `backup_items` document goes through the escape in + `extenddb_storage::util::escape_control`: U+0000 is stored as U+0001 U+0001 + and U+0001 as U+0001 U+0002; every other character is unchanged, so ordinary + data is stored exactly as before. The mapping preserves UTF-8 byte order and + prefixes, so `COLLATE "C"` comparisons, `BETWEEN`, `begins_with`, and + row-comparison cursors on escaped columns return the rows the raw text + would. SQL that addresses an attribute by name inside `item_data` (the TTL + index and sweep, vector search filters) uses the escaped name. The MongoDB + backend applies the same escape to BSON field names only, since BSON string + values are length-prefixed and hold the byte as is. Rows written before the + escape existed are re-encoded once by the operator migration + `004_escape_control_chars` (see the upgrade manual). + - **GSI tables**: GSI tables include base table primary key columns (`base_pk`, `base_sk_*`) as actual SQL columns (not just inside `item_data` JSONB). This is required because: (1) GSI keys are not unique — two base table items can diff --git a/docs/manuals/07-upgrade-manual.md b/docs/manuals/07-upgrade-manual.md index fe2bf5eb2..3445f26aa 100755 --- a/docs/manuals/07-upgrade-manual.md +++ b/docs/manuals/07-upgrade-manual.md @@ -212,6 +212,66 @@ The first release, so all 0.0.2 deployments were fresh installs. --- +## Data Migrations + +Some changes rewrite rows in the data database rather than the catalog schema. These are +implemented in Rust (`DATA_CODE_MIGRATIONS` in `crates/storage-postgres/src/migrations.rs`), +tracked in the data database's own `schema_history` table, and applied once by +`extenddb migrate --yes`. Running `extenddb migrate` without `--yes` lists them as pending +and changes nothing. + +### 004_escape_control_chars + +Strings are stored through an order-preserving escape so that U+0000, which DynamoDB +accepts anywhere a string appears, can be held in PostgreSQL `TEXT` and `jsonb` columns: +U+0000 is stored as U+0001 U+0001, and U+0001 as U+0001 U+0002. Rows written before the +escape existed hold their strings raw, and a raw U+0001 in such a row could be misread by +the decoder. This migration re-encodes those rows. + +What it rewrites: every data table and index table (the text key columns and the +`item_data` document), vector index tables, `gsi_pending`, `stream_records`, and +`backup_items`. Only rows that contain U+0001 are touched; a database with none is left as +it is, apart from the `schema_history` row. + +The migration scans every data and index table once (each row's `jsonb` document is +rendered to text to find candidates), so on a large deployment it takes time proportional +to the size of the data. It processes one table per transaction and keeps a per-table +progress marker in `data_code_migration_progress`, so an interrupted run resumes with the +tables that are not yet done and never rewrites a table twice. + +Order of operations. A server built with the escape refuses to start on a data database +that does not record this migration, with a message naming it, because rows the new build +writes are indistinguishable from legacy rows to the migration (both can contain U+0001) and +would be rewritten by it. Upgrade a PostgreSQL deployment in this order: + +1. Stop every ExtendDB server that uses the data database. A server from an earlier release + still writing during the migration would leave unescaped rows behind. +2. Run `extenddb migrate --yes` with the new binary. +3. Start the new servers. + +A fresh `extenddb init` records the migration, so new deployments start without this step. +A rolling deployment that mixes the two releases against one data database is not +supported for this release: the new server will not start until the migration has run, +and the old server must be stopped before it runs. + +`extenddb migrate --yes` refuses to run this migration while any other session is connected +to the data database, and lists them, because a server still writing during the rewrite +would leave rows in the wrong encoding. If the only other connections are idle ones held +by a connection pooler, set `EXTENDDB_MIGRATE_IGNORE_CONNECTIONS=1` to skip that check. +Nothing stops a server from an earlier release from being started against a migrated +data database afterwards; it would read escaped rows as literal data. Do not downgrade +past this release without restoring the data database from a backup taken before the +migration. + +MongoDB. The MongoDB backend applies the same escape to BSON field names (attribute +names and map keys) and has no migration and no startup check. Field names written by +an earlier release are stored raw. A name containing a lone U+0001 reads back unchanged, +as do all string values and all key strings; a name containing U+0001 immediately +followed by U+0001 or U+0002 reads back with that pair decoded as U+0000 or U+0001. +Deployments whose attribute names never contain U+0001 are unaffected. + +--- + ## Behavior Changes by Release Catalog upgrades change the schema; behavior changes alter how the running server diff --git a/tests/test_nul_strings.py b/tests/test_nul_strings.py new file mode 100644 index 000000000..0e1d8c511 --- /dev/null +++ b/tests/test_nul_strings.py @@ -0,0 +1,378 @@ +# SPDX-License-Identifier: Apache-2.0 +"""U+0000 and U+0001 inside strings: keys, index keys, values, and names. + +DynamoDB places no character restriction on strings. Measured 2026-09-18 +against the service on a scratch table: PutItem accepts the NUL character +inside a partition key, a sort key, a GSI key, a non-key value, a string in a +list, a map key, and a top-level attribute name, and GetItem returns every one +byte-identical. NUL sorts as the byte 0x00, before every other character; it is +a different character from the six-character text ``\\u0000``; range and prefix +conditions treat it as an ordinary byte; the GSI holds the item; DeleteItem on +the key removes it. These tests pin exactly that, so every backend has to store +the character rather than its escape text and has to keep it out of the way of +its own storage encoding. + +U+0001 is included because a storage encoding that handles NUL has to leave the +next code point alone as well. +""" + +from __future__ import annotations + +import time + +import pytest +from botocore.exceptions import ClientError + +from conftest import scoped_table, wait_for_gsi_items + +NUL = "\u0000" +SOH = "\u0001" +ESCAPE_TEXT = "\\u0000" # backslash, u, 0, 0, 0, 0: six characters, not a NUL + +PK = f"a{NUL}b" +PK_WITHOUT_NUL = "ab" +PK_SOH = f"a{SOH}b" + +# One partition, five sort keys whose byte order is: NUL, U+0001, space, the +# escape text (starts with 0x5c), then "a" (0x61). +SORT_KEYS = [NUL, SOH, " ", ESCAPE_TEXT, "a"] + +# The item under the NUL sort key carries the character in every kind of +# position the item model has. +RICH_ITEM = { + "pk": {"S": PK}, + "sk": {"S": NUL}, + "g": {"S": f"g{NUL}"}, + "v": {"S": f"v{NUL}"}, + f"n{NUL}n": {"S": "x"}, + "m": {"M": {f"k{NUL}": {"S": f"{NUL}"}, f"k{SOH}{SOH}": {"N": "1"}}}, + "l": {"L": [{"S": NUL}, {"S": f"{SOH}{NUL}{SOH}"}]}, + "ss": {"SS": [NUL, "a"]}, +} + + +def _query(client, table, **kwargs): + return client.query( + TableName=table, + KeyConditionExpression=kwargs.pop("cond", "pk = :p"), + ExpressionAttributeValues={":p": {"S": PK}, **kwargs.pop("values", {})}, + **kwargs, + ) + + +def _sks(resp): + return [i["sk"]["S"] for i in resp["Items"]] + + +@pytest.fixture(scope="class") +def nul_table(dynamodb_client): + """(S, S) table with GSI gsi1 on g; the partition PK holds SORT_KEYS, plus + a control item in the partition without the NUL and one in a partition + keyed with U+0001.""" + with scoped_table( + dynamodb_client, + attribute_definitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + {"AttributeName": "g", "AttributeType": "S"}, + ], + key_schema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + GlobalSecondaryIndexes=[ + { + "IndexName": "gsi1", + "KeySchema": [{"AttributeName": "g", "KeyType": "HASH"}], + "Projection": {"ProjectionType": "ALL"}, + } + ], + ) as name: + dynamodb_client.put_item(TableName=name, Item=RICH_ITEM) + for sk in SORT_KEYS[1:]: + dynamodb_client.put_item( + TableName=name, Item={"pk": {"S": PK}, "sk": {"S": sk}, "g": {"S": "plain"}} + ) + dynamodb_client.put_item( + TableName=name, Item={"pk": {"S": PK_WITHOUT_NUL}, "sk": {"S": "x"}, "g": {"S": "plain"}} + ) + dynamodb_client.put_item( + TableName=name, + Item={"pk": {"S": PK_SOH}, "sk": {"S": SOH}, "g": {"S": f"g{SOH}"}, "v": {"S": SOH}}, + ) + yield name + + +class TestNulInStrings: + def test_item_with_nul_everywhere_round_trips(self, dynamodb_client, nul_table): + resp = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK}, "sk": {"S": NUL}}, ConsistentRead=True + ) + assert resp["Item"] == RICH_ITEM + + def test_soh_item_round_trips(self, dynamodb_client, nul_table): + resp = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK_SOH}, "sk": {"S": SOH}}, ConsistentRead=True + ) + assert resp["Item"] == { + "pk": {"S": PK_SOH}, + "sk": {"S": SOH}, + "g": {"S": f"g{SOH}"}, + "v": {"S": SOH}, + } + + def test_partition_order_is_byte_order_with_nul_first(self, dynamodb_client, nul_table): + assert _sks(_query(dynamodb_client, nul_table, ConsistentRead=True)) == SORT_KEYS + assert ( + _sks(_query(dynamodb_client, nul_table, ScanIndexForward=False, ConsistentRead=True)) + == SORT_KEYS[::-1] + ) + + def test_nul_and_its_escape_text_are_different_keys(self, dynamodb_client, nul_table): + by_nul = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK}, "sk": {"S": NUL}}, ConsistentRead=True + )["Item"] + by_text = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK}, "sk": {"S": ESCAPE_TEXT}}, ConsistentRead=True + )["Item"] + assert by_nul["sk"]["S"] == NUL and len(by_nul["sk"]["S"]) == 1 + assert by_text["sk"]["S"] == ESCAPE_TEXT and len(by_text["sk"]["S"]) == 6 + + def test_partition_key_with_nul_is_distinct_from_the_key_without_it( + self, dynamodb_client, nul_table + ): + resp = dynamodb_client.get_item( + TableName=nul_table, Key={"pk": {"S": PK_WITHOUT_NUL}, "sk": {"S": NUL}}, ConsistentRead=True + ) + assert "Item" not in resp + resp = _query(dynamodb_client, nul_table, values={":p": {"S": PK_WITHOUT_NUL}}, ConsistentRead=True) + assert _sks(resp) == ["x"] + + def test_range_conditions_treat_nul_as_a_byte(self, dynamodb_client, nul_table): + below_space = _query( + dynamodb_client, nul_table, cond="pk = :p AND sk < :s", values={":s": {"S": " "}}, + ConsistentRead=True, + ) + assert _sks(below_space) == [NUL, SOH] + between = _query( + dynamodb_client, nul_table, cond="pk = :p AND sk BETWEEN :lo AND :hi", + values={":lo": {"S": NUL}, ":hi": {"S": SOH}}, ConsistentRead=True, + ) + assert _sks(between) == [NUL, SOH] + prefix = _query( + dynamodb_client, nul_table, cond="pk = :p AND begins_with(sk, :pre)", + values={":pre": {"S": NUL}}, ConsistentRead=True, + ) + assert _sks(prefix) == [NUL] + text_prefix = _query( + dynamodb_client, nul_table, cond="pk = :p AND begins_with(sk, :pre)", + values={":pre": {"S": "\\"}}, ConsistentRead=True, + ) + assert _sks(text_prefix) == [ESCAPE_TEXT] + + def test_pages_of_one_visit_every_key_once(self, dynamodb_client, nul_table): + """The cursor for a NUL key travels through LastEvaluatedKey and back + through ExclusiveStartKey; the walk must neither skip nor repeat.""" + seen: list[str] = [] + kwargs: dict = {"Limit": 1, "ConsistentRead": True} + while True: + resp = _query(dynamodb_client, nul_table, **kwargs) + seen += _sks(resp) + if "LastEvaluatedKey" not in resp: + break + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + assert seen == SORT_KEYS + + def test_gsi_holds_the_item_with_a_nul_key(self, dynamodb_client, nul_table): + def walk(): + return dynamodb_client.query( + TableName=nul_table, + IndexName="gsi1", + KeyConditionExpression="g = :g", + ExpressionAttributeValues={":g": {"S": f"g{NUL}"}}, + )["Items"] + + items = wait_for_gsi_items(walk, 1) + assert items == [RICH_ITEM] + + def test_filter_on_a_nul_named_attribute(self, dynamodb_client, nul_table): + resp = dynamodb_client.scan( + TableName=nul_table, + FilterExpression="#n = :x", + ExpressionAttributeNames={"#n": f"n{NUL}n"}, + ExpressionAttributeValues={":x": {"S": "x"}}, + ConsistentRead=True, + ) + assert resp["Items"] == [RICH_ITEM] + resp = _query( + dynamodb_client, nul_table, FilterExpression="v = :v", + values={":v": {"S": f"v{NUL}"}}, ConsistentRead=True, + ) + assert resp["Items"] == [RICH_ITEM] + + def test_condition_expression_sees_the_nul_named_attribute(self, dynamodb_client, nul_table): + # Exists on the rich item: the conditional write goes through. + dynamodb_client.update_item( + TableName=nul_table, + Key={"pk": {"S": PK}, "sk": {"S": NUL}}, + UpdateExpression="SET touched = :t", + ConditionExpression="attribute_exists(#n)", + ExpressionAttributeNames={"#n": f"n{NUL}n"}, + ExpressionAttributeValues={":t": {"BOOL": True}}, + ) + with pytest.raises(ClientError) as exc: + dynamodb_client.update_item( + TableName=nul_table, + Key={"pk": {"S": PK}, "sk": {"S": "a"}}, + UpdateExpression="SET touched = :t", + ConditionExpression="attribute_exists(#n)", + ExpressionAttributeNames={"#n": f"n{NUL}n"}, + ExpressionAttributeValues={":t": {"BOOL": True}}, + ) + assert exc.value.response["Error"]["Code"] == "ConditionalCheckFailedException" + dynamodb_client.update_item( + TableName=nul_table, + Key={"pk": {"S": PK}, "sk": {"S": NUL}}, + UpdateExpression="REMOVE touched", + ) + + def test_update_item_sets_nul_named_attribute_and_nul_value(self, dynamodb_client, nul_table): + key = {"pk": {"S": PK}, "sk": {"S": " "}} + resp = dynamodb_client.update_item( + TableName=nul_table, + Key=key, + UpdateExpression="SET #a = :v, m2 = :m", + ExpressionAttributeNames={"#a": f"u{NUL}"}, + ExpressionAttributeValues={ + ":v": {"S": f"{NUL}z"}, + ":m": {"M": {f"{NUL}": {"L": [{"S": SOH}]}}}, + }, + ReturnValues="ALL_NEW", + ) + assert resp["Attributes"][f"u{NUL}"] == {"S": f"{NUL}z"} + assert resp["Attributes"]["m2"] == {"M": {f"{NUL}": {"L": [{"S": SOH}]}}} + got = dynamodb_client.get_item(TableName=nul_table, Key=key, ConsistentRead=True)["Item"] + assert got[f"u{NUL}"] == {"S": f"{NUL}z"} + assert got["m2"] == {"M": {f"{NUL}": {"L": [{"S": SOH}]}}} + resp = dynamodb_client.update_item( + TableName=nul_table, + Key=key, + UpdateExpression="REMOVE #a, m2", + ExpressionAttributeNames={"#a": f"u{NUL}"}, + ReturnValues="ALL_NEW", + ) + assert f"u{NUL}" not in resp["Attributes"] and "m2" not in resp["Attributes"] + + def test_batch_and_transact_writes_and_reads(self, dynamodb_client, nul_table): + b1 = {"pk": {"S": f"batch{NUL}"}, "sk": {"S": NUL}, "g": {"S": "plain"}, f"{NUL}": {"S": NUL}} + t1 = {"pk": {"S": f"tx{NUL}"}, "sk": {"S": f"{NUL}{NUL}"}, "g": {"S": "plain"}} + dynamodb_client.batch_write_item(RequestItems={nul_table: [{"PutRequest": {"Item": b1}}]}) + dynamodb_client.transact_write_items( + TransactItems=[{"Put": {"TableName": nul_table, "Item": t1}}] + ) + got = dynamodb_client.batch_get_item( + RequestItems={ + nul_table: { + "Keys": [{"pk": b1["pk"], "sk": b1["sk"]}, {"pk": t1["pk"], "sk": t1["sk"]}], + "ConsistentRead": True, + } + } + )["Responses"][nul_table] + assert sorted(got, key=lambda i: i["pk"]["S"]) == sorted([b1, t1], key=lambda i: i["pk"]["S"]) + tg = dynamodb_client.transact_get_items( + TransactItems=[ + {"Get": {"TableName": nul_table, "Key": {"pk": b1["pk"], "sk": b1["sk"]}}}, + {"Get": {"TableName": nul_table, "Key": {"pk": t1["pk"], "sk": t1["sk"]}}}, + ] + )["Responses"] + assert [r["Item"] for r in tg] == [b1, t1] + dynamodb_client.batch_write_item( + RequestItems={ + nul_table: [ + {"DeleteRequest": {"Key": {"pk": b1["pk"], "sk": b1["sk"]}}}, + {"DeleteRequest": {"Key": {"pk": t1["pk"], "sk": t1["sk"]}}}, + ] + } + ) + + def test_scan_counts_every_item(self, dynamodb_client, nul_table): + resp = dynamodb_client.scan(TableName=nul_table, Select="COUNT", ConsistentRead=True) + assert resp["Count"] == len(SORT_KEYS) + 2 + + def test_delete_item_on_a_nul_key(self, dynamodb_client, nul_table): + key = {"pk": {"S": f"del{NUL}"}, "sk": {"S": NUL}} + dynamodb_client.put_item(TableName=nul_table, Item={**key, "g": {"S": "plain"}}) + resp = dynamodb_client.delete_item(TableName=nul_table, Key=key, ReturnValues="ALL_OLD") + assert resp["Attributes"] == {**key, "g": {"S": "plain"}} + assert "Item" not in dynamodb_client.get_item(TableName=nul_table, Key=key, ConsistentRead=True) + + +@pytest.fixture(scope="module") +def nul_streams_client(endpoint_url): + """Streams client on the ExtendDB endpoint. The service routes Streams to a + separate endpoint boto3 resolves on its own, so this test is ExtendDB only, + as the stream tests in ``test_streams.py`` are.""" + if not endpoint_url: + pytest.skip("stream record fidelity is checked against ExtendDB servers only") + import os + + import boto3 + + kwargs: dict = dict( + service_name="dynamodbstreams", + region_name=os.environ.get("AWS_DEFAULT_REGION", "us-east-1"), + endpoint_url=endpoint_url, + ) + if endpoint_url.startswith("https://"): + kwargs["verify"] = False + return boto3.client(**kwargs) + + +def test_stream_records_carry_nul_keys_names_and_values(dynamodb_client, nul_streams_client): + """INSERT, MODIFY, and REMOVE records for an item whose key, attribute + names, map keys, and values hold U+0000 and U+0001 come back byte-identical + in Keys, NewImage, and OldImage.""" + from test_streams import _drain_all_shards + + with scoped_table( + dynamodb_client, + attribute_definitions=[ + {"AttributeName": "pk", "AttributeType": "S"}, + {"AttributeName": "sk", "AttributeType": "S"}, + ], + key_schema=[ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}, + ], + StreamSpecification={"StreamEnabled": True, "StreamViewType": "NEW_AND_OLD_IMAGES"}, + ) as name: + stream_arn = dynamodb_client.describe_table(TableName=name)["Table"]["LatestStreamArn"] + key = {"pk": {"S": PK}, "sk": {"S": NUL}} + first = {**RICH_ITEM} + dynamodb_client.put_item(TableName=name, Item=first) + second = dynamodb_client.update_item( + TableName=name, + Key=key, + UpdateExpression="SET #a = :v REMOVE #n", + ExpressionAttributeNames={"#a": f"u{NUL}", "#n": f"n{NUL}n"}, + ExpressionAttributeValues={":v": {"S": f"{SOH}{NUL}"}}, + ReturnValues="ALL_NEW", + )["Attributes"] + dynamodb_client.delete_item(TableName=name, Key=key) + + deadline = time.monotonic() + 30 + records: list[dict] = [] + while time.monotonic() < deadline: + records = _drain_all_shards(nul_streams_client, stream_arn) + if len(records) >= 3: + break + time.sleep(0.5) + assert [r["eventName"] for r in records] == ["INSERT", "MODIFY", "REMOVE"] + for r in records: + assert r["dynamodb"]["Keys"] == key + assert records[0]["dynamodb"]["NewImage"] == first + assert "OldImage" not in records[0]["dynamodb"] + assert records[1]["dynamodb"]["OldImage"] == first + assert records[1]["dynamodb"]["NewImage"] == second + assert records[2]["dynamodb"]["OldImage"] == second + assert "NewImage" not in records[2]["dynamodb"]