fix: all_pks() for complex keys (#471)

* fix: all_pks for complex keys

* fix tests

* more fixes

* support Python below 3.9+

* black

* linter again
This commit is contained in:
Yaraslau Zhylko 2023-02-12 13:57:27 +02:00 committed by GitHub
parent 250d29de30
commit 412bdd6401
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 76 additions and 16 deletions

View file

@ -151,6 +151,14 @@ def decode_redis_value(
return obj.decode(encoding)
# TODO: replace with `str.removeprefix()` when only Python 3.9+ is supported
def remove_prefix(value: str, prefix: str) -> str:
"""Remove a prefix from a string."""
if value.startswith(prefix):
value = value[len(prefix) :] # noqa: E203
return value
class PipelineError(Exception):
"""A Redis pipeline error."""
@ -1350,16 +1358,12 @@ class HashModel(RedisModel, abc.ABC):
@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
# TODO: We need to decide how we want to handle the lack of
# decode_responses=True...
return (
key.split(":")[-1]
remove_prefix(key, key_prefix)
if isinstance(key, str)
else key.decode(cls.Meta.encoding).split(":")[-1]
else remove_prefix(key.decode(cls.Meta.encoding), key_prefix)
async for key in cls.db().scan_iter(f"{key_prefix}*", _type="HASH")
)
@ -1521,16 +1525,12 @@ class JsonModel(RedisModel, abc.ABC):
@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
# TODO: We need to decide how we want to handle the lack of
# decode_responses=True...
return (
key.split(":")[-1]
remove_prefix(key, key_prefix)
if isinstance(key, str)
else key.decode(cls.Meta.encoding).split(":")[-1]
else remove_prefix(key.decode(cls.Meta.encoding), key_prefix)
async for key in cls.db().scan_iter(f"{key_prefix}*", _type="ReJSON-RL")
)