diff --git a/crates/storage-postgres/src/data/query.rs b/crates/storage-postgres/src/data/query.rs index c97497e5..6a84dc1b 100755 --- a/crates/storage-postgres/src/data/query.rs +++ b/crates/storage-postgres/src/data/query.rs @@ -24,9 +24,13 @@ pub(crate) enum PaginationBinds { None, /// Index query where base table has no SK — only `base_pk` as tie-breaker. BasePkOnly { pk_text: String }, - /// Index query where base table has a SK — `base_sk` as tie-breaker. + /// LSI query. `base_sk` alone is a complete tie-breaker, because every + /// row shares the queried partition key. BaseSkOnly { sk: SortKeyValue }, - /// Hash-only index where base table has a SK — both `base_pk` and `base_sk`. + /// Full base primary key as tie-breaker. Required by a GSI query when the + /// base table has a SK: rows sharing an index SK can come from different + /// base partitions and can also share a base SK, so `base_sk` alone does + /// not identify a row. Also used by a hash-only index on such a table. BasePkAndSk { pk_text: String, sk: SortKeyValue }, } diff --git a/crates/storage-postgres/src/data/query_scan.rs b/crates/storage-postgres/src/data/query_scan.rs index 6d61e9af..0c74ad51 100644 --- a/crates/storage-postgres/src/data/query_scan.rs +++ b/crates/storage-postgres/src/data/query_scan.rs @@ -46,17 +46,34 @@ fn build_pagination_where( } else { "" }; - // LSI: base SK follows ScanIndexForward because items share the - // same partition and the base SK is part of the composite sort order. - // GSI: base SK is always ">" (ascending) because it's only a - // uniqueness tie-breaker, not a user-visible sort dimension. - let base_cmp = if is_lsi { cmp } else { ">" }; - format!( - " AND ({sk_col}{collate} {cmp} ${p1} OR \ - ({sk_col}{collate} = ${p1} AND {base_col}{base_collate} {base_cmp} ${p2}))", - p1 = param_idx, - p2 = param_idx + 1 - ) + if is_lsi { + // LSI: every row shares the queried partition key, so the base + // table's sort key alone identifies a row uniquely. It is also a + // user-visible sort dimension, so it follows ScanIndexForward. + format!( + " AND ({sk_col}{collate} {cmp} ${p1} OR \ + ({sk_col}{collate} = ${p1} AND {base_col}{base_collate} {cmp} ${p2}))", + p1 = param_idx, + p2 = param_idx + 1 + ) + } else { + // GSI: the tie-breaker must be the FULL base primary key. Rows in + // a GSI partition are unique on (index SK, base PK, base SK), not + // on (index SK, base SK): many base partitions can project the + // same index SK and the same base SK. Comparing base SK alone + // made a page-two query return nothing whenever the rows sharing + // an index SK also shared a base SK, so a paginating client + // silently stopped after page one. The base key stays ascending + // because it is a uniqueness tie-breaker, not a sort dimension. + format!( + " AND ({sk_col}{collate} {cmp} ${p1} OR \ + ({sk_col}{collate} = ${p1} AND (base_pk COLLATE \"C\" > ${p2} OR \ + (base_pk = ${p2} AND {base_col}{base_collate} > ${p3}))))", + p1 = param_idx, + p2 = param_idx + 1, + p3 = param_idx + 2 + ) + } } else if is_index { // Index with no base SK — use base_pk as tie-breaker format!( @@ -227,20 +244,32 @@ impl PostgresEngine { }; let dir = if forward { "ASC" } else { "DESC" }; if let Some((_, base_sk_type)) = &base_sk_info { - // Index queries sub-sort by base table SK when index sort keys are equal. + // Index queries sub-sort by the base table key when index sort + // keys are equal. // LSI: base SK follows ScanIndexForward (same partition, composite sort). - // GSI: base SK is always ASC (just a uniqueness tie-breaker). + // GSI: the full base primary key, ascending, because it is only a + // uniqueness tie-breaker. It must match the pagination predicate + // exactly; ordering by base SK alone leaves rows that share an + // index SK and a base SK in an arbitrary order, which no + // ExclusiveStartKey can resume from deterministically. let base_col = format!("base_{}", sk_column(*base_sk_type)); let base_collate = if *base_sk_type == ScalarAttributeType::S { " COLLATE \"C\"" } else { "" }; - let base_dir = if is_lsi { dir } else { "ASC" }; - let _ = write!( - sql, - " ORDER BY {sk_col}{collate} {dir}, {base_col}{base_collate} {base_dir}" - ); + if is_lsi { + let _ = write!( + sql, + " ORDER BY {sk_col}{collate} {dir}, {base_col}{base_collate} {dir}" + ); + } else { + let _ = write!( + sql, + " ORDER BY {sk_col}{collate} {dir}, base_pk COLLATE \"C\" ASC, \ + {base_col}{base_collate} ASC" + ); + } } else if index_name.is_some() { // Index with SK but no base SK: use base_pk as secondary sort let _ = write!( @@ -280,16 +309,32 @@ impl PostgresEngine { if sk_info_val.is_some() && let Some((ref base_sk_name, base_sk_type)) = base_sk_info { - // Index that has its own SK, with a base-table SK tie-breaker. - // The index SK is bound separately (see execute_query_sql); here - // we bind only the base SK. (SQL has $N for base_sk.) - // A hash-only index falls through to the BasePkAndSk arm below, - // because its SQL binds base_pk AND base_sk, not base_sk alone. - if let Some(base_sk_val) = start_key.get(base_sk_name.as_str()) { - let sk = parse_sk(base_sk_val, base_sk_type)?; - PaginationBinds::BaseSkOnly { sk } + // Index that has its own SK, with a base-table tie-breaker. The + // index SK is bound separately (see execute_query_sql); the binds + // here supply the tie-breaker. + // + // An LSI needs the base SK alone, because every row shares the + // queried partition key. A GSI needs the full base primary key, + // because rows sharing an index SK can come from different base + // partitions and can also share a base SK. + let base_sk = start_key + .get(base_sk_name.as_str()) + .map(|v| parse_sk(v, base_sk_type)) + .transpose()?; + if is_lsi { + match base_sk { + Some(sk) => PaginationBinds::BaseSkOnly { sk }, + None => PaginationBinds::None, + } } else { - PaginationBinds::None + let base_pk_attr = &key_info.base_key_schema[0].attribute_name; + match (start_key.get(base_pk_attr.as_str()), base_sk) { + (Some(pk_val), Some(sk)) => PaginationBinds::BasePkAndSk { + pk_text: pk_to_text(pk_val)?.into_owned(), + sk, + }, + _ => PaginationBinds::None, + } } } else if index_name.is_some() && sk_info_val.is_some() { // Index with SK but no base SK — SQL has $N for base_pk diff --git a/crates/storage-sqlite/src/data/query_scan.rs b/crates/storage-sqlite/src/data/query_scan.rs index d9ebbf12..b7c03f50 100644 --- a/crates/storage-sqlite/src/data/query_scan.rs +++ b/crates/storage-sqlite/src/data/query_scan.rs @@ -126,8 +126,16 @@ impl SqliteEngine { let sk_col = sk_column(sk_type); if let Some((_, base_type)) = &base_sk_info { let base_col = format!("base_{}", sk_column(*base_type)); - let base_dir = if is_lsi { dir } else { "ASC" }; - let _ = write!(sql, " ORDER BY {sk_col} {dir}, {base_col} {base_dir}"); + if is_lsi { + let _ = write!(sql, " ORDER BY {sk_col} {dir}, {base_col} {dir}"); + } else { + // GSI: order by the full base primary key after the index SK + // so the ordering matches the pagination tie-breaker exactly. + // Ordering by base SK alone leaves rows that share an index SK + // and a base SK in an arbitrary order, which no + // ExclusiveStartKey can resume from deterministically. + let _ = write!(sql, " ORDER BY {sk_col} {dir}, base_pk ASC, {base_col} ASC"); + } } else if index_name.is_some() { let _ = write!(sql, " ORDER BY {sk_col} {dir}, base_pk ASC"); } else { @@ -356,18 +364,43 @@ fn append_query_pagination( if let Some((base_name, base_type)) = base_sk_info { let base_col = format!("base_{}", sk_column(*base_type)); - let base_cmp = if is_lsi { cmp } else { ">" }; - let _ = write!( - sql, - " AND ({sk_col} {cmp} ? OR ({sk_col} = ? AND {base_col} {base_cmp} ?))" - ); let sk_bv = sk_bv.unwrap_or(BoundValue::Text(String::new())); - binds.push(sk_bv.clone()); - binds.push(sk_bv); - if let Some(v) = start_key.get(base_name.as_str()) { - binds.push(sk_bound(&parse_sk(v, *base_type)?)); + let base_sk_bv = if let Some(v) = start_key.get(base_name.as_str()) { + sk_bound(&parse_sk(v, *base_type)?) } else { - binds.push(BoundValue::Text(String::new())); + BoundValue::Text(String::new()) + }; + + if is_lsi { + // LSI: every row shares the queried partition key, so the base + // table's sort key alone identifies a row uniquely, and it is a + // user-visible sort dimension so it follows ScanIndexForward. + let _ = write!( + sql, + " AND ({sk_col} {cmp} ? OR ({sk_col} = ? AND {base_col} {cmp} ?))" + ); + binds.push(sk_bv.clone()); + binds.push(sk_bv); + binds.push(base_sk_bv); + } else { + // GSI: the tie-breaker must be the FULL base primary key. Rows + // in a GSI partition are unique on (index SK, base PK, base SK), + // not on (index SK, base SK): many base partitions can project + // the same index SK and the same base SK. Comparing base SK + // alone made a page-two query return nothing whenever the rows + // sharing an index SK also shared a base SK, so a paginating + // client silently stopped after page one. + let base_pk_bv = BoundValue::Text(base_pk_from_start_key(start_key, key_info)?); + let _ = write!( + sql, + " AND ({sk_col} {cmp} ? OR ({sk_col} = ? AND (base_pk > ? \ + OR (base_pk = ? AND {base_col} > ?))))" + ); + binds.push(sk_bv.clone()); + binds.push(sk_bv); + binds.push(base_pk_bv.clone()); + binds.push(base_pk_bv); + binds.push(base_sk_bv); } } else if is_index { let _ = write!( diff --git a/tests/test_query_scan.py b/tests/test_query_scan.py index ed7a57eb..d08d35d3 100755 --- a/tests/test_query_scan.py +++ b/tests/test_query_scan.py @@ -1310,6 +1310,153 @@ def test_paginate_reverse_gsi_with_tied_sort_keys( assert len(ids) == len(set(ids)), "Duplicate items in reverse GSI pagination" +@pytest.fixture(scope="class") +def gsi_on_composite_base_table(dynamodb_client): + """Base table with a range key, plus a GSI with a sort key. + + Two partitions in the index: + "tied" - every item shares the GSI sort key AND the base range key, so + the only thing distinguishing rows is the base partition key. + "spread" - items share the GSI sort key but have distinct base range keys. + + The "tied" partition is the shape that broke: with the base range key equal + across rows, a tie-breaker built from the base range key alone cannot make + progress, so page two came back empty. + """ + with scoped_table( + dynamodb_client, + attribute_definitions=[ + {"AttributeName": "itemId", "AttributeType": "S"}, + {"AttributeName": "sortId", "AttributeType": "S"}, + {"AttributeName": "category", "AttributeType": "S"}, + {"AttributeName": "priority", "AttributeType": "N"}, + ], + key_schema=[ + {"AttributeName": "itemId", "KeyType": "HASH"}, + {"AttributeName": "sortId", "KeyType": "RANGE"}, + ], + GlobalSecondaryIndexes=[ + { + "IndexName": "CategoryPriorityGSI", + "KeySchema": [ + {"AttributeName": "category", "KeyType": "HASH"}, + {"AttributeName": "priority", "KeyType": "RANGE"}, + ], + "Projection": {"ProjectionType": "ALL"}, + }, + ], + ) as name: + for i in range(1, 8): + dynamodb_client.put_item( + TableName=name, + Item={ + "itemId": {"S": f"item-{i}"}, + "sortId": {"S": "S"}, + "category": {"S": "tied"}, + "priority": {"N": "1"}, + }, + ) + dynamodb_client.put_item( + TableName=name, + Item={ + "itemId": {"S": f"item-{i}"}, + "sortId": {"S": f"s-{i}"}, + "category": {"S": "spread"}, + "priority": {"N": "1"}, + }, + ) + yield name + + +class TestGSIOnCompositeBaseTable: + """GSI pagination works when the base table has a range key. + + The hash-only-base equivalents live in TestGSIOnHashOnlyBaseTable and passed + throughout: a GSI tie-breaker built from the base partition key alone is + complete when that is the whole base primary key. Adding a base range key is + what exposed the incomplete tie-breaker. + """ + + def _collect(self, client, table, category, limit=2, forward=True): + def run(): + items = [] + start = None + # Bound the loop so a tie-breaker that fails to advance surfaces as a + # short result rather than spinning here forever. + for _ in range(50): + kwargs = { + "TableName": table, + "IndexName": "CategoryPriorityGSI", + "KeyConditionExpression": "category = :c", + "ExpressionAttributeValues": {":c": {"S": category}}, + "Limit": limit, + } + if not forward: + kwargs["ScanIndexForward"] = False + if start: + kwargs["ExclusiveStartKey"] = start + resp = client.query(**kwargs) + items.extend(resp["Items"]) + if "LastEvaluatedKey" not in resp: + break + start = resp["LastEvaluatedKey"] + return items + + return run + + def test_paginate_duplicate_gsi_and_base_sort_keys( + self, dynamodb_client, gsi_on_composite_base_table + ): + """All items are returned when the GSI sort key and base range key are tied. + + Every row shares category, priority and sortId, so only the base partition + key separates them. Page one previously returned 2 of 7 with a + LastEvaluatedKey, and resuming from it returned nothing, so a paginating + client silently saw 2 items and stopped. + """ + items = wait_for_gsi_items( + self._collect(dynamodb_client, gsi_on_composite_base_table, "tied"), 7 + ) + assert len(items) == 7, f"expected 7 items, got {len(items)}" + ids = sorted(i["itemId"]["S"] for i in items) + assert ids == sorted(f"item-{n}" for n in range(1, 8)) + + def test_paginate_duplicate_gsi_sort_keys_distinct_base_sort_keys( + self, dynamodb_client, gsi_on_composite_base_table + ): + """The base range key differing is not required for correctness. + + Guards against a fix that only works when the base range key happens to + break the tie. + """ + items = wait_for_gsi_items( + self._collect(dynamodb_client, gsi_on_composite_base_table, "spread"), 7 + ) + assert len(items) == 7, f"expected 7 items, got {len(items)}" + assert len({i["sortId"]["S"] for i in items}) == 7 + + def test_paginate_reverse_with_tied_sort_keys( + self, dynamodb_client, gsi_on_composite_base_table + ): + """Reverse pagination is complete and non-duplicating on tied sort keys. + + ScanIndexForward reverses the index sort key only; the base primary key + tie-breaker stays ascending, so the resume predicate and the ORDER BY have + to agree on that. If the tie-breaker followed ScanIndexForward instead, + this would skip items or repeat them. + """ + items = wait_for_gsi_items( + self._collect( + dynamodb_client, gsi_on_composite_base_table, "tied", limit=3, forward=False + ), + 7, + ) + assert len(items) == 7, f"expected 7 items, got {len(items)}" + ids = [i["itemId"]["S"] for i in items] + assert len(ids) == len(set(ids)), f"duplicates in reverse pagination: {ids}" + + + @pytest.fixture(scope="class") def base_key_schema_table(dynamodb_client): """Table with GSI for testing base_key_schema propagation.