diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index b0d641a98..421a28e08 100755 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -341,7 +341,7 @@ SQLite, PostgreSQL, Cassandra and MongoDB migrate through this framework. ### Example: a nullable column on `tables` -Migration 6, `table_note`, adds a nullable `note TEXT` column to the catalog's +Migration 7, `table_note`, adds a nullable `note TEXT` column to the catalog's `tables` table on SQLite and PostgreSQL. In `crates/storage/src/migrations/registry.rs`, at the end of the @@ -349,18 +349,18 @@ In `crates/storage/src/migrations/registry.rs`, at the end of the ```rust /// Adds a nullable `note` column to the `tables` catalog table. - TableNote = 6, "table_note", + TableNote = 7, "table_note", "add a nullable note column to the tables catalog"; ``` PostgreSQL has `ADD COLUMN IF NOT EXISTS`, so its body is a file, -`crates/storage-postgres/migrations/0006_table_note.sql`: +`crates/storage-postgres/migrations/0007_table_note.sql`: ```sql -- Copyright 2026 ExtendDB contributors -- SPDX-License-Identifier: Apache-2.0 -- --- Migration 6, table_note, PostgreSQL catalog database: a nullable note +-- Migration 7, table_note, PostgreSQL catalog database: a nullable note -- column on tables. IF NOT EXISTS makes a second run change nothing. ALTER TABLE tables ADD COLUMN IF NOT EXISTS note TEXT; @@ -375,7 +375,7 @@ In `crates/storage-postgres/src/migrations/mod.rs`, the arm in `migration()`: bodies: vec![file_body!( StoreKind::Catalog, "migrations", - "0006_table_note.sql" + "0007_table_note.sql" )], }, ``` @@ -393,13 +393,13 @@ and the id in the transactional arm of `runs_in()`: SQLite has no `ADD COLUMN IF NOT EXISTS`, and a file with a plain `ADD COLUMN` fails when it runs a second time, so the SQLite body is code that checks for the column first. -`crates/storage-sqlite/src/migrations/m0006_table_note.rs`: +`crates/storage-sqlite/src/migrations/m0007_table_note.rs`: ```rust // Copyright 2026 ExtendDB contributors // SPDX-License-Identifier: Apache-2.0 -//! Migration 6, `table_note`, SQLite: a nullable `note` column on `tables`. +//! Migration 7, `table_note`, SQLite: a nullable `note` column on `tables`. //! //! SQLite has no `ADD COLUMN IF NOT EXISTS`, so the body reads the table's //! columns first and adds the column only when it is missing. A second run @@ -411,10 +411,10 @@ use extenddb_storage::migrations::{CodeBody, Done, MigrationError}; use super::SqliteStores; /// This file, whose SHA-256 the checksum test pins. -pub(crate) const SOURCE: &str = include_str!("m0006_table_note.rs"); +pub(crate) const SOURCE: &str = include_str!("m0007_table_note.rs"); /// Recorded as this migration's checksum. Changes only with a new migration. -const FINGERPRINT: &str = "0006_table_note/sqlite/1"; +const FINGERPRINT: &str = "0007_table_note/sqlite/1"; pub(crate) struct TableNote; @@ -453,7 +453,7 @@ beside the others: ```rust #[rustfmt::skip] #[allow(clippy::all, deprecated)] -mod m0006_table_note; +mod m0007_table_note; ``` and the arm in `migration()`: @@ -464,8 +464,8 @@ and the arm in `migration()`: convergent: true, bodies: vec![StoreBody::code( StoreKind::Catalog, - &m0006_table_note::TableNote, - m0006_table_note::SOURCE, + &m0007_table_note::TableNote, + m0007_table_note::SOURCE, )], }, ``` @@ -477,8 +477,8 @@ rows of `tables`: ```rust #[tokio::test] -async fn migration_6_adds_the_note_column_to_tables() { - let db = temp_db("m6"); +async fn migration_7_adds_the_note_column_to_tables() { + let db = temp_db("m7"); load_dump(&db, DEV_0_1_14).await; run(&db).await.expect("migrate"); let mut text = dump( @@ -488,12 +488,12 @@ async fn migration_6_adds_the_note_column_to_tables() { ) .await; text.push_str(&dump(&db, "SELECT table_name, note FROM tables ORDER BY table_name").await); - golden("0006_table_note", &text); + golden("0007_table_note", &text); } ``` -`EXTENDDB_BLESS_EXPECTED=1 cargo test -p extenddb-storage-sqlite --lib migration_6` -writes `crates/storage-sqlite/tests/fixtures/golden/0006_table_note.txt`, one +`EXTENDDB_BLESS_EXPECTED=1 cargo test -p extenddb-storage-sqlite --lib migration_7` +writes `crates/storage-sqlite/tests/fixtures/golden/0007_table_note.txt`, one tab-separated line per row under each query: ``` @@ -508,15 +508,15 @@ The PostgreSQL body is a file, so it has no golden test. The example declares `Kind::Breaking`. A nullable column can be additive by the rule in step 4, but declaring it additive needs the old-binary job, which this example does not add. As a breaking migration it makes a 1.x binary -without migration 6 refuse a catalog that has it, and a rollback restores the +without migration 7 refuse a catalog that has it, and a rollback restores the backup. The same PR changes the code that reads and writes `tables.note` in each backend's catalog store, and commits the `docs/schema/sqlite.txt` and `docs/schema/postgres.txt` that the integration jobs regenerate. It changes no per-table DDL, because `tables` is a catalog table, and adds no `RELEASED` -entry. `extenddb migrate --list` then shows migration 6, and on a catalog at -migration 5 `extenddb migrate` lists migration 6 as pending and exits 3. +entry. `extenddb migrate --list` then shows migration 7, and on a catalog at +migration 6 `extenddb migrate` lists migration 7 as pending and exits 3. ### Running the migration tests diff --git a/crates/storage-cassandra/src/migrations/mod.rs b/crates/storage-cassandra/src/migrations/mod.rs index aae33b64f..44a8b2229 100644 --- a/crates/storage-cassandra/src/migrations/mod.rs +++ b/crates/storage-cassandra/src/migrations/mod.rs @@ -127,6 +127,10 @@ pub fn migration(id: MigrationId) -> Migration { "0005_drop_continuous_backups.cql" )], }, + MigrationId::EscapeControlChars => Migration::NotApplicable( + "Cassandra stores U+0000 in strings as written, so its stored form of strings does \ + not change", + ), } } diff --git a/crates/storage-mongodb/src/data/mod.rs b/crates/storage-mongodb/src/data/mod.rs index d1678862a..4ee910a3b 100644 --- a/crates/storage-mongodb/src/data/mod.rs +++ b/crates/storage-mongodb/src/data/mod.rs @@ -15,7 +15,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 { @@ -34,6 +36,53 @@ 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. +/// +/// Documents written by a release before the escape hold their field names +/// raw, and there is no migration and no startup check for them: unlike +/// `PostgreSQL`, nothing here can fail to store (BSON strings already hold +/// U+0000, and a raw field name reads back through the driver), so the only +/// effect is on a legacy name that happens to contain U+0001 immediately +/// followed by U+0001 or U+0002, which this decoder reads as U+0000 or +/// U+0001. A lone U+0001 and every string value are unchanged. That exposure +/// was judged too narrow to make every `MongoDB` upgrade acknowledge a gate +/// with nothing to rewrite; the upgrade manual's "Data Migrations" section +/// states the limitation for operators. +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 }` @@ -44,10 +93,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(); @@ -194,9 +242,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) @@ -256,8 +302,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}")))?; @@ -436,6 +481,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 201f48333..518e9133e 100644 --- a/crates/storage-mongodb/src/data_engine.rs +++ b/crates/storage-mongodb/src/data_engine.rs @@ -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; @@ -802,8 +803,9 @@ impl MongoEngine { // common (no-GSI) case; when the cache is stale-or-unknown // we fall through to the slow path which is authoritative. Key // attribute names that are not plain path names (a dot, a leading - // `$`) also take the slow path, because the upsert writes the key - // attributes as `item_data.` (see `update_create`). + // `$`, U+0000 or U+0001) also take the slow path, because the + // upsert writes the key attributes as `item_data.` (see + // `update_create`). if condition.is_none() && !return_old && stream.is_none() @@ -1725,8 +1727,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 } => { @@ -2107,8 +2112,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 @@ -4479,6 +4486,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/data_engine_update_create.rs b/crates/storage-mongodb/src/data_engine_update_create.rs index a7f4e7e01..f8ca78d86 100644 --- a/crates/storage-mongodb/src/data_engine_update_create.rs +++ b/crates/storage-mongodb/src/data_engine_update_create.rs @@ -29,13 +29,18 @@ //! step into a sub-document. A key attribute named `a.b` would be inserted as //! `item_data: { a: { b: ... } }`, where PutItem stores the field `"a.b"`, //! and the same path conflicts with an update of an attribute named `a`. -//! A table with a key attribute name that `key_names_are_plain_paths` -//! rejects therefore does not take the native upsert; it takes the session -//! path, which writes the document `item_to_document` builds. +//! A key attribute name containing U+0000 or U+0001 has the same problem in +//! another form: `item_to_document` stores it escaped +//! (`extenddb_storage::util::control_chars`), and the path would name the +//! raw text. A table with a key attribute name that +//! `key_names_are_plain_paths` rejects therefore does not take the native +//! upsert; it takes the session path, which writes the document +//! `item_to_document` builds. use bson::Document; use extenddb_core::types::{Item, KeySchemaElement, TableKeyInfo}; use extenddb_storage::error::StorageError; +use extenddb_storage::util::needs_escape; use mongodb::options::{FindOneAndUpdateOptions, ReturnDocument}; use super::{TRANSIENT_RETRY_ATTEMPTS, backoff_sleep}; @@ -44,9 +49,11 @@ use crate::data::key_fields::insert_only_key_fields; /// Whether `name` is usable as one field name in a MongoDB update path. A dot /// would split it into nested fields and a leading `$` would read as an -/// operator or positional path. +/// operator or positional path. A name containing U+0000 or U+0001 is stored +/// in escaped form, so a path built from the raw name would address a +/// different field (and BSON refuses U+0000 in a field name). pub(super) fn is_plain_path_name(name: &str) -> bool { - !name.contains('.') && !name.starts_with('$') + !name.contains('.') && !name.starts_with('$') && !needs_escape(name) } /// Whether every key attribute of `key_schema` is a plain path name, so that @@ -173,4 +180,25 @@ mod tests { ); } } + + #[test] + fn a_key_name_with_u0000_or_u0001_leaves_the_native_upsert() { + // item_to_document stores these names escaped, so the + // `item_data.` path of the upsert would name another field. + assert!(key_names_are_plain_paths(&schema(&[( + "a\u{2}b", + KeyType::Hash + )]))); + for name in ["a\u{0}b", "a\u{1}b", "\u{1}\u{2}"] { + assert!( + !key_names_are_plain_paths(&schema(&[(name, KeyType::Hash)])), + "hash key {name:?}" + ); + let keys = [("pk", KeyType::Hash), (name, KeyType::Range)]; + assert!( + !key_names_are_plain_paths(&schema(&keys)), + "range key {name:?}" + ); + } + } } 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/migrations/mod.rs b/crates/storage-mongodb/src/migrations/mod.rs index 9a3c88598..d5a41da40 100644 --- a/crates/storage-mongodb/src/migrations/mod.rs +++ b/crates/storage-mongodb/src/migrations/mod.rs @@ -137,6 +137,11 @@ pub fn migration(id: MigrationId) -> Migration { m0005_drop_continuous_backups::SOURCE, )], }, + MigrationId::EscapeControlChars => Migration::NotApplicable( + "string values already held U+0000; field names an earlier release stored are read \ + back as written unless they hold U+0001 followed by U+0001 or U+0002, and are not \ + rewritten (see the upgrade manual)", + ), } } 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_record_docs.rs b/crates/storage-mongodb/src/stream_record_docs.rs index c84bbadf9..4fa57587c 100644 --- a/crates/storage-mongodb/src/stream_record_docs.rs +++ b/crates/storage-mongodb/src/stream_record_docs.rs @@ -12,7 +12,7 @@ use extenddb_core::types::StreamRecord; use extenddb_storage::error::StorageError; -use mongodb::bson::{self, Document}; +use mongodb::bson::Document; /// The records in `docs`, in order. A document whose `has_record` is false /// carries none. One whose `has_record` is true must carry `record_data`. @@ -23,8 +23,9 @@ pub(crate) fn records_from_docs(docs: Vec) -> Result let record_bson = d .remove("record_data") .ok_or_else(|| StorageError::Internal("Missing record_data".to_owned()))?; - let json_val: serde_json::Value = - bson::from_bson(record_bson).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())) }) .collect() diff --git a/crates/storage-mongodb/tests/migrations_fresh.rs b/crates/storage-mongodb/tests/migrations_fresh.rs index 12cf1838e..ea0859cdf 100644 --- a/crates/storage-mongodb/tests/migrations_fresh.rs +++ b/crates/storage-mongodb/tests/migrations_fresh.rs @@ -303,7 +303,7 @@ async fn a_fresh_store_gets_every_migration_and_the_expected_definitions() { let Some(uri) = server() else { return }; let d = Deployment::new(&uri, "fresh").await; let report = d.create().await; - assert_eq!((report.from, report.to), (None, Some(5))); + assert_eq!((report.from, report.to), (None, Some(6))); let applied: Vec<(u32, String)> = report .applied .iter() @@ -318,6 +318,7 @@ async fn a_fresh_store_gets_every_migration_and_the_expected_definitions() { (3, d.target.catalog.clone()), (4, d.target.catalog.clone()), (5, d.target.catalog.clone()), + (6, d.target.catalog.clone()), ] ); let not_applicable: Vec = report @@ -326,8 +327,8 @@ async fn a_fresh_store_gets_every_migration_and_the_expected_definitions() { .filter(|s| s.outcome == Outcome::NotApplicable) .map(|s| s.id) .collect(); - assert_eq!(not_applicable, [3, 4], "{report:?}"); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(not_applicable, [3, 4, 6], "{report:?}"); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); let rows = d.catalog_ledger().await; let outcomes: Vec<(u32, Outcome)> = rows.iter().map(|r| (r.id, r.outcome)).collect(); assert_eq!( @@ -338,6 +339,7 @@ async fn a_fresh_store_gets_every_migration_and_the_expected_definitions() { (3, Outcome::NotApplicable), (4, Outcome::NotApplicable), (5, Outcome::Applied), + (6, Outcome::NotApplicable), ] ); assert!( @@ -347,7 +349,7 @@ async fn a_fresh_store_gets_every_migration_and_the_expected_definitions() { let data_rows = Deployment::documents(&d.data(), "extenddb_migrations").await; assert_eq!(data_rows.len(), 1); assert_eq!(data_rows[0].get_i64("_id"), Ok(1)); - assert_eq!(d.setting("catalog_version").await.as_deref(), Some("5")); + assert_eq!(d.setting("catalog_version").await.as_deref(), Some("6")); assert_eq!(d.setting("resource_tag_migration").await, None); let history = Deployment::documents(&d.catalog(), "schema_history").await; assert_eq!(history.len(), 1); @@ -389,7 +391,7 @@ async fn a_second_run_applies_nothing_and_reconverge_keeps_the_rows() { let before = d.catalog_ledger().await; let again = d.migrate().await.expect("migrate"); assert!(again.applied.is_empty()); - assert_eq!((again.from, again.to), (Some(5), Some(5))); + assert_eq!((again.from, again.to), (Some(6), Some(6))); let reconverged = d.reconverge().await.expect("reconverge"); // Migration 1's steps in the catalog and the data database, and migration @@ -421,7 +423,7 @@ async fn a_second_run_applies_nothing_and_reconverge_keeps_the_rows() { before, "reconverge leaves the rows as they are" ); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); d.drop_databases().await; } @@ -456,7 +458,7 @@ async fn forget_clears_an_unreleased_id_and_migrate_applies_it_again() { assert_eq!(pending.iter().map(|p| p.id).collect::>(), [1, 1]); let report = d.migrate().await.expect("migrate"); assert_eq!(report.applied.len(), 2); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); d.drop_databases().await; } @@ -473,19 +475,19 @@ async fn a_changed_checksum_a_newer_catalog_and_an_older_release_are_reported() doc! { "_id": id, "name": "later", "kind": kind, "checksum": "c", "outcome": "applied", "notes": "", "applied_at": bson::DateTime::now(), "applied_by": "9.9.9" } }; - ledger.insert_one(row(6, "breaking")).await.expect("insert"); + ledger.insert_one(row(7, "breaking")).await.expect("insert"); let status = d.status().await; assert_eq!(status.label(), "AheadIncompatible"); assert_eq!(status.exit_code(), exit::AHEAD_INCOMPATIBLE); ledger - .delete_one(doc! { "_id": 6_i64 }) + .delete_one(doc! { "_id": 7_i64 }) .await .expect("delete"); - ledger.insert_one(row(6, "additive")).await.expect("insert"); + ledger.insert_one(row(7, "additive")).await.expect("insert"); assert_eq!(d.status().await.label(), "AheadCompatible"); assert!(d.migrate().await.expect("migrate").applied.is_empty()); ledger - .delete_one(doc! { "_id": 6_i64 }) + .delete_one(doc! { "_id": 7_i64 }) .await .expect("delete"); @@ -500,7 +502,7 @@ async fn a_changed_checksum_a_newer_catalog_and_an_older_release_are_reported() assert_eq!(refused.exit_code(), exit::ATTENTION); d.reconverge().await.expect("reconverge"); assert_eq!(d.setting("resource_tag_migration").await, None); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); settings .update_one( doc! { "_id": "catalog_version" }, @@ -510,7 +512,7 @@ async fn a_changed_checksum_a_newer_catalog_and_an_older_release_are_reported() .expect("version"); assert_eq!(d.status().await.label(), "WrittenByOlderRelease"); d.reconverge().await.expect("reconverge"); - assert_eq!(d.setting("catalog_version").await.as_deref(), Some("5")); + assert_eq!(d.setting("catalog_version").await.as_deref(), Some("6")); // A recorded checksum this binary's body does not have. ledger @@ -558,6 +560,6 @@ async fn a_ledger_document_this_binary_cannot_read_stops_the_status() { .expect_err("an error"); migrations.close().await; assert!(error.to_string().contains("\"sideways\""), "{error}"); - assert_eq!(MigrationId::LATEST.number(), 5); + assert_eq!(MigrationId::LATEST.number(), 6); d.drop_databases().await; } diff --git a/crates/storage-mongodb/tests/migrations_guards.rs b/crates/storage-mongodb/tests/migrations_guards.rs index 793f65361..55284c96f 100644 --- a/crates/storage-mongodb/tests/migrations_guards.rs +++ b/crates/storage-mongodb/tests/migrations_guards.rs @@ -207,7 +207,7 @@ async fn migrate_creates_the_indexes_a_deployment_lacks() { .await .contains(&"table_id_1".to_owned()) ); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); assert_fresh_definitions(&d).await; d.drop_databases().await; } @@ -232,7 +232,7 @@ async fn stop_at_every_request(uri: &str, state: Option<&str>) -> usize { } assert_eq!( d.status().await, - CatalogStatus::Current { version: 5 }, + CatalogStatus::Current { version: 6 }, "{requests}" ); assert_fresh_definitions(&d).await; diff --git a/crates/storage-mongodb/tests/migrations_holders.rs b/crates/storage-mongodb/tests/migrations_holders.rs index ed0cd866c..5eba917e5 100644 --- a/crates/storage-mongodb/tests/migrations_holders.rs +++ b/crates/storage-mongodb/tests/migrations_holders.rs @@ -113,7 +113,7 @@ async fn migration_2_keeps_the_lock_through_a_scan_longer_than_the_idle_limit() applied.expect("every migration applies"); // A shorter run would apply without the progress the scan records. assert!(took > timing.idle_limit * 4, "the run took {took:?}"); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); assert!(d.catalog_ledger().await.iter().any(|row| row.id == 2)); d.drop_databases().await; } @@ -151,7 +151,7 @@ async fn the_status_read_under_the_lock_keeps_it_through_a_long_scan() { "the status read took {took:?}" ); applied.expect("every migration applies"); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); d.drop_databases().await; } diff --git a/crates/storage-mongodb/tests/migrations_locks.rs b/crates/storage-mongodb/tests/migrations_locks.rs index 84727523a..3f249d638 100644 --- a/crates/storage-mongodb/tests/migrations_locks.rs +++ b/crates/storage-mongodb/tests/migrations_locks.rs @@ -53,10 +53,10 @@ async fn two_migrators_started_together_apply_once_and_the_other_finds_nothing() applied.sort_unstable(); assert_eq!( applied, - [0, 6], + [0, 7], "one applies every step and the other finds nothing" ); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); d.drop_databases().await; } @@ -204,7 +204,7 @@ async fn migrate_waits_for_a_running_server_and_a_server_refuses_while_a_migrati .is_none() ); // The preview and verify read the status beside the server. - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); serving.release().await; let (held, _) = lock(&d.target, false, Duration::from_secs(2), Timing::PRODUCTION) .await diff --git a/crates/storage-mongodb/tests/migrations_released.rs b/crates/storage-mongodb/tests/migrations_released.rs index c636de7a9..19db70d96 100644 --- a/crates/storage-mongodb/tests/migrations_released.rs +++ b/crates/storage-mongodb/tests/migrations_released.rs @@ -164,12 +164,12 @@ async fn every_released_catalog_upgrades_to_the_definitions_of_a_fresh_install() panic!("{state} is not behind"); }; assert_eq!(version, None, "{state} has no ledger"); - assert_eq!(pending.len(), 6, "{state}"); + assert_eq!(pending.len(), 7, "{state}"); let report = d.migrate().await.expect("migrate"); - assert_eq!((report.from, report.to), (None, Some(5)), "{state}"); + assert_eq!((report.from, report.to), (None, Some(6)), "{state}"); assert_eq!( d.status().await, - CatalogStatus::Current { version: 5 }, + CatalogStatus::Current { version: 6 }, "{state}" ); @@ -401,7 +401,7 @@ async fn a_reconverge_leaves_policies_written_after_migration_2_as_they_are() { .any(|n| n.starts_with("Skipped 2 resource_tag_condition_keys:")), "{report:?}" ); - assert_eq!(d.status().await, CatalogStatus::Current { version: 5 }); + assert_eq!(d.status().await, CatalogStatus::Current { version: 6 }); assert_eq!( documents(Deployment::documents(&d.catalog(), "iam_policies").await), before diff --git a/crates/storage-postgres/src/backup_engine.rs b/crates/storage-postgres/src/backup_engine.rs index 6d93da862..3303c31de 100755 --- a/crates/storage-postgres/src/backup_engine.rs +++ b/crates/storage-postgres/src/backup_engine.rs @@ -17,8 +17,7 @@ use extenddb_storage::backup_definition::{ }; use extenddb_storage::error::StorageError; use extenddb_storage::util::{ - SortKeyValue, backup_arn, backup_id_epoch_millis, composite_pk_to_text, new_backup_id, - parse_sk, sk_column_n, + SortKeyValue, backup_arn, backup_id_epoch_millis, new_backup_id, sk_column_n, }; use extenddb_storage::{BackupEngine, BackupRestoreSource, RestoreRequest}; use futures::future::BoxFuture; @@ -26,9 +25,10 @@ use tokio::sync::{OwnedSemaphorePermit, Semaphore}; use crate::PostgresEngine; use crate::create_table::RestoreRegistration; +use crate::data::key_text::{composite_pk_to_text, parse_sk}; use crate::data::{ all_sort_key_info, data_table_name, index_table_name, insert_index_rows_multi, - item_has_index_keys, project_item_for_index, + item_has_index_keys, json_to_item, project_item_for_index, }; /// Convert a `PostgreSQL` `TIMESTAMPTZ` to epoch seconds as `f64`, keeping the @@ -530,8 +530,14 @@ impl PostgresEngine { let mut batch_bytes: usize = 0; let mut item_count: i64 = 0; while let Some(item_json) = rows.try_next().await.map_err(db_err)? { - let item: Item = serde_json::from_str(&item_json) - .map_err(|e| StorageError::Internal(format!("Parse backup item: {e}")))?; + // The stored text carries the control-character escape; it is + // copied into the target's `item_data` as it is, and decoded here + // so that the key columns and index rows derived from it are + // escaped exactly once by the same helpers the write path uses. + let item = json_to_item( + serde_json::from_str(&item_json) + .map_err(|e| StorageError::Internal(format!("Parse backup item: {e}")))?, + )?; // The parsed item is kept for its keys and index projections; the // budget counts its stored text, which bounds both forms. batch_bytes += item_json.len(); @@ -1220,8 +1226,10 @@ impl BackupEngine for PostgresEngine { let mut rows = sqlx::query_as::<_, (String, String)>(&select_sql).fetch(&mut *snapshot); while let Some((pk, data)) = rows.try_next().await.map_err(db_err)? { - let item: Item = serde_json::from_str(&data) - .map_err(|e| StorageError::Internal(format!("Parse item: {e}")))?; + let item = + json_to_item(serde_json::from_str(&data).map_err(|e| { + StorageError::Internal(format!("Parse item: {e}")) + })?)?; let item_bytes = extenddb_core::types::item_size_bytes(&item); drop(item); size += i64::try_from(item_bytes).unwrap_or(i64::MAX); diff --git a/crates/storage-postgres/src/data/data_engine.rs b/crates/storage-postgres/src/data/data_engine.rs index 88dd03794..de71c0f2b 100755 --- a/crates/storage-postgres/src/data/data_engine.rs +++ b/crates/storage-postgres/src/data/data_engine.rs @@ -7,10 +7,10 @@ use extenddb_core::expression::{Expr, ExpressionMaps, KeyCondition, UpdateAction}; use extenddb_core::types::{Item, TableKeyInfo}; use extenddb_storage::error::StorageError; -use extenddb_storage::util::pk_to_text; use extenddb_storage::{DataEngine, IdempotencyKey, StreamCapture, TransactGetOp, TransactWriteOp}; use futures::future::BoxFuture; +use super::key_text::pk_to_text; use crate::PostgresEngine; impl DataEngine for PostgresEngine { @@ -201,8 +201,10 @@ impl DataEngine for PostgresEngine { // ((hashtext(pk)::bigint & 2147483647) % total = seg), evaluated for // the candidate key. `key_info` is scan-shaped, so `key_schema[0]` is // the pk of whichever table (data or index) the scan targets, matching - // the `pk` column the predicate hashes. Bigint bitmask instead of - // abs() per CB-20 / SP-SCN-002 (hashtext() == i32::MIN). + // the `pk` column the predicate hashes, so the key text is the stored + // (escaped) form, as `pk_to_text` from `key_text` gives it. Bigint + // bitmask instead of abs() per CB-20 / SP-SCN-002 (hashtext() == + // i32::MIN). let pk_name = key_info.key_schema[0].attribute_name.clone(); let key = key.clone(); Box::pin(async move { diff --git a/crates/storage-postgres/src/data/delete_item.rs b/crates/storage-postgres/src/data/delete_item.rs index 9a6dfda93..63a29e7b3 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::gsi_build::fix_up_stale_write; use super::index::{enqueue_async_indexes, fetch_write_path_indexes, sync_indexes}; diff --git a/crates/storage-postgres/src/data/gsi_build.rs b/crates/storage-postgres/src/data/gsi_build.rs index 22ee5bf6f..2e84dd617 100644 --- a/crates/storage-postgres/src/data/gsi_build.rs +++ b/crates/storage-postgres/src/data/gsi_build.rs @@ -58,13 +58,14 @@ use extenddb_core::types::{ TableKeyInfo, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{composite_pk_to_text, parse_sk, sk_column_n}; +use extenddb_storage::util::sk_column_n; use sqlx::Row as _; use super::index::{ IndexMeta, db_error, delete_index_row_multi, fetch_indexes_for_table, insert_index_row_multi, project_item_for_index, }; +use super::key_text::{composite_pk_to_text, parse_sk}; use super::{all_sort_key_info, data_table_name, index_table_name, json_to_item}; use crate::gsi_queue::{GsiApplyContext, GsiIndexDef, PendingApplyContext, enqueue_gsi_pending}; diff --git a/crates/storage-postgres/src/data/index.rs b/crates/storage-postgres/src/data/index.rs index 09c5627d7..c8d48595b 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 @@ -571,9 +573,7 @@ pub(crate) async fn insert_index_rows_multi( for &(sk_name, sk_type) in base_sks { query = bind_index_sk(query, idx_table, item, sk_name, sk_type)?; } - query = query.bind( - serde_json::to_value(projected).map_err(|e| StorageError::Internal(e.to_string()))?, - ); + query = query.bind(item_to_json(projected)?); } query.execute(&mut **tx).await.map_err(db_error)?; 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..3e5bfceac --- /dev/null +++ b/crates/storage-postgres/src/data/key_text.rs @@ -0,0 +1,247 @@ +// 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(util::pk_to_text(val)?); + } + Ok(join_escaped_parts(&parts)) +} + +/// The stored text of a composite partition key whose parts, as +/// `extenddb_storage::util::pk_to_text` renders each HASH attribute, are +/// `raw_parts`: each part escaped, then netstring-joined. The netstring length +/// is that of the escaped part. +/// +/// [`composite_pk_to_text`] builds every composite key column through this +/// function, and catalog migration 6, `escape_control_chars`, rebuilds legacy +/// composite key columns through it ([`reescape_legacy_composite`]). +pub(crate) fn join_escaped_parts>(raw_parts: &[S]) -> String { + let parts: Vec = raw_parts + .iter() + .map(|p| escape_control(p.as_ref()).into_owned()) + .collect(); + encode_netstring_composite(&parts) +} + +/// Split a netstring-joined composite key (`:,` per part, `len` +/// the part's length in bytes) into its parts. `None` when `text` is not a +/// sequence of well-formed netstrings. +pub(crate) fn split_netstring_composite(text: &str) -> Option> { + let mut parts = Vec::new(); + let mut rest = text; + while !rest.is_empty() { + let colon = rest.find(':')?; + let digits = &rest[..colon]; + if digits.is_empty() || !digits.bytes().all(|b| b.is_ascii_digit()) { + return None; + } + let len: usize = digits.parse().ok()?; + let start = colon + 1; + let end = start.checked_add(len)?; + let part = rest.get(start..end)?; + if rest.as_bytes().get(end) != Some(&b',') { + return None; + } + parts.push(part); + rest = &rest[end + 1..]; + } + (!parts.is_empty()).then_some(parts) +} + +/// The stored form of a composite key column that a build without the escape +/// wrote: its netstring parts hold raw text, so each is escaped and the parts +/// are joined again with the new lengths, as [`composite_pk_to_text`] writes the +/// same key. Escaping the whole text instead would leave every length prefix +/// counting the raw part. `None` when `legacy` is not a netstring. +pub(crate) fn reescape_legacy_composite(legacy: &str) -> Option { + split_netstring_composite(legacy).map(|parts| join_escaped_parts(&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 netstring_split_inverts_the_join() { + let cases: [&[&str]; 5] = [ + &["a", "b"], + &["", "x"], + &["1:a,", ",:", "z"], + &["a\u{1}\u{2}", "\u{e9}t\u{e9}", "12"], + &["only"], + ]; + for parts in cases { + let owned: Vec = parts.iter().map(|p| (*p).to_owned()).collect(); + let joined = encode_netstring_composite(&owned); + assert_eq!( + split_netstring_composite(&joined).as_deref(), + Some(parts), + "{joined:?}" + ); + } + for bad in [ + "", + "a", + "1:ab,", + "2:a,", + "x:ab,", + "+1:a,", + "1:a", + "1:a,2", + "1:\u{e9},", + ] { + assert_eq!(split_netstring_composite(bad), None, "{bad:?}"); + } + } + + #[test] + fn a_legacy_composite_key_is_rebuilt_as_the_write_path_builds_it() { + // The case from the review: a raw U+0001 in the first of two HASH + // parts. Escaping the whole text would give 2:a\x01\x02,1:b, with a + // length prefix that no longer matches its part. + let legacy = "2:a\u{1},1:b,"; + let rebuilt = reescape_legacy_composite(legacy).expect("a netstring"); + assert_eq!(rebuilt, "3:a\u{1}\u{2},1:b,"); + + let mut item = Item::new(); + item.insert("h1".to_owned(), AttributeValue::S("a\u{1}".into())); + item.insert("h2".to_owned(), AttributeValue::S("b".into())); + assert_eq!( + composite_pk_to_text(&item, &[ks("h1"), ks("h2")]).unwrap(), + rebuilt + ); + + // Parts without U+0001 keep their text and their lengths. + assert_eq!( + reescape_legacy_composite("1:x,2:yz,").as_deref(), + Some("1:x,2:yz,") + ); + assert_eq!(reescape_legacy_composite("a\u{1}b"), None); + } + + #[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 d8056ea14..c8b2fdcf9 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. /// @@ -63,8 +64,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. @@ -137,6 +174,7 @@ mod ddl; mod delete_item; pub(crate) mod gsi_build; 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 2fbc3102b..1e06bfd72 100755 --- a/crates/storage-postgres/src/data/put_item.rs +++ b/crates/storage-postgres/src/data/put_item.rs @@ -7,13 +7,15 @@ 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::gsi_build::fix_up_stale_write; 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 { @@ -31,8 +33,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 a57b69bec..a9537deb9 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 be6de8a6f..293718995 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 d17d0c437..33e545333 100755 --- a/crates/storage-postgres/src/data/tx_helpers.rs +++ b/crates/storage-postgres/src/data/tx_helpers.rs @@ -12,10 +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::index::db_error; -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( @@ -106,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) { @@ -158,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) @@ -371,8 +370,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, created_at) \ diff --git a/crates/storage-postgres/src/data/update_item.rs b/crates/storage-postgres/src/data/update_item.rs index 759990b40..b9917b021 100755 --- a/crates/storage-postgres/src/data/update_item.rs +++ b/crates/storage-postgres/src/data/update_item.rs @@ -8,10 +8,13 @@ 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::gsi_build::fix_up_stale_write; 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}; @@ -221,8 +224,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 b48170e24..1a2f2bfa2 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 f068089ed..435736bd8 100644 --- a/crates/storage-postgres/src/gsi_queue.rs +++ b/crates/storage-postgres/src/gsi_queue.rs @@ -26,11 +26,11 @@ use std::sync::Arc; +use crate::data::key_text::composite_pk_to_text; use extenddb_core::types::{ AttributeDefinition, AttributeValue, 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; @@ -157,7 +157,7 @@ const PG_UNDEFINED_TABLE: &str = "SQLSTATE 42P01"; /// order. A local /// hash (not `std`'s, which is not guaranteed stable across builds) keeps the /// mapping fixed for a given key forever. -fn partition_for(base_pk_text: &str) -> i32 { +pub(crate) fn partition_for(base_pk_text: &str) -> i32 { let mut hash: u64 = 0xcbf2_9ce4_8422_2325; // FNV-1a offset basis for byte in base_pk_text.as_bytes() { hash ^= u64::from(*byte); @@ -166,6 +166,24 @@ fn partition_for(base_pk_text: &str) -> i32 { (hash % NUM_WORKERS) as i32 } +/// The `worker_partition` of a queued row: the partition of its base key, +/// which `new_item` carries for puts and updates and `old_item` for deletes +/// (the base key is immutable over an item's lifetime). The key text is the +/// stored (escaped) form, so a key holding U+0001 hashes differently from the +/// raw text a build without the escape hashed; catalog migration 6, +/// `escape_control_chars`, recomputes the column of such rows through this +/// function. +pub(crate) fn worker_partition( + old_item: Option<&Item>, + new_item: Option<&Item>, + context: &PendingApplyContext, +) -> Result { + Ok(match new_item.or(old_item) { + Some(item) => partition_for(&composite_pk_to_text(item, context.base_key_schema())?), + None => 0, + }) +} + /// Check if a `StorageError` is caused by an undefined table (SQLSTATE 42P01). /// /// This occurs when an index table is dropped (table deleted) while an async @@ -322,24 +340,10 @@ 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()))?; - - // Route all updates for a given base item to one worker. - // The base key is immutable over an item's lifetime; `new_item` carries it - // for puts/updates, `old_item` for deletes. - let worker_partition = match new_item.or(old_item) { - Some(item) => partition_for(&composite_pk_to_text(item, context.base_key_schema())?), - None => 0, - }; + 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)?; + let worker_partition = worker_partition(old_item, new_item, context)?; let delay_interval = jitter_delay_ms(delay_ms) as f64 / 1000.0; sqlx::query( @@ -584,16 +588,9 @@ fn decode_row( new_json: Option, ctx_json: serde_json::Value, ) -> Result<(Option, Option, PendingApplyContext), 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)?; Ok((old_item, new_item, context)) } diff --git a/crates/storage-postgres/src/metadata_engine.rs b/crates/storage-postgres/src/metadata_engine.rs index b2f958b3b..6ef499c04 100755 --- a/crates/storage-postgres/src/metadata_engine.rs +++ b/crates/storage-postgres/src/metadata_engine.rs @@ -292,7 +292,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( @@ -408,7 +410,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/backend.rs b/crates/storage-postgres/src/migrations/backend.rs index 2924fa23f..412945dbc 100644 --- a/crates/storage-postgres/src/migrations/backend.rs +++ b/crates/storage-postgres/src/migrations/backend.rs @@ -57,6 +57,10 @@ pub struct PostgresStores { lock: Option, /// The `application_name` of this runner's sessions. pub(crate) application: String, + /// Whether migration 6's `detect` looks for other sessions: in a run, not + /// in a status read, and not for a target that allows them + /// ([`PostgresTarget::with_other_sessions_allowed`]). + pub(crate) checks_sessions: bool, } impl PostgresStores { @@ -277,6 +281,7 @@ impl PostgresMigrations { data_watch, lock, application, + checks_sessions: !read_only && !target.other_sessions_allowed, }, catalog_name, data_name, diff --git a/crates/storage-postgres/src/migrations/m0006_escape_control_chars/guard.rs b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/guard.rs new file mode 100644 index 000000000..a7a956638 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/guard.rs @@ -0,0 +1,131 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! The sessions check of migration 6's `detect`: the run stops, before its +//! first write, while another session holds a database a body rewrites. +//! +//! The bodies rewrite rows in place and cannot tell a row a running server +//! writes from a legacy row: a server of an earlier release still writing +//! leaves rows in the earlier encoding behind it. `migrate --yes` waits for +//! the servers of releases that take the serving key, and the upgrade manual +//! says to stop every server first; this check finds a server of an earlier +//! release, or any other client, that was not stopped. The runner's own +//! sessions, the migration lock's session and read-only status sessions are +//! left out by `application_name`. A store with nothing to rewrite (no +//! table, no backup, no queued or stream row), such as the one `init` +//! creates, needs no check. Setting [`IGNORE_CONNECTIONS_ENV`] skips the +//! check, for a database whose only other sessions are idle connections a +//! connection pooler holds. +//! +//! The migrator connects as the application role, which is not a superuser. +//! For sessions of other roles PostgreSQL then shows `datname`, `usename` +//! and `application_name` but hides `backend_type` and `client_addr`, so a +//! filter on `backend_type` alone would miss an operator's `psql` or a server +//! running as another role. A session is counted when it has a role +//! (`usename`), which background workers such as autovacuum never have, and +//! its type is `client backend` or hidden. + +use extenddb_storage::migrations::MigrationError; +use sqlx::PgPool; + +use super::super::connect::{RUNNER_APPLICATION_PREFIX, STATUS_APPLICATION}; +use super::db_error; + +/// Skips the check when set, to any value. +pub(crate) const IGNORE_CONNECTIONS_ENV: &str = "EXTENDDB_MIGRATE_IGNORE_CONNECTIONS"; + +/// `application_name` of the migration lock's session (`locks.rs`). +const LOCK_APPLICATION: &str = "extenddb migrate"; + +/// Whether the public schema of the database `pool` reaches has `table` and +/// the table holds a row. +async fn has_row(pool: &PgPool, table: &str) -> Result { + let exists: bool = sqlx::query_scalar("SELECT to_regclass($1) IS NOT NULL") + .bind(format!("public.{table}")) + .fetch_one(pool) + .await + .map_err(|e| db_error(&format!("look for {table}"), &e))?; + if !exists { + return Ok(false); + } + sqlx::query_scalar(&format!("SELECT EXISTS (SELECT 1 FROM {table})")) + .fetch_one(pool) + .await + .map_err(|e| db_error(&format!("read {table}"), &e)) +} + +/// Whether the catalog body has rows to rewrite: `backup_items` holds one. +pub(crate) async fn catalog_has_rows(catalog: &PgPool) -> Result { + if std::env::var_os(IGNORE_CONNECTIONS_ENV).is_some() { + return Ok(false); + } + has_row(catalog, "backup_items").await +} + +/// Whether the data body may have rows to rewrite: the catalog lists a table, +/// or `gsi_pending` or `stream_records` holds a row. +pub(crate) async fn data_has_rows(catalog: &PgPool, data: &PgPool) -> Result { + if std::env::var_os(IGNORE_CONNECTIONS_ENV).is_some() { + return Ok(false); + } + Ok(has_row(catalog, "tables").await? + || has_row(data, "gsi_pending").await? + || has_row(data, "stream_records").await?) +} + +/// One item naming the other sessions on the database `pool` reaches, or none; +/// `which` names the database (`"data"` or `"catalog"`). +pub(crate) async fn other_sessions(pool: &PgPool, which: &str) -> Result, MigrationError> { + type Row = (String, Option, Option); + let (database, others): (String, Vec) = { + let database: String = sqlx::query_scalar("SELECT current_database()::text") + .fetch_one(pool) + .await + .map_err(|e| db_error("name the database", &e))?; + let others = sqlx::query_as( + "SELECT usename::text, application_name, client_addr::text \ + FROM pg_stat_activity \ + WHERE datname = current_database() \ + AND usename IS NOT NULL \ + AND coalesce(backend_type, 'client backend') = 'client backend' \ + AND pid <> pg_backend_pid() \ + AND NOT starts_with(coalesce(application_name, ''), $1) \ + AND coalesce(application_name, '') NOT IN ($2, $3)", + ) + .bind(RUNNER_APPLICATION_PREFIX) + .bind(LOCK_APPLICATION) + .bind(STATUS_APPLICATION) + .fetch_all(pool) + .await + .map_err(|e| db_error("look for other sessions", &e))?; + (database, others) + }; + if others.is_empty() { + return Ok(Vec::new()); + } + let mut described: Vec = others + .iter() + .map(|(user, app, addr)| { + format!( + "{user} from {} ({})", + addr.as_deref().unwrap_or("no client address shown"), + match app.as_deref() { + Some("") | None => "no application name", + Some(a) => a, + } + ) + }) + .collect(); + described.sort(); + described.dedup(); + Ok(vec![format!( + "{} other session(s) hold the {which} database {database}: {}. Migration 6 rewrites \ + stored strings and runs with every ExtendDB server stopped and every other session \ + on the database closed, because a server still writing would leave rows in the \ + earlier encoding. Stop them and run `extenddb migrate --yes` again. If the only \ + other sessions are idle connections a connection pooler holds, set \ + {IGNORE_CONNECTIONS_ENV}=1 to skip this check", + others.len(), + described.join("; ") + )]) +} diff --git a/crates/storage-postgres/src/migrations/m0006_escape_control_chars/keys.rs b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/keys.rs new file mode 100644 index 000000000..4e6658301 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/keys.rs @@ -0,0 +1,135 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Two rewrites of migration 6 that cannot be done by escaping a stored text +//! where it lies, because the text is derived from the key strings: +//! +//! - A partition key column of a key schema with more than one HASH attribute +//! holds the netstring join of the parts (`:,` each), and the +//! escape changes a part's length. [`rebuild_composite_rows`] splits each +//! legacy value, escapes the parts and joins them again with the new +//! lengths, through the function the write path uses +//! (`key_text::join_escaped_parts`). +//! - `gsi_pending.worker_partition` is a hash of the stored base key text +//! (`gsi_queue::partition_for`), so it changes with the key text. +//! [`recompute_worker_partitions`] recomputes it with the function the +//! enqueue uses (`gsi_queue::worker_partition`). +//! +//! Both run in the data body's transaction, after the table's other columns. + +use extenddb_storage::migrations::MigrationError; +use sqlx::PgConnection; + +use super::{db_error, rewrite_longest_first}; +use crate::data::key_text::reescape_legacy_composite; +use crate::gsi_queue::{PendingApplyContext, worker_partition}; + +/// The `backup_items` rows whose `pk` holds a composite key: the rows of a +/// backup whose table's key schema has more than one HASH attribute. A boolean +/// SQL expression over a `backup_items` row; `backups` is in the same catalog +/// database. +pub(crate) const BACKUP_ITEMS_COMPOSITE_ROWS: &str = "backup_arn IN (\ + SELECT b.backup_arn FROM backups b WHERE (\ + SELECT COUNT(*) FROM jsonb_array_elements(b.key_schema) e \ + WHERE e->>'KeyType' = 'HASH') > 1)"; + +/// Whether a catalog key schema (`[{"AttributeName", "KeyType"}, ...]`) has +/// more than one HASH attribute, so its partition key column holds a netstring +/// join. +pub(crate) fn has_composite_partition_key(key_schema: &serde_json::Value) -> bool { + key_schema.as_array().is_some_and(|elements| { + elements + .iter() + .filter(|e| e.get("KeyType").and_then(serde_json::Value::as_str) == Some("HASH")) + .count() + > 1 + }) +} + +/// Rebuild the composite key text in `col` of the rows that `rows` (a boolean +/// SQL expression over the row) selects and that hold U+0001, one row at a +/// time, longest legacy text first and in batches ([`rewrite_longest_first`]). +/// A SQL `replace()` over the whole text would leave each length prefix +/// counting the raw part (`2:a\x01,1:b,` would become `2:a\x01\x02,1:b,`, +/// where this release writes and looks up `3:a\x01\x02,1:b,`). A value that is +/// not a netstring fails the body, naming the table and column (`what`), and +/// the store's transaction rolls back. +pub(crate) async fn rebuild_composite_rows( + conn: &mut PgConnection, + quoted_table: &str, + col: &str, + rows: &str, + what: &str, +) -> Result { + rewrite_longest_first(conn, quoted_table, col, rows, what, |legacy| { + reescape_legacy_composite(legacy).ok_or_else(|| { + MigrationError::Store(format!( + "re-encode {what}: the composite key {legacy:?} is not a netstring" + )) + }) + }) + .await +} + +/// The partition a row queued with these columns gets from this release's +/// enqueue, or `None` when a column does not decode (a worker handles such a +/// row as a poison row, whatever partition holds it). +fn current_partition( + old_item: Option, + new_item: Option, + index_context: serde_json::Value, +) -> Option { + let old_item = old_item.map(crate::data::json_to_item).transpose().ok()?; + let new_item = new_item.map(crate::data::json_to_item).transpose().ok()?; + let context: PendingApplyContext = crate::data::stored_json_to(index_context).ok()?; + worker_partition(old_item.as_ref(), new_item.as_ref(), &context).ok() +} + +/// Recompute `worker_partition` of the queued rows whose item holds U+0001, +/// after their jsonb columns were re-encoded in the same transaction. +/// +/// A release without the escape hashed the raw base key; this release hashes +/// the escaped one, so for a key holding U+0001 a row queued before the +/// upgrade and the next write of the same item would be claimed by different +/// workers, in either order, which breaks the queue's one item, one worker +/// rule. Every row of one item moves to the same new partition and keeps its +/// `id` and `ready_at`, so the item's rows are still claimed in write order, +/// and a row enqueued after the upgrade gets a `ready_at` no earlier than +/// theirs from the partition clamp in `enqueue_gsi_pending`. +pub(crate) async fn recompute_worker_partitions(conn: &mut PgConnection) -> Result { + type Row = ( + i64, + i32, + Option, + Option, + serde_json::Value, + ); + // '\u0001' is how jsonb::text renders U+0001; chr(92) spells the backslash + // as in `escape_jsonb_rows`. The base key is in new_item, or in old_item for + // a delete, as `worker_partition` reads it. + let rows: Vec = sqlx::query_as( + "SELECT id, worker_partition, old_item, new_item, index_context FROM gsi_pending \ + WHERE position((chr(92) || 'u0001') IN COALESCE(new_item, old_item)::text) > 0 \ + ORDER BY id FOR UPDATE", + ) + .fetch_all(&mut *conn) + .await + .map_err(|e| db_error("find the gsi_pending rows to repartition", &e))?; + let mut changed: u64 = 0; + for (id, stored, old_item, new_item, index_context) in rows { + let Some(partition) = current_partition(old_item, new_item, index_context) else { + continue; + }; + if partition == stored { + continue; + } + sqlx::query("UPDATE gsi_pending SET worker_partition = $1 WHERE id = $2") + .bind(partition) + .bind(id) + .execute(&mut *conn) + .await + .map_err(|e| db_error(&format!("repartition gsi_pending row {id}"), &e))?; + changed += 1; + } + Ok(changed) +} diff --git a/crates/storage-postgres/src/migrations/m0006_escape_control_chars/mod.rs b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/mod.rs new file mode 100644 index 000000000..5d7e5c5f5 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/mod.rs @@ -0,0 +1,578 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Migration 6, `escape_control_chars`, PostgreSQL: re-encode the rows an +//! earlier release wrote, whose strings are stored raw. +//! +//! This release stores every string through the order-preserving escape in +//! `extenddb_storage::util::escape_control` (U+0000 as U+0001 U+0001, U+0001 +//! as U+0001 U+0002), so that U+0000, which DynamoDB accepts anywhere a string +//! appears, can be held in `TEXT` and `jsonb` columns. A row an earlier +//! release wrote holds its strings raw, and a raw U+0001 followed by U+0001 or +//! U+0002 reads back differently under the decoder (a lone U+0001 reads back +//! unchanged). Earlier releases could not store U+0000, so for their rows the +//! escape of the stored text is the re-encode, and a row without U+0001 +//! encodes to itself and is left alone. +//! +//! Two bodies, one per database. The catalog body rewrites `backup_items`, +//! which lives in the catalog database; the data body rewrites every data +//! table, index table and vector index table (the text key columns and +//! `item_data`), `gsi_pending` and `stream_records`, and recomputes +//! `gsi_pending.worker_partition` for the rows it changes. Tables are listed +//! from the catalog rows; a row whose table is `DELETING`, or whose physical +//! table is missing, is skipped and noted. +//! +//! The escape is not idempotent: an escaped row also holds U+0001, and a +//! legacy U+0001 U+0002 cannot be told from an escaped U+0001. Each body runs +//! in the runner's transaction and commits with its ledger row, so a run that +//! stops leaves the store as it was, and a store whose row is recorded is +//! never rewritten again; the migration is not convergent. A server of this +//! release refuses a store where the migration is pending, so the bodies see +//! only rows earlier releases wrote, provided those servers are stopped: +//! each body's `detect` stops a run, before the run's first write, while +//! another session holds the database the body rewrites (see [`guard`]). A +//! status read does not look for sessions, so the preview runs beside a +//! server as for every other migration. +//! +//! The key text and document rewrites call the write path's escape and key +//! text functions (`escape_control`, `escape_json_strings`, +//! `key_text::reescape_legacy_composite`, `gsi_queue::worker_partition`), so +//! a migrated row is the row this release writes for the same item. +//! +//! This module's checksum covers `mod.rs`, `keys.rs`, `order.rs` and +//! `guard.rs`, and is pinned once the migration is released. They are +//! excluded from rustfmt and clippy so neither can require an edit. + +mod guard; +mod keys; +mod order; + +use async_trait::async_trait; +use extenddb_storage::migrations::{CodeBody, Done, MigrationError}; +use sqlx::{PgConnection, PgPool}; + +use super::PostgresStores; + +pub(crate) use guard::other_sessions; +#[cfg(test)] +pub(crate) use guard::{IGNORE_CONNECTIONS_ENV, catalog_has_rows, data_has_rows}; +pub(crate) use keys::{BACKUP_ITEMS_COMPOSITE_ROWS, has_composite_partition_key}; +pub(crate) use order::{BATCH, rewrite_longest_first, with_detail}; + +/// The text of every file of this module, whose SHA-256 the checksum test pins. +pub(crate) const SOURCE: &str = concat!( + include_str!("mod.rs"), + include_str!("keys.rs"), + include_str!("order.rs"), + include_str!("guard.rs"), +); + +/// Recorded as this migration's checksum. Changes only with a new migration. +const FINGERPRINT: &str = "0006_escape_control_chars/postgres/1"; + +/// The catalog body: `backup_items`. +pub(crate) struct EscapeBackupItems; + +/// The data body: the data, index and vector index tables, `gsi_pending` and +/// `stream_records`. +pub(crate) struct EscapeDataRows; + +#[async_trait] +impl CodeBody for EscapeBackupItems { + fn fingerprint(&self) -> &'static str { + FINGERPRINT + } + + /// Another session on the catalog database, while `backup_items` holds a + /// row. + async fn detect(&self, stores: &PostgresStores) -> Result, MigrationError> { + if !stores.checks_sessions || !guard::catalog_has_rows(&stores.catalog).await? { + return Ok(Vec::new()); + } + other_sessions(&stores.catalog, "catalog").await + } + + async fn run(&self, stores: &PostgresStores) -> Result { + let mut conn = stores + .catalog + .acquire() + .await + .map_err(|e| MigrationError::store("reach the catalog database", e))?; + escape_catalog(&mut conn).await.map(|notes| Done { notes }) + } +} + +#[async_trait] +impl CodeBody for EscapeDataRows { + fn fingerprint(&self) -> &'static str { + FINGERPRINT + } + + /// Another session on the data database, while the catalog lists a table + /// or `gsi_pending` or `stream_records` holds a row. + async fn detect(&self, stores: &PostgresStores) -> Result, MigrationError> { + if !stores.checks_sessions || !guard::data_has_rows(&stores.catalog, &stores.data).await? { + return Ok(Vec::new()); + } + other_sessions(&stores.data, "data").await + } + + async fn run(&self, stores: &PostgresStores) -> Result { + let mut conn = stores + .data + .acquire() + .await + .map_err(|e| MigrationError::store("reach the data database", e))?; + escape_data(&stores.catalog, &mut conn).await.map(|notes| Done { notes }) + } +} + +/// A failed statement as a migration error: `what`, the database's message +/// with its DETAIL ([`with_detail`]), and the SQLSTATE, which the runner reads +/// to run the body again after a lock timeout. +pub(crate) fn db_error(what: &str, e: &sqlx::Error) -> MigrationError { + let message = format!("{what}: {}", with_detail(e)); + match e.as_database_error().and_then(|d| d.code()) { + Some(code) => MigrationError::Database { + message, + sqlstate: code.into_owned(), + }, + None => MigrationError::Store(message), + } +} + +/// One physical table's rewrite: its text key columns, 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, + /// `gsi_pending` only: recompute `worker_partition` after the re-encode + /// ([`keys::recompute_worker_partitions`]). + recompute_worker_partition: bool, + /// The DynamoDB table (and index) the physical table holds, for errors + /// and notes. + ddb_name: Option, +} + +impl EscapeWork { + /// `.`, with the DynamoDB name when known: how errors name a + /// column. + fn column(&self, col: &str) -> String { + match &self.ddb_name { + Some(ddb) => format!("{}.{col} ({ddb})", self.table), + None => format!("{}.{col}", self.table), + } + } + + /// `
`, with the DynamoDB name when known: how notes name a table. + fn label(&self) -> String { + match &self.ddb_name { + Some(ddb) => format!("{} ({ddb})", self.table), + None => self.table.clone(), + } + } +} + +/// A text key column and the rows of it that hold a composite partition key. +struct TextKeyColumn { + name: String, + /// SQL boolean expression over the row, true where the column holds the + /// netstring join of several HASH attributes. Those rows are rebuilt one + /// by one ([`keys::rebuild_composite_rows`]); the others hold one + /// attribute's text and get the single-statement SQL replace. + composite_rows: String, +} + +impl TextKeyColumn { + fn new(name: String, composite: bool) -> Self { + let composite_rows = if composite { "TRUE" } else { "FALSE" }.to_owned(); + Self { + name, + composite_rows, + } + } +} + +/// The columns of `table` in the public schema of type `data_type` among +/// `names` (or, for `None`, the key columns `pk`, `base_pk`, `sk*_s` and +/// `base_sk*_s`), in column order. +async fn columns( + conn: &mut PgConnection, + table: &str, + data_type: &str, + names: Option<&[&str]>, +) -> Result, MigrationError> { + let names: Option> = names.map(|n| n.iter().map(|c| (*c).to_owned()).collect()); + sqlx::query_scalar( + "SELECT column_name::text FROM information_schema.columns \ + WHERE table_schema = 'public' AND table_name = $1 AND data_type = $2 \ + AND CASE WHEN $3::text[] IS NULL \ + THEN column_name IN ('pk', 'base_pk') OR column_name::text ~ '^(base_)?sk[0-9]*_s$' \ + ELSE column_name::text = ANY($3) END \ + ORDER BY ordinal_position", + ) + .bind(table) + .bind(data_type) + .bind(names) + .fetch_all(&mut *conn) + .await + .map_err(|e| db_error(&format!("list the {data_type} columns of {table}"), &e)) +} + +/// The rewrite of a `_ddb_*` table, whose columns are read from +/// `information_schema.columns` so that exactly the sort key columns this +/// table has (`sk_s`, `sk_s`, `base_sk_s`, ...) are covered. `None` +/// when the physical table does not exist. `composite` names the columns +/// whose key schema has more than one HASH attribute, so they hold a +/// netstring join on every row. +async fn dynamic_table_work( + conn: &mut PgConnection, + table: String, + composite: &[&str], +) -> Result, MigrationError> { + let text_columns = columns(conn, &table, "text", None).await?; + let jsonb_columns = columns(conn, &table, "jsonb", Some(&["item_data"])).await?; + if text_columns.is_empty() && jsonb_columns.is_empty() { + return Ok(None); + } + let text_columns = text_columns + .into_iter() + .map(|name| { + let composite = composite.contains(&name.as_str()); + TextKeyColumn::new(name, composite) + }) + .collect(); + Ok(Some(EscapeWork { + table, + text_columns, + jsonb_columns, + recompute_worker_partition: false, + ddb_name: None, + })) +} + +/// The rewrite of a fixed-name table, with the candidate columns filtered +/// through `information_schema.columns`. `None` when the table is absent. +async fn fixed_table_work( + conn: &mut PgConnection, + table: &str, + text: &[&str], + jsonb: &[&str], +) -> Result, MigrationError> { + let text_columns = if text.is_empty() { + Vec::new() + } else { + columns(conn, table, "text", Some(text)).await? + }; + let jsonb_columns = columns(conn, table, "jsonb", Some(jsonb)).await?; + if text_columns.is_empty() && jsonb_columns.is_empty() { + return Ok(None); + } + Ok(Some(EscapeWork { + table: table.to_owned(), + text_columns: text_columns + .into_iter() + .map(|name| TextKeyColumn::new(name, false)) + .collect(), + jsonb_columns, + recompute_worker_partition: false, + ddb_name: None, + })) +} + +/// 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 hold the literal six-character text backslash-u-0001 both come +/// back equal and are skipped. +pub(crate) 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. +/// +/// The candidate ctids are read once, before any row of this column is +/// touched, so no row is visited twice; an update moves a row (a ctid names a +/// row version), always after its value was read. The values are then read +/// in batches of [`BATCH`] by ctid, so a large table never loads into memory. +async fn escape_jsonb_rows( + conn: &mut PgConnection, + quoted_table: &str, + col: &str, +) -> Result { + // '\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 replace() is not used here: it would corrupt 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 *conn) + .await + .map_err(|e| db_error(&format!("find the 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(BATCH) { + let rows: Vec<(String, serde_json::Value)> = sqlx::query_as(&fetch_batch) + .bind(batch.to_vec()) + .fetch_all(&mut *conn) + .await + .map_err(|e| db_error(&format!("read rows of {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 *conn) + .await + .map_err(|e| db_error(&format!("re-encode a row of {quoted_table}"), &e))?; + changed += 1; + } + } + } + Ok(changed) +} + +/// Rewrite one table; returns the number of row updates. +async fn escape_rows_in_table( + conn: &mut PgConnection, + work: &EscapeWork, +) -> Result { + let quoted = format!("\"{}\"", work.table); + let mut rewritten: u64 = 0; + + // Key text columns holding one attribute's text: TEXT has no escaping + // layer, so a single SQL statement per column is exact. NULL columns never + // match (position() is NULL). The statement can fail on the primary key + // when two legacy keys differ only by a U+0002 after a U+0001, depending + // on the order it meets them (`order`); it runs under a savepoint, and on + // a unique violation the column is rewritten again row by row, longest + // legacy text first. Rows holding a composite key are rebuilt part by + // part instead; the other columns of the row keep the replace(). + for col in &work.text_columns { + let TextKeyColumn { + name, + composite_rows, + } = col; + let what = work.column(name); + let single_rows = format!("NOT ({composite_rows})"); + let sql = format!( + "UPDATE {quoted} SET \"{name}\" = replace(\"{name}\", chr(1), chr(1) || chr(2)) \ + WHERE position(chr(1) IN \"{name}\") > 0 AND {single_rows}" + ); + sqlx::query("SAVEPOINT escape_column") + .execute(&mut *conn) + .await + .map_err(|e| db_error(&format!("savepoint before {what}"), &e))?; + match sqlx::query(&sql).execute(&mut *conn).await { + Ok(done) => { + sqlx::query("RELEASE SAVEPOINT escape_column") + .execute(&mut *conn) + .await + .map_err(|e| db_error(&format!("release the savepoint of {what}"), &e))?; + rewritten += done.rows_affected(); + } + Err(e) if crate::pg_util::is_unique_violation(&e) => { + sqlx::query("ROLLBACK TO SAVEPOINT escape_column") + .execute(&mut *conn) + .await + .map_err(|e| db_error(&format!("roll back {what}"), &e))?; + // Legacy text holds no U+0000, so the escape is the replace() + // above. + rewritten += rewrite_longest_first(conn, "ed, name, &single_rows, &what, |legacy| { + Ok(extenddb_storage::util::escape_control(legacy).into_owned()) + }) + .await?; + } + Err(e) => return Err(db_error(&format!("re-encode {what}"), &e)), + } + if composite_rows != "FALSE" { + rewritten += + keys::rebuild_composite_rows(conn, "ed, name, composite_rows, &what).await?; + } + } + + for col in &work.jsonb_columns { + rewritten += escape_jsonb_rows(conn, "ed, col).await?; + } + if work.recompute_worker_partition { + rewritten += keys::recompute_worker_partitions(conn).await?; + } + Ok(rewritten) +} + +/// Rewrite every table of `work` in order; one note per table with a change. +async fn escape_all( + conn: &mut PgConnection, + work: &[EscapeWork], + notes: &mut Vec, +) -> Result<(), MigrationError> { + for w in work { + let rewritten = escape_rows_in_table(conn, w).await?; + if rewritten > 0 { + notes.push(format!("{}: {rewritten} row update(s)", w.label())); + } + } + Ok(()) +} + +/// The catalog body on `conn`, the catalog database in the runner's +/// transaction. `backup_items.pk` is the data table's `pk` as the backup +/// copied it, so a backup of a table with several HASH attributes holds a +/// netstring join there. +pub(crate) async fn escape_catalog(conn: &mut PgConnection) -> Result, MigrationError> { + let mut notes = Vec::new(); + let mut work = Vec::new(); + if let Some(mut w) = fixed_table_work(conn, "backup_items", &["pk", "sk"], &["item_data"]).await? { + for col in &mut w.text_columns { + if col.name == "pk" { + col.composite_rows = BACKUP_ITEMS_COMPOSITE_ROWS.to_owned(); + } + } + work.push(w); + } + escape_all(conn, &work, &mut notes).await?; + Ok(notes) +} + +/// The data body on `conn`, the data database in the runner's transaction, +/// with the tables listed from `catalog`. +/// +/// A partition key column whose key schema has more than one HASH attribute +/// holds a netstring join (`key_text::composite_pk_to_text`): `pk` of a table +/// or index with such a key schema, and `base_pk` of the index and vector +/// tables of a table with one. +pub(crate) async fn escape_data( + catalog: &PgPool, + conn: &mut PgConnection, +) -> Result, MigrationError> { + let mut notes = Vec::new(); + let mut work: Vec = Vec::new(); + + type TableRow = (String, String, String, String, serde_json::Value); + let tables: Vec = sqlx::query_as( + "SELECT table_id, table_name, account_id, table_status, key_schema FROM tables \ + ORDER BY table_id COLLATE \"C\"", + ) + .fetch_all(catalog) + .await + .map_err(|e| db_error("list tables", &e))?; + let deleting: std::collections::HashSet<&str> = tables + .iter() + .filter(|(_, _, _, status, _)| status == "DELETING") + .map(|(id, _, _, _, _)| id.as_str()) + .collect(); + let composite_tables: std::collections::HashSet<&str> = tables + .iter() + .filter(|(_, _, _, _, ks)| has_composite_partition_key(ks)) + .map(|(id, _, _, _, _)| id.as_str()) + .collect(); + // How errors and notes name the DynamoDB table of a physical table. + let ddb_tables: std::collections::HashMap<&str, String> = tables + .iter() + .map(|(id, name, account, _, _)| (id.as_str(), format!("table {name:?} of account {account}"))) + .collect(); + let ddb_table = |table_id: &str| { + ddb_tables + .get(table_id) + .cloned() + .unwrap_or_else(|| format!("table id {table_id}")) + }; + let skip = |label: String, physical: &str, notes: &mut Vec, table_id: &str| { + if deleting.contains(table_id) { + notes.push(format!("{label}: skipped, the table is being deleted")); + } else { + notes.push(format!("{label}: skipped, its physical table {physical} does not exist")); + } + }; + + for (id, _, _, _, ks) in &tables { + let label = ddb_table(id); + let physical = format!("_ddb_{id}"); + let cols: &[&str] = if has_composite_partition_key(ks) { &["pk"] } else { &[] }; + match dynamic_table_work(conn, physical.clone(), cols).await? { + Some(mut w) if !deleting.contains(id.as_str()) => { + w.ddb_name = Some(label); + work.push(w); + } + _ => skip(label, &physical, &mut notes, id), + } + } + let indexes: Vec<(String, String, String, serde_json::Value)> = sqlx::query_as( + "SELECT index_id, table_id, index_name, key_schema FROM indexes \ + ORDER BY index_id COLLATE \"C\"", + ) + .fetch_all(catalog) + .await + .map_err(|e| db_error("list indexes", &e))?; + for (id, table_id, name, ks) in &indexes { + let label = format!("index {name:?} of {}", ddb_table(table_id)); + let physical = format!("_ddb_{id}"); + let mut cols: Vec<&str> = Vec::new(); + if has_composite_partition_key(ks) { + cols.push("pk"); + } + if composite_tables.contains(table_id.as_str()) { + cols.push("base_pk"); + } + match dynamic_table_work(conn, physical.clone(), &cols).await? { + Some(mut w) if !deleting.contains(table_id.as_str()) => { + w.ddb_name = Some(label); + work.push(w); + } + _ => skip(label, &physical, &mut notes, table_id), + } + } + let vectors: Vec<(String, String, String)> = sqlx::query_as( + "SELECT index_id, table_id, index_name FROM vector_indexes \ + ORDER BY index_id COLLATE \"C\"", + ) + .fetch_all(catalog) + .await + .map_err(|e| db_error("list vector indexes", &e))?; + for (id, table_id, name) in &vectors { + let label = format!("vector index {name:?} of {}", ddb_table(table_id)); + let physical = format!("_ddb_vec_{id}"); + let cols: &[&str] = if composite_tables.contains(table_id.as_str()) { + &["base_pk"] + } else { + &[] + }; + match dynamic_table_work(conn, physical.clone(), cols).await? { + Some(mut w) if !deleting.contains(table_id.as_str()) => { + w.ddb_name = Some(label); + work.push(w); + } + _ => skip(label, &physical, &mut notes, table_id), + } + } + + // Fixed tables holding item documents. + if let Some(mut w) = + fixed_table_work(conn, "gsi_pending", &[], &["old_item", "new_item", "index_context"]) + .await? + { + w.recompute_worker_partition = true; + work.push(w); + } + if let Some(w) = fixed_table_work(conn, "stream_records", &[], &["record_data"]).await? { + work.push(w); + } + + escape_all(conn, &work, &mut notes).await?; + Ok(notes) +} diff --git a/crates/storage-postgres/src/migrations/m0006_escape_control_chars/order.rs b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/order.rs new file mode 100644 index 000000000..b8876f7fd --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_escape_control_chars/order.rs @@ -0,0 +1,119 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Row-by-row rewrites of one key column for migration 6, in an order that +//! keeps the table's primary key valid after every row. +//! +//! `PostgreSQL` checks a primary key as each row changes, not at the end of the +//! statement. Two legacy keys that differ only by a U+0002 after a U+0001, such +//! as `a\x01` and `a\x01\x02` with the same other key columns, escape to two +//! different texts, `a\x01\x02` and `a\x01\x02\x02`, but the first one's new +//! text is the second one's legacy text. If the shorter one is rewritten first, +//! the table holds `a\x01\x02` twice until the longer one is rewritten, and the +//! update fails with SQLSTATE 23505. A one-statement `UPDATE` meets the two rows +//! in heap order, so whether it fails depends on where they are stored. +//! +//! [`rewrite_longest_first`] rewrites the rows in descending byte length of +//! their legacy text, which never produces such a duplicate: +//! +//! - Every rewrite makes the text longer. Every row it rewrites holds a U+0001, +//! and the escape adds one byte for each; a rebuilt composite key also gives +//! each part its new length, which never takes fewer digits than the old one. +//! - Suppose rewriting row `r` gives the text that row `s` holds, and the two +//! rows agree on the other key columns. If `s` has not been rewritten yet, it +//! holds its legacy text, and that text is the new text of `r`, so it is +//! longer than the legacy text of `r` and `s` was rewritten earlier. If `s` +//! has been rewritten, the two rows have the same new text, and since the +//! escape is injective their legacy texts were equal. +//! - Each other key column holds the legacy texts of every row, or, when an +//! earlier pass of the same transaction rewrote it, their escapes. Either +//! way two rows agree on it now exactly when they agreed before the +//! migration. Two rows that agreed on every key column before the migration +//! would have broken the primary key, so no rewrite meets a duplicate. + +use extenddb_storage::migrations::MigrationError; +use sqlx::PgConnection; + +use super::db_error; + +/// Rows read from the cursor per round trip, and the batch `escape_jsonb_rows` +/// reads. +pub(crate) const BATCH: usize = 500; + +/// The cursor [`rewrite_longest_first`] declares and closes before returning. +const CURSOR: &str = "escape_0006_longest_first"; + +/// Rewrite `col` of the rows that hold U+0001 and that `rows` (a boolean SQL +/// expression over the row) selects, one row at a time, longest legacy text +/// first; `rewrite` gives a row's new text from its legacy text. Runs in the +/// caller's transaction and returns the number of rows changed. `what` names +/// the column in errors, with the DynamoDB table when the caller knows it. +/// +/// One cursor is declared over the matching rows in that order and read +/// [`BATCH`] rows at a time, so this process holds one batch, and `PostgreSQL` +/// sorts the matching rows in `work_mem`, spilling to temporary files beyond +/// it. The cursor's rows are fixed when it is declared and each is returned +/// once, so no row is rewritten twice. Each row is locked as the cursor +/// returns it, as the single-statement path locks the rows it updates. +pub(crate) async fn rewrite_longest_first( + conn: &mut PgConnection, + quoted_table: &str, + col: &str, + rows: &str, + what: &str, + rewrite: impl Fn(&str) -> Result, +) -> Result { + let failed = |step: &str, e: &sqlx::Error| db_error(&format!("{step} {what}"), e); + sqlx::query(&format!( + "DECLARE {CURSOR} NO SCROLL CURSOR FOR \ + SELECT ctid::text, \"{col}\" FROM {quoted_table} \ + WHERE position(chr(1) IN \"{col}\") > 0 AND ({rows}) \ + ORDER BY octet_length(\"{col}\") DESC FOR UPDATE" + )) + .execute(&mut *conn) + .await + .map_err(|e| failed("find the rows to rewrite in", &e))?; + let fetch = format!("FETCH {BATCH} FROM {CURSOR}"); + let update = format!("UPDATE {quoted_table} SET \"{col}\" = $1 WHERE ctid = $2::tid"); + let mut changed: u64 = 0; + loop { + let batch: Vec<(String, String)> = sqlx::query_as(&fetch) + .fetch_all(&mut *conn) + .await + .map_err(|e| failed("read the rows to rewrite in", &e))?; + for (ctid, legacy) in &batch { + let done = sqlx::query(&update) + .bind(rewrite(legacy)?) + .bind(ctid) + .execute(&mut *conn) + .await + .map_err(|e| failed("re-encode", &e))?; + changed += done.rows_affected(); + } + if batch.len() < BATCH { + break; + } + } + sqlx::query(&format!("CLOSE {CURSOR}")) + .execute(&mut *conn) + .await + .map_err(|e| failed("close the cursor over", &e))?; + Ok(changed) +} + +/// A database error as text, with `PostgreSQL`'s DETAIL when the error has one; +/// for a unique violation the DETAIL names the key already present. sqlx's +/// `Display` writes the message alone. The DETAIL is quoted with its control +/// characters escaped, since the keys this migration rewrites hold them. +pub(crate) fn with_detail(e: &sqlx::Error) -> String { + let detail = match e { + sqlx::Error::Database(db) => db + .try_downcast_ref::() + .and_then(sqlx::postgres::PgDatabaseError::detail), + _ => None, + }; + match detail { + Some(detail) => format!("{e} (DETAIL: {detail:?})"), + None => e.to_string(), + } +} diff --git a/crates/storage-postgres/src/migrations/m0006_guard_tests.rs b/crates/storage-postgres/src/migrations/m0006_guard_tests.rs new file mode 100644 index 000000000..2552ea163 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_guard_tests.rs @@ -0,0 +1,195 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Live checks of migration 6's sessions check (`other_sessions`, which both +//! bodies' `detect` call), with the scratch database and the skip convention +//! of `m0006_tests`. That the runner stops on a `detect` item before its first +//! write is checked over two databases in `tests/migrations_golden.rs`. + +use sqlx::PgPool; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + +use super::m0006_escape_control_chars::{ + IGNORE_CONNECTIONS_ENV, catalog_has_rows, data_has_rows, other_sessions, +}; +use super::m0006_tests::{base_conn, scratch, seed, skip, test_application}; + +/// `PgPool::close` returns before PostgreSQL removes the closed session from +/// `pg_stat_activity`, so a check right after a close can still count it. +/// Read the items again until there are none, for at most five seconds, and +/// return the last answer. +async fn after_close(pool: &PgPool, which: &str) -> Vec { + let mut items = other_sessions(pool, which) + .await + .expect("look for sessions"); + for _ in 0..50 { + if items.is_empty() { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + items = other_sessions(pool, which) + .await + .expect("look for sessions"); + } + items +} + +/// A pool of one unnamed connection to `database`, as a server or an +/// operator's `psql` would hold. +async fn unnamed(base: &str, database: &str) -> PgPool { + let pool = PgPoolOptions::new() + .max_connections(1) + .connect(&format!("{base}/{database}")) + .await + .expect("connect as another client would"); + sqlx::query("SELECT 1").execute(&pool).await.unwrap(); + pool +} + +/// A server (or any other session) still connected to the data database is +/// one item; the runner's own sessions are not. +#[tokio::test] +async fn another_session_on_the_data_database_is_reported() { + if base_conn().is_none() { + skip("another_session_on_the_data_database_is_reported"); + return; + } + let s = scratch().await; + seed(&s.db).await; + let base = base_conn().unwrap(); + assert_eq!(after_close(&s.db, "data").await, Vec::::new()); + + let server = unnamed(&base, &s.db_name).await; + let items = other_sessions(&s.db, "data").await.unwrap(); + assert_eq!(items.len(), 1, "{items:?}"); + assert!( + items[0].contains(&format!( + "1 other session(s) hold the data database {}", + s.db_name + )), + "{items:?}" + ); + assert!(items[0].contains("no application name"), "{items:?}"); + assert!(items[0].contains(IGNORE_CONNECTIONS_ENV), "{items:?}"); + + server.close().await; + assert_eq!(after_close(&s.db, "data").await, Vec::::new()); + s.cleanup().await; +} + +/// With the catalog in its own database, a session holding only that database +/// is reported by the catalog body's check, which names the catalog database, +/// and not by the data body's. +#[tokio::test] +async fn a_session_on_a_separate_catalog_is_reported_for_the_catalog() { + if base_conn().is_none() { + skip("a_session_on_a_separate_catalog_is_reported_for_the_catalog"); + return; + } + let data = scratch().await; + let catalog = scratch().await; + let base = base_conn().unwrap(); + + let client = unnamed(&base, &catalog.db_name).await; + assert_eq!(after_close(&data.db, "data").await, Vec::::new()); + let items = other_sessions(&catalog.db, "catalog").await.unwrap(); + assert_eq!(items.len(), 1, "{items:?}"); + assert!( + items[0].contains(&format!( + "other session(s) hold the catalog database {}", + catalog.db_name + )), + "{items:?}" + ); + + client.close().await; + assert_eq!( + after_close(&catalog.db, "catalog").await, + Vec::::new() + ); + data.cleanup().await; + catalog.cleanup().await; +} + +/// The migrator connects as the application role, which is not a superuser, +/// and `pg_stat_activity` hides `backend_type` of other roles' sessions from +/// it. A session of another role (an operator's `psql`, a server under a +/// different role) still counts. +#[tokio::test] +async fn another_roles_session_counts_from_a_non_superuser_pool() { + if base_conn().is_none() { + skip("another_roles_session_counts_from_a_non_superuser_pool"); + return; + } + let s = scratch().await; + let base = base_conn().unwrap(); + let role = format!( + "eddb_guard_{}", + &uuid::Uuid::new_v4().simple().to_string()[..12] + ); + sqlx::query(&format!( + "CREATE ROLE \"{role}\" LOGIN PASSWORD 'guard-test' NOSUPERUSER" + )) + .execute(&s.admin) + .await + .expect("create a plain login role"); + + // The migrator's view: the plain role, named as a runner. + let options = format!("{base}/{}", s.db_name) + .parse::() + .expect("scratch connection string") + .username(&role) + .password("guard-test") + .application_name(&test_application()); + let migrator = PgPoolOptions::new() + .max_connections(1) + .connect_with(options) + .await + .expect("connect as the plain role"); + // Only the scratch pool is connected, and it is named as a runner. + assert_eq!(after_close(&migrator, "data").await, Vec::::new()); + + // A superuser session with no application name. The plain role sees its + // datname and usename but not its backend_type. + let holder = unnamed(&base, &s.db_name).await; + let hidden: Option = sqlx::query_scalar( + "SELECT backend_type FROM pg_stat_activity \ + WHERE datname = current_database() AND usename <> current_user LIMIT 1", + ) + .fetch_one(&migrator) + .await + .expect("read another role's row"); + assert!( + hidden.is_none(), + "this test needs a role that cannot see other sessions' backend_type" + ); + let items = other_sessions(&migrator, "data").await.unwrap(); + assert_eq!(items.len(), 1, "{items:?}"); + assert!(items[0].contains("no client address shown"), "{items:?}"); + + holder.close().await; + assert_eq!(after_close(&migrator, "data").await, Vec::::new()); + migrator.close().await; + sqlx::query(&format!("DROP ROLE \"{role}\"")) + .execute(&s.admin) + .await + .expect("remove the test role"); + s.cleanup().await; +} + +/// A store with nothing to rewrite needs no check: a fresh store, such as the +/// one `init` creates, has no table, no backup and no queued or stream row. +#[tokio::test] +async fn only_a_store_with_rows_is_checked() { + if base_conn().is_none() { + skip("only_a_store_with_rows_is_checked"); + return; + } + let s = scratch().await; + assert!(!catalog_has_rows(&s.db).await.unwrap()); + assert!(!data_has_rows(&s.db, &s.db).await.unwrap()); + seed(&s.db).await; + assert!(catalog_has_rows(&s.db).await.unwrap()); + assert!(data_has_rows(&s.db, &s.db).await.unwrap()); + s.cleanup().await; +} diff --git a/crates/storage-postgres/src/migrations/m0006_keys_tests.rs b/crates/storage-postgres/src/migrations/m0006_keys_tests.rs new file mode 100644 index 000000000..9d5ae9812 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_keys_tests.rs @@ -0,0 +1,332 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Live checks of migration 6's composite key rebuild and `worker_partition` +//! recompute, with the scratch database and the skip convention of +//! `m0006_tests`: without `EXTENDDB_TEST_PG_CONNECTION_STRING` every live test +//! reports a skip and passes. Each seeds rows the way a release without the +//! escape wrote them. + +//! Live checks against a scratch database, with the helpers and the +//! skip convention of `migrations::live_reencode`: without +//! `EXTENDDB_TEST_PG_CONNECTION_STRING` every test reports a skip and +//! passes. Each seeds rows the way a build without the escape wrote them. + +use extenddb_core::types::{ + AttributeDefinition, AttributeValue, Item, KeySchemaElement, KeyType, Projection, + ProjectionType, ScalarAttributeType, +}; +use serde_json::json; +use sqlx::PgPool; + +use super::m0006_tests::{base_conn, run_pool, scratch, skip}; +use crate::data::key_text::composite_pk_to_text; +use crate::gsi_queue::{ + GsiApplyContext, GsiIndexDef, PendingApplyContext, enqueue_gsi_pending, partition_for, + worker_partition, +}; + +fn element(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } +} + +fn text(v: &str) -> AttributeValue { + AttributeValue::S(v.to_owned()) +} + +async fn exec(db: &PgPool, sql: &str, binds: &[serde_json::Value]) { + let mut q = sqlx::query(sql); + for b in binds { + q = q.bind(b.clone()); + } + q.execute(db).await.unwrap_or_else(|e| panic!("{sql}: {e}")); +} + +async fn texts(db: &PgPool, sql: &str) -> Vec { + sqlx::query_scalar(sql) + .fetch_all(db) + .await + .unwrap_or_else(|e| panic!("{sql}: {e}")) +} + +/// A table with two HASH attributes (`t2`), a two-HASH GSI on a one-HASH +/// table (`i2` on `t1`), a one-HASH GSI on `t2` (`i3`), and a backup of +/// each table, each holding a legacy row with a raw U+0001 in a HASH part. +async fn seed_composite(db: &PgPool) { + exec( + db, + "INSERT INTO accounts (account_id, account_name) VALUES ('111122223333', 'esc')", + &[], + ) + .await; + let one = json!([{"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}]); + let two = json!([{"AttributeName": "h1", "KeyType": "HASH"}, + {"AttributeName": "h2", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"}]); + for (id, ks) in [("t1", &one), ("t2", &two)] { + exec( + db, + &format!( + "INSERT INTO tables (account_id, table_name, key_schema, \ + attribute_definitions, table_status, table_arn, table_id) \ + VALUES ('111122223333', '{id}', $1, '[]', 'ACTIVE', \ + 'arn:aws:dynamodb:us-east-1:111122223333:table/{id}', '{id}')" + ), + std::slice::from_ref(ks), + ) + .await; + } + for (table, index, ks) in [ + ( + "t1", + "i2", + json!([{"AttributeName": "g1", "KeyType": "HASH"}, + {"AttributeName": "g2", "KeyType": "HASH"}]), + ), + ( + "t2", + "i3", + json!([{"AttributeName": "g", "KeyType": "HASH"}]), + ), + ] { + exec( + db, + &format!( + "INSERT INTO indexes (table_id, index_id, index_name, index_type, \ + key_schema, projection) VALUES ('{table}', '{index}', '{index}', 'GSI', \ + $1, '{{\"ProjectionType\": \"ALL\"}}')" + ), + &[ks], + ) + .await; + } + sqlx::raw_sql( + r#"CREATE TABLE "_ddb_t2" ( + pk TEXT COLLATE "C" NOT NULL, sk_s TEXT COLLATE "C", + item_data JSONB NOT NULL, PRIMARY KEY (pk, sk_s)); + CREATE TABLE "_ddb_i2" ( + pk TEXT COLLATE "C" NOT NULL, sk_s TEXT COLLATE "C", + base_pk TEXT COLLATE "C" NOT NULL, base_sk_s TEXT COLLATE "C", + item_data JSONB NOT NULL, PRIMARY KEY (pk, base_pk, base_sk_s)); + CREATE TABLE "_ddb_i3" ( + pk TEXT COLLATE "C" NOT NULL, sk_s TEXT COLLATE "C", + base_pk TEXT COLLATE "C" NOT NULL, base_sk_s TEXT COLLATE "C", + item_data JSONB NOT NULL, PRIMARY KEY (pk, base_pk, base_sk_s)); + INSERT INTO "_ddb_t2" VALUES + (E'2:a\x01,1:b,', 's', + '{"h1": {"S": "a\u0001"}, "h2": {"S": "b"}, "sk": {"S": "s"}}'), + ('1:x,1:y,', 's', '{"h1": {"S": "x"}, "h2": {"S": "y"}, "sk": {"S": "s"}}'); + INSERT INTO "_ddb_i2" VALUES + (E'2:g\x01,1:h,', NULL, E'p\x01', E's\x01', + '{"g1": {"S": "g\u0001"}, "g2": {"S": "h"}, "pk": {"S": "p\u0001"}}'); + INSERT INTO "_ddb_i3" VALUES + (E'g\x01', NULL, E'2:a\x01,1:b,', 's', + '{"g": {"S": "g\u0001"}, "h1": {"S": "a\u0001"}, "h2": {"S": "b"}}'); + INSERT INTO backups (backup_arn, backup_name, table_id, table_name, account_id, + key_schema, attribute_definitions) VALUES + ('b1', 'b1', 't1', 't1', '111122223333', + '[{"AttributeName": "pk", "KeyType": "HASH"}]', '[]'), + ('b2', 'b2', 't2', 't2', '111122223333', + '[{"AttributeName": "h1", "KeyType": "HASH"}, + {"AttributeName": "h2", "KeyType": "HASH"}]', '[]'); + INSERT INTO backup_items (backup_arn, pk, sk, item_data) VALUES + ('b1', E'p\x01', NULL, '{"pk": {"S": "p\u0001"}}'), + ('b2', E'2:a\x01,1:b,', NULL, '{"h1": {"S": "a\u0001"}, "h2": {"S": "b"}}');"#, + ) + .execute(db) + .await + .expect("seed the legacy rows"); +} + +#[tokio::test] +async fn rebuilds_composite_partition_keys_part_by_part() { + if base_conn().is_none() { + skip("rebuilds_composite_partition_keys_part_by_part"); + return; + } + let s = scratch().await; + seed_composite(&s.db).await; + run_pool(&s.db).await.expect("run the data code migrations"); + + // What this build writes and looks up for the legacy key + // h1 = "a\x01", h2 = "b": each part escaped, then joined. + let mut item = Item::new(); + item.insert("h1".to_owned(), text("a\u{1}")); + item.insert("h2".to_owned(), text("b")); + let ks = [element("h1", KeyType::Hash), element("h2", KeyType::Hash)]; + let expected = composite_pk_to_text(&item, &ks).expect("composite key text"); + assert_eq!(expected, "3:a\u{1}\u{2},1:b,"); + + // The table with two HASH attributes: the legacy row is rebuilt, the + // plain row keeps its text. + assert_eq!( + texts(&s.db, "SELECT pk FROM \"_ddb_t2\" ORDER BY pk").await, + vec!["1:x,1:y,".to_owned(), expected.clone()] + ); + // A two-HASH GSI: pk rebuilt, the one-HASH base key replaced. + let row: (String, String, String) = + sqlx::query_as("SELECT pk, base_pk, base_sk_s FROM \"_ddb_i2\"") + .fetch_one(&s.db) + .await + .expect("read _ddb_i2"); + assert_eq!( + row, + ( + "3:g\u{1}\u{2},1:h,".to_owned(), + "p\u{1}\u{2}".to_owned(), + "s\u{1}\u{2}".to_owned() + ) + ); + // A one-HASH GSI on the two-HASH table: pk replaced, base_pk rebuilt. + let row: (String, String) = sqlx::query_as("SELECT pk, base_pk FROM \"_ddb_i3\"") + .fetch_one(&s.db) + .await + .expect("read _ddb_i3"); + assert_eq!(row, ("g\u{1}\u{2}".to_owned(), expected.clone())); + // backup_items: the backup of the two-HASH table is rebuilt, the + // other replaced. + assert_eq!( + texts(&s.db, "SELECT pk FROM backup_items ORDER BY backup_arn").await, + vec!["p\u{1}\u{2}".to_owned(), expected] + ); + s.cleanup().await; +} + +fn context() -> PendingApplyContext { + PendingApplyContext::Gsi(GsiApplyContext { + base_key_schema: vec![element("pk", KeyType::Hash)], + attribute_definitions: vec![AttributeDefinition { + attribute_name: "pk".to_owned(), + attribute_type: ScalarAttributeType::S, + }], + index: GsiIndexDef { + index_id: "i1".to_owned(), + key_schema: vec![element("g", KeyType::Hash)], + projection: Projection { + projection_type: ProjectionType::All, + non_key_attributes: None, + }, + }, + }) +} + +/// The partition a build without the escape gave a row for `item`: the +/// hash of the raw key text. +fn legacy_partition(item: &Item) -> i32 { + let raw = extenddb_storage::util::composite_pk_to_text(item, &[element("pk", KeyType::Hash)]) + .expect("raw key text"); + partition_for(&raw) +} + +#[tokio::test] +async fn recomputes_the_worker_partition_of_rows_queued_before_the_upgrade() { + if base_conn().is_none() { + skip("recomputes_the_worker_partition_of_rows_queued_before_the_upgrade"); + return; + } + let ctx = context(); + // A key holding U+0001 whose raw and escaped texts hash to different + // workers, so the test sees the move. + let item = [ + "a\u{1}b", + "b\u{1}", + "\u{1}c", + "key\u{1}1", + "key\u{1}2", + "key\u{1}3", + ] + .into_iter() + .map(|k| { + let mut item = Item::new(); + item.insert("pk".to_owned(), text(k)); + item.insert("g".to_owned(), text("v")); + item + }) + .find(|item| { + legacy_partition(item) != worker_partition(None, Some(item), &ctx).expect("partition") + }) + .expect("a key whose partition moves"); + let current = worker_partition(None, Some(&item), &ctx).expect("partition"); + let legacy = legacy_partition(&item); + + // A key without U+0001 and a non-key string with it: same text both ways. + let mut plain = Item::new(); + plain.insert("pk".to_owned(), text("plain")); + plain.insert("note".to_owned(), text("n\u{1}")); + let plain_partition = legacy_partition(&plain); + + let s = scratch().await; + let raw_ctx = serde_json::to_value(&ctx).expect("context json"); + let raw = |i: &Item| serde_json::to_value(i).expect("item json"); + for (old_item, new_item, partition) in [ + (None, Some(raw(&item)), legacy), + (Some(raw(&item)), None, legacy), + (None, Some(raw(&plain)), plain_partition), + ] { + sqlx::query( + "INSERT INTO gsi_pending (table_id, worker_partition, old_item, new_item, \ + index_context) VALUES ('t1', $1, $2, $3, $4)", + ) + .bind(partition) + .bind(old_item) + .bind(new_item) + .bind(&raw_ctx) + .execute(&s.db) + .await + .expect("seed a queued row"); + } + + run_pool(&s.db).await.expect("run the data code migrations"); + + let partitions: Vec = + sqlx::query_scalar("SELECT worker_partition FROM gsi_pending ORDER BY id") + .fetch_all(&s.db) + .await + .expect("read gsi_pending"); + assert_eq!(partitions, vec![current, current, plain_partition]); + + // The next write of the item is queued on the same worker. + let mut tx = s.db.begin().await.expect("begin"); + enqueue_gsi_pending(&mut tx, "t1", None, Some(&item), 1, &ctx) + .await + .expect("enqueue the next write"); + tx.commit().await.expect("commit"); + let next: i32 = + sqlx::query_scalar("SELECT worker_partition FROM gsi_pending ORDER BY id DESC LIMIT 1") + .fetch_one(&s.db) + .await + .expect("read the new row"); + assert_eq!(next, current); + s.cleanup().await; +} + +mod unit { + use super::super::m0006_escape_control_chars::has_composite_partition_key; + use serde_json::json; + + #[test] + fn only_a_key_schema_with_two_or_more_hash_attributes_is_composite() { + let one = json!([ + {"AttributeName": "pk", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ]); + let two = json!([ + {"AttributeName": "h1", "KeyType": "HASH"}, + {"AttributeName": "h2", "KeyType": "HASH"}, + {"AttributeName": "sk", "KeyType": "RANGE"} + ]); + let ranges = json!([ + {"AttributeName": "h", "KeyType": "HASH"}, + {"AttributeName": "r1", "KeyType": "RANGE"}, + {"AttributeName": "r2", "KeyType": "RANGE"} + ]); + assert!(!has_composite_partition_key(&one)); + assert!(has_composite_partition_key(&two)); + assert!(!has_composite_partition_key(&ranges)); + assert!(!has_composite_partition_key(&json!({}))); + } +} diff --git a/crates/storage-postgres/src/migrations/m0006_order_tests.rs b/crates/storage-postgres/src/migrations/m0006_order_tests.rs new file mode 100644 index 000000000..e6c6691a4 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_order_tests.rs @@ -0,0 +1,389 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Live checks of migration 6's ordered key rewrite, with the scratch database +//! and the skip convention of `m0006_tests`. The tables are created by the DDL +//! `CreateTable` runs, so their primary keys are the production ones. Each +//! test stores legacy keys that differ only by a U+0002 after a U+0001, the +//! shorter one first in the heap, the order a one-statement rewrite fails on. +//! `item_data` is `{}`: the tests check the key columns. + +//! Live checks with the scratch database and the skip convention of +//! `migrations::live_reencode`. The tables are created by the DDL +//! `CreateTable` runs, so their primary keys are the production ones. Each +//! test stores legacy keys that differ only by a U+0002 after a U+0001, +//! the shorter one first in the heap, the order a one-statement rewrite +//! fails on. `item_data` is `{}`: the tests check the key columns. + +use extenddb_core::types::{AttributeDefinition, KeySchemaElement, KeyType, ScalarAttributeType}; +use extenddb_storage::util::escape_control; +use sqlx::PgPool; + +use super::m0006_escape_control_chars::{BATCH, rewrite_longest_first}; +use super::m0006_tests::{base_conn, run_pool, scratch, skip}; +use crate::PostgresEngine; +use crate::data::key_text::reescape_legacy_composite; + +fn key(name: &str, key_type: KeyType) -> KeySchemaElement { + KeySchemaElement { + attribute_name: name.to_owned(), + key_type, + } +} + +fn string_attributes(key_schema: &[KeySchemaElement]) -> Vec { + key_schema + .iter() + .map(|k| AttributeDefinition { + attribute_name: k.attribute_name.clone(), + attribute_type: ScalarAttributeType::S, + }) + .collect() +} + +fn json(value: &T) -> serde_json::Value { + serde_json::to_value(value).expect("json") +} + +/// The account, the `tables` row and the data table of table `id`. +async fn table(db: &PgPool, id: &str, key_schema: &[KeySchemaElement]) { + let attributes = string_attributes(key_schema); + sqlx::query( + "INSERT INTO accounts (account_id, account_name) VALUES ('111122223333', 'order') \ + ON CONFLICT DO NOTHING", + ) + .execute(db) + .await + .expect("seed the account"); + sqlx::query( + "INSERT INTO tables (account_id, table_name, key_schema, attribute_definitions, \ + table_status, table_arn, table_id) VALUES ('111122223333', $1, $2, $3, 'ACTIVE', \ + 'arn:aws:dynamodb:us-east-1:111122223333:table/' || $1, $1)", + ) + .bind(id) + .bind(json(&key_schema)) + .bind(json(&attributes)) + .execute(db) + .await + .expect("seed the tables row"); + let mut tx = db.begin().await.expect("begin"); + PostgresEngine::create_data_table(&mut tx, id, key_schema, &attributes) + .await + .expect("create the data table"); + tx.commit().await.expect("commit"); +} + +/// The `indexes` row and the data table of GSI `id` on table `table_id`. +async fn index( + db: &PgPool, + table_id: &str, + base_key_schema: &[KeySchemaElement], + id: &str, + key_schema: &[KeySchemaElement], +) { + sqlx::query( + "INSERT INTO indexes (table_id, index_id, index_name, index_type, key_schema, \ + projection) VALUES ($1, $2, $2, 'GSI', $3, '{\"ProjectionType\": \"ALL\"}')", + ) + .bind(table_id) + .bind(id) + .bind(json(&key_schema)) + .execute(db) + .await + .expect("seed the indexes row"); + let mut tx = db.begin().await.expect("begin"); + PostgresEngine::create_index_data_table( + &mut tx, + id, + key_schema, + &string_attributes(key_schema), + base_key_schema, + &string_attributes(base_key_schema), + ) + .await + .expect("create the index data table"); + tx.commit().await.expect("commit"); +} + +/// Insert rows into `_ddb_` one at a time, in the order given, which +/// is their heap order in a new table. +async fn insert(db: &PgPool, id: &str, columns: &[&str], rows: &[Vec]) { + let names: Vec = columns.iter().map(|c| format!("\"{c}\"")).collect(); + let params: Vec = (1..=columns.len()).map(|i| format!("${i}")).collect(); + let sql = format!( + "INSERT INTO \"_ddb_{id}\" ({}, item_data) VALUES ({}, '{{}}')", + names.join(", "), + params.join(", ") + ); + for row in rows { + let mut q = sqlx::query(&sql); + for value in row { + q = q.bind(value); + } + q.execute(db).await.unwrap_or_else(|e| panic!("{sql}: {e}")); + } +} + +fn row(values: &[&str]) -> Vec { + values.iter().map(|v| (*v).to_owned()).collect() +} + +/// Every value of `col` in `_ddb_`, sorted. +async fn column(db: &PgPool, id: &str, col: &str) -> Vec { + let sql = format!("SELECT \"{col}\" FROM \"_ddb_{id}\""); + let mut values: Vec = sqlx::query_scalar(&sql) + .fetch_all(db) + .await + .unwrap_or_else(|e| panic!("{sql}: {e}")); + values.sort(); + values +} + +/// The escape of each legacy text, sorted: what 004 must leave. +fn escaped>(legacy: &[S]) -> Vec { + let mut values: Vec = legacy + .iter() + .map(|t| escape_control(t.as_ref()).into_owned()) + .collect(); + values.sort(); + values +} + +fn rebuilt>(legacy: &[S]) -> Vec { + let mut values: Vec = legacy + .iter() + .map(|t| reescape_legacy_composite(t.as_ref()).expect("a netstring")) + .collect(); + values.sort(); + values +} + +async fn migrate(db: &PgPool) { + run_pool(db).await.expect("004 completes"); +} + +#[tokio::test] +async fn rewrites_a_partition_key_pair_longest_first() { + if base_conn().is_none() { + skip("rewrites_a_partition_key_pair_longest_first"); + return; + } + let s = scratch().await; + table( + &s.db, + "tp", + &[key("pk", KeyType::Hash), key("sk", KeyType::Range)], + ) + .await; + let legacy = ["a\u{1}", "a\u{1}\u{2}"]; + let rows: Vec> = legacy.iter().map(|pk| row(&[pk, "s"])).collect(); + insert(&s.db, "tp", &["pk", "sk_s"], &rows).await; + migrate(&s.db).await; + assert_eq!(column(&s.db, "tp", "pk").await, escaped(&legacy)); + s.cleanup().await; +} + +#[tokio::test] +async fn rewrites_a_sort_key_pair_longest_first() { + if base_conn().is_none() { + skip("rewrites_a_sort_key_pair_longest_first"); + return; + } + let s = scratch().await; + table( + &s.db, + "ts", + &[key("pk", KeyType::Hash), key("sk", KeyType::Range)], + ) + .await; + let legacy = ["s\u{1}", "s\u{1}\u{2}"]; + let rows: Vec> = legacy.iter().map(|sk| row(&["p", sk])).collect(); + insert(&s.db, "ts", &["pk", "sk_s"], &rows).await; + migrate(&s.db).await; + assert_eq!(column(&s.db, "ts", "sk_s").await, escaped(&legacy)); + assert_eq!(column(&s.db, "ts", "pk").await, vec!["p", "p"]); + s.cleanup().await; +} + +#[tokio::test] +async fn rebuilds_a_composite_partition_key_pair_longest_first() { + if base_conn().is_none() { + skip("rebuilds_a_composite_partition_key_pair_longest_first"); + return; + } + let s = scratch().await; + let ks = [ + key("h1", KeyType::Hash), + key("h2", KeyType::Hash), + key("sk", KeyType::Range), + ]; + table(&s.db, "tc", &ks).await; + // h1 = "a\x01" and h1 = "a\x01\x02", h2 = "b": the first one's rebuilt + // text is the second one's legacy text. + let legacy = ["2:a\u{1},1:b,", "3:a\u{1}\u{2},1:b,"]; + let rows: Vec> = legacy.iter().map(|pk| row(&[pk, "s"])).collect(); + insert(&s.db, "tc", &["pk", "sk_s"], &rows).await; + migrate(&s.db).await; + let expected = rebuilt(&legacy); + assert_eq!( + expected, + vec!["3:a\u{1}\u{2},1:b,", "4:a\u{1}\u{2}\u{2},1:b,"] + ); + assert_eq!(column(&s.db, "tc", "pk").await, expected); + s.cleanup().await; +} + +#[tokio::test] +async fn rewrites_an_index_base_key_pair_longest_first() { + if base_conn().is_none() { + skip("rewrites_an_index_base_key_pair_longest_first"); + return; + } + let s = scratch().await; + let base = [key("pk", KeyType::Hash), key("sk", KeyType::Range)]; + table(&s.db, "tb", &base).await; + index(&s.db, "tb", &base, "ib", &[key("g", KeyType::Hash)]).await; + // Two base items with the same index key and base sort key. The base + // table's own pair is the partition key case above; it is left empty + // here so the index table is the one that is checked. + let legacy = ["p\u{1}", "p\u{1}\u{2}"]; + let rows: Vec> = legacy.iter().map(|p| row(&["g", p, "s"])).collect(); + insert(&s.db, "ib", &["pk", "base_pk", "base_sk_s"], &rows).await; + migrate(&s.db).await; + assert_eq!(column(&s.db, "ib", "base_pk").await, escaped(&legacy)); + s.cleanup().await; +} + +/// More rows than one batch on both paths: a single-attribute column whose +/// one-statement rewrite fails, and a composite partition key column. The +/// pair of each table is three U+0001 long, so the rows between its two +/// lengths put the longer one in the first batch and the shorter one in the +/// second. Every final text is the escape of its legacy text; a row +/// rewritten twice would hold a longer one. +#[tokio::test] +async fn rewrites_each_row_once_across_batches() { + if base_conn().is_none() { + skip("rewrites_each_row_once_across_batches"); + return; + } + let fillers = 600; + assert!(fillers > BATCH); + let s = scratch().await; + + table( + &s.db, + "tf", + &[key("pk", KeyType::Hash), key("sk", KeyType::Range)], + ) + .await; + let short = "a\u{1}\u{1}\u{1}"; + let long = "a\u{1}\u{2}\u{1}\u{2}\u{1}\u{2}"; + assert_eq!(escape_control(short), long); + let mut legacy = vec![short.to_owned()]; + legacy.extend((0..fillers).map(|i| format!("k\u{1}{i:04}"))); + legacy.push(long.to_owned()); + let rows: Vec> = legacy.iter().map(|pk| row(&[pk, "s"])).collect(); + insert(&s.db, "tf", &["pk", "sk_s"], &rows).await; + + let ks = [ + key("h1", KeyType::Hash), + key("h2", KeyType::Hash), + key("sk", KeyType::Range), + ]; + table(&s.db, "tk", &ks).await; + let short = "4:a\u{1}\u{1}\u{1},1:b,"; + let long = "7:a\u{1}\u{2}\u{1}\u{2}\u{1}\u{2},1:b,"; + assert_eq!(reescape_legacy_composite(short).as_deref(), Some(long)); + let mut composite = vec![short.to_owned()]; + composite.extend((0..fillers).map(|i| format!("5:k\u{1}{i:03},1:b,"))); + composite.push(long.to_owned()); + let rows: Vec> = composite.iter().map(|pk| row(&[pk, "s"])).collect(); + insert(&s.db, "tk", &["pk", "sk_s"], &rows).await; + + migrate(&s.db).await; + assert_eq!(column(&s.db, "tf", "pk").await, escaped(&legacy)); + assert_eq!(column(&s.db, "tk", "pk").await, rebuilt(&composite)); + s.cleanup().await; +} + +/// A pair whose byte order is the reverse of its length order: the escape +/// of `a\x01b` is `a\x01\x02b`, which sorts below it, so an order by the +/// text itself, highest first, would rewrite the shorter key first. +#[tokio::test] +async fn rewrites_a_pair_whose_byte_order_is_shortest_first() { + if base_conn().is_none() { + skip("rewrites_a_pair_whose_byte_order_is_shortest_first"); + return; + } + let s = scratch().await; + table( + &s.db, + "tv", + &[key("pk", KeyType::Hash), key("sk", KeyType::Range)], + ) + .await; + let legacy = ["a\u{1}b", "a\u{1}\u{2}b"]; + let rows: Vec> = legacy.iter().map(|pk| row(&[pk, "s"])).collect(); + insert(&s.db, "tv", &["pk", "sk_s"], &rows).await; + migrate(&s.db).await; + assert_eq!(column(&s.db, "tv", "pk").await, escaped(&legacy)); + s.cleanup().await; +} + +/// An error from a rewrite names the DynamoDB table of the physical table: +/// here a two-HASH table whose `pk` is not a netstring. +#[tokio::test] +async fn an_error_names_the_dynamodb_table() { + if base_conn().is_none() { + skip("an_error_names_the_dynamodb_table"); + return; + } + let s = scratch().await; + let ks = [ + key("h1", KeyType::Hash), + key("h2", KeyType::Hash), + key("sk", KeyType::Range), + ]; + table(&s.db, "tn", &ks).await; + insert(&s.db, "tn", &["pk", "sk_s"], &[row(&["x\u{1}", "s"])]).await; + let err = run_pool(&s.db) + .await + .expect_err("a composite key that is not a netstring"); + let text = err.to_string(); + assert!( + text.contains("_ddb_tn.pk (table \"tn\" of account 111122223333)"), + "{text}" + ); + s.cleanup().await; +} + +#[tokio::test] +async fn a_duplicate_that_remains_is_reported_with_its_key() { + if base_conn().is_none() { + skip("a_duplicate_that_remains_is_reported_with_its_key"); + return; + } + let s = scratch().await; + sqlx::raw_sql( + r#"CREATE TABLE "dup" (pk TEXT COLLATE "C" PRIMARY KEY); + INSERT INTO "dup" VALUES (E'x\x01'), (E'y\x01');"#, + ) + .execute(&s.db) + .await + .expect("seed"); + let mut tx = s.db.begin().await.expect("begin"); + // A rewrite that gives every row one text: the second update fails. + let err = rewrite_longest_first(&mut tx, "\"dup\"", "pk", "TRUE", "dup.pk (t)", |_| { + Ok("z\u{1}".to_owned()) + }) + .await + .expect_err("a duplicate"); + let text = err.to_string(); + assert!(text.starts_with("re-encode dup.pk (t): "), "{text}"); + assert!( + text.contains("(DETAIL: \"Key (pk)=(z\\u{1}) already exists.\")"), + "{text}" + ); + tx.rollback().await.expect("rollback"); + s.cleanup().await; +} diff --git a/crates/storage-postgres/src/migrations/m0006_tests.rs b/crates/storage-postgres/src/migrations/m0006_tests.rs new file mode 100644 index 000000000..cecc77753 --- /dev/null +++ b/crates/storage-postgres/src/migrations/m0006_tests.rs @@ -0,0 +1,493 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Live checks of migration 6's rewrites against a scratch database. +//! +//! Each test needs `EXTENDDB_TEST_PG_CONNECTION_STRING` (a base URL with no +//! database name, for example `postgresql://postgres@127.0.0.1:5432`), builds +//! a throwaway database, and drops it when it passes. Without the variable +//! every test reports a skip and passes. One database serves as both catalog +//! and data, created by the file bodies of every migration in id order, as +//! `apply_all_migrations` creates the one-database layout; the bodies' +//! rewrites ([`escape_catalog`], [`escape_data`]) then run in one +//! transaction, as the runner runs each body in its store's transaction. The +//! runner, the ledger and the sessions check over two databases are covered +//! by `tests/migrations_golden.rs`. + +use extenddb_storage::migrations::{Body, MigrationError, MigrationId}; +use serde_json::json; +use sqlx::PgPool; +use sqlx::postgres::{PgConnectOptions, PgPoolOptions}; + +use super::connect::RUNNER_APPLICATION_PREFIX; +use super::m0006_escape_control_chars::{escape_catalog, escape_data}; + +/// The `application_name` of the scratch pools: a runner's, so the sessions +/// check of migration 6 treats them as the migration's own. +pub(crate) fn test_application() -> String { + format!("{RUNNER_APPLICATION_PREFIX}test") +} + +pub(crate) struct Scratch { + pub(crate) db: PgPool, + pub(crate) admin: PgPool, + pub(crate) db_name: String, +} + +impl Scratch { + pub(crate) 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; + } +} + +pub(crate) 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()) +} + +pub(crate) 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." + ); +} + +/// A scratch database holding the catalog and the data tables, with the file +/// bodies of every migration applied. +pub(crate) 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"); + let db = PgPoolOptions::new() + .max_connections(2) + .connect_with( + format!("{base}/{db_name}") + .parse::() + .expect("scratch connection string") + .application_name(&test_application()), + ) + .await + .expect("connect to the scratch database"); + for id in MigrationId::ALL { + for body in super::migration(id).bodies() { + if let Body::File { name, text } = &body.body { + sqlx::raw_sql(text) + .execute(&db) + .await + .unwrap_or_else(|e| panic!("{name}: {e}")); + } + } + } + Scratch { db, admin, db_name } +} + +/// Run the catalog body's and the data body's rewrites on the connection of +/// `one` (a pool of one connection to the scratch database), in one +/// transaction that commits when both succeed, reading the catalog through +/// `catalog`; the notes of both, catalog first. +pub(crate) async fn run_on(catalog: &PgPool, one: &PgPool) -> Result, MigrationError> { + let mut conn = one.acquire().await.expect("acquire the connection"); + sqlx::raw_sql("BEGIN") + .execute(&mut *conn) + .await + .expect("BEGIN"); + let result = async { + let mut notes = escape_catalog(&mut conn).await?; + notes.extend(escape_data(catalog, &mut conn).await?); + Ok(notes) + } + .await; + let end = if result.is_ok() { "COMMIT" } else { "ROLLBACK" }; + sqlx::raw_sql(end).execute(&mut *conn).await.expect(end); + result +} + +/// [`run_on`] through a connection of its own to the database `db` reaches. +pub(crate) async fn run_pool(db: &PgPool) -> Result, MigrationError> { + let one = PgPoolOptions::new() + .max_connections(1) + .connect_with((*db.connect_options()).clone()) + .await + .expect("connect a second session"); + let result = run_on(db, &one).await; + one.close().await; + result +} + +/// [`run_pool`] on the scratch database. +pub(crate) async fn run(s: &Scratch) -> Result, MigrationError> { + run_pool(&s.db).await +} + +/// 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. +pub(crate) 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() { + if base_conn().is_none() { + skip("reencodes_legacy_rows"); + return; + } + let s = scratch().await; + seed(&s.db).await; + + let notes = run(&s).await.expect("run migration 6"); + assert_eq!( + notes, + vec![ + "backup_items: 3 row update(s)", + "_ddb_t1 (table \"t\" of account 111122223333): 3 row update(s)", + "_ddb_i1 (index \"gsi1\" of table \"t\" of account 111122223333): 5 row update(s)", + "gsi_pending: 2 row update(s)", + "stream_records: 1 row update(s)", + ] + ); + + // 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"}}),), + ] + ); + + 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(1) + .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(&test_application()), + ) + .await + .expect("connect with standard_conforming_strings off"); + + run_on(&s.db, &scs_off).await.expect("run migration 6"); + + 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; +} diff --git a/crates/storage-postgres/src/migrations/mod.rs b/crates/storage-postgres/src/migrations/mod.rs index 740c3469a..c1bafcfa5 100644 --- a/crates/storage-postgres/src/migrations/mod.rs +++ b/crates/storage-postgres/src/migrations/mod.rs @@ -35,6 +35,9 @@ mod m0003_sort_key_definitions; #[rustfmt::skip] #[allow(clippy::all, deprecated)] mod m0004_index_table_base_key; +#[rustfmt::skip] +#[allow(clippy::all, deprecated)] +mod m0006_escape_control_chars; mod postcheck; mod registrar; mod rows; @@ -127,6 +130,24 @@ pub fn migration(id: MigrationId) -> Migration { "0005_drop_continuous_backups.sql" )], }, + MigrationId::EscapeControlChars => Migration::Applies { + kind: Kind::Breaking, + // The escape is not idempotent: running the bodies again would + // escape the strings written through the API since they ran. + convergent: false, + bodies: vec![ + StoreBody::code( + StoreKind::Catalog, + &m0006_escape_control_chars::EscapeBackupItems, + m0006_escape_control_chars::SOURCE, + ), + StoreBody::code( + StoreKind::Data, + &m0006_escape_control_chars::EscapeDataRows, + m0006_escape_control_chars::SOURCE, + ), + ], + }, } } @@ -148,7 +169,8 @@ pub(crate) const fn runs_in(id: MigrationId) -> RunsIn { MigrationId::Baseline | MigrationId::ResourceTagConditionKeys | MigrationId::SortKeyDefinitions - | MigrationId::DropContinuousBackups => RunsIn::Transaction, + | MigrationId::DropContinuousBackups + | MigrationId::EscapeControlChars => RunsIn::Transaction, MigrationId::IndexTableBaseKey => RunsIn::Autocommit, } } @@ -174,6 +196,9 @@ pub struct PostgresTarget { pub data_connection_row: Option, /// Timeout for each connection attempt. pub(crate) connect_timeout: Duration, + /// Whether a run leaves out migration 6's check for other sessions + /// ([`PostgresTarget::with_other_sessions_allowed`]). + pub(crate) other_sessions_allowed: bool, } impl PostgresTarget { @@ -187,9 +212,21 @@ impl PostgresTarget { data_database_given: false, data_connection_row: None, connect_timeout: Duration::from_secs(10), + other_sessions_allowed: false, } } + /// Leave out migration 6's check for other sessions on the databases it + /// rewrites, for a test harness whose own sessions hold its scratch + /// databases (lock holders, observers). Every path an operator runs keeps + /// the check, apart from `EXTENDDB_MIGRATE_IGNORE_CONNECTIONS`. + #[cfg(any(feature = "test-util", feature = "test-hooks"))] + #[must_use] + pub fn with_other_sessions_allowed(mut self) -> Self { + self.other_sessions_allowed = true; + self + } + /// Mark the data database as given on the command line. #[must_use] pub fn with_given_data_database(mut self) -> Self { @@ -351,8 +388,9 @@ pub async fn apply_all_migrations( return Ok(()); } let data_database = data_options.get_database().unwrap_or("postgres").to_owned(); + // The harness's own pools hold the scratch databases. migrate( - &PostgresTarget::new(catalog_options, data_database), + &PostgresTarget::new(catalog_options, data_database).with_other_sessions_allowed(), "test", Duration::ZERO, ) @@ -361,5 +399,13 @@ pub async fn apply_all_migrations( .map_err(|e| format!("{e:?}")) } +#[cfg(test)] +mod m0006_guard_tests; +#[cfg(test)] +mod m0006_keys_tests; +#[cfg(test)] +mod m0006_order_tests; +#[cfg(test)] +mod m0006_tests; #[cfg(test)] mod tests; diff --git a/crates/storage-postgres/src/snapshot_engine.rs b/crates/storage-postgres/src/snapshot_engine.rs index abb066cf3..35644bbf4 100644 --- a/crates/storage-postgres/src/snapshot_engine.rs +++ b/crates/storage-postgres/src/snapshot_engine.rs @@ -643,14 +643,11 @@ fn cursor_stream( } let mut items = Vec::with_capacity(rows.len()); for (data,) in rows { - match serde_json::from_str::(&data) { + match decode_row(&data) { Ok(item) => items.push(item), Err(e) => { state.abandon(); - return Some(( - Err(StorageError::Internal(format!("Parse item: {e}"))), - state, - )); + return Some((Err(e), state)); } } } @@ -660,6 +657,16 @@ fn cursor_stream( )) } +/// One cursor row, `item_data::text`, as the item it stores. `item_data` holds +/// every string through the escape of U+0000 and U+0001 (`data::item_to_json`), +/// so the row is decoded by `data::json_to_item`, as every other reader of the +/// column decodes it. +fn decode_row(data: &str) -> Result { + let value: serde_json::Value = serde_json::from_str(data) + .map_err(|e| StorageError::Internal(format!("Parse item: {e}")))?; + crate::data::json_to_item(value) +} + /// Two SQLSTATEs mean the server discarded the snapshot's rows before the /// stream drained them: 40001 (`serialization_failure`), which on a /// `REPEATABLE READ READ ONLY` transaction can only come from a hot standby @@ -679,3 +686,27 @@ pub fn map_fetch_error(e: &sqlx::Error) -> StorageError { } db_err(e) } + +#[cfg(test)] +mod tests { + use extenddb_core::types::{AttributeValue, Item}; + + use super::decode_row; + + #[test] + fn a_row_decodes_to_the_item_it_stores() { + let mut map = std::collections::BTreeMap::new(); + map.insert( + "m\u{1}".to_owned(), + AttributeValue::S("\u{1}\u{2}".to_owned()), + ); + let mut item = Item::new(); + item.insert("pk".to_owned(), AttributeValue::S("a\u{0}b".to_owned())); + item.insert("k\u{0}".to_owned(), AttributeValue::M(map)); + let stored = crate::data::item_to_json(&item) + .expect("stored form") + .to_string(); + assert!(!stored.contains("\\u0000"), "{stored}"); + assert_eq!(decode_row(&stored).expect("decode"), item); + } +} diff --git a/crates/storage-postgres/src/stream_engine.rs b/crates/storage-postgres/src/stream_engine.rs index da7fd122d..6a4bd4b65 100755 --- a/crates/storage-postgres/src/stream_engine.rs +++ b/crates/storage-postgres/src/stream_engine.rs @@ -165,8 +165,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", @@ -277,9 +276,7 @@ impl StreamEngine for PostgresEngine { let records: Vec = rows .into_iter() .filter_map(|(_, data)| data) - .map(|data| { - serde_json::from_value(data).map_err(|e| StorageError::Internal(e.to_string())) - }) + .map(crate::data::stored_json_to) .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-postgres/tests/backup_restore.rs b/crates/storage-postgres/tests/backup_restore.rs index 5c0a6f588..5d102fd94 100644 --- a/crates/storage-postgres/tests/backup_restore.rs +++ b/crates/storage-postgres/tests/backup_restore.rs @@ -2626,3 +2626,135 @@ async fn a_definition_read_that_loses_its_connection_is_transient() { assert!(recovered, "the pool reconnects after the failover"); s.cleanup().await; } + +/// Every stored row of one physical table, each as the JSON text of the +/// whole row, in a stable order. Two tables with the same stored form compare +/// equal; a difference in any column (an escaped key text, an `item_data` +/// document) shows up. +async fn stored_rows(s: &Scratch, physical_id: &str) -> Vec { + sqlx::query_scalar(&format!( + "SELECT to_jsonb(t)::text FROM \"_ddb_{physical_id}\" t ORDER BY 1" + )) + .fetch_all(&s.db) + .await + .expect("read stored rows") +} + +/// Items holding U+0000 and U+0001 in key values, index key values, nested +/// values, and attribute names are stored through the control-character +/// escape. A restore copies `backup_items` text into the target as it is and +/// derives the key columns and index rows from it, so those must come out in +/// the same stored form as the source's (escaped once, not twice), and the +/// restored items must read back equal to what was written. +#[tokio::test] +async fn restore_keeps_control_characters_in_keys_indexes_and_items() { + if base_conn().is_none() { + eprintln!("SKIP restore_keeps_control_characters_in_keys_indexes_and_items: no PostgreSQL"); + return; + } + let s = scratch().await; + let src = indexed_source_with(&s, "nul_src", 0).await; + let key_info = |table_name: &str, table_id: &str, desc: &TableDescription| TableKeyInfo { + table_name: table_name.to_owned(), + account_id: ACCOUNT.to_owned(), + table_id: table_id.to_owned(), + key_schema: desc.key_schema.clone(), + base_key_schema: desc.key_schema.clone(), + attribute_definitions: desc.attribute_definitions.clone(), + ..Default::default() + }; + let src_desc = s + .engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: "nul_src".to_owned(), + }, + ) + .await + .expect("describe the source"); + let src_key_info = key_info("nul_src", &src, &src_desc); + let items: Vec = [ + serde_json::json!({ + "pk": {"S": "p\u{0}a"}, "sk": {"N": "1"}, + "g": {"S": "g\u{0}\u{1}"}, "l": {"S": "l\u{1}\u{1}x"}, + "n\u{0}ame": {"S": "v\u{0}"}, + "m": {"M": {"k\u{1}": {"S": "\u{0}"}, "plain": {"N": "2"}}}, + "list": {"L": [{"S": "\u{1}\u{2}"}, {"SS": ["a\u{0}", "b"]}]} + }), + serde_json::json!({ + "pk": {"S": "p\u{0}a"}, "sk": {"N": "2"}, + "g": {"S": "g\u{0}\u{1}"}, "l": {"S": "l\u{0}"} + }), + serde_json::json!({"pk": {"S": "\u{1}"}, "sk": {"N": "3"}, "l": {"S": "\u{1}\u{2}"}}), + serde_json::json!({"pk": {"S": "plain"}, "sk": {"N": "4"}, "g": {"S": "g"}}), + ] + .into_iter() + .map(|v| serde_json::from_value(v).expect("item")) + .collect(); + for item in &items { + s.engine + .put_item( + &src_key_info, + item.clone(), + false, + None, + &extenddb_core::expression::ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + } + + let backup = s + .engine + .create_backup(ACCOUNT, "nul_src", "b") + .await + .expect("backup"); + let desc = restore(&s, "nul_dst", &backup.backup_arn) + .await + .expect("restore"); + assert_eq!(table_status(&s, "nul_dst").await.as_deref(), Some("ACTIVE")); + + // The base table and each index table hold the same stored rows as the + // source's, so every key column and document was escaped exactly once. + assert_eq!( + stored_rows(&s, &desc.table_id).await, + stored_rows(&s, &src).await + ); + let src_idx = index_ids(&s, &src).await; + let dst_idx = index_ids(&s, &desc.table_id).await; + assert_eq!(src_idx.len(), 2); + for ((name, a), (_, b)) in src_idx.iter().zip(&dst_idx) { + let want = stored_rows(&s, a).await; + assert!(!want.is_empty(), "index {name} has rows at the source"); + assert_eq!(stored_rows(&s, b).await, want, "index {name}"); + } + + // The restored items read back equal to what was written. + let dst_desc = s + .engine + .describe_table( + ACCOUNT, + DescribeTableInput { + table_name: "nul_dst".to_owned(), + }, + ) + .await + .expect("describe the restored table"); + let dst_key_info = key_info("nul_dst", &desc.table_id, &dst_desc); + for item in &items { + let key: Item = item + .iter() + .filter(|(k, _)| *k == "pk" || *k == "sk") + .map(|(k, v)| (k.clone(), v.clone())) + .collect(); + let got = s + .engine + .get_item(&dst_key_info, &key) + .await + .expect("point read"); + assert_eq!(got.as_ref(), Some(item)); + } + s.cleanup().await; +} diff --git a/crates/storage-postgres/tests/fixtures/golden/0006_escape_control_chars.txt b/crates/storage-postgres/tests/fixtures/golden/0006_escape_control_chars.txt new file mode 100644 index 000000000..fafe378b2 --- /dev/null +++ b/crates/storage-postgres/tests/fixtures/golden/0006_escape_control_chars.txt @@ -0,0 +1,66 @@ +_ddb_I:seed_tbl/gsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s000\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s003\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s006\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s009\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p1\u{1}\u{2}" sk=Some("s001\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p1\u{1}\u{2}" sk=Some("s004\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p1\u{1}\u{2}" sk=Some("s007\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p2\u{1}\u{2}" sk=Some("s002\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p2\u{1}\u{2}" sk=Some("s005\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 base pk="seed-p2\u{1}\u{2}" sk=Some("s008\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g0\u{1}\u{2}" sk=Some("s000\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g0\u{1}\u{2}" sk=Some("s002\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g0\u{1}\u{2}" sk=Some("s004\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g0\u{1}\u{2}" sk=Some("s006\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g0\u{1}\u{2}" sk=Some("s008\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g1\u{1}\u{2}" sk=Some("s001\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g1\u{1}\u{2}" sk=Some("s003\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g1\u{1}\u{2}" sk=Some("s005\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g1\u{1}\u{2}" sk=Some("s007\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/gsi1 pk="g1\u{1}\u{2}" sk=Some("s009\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s000\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s003\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s006\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p0\u{1}\u{2}" sk=Some("s009\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p1\u{1}\u{2}" sk=Some("s001\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p1\u{1}\u{2}" sk=Some("s004\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p1\u{1}\u{2}" sk=Some("s007\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p2\u{1}\u{2}" sk=Some("s002\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p2\u{1}\u{2}" sk=Some("s005\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 base pk="seed-p2\u{1}\u{2}" sk=Some("s008\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p1\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p1\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p1\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p2\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p2\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_I:seed_tbl/lsi1 pk="seed-p2\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p0\u{1}\u{2}" sk=Some("s000\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p0\u{1}\u{2}" sk=Some("s003\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p0\u{1}\u{2}" sk=Some("s006\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p0\u{1}\u{2}" sk=Some("s009\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p1\u{1}\u{2}" sk=Some("s001\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p1\u{1}\u{2}" sk=Some("s004\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p1\u{1}\u{2}" sk=Some("s007\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p2\u{1}\u{2}" sk=Some("s002\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p2\u{1}\u{2}" sk=Some("s005\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +_ddb_T:seed_tbl pk="seed-p2\u{1}\u{2}" sk=Some("s008\u{1}\u{2}") Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p0\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p1\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p1\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p1\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p2\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p2\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +backup_items pk="seed-p2\u{1}\u{2}" sk=None Some("note\u{1}\u{2}=v\u{1}\u{2}\u{2}") +catalog notes: +backup_items: 20 row update(s) +data notes: +_ddb_T:seed_tbl (table "seed_tbl" of account 123456789012): 30 row update(s) +_ddb_I:seed_tbl/lsi1 (index "lsi1" of table "seed_tbl" of account 123456789012): 40 row update(s) +_ddb_I:seed_tbl/gsi1 (index "gsi1" of table "seed_tbl" of account 123456789012): 50 row update(s) diff --git a/crates/storage-postgres/tests/migrations_golden.rs b/crates/storage-postgres/tests/migrations_golden.rs index 47d36bb25..73f6e9b21 100644 --- a/crates/storage-postgres/tests/migrations_golden.rs +++ b/crates/storage-postgres/tests/migrations_golden.rs @@ -12,7 +12,7 @@ use std::fmt::Write as _; use std::time::Duration; use extenddb_storage::migrations::{CatalogStatus, MigrationError, MigrationId, Runner}; -use extenddb_storage_postgres::migrations::{PostgresMigrations, migrate}; +use extenddb_storage_postgres::migrations::{PostgresMigrations, PostgresTarget, migrate}; use migrations_support::{ACCOUNT, LATEST, Scratch, base_conn, id_names, skip}; use sqlx::PgPool; @@ -422,3 +422,217 @@ async fn migration_4_adds_the_base_key_index_and_skips_what_it_cannot_reach() { catalog.close().await; data.close().await; } + +/// The physical tables of `seed_tbl` and their text key columns. +async fn seed_tbl_tables(catalog: &PgPool) -> Vec<(String, &'static [&'static str])> { + let table: String = + sqlx::query_scalar("SELECT table_id FROM tables WHERE table_name = 'seed_tbl'") + .fetch_one(catalog) + .await + .unwrap(); + let indexes: Vec = + sqlx::query_scalar("SELECT index_id FROM indexes WHERE table_id = $1 ORDER BY index_name") + .bind(&table) + .fetch_all(catalog) + .await + .unwrap(); + let mut tables: Vec<(String, &'static [&'static str])> = + vec![(format!("_ddb_{table}"), &["pk", "sk_s"])]; + for index in indexes { + tables.push(( + format!("_ddb_{index}"), + &["pk", "sk_s", "base_pk", "base_sk_s"], + )); + } + tables +} + +/// The attribute every legacy row gets: a name and a value holding U+0001 +/// as an earlier release stored it, the value followed by U+0002. +const LEGACY_NOTE: &str = "item_data || jsonb_build_object('note' || chr(1), \ + jsonb_build_object('S', 'v' || chr(1) || chr(2)))"; + +/// Each row of `table` as `
pk=.. sk=.. `, with the +/// strings in Rust's escaped form, so a control character is visible. +async fn rows_of(pool: &PgPool, table: &str, label: &str, key: [&str; 2]) -> Vec { + let rows: Vec<(String, Option, Option)> = sqlx::query_as(&format!( + "SELECT {}, {}, (SELECT string_agg(key || '=' || (value->>'S'), ',' ORDER BY key) \ + FROM jsonb_each(item_data) WHERE key LIKE 'note%') FROM \"{table}\"", + key[0], key[1] + )) + .fetch_all(pool) + .await + .unwrap(); + rows.into_iter() + .map(|(pk, sk, note)| format!("{label} pk={pk:?} sk={sk:?} {note:?}")) + .collect() +} + +/// Migration 6 on the 0.1.14 fixture whose strings an earlier release +/// stored with a raw U+0001: in the key columns of `seed_tbl`'s table, GSI +/// and LSI, in their documents, and in the backup's `backup_items`. +#[tokio::test] +async fn migration_6_escapes_the_strings_an_earlier_release_stored() { + let Some(base) = base_conn() else { + return skip("golden 6"); + }; + let s = Scratch::new(&base, "pgm_g6").await; + s.load("0.1.14").await; + s.migrate_through(MigrationId::DropContinuousBackups).await; + let catalog = s.catalog().await; + let data = s.data().await; + let tables = seed_tbl_tables(&catalog).await; + for (table, columns) in &tables { + let keys: Vec = columns + .iter() + .map(|c| format!("{c} = {c} || chr(1)")) + .collect(); + sqlx::query(&format!( + "UPDATE \"{table}\" SET {}, item_data = {LEGACY_NOTE}", + keys.join(", ") + )) + .execute(&data) + .await + .unwrap(); + } + sqlx::query(&format!( + "UPDATE backup_items SET pk = pk || chr(1), sk = sk || chr(1), item_data = {LEGACY_NOTE}" + )) + .execute(&catalog) + .await + .unwrap(); + // A status read does not look for other sessions: this test's own are + // open. + assert_eq!(s.status().await.label(), "Behind"); + catalog.close().await; + data.close().await; + s.migrate_through(MigrationId::EscapeControlChars).await; + assert_eq!(s.status().await, CatalogStatus::Current { version: LATEST }); + + let catalog = s.catalog().await; + let data = s.data().await; + let names = id_names(&catalog).await; + let mut lines = Vec::new(); + for (table, columns) in &tables { + let mut label = table.clone(); + for (id, name) in &names { + label = label.replace(id.as_str(), name); + } + let mut rows = rows_of(&data, table, &label, [columns[0], columns[1]]).await; + if columns.len() > 2 { + // An index row also names its base item. + rows = rows_of( + &data, + table, + &format!("{label} base"), + ["base_pk", "base_sk_s"], + ) + .await + .into_iter() + .chain(rows) + .collect(); + } + lines.extend(rows); + } + lines.extend(rows_of(&catalog, "backup_items", "backup_items", ["pk", "sk"]).await); + lines.sort(); + let mut out = lines.join("\n"); + for (pool, store) in [(&catalog, "catalog"), (&data, "data")] { + let mut note = notes(pool, 6).await; + for (id, name) in &names { + note = note.replace(id.as_str(), name); + } + write!(out, "\n{store} notes:\n{note}").unwrap(); + } + out.push('\n'); + golden("0006_escape_control_chars.txt", &out); + catalog.close().await; + data.close().await; +} + +/// The id-6 rows of the ledger of the database `pool` reaches. +async fn ledger_rows_of_6(pool: &PgPool) -> i64 { + sqlx::query_scalar("SELECT count(*) FROM extenddb_migrations WHERE id = 6") + .fetch_one(pool) + .await + .unwrap() +} + +/// While another session holds the data database of a store with tables, +/// `migrate --yes` stops with exit 5 before its first write, naming the +/// session, and migrates once the session is gone. +#[tokio::test] +async fn migration_6_stops_before_any_write_while_another_session_holds_the_data() { + let Some(base) = base_conn() else { + return skip("migration 6 sessions check"); + }; + let s = Scratch::new(&base, "pgm_g6s").await; + s.load("0.1.14").await; + s.migrate_through(MigrationId::DropContinuousBackups).await; + // A server of an earlier release, still connected. + let server = migrations_support::pool(&s.data_url(), 1).await; + sqlx::query("SELECT 1").execute(&server).await.unwrap(); + // The preview runs beside it. + assert_eq!(s.status().await.label(), "Behind"); + let err = migrate( + &s.target_checking_sessions(), + "test", + Duration::from_secs(5), + ) + .await + .expect_err("refused"); + assert_eq!(err.exit_code(), 5, "{err:?}"); + let text = format!("{err:?}"); + assert!( + text.contains(&format!( + "1 other session(s) hold the data database {}", + s.data_name + )), + "{text}" + ); + let catalog = s.catalog().await; + let data = s.data().await; + assert_eq!( + ledger_rows_of_6(&catalog).await, + 0, + "the refused run wrote the catalog ledger" + ); + assert_eq!( + ledger_rows_of_6(&data).await, + 0, + "the refused run wrote the data ledger" + ); + catalog.close().await; + data.close().await; + server.close().await; + // Connection strings that name an application, as an operator's can: the + // runner names its own sessions over them, so the check does not count + // them. + let catalog = s.catalog().await; + sqlx::query("UPDATE settings SET value = $1 WHERE key = 'data_database_connection_string'") + .bind(format!("{}?application_name=psql", s.data_url())) + .execute(&catalog) + .await + .unwrap(); + catalog.close().await; + let named = || { + PostgresTarget::new( + format!("{}?application_name=psql", s.catalog_url()) + .parse() + .unwrap(), + &s.data_name, + ) + }; + // PostgreSQL removes a closed session from pg_stat_activity a moment + // after the close returns. + let mut result = migrate(&named(), "test", Duration::from_secs(5)).await; + for _ in 0..50 { + if result.is_ok() { + break; + } + tokio::time::sleep(Duration::from_millis(100)).await; + result = migrate(&named(), "test", Duration::from_secs(5)).await; + } + result.expect("migrate once the session is gone"); + assert_eq!(s.status().await, CatalogStatus::Current { version: LATEST }); +} diff --git a/crates/storage-postgres/tests/migrations_support/mod.rs b/crates/storage-postgres/tests/migrations_support/mod.rs index 31646f90e..752ab2c85 100644 --- a/crates/storage-postgres/tests/migrations_support/mod.rs +++ b/crates/storage-postgres/tests/migrations_support/mod.rs @@ -214,7 +214,17 @@ impl Scratch { } /// The runner's target, which writes the data connection rows as `init`. + /// The tests hold sessions of their own on the scratch databases (lock + /// holders, observers), so the target leaves out migration 6's check for + /// other sessions; [`Scratch::target_checking_sessions`] keeps it. pub fn target(&self) -> PostgresTarget { + self.target_checking_sessions() + .with_other_sessions_allowed() + } + + /// [`Scratch::target`] with migration 6's check for other sessions, as + /// `extenddb migrate --yes` runs it. + pub fn target_checking_sessions(&self) -> PostgresTarget { PostgresTarget::new(self.options(), &self.data_name) .with_data_connection_row(self.data_url()) } diff --git a/crates/storage-postgres/tests/table_snapshot.rs b/crates/storage-postgres/tests/table_snapshot.rs index 7a985c88d..3a4cab390 100644 --- a/crates/storage-postgres/tests/table_snapshot.rs +++ b/crates/storage-postgres/tests/table_snapshot.rs @@ -521,6 +521,51 @@ async fn legacy_collation_table_streams_unordered_and_complete(shared: bool) { } on_both_layouts!(legacy_collation_table_streams_unordered_and_complete); +/// An item with U+0000 in an attribute name and U+0001 in a value is stored +/// through the escape (`data::item_to_json`), and the stream yields it as it +/// was written, not in its stored form. +#[tokio::test] +async fn an_item_with_control_characters_streams_as_written() { + use extenddb_core::expression::ExpressionMaps; + use extenddb_core::types::{AttributeValue, Item}; + use extenddb_storage::DataEngine; + use futures::StreamExt; + if base_conn().is_none() { + eprintln!("SKIP an_item_with_control_characters_streams_as_written: no PostgreSQL"); + return; + } + let s = scratch().await; + create_active(&s, plain_input("control")).await; + let item: Item = [ + ("pk".to_owned(), AttributeValue::S("p".to_owned())), + ( + "a\u{0}b".to_owned(), + AttributeValue::S("v\u{1}w".to_owned()), + ), + ] + .into_iter() + .collect(); + let ki = key_info(&s, "control").await; + s.engine + .put_item( + &ki, + item.clone(), + false, + None, + &ExpressionMaps::default(), + None, + ) + .await + .expect("put"); + let snapshot = s + .engine + .open_table_snapshot(ACCOUNT, "control") + .await + .expect("open"); + let streamed: Vec = snapshot.items.map(|r| r.expect("item")).collect().await; + assert_eq!(streamed, vec![item]); +} + /// The snapshot bound is half the pool, with a separate data database or one /// shared pool: five snapshots holding their data connections on a /// ten-connection pool, the sixth open refused at once with `LimitExceeded`, diff --git a/crates/storage-sqlite/src/migrations/mod.rs b/crates/storage-sqlite/src/migrations/mod.rs index 06647fbb4..245e28ccc 100644 --- a/crates/storage-sqlite/src/migrations/mod.rs +++ b/crates/storage-sqlite/src/migrations/mod.rs @@ -106,6 +106,10 @@ pub fn migration(id: MigrationId) -> Migration { "0005_drop_continuous_backups.sql" )], }, + MigrationId::EscapeControlChars => Migration::NotApplicable( + "SQLite stores U+0000 in strings as written, so its stored form of strings does not \ + change", + ), } } diff --git a/crates/storage-sqlite/src/migrations/released_tests.rs b/crates/storage-sqlite/src/migrations/released_tests.rs index 71dcee691..9e3959272 100644 --- a/crates/storage-sqlite/src/migrations/released_tests.rs +++ b/crates/storage-sqlite/src/migrations/released_tests.rs @@ -390,11 +390,14 @@ async fn a_catalog_an_older_release_wrote_to_is_reported_and_reconverged() { next_second().await; let report = reconverge(&db).await.expect("reconverge"); - // Every convergent migration ran again; migration 2, which is not - // convergent, is named as skipped. + // Every convergent migration that applies on SQLite ran again; + // migration 2, which is not convergent, is named as skipped. let convergent = MigrationId::ALL .into_iter() - .filter(|id| super::migration(*id).convergent()) + .filter(|id| { + let m = super::migration(*id); + m.convergent() && !m.bodies().is_empty() + }) .count(); assert_eq!(report.applied.len(), convergent, "{trace}: {report:?}"); let skipped = format!("Skipped {}:", MigrationId::ResourceTagConditionKeys); diff --git a/crates/storage-sqlite/src/migrations/tests.rs b/crates/storage-sqlite/src/migrations/tests.rs index a837e802e..564a2f6bb 100644 --- a/crates/storage-sqlite/src/migrations/tests.rs +++ b/crates/storage-sqlite/src/migrations/tests.rs @@ -270,7 +270,15 @@ async fn a_fresh_database_is_current_after_one_run_and_seeds_nothing() { "SELECT COUNT(*) FROM extenddb_migrations WHERE outcome = 'applied'" ) .await, - i64::from(LATEST) + // One row for each migration that applies on SQLite; the others are + // recorded as not applicable. + i64::try_from( + MigrationId::ALL + .into_iter() + .filter(|id| !super::migration(*id).bodies().is_empty()) + .count() + ) + .unwrap() ); assert_eq!(version_row(&db).await, Some(LATEST.to_string())); assert_eq!( diff --git a/crates/storage/src/migrations/registry.rs b/crates/storage/src/migrations/registry.rs index dc426fd18..fc70c947f 100644 --- a/crates/storage/src/migrations/registry.rs +++ b/crates/storage/src/migrations/registry.rs @@ -67,6 +67,10 @@ migration_ids! { /// it. DropContinuousBackups = 5, "drop_continuous_backups", "drop the continuous_backups table, which no request reads or writes since #335"; + /// Re-encodes the strings a PostgreSQL store holds from an earlier + /// release into the escape that stores U+0000 (#364). + EscapeControlChars = 6, "escape_control_chars", + "re-encode the strings PostgreSQL stores so that they can hold U+0000"; } /// Highest id that has shipped in a release. Bodies at or below it are frozen: diff --git a/crates/storage/src/util/control_chars.rs b/crates/storage/src/util/control_chars.rs new file mode 100644 index 000000000..e1d708d9c --- /dev/null +++ b/crates/storage/src/util/control_chars.rs @@ -0,0 +1,333 @@ +// 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; on PostgreSQL, catalog migration 6, +//! `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 8177f6e50..b72d1f202 100755 --- a/crates/storage/src/util/key.rs +++ b/crates/storage/src/util/key.rs @@ -21,9 +21,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], @@ -341,7 +342,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 bb79c2103..11414e696 100755 --- a/crates/storage/src/util/mod.rs +++ b/crates/storage/src/util/mod.rs @@ -9,6 +9,7 @@ mod arn; mod backup_id; mod catalog_version; +mod control_chars; mod key; mod shard_stream; @@ -23,6 +24,10 @@ pub use backup_id::{ is_backup_resource_arn, new_backup_id, parse_backup_arn, }; pub use catalog_version::check_stored_catalog_version; +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 6aad21e84..d8d59976b 100755 --- a/docs/design/04-component-storage.md +++ b/docs/design/04-component-storage.md @@ -406,6 +406,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 on PostgreSQL by catalog migration 6, + `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/getting-started.md b/docs/getting-started.md index 09d382c2f..8d4154351 100755 --- a/docs/getting-started.md +++ b/docs/getting-started.md @@ -269,7 +269,7 @@ Config: extenddb.toml Catalog database: extenddb_catalog --- Checking catalog migrations... - OK: Current: the catalog is at migration 5 + OK: Current: the catalog is at migration 6 --- Checking catalog connection... OK: Connected to catalog. --- Checking data database... @@ -287,7 +287,7 @@ extenddb runs as a daemon (background process) and logs to syslog. On startup it ```bash ./target/release/extenddb serve --config extenddb.toml -# extenddb 0.1.14 (catalog migration 5) starting on 127.0.0.1:18443 +# extenddb 0.1.14 (catalog migration 6) starting on 127.0.0.1:18443 # storage: postgres (postgresql://extenddb:***@localhost:5432/extenddb_catalog) ``` @@ -1366,7 +1366,7 @@ Each runner requires its tools to be installed. The runner checks prerequisites ```bash ./target/release/extenddb version # extenddb 0.1.14 -# catalog migration 5 (postgres) +# catalog migration 6 (postgres) # commit abc1234 # built 2026-04-17T12:00:00Z ``` diff --git a/docs/manuals/04-quickstart-setup-guide.md b/docs/manuals/04-quickstart-setup-guide.md index 6e8339c6b..08333d9dd 100755 --- a/docs/manuals/04-quickstart-setup-guide.md +++ b/docs/manuals/04-quickstart-setup-guide.md @@ -130,7 +130,7 @@ Check the version: ```bash ./target/release/extenddb version # extenddb 0.1.14 -# catalog migration 5 (postgres) +# catalog migration 6 (postgres) # commit abc1234 # built 2026-04-17T12:00:00Z ``` @@ -172,7 +172,7 @@ Config: extenddb.toml Catalog database: extenddb_catalog --- Checking catalog migrations... - OK: Current: the catalog is at migration 5 + OK: Current: the catalog is at migration 6 --- Checking catalog connection... OK: Connected to catalog. --- Checking data database... diff --git a/docs/manuals/07-upgrade-manual.md b/docs/manuals/07-upgrade-manual.md index 04e3111a6..95424ba58 100755 --- a/docs/manuals/07-upgrade-manual.md +++ b/docs/manuals/07-upgrade-manual.md @@ -4,7 +4,7 @@ ## Current Status -This release versions the PostgreSQL, SQLite, Cassandra and MongoDB catalogs by a migration ledger and carries catalog migrations 1 to 5, which `extenddb migrate --list` prints. **Every deployment created by a 0.1.x release runs `extenddb migrate --yes` once** before a server of this release starts on it: the server refuses a catalog with pending migrations. Migration 1 brings the catalog of every 0.1.x release to the v1.0 schema, and migration 5 drops `continuous_backups`, a table and on MongoDB a collection (see [Catalog 0.0.5](#catalog-005)). +This release versions the PostgreSQL, SQLite, Cassandra and MongoDB catalogs by a migration ledger and carries catalog migrations 1 to 6, which `extenddb migrate --list` prints. **Every deployment created by a 0.1.x release runs `extenddb migrate --yes` once** before a server of this release starts on it: the server refuses a catalog with pending migrations. Migration 1 brings the catalog of every 0.1.x release to the v1.0 schema, and migration 5 drops `continuous_backups`, a table and on MongoDB a collection (see [Catalog 0.0.5](#catalog-005)). On PostgreSQL, migration 6 re-encodes the strings an earlier release stored and runs only with every server of an earlier release stopped (see [U+0000 in strings on PostgreSQL and MongoDB](#u0000-in-strings-on-postgresql-and-mongodb-catalog-migration-6)). A backup that a 0.1.x release took can restore without its global and local secondary indexes: a SQLite backup taken by 0.1.14 records only the table's key schema and attribute definitions, and its restore creates the table with its items and no index. Take a new backup of each table after upgrading to v1.0, and keep the 0.1.x backups for their items only. @@ -39,7 +39,7 @@ Exit status 6 means the migration lock, or a running server, was still held when `verify` also checks the connection to the data database. When that check fails, `verify` exits 1 on a Behind catalog instead of 3, because `migrate --yes` cannot apply the pending migrations while the data database cannot be reached. It still exits 4 or 5 for the statuses that need a person. -A migration is additive when the previous release serves the catalog correctly with it applied, and breaking otherwise. Migrations 1 to 5 are breaking: do not run a 0.1.x server against a migrated catalog. A 0.1.x binary that writes to a migrated catalog (its `migrate --yes`, or a SQLite dev server) leaves it WrittenByOlderRelease. `migrate --yes --reconverge` repairs it by running the convergent migrations again. Migration 2, which rewrites stored policies, is not convergent: the repair skips it and prints a line that says so (see the v0.1.14 notes below). +A migration is additive when the previous release serves the catalog correctly with it applied, and breaking otherwise. Migrations 1 to 6 are breaking: do not run a 0.1.x server against a migrated catalog. A 0.1.x binary that writes to a migrated catalog (its `migrate --yes`, or a SQLite dev server) leaves it WrittenByOlderRelease. `migrate --yes --reconverge` repairs it by running the convergent migrations again. Migration 2, which rewrites stored policies, and migration 6, which re-encodes stored strings on PostgreSQL, are not convergent: the repair skips them and prints a line for each that says so (see the v0.1.14 notes below). ## Writing a New Migration @@ -311,6 +311,24 @@ Catalog upgrades change the schema; behavior changes alter how the running serve interprets existing data or configuration. Review these before upgrading a live deployment, even when no catalog migration is required. +### U+0000 in strings on PostgreSQL and MongoDB (catalog migration 6) + +DynamoDB accepts U+0000 anywhere a string appears. The PostgreSQL backend now stores every string through an order-preserving escape: U+0000 as U+0001 U+0001, U+0001 as U+0001 U+0002, and every other character as it is, so a string that holds neither character is stored as before. The MongoDB backend applies the same escape to BSON field names (attribute names and map keys); BSON string values already held U+0000. SQLite and Cassandra store the character as written, and their stored data does not change. + +**Migration 6, `escape_control_chars`, PostgreSQL.** Rows an earlier release wrote hold their strings raw, and a raw U+0001 immediately followed by U+0001 or U+0002 reads back differently under the new decoder; a lone U+0001 reads back unchanged. Migration 6 re-encodes those rows: in the data database, every data table, index table and vector index table (the text key columns and `item_data`), `gsi_pending` and `stream_records`; in the catalog database, `backup_items`. Only rows that contain U+0001 change. A partition key column of a key schema with more than one HASH attribute holds the netstring join of the parts, so it is rebuilt part by part, and `gsi_pending.worker_partition` is recomputed for the queued rows the migration changes. The migration reads every row of every table once, so on a large deployment it takes time proportional to the size of the data. Each database's rewrite runs in one transaction with its ledger row: a run that stops leaves that database as it was, and the next `migrate --yes` runs it again. The ledger notes, which `migrate` and `verify` print, name each table the migration changed with its number of row updates. + +Upgrade in this order: + +1. Stop every ExtendDB server that uses the deployment, including the servers of earlier releases. +2. Run `extenddb migrate --yes` with the new binary. +3. Start the new servers. + +The migration cannot tell a row a running server writes from a legacy row, so it runs only while no server writes. `migrate --yes` waits for the servers that hold the serving key, and a server of a 0.1.x release does not hold it. While the catalog lists a table or `gsi_pending` or `stream_records` holds a row, migration 6 therefore stops the run before its first write when any other session is connected to the data database, and while `backup_items` holds a row it does the same for the catalog database: `migrate` names each such session by role, client address and application name, and exits 5. The migrator's own sessions and read-only status sessions are not counted. Stop those sessions and run again. If the only other sessions are idle connections that a connection pooler holds, set `EXTENDDB_MIGRATE_IGNORE_CONNECTIONS=1` to skip the check. The preview, `extenddb migrate` without `--yes`, does not look for sessions, so it runs beside a server and reports migration 6 as pending. + +A server of this release refuses a deployment where migration 6 is pending, and `extenddb init` applies it to a new deployment, which has nothing to rewrite. The migration is breaking and not convergent: a server of an earlier release refuses the migrated catalog, `migrate --yes --reconverge` skips migration 6, and going back to an earlier release means restoring both databases from the backups taken before the upgrade. A 0.1.x binary given a restored catalog and a migrated data database would read escaped strings as data. + +**MongoDB.** The MongoDB backend has no migration for this change; its ledger records migration 6 as not applicable. Field names an earlier release wrote are stored raw. A name that holds a lone U+0001 reads back unchanged, as do all string values and all key strings; a name that holds 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 not affected. + ### v0.1.15: Point-in-time recovery is reported as unsupported (catalog migration required) **What changed.** `UpdateContinuousBackups` no longer accepts an enable request, and `DescribeContinuousBackups` no longer reports point-in-time recovery `ENABLED` with a 35-day window. Point-in-time recovery is reported as unsupported: describe answers `PointInTimeRecoveryStatus: DISABLED`, an enable request returns `ContinuousBackupsUnavailableException`, and `RestoreTableToPointInTime` checks the source table and then returns `PointInTimeRecoveryUnavailableException`. A table that does not exist is `TableNotFoundException` on all three operations, where `DescribeContinuousBackups` and `UpdateContinuousBackups` answered `ResourceNotFoundException`, and `UpdateContinuousBackups` without `PointInTimeRecoverySpecification` is a `ValidationException`, where it was taken as disabling. See [Differences from DynamoDB](../differences-from-dynamodb.md) for the full surface. The catalog migration removes the storage of a per-table enable flag that only the description read. diff --git a/docs/manuals/08-install-linux.md b/docs/manuals/08-install-linux.md index 928ba957d..2b868522e 100755 --- a/docs/manuals/08-install-linux.md +++ b/docs/manuals/08-install-linux.md @@ -118,7 +118,7 @@ Expected: ``` === extenddb verify === ... - OK: Current: the catalog is at migration 5 + OK: Current: the catalog is at migration 6 ... === HEALTHY: All checks passed === ``` diff --git a/docs/manuals/09-install-macos.md b/docs/manuals/09-install-macos.md index 85ba0f1e8..b1ca6de2d 100755 --- a/docs/manuals/09-install-macos.md +++ b/docs/manuals/09-install-macos.md @@ -96,7 +96,7 @@ Expected: ``` === extenddb verify === ... - OK: Current: the catalog is at migration 5 + OK: Current: the catalog is at migration 6 ... === HEALTHY: All checks passed === ``` diff --git a/tests/test_cli_vector_catalog_migration.py b/tests/test_cli_vector_catalog_migration.py index f5c4b4ad4..f24b23337 100644 --- a/tests/test_cli_vector_catalog_migration.py +++ b/tests/test_cli_vector_catalog_migration.py @@ -9,7 +9,9 @@ `continuous_backups` table. These tests load a deployment of a release from before the vector index schema, from its committed fixture, and check that the binary refuses to serve it, that `extenddb migrate` upgrades it, and that -the upgraded deployment then serves. They also check that +the upgraded deployment then serves. They also check that `migrate --yes` +applies migration 6 to a deployment holding a table and a backup and stops +while another session holds the catalog, that `--reconverge` runs every convergent migration again on a current catalog, and that a catalog ledger row for a migration this binary does not carry makes `serve` refuse. @@ -176,6 +178,76 @@ def _init(cli_env): assert result.returncode == 0, result.stdout + result.stderr +def _seed_table_and_backup(cli_env): + """Create a table, an item and a backup through the running server, so + that both bodies of migration 6 check for other sessions: the data body + while the catalog lists a table, the catalog body while `backup_items` + holds a row.""" + import uuid + + import boto3 + from conftest import wait_for_active + from management_helpers import ManagementClient + + endpoint = f"https://127.0.0.1:{cli_env['port']}" + mgmt = ManagementClient( + endpoint, os.environ.get("EXTENDDB_ADMIN_USER") or "admin", "TestPass1!" + ) + account_id = f"{uuid.uuid4().int % 10**12:012d}" + assert mgmt.create_account(account_id, "escape").status_code == 201 + assert mgmt.create_user(account_id, "seed", "TestPass123!").status_code == 201 + keys = mgmt.create_access_key(account_id, "seed").json() + assert mgmt.put_user_policy( + account_id, "seed", "full", + {"Version": "2012-10-17", "Statement": [ + {"Effect": "Allow", "Action": "dynamodb:*", "Resource": "*"}]}, + ).status_code == 204 + client = boto3.client( + "dynamodb", + endpoint_url=endpoint, + aws_access_key_id=keys["access_key_id"], + aws_secret_access_key=keys["secret_access_key"], + region_name="us-east-1", + verify=False, + ) + client.create_table( + TableName="escape", + AttributeDefinitions=[{"AttributeName": "pk", "AttributeType": "S"}], + KeySchema=[{"AttributeName": "pk", "KeyType": "HASH"}], + BillingMode="PAY_PER_REQUEST", + ) + wait_for_active(client, "escape") + client.put_item(TableName="escape", Item={"pk": {"S": "a"}, "v": {"S": "b"}}) + arn = client.create_backup(TableName="escape", BackupName="escape")["BackupDetails"][ + "BackupArn" + ] + deadline = time.monotonic() + 30 + while True: + details = client.describe_backup(BackupArn=arn)["BackupDescription"]["BackupDetails"] + if details["BackupStatus"] == "AVAILABLE": + return + assert time.monotonic() < deadline, details + time.sleep(0.5) + + +def _forget_migration_6(cli_env): + """Take migration 6 out of both ledgers and set the catalog version back + to 5, as a deployment of the release before it holds them. The rows this + build wrote hold no U+0001, so the migration rewrites none of them when it + runs again.""" + for database in (cli_env["db_name"], _data_db(cli_env)): + conn = _conn(database) + try: + with conn.cursor() as cur: + cur.execute("DELETE FROM extenddb_migrations WHERE id = 6") + if database == cli_env["db_name"]: + cur.execute( + "UPDATE settings SET value = '5' WHERE key = 'catalog_version'" + ) + finally: + conn.close() + + def _serve_refused(cli_env): """Run `serve --foreground`; it must exit on its own, and the result is returned. A server that starts serving is stopped and fails the test.""" @@ -340,6 +412,76 @@ def test_migrate_upgrades_a_catalog_a_release_before_1_0_created(self, cli_env): assert again.returncode == 0, again_output assert "up to date" in again_output.lower(), again_output + def test_migrate_applies_migration_6_beside_its_own_sessions(self, cli_env): + """`migrate --yes` applies migration 6 to a deployment holding a table + and a backup, and stops while another session holds the catalog. + + Migration 6 rewrites stored strings, so each body stops the run while a + session other than the migrate command's own holds the database it + rewrites and that database has rows to rewrite. While the check runs, + the command holds sessions of its own: the migration lock's session on + the catalog database and the runner's sessions on both databases. The + check leaves them out by their application_name; if it counted them, + it would refuse every upgrade. The deployment comes from this binary's + init, with migration 6 then taken out of both ledgers, as an upgrade + from the release before it finds them. + """ + _init(cli_env) + _patch_config_port(cli_env["config_path"], cli_env["port"]) + served = _run_extenddb("serve", config=cli_env["config_path"]) + assert served.returncode == 0, served.stdout + served.stderr + try: + assert _wait_for_server(cli_env["port"]), "the deployment did not serve" + _seed_table_and_backup(cli_env) + finally: + _run_extenddb("stop", config=cli_env["config_path"], check=False) + time.sleep(1) + _forget_migration_6(cli_env) + catalog, data = cli_env["db_name"], _data_db(cli_env) + + pending = _run_extenddb( + "migrate", *_pg_args(), config=cli_env["config_path"], check=False + ) + pending_output = pending.stdout + pending.stderr + assert pending.returncode == 3, pending_output + assert "6 escape_control_chars (catalog, breaking)" in pending_output, pending_output + assert "6 escape_control_chars (data, breaking)" in pending_output, pending_output + + # Another client's session on the catalog database, which holds the + # backup's items: the run stops before any write, in either store. + holder = _conn(catalog) + try: + refused = _run_extenddb( + "migrate", "--yes", *_pg_args(), + config=cli_env["config_path"], + check=False, + timeout=90, + ) + finally: + holder.close() + refused_output = refused.stdout + refused.stderr + assert refused.returncode == 5, refused_output + assert ( + f"other session(s) hold the catalog database {catalog}" in refused_output + ), refused_output + assert 6 not in _ledger_ids(catalog) + assert 6 not in _ledger_ids(data) + + applied = _run_extenddb( + "migrate", "--yes", *_pg_args(), + config=cli_env["config_path"], + check=False, + timeout=90, + ) + output = applied.stdout + applied.stderr + assert applied.returncode == 0, output + assert "other session(s) hold" not in output, output + assert f"Applied 6 escape_control_chars in {catalog}" in output, output + assert f"Applied 6 escape_control_chars in {data}" in output, output + assert 6 in _ledger_ids(catalog) + assert 6 in _ledger_ids(data) + assert _catalog_version(cli_env) == str(_latest()) + def test_reconverge_runs_every_convergent_migration_again_on_a_current_catalog( self, cli_env ): 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"]