Skip to content

fix(storage): tie-break GSI pagination on the full base primary key #239

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
LeeroyHannigan wants to merge 1 commit into
base: main
Choose a base branch
from
Open

fix(storage): tie-break GSI pagination on the full base primary key #239

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 6 additions & 2 deletions crates/storage-postgres/src/data/query.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 },
}

Expand Down
99 changes: 72 additions & 27 deletions crates/storage-postgres/src/data/query_scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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!(
Expand Down Expand Up @@ -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!(
Expand Down Expand Up @@ -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
Expand Down
57 changes: 45 additions & 12 deletions crates/storage-sqlite/src/data/query_scan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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!(
Expand Down
147 changes: 147 additions & 0 deletions tests/test_query_scan.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down
Loading