DynamoDB: query() should handle GSI's without range key (#7729)

Author: bblommers

moto/dynamodb/models/table.py

@@ -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}"

tests/test_dynamodb/test_dynamodb_query.py

@@ -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):