Adds all_pks() method and test.

This commit is contained in:
Simon Prickett 2022-04-12 11:31:09 +01:00
parent 6453867eea
commit 274ff788b1
2 changed files with 45 additions and 0 deletions

View file

@ -1483,6 +1483,22 @@ class JsonModel(RedisModel, abc.ABC):
await db.execute_command("JSON.SET", self.key(), ".", self.json()) await db.execute_command("JSON.SET", self.key(), ".", self.json())
return self 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): async def update(self, **field_values):
validate_model_fields(self.__class__, field_values) validate_model_fields(self.__class__, field_values)
for field, value in field_values.items(): for field, value in field_values.items():

View file

@ -186,6 +186,35 @@ async def test_saves_model_and_creates_pk(address, m, redis):
assert member2 == member assert member2 == member
assert member2.address == address 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 @pytest.mark.asyncio
async def test_delete(address, m, redis): async def test_delete(address, m, redis):