Merge pull request #202 from redis/add-all-pks-for-json-model

Adds `all_pks()` method for `JsonModel` and supporting test.

Author: Simon Prickett

aredis_om/model/model.py

@@ -1483,6 +1483,22 @@ async def save(self, pipeline: Optional[Pipeline] = None) -> "JsonModel":
         await db.execute_command("JSON.SET", self.key(), ".", self.json())
         return self
 
+    @classmethod
+    async def all_pks(cls):  # type: ignore
+        key_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
+        # TODO: We assume the key ends with the default separator, ":" -- when
+        #  we make the separator configurable, we need to update this as well.
+        #  ... And probably lots of other places ...
+        #
+        # TODO: Also, we need to decide how we want to handle the lack of
+        #  decode_responses=True...
+        return (
+            key.split(":")[-1]
+            if isinstance(key, str)
+            else key.decode(cls.Meta.encoding).split(":")[-1]
+            async for key in cls.db().scan_iter(f"{key_prefix}*", _type="ReJSON-RL")
+        )
+
     async def update(self, **field_values):
         validate_model_fields(self.__class__, field_values)
         for field, value in field_values.items():

tests/test_hash_model.py

@@ -374,6 +374,35 @@ async def test_saves_model_and_creates_pk(m):
     member2 = await m.Member.get(member.pk)
     assert member2 == member
 
+@pytest.mark.asyncio
+async def test_all_pks(m):
+    member = m.Member(
+        first_name="Simon",
+        last_name="Prickett",
+        email="[email protected]",
+        join_date=today,
+        age=97,
+        bio="This is a test user to be deleted.",
+    )
+
+    await member.save()
+
+    member1 = m.Member(
+        first_name="Andrew",
+        last_name="Brookins",
+        email="[email protected]",
+        join_date=today,
+        age=38,
+        bio="This is a test user to be deleted.",
+    )
+
+    await member1.save()   
+
+    pk_list = []
+    async for pk in await m.Member.all_pks():
+        pk_list.append(pk)
+
+    assert len(pk_list) == 2
 
 @pytest.mark.asyncio
 async def test_delete(m):

tests/test_json_model.py

@@ -186,6 +186,35 @@ async def test_saves_model_and_creates_pk(address, m, redis):
     assert member2 == member
     assert member2.address == address
 
+@pytest.mark.asyncio
+async def test_all_pks(address, m, redis):
+    member = m.Member(
+        first_name="Andrew",
+        last_name="Brookins",
+        email="[email protected]",
+        join_date=today,
+        age=38,
+        address=address,
+    )
+
+    await member.save()  
+
+    member1 = m.Member(
+        first_name="Simon",
+        last_name="Prickett",
+        email="[email protected]",
+        join_date=today,
+        age=99,
+        address=address,
+    )
+
+    await member1.save()
+
+    pk_list = []
+    async for pk in await m.Member.all_pks():
+        pk_list.append(pk)
+
+    assert len(pk_list) == 2
 
 @pytest.mark.asyncio
 async def test_delete(address, m, redis):