diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py index dafd41a..9f2c3a7 100644 --- a/aredis_om/model/model.py +++ b/aredis_om/model/model.py @@ -1483,6 +1483,22 @@ class JsonModel(RedisModel, abc.ABC): 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(): diff --git a/tests/test_hash_model.py b/tests/test_hash_model.py index ddcf905..e8f80e0 100644 --- a/tests/test_hash_model.py +++ b/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="s@example.com", + 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="a@example.com", + 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): diff --git a/tests/test_json_model.py b/tests/test_json_model.py index d920a40..13c2834 100644 --- a/tests/test_json_model.py +++ b/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="a@example.com", + join_date=today, + age=38, + address=address, + ) + + await member.save() + + member1 = m.Member( + first_name="Simon", + last_name="Prickett", + email="s@example.com", + 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):