Add all_pks() method to HashModel

This commit is contained in:
Andrew Brookins 2021-11-08 16:28:58 -08:00
parent db7b8d19ad
commit bc441143de
3 changed files with 19 additions and 1 deletions

View file

@ -1,6 +1,6 @@
[tool.poetry] [tool.poetry]
name = "redis-om" name = "redis-om"
version = "0.0.8" version = "0.0.9"
description = "A high-level library containing useful Redis abstractions and tools, like an ORM and leaderboard." description = "A high-level library containing useful Redis abstractions and tools, like an ORM and leaderboard."
authors = ["Andrew Brookins <andrew.brookins@redis.com>"] authors = ["Andrew Brookins <andrew.brookins@redis.com>"]
maintainers = ["Andrew Brookins <andrew.brookins@redis.com>"] maintainers = ["Andrew Brookins <andrew.brookins@redis.com>"]

View file

@ -1014,6 +1014,7 @@ class ModelMeta(ModelMetaclass):
new_class._meta.primary_key_creator_cls = getattr( new_class._meta.primary_key_creator_cls = getattr(
base_meta, "primary_key_creator_cls", UlidPrimaryKey base_meta, "primary_key_creator_cls", UlidPrimaryKey
) )
# TODO: Configurable key separate, defaults to ":"
if not getattr(new_class._meta, "index_name", None): if not getattr(new_class._meta, "index_name", None):
new_class._meta.index_name = ( new_class._meta.index_name = (
f"{new_class._meta.global_key_prefix}:" f"{new_class._meta.global_key_prefix}:"
@ -1203,6 +1204,17 @@ class HashModel(RedisModel, abc.ABC):
db.hset(self.key(), mapping=document) db.hset(self.key(), mapping=document)
return self return self
@classmethod
def all_pks(cls):
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 ...
return (
key.split(":")[-1]
for key in cls.db().scan_iter(f"{key_prefix}*", _type="HASH")
)
@classmethod @classmethod
def get(cls, pk: Any) -> "HashModel": def get(cls, pk: Any) -> "HashModel":
document = cls.db().hgetall(cls.make_primary_key(pk)) document = cls.db().hgetall(cls.make_primary_key(pk))

View file

@ -411,6 +411,12 @@ def test_sorting(members, m):
m.Member.find().sort_by("join_date").all() m.Member.find().sort_by("join_date").all()
def test_all_keys(members, m):
pks = sorted(list(m.Member.all_pks()))
assert len(pks) == 3
assert pks == sorted([m.pk for m in members])
def test_not_found(m): def test_not_found(m):
with pytest.raises(NotFoundError): with pytest.raises(NotFoundError):
# This ID does not exist. # This ID does not exist.