Skip to content

Commit bc44114

Browse files
author
Andrew Brookins
committed
Add all_pks() method to HashModel
1 parent db7b8d1 commit bc44114

File tree

3 files changed

+19
-1
lines changed

3 files changed

+19
-1
lines changed

pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tool.poetry]
22
name = "redis-om"
3-
version = "0.0.8"
3+
version = "0.0.9"
44
description = "A high-level library containing useful Redis abstractions and tools, like an ORM and leaderboard."
55
authors = ["Andrew Brookins <[email protected]>"]
66
maintainers = ["Andrew Brookins <[email protected]>"]

redis_om/model/model.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1014,6 +1014,7 @@ def __new__(cls, name, bases, attrs, **kwargs): # noqa C901
10141014
new_class._meta.primary_key_creator_cls = getattr(
10151015
base_meta, "primary_key_creator_cls", UlidPrimaryKey
10161016
)
1017+
# TODO: Configurable key separate, defaults to ":"
10171018
if not getattr(new_class._meta, "index_name", None):
10181019
new_class._meta.index_name = (
10191020
f"{new_class._meta.global_key_prefix}:"
@@ -1203,6 +1204,17 @@ def save(self, pipeline: Optional[Pipeline] = None) -> "HashModel":
12031204
db.hset(self.key(), mapping=document)
12041205
return self
12051206

1207+
@classmethod
1208+
def all_pks(cls):
1209+
key_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
1210+
# TODO: We assume the key ends with the default separator, ":" -- when
1211+
# we make the separator configurable, we need to update this as well.
1212+
# ... And probably lots of other places ...
1213+
return (
1214+
key.split(":")[-1]
1215+
for key in cls.db().scan_iter(f"{key_prefix}*", _type="HASH")
1216+
)
1217+
12061218
@classmethod
12071219
def get(cls, pk: Any) -> "HashModel":
12081220
document = cls.db().hgetall(cls.make_primary_key(pk))

tests/test_hash_model.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -411,6 +411,12 @@ def test_sorting(members, m):
411411
m.Member.find().sort_by("join_date").all()
412412

413413

414+
def test_all_keys(members, m):
415+
pks = sorted(list(m.Member.all_pks()))
416+
assert len(pks) == 3
417+
assert pks == sorted([m.pk for m in members])
418+
419+
414420
def test_not_found(m):
415421
with pytest.raises(NotFoundError):
416422
# This ID does not exist.

0 commit comments

Comments
 (0)