From d3dbf15a51ae77df513a6562366f27f37a4c2665 Mon Sep 17 00:00:00 2001
From: Bert Blommers <info@bertblommers.nl>
Date: Wed, 29 May 2024 20:45:43 +0000
Subject: [PATCH] DynamoDB: query() should handle GSI's without range key
 (#7729)

---
 moto/dynamodb/models/table.py              |  9 ++---
 tests/test_dynamodb/test_dynamodb_query.py | 40 ++++++++++++++++++++++
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/moto/dynamodb/models/table.py b/moto/dynamodb/models/table.py
index e64f5eecb0f2..32a4d3c81c32 100644
--- a/moto/dynamodb/models/table.py
+++ b/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}"
diff --git a/tests/test_dynamodb/test_dynamodb_query.py b/tests/test_dynamodb/test_dynamodb_query.py
index 71b4b8fc76c1..79b07b368e0b 100644
--- a/tests/test_dynamodb/test_dynamodb_query.py
+++ b/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):