Allow users to define a new primary key. (#347)

* make primary key programmable

Signed-off-by: wiseaidev <business@wiseai.dev>

* get primary key field using the `key` method

Signed-off-by: wiseaidev <business@wiseai.dev>

* adjust delete_many & expire methods

Signed-off-by: wiseaidev <business@wiseai.dev>

* fix query for int primary key

Signed-off-by: wiseaidev <business@wiseai.dev>

* fix grammar

Signed-off-by: wiseaidev <business@wiseai.dev>

* add unit tests

Signed-off-by: wiseaidev <business@wiseai.dev>

Signed-off-by: wiseaidev <business@wiseai.dev>
Co-authored-by: Chayim <chayim@users.noreply.github.com>
Co-authored-by: dvora-h <67596500+dvora-h@users.noreply.github.com>
This commit is contained in:
Mahmoud Harmouch 2022-09-08 14:38:49 +03:00 committed by GitHub
parent 2e09234b68
commit 551429c01a
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 74 additions and 13 deletions

View file

@ -564,7 +564,10 @@ class FindQuery:
separator_char,
)
return ""
if separator_char in value:
if isinstance(value, int):
# This if will hit only if the field is a primary key of type int
result = f"@{field_name}:[{value} {value}]"
elif separator_char in value:
# The value contains the TAG field separator. We can work
# around this by breaking apart the values and unioning them
# with multiple field:{} queries.
@ -1106,12 +1109,12 @@ class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
extra = "allow"
def __init__(__pydantic_self__, **data: Any) -> None:
super().__init__(**data)
__pydantic_self__.validate_primary_key()
super().__init__(**data)
def __lt__(self, other):
"""Default sort: compare primary key of models."""
return self.pk < other.pk
return self.key() < other.key()
def key(self):
"""Return the Redis key for this model."""
@ -1150,7 +1153,7 @@ class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
db = self._get_db(pipeline)
# TODO: Wrap any Redis response errors in a custom exception?
await db.expire(self.make_primary_key(self.pk), num_seconds)
await db.expire(self.key(), num_seconds)
@validator("pk", always=True, allow_reuse=True)
def validate_pk(cls, v):
@ -1167,7 +1170,9 @@ class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
primary_keys += 1
if primary_keys == 0:
raise RedisModelError("You must define a primary key for the model")
elif primary_keys > 1:
elif primary_keys == 2:
cls.__fields__.pop('pk')
elif primary_keys > 2:
raise RedisModelError("You must define only one primary key for a model")
@classmethod
@ -1275,7 +1280,7 @@ class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
db = cls._get_db(pipeline)
for chunk in ichunked(models, 100):
pks = [cls.make_primary_key(model.pk) for model in chunk]
pks = [model.key() for model in chunk]
await cls._delete(db, *pks)
return len(models)