Skip to content

Commit 274ff78

Browse files
author
Simon Prickett
committed
Adds all_pks() method and test.
1 parent 6453867 commit 274ff78

File tree

2 files changed

+45
-0
lines changed

2 files changed

+45
-0
lines changed

aredis_om/model/model.py

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1483,6 +1483,22 @@ async def save(self, pipeline: Optional[Pipeline] = None) -> "JsonModel":
14831483
await db.execute_command("JSON.SET", self.key(), ".", self.json())
14841484
return self
14851485

1486+
@classmethod
1487+
async def all_pks(cls): # type: ignore
1488+
key_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
1489+
# TODO: We assume the key ends with the default separator, ":" -- when
1490+
# we make the separator configurable, we need to update this as well.
1491+
# ... And probably lots of other places ...
1492+
#
1493+
# TODO: Also, we need to decide how we want to handle the lack of
1494+
# decode_responses=True...
1495+
return (
1496+
key.split(":")[-1]
1497+
if isinstance(key, str)
1498+
else key.decode(cls.Meta.encoding).split(":")[-1]
1499+
async for key in cls.db().scan_iter(f"{key_prefix}*", _type="ReJSON-RL")
1500+
)
1501+
14861502
async def update(self, **field_values):
14871503
validate_model_fields(self.__class__, field_values)
14881504
for field, value in field_values.items():

tests/test_json_model.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -186,6 +186,35 @@ async def test_saves_model_and_creates_pk(address, m, redis):
186186
assert member2 == member
187187
assert member2.address == address
188188

189+
@pytest.mark.asyncio
190+
async def test_all_pks(address, m, redis):
191+
member = m.Member(
192+
first_name="Andrew",
193+
last_name="Brookins",
194+
195+
join_date=today,
196+
age=38,
197+
address=address,
198+
)
199+
200+
await member.save()
201+
202+
member1 = m.Member(
203+
first_name="Simon",
204+
last_name="Prickett",
205+
206+
join_date=today,
207+
age=99,
208+
address=address,
209+
)
210+
211+
await member1.save()
212+
213+
pk_list = []
214+
async for pk in await m.Member.all_pks():
215+
pk_list.append(pk)
216+
217+
assert len(pk_list) == 2
189218

190219
@pytest.mark.asyncio
191220
async def test_delete(address, m, redis):

0 commit comments

Comments
 (0)