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/6] 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/6] 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] From ec5be5c32ae7f35809210c61a668a8b766e47774 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: Fri, 31 Jul 2026 11:15:06 +0900 Subject: [PATCH 3/6] Fix v1 nested update paths --- aredis_om/model/model.py | 5 ++++- tests/test_find_query_update.py | 16 ++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py index 895c21c2..56c35ac1 100644 --- a/aredis_om/model/model.py +++ b/aredis_om/model/model.py @@ -2038,7 +2038,10 @@ def _get_update_field(self, field_name: str): candidate for candidate in (annotation, *get_args(annotation)) if isinstance(candidate, type) - and hasattr(candidate, "model_fields") + and ( + hasattr(candidate, "model_fields") + or hasattr(candidate, "__fields__") + ) ), None, ) diff --git a/tests/test_find_query_update.py b/tests/test_find_query_update.py index 1e433881..439274a3 100644 --- a/tests/test_find_query_update.py +++ b/tests/test_find_query_update.py @@ -1,4 +1,5 @@ import abc +from types import SimpleNamespace from typing import Any, Dict, List import pytest @@ -79,6 +80,21 @@ class Document(JsonModel, abc.ABC, index=True): return Document +def test_query_update_resolves_pydantic_v1_nested_models(): + class LegacyNestedModel: + __fields__ = {"city": object()} + + class RootModel: + model_fields = {"address": SimpleNamespace(annotation=LegacyNestedModel)} + + query = object.__new__(model_module.FindQuery) + query.model = RootModel + + assert ( + query._get_update_field("address__city") is LegacyNestedModel.__fields__["city"] + ) + + @py_test_mark_asyncio async def test_query_update_uses_key_only_search_and_partial_hset(monkeypatch): enable_search(monkeypatch) From d2a4022faf1855233af94ab066f1faba58c19e71 Mon Sep 17 00:00:00 2001 From: Andrew Brookins Date: Fri, 31 Jul 2026 17:24:12 -0700 Subject: [PATCH 4/6] Preserve query update field TTLs --- aredis_om/model/model.py | 29 +++++++++++++++++++++++++++-- tests/test_find_query_update.py | 20 ++++++++++++++++++++ tests/test_hash_field_expiration.py | 23 +++++++++++++++++++++++ 3 files changed, 70 insertions(+), 2 deletions(-) diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py index 56c35ac1..c9469515 100644 --- a/aredis_om/model/model.py +++ b/aredis_om/model/model.py @@ -2032,7 +2032,7 @@ def _get_update_field(self, field_name: str): if index == len(parts) - 1: return field - annotation = getattr(field, "annotation", None) + annotation = field.annotation if PYDANTIC_V2 else field.outer_type_ nested_model = next( ( candidate @@ -2061,7 +2061,8 @@ def _validated_update_values(self, field_values: 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) + annotation = field.annotation if PYDANTIC_V2 else field.outer_type_ + field_definitions[field_name] = (annotation, field_info) update_model = create_model( # type: ignore[call-overload] f"{self.model.__name__}Update", **field_definitions @@ -2122,6 +2123,25 @@ async def _search_keys_only(self) -> List[Union[str, bytes]]: return keys + async def _get_preserved_hash_field_ttls( + self, keys: List[Union[str, bytes]], field_names: List[str] + ) -> Dict[Union[str, bytes], Dict[str, int]]: + """Return positive TTLs for HashModel fields that will be updated.""" + if not field_names or not supports_hash_field_expiration(): + return {} + + ttl_pipeline = self.model.db().pipeline(transaction=False) + for key in keys: + ttl_pipeline.httl(key, *field_names) + + ttl_results = await ttl_pipeline.execute() + return { + key: { + field_name: ttl for field_name, ttl in zip(field_names, ttls) if ttl > 0 + } + for key, ttls in zip(keys, ttl_results) + } + async def update(self, use_transaction=True, **field_values) -> int: """ Update models that match this query to the given field-value pairs. @@ -2140,8 +2160,13 @@ async def update(self, use_transaction=True, **field_values) -> int: pipeline = self.model.db().pipeline(transaction=use_transaction) if issubclass(self.model, HashModel): + preserved_ttls = await self._get_preserved_hash_field_ttls( + keys, list(serialized_values) + ) for key in keys: pipeline.hset(key, mapping=serialized_values) + for field_name, ttl in preserved_ttls.get(key, {}).items(): + pipeline.hexpire(key, ttl, field_name) else: for key in keys: for field, value in serialized_values.items(): diff --git a/tests/test_find_query_update.py b/tests/test_find_query_update.py index 439274a3..accb69c0 100644 --- a/tests/test_find_query_update.py +++ b/tests/test_find_query_update.py @@ -55,6 +55,7 @@ def pipeline(self, transaction: bool = True): def enable_search(monkeypatch): monkeypatch.setattr(model_module, "has_redisearch", lambda db: True) + monkeypatch.setattr(model_module, "supports_hash_field_expiration", lambda: False) def hash_model(database): @@ -95,6 +96,25 @@ class RootModel: ) +def test_query_update_resolves_pydantic_v1_field_types(monkeypatch): + class LegacyNestedModel: + __fields__ = {"city": object()} + + class LegacyField: + outer_type_ = LegacyNestedModel + + class RootModel: + __fields__ = {"address": LegacyField()} + + monkeypatch.setattr(model_module, "PYDANTIC_V2", False) + query = object.__new__(model_module.FindQuery) + query.model = RootModel + + assert ( + query._get_update_field("address__city") is LegacyNestedModel.__fields__["city"] + ) + + @py_test_mark_asyncio async def test_query_update_uses_key_only_search_and_partial_hset(monkeypatch): enable_search(monkeypatch) diff --git a/tests/test_hash_field_expiration.py b/tests/test_hash_field_expiration.py index 300ebc9e..7579d118 100644 --- a/tests/test_hash_field_expiration.py +++ b/tests/test_hash_field_expiration.py @@ -280,6 +280,29 @@ async def test_update_preserves_field_expiration(models, redis): assert updated_ttl <= initial_ttl +@py_test_mark_asyncio +async def test_find_query_update_preserves_updated_field_expiration(models, redis): + """FindQuery.update() should preserve the TTL of an updated hash field.""" + session = models.Session( + user_id="user123", + token="abc123", + refresh_token="refresh456", + ) + await session.save() + + initial_ttl = await session.field_ttl("token") + assert initial_ttl > 0 + + updated_count = await models.Session.find(models.Session.pk == session.pk).update( + token="updated-token" + ) + + assert updated_count == 1 + updated_ttl = await session.field_ttl("token") + assert updated_ttl > 0 + assert updated_ttl <= initial_ttl + + @py_test_mark_asyncio async def test_save_preserves_manually_set_ttl(models, redis): """ From 256a3c879294c5485ca2e5e35fee645e715bb6c0 Mon Sep 17 00:00:00 2001 From: Andrew Brookins Date: Fri, 31 Jul 2026 17:25:58 -0700 Subject: [PATCH 5/6] Check query TTL support per connection --- aredis_om/model/model.py | 4 +++- tests/test_find_query_update.py | 10 +++++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py index 1767f3ae..10b6211d 100644 --- a/aredis_om/model/model.py +++ b/aredis_om/model/model.py @@ -2145,7 +2145,9 @@ async def _get_preserved_hash_field_ttls( self, keys: List[Union[str, bytes]], field_names: List[str] ) -> Dict[Union[str, bytes], Dict[str, int]]: """Return positive TTLs for HashModel fields that will be updated.""" - if not field_names or not supports_hash_field_expiration(): + if not field_names: + return {} + if not await supports_hash_field_expiration(self.model.db()): return {} ttl_pipeline = self.model.db().pipeline(transaction=False) diff --git a/tests/test_find_query_update.py b/tests/test_find_query_update.py index accb69c0..bb9813a9 100644 --- a/tests/test_find_query_update.py +++ b/tests/test_find_query_update.py @@ -55,7 +55,15 @@ def pipeline(self, transaction: bool = True): def enable_search(monkeypatch): monkeypatch.setattr(model_module, "has_redisearch", lambda db: True) - monkeypatch.setattr(model_module, "supports_hash_field_expiration", lambda: False) + + async def field_expiration_is_unsupported(_conn): + return False + + monkeypatch.setattr( + model_module, + "supports_hash_field_expiration", + field_expiration_is_unsupported, + ) def hash_model(database): From 4ca575fac56e5a1e7fb13c21f101632415bf51df 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: Sun, 2 Aug 2026 20:34:08 +0900 Subject: [PATCH 6/6] Skip empty hash query updates --- aredis_om/model/model.py | 3 +++ tests/test_find_query_update.py | 23 ++++++++++++++++++++++- 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/aredis_om/model/model.py b/aredis_om/model/model.py index 10b6211d..a8dc46a4 100644 --- a/aredis_om/model/model.py +++ b/aredis_om/model/model.py @@ -2174,6 +2174,9 @@ async def update(self, use_transaction=True, **field_values) -> int: The number of matching records updated. """ serialized_values = self._serialize_update_values(field_values) + if not serialized_values: + return 0 + keys = await self._search_keys_only() if not keys: return 0 diff --git a/tests/test_find_query_update.py b/tests/test_find_query_update.py index bb9813a9..77de1919 100644 --- a/tests/test_find_query_update.py +++ b/tests/test_find_query_update.py @@ -1,6 +1,6 @@ import abc from types import SimpleNamespace -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional import pytest from pydantic import ValidationError @@ -76,6 +76,14 @@ class User(HashModel, abc.ABC, index=True): return User +def nullable_hash_model(database): + class User(HashModel, abc.ABC, index=True): + nickname: Optional[str] = None + + User.Meta.database = database + return User + + def json_model(database): class Address(EmbeddedJsonModel): city: str @@ -155,6 +163,19 @@ async def test_query_update_validates_values_before_search(monkeypatch): assert database.pipelines == [] +@py_test_mark_asyncio +async def test_query_update_skips_empty_hash_mapping(monkeypatch): + enable_search(monkeypatch) + database = FakeDatabase({0: [1, "user:1"]}) + User = nullable_hash_model(database) + + updated = await User.find().update(nickname=None) + + assert updated == 0 + assert database.search_calls == [] + assert database.pipelines == [] + + @py_test_mark_asyncio async def test_query_update_paginates_key_only_results(monkeypatch): enable_search(monkeypatch)