Skip to content

Commit

Permalink
DynamoDB: query() should handle GSI's without range key (#7729)
Browse files Browse the repository at this point in the history
  • Loading branch information
bblommers authored May 29, 2024
1 parent b19afe8 commit d3dbf15
Show file tree
Hide file tree
Showing 2 changed files with 45 additions and 4 deletions.
9 changes: 5 additions & 4 deletions moto/dynamodb/models/table.py
Original file line number Diff line number Diff line change
Expand Up @@ -702,11 +702,12 @@ def query(
key for key in index.schema if key["KeyType"] == "RANGE"
][0]
except IndexError:
if isinstance(index, GlobalSecondaryIndex):
# If we're querying a GSI that does not have an index, the main hash key acts as a range key
index_range_key = {"AttributeName": self.hash_key_attr}
if isinstance(index, GlobalSecondaryIndex) and self.range_key_attr:
# If we're querying a GSI that does not have a range key, the main range key acts as a range key
index_range_key = {"AttributeName": self.range_key_attr}
else:
index_range_key = None
# If we don't have a range key on the main table either, the hash key acts as a range key
index_range_key = {"AttributeName": self.hash_key_attr}
if range_comparison:
raise ValueError(
f"Range Key comparison but no range key found for index: {index_name}"
Expand Down
40 changes: 40 additions & 0 deletions tests/test_dynamodb/test_dynamodb_query.py
Original file line number Diff line number Diff line change
Expand Up @@ -240,6 +240,46 @@ def test_query_gsi_pagination(table_name=None):
]


@pytest.mark.aws_verified
@dynamodb_aws_verified(add_range=True, add_gsi=True)
def test_query_gsi_pagination_with_string_range(table_name=None):
dynamodb = boto3.resource("dynamodb", region_name="us-east-1")
table = dynamodb.Table(table_name)

for i in range(3, 7):
table.put_item(Item={"pk": "the-key", "sk": f"{i}", "gsi_pk": "johndoe"})

for i in range(9, 6, -1):
table.put_item(Item={"pk": "the-key", "sk": f"{i}", "gsi_pk": "johndoe"})

for i in range(3):
table.put_item(Item={"pk": "the-key", "sk": f"{i}", "gsi_pk": "johndoe"})

page1 = table.query(
KeyConditionExpression=Key("gsi_pk").eq("johndoe"),
IndexName="test_gsi",
Limit=6,
)
assert page1["Count"] == 6
assert page1["ScannedCount"] == 6
assert len(page1["Items"]) == 6

page2 = table.query(
KeyConditionExpression=Key("gsi_pk").eq("johndoe"),
IndexName="test_gsi",
Limit=6,
ExclusiveStartKey=page1["LastEvaluatedKey"],
)
assert page2["Count"] == 4
assert page2["ScannedCount"] == 4
assert len(page2["Items"]) == 4
assert "LastEvaluatedKey" not in page2

results = page1["Items"] + page2["Items"]
subjects = set([int(r["sk"]) for r in results])
assert subjects == set(range(10))


@pytest.mark.aws_verified
@dynamodb_aws_verified(add_range=True, numeric_gsi_range=True)
def test_query_gsi_pagination_with_numeric_range(table_name=None):
Expand Down

0 comments on commit d3dbf15

Please sign in to comment.