From ce1845ae2abd24213f4d77a3b9a1577631ebdab5 Mon Sep 17 00:00:00 2001
From: Rico Hermans <rix0rrr@gmail.com>
Date: Tue, 10 Dec 2024 10:16:56 +0100
Subject: [PATCH] Make sure that we can't condition-query on PKs, ever

---
 tests/test_dynamo.py | 12 ++++++++++++
 website/dynamo.py    |  4 +++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/tests/test_dynamo.py b/tests/test_dynamo.py
index 2539b3c190b..a6cb521208d 100644
--- a/tests/test_dynamo.py
+++ b/tests/test_dynamo.py
@@ -690,6 +690,18 @@ def test_may_not_use_condition_on_index_partition_key(self):
                 'id': dynamo.BeginsWith('k'),
             })
 
+    def test_may_not_use_condition_on_table_with_only_partition_key(self):
+        self.table = dynamo.Table(
+            dynamo.MemoryStorage(),
+            'table',
+            partition_key='id',
+        )
+
+        with self.assertRaises(ValueError):
+            self.table.get_many({
+                'id': dynamo.BeginsWith('k'),
+            })
+
     def test_can_disambiguate_between_indexes_with_same_pk(self):
         self.table = dynamo.Table(
             dynamo.MemoryStorage(),
diff --git a/website/dynamo.py b/website/dynamo.py
index 9a1447b9c2b..a194de5d863 100644
--- a/website/dynamo.py
+++ b/website/dynamo.py
@@ -1632,7 +1632,9 @@ def decode_object(obj):
 def validate_only_sort_key(conds, sort_key):
     """Check that non-Equals conditions are only used on the sort key."""
     non_equals_fields = [k for k, v in conds.items() if not isinstance(v, Equals)]
-    if sort_key and set(non_equals_fields) - {sort_key}:
+    sort_key_set = {sort_key} if sort_key else set()
+
+    if set(non_equals_fields) - sort_key_set:
         raise ValueError(f"Non-Equals conditions only allowed on sort key {sort_key}, got: {list(conds)}")