From b547ac7c312d8fcaa0d89c97ea6909e9dedd1b42 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=97=84=EC=9C=A4=EC=84=AD?= <62834176+lh0156@users.noreply.github.com> Date: Mon, 27 Jul 2026 03:38:05 +0900 Subject: [PATCH 1/2] Optimize bulk query updates Refs #777 --- aredis_om/model/model.py | 146 ++++++++++++++++++++++++---- docs/querying.md | 3 +- tests/test_find_query_update.py | 164 ++++++++++++++++++++++++++++++++ 3 files changed, 295 insertions(+), 18 deletions(-) create mode 100644 tests/test_find_query_update.py diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py index 07f6de14..895c21c2 100644 --- a/aredis_om/model/model.py +++ b/aredis_om/model/model.py @@ -27,7 +27,7 @@ from typing import get_args as typing_get_args from more_itertools import ichunked -from pydantic import BaseModel +from pydantic import BaseModel, create_model try: @@ -978,6 +978,7 @@ def dict(self) -> Dict[str, Any]: page_size=self.page_size, limit=self.limit, expressions=copy(self.expressions), + knn=self.knn, sort_fields=copy(self.sort_fields), projected_fields=copy(self.projected_fields), nocontent=self.nocontent, @@ -2012,29 +2013,140 @@ def only(self, *fields: str): raise ValueError("only() requires at least one field name") return self.copy(projected_fields=list(fields)) - async def update(self, use_transaction=True, **field_values): + def _get_update_field(self, field_name: str): + """Return the Pydantic field targeted by an update path.""" + current_model = self.model + parts = field_name.split("__") + + for index, part in enumerate(parts): + model_fields = getattr(current_model, "model_fields", None) + if model_fields is None: + model_fields = getattr(current_model, "__fields__", {}) + if part not in model_fields: + raise QuerySyntaxError( + f"The field {field_name} does not exist on the model " + f"{self.model.__name__}" + ) + + field = model_fields[part] + if index == len(parts) - 1: + return field + + annotation = getattr(field, "annotation", None) + nested_model = next( + ( + candidate + for candidate in (annotation, *get_args(annotation)) + if isinstance(candidate, type) + and hasattr(candidate, "model_fields") + ), + None, + ) + if nested_model is None: + raise QuerySyntaxError( + f"The update path {field_name} does not contain a nested model " + f"at {part}" + ) + current_model = nested_model + + raise QuerySyntaxError(f"The update path {field_name} is empty") + + def _validated_update_values(self, field_values: Dict[str, Any]): + """Validate update values once using the target fields' constraints.""" + validate_model_fields(self.model, field_values) + field_definitions: Dict[str, Any] = {} + for field_name in field_values: + field = self._get_update_field(field_name) + field_info = field if PYDANTIC_V2 else field.field_info + field_definitions[field_name] = (field.annotation, field_info) + + update_model = create_model( # type: ignore[call-overload] + f"{self.model.__name__}Update", **field_definitions + ) + if PYDANTIC_V2: + return update_model.model_validate(field_values).model_dump() + return update_model.parse_obj(field_values).dict() + + def _serialize_update_values(self, field_values: Dict[str, Any]) -> Dict[str, Any]: + """Validate and encode values using the storage model's conventions.""" + validated_values = self._validated_update_values(field_values) + + if issubclass(self.model, HashModel): + document = convert_datetime_to_timestamp(validated_values) + model_fields = getattr(self.model, "model_fields", None) + if model_fields is None: + model_fields = getattr(self.model, "__fields__", {}) + document = convert_vector_to_bytes(document, model_fields) + document = convert_bytes_to_base64(document) + document = jsonable_encoder(document) + document = { + key: value for key, value in document.items() if value is not None + } + return { + key: ("1" if value else "0") if isinstance(value, bool) else value + for key, value in document.items() + } + + document = convert_datetime_to_timestamp(validated_values) + document = convert_bytes_to_base64(document) + return jsonable_encoder(document) + + async def _search_keys_only(self) -> List[Union[str, bytes]]: + """Return all matching Redis keys without loading document contents.""" + query = self.copy( + limit=self.page_size, + nocontent=True, + projected_fields=[], + return_as_dict=False, + ) + keys: List[Union[str, bytes]] = [] + + while True: + raw_result = await query.execute( + exhaust_results=False, + return_raw_result=True, + ) + if not raw_result: + break + + count = raw_result[0] + page_keys = raw_result[1:] + keys.extend(page_keys) + if not page_keys or query.offset + len(page_keys) >= count: + break + + query = query.copy(offset=query.offset + len(page_keys)) + + return keys + + async def update(self, use_transaction=True, **field_values) -> int: """ Update models that match this query to the given field-value pairs. Keys and values given as keyword arguments are interpreted as fields on the target model and the values as the values to which to set the given fields. + + Returns: + The number of matching records updated. """ - validate_model_fields(self.model, field_values) - pipeline = await self.model.db().pipeline() if use_transaction else None - - # TODO: async for here? - for model in await self.all(): - for field, value in field_values.items(): - setattr(model, field, value) - # TODO: In the non-transaction case, can we do more to detect - # failure responses from Redis? - await model.save(pipeline=pipeline) - - if pipeline: - # TODO: Response type? - # TODO: Better error detection for transactions. - await pipeline.execute() + serialized_values = self._serialize_update_values(field_values) + keys = await self._search_keys_only() + if not keys: + return 0 + + pipeline = self.model.db().pipeline(transaction=use_transaction) + if issubclass(self.model, HashModel): + for key in keys: + pipeline.hset(key, mapping=serialized_values) + else: + for key in keys: + for field, value in serialized_values.items(): + path = "$." + field.replace("__", ".") + pipeline.json().set(key, path, value) + + await pipeline.execute() + return len(keys) async def delete(self): """Delete all matching records in this query.""" diff --git a/docs/querying.md b/docs/querying.md index 184b02c4..73c0d2a4 100644 --- a/docs/querying.md +++ b/docs/querying.md @@ -306,9 +306,10 @@ Update all matching records with new field values: ```python # Give everyone in the "premium" tier a discount -await Customer.find( +updated_count = await Customer.find( Customer.tier == "premium" ).update(discount_percent=20) +# `updated_count` is the number of matching records updated. ``` ### `.delete()` - Delete Multiple Records diff --git a/tests/test_find_query_update.py b/tests/test_find_query_update.py new file mode 100644 index 00000000..876ad3ee --- /dev/null +++ b/tests/test_find_query_update.py @@ -0,0 +1,164 @@ +from typing import Any, Dict, List + +import pytest +from pydantic import ValidationError + +from aredis_om import EmbeddedJsonModel, Field, HashModel, JsonModel +from aredis_om.model import model as model_module + +from .conftest import py_test_mark_asyncio + + +class FakePipeline: + def __init__(self): + self.hash_updates: List[tuple[str, Dict[str, Any]]] = [] + self.json_updates: List[tuple[str, str, Any]] = [] + self.execute_count = 0 + + def hset(self, key: str, mapping: Dict[str, Any]): + self.hash_updates.append((key, mapping)) + return self + + def json(self): + return self + + def set(self, key: str, path: str, value: Any): + self.json_updates.append((key, path, value)) + return self + + async def execute(self): + self.execute_count += 1 + return [] + + +class FakeDatabase: + def __init__(self, search_results: Dict[int, List[Any]]): + self.search_results = search_results + self.search_calls: List[tuple[Any, ...]] = [] + self.pipelines: List[FakePipeline] = [] + self.pipeline_transactions: List[bool] = [] + + async def execute_command(self, *args): + self.search_calls.append(args) + limit_index = args.index("LIMIT") + offset = args[limit_index + 1] + return self.search_results[offset] + + def pipeline(self, transaction: bool = True): + self.pipeline_transactions.append(transaction) + pipeline = FakePipeline() + self.pipelines.append(pipeline) + return pipeline + + +def enable_search(monkeypatch): + monkeypatch.setattr(model_module, "has_redisearch", lambda db: True) + + +def hash_model(database): + class User(HashModel, index=True): + status: str + age: int = Field(ge=0) + payload: str + + User.Meta.database = database + return User + + +def json_model(database): + class Address(EmbeddedJsonModel): + city: str + + class Document(JsonModel, index=True): + address: Address + status: str + payload: Dict[str, str] + + Document.Meta.database = database + return Document + + +@py_test_mark_asyncio +async def test_query_update_uses_key_only_search_and_partial_hset(monkeypatch): + enable_search(monkeypatch) + database = FakeDatabase({0: [2, "user:1", "user:2"]}) + User = hash_model(database) + + updated = await User.find().only("status").update(status="active", age="42") + + assert updated == 2 + assert database.search_calls[0][-1] == "NOCONTENT" + assert all("RETURN" not in call for call in database.search_calls) + assert database.pipeline_transactions == [True] + assert database.pipelines[0].hash_updates == [ + ("user:1", {"status": "active", "age": 42}), + ("user:2", {"status": "active", "age": 42}), + ] + assert database.pipelines[0].execute_count == 1 + + +@py_test_mark_asyncio +async def test_query_update_validates_values_before_search(monkeypatch): + enable_search(monkeypatch) + database = FakeDatabase({0: [1, "user:1"]}) + User = hash_model(database) + + with pytest.raises(ValidationError): + await User.find().update(age=-1) + + assert database.search_calls == [] + assert database.pipelines == [] + + +@py_test_mark_asyncio +async def test_query_update_paginates_key_only_results(monkeypatch): + enable_search(monkeypatch) + database = FakeDatabase( + { + 0: [3, "user:1"], + 1: [3, "user:2"], + 2: [3, "user:3"], + } + ) + User = hash_model(database) + + query = User.find() + query.page_size = 1 + query.limit = 1 + updated = await query.update(status="active") + + assert updated == 3 + assert [call[call.index("LIMIT") + 1] for call in database.search_calls] == [ + 0, + 1, + 2, + ] + assert len(database.pipelines[0].hash_updates) == 3 + + +@py_test_mark_asyncio +async def test_query_update_uses_json_paths_and_non_transactional_pipeline(monkeypatch): + enable_search(monkeypatch) + database = FakeDatabase({0: [1, "document:1"]}) + Document = json_model(database) + + updated = await Document.find().update(use_transaction=False, status="active") + + assert updated == 1 + assert database.pipeline_transactions == [False] + assert database.pipelines[0].json_updates == [("document:1", "$.status", "active")] + assert database.pipelines[0].execute_count == 1 + + +@py_test_mark_asyncio +async def test_query_update_uses_nested_json_path(monkeypatch): + enable_search(monkeypatch) + database = FakeDatabase({0: [1, "document:1"]}) + Document = json_model(database) + + updated = await Document.find().update(address__city="Seoul") + + assert updated == 1 + assert database.pipelines[0].json_updates == [ + ("document:1", "$.address.city", "Seoul") + ] From 408cc9f55653deb69c5a302dafd7f7a81d3628cf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=EC=97=84=EC=9C=A4=EC=84=AD?= <62834176+lh0156@users.noreply.github.com> Date: Mon, 27 Jul 2026 03:49:22 +0900 Subject: [PATCH 2/2] Keep query update test models isolated --- tests/test_find_query_update.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_find_query_update.py b/tests/test_find_query_update.py index 876ad3ee..1e433881 100644 --- a/tests/test_find_query_update.py +++ b/tests/test_find_query_update.py @@ -1,3 +1,4 @@ +import abc from typing import Any, Dict, List import pytest @@ -56,7 +57,7 @@ def enable_search(monkeypatch): def hash_model(database): - class User(HashModel, index=True): + class User(HashModel, abc.ABC, index=True): status: str age: int = Field(ge=0) payload: str @@ -69,7 +70,7 @@ def json_model(database): class Address(EmbeddedJsonModel): city: str - class Document(JsonModel, index=True): + class Document(JsonModel, abc.ABC, index=True): address: Address status: str payload: Dict[str, str]