From a7fee4bb3b42dcc75fb708ae380e89b557872f1f Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Sun, 25 Oct 2026 22:43:00 +0000 Subject: [PATCH 1/2] fix(query): bound begins_with on a string sort key by the next prefix Query with begins_with(sk, :p) on a string sort key read the range sk >= p AND sk < p + U+10FFFF on Cassandra, SQLite and PostgreSQL. A key that starts with p followed by U+10FFFF (p + U+10FFFF, p + U+10FFFF + "b") sorts at or above that bound and was left out. Amazon DynamoDB returns those keys on base tables, GSIs and LSIs (measured 2026-10-25 at 22:31, 22:35 and 23:20 UTC, each request twice), ordered by their UTF-8 bytes, which is how the three backends compare string sort keys. The new extenddb_storage::util::string_prefix_upper_bound returns the prefix with its last character replaced by the next scalar value: trailing U+10FFFF characters are dropped first, and U+D7FF is followed by U+E000. A prefix that is empty or all U+10FFFF has no upper bound, and the three backends then read from the prefix up. MongoDB already computed the same bound in its own next_string_prefix and is unchanged. On PostgreSQL the bound is now a bound parameter instead of `$p || chr(1114111)`, so build_sk_sql takes the expression maps to know whether the placeholder exists, and bind_sk_condition binds the bound it computes from the same function. A numeric prefix on a numeric sort key is refused in build_sk_sql before any SQL is built, so the placeholder and bind counts cannot disagree; it answered InternalServerError before and still does. On Cassandra the old string_upper_bound leaves query.rs. --- crates/storage-cassandra/src/data/query.rs | 6 - crates/storage-cassandra/src/data/sk_range.rs | 74 ++++++++++-- crates/storage-postgres/src/data/query.rs | 109 +++++++++++++++--- .../storage-postgres/src/data/query_scan.rs | 7 +- crates/storage-sqlite/src/data/query.rs | 24 ++-- crates/storage/src/util/mod.rs | 2 + crates/storage/src/util/prefix.rs | 102 ++++++++++++++++ 7 files changed, 282 insertions(+), 42 deletions(-) create mode 100644 crates/storage/src/util/prefix.rs diff --git a/crates/storage-cassandra/src/data/query.rs b/crates/storage-cassandra/src/data/query.rs index 1afe7e2c2..f418a7b6a 100644 --- a/crates/storage-cassandra/src/data/query.rs +++ b/crates/storage-cassandra/src/data/query.rs @@ -49,12 +49,6 @@ enum PaginationBinds { BasePkAndSk { pk_text: String, sk: SortKeyValue }, } -/// Compute upper bound for `begins_with` on strings. -/// Appends the maximum Unicode codepoint to create an exclusive upper bound. -pub(super) fn string_upper_bound(prefix: &str) -> String { - format!("{prefix}\u{10FFFF}") -} - /// Compute upper bound for `begins_with` on binary data. /// Returns `None` when the prefix is empty or all-0xFF (no finite upper bound). pub(super) fn binary_upper_bound(prefix: &[u8]) -> Option> { diff --git a/crates/storage-cassandra/src/data/sk_range.rs b/crates/storage-cassandra/src/data/sk_range.rs index 8973daf41..5a37faa07 100644 --- a/crates/storage-cassandra/src/data/sk_range.rs +++ b/crates/storage-cassandra/src/data/sk_range.rs @@ -8,9 +8,9 @@ use std::cmp::Ordering; use extenddb_core::expression::{CompareOp, ExpressionMaps, SortKeyCondition}; use extenddb_core::types::ScalarAttributeType; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{SortKeyValue, parse_sk}; +use extenddb_storage::util::{SortKeyValue, parse_sk, string_prefix_upper_bound}; -use super::query::{binary_upper_bound, string_upper_bound}; +use super::query::binary_upper_bound; use super::resolve_expr_to_av; /// One side of a sort key range: the CQL operator and the bound value. @@ -66,10 +66,10 @@ impl SkRange { SortKeyCondition::BeginsWith { prefix, .. } => { let prefix_av = resolve_expr_to_av(prefix, expression_maps)?; let prefix_sk = parse_sk(&prefix_av, sk_type)?; - // Exclusive upper bound by type; None when the prefix is empty - // or all 0xFF (no finite upper bound). + // Exclusive upper bound by type; None when the prefix is empty, + // all U+10FFFF or all 0xFF (no finite upper bound). let upper_sk_opt = match &prefix_sk { - SortKeyValue::S(s) => Some(SortKeyValue::S(string_upper_bound(s))), + SortKeyValue::S(s) => string_prefix_upper_bound(s).map(SortKeyValue::S), SortKeyValue::B(b) => binary_upper_bound(b).map(SortKeyValue::B), SortKeyValue::N(_) => { return Err(StorageError::Validation( @@ -261,11 +261,57 @@ mod tests { ); } + #[test] + fn begins_with_range_admits_every_key_with_the_prefix_and_no_other() { + // Sort keys measured on the service on 2026-10-25: begins_with(sk, p) + // returns exactly the keys that start with p, in UTF-8 byte order. + let keys = [ + "a", + "ab", + "a\u{D7FF}", + "a\u{E000}", + "a\u{FFFF}", + "a\u{10000}", + "a\u{10FFFF}", + "a\u{10FFFF}b", + "a\u{10FFFF}\u{10FFFF}", + "b", + "\u{10FFFF}", + "\u{10FFFF}x", + "\u{10FFFF}\u{10FFFF}", + "z", + ]; + for prefix in [ + "a", + "a\u{10FFFF}", + "a\u{10FFFF}\u{10FFFF}", + "a\u{D7FF}", + "\u{10FFFF}", + ] { + let range = SkRange { + lower: Some((">=", s(prefix))), + upper: string_prefix_upper_bound(prefix).map(|u| ("<", SortKeyValue::S(u))), + }; + for key in keys { + assert_eq!( + range.contains(&s(key)), + key.starts_with(prefix), + "prefix {} key {}", + prefix.escape_unicode(), + key.escape_unicode() + ); + } + } + } + #[test] fn begins_with_in_reverse_replaces_the_exclusive_upper_bound() { let mut range = SkRange { lower: Some((">=", s("ab"))), - upper: Some(("<", SortKeyValue::S(string_upper_bound("ab")))), + upper: Some(( + "<", + SortKeyValue::S(string_prefix_upper_bound("ab").unwrap()), + )), }; range.advance_past(s("abz"), false); assert_eq!( @@ -370,17 +416,23 @@ mod tests { ), (" AND sk >= ? AND sk <= ?".to_owned(), strs(&["aa", "m"])) ); - // begins_with on a string: inclusive prefix, exclusive prefix plus U+10FFFF. + // begins_with on a string: inclusive prefix, exclusive next prefix. assert_eq!( range( "pk = :p AND begins_with(sk, :v)", &[(":v", s_av("ab"))], ScalarAttributeType::S ), - ( - " AND sk >= ? AND sk < ?".to_owned(), - strs(&["ab", "ab\u{10FFFF}"]) - ) + (" AND sk >= ? AND sk < ?".to_owned(), strs(&["ab", "ac"])) + ); + // A prefix of U+10FFFF alone has no finite upper bound. + assert_eq!( + range( + "pk = :p AND begins_with(sk, :v)", + &[(":v", s_av("\u{10FFFF}"))], + ScalarAttributeType::S + ), + (" AND sk >= ?".to_owned(), strs(&["\u{10FFFF}"])) ); // begins_with on binary: the last byte below 0xFF is incremented and the // rest dropped; an all-0xFF prefix has no finite upper bound. diff --git a/crates/storage-postgres/src/data/query.rs b/crates/storage-postgres/src/data/query.rs index 6a84dc1b7..69e16748d 100755 --- a/crates/storage-postgres/src/data/query.rs +++ b/crates/storage-postgres/src/data/query.rs @@ -12,7 +12,7 @@ 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::{parse_sk, pk_to_text, sk_info, string_prefix_upper_bound}; /// Extra pagination bind values for index queries. /// @@ -82,8 +82,9 @@ pub(crate) struct SkSqlInfo { pub(crate) fn build_sk_sql( sk_cond: &SortKeyCondition, sk_col: &str, + maps: &ExpressionMaps, param_idx: &mut u32, -) -> SkSqlInfo { +) -> Result { // Apply COLLATE "C" for string sort key columns to get byte-order comparison. let collate = if sk_col == "sk_s" || sk_col.ends_with("_s") { " COLLATE \"C\"" @@ -102,7 +103,7 @@ pub(crate) fn build_sk_sql( }; let frag = format!(" AND {sk_col}{collate} {sql_op} ${param_idx}"); *param_idx += 1; - SkSqlInfo { fragment: frag } + Ok(SkSqlInfo { fragment: frag }) } SortKeyCondition::Between { .. } => { let frag = format!( @@ -111,9 +112,9 @@ pub(crate) fn build_sk_sql( hi = *param_idx + 1 ); *param_idx += 2; - SkSqlInfo { fragment: frag } + Ok(SkSqlInfo { fragment: frag }) } - SortKeyCondition::BeginsWith { .. } => { + SortKeyCondition::BeginsWith { prefix, .. } => { let is_binary = sk_col == "sk_b" || sk_col.ends_with("_b"); if is_binary { // For binary sort keys, bind both prefix and incremented-prefix @@ -125,22 +126,50 @@ pub(crate) fn build_sk_sql( hi = *param_idx + 1 ); *param_idx += 2; - SkSqlInfo { fragment: frag } - } else { - // For string sort keys, append the max Unicode code point. - // With COLLATE "C" (byte order), prefix || chr(1114111) is - // strictly greater than any string starting with prefix. + Ok(SkSqlInfo { fragment: frag }) + } else if string_prefix_bounded(prefix, sk_col, maps)? { + // For string sort keys the exclusive upper bound is the next + // prefix in code point order, computed in Rust (see + // bind_sk_condition). With COLLATE "C" the comparison is in + // byte order, which is code point order. let frag = format!( - " AND {sk_col}{collate} >= ${p} AND {sk_col}{collate} < (${p} || chr(1114111))", - p = *param_idx + " AND {sk_col}{collate} >= ${lo} AND {sk_col}{collate} < ${hi}", + lo = *param_idx, + hi = *param_idx + 1 ); + *param_idx += 2; + Ok(SkSqlInfo { fragment: frag }) + } else { + // A prefix of only U+10FFFF has no upper bound. + let frag = format!(" AND {sk_col}{collate} >= ${p}", p = *param_idx); *param_idx += 1; - SkSqlInfo { fragment: frag } + Ok(SkSqlInfo { fragment: frag }) } } } } +/// Whether a string `begins_with` prefix has a finite upper bound, so the SQL +/// carries a placeholder for it. [`bind_sk_condition`] binds that bound. +/// +/// A numeric prefix on a numeric column is refused here: it passes `parse_sk` +/// in [`bind_sk_condition`] and would bind one value for two placeholders. +/// Other non-string prefixes are refused by `parse_sk` before anything is +/// bound. +fn string_prefix_bounded( + prefix: &Expr, + sk_col: &str, + maps: &ExpressionMaps, +) -> Result { + match resolve_expr_to_av(prefix, maps)? { + AttributeValue::S(s) => Ok(string_prefix_upper_bound(&s).is_some()), + AttributeValue::N(_) if sk_col.ends_with("_n") => Err(StorageError::Internal( + "begins_with is not supported on numeric sort keys".to_owned(), + )), + _ => Ok(true), + } +} + /// Execute a query SQL statement with dynamic parameter binding. /// /// # Parameter binding order: @@ -245,7 +274,14 @@ fn bind_sk_condition<'q>( let q = bind_sk_value(query, &sk); Ok(q.bind(upper)) } else { - Ok(bind_sk_value(query, &sk)) + let q = bind_sk_value(query, &sk); + match &sk { + SortKeyValue::S(p) => match string_prefix_upper_bound(p) { + Some(upper) => Ok(q.bind(upper)), + None => Ok(q), + }, + _ => Ok(q), + } } } SortKeyCondition::Between { low, high, .. } => { @@ -355,3 +391,48 @@ pub(crate) async fn execute_scan_sql( pub(crate) fn build_key(item: &Item, key_schema: &[KeySchemaElement]) -> Item { extract_key(item, key_schema) } + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use extenddb_core::expression::{Expr, ExpressionMaps, PathElement, SortKeyCondition}; + use extenddb_core::types::AttributeValue; + + use super::build_sk_sql; + + fn begins_with(value: AttributeValue) -> (SortKeyCondition, ExpressionMaps) { + let cond = SortKeyCondition::BeginsWith { + path: vec![PathElement::Attribute("sk".to_owned())], + prefix: Expr::Placeholder(":v".to_owned()), + }; + let maps = ExpressionMaps::new(HashMap::new(), HashMap::from([(":v".to_owned(), value)])); + (cond, maps) + } + + #[test] + fn a_string_prefix_gets_an_upper_bound_placeholder_unless_it_has_none() { + let (cond, maps) = begins_with(AttributeValue::S("ab".to_owned())); + let mut idx = 2; + let info = build_sk_sql(&cond, "sk_s", &maps, &mut idx).unwrap(); + assert_eq!( + info.fragment, + " AND sk_s COLLATE \"C\" >= $2 AND sk_s COLLATE \"C\" < $3" + ); + assert_eq!(idx, 4); + + let (cond, maps) = begins_with(AttributeValue::S("\u{10FFFF}".to_owned())); + let mut idx = 2; + let info = build_sk_sql(&cond, "sk_s", &maps, &mut idx).unwrap(); + assert_eq!(info.fragment, " AND sk_s COLLATE \"C\" >= $2"); + assert_eq!(idx, 3); + } + + #[test] + fn a_numeric_prefix_on_a_numeric_column_is_refused_before_any_sql() { + let (cond, maps) = begins_with(AttributeValue::N("1".to_owned())); + let mut idx = 2; + assert!(build_sk_sql(&cond, "sk_n", &maps, &mut idx).is_err()); + assert_eq!(idx, 2); + } +} diff --git a/crates/storage-postgres/src/data/query_scan.rs b/crates/storage-postgres/src/data/query_scan.rs index c0f0c8b05..edd08da07 100644 --- a/crates/storage-postgres/src/data/query_scan.rs +++ b/crates/storage-postgres/src/data/query_scan.rs @@ -161,7 +161,12 @@ impl PostgresEngine { let sk_sql_info = if let (Some(sk_cond), Some((_, sk_type))) = (&key_condition.sk_condition, sk_info_val) { - Some(build_sk_sql(sk_cond, sk_column(sk_type), &mut param_idx)) + Some(build_sk_sql( + sk_cond, + sk_column(sk_type), + maps, + &mut param_idx, + )?) } else { None }; diff --git a/crates/storage-sqlite/src/data/query.rs b/crates/storage-sqlite/src/data/query.rs index c1139e85b..e51d2b21a 100644 --- a/crates/storage-sqlite/src/data/query.rs +++ b/crates/storage-sqlite/src/data/query.rs @@ -6,16 +6,17 @@ //! //! Sort-key comparisons bind through [`super::sk_bound`], so `N` keys are //! compared as their order-preserving TEXT encoding (D2): `>`, `<`, `BETWEEN` -//! all remain numerically correct. `begins_with` applies only to `S` (string -//! prefix via the maximal code point `char(1114111)`) and `B` (byte-range via -//! an incremented upper bound); DynamoDB rejects `begins_with` on `N`. +//! all remain numerically correct. `begins_with` applies only to `S` and `B`, +//! each read as a range up to the next prefix (see +//! [`string_prefix_upper_bound`] and [`binary_prefix_upper_bound`]); DynamoDB +//! rejects `begins_with` on `N`. use extenddb_core::expression::{self, CompareOp, Expr, ExpressionMaps, SortKeyCondition}; use extenddb_core::types::{ AttributeValue, Item, KeySchemaElement, ScalarAttributeType, extract_key, }; use extenddb_storage::error::StorageError; -use extenddb_storage::util::{SortKeyValue, parse_sk}; +use extenddb_storage::util::{SortKeyValue, parse_sk, string_prefix_upper_bound}; use super::BoundValue; @@ -103,12 +104,15 @@ pub(super) fn build_sk_sql_and_binds( // rows. None => Ok((format!(" AND {sk_col} >= ?"), vec![SortKeyValue::B(b)])), }, - // For strings the SQL upper bound is `? || char(1114111)`, - // so the same prefix is bound twice. - SortKeyValue::S(s) => Ok(( - format!(" AND {sk_col} >= ? AND {sk_col} < (? || char(1114111))"), - vec![SortKeyValue::S(s.clone()), SortKeyValue::S(s)], - )), + // The exclusive upper bound is the next prefix in code point + // order; a prefix of only U+10FFFF has none. + SortKeyValue::S(s) => match string_prefix_upper_bound(&s) { + Some(upper) => Ok(( + format!(" AND {sk_col} >= ? AND {sk_col} < ?"), + vec![SortKeyValue::S(s), SortKeyValue::S(upper)], + )), + None => Ok((format!(" AND {sk_col} >= ?"), vec![SortKeyValue::S(s)])), + }, SortKeyValue::N(_) => Err(StorageError::Validation( "begins_with is not supported on numeric sort keys".to_owned(), )), diff --git a/crates/storage/src/util/mod.rs b/crates/storage/src/util/mod.rs index b9c68a0f3..d95529b81 100755 --- a/crates/storage/src/util/mod.rs +++ b/crates/storage/src/util/mod.rs @@ -10,6 +10,7 @@ mod arn; mod backup_id; mod catalog_version; mod key; +mod prefix; pub use arn::{ ARN_PARTITION, format_stream_label, index_arn, new_stream_label, parse_stream_arn, stream_arn, @@ -27,3 +28,4 @@ pub use key::{ merge_attribute_definitions, parse_sk, pk_to_text, recover_sort_key_definitions, sk_column, sk_column_n, sk_info, }; +pub use prefix::string_prefix_upper_bound; diff --git a/crates/storage/src/util/prefix.rs b/crates/storage/src/util/prefix.rs new file mode 100644 index 000000000..8b281aaf7 --- /dev/null +++ b/crates/storage/src/util/prefix.rs @@ -0,0 +1,102 @@ +// Copyright 2026 ExtendDB contributors +// SPDX-License-Identifier: Apache-2.0 + +//! Range bounds for `begins_with` on a string sort key. + +/// Exclusive upper bound for `begins_with(sk, prefix)` on a string sort key. +/// +/// Sort keys are ordered by their UTF-8 bytes, which is code point order, so +/// the smallest string above every string that starts with `prefix` is the +/// prefix with its last character replaced by the next scalar value. A +/// trailing U+10FFFF has no next value and is dropped first. Returns `None` +/// when the prefix is empty or all U+10FFFF: every string at or above the +/// prefix then starts with it, and the range has no upper bound. +/// +/// Appending U+10FFFF to the prefix, as an upper bound, leaves out the keys +/// that continue with U+10FFFF. +#[must_use] +pub fn string_prefix_upper_bound(prefix: &str) -> Option { + let mut chars: Vec = prefix.chars().collect(); + while let Some(last) = chars.pop() { + if last == char::MAX { + continue; + } + // The only value in range without a successor is U+D7FF: the scalar + // values skip the surrogates U+D800 to U+DFFF. + chars.push(char::from_u32(u32::from(last) + 1).unwrap_or('\u{E000}')); + return Some(chars.into_iter().collect()); + } + None +} + +#[cfg(test)] +mod tests { + use super::string_prefix_upper_bound; + + #[test] + fn the_bound_is_the_next_prefix_in_code_point_order() { + let cases: [(&str, Option<&str>); 10] = [ + ("ab", Some("ac")), + ("a\u{7F}", Some("a\u{80}")), + ("a\u{D7FF}", Some("a\u{E000}")), + ("a\u{FFFF}", Some("a\u{10000}")), + ("a\u{10FFFF}", Some("b")), + ("a\u{10FFFF}\u{10FFFF}", Some("b")), + ("\u{10FFFF}a", Some("\u{10FFFF}b")), + ("\u{10FFFF}", None), + ("\u{10FFFF}\u{10FFFF}", None), + ("", None), + ]; + for (prefix, want) in cases { + assert_eq!( + string_prefix_upper_bound(prefix).as_deref(), + want, + "{}", + prefix.escape_unicode() + ); + } + } + + #[test] + fn the_range_holds_every_key_with_the_prefix_and_no_other() { + // Sort keys measured on the service on 2026-10-25: begins_with(sk, p) + // returns exactly the keys that start with p, in UTF-8 byte order. + let keys = [ + "a", + "ab", + "a\u{D7FF}", + "a\u{E000}", + "a\u{FFFF}", + "a\u{10000}", + "a\u{10FFFF}", + "a\u{10FFFF}b", + "a\u{10FFFF}\u{10FFFF}", + "b", + "\u{10FFFF}", + "\u{10FFFF}x", + "\u{10FFFF}\u{10FFFF}", + "z", + ]; + for prefix in [ + "a", + "a\u{10FFFF}", + "a\u{10FFFF}\u{10FFFF}", + "a\u{D7FF}", + "a\u{FFFF}", + "\u{10FFFF}", + ] { + let upper = string_prefix_upper_bound(prefix); + for key in keys { + let in_range = key.as_bytes() >= prefix.as_bytes() + && upper.as_ref().is_none_or(|u| key.as_bytes() < u.as_bytes()); + assert_eq!( + in_range, + key.starts_with(prefix), + "prefix {} key {}", + prefix.escape_unicode(), + key.escape_unicode() + ); + } + } + } +} From 38d70bfcf09a866a9ba046d0f75f29b3a416c607 Mon Sep 17 00:00:00 2001 From: Lee Hannigan Date: Sun, 25 Oct 2026 22:43:00 +0000 Subject: [PATCH 2/2] test(query): begins_with on sort keys that continue with U+10FFFF Dual-target wire tests from measurements on 2026-10-25 at 22:31 and 22:35 UTC: begins_with with prefixes ending in a, U+10FFFF, U+D7FF and U+FFFF, forward and reverse; Limit 1 pages over the prefixes a, a + U+10FFFF and U+10FFFF; and two BETWEEN ranges whose bounds sit next to U+10FFFF. The service returns every key that starts with the prefix, in UTF-8 byte order. --- tests/test_query_begins_with_max_character.py | 140 ++++++++++++++++++ 1 file changed, 140 insertions(+) create mode 100644 tests/test_query_begins_with_max_character.py diff --git a/tests/test_query_begins_with_max_character.py b/tests/test_query_begins_with_max_character.py new file mode 100644 index 000000000..84698a437 --- /dev/null +++ b/tests/test_query_begins_with_max_character.py @@ -0,0 +1,140 @@ +# Copyright 2026 ExtendDB contributors +# SPDX-License-Identifier: Apache-2.0 + +"""begins_with on a string sort key whose next character is at or near U+10FFFF. + +Amazon DynamoDB returns every sort key that starts with the prefix, in the +order of their UTF-8 bytes, including keys that continue with U+10FFFF +(measured 2026-10-25 at 22:31 UTC, us-east-1, each request twice, both scan +directions; Limit 1 pages over the prefixes a, a + U+10FFFF and U+10FFFF at +22:35 UTC). + +Dual-target against Amazon DynamoDB and extenddb. +""" + +from __future__ import annotations + +import pytest + +from conftest import scoped_table + +MAX = "\U0010ffff" + +# Every sort key the table holds, in the service's order (UTF-8 byte order). +KEYS = [ + "a", + "ab", + "a\ud7ff", + "a\ue000", + "a\uffff", + "a\U00010000", + "a" + MAX, + "a" + MAX + "b", + "a" + MAX + MAX, + "b", + "z", + MAX, + MAX + "x", + MAX + MAX, +] + +PK = "begins-with-max" + + +@pytest.fixture(scope="module") +def table(dynamodb_client): + 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"}, + ], + ) as name: + for sk in KEYS: + dynamodb_client.put_item(TableName=name, Item={"pk": {"S": PK}, "sk": {"S": sk}}) + yield name + + +def _sort_keys(dynamodb_client, table, key_condition, values, forward): + items = [] + pages = 0 + kwargs = { + "TableName": table, + "KeyConditionExpression": key_condition, + "ExpressionAttributeValues": {":p": {"S": PK}, **values}, + "ScanIndexForward": forward, + "ConsistentRead": True, + } + while True: + pages += 1 + assert pages <= len(KEYS) + 1, f"no last page after {pages - 1} pages: {items}" + resp = dynamodb_client.query(**kwargs) + items.extend(i["sk"]["S"] for i in resp["Items"]) + if "LastEvaluatedKey" not in resp: + return items + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + + +@pytest.mark.parametrize("forward", [True, False], ids=["forward", "reverse"]) +@pytest.mark.parametrize( + "prefix", + ["a", "a" + MAX, "a" + MAX + MAX, MAX, MAX + MAX, "a\ud7ff", "a\uffff"], + ids=["a", "a_max", "a_max_max", "max", "max_max", "a_d7ff", "a_ffff"], +) +def test_begins_with_returns_every_key_with_the_prefix(dynamodb_client, table, prefix, forward): + got = _sort_keys( + dynamodb_client, + table, + "pk = :p AND begins_with(sk, :s)", + {":s": {"S": prefix}}, + forward, + ) + want = [k for k in KEYS if k.startswith(prefix)] + assert got == (want if forward else want[::-1]) + + +@pytest.mark.parametrize("prefix", ["a", "a" + MAX, MAX], ids=["a", "a_max", "max"]) +def test_begins_with_pages_through_keys_after_the_prefix(dynamodb_client, table, prefix): + """Limit 1 makes every page end inside the range, so each page's start key + is folded into the same upper bound. U+10FFFF alone has no upper bound.""" + want = [k for k in KEYS if k.startswith(prefix)] + for forward in (True, False): + items = [] + kwargs = { + "TableName": table, + "KeyConditionExpression": "pk = :p AND begins_with(sk, :s)", + "ExpressionAttributeValues": {":p": {"S": PK}, ":s": {"S": prefix}}, + "ScanIndexForward": forward, + "ConsistentRead": True, + "Limit": 1, + } + for _ in range(len(KEYS) + 1): + resp = dynamodb_client.query(**kwargs) + items.extend(i["sk"]["S"] for i in resp["Items"]) + if "LastEvaluatedKey" not in resp: + break + kwargs["ExclusiveStartKey"] = resp["LastEvaluatedKey"] + else: + pytest.fail(f"no last page after {len(KEYS) + 1} pages: {items}") + assert items == (want if forward else want[::-1]) + + +@pytest.mark.parametrize( + ("low", "high"), + [("a", "a" + MAX), ("a" + MAX, "b")], + ids=["a_to_a_max", "a_max_to_b"], +) +def test_between_includes_both_bounds(dynamodb_client, table, low, high): + got = _sort_keys( + dynamodb_client, + table, + "pk = :p AND sk BETWEEN :l AND :u", + {":l": {"S": low}, ":u": {"S": high}}, + True, + ) + lo, hi = low.encode(), high.encode() + assert got == [k for k in KEYS if lo <= k.encode() <= hi]