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:
parent
2e09234b68
commit
551429c01a
2 changed files with 74 additions and 13 deletions
|
@ -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)
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue